Spaces:
Runtime error
Runtime error
File size: 5,093 Bytes
0a19322 | 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 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 | """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}
|