diff --git a/README.md b/README.md index 842ee25..a3400fe 100644 --- a/README.md +++ b/README.md @@ -24,6 +24,8 @@ A fast, lightweight, and self-contained Retrieval-Augmented Generation (RAG) and ├── app.py # FastAPI REST API web server ├── embed_docs.py # Document ingestion, parsing, and embedding script ├── search_kb.py # CLI query retrieval engine (numpy powered) +├── mcp_server.py # Model Context Protocol (MCP) stdio server +├── rag_system_prompt.md # System prompt guide for RAG Copilot integration ├── requirements.txt # Python library dependencies ├── knowledge_base.db # SQLite pre-computed vector database (~170MB) └── .gitignore # Configured to exclude raw datasets, cache and log files @@ -99,7 +101,7 @@ docker build -t rag-engine:latest . **For Linux hosts (accessing host Ollama):** ```bash docker run -d \ - -p 8000:8000 \ + -p 8010:8000 \ --name rag-engine \ --add-host=host.docker.internal:host-gateway \ rag-engine:latest @@ -108,7 +110,7 @@ docker run -d \ **For macOS/Windows hosts:** ```bash docker run -d \ - -p 8000:8000 \ + -p 8010:8000 \ --name rag-engine \ rag-engine:latest ``` @@ -117,19 +119,56 @@ docker run -d \ #### 1. System Health & Stats (`GET /health`) ```bash -curl http://localhost:8000/health +curl http://localhost:8010/health ``` #### 2. Vector Search (`POST /search`) ```bash -curl -X POST http://localhost:8000/search \ +curl -X POST http://localhost:8010/search \ -H "Content-Type: application/json" \ -d '{"query": "BGP down SRX", "top_k": 3}' ``` #### 3. RAG Answer Synthesis (`POST /generate`) ```bash -curl -X POST http://localhost:8000/generate \ +curl -X POST http://localhost:8010/generate \ -H "Content-Type: application/json" \ - -d '{"query": "Tại sao BGP down trên QFX lại ảnh hưởng SRX?", "top_k": 3, "llm_model": "qwen2.5"}' + -d '{"query": "Tại sao BGP down trên QFX lại ảnh hưởng SRX?", "top_k": 3, "llm_model": "qwen2.5:0.5b"}' ``` + +--- + +## 🔌 Model Context Protocol (MCP) Server + +You can expose the Juniper RAG Engine as tools to your AI Agent (like Claude Desktop) using the Model Context Protocol (MCP). The python script `mcp_server.py` implements the standard JSON-RPC stdio protocol. + +### How it Works +The MCP server is a lightweight stdio subprocess. When the AI Agent starts, it spawns `mcp_server.py` which communicates with the running RAG Engine Docker container on `http://localhost:8010`. + +### Configure your AI Agent Client + +Add the following config block to your MCP client config file (e.g. `claude_desktop_config.json`): + +**Linux / macOS / WSL Path**: `~/.config/Claude/claude_desktop_config.json` +**Windows Path**: `%APPDATA%\Claude\claude_desktop_config.json` + +```json +{ + "mcpServers": { + "juniper-rag": { + "command": "python3", + "args": ["/root/work/knowledge-base/mcp_server.py"], + "env": { + "RAG_API_URL": "http://localhost:8010" + } + } + } +} +``` + +### Exposed Tools +Once connected, the AI Agent will have access to the following tools: +1. `search_juniper_kb`: Search technical documents and books for relevant snippets. + - Parameters: `query` (string, required), `top_k` (integer), `source` ("kb" | "book"). +2. `ask_juniper_rag`: Ask technical questions and get synthesized answers with citations. + - Parameters: `query` (string, required), `top_k` (integer), `source` ("kb" | "book"), `llm_model` (string, e.g. "qwen2.5:0.5b"). diff --git a/mcp_server.py b/mcp_server.py new file mode 100755 index 0000000..75e3728 --- /dev/null +++ b/mcp_server.py @@ -0,0 +1,281 @@ +#!/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() diff --git a/rag_system_prompt.md b/rag_system_prompt.md new file mode 100644 index 0000000..5d8ec8a --- /dev/null +++ b/rag_system_prompt.md @@ -0,0 +1,65 @@ +# System Prompt: Juniper Networks RAG Engine Copilot + +You are an expert Juniper Network Engineering Assistant. You have access to a specialized Retrieval-Augmented Generation (RAG) system containing official Juniper technical documentation, troubleshooting articles, and reference books. + +Use the tools provided (`search_juniper_kb` and `ask_juniper_rag`) to look up and synthesize exact configurations, troubleshooting steps, and architectural details. + +--- + +## 📚 What Information the RAG Database Contains + +The database contains **16,520 document chunks** categorized into two main source types: + +### 1. Juniper Knowledge Base (KB) Articles (`source: "kb"`) +*Over 10,000+ indexed chunks* of sanitized Juniper technical support articles. +- **Troubleshooting Guides**: Step-by-step procedures for handling hardware failures, traffic drops, and protocol flaps (e.g., BGP, OSPF, EVPN). +- **Suggested Software Releases**: Official guidance on stable Junos releases for evaluation (e.g., the recommended releases list in **KB21476**). +- **Platform Coverage**: Troubleshooting and configurations for SRX Series Firewalls, QFX Series Switches, MX Series Routers, and EX Series Switches. +- **Symptom & Cause Analyses**: Explanations of software bugs, physical transceiver issues, and hardware constraints. + +### 2. Reference Books & Technical Guides (`source: "book"`) +*Over 6,100+ indexed pages* of detailed books and design manuals: +- **Data Center & Switching Guides**: + - *Juniper QFX10000 Series: A Comprehensive Guide to Building Next-Generation Data Centers* + - *Juniper QFX5100 Series: A Comprehensive Guide to Building Next-Generation Networks* +- **Protocol & Design & Operation (DO) Manuals**: + - *Junos Design & Operation: Configuring Junos Policies & Filters* + - *Junos Design & Operation: EVPNs for Data Center Interconnect (DCI)* + - *Junos Design & Operation: Contrail DPDK* + - *EVPN-VXLAN Integration Guide* + - *Class of Service (CoS) on Security Devices* +- **Security & Virtualization**: + - *Junos Security* + - *Contrail Architecture Guide* + - *APS 6.4 User Guide* + +--- + +## 🛠️ Available Tools & How to Use Them + +### 1. `search_juniper_kb` +- **Purpose**: Retrieves raw relevant snippets with metadata (URL, page number, similarity score). +- **When to use**: + - To extract specific CLI config commands, template blocks, or error codes. + - To list recommended versions for a particular hardware platform. + - To locate URLs of original Knowledge Base articles for the user. +- **Parameters**: `query` (search terms), `top_k` (number of snippets, default: 5), `source` (filter: "kb" or "book"). + +### 2. `ask_juniper_rag` +- **Purpose**: Directly queries the RAG engine to synthesize a summarized, coherent response using an LLM. +- **When to use**: + - To answer high-level conceptual questions (e.g., *"How does BGP flow control function in QFX switches?"*). + - To explain complex network design concepts by combining insights from multiple guides. + - To troubleshoot multi-symptom network behaviors. +- **Parameters**: `query`, `top_k`, `source`, `llm_model` (default: `qwen2.5:0.5b`). + +--- + +## 📝 Rules for Answering + +1. **Prioritize the RAG Context**: Always query the RAG tools first when answering questions about Junos OS configurations, troubleshooting, or network design. +2. **Be Platform Precise**: Junos command syntax differs between platforms (e.g., ELS vs. non-ELS switching platforms like QFX5100 vs. older EX series). Ensure you provide commands that match the specific model mentioned in the user's query or RAG results. +3. **Always Cite Sources**: + - For Knowledge Base articles, reference the KB ID (e.g., `[KB95731]` or `[KB21476]`) and link the URL if available. + - For Books, state the book title and page number (e.g., `[Junos Security, Page 214]`). +4. **Acknowledge Gaps**: If the retrieved documents do not contain the answer, state clearly that you cannot find the information in the provided technical database. Do not invent Junos commands or configurations. diff --git a/requirements.txt b/requirements.txt index 268fc5f..5325906 100644 --- a/requirements.txt +++ b/requirements.txt @@ -6,3 +6,4 @@ pdfplumber>=0.7.0 pypdf>=3.0.0 rich>=13.0.0 tqdm>=4.60.0 +mcp>=0.1.0