apunkt/plano
1
0
Fork 0
mirror of https://github.com/katanemo/plano.git synced 2026-07-05 15:52:12 +02:00
Code Issues Projects Releases Packages Wiki Activity

more changes

Browse source
This commit is contained in:
Adil Hafeez 2025-11-28 11:34:43 -08:00
parent dcfc85ca74
commit 4140c1cde4
No known key found for this signature in database
GPG key ID: 9B18EF7691369645
15 changed files with 883 additions and 109 deletions
Download patch file Download diff file Expand all files Collapse all files

148
demos/use_cases/rag_agent/README.md
View file

@ -1,28 +1,148 @@
# RAG Agent Query Parser
# RAG Agent with MCP Protocol
A FastAPI service that rewrites user queries using archgw and gpt-4o-mini for better retrieval accuracy.
A multi-agent RAG system using the Model Context Protocol (MCP) for agent communication.
## How it Works
## Architecture
1. Receives a chat completion request with conversation history
2. Calls archgw's LLM gateway with gpt-4o-mini to rewrite the last user query
3. Returns the rewritten query as the assistant response
This demo consists of three MCP agents:
1. **Query Rewriter** - Rewrites user queries for better retrieval
2. **Context Builder** - Retrieves relevant context from knowledge base
3. **Response Generator** - Generates final responses with context
Each agent runs as an independent MCP server and exposes tools that can be called via the MCP protocol.
## MCP Tools
### Query Rewriter Agent
- **Tool**: `rewrite_query_with_archgw`
- **Description**: Rewrites user queries using LLM for better retrieval
- **Port**: 10500
### Context Builder Agent
- **Tool**: `chat_completions`
- **Description**: Augments queries with relevant context from knowledge base
- **Port**: 10501
### Response Generator Agent
- **Port**: 10502
## Setup and Running
1. **Start archgw**:
```bash
archgw up --foreground
```
### 1. Start archgw
```bash
archgw up --foreground
```
2. **Start the query parser service**:
```bash
uv run python -m rag_agent.query_parser
```
### 2. Start Individual Agents
**Query Rewriter:**
```bash
uv run python -m rag_agent \
--agent query_rewriter \
--host 0.0.0.0 \
--port 10500 \
--transport sse
```
**Context Builder:**
```bash
uv run python -m rag_agent \
--agent context_builder \
--host 0.0.0.0 \
--port 10501 \
--transport sse
```
**Response Generator:**
```bash
uv run python -m rag_agent \
--agent response_generator \
--host 0.0.0.0 \
--port 10502 \
--transport sse
```
### 3. Start All Agents at Once
```bash
./start_agents.sh
```
## Configuration
The `arch_config.yaml` defines how agents are connected:
```yaml
agent_filters:
- id: query_rewriter
url: mcp://host.docker.internal:10500
tool: rewrite_query_with_archgw # MCP tool name
- id: context_builder
url: mcp://host.docker.internal:10501
tool: chat_completions
```
### MCP Tool Invocation Patterns
The config supports different ways to specify MCP tools:
**1. Separate tool field (recommended):**
```yaml
- id: query_rewriter
url: mcp://host.docker.internal:10500
tool: rewrite_query_with_archgw
```
**2. Tool in URL path:**
```yaml
- id: query_rewriter
url: mcp://host.docker.internal:10500/rewrite_query_with_archgw
```
**3. Tool as query parameter:**
```yaml
- id: query_rewriter
url: mcp://host.docker.internal:10500?tool=rewrite_query_with_archgw
```
## CLI Options
```bash
uv run python -m rag_agent --help
Options:
--transport TEXT Transport type: stdio or sse (default: sse)
--host TEXT Host to bind MCP server to (default: localhost)
--port INTEGER Port for MCP server (default: 10500)
--agent TEXT Agent name: query_rewriter, context_builder, or response_generator (required)
--name TEXT Custom MCP server name (optional)
```
## Environment Variables
```bash
# archgw LLM Gateway base URL (default: http://localhost:12000/v1)
export LLM_GATEWAY_ENDPOINT="http://localhost:12000/v1"
# OpenAI API Key for model providers
export OPENAI_API_KEY="your-key-here"
```
## Testing
See `sample_queries.md` for example queries to test the RAG system.
Example request:
```bash
curl -X POST http://localhost:8001/v1/chat/completions \
-H "Content-Type: application/json" \
-d '{
"model": "gpt-4o",
"messages": [
{
"role": "user",
"content": "What is the guaranteed uptime for TechCorp?"
}
]
}'
```

