Spaces:
Sleeping
Sleeping
File size: 2,631 Bytes
c16e1c9 aa63765 c16e1c9 aa63765 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 |
import os
import httpx
from dotenv import load_dotenv
load_dotenv()
class RAGClient:
"""
Communicates with the RAG MCP server.
"""
def __init__(self):
self.base_url = os.getenv("RAG_MCP_URL")
self.search_endpoint = f"{self.base_url}/search"
self.ingest_endpoint = f"{self.base_url}/ingest"
async def search(self, query: str, tenant_id: str):
"""
Sends the query to the RAG server and returns document chunks.
"""
try:
async with httpx.AsyncClient() as client:
response = await client.post(
self.search_endpoint,
json={
"query": query,
"tenant_id": tenant_id
}
)
if response.status_code != 200:
return []
data = response.json()
return data.get("results", [])
except Exception as e:
print("RAG Client Error:", e)
return []
async def ingest(self, content: str, tenant_id: str):
"""
Sends content to the RAG server for ingestion.
"""
try:
async with httpx.AsyncClient() as client:
response = await client.post(
self.ingest_endpoint,
json={
"tenant_id": tenant_id,
"content": content
}
)
if response.status_code != 200:
return {"error": f"HTTP {response.status_code}"}
data = response.json()
return data
except Exception as e:
print("RAG Ingest Error:", e)
return {"error": str(e)}
async def list_documents(self, tenant_id: str, limit: int = 1000, offset: int = 0):
"""
List all documents for a tenant.
"""
try:
async with httpx.AsyncClient() as client:
response = await client.get(
f"{self.base_url}/list",
params={
"tenant_id": tenant_id,
"limit": limit,
"offset": offset
}
)
if response.status_code != 200:
return {"documents": [], "total": 0, "limit": limit, "offset": offset}
data = response.json()
return data
except Exception as e:
print("RAG List Error:", e)
return {"documents": [], "total": 0, "limit": limit, "offset": offset}
|