Spaces:
Sleeping
Sleeping
File size: 2,285 Bytes
46e770f | 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 | """RAG Engine: Retrieve from ChromaDB + Generate via Agentic LLM."""
import logging
from typing import Dict, Any, List
from datetime import datetime
from app.core.database import query_documents, save_chat
from app.agent_workflow import agent
logger = logging.getLogger("deltamind.rag")
class RAGEngine:
def __init__(self):
self.counter = 0
def _session_id(self) -> str:
self.counter += 1
return f"s_{datetime.now().strftime('%Y%m%d%H%M%S')}_{self.counter}"
def retrieve(self, query: str, n: int = 4) -> Dict:
results = query_documents(query, n)
docs = results.get("documents", [[]])[0]
metas = results.get("metadatas", [[]])[0]
dists = results.get("distances", [[]])[0]
parts, citations = [], []
for i, (doc, meta, dist) in enumerate(zip(docs, metas, dists)):
rel = max(0, round(1 - dist, 2))
src = meta.get("source", "unknown")
parts.append(f"[Source {i+1}: {src} | Relevance: {rel}]\n{doc}")
citations.append({"index": i+1, "source": src, "relevance": rel})
return {"context": "\n\n---\n\n".join(parts), "citations": citations, "count": len(docs)}
async def chat(self, query: str, session_id: str = None, history: List[Dict] = None) -> Dict:
sid = session_id or self._session_id()
# 1. Retrieve
ret = self.retrieve(query)
# 2. Agent Generation (with potential tool calling)
result = await agent.run(query, context=ret["context"], history=history)
resp = {
"session_id": sid, "query": query, "response": result.get("content",""),
"provider": result.get("provider",""), "model": result.get("model",""),
"route": result.get("route",""), "elapsed": result.get("elapsed",0),
"citations": ret["citations"], "docs_retrieved": ret["count"],
"timestamp": datetime.now().isoformat()
}
# 3. Save to history
save_chat(sid, "user", query, result.get("provider",""))
save_chat(sid, "assistant", result.get("content",""), result.get("provider",""))
return resp
rag_engine = RAGEngine()
|