Spaces:
Runtime error
Runtime error
| """Orchestrates: understand question -> retrieve -> ask the LLM -> return answer + sources.""" | |
| import logging | |
| from . import llm, query, retrieval | |
| from .config import settings | |
| from .registry import get_registry | |
| log = logging.getLogger("rag.rag") | |
| SYSTEM_PROMPT = """You answer questions about the user's uploaded documents using ONLY the provided context. | |
| - Cite sources inline as [1], [2] using the numbers in the context. | |
| - If the context does not contain the answer, say clearly that you could not find it in the documents. Never guess. | |
| - If the user asks what is on a specific page, describe or summarise that page from the context. | |
| - If the context lists notes about missing pages, mention them briefly. | |
| - Reply in the same language the user wrote in (English, Hindi or Hinglish). | |
| - Be concise and well organised.""" | |
| class NoReadyDocuments(Exception): | |
| pass | |
| def _merge_pages(chunks: list[dict]) -> list[dict]: | |
| """For page lookups: one source per (document, page) instead of one per chunk.""" | |
| merged: dict[tuple, dict] = {} | |
| for c in chunks: | |
| key = (c["doc_id"], c["page"]) | |
| if key in merged: | |
| merged[key]["text"] += "\n\n" + c["text"] | |
| else: | |
| merged[key] = {**c} | |
| return list(merged.values()) | |
| def _page_notes(scope: list[dict], pages: list[int], found: set[tuple]) -> list[str]: | |
| notes = [] | |
| for d in scope: | |
| for p in pages: | |
| if (d["doc_id"], p) in found: | |
| continue | |
| if d.get("pages") and p > d["pages"]: | |
| notes.append(f'"{d["filename"]}" has only {d["pages"]} pages, so page {p} does not exist.') | |
| else: | |
| notes.append( | |
| f'Page {p} of "{d["filename"]}" has no extractable text (blank page or a scanned image).' | |
| ) | |
| return notes | |
| def _build_context(sources: list[dict]) -> str: | |
| parts, used = [], 0 | |
| for i, s in enumerate(sources): | |
| block = f"[{i + 1}] ({s['filename']}, page {s['page']})\n{s['text']}" | |
| room = settings.max_context_chars - used | |
| if room <= 200: | |
| break | |
| block = block[:room] | |
| parts.append(block) | |
| used += len(block) | |
| return "\n\n".join(parts) | |
| def _shape(sources: list[dict]) -> list[dict]: | |
| return [ | |
| { | |
| "doc_id": s["doc_id"], | |
| "filename": s["filename"], | |
| "page": s["page"], | |
| "score": round(float(s.get("score", 0.0)), 4), | |
| "text": s["text"], | |
| } | |
| for s in sources | |
| ] | |
| def answer_question(question: str, doc_ids: list[str] | None = None, history: list[dict] | None = None) -> dict: | |
| history = history or [] | |
| ready = get_registry().ready() | |
| if not ready: | |
| raise NoReadyDocuments() | |
| scope = [d for d in ready if d["doc_id"] in doc_ids] if doc_ids else ready | |
| if not scope: | |
| scope = ready | |
| plan = query.analyze(question, scope, history) | |
| explicit = len(scope) == 1 | |
| if plan.doc_ids: # the question names a document -> narrow to it | |
| scope = [d for d in scope if d["doc_id"] in plan.doc_ids] or scope | |
| explicit = True | |
| scope_ids = [d["doc_id"] for d in scope] | |
| notes: list[str] = [] | |
| note_for_user = "" | |
| if plan.intent == "page_lookup" and plan.pages: | |
| pages = plan.pages | |
| if len(pages) > settings.max_lookup_pages: | |
| note_for_user = f"\n\n(I read the first {settings.max_lookup_pages} of the {len(pages)} pages you asked about.)" | |
| pages = pages[: settings.max_lookup_pages] | |
| chunks = retrieval.fetch_pages(scope_ids, pages) | |
| sources = _merge_pages(chunks) | |
| found = {(s["doc_id"], s["page"]) for s in sources} | |
| notes = _page_notes(scope, pages, found) | |
| if sources and len(scope) > 1 and not explicit: | |
| notes = [] # don't nag about documents that simply don't have that page | |
| mode = "page_lookup" | |
| shown_pages = pages | |
| else: | |
| sources = retrieval.hybrid_search(plan.search_query, scope_ids, plan.pages or None) | |
| mode = "search" | |
| shown_pages = plan.pages | |
| meta = { | |
| "mode": mode, | |
| "search_query": plan.search_query if mode == "search" else None, | |
| "pages": shown_pages or None, | |
| "documents": [d["filename"] for d in scope], | |
| } | |
| if not sources: | |
| msg = " ".join(notes) if notes else "I couldn't find anything relevant to this in the selected documents." | |
| return {"answer": msg, "sources": [], **meta} | |
| context = _build_context(sources) | |
| note_block = ("\n\nNotes:\n- " + "\n- ".join(notes)) if notes else "" | |
| messages = [{"role": "system", "content": SYSTEM_PROMPT}] | |
| for h in history[-4:]: | |
| messages.append({"role": h["role"], "content": h["content"][:600]}) | |
| messages.append({"role": "user", "content": f"Context:\n{context}{note_block}\n\nQuestion: {question}"}) | |
| answer = llm.chat(messages) | |
| if not answer: | |
| answer = "The model returned an empty reply. Please try asking again." | |
| return {"answer": answer + note_for_user, "sources": _shape(sources), **meta} | |