31
demos/use_cases/rag_agent/arch_config.yaml
View file

@ -1,12 +1,24 @@
version: v0.3.0
agents:
- id: query_rewriter
url: http://host.docker.internal:10500/v1/chat/completions
- id: context_builder
url: http://host.docker.internal:10501/v1/chat/completions
- id: rag_agent
url: http://host.docker.internal:10502/v1/chat/completions
url: mcp://host.docker.internal:10501
# only sse is supported
# transport: sse or stdio
# optional tool name, defaults to "invoke"
# tool: invoke
- id: travel_agent
url: mcp://host.docker.internal:10502
agent_filters:
- id: query_rewriter
url: mcp://host.docker.internal:10500
# tool is optional, defaults to id
# tool: query_rewriter
- id: context_builder
url: mcp://host.docker.internal:10500
- id: input_guards
url: mcp://host.docker.internal:10500
model_providers:
- model: openai/gpt-4o-mini
@ -23,15 +35,20 @@ model_aliases:
listeners:
- type: agent
name: agent_1
port: 8001
router: arch_agent_router
agents:
- id: rag_agent
description: virtual assistant for device contracts for simple queries
description: virtual assistant for retrieval augmented generation tasks
filter_chain:
- input_guards
- query_rewriter
- context_builder
- id: travel_agent
description: virtual assistant for travel bookings and recommendations
filter_chain:
- input_guards
tracing:
random_sampling: 100

88
demos/use_cases/rag_agent/src/rag_agent/__init__.py
View file

@ -5,57 +5,45 @@ mcp = None
@click.command()
@click.option("--transport", "transport", default="stdio")
@click.option("--host", "host", default="localhost")
@click.option("--port", "port", default=10101)
@click.option("--agent", "agent", default=None)
@click.option(
"--rest-server",
"rest_server",
is_flag=True,
help="Start REST server instead of MCP server",
)
@click.option("--rest-port", "rest_port", default=8000, help="Port for REST server")
def main(host, port, agent, transport, rest_server, rest_port):
if rest_server:
print(f"Starting REST server on {host}:{rest_port} for agent: {agent}")
if agent == "query_parser":
from rag_agent.query_rewriter_agent import start_server
start_server(host=host, port=rest_port)
return
elif agent == "context_builder":
from rag_agent.context_builder_agent import (
start_server,
)
start_server(host=host, port=rest_port)
return
elif agent == "response_generator":
from rag_agent.response_generator_agent import start_server
start_server(host=host, port=rest_port)
return
else:
print("Please specify an agent to start with --agent option.")
return
print(f"Starting agent(s): {agent if agent else 'all'}")
@click.option("--transport", "transport", default="sse", help="Transport type: stdio or sse")
@click.option("--host", "host", default="localhost", help="Host to bind MCP server to")
@click.option("--port", "port", type=int, default=10500, help="Port for MCP server")
@click.option("--agent", "agent", required=True, help="Agent name: query_rewriter, context_builder, or response_generator")
@click.option("--name", "agent_name", default=None, help="Custom MCP server name (defaults to agent type)")
def main(host, port, agent, transport, agent_name):
"""Start a RAG agent as an MCP server."""
# Map friendly names to agent modules
agent_map = {
"query_rewriter": ("rag_agent.query_rewriter", "Query Rewriter Agent"),
"context_builder": ("rag_agent.context_builder_agent", "Context Builder Agent"),
"response_generator": ("rag_agent.response_generator", "Response Generator Agent"),
}
if agent not in agent_map:
print(f"Error: Unknown agent '{agent}'")
print(f"Available agents: {', '.join(agent_map.keys())}")
return
module_name, default_name = agent_map[agent]
mcp_name = agent_name or default_name
print(f"Starting MCP server: {mcp_name}")
print(f" Agent: {agent}")
print(f" Transport: {transport}")
print(f" Host: {host}")
print(f" Port: {port}")
global mcp
mcp = FastMCP("RAG Agent Demo", host=host, port=port)
if agent == "query_parser":
import rag_agent.query_parser
elif agent == "document_store":
import rag_agent.document_store
elif agent == "response_generator":
import rag_agent.response_generator
else:
import rag_agent.query_parser
import rag_agent.document_store
import rag_agent.response_generator
print("All agents loaded.")
mcp = FastMCP(mcp_name, host=host, port=port)
# Import the agent module to register its tools
import importlib
importlib.import_module(module_name)
print(f"Agent '{agent}' loaded successfully")
print(f"MCP server ready on {transport}://{host}:{port}")
mcp.run(transport=transport)

