nothingworry's picture
imporve RAG
9d50a01
import httpx
from dataclasses import dataclass, field
@dataclass
class MCPClient:
rag_url: str
web_url: str
admin_url: str
client: httpx.AsyncClient = field(default_factory=lambda: httpx.AsyncClient(timeout=30))
async def call_rag(self, tenant_id: str, query: str, threshold: float = 0.3):
"""
Calls the RAG search endpoint and returns the unwrapped results.
The MCP server wraps responses in a 'data' field, so we extract it.
Uses a lower threshold (0.3) by default to ensure we find relevant results
even if semantic similarity is moderate.
"""
r = await self.client.post(
f"{self.rag_url}/search",
json={
"tenant_id": tenant_id,
"query": query,
"threshold": threshold # Lower threshold for better recall
}
)
if r.status_code != 200:
return {"results": [], "error": f"HTTP {r.status_code}"}
data = r.json()
# MCP server wraps response in a 'data' field
# Extract the actual result data
if isinstance(data, dict) and "data" in data:
return data["data"]
# If not wrapped, return as-is (backward compatibility)
return data
async def call_web(self, tenant_id: str, query: str):
r = await self.client.post(f"{self.web_url}/search", json={"tenant_id":tenant_id,"query":query})
return r.json()
async def call_admin(self, tenant_id: str, query: str):
r = await self.client.post(f"{self.admin_url}/eval", json={"tenant_id":tenant_id,"query":query})
return r.json()