#!/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()