17
demos/use_cases/rag_agent/src/rag_agent/context_builder_agent.py → demos/use_cases/rag_agent/src/rag_agent/context_builder.py
View file

@ -10,7 +10,8 @@ from pathlib import Path
import uvicorn
from .api import ChatMessage, ChatCompletionRequest, ChatCompletionResponse
from . import mcp
from fastmcp.server.dependencies import get_http_headers
# Set up logging
logging.basicConfig(
@ -190,12 +191,12 @@ class Response(BaseModel):
# FastAPI app for REST server
app = FastAPI(title="RAG Content Builder Agent", version="1.0.0")
@mcp.tool()
@app.post("/v1/chat/completions")
async def chat_completions(
request_body: ChatCompletionRequest, request: Request
async def context_builder(
request_body: ChatCompletionRequest
) -> ChatCompletionResponse:
"""Chat completions endpoint that augments user queries with relevant context from the knowledge base."""
""" chat completions endpoint that augments user queries with relevant context from the knowledge base."""
import time
import uuid
@ -203,8 +204,10 @@ async def chat_completions(
f"Received chat completion request with {len(request_body.messages)} messages"
)
# Read traceparent header if present
traceparent_header = request.headers.get("traceparent")
# Get traceparent header from HTTP request using FastMCP's dependency function
headers = get_http_headers()
traceparent_header = headers.get("traceparent")
if traceparent_header:
logger.info(f"Received traceparent header: {traceparent_header}")
else:

15
demos/use_cases/rag_agent/src/rag_agent/query_rewriter_agent.py → demos/use_cases/rag_agent/src/rag_agent/query_rewriter.py
View file

@ -8,7 +8,8 @@ import logging
import uvicorn
from .api import ChatMessage, ChatCompletionRequest, ChatCompletionResponse
from . import mcp
from fastmcp.server.dependencies import get_http_headers
# Set up logging
logging.basicConfig(
@ -28,11 +29,10 @@ archgw_client = AsyncOpenAI(
api_key="EMPTY", # archgw doesn't require a real API key
)
async def rewrite_query_with_archgw(
messages: List[ChatMessage], traceparent_header: str
) -> str:
# Prepare the system prompt for query rewriting
""" Rewrite the user query using LLM for better retrieval. """
system_prompt = """You are a query rewriter that improves user queries for better retrieval.
Given a conversation history, rewrite the last user message to be more specific and context-aware.
@ -90,7 +90,8 @@ app = FastAPI(title="RAG Agent Query Parser", version="1.0.0")
@app.post("/v1/chat/completions")
async def chat_completions(request_body: ChatCompletionRequest, request: Request):
@mcp.tool()
async def query_rewriter(request_body: ChatCompletionRequest):
"""Chat completions endpoint that rewrites the last user query using archgw."""
import time
import uuid
@ -99,8 +100,10 @@ async def chat_completions(request_body: ChatCompletionRequest, request: Request
f"Received chat completion request with {len(request_body.messages)} messages"
)
# Read traceparent header if present
traceparent_header = request.headers.get("traceparent")
# Get traceparent header from HTTP request using FastMCP's dependency function
headers = get_http_headers()
traceparent_header = headers.get("traceparent")
if traceparent_header:
logger.info(f"Received traceparent header: {traceparent_header}")
else:

12
demos/use_cases/rag_agent/src/rag_agent/response_generator_agent.py → demos/use_cases/rag_agent/src/rag_agent/rag_agent.py
View file

@ -15,6 +15,9 @@ from .api import (
ChatCompletionStreamResponse,
)
from . import mcp
from fastmcp.server.dependencies import get_http_headers
# Set up logging
logging.basicConfig(
level=logging.INFO,
@ -60,14 +63,17 @@ def prepare_response_messages(request_body: ChatCompletionRequest):
@app.post("/v1/chat/completions")
async def chat_completions(request_body: ChatCompletionRequest, request: Request):
@mcp.tool(name="invoke")
async def chat_completion(request_body: ChatCompletionRequest):
"""Chat completions endpoint that generates a coherent response based on all context."""
logger.info(
f"Received chat completion request with {len(request_body.messages)} messages"
)
# Read traceparent header if present
traceparent_header = request.headers.get("traceparent")
# Get traceparent header from HTTP request using FastMCP's dependency function
headers = get_http_headers()
traceparent_header = headers.get("traceparent")
if traceparent_header:
logger.info(f"Received traceparent header: {traceparent_header}")
else: