Spaces:
Sleeping
Sleeping
| """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() | |