282 lines
10 KiB
Python
Executable File
282 lines
10 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
import sys
|
|
import json
|
|
import os
|
|
import urllib.request
|
|
import urllib.error
|
|
|
|
# Configuration
|
|
RAG_API_URL = os.environ.get("RAG_API_URL", "http://localhost:8010")
|
|
|
|
def log(msg):
|
|
"""Logs a message to stderr (since stdout is reserved for JSON-RPC messages)."""
|
|
sys.stderr.write(f"[MCP-Server] {msg}\n")
|
|
sys.stderr.flush()
|
|
|
|
def make_api_request(endpoint, data):
|
|
"""Sends a POST request to the RAG Engine FastAPI server using built-in urllib."""
|
|
url = f"{RAG_API_URL.rstrip('/')}/{endpoint.lstrip('/')}"
|
|
headers = {"Content-Type": "application/json"}
|
|
req_body = json.dumps(data).encode("utf-8")
|
|
|
|
log(f"Calling endpoint: {url} with parameters: {data}")
|
|
|
|
try:
|
|
req = urllib.request.Request(url, data=req_body, headers=headers, method="POST")
|
|
with urllib.request.urlopen(req, timeout=90) as response:
|
|
res_body = response.read().decode("utf-8")
|
|
return json.loads(res_body), None
|
|
except urllib.error.URLError as e:
|
|
err_msg = f"Failed to connect to RAG Engine API at {url}. Error: {e.reason if hasattr(e, 'reason') else e}"
|
|
log(f"API Error: {err_msg}")
|
|
return None, err_msg
|
|
except Exception as e:
|
|
err_msg = f"Unexpected error connecting to RAG Engine API: {str(e)}"
|
|
log(f"API Error: {err_msg}")
|
|
return None, err_msg
|
|
|
|
def handle_initialize(request_id):
|
|
"""Handles the initialize request from the MCP client."""
|
|
response = {
|
|
"jsonrpc": "2.0",
|
|
"id": request_id,
|
|
"result": {
|
|
"protocolVersion": "2024-11-05",
|
|
"capabilities": {
|
|
"tools": {}
|
|
},
|
|
"serverInfo": {
|
|
"name": "juniper-rag",
|
|
"version": "1.0.0"
|
|
}
|
|
}
|
|
}
|
|
return response
|
|
|
|
def handle_tools_list(request_id):
|
|
"""Returns the list of available tools to the MCP client."""
|
|
tools = [
|
|
{
|
|
"name": "search_juniper_kb",
|
|
"description": "Semantic search in Juniper Technical Documents database for relevant information. Returns a list of matching chunks with similarity scores, source paths/URLs, and page numbers.",
|
|
"inputSchema": {
|
|
"type": "object",
|
|
"properties": {
|
|
"query": {
|
|
"type": "string",
|
|
"description": "The search query (e.g. 'BGP down on SRX')"
|
|
},
|
|
"top_k": {
|
|
"type": "integer",
|
|
"description": "Number of matching chunks to retrieve (default: 5)",
|
|
"default": 5
|
|
},
|
|
"source": {
|
|
"type": "string",
|
|
"description": "Filter by source type: 'kb' (Juniper KB articles) or 'book' (Networking books/PDFs)",
|
|
"enum": ["kb", "book"]
|
|
}
|
|
},
|
|
"required": ["query"]
|
|
}
|
|
},
|
|
{
|
|
"name": "ask_juniper_rag",
|
|
"description": "Ask a question about Juniper network setups. Retrieves relevant documentation and uses Ollama to synthesize a detailed answer with citations.",
|
|
"inputSchema": {
|
|
"type": "object",
|
|
"properties": {
|
|
"query": {
|
|
"type": "string",
|
|
"description": "The question to answer"
|
|
},
|
|
"top_k": {
|
|
"type": "integer",
|
|
"description": "Number of chunks to retrieve for LLM context (default: 5)",
|
|
"default": 5
|
|
},
|
|
"source": {
|
|
"type": "string",
|
|
"description": "Filter by source type: 'kb' (Juniper KB articles) or 'book' (Networking books/PDFs)",
|
|
"enum": ["kb", "book"]
|
|
},
|
|
"llm_model": {
|
|
"type": "string",
|
|
"description": "Ollama LLM model name to use for synthesis (default: 'qwen2.5' or fallback model)",
|
|
"default": "qwen2.5"
|
|
}
|
|
},
|
|
"required": ["query"]
|
|
}
|
|
}
|
|
]
|
|
|
|
response = {
|
|
"jsonrpc": "2.0",
|
|
"id": request_id,
|
|
"result": {
|
|
"tools": tools
|
|
}
|
|
}
|
|
return response
|
|
|
|
def handle_tools_call(request_id, params):
|
|
"""Executes a tool and returns the response."""
|
|
tool_name = params.get("name")
|
|
arguments = params.get("arguments", {})
|
|
|
|
if tool_name == "search_juniper_kb":
|
|
# Prepare parameters for FastAPI /search endpoint
|
|
query = arguments.get("query")
|
|
top_k = arguments.get("top_k", 5)
|
|
source = arguments.get("source")
|
|
|
|
payload = {"query": query, "top_k": top_k}
|
|
if source:
|
|
payload["source"] = source
|
|
|
|
data, err = make_api_request("search", payload)
|
|
if err:
|
|
text_content = f"Error performing search: {err}\nMake sure RAG Engine is running."
|
|
else:
|
|
# Format results in a nice text block
|
|
blocks = []
|
|
for idx, item in enumerate(data, 1):
|
|
block = f"[{idx}] Title: {item.get('title')}\n"
|
|
block += f" Source Type: {item.get('source_type').upper()}\n"
|
|
if item.get("page_num"):
|
|
block += f" Page: {item.get('page_num')}\n"
|
|
if item.get("url"):
|
|
block += f" URL: {item.get('url')}\n"
|
|
block += f" Similarity Score: {item.get('score'):.4f}\n"
|
|
block += f" Content:\n{item.get('text')}\n"
|
|
blocks.append(block)
|
|
|
|
if not blocks:
|
|
text_content = "No matching document chunks found."
|
|
else:
|
|
text_content = "\n---\n".join(blocks)
|
|
|
|
return {
|
|
"jsonrpc": "2.0",
|
|
"id": request_id,
|
|
"result": {
|
|
"content": [
|
|
{"type": "text", "text": text_content}
|
|
]
|
|
}
|
|
}
|
|
|
|
elif tool_name == "ask_juniper_rag":
|
|
# Prepare parameters for FastAPI /generate endpoint
|
|
query = arguments.get("query")
|
|
top_k = arguments.get("top_k", 5)
|
|
source = arguments.get("source")
|
|
llm_model = arguments.get("llm_model", "qwen2.5")
|
|
|
|
payload = {"query": query, "top_k": top_k, "llm_model": llm_model}
|
|
if source:
|
|
payload["source"] = source
|
|
|
|
data, err = make_api_request("generate", payload)
|
|
if err:
|
|
text_content = f"Error generating answer: {err}\nMake sure RAG Engine and Ollama are running."
|
|
else:
|
|
answer = data.get("answer", "")
|
|
sources = data.get("sources", [])
|
|
|
|
text_content = f"### Answer:\n{answer}\n\n"
|
|
text_content += "### Sources:\n"
|
|
|
|
for idx, item in enumerate(sources, 1):
|
|
source_str = f"- [{idx}] {item.get('title')} ({item.get('source_type').upper()})"
|
|
if item.get("page_num"):
|
|
source_str += f" - Page {item.get('page_num')}"
|
|
if item.get("url"):
|
|
source_str += f" - {item.get('url')}"
|
|
text_content += f"{source_str}\n"
|
|
|
|
return {
|
|
"jsonrpc": "2.0",
|
|
"id": request_id,
|
|
"result": {
|
|
"content": [
|
|
{"type": "text", "text": text_content}
|
|
]
|
|
}
|
|
}
|
|
|
|
else:
|
|
return {
|
|
"jsonrpc": "2.0",
|
|
"id": request_id,
|
|
"error": {
|
|
"code": -32601,
|
|
"message": f"Method not found: {tool_name}"
|
|
}
|
|
}
|
|
|
|
def main():
|
|
"""Main input loop reading from stdin and writing to stdout."""
|
|
log("Starting Juniper RAG MCP Stdio Server...")
|
|
log(f"Configured RAG Engine API URL: {RAG_API_URL}")
|
|
|
|
while True:
|
|
try:
|
|
line = sys.stdin.readline()
|
|
if not line:
|
|
break
|
|
|
|
line = line.strip()
|
|
if not line:
|
|
continue
|
|
|
|
log(f"Received raw line: {line[:200]}...")
|
|
request = json.loads(line)
|
|
method = request.get("method")
|
|
request_id = request.get("id")
|
|
|
|
# MCP initialization workflow and routing
|
|
if method == "initialize":
|
|
response = handle_initialize(request_id)
|
|
elif method == "notifications/initialized":
|
|
# client notification, no response required
|
|
continue
|
|
elif method == "tools/list":
|
|
response = handle_tools_list(request_id)
|
|
elif method == "tools/call":
|
|
params = request.get("params", {})
|
|
response = handle_tools_call(request_id, params)
|
|
elif method == "ping":
|
|
response = {
|
|
"jsonrpc": "2.0",
|
|
"id": request_id,
|
|
"result": {}
|
|
}
|
|
else:
|
|
# Unknown method or notification
|
|
if request_id is not None:
|
|
response = {
|
|
"jsonrpc": "2.0",
|
|
"id": request_id,
|
|
"error": {
|
|
"code": -32601,
|
|
"message": f"Method not found: {method}"
|
|
}
|
|
}
|
|
else:
|
|
continue
|
|
|
|
# Send response back to stdout
|
|
sys.stdout.write(json.dumps(response) + "\n")
|
|
sys.stdout.flush()
|
|
log(f"Sent response for request id: {request_id}")
|
|
|
|
except json.JSONDecodeError:
|
|
log("Received invalid JSON data.")
|
|
except Exception as e:
|
|
log(f"Error in main loop: {str(e)}")
|
|
|
|
if __name__ == "__main__":
|
|
main()
|