"""LangGraph ingestion, evidence-checked Q&A, and independent reports.""" from typing import Any, TypedDict from langgraph.graph import END, START, StateGraph from .ingest import read_document from .retrieval import HybridIndex, evidence_for from .schema import Claim, Draft, Result, Rewrite, Summary from .config import FULL_CONTEXT_CHAR_LIMIT class GenerateState(TypedDict, total=False): task: str evidence: list summary: bool attempts: int errors: list[str] claims: list result: Any class QuestionState(TypedDict, total=False): question: str doc_ids: list[str] per_document: bool docs: list standalone: str evidence: list result: Any class IngestState(TypedDict, total=False): name: str data: bytes document: Any class ReportState(TypedDict, total=False): doc_id: str document: Any evidence: list fields: dict notes: list[str] summary: Any report: dict def build_generation_graph(engine, rules): def draft(state): if not state["evidence"]: return {"result": Result(status="abstained", notes=["No usable source evidence."])} instruction = rules + ("\nWrite a neutral 1–6 sentence summary of the complete document. " "One sentence per claim. Cover identity, key facts, reasoning and operative outcome " "where stated. Keep separate orders separate; do not pad with invented facts." if state["summary"] else "\nAnswer the question directly in 1–3 short claims. If evidence does not establish an " "answer, return answerable=false and claims=[]. Do not guess.") response = engine.models.json(Summary if state["summary"] else Draft, instruction, {"task": state["task"], "evidence": engine.payload(state["evidence"]), "corrections_needed": state.get("errors", [])}, max_tokens=1800 if state["summary"] else 1200) if not state["summary"] and not response.answerable: return {"result": Result(status="abstained", notes=["The supplied evidence does not establish an answer."])} sources = {e.id: e for e in state["evidence"]} update = {"attempts": state.get("attempts", 0) + 1, "result": None, "errors": []} if any(c.source_id not in sources for c in response.claims): update.update(claims=[], errors=["Choose source IDs from the supplied evidence."]) else: update["claims"] = [Claim(text=c.text, source_id=c.source_id, quote=sources[c.source_id].text) for c in response.claims] return update def verify(state): return {"errors": state["errors"] or engine.verify(state["claims"], state["evidence"], state["task"])} def finish(state): if state["errors"]: return {"result": Result(status="abstained", notes=state["errors"])} used = {c.source_id for c in state["claims"]} return {"result": Result(status="checked", claims=state["claims"], sources=[e for e in state["evidence"] if e.id in used], displayed=[c.text for c in state["claims"]])} graph = StateGraph(GenerateState) graph.add_node("draft", draft) graph.add_node("verify", verify) graph.add_node("finish", finish) graph.add_edge(START, "draft") graph.add_conditional_edges("draft", lambda s: "done" if s.get("result") is not None else "verify", {"done": END, "verify": "verify"}) graph.add_conditional_edges("verify", lambda s: "retry" if s["errors"] and s["attempts"] < 2 else "finish", {"retry": "draft", "finish": "finish"}) graph.add_edge("finish", END) return graph.compile(name="evidence_checked_generation") def build_question_graph(engine): import re def conversation(state): # Match the entire message: a greeting attached to a document question # must still go through retrieval and verification. text = re.sub(r"[^\w\s]", " ", state["question"].casefold()) text = " ".join(text.split()) greeting = r"(?:hi|hello|hey|hii|hiya|namaste|good morning|good afternoon|good evening)" wellbeing = r"(?:how are (?:you|u)(?: doing)?|how (?:r|re) (?:you|u)|how s it going)" thanks = r"(?:thanks(?: a lot| so much)?|thank (?:you|u)(?: so much)?|thx)" help_request = r"(?:help|help me|can you help(?: me)?|what can you do|who are you)" goodbye = r"(?:bye|goodbye|see you|good night)" acknowledgement = r"(?:ok|okay|great|cool|you re welcome|i am (?:fine|good)|i m (?:fine|good))" phrase = f"(?:{greeting}|{wellbeing}|{thanks}|{help_request}|{goodbye}|{acknowledgement})" if not re.fullmatch(rf"{phrase}(?:\s+{phrase})*", text): return {"result": None} if re.search(help_request, text): reply = "I'm LexQuery. I can answer questions about your PDFs and summarize them." reply += " Upload a PDF, then ask a question or open the Summaries tab." elif re.search(wellbeing, text): reply = "I'm here and ready to help! How can I help you with your documents?" elif re.search(thanks, text): reply = "You're welcome! Let me know if you have another question." elif re.search(goodbye, text): reply = "Goodbye! Come back whenever you need help with a document." elif re.search(greeting, text): reply = "Hi! What would you like to know about your documents?" if state["doc_ids"] else ( "Hi! Upload a PDF when you're ready. I can answer questions about it or give you a summary.") else: reply = "Sure! Let me know what you'd like to know about your documents." return {"result": Result(status="conversational", displayed=[reply])} def scope(state): if not state["question"].strip() or not state["doc_ids"]: return {"result": Result(status="abstained", notes=["Upload and select a PDF, then ask a question about it."])} current = tuple(sorted(state["doc_ids"])) if current != engine.chat_scope: engine.history.clear() engine.chat_scope = current docs = [engine.documents[i] for i in state["doc_ids"]] if any(not d.ready for d in docs): return {"result": Result(status="needs_review", notes=["Some selected documents contain no usable text. Remove them and upload a clearer PDF."])} return {"docs": docs, "standalone": state["question"], "result": None} def rewrite(state): if not engine.history: return {} response = engine.models.json(Rewrite, "Rewrite the English follow-up as a standalone English question. History only resolves " "references; it is never factual evidence. Do not add an assumed outcome. If the case " "or person is ambiguous, needs_clarification=true.", {"question": state["question"], "history": engine.history[-4:]}, max_tokens=300) if response.needs_clarification: return {"result": Result(status="abstained", notes=["Please name the case or document you mean."])} return {"standalone": response.question} def route(state): if state.get("result") is not None: return "done" if state["per_document"] or re.search(r"\b(which cases|which documents|all cases|how many)\b", state["standalone"], re.I): return "collection" return "retrieve" def retrieve(state): all_evidence = [e for d in state["docs"] for e in evidence_for(d)] # Small/medium collections fit directly: avoid an embedding download becoming # a prerequisite for the first answer, and include every page's context. if sum(len(e.text) for e in all_evidence) <= FULL_CONTEXT_CHAR_LIMIT: return {"evidence": all_evidence} try: if engine.index is None: engine.index = HybridIndex(list(engine.documents.values()), engine.models.embedder()) evidence = engine.index.search(state["standalone"], state["doc_ids"]) except Exception: # Keep retrieval available if local weights cannot load. Same page citations. from .retrieval import terms from rank_bm25 import BM25Okapi corpus = [terms(e.text) for e in all_evidence] scores = BM25Okapi(corpus).get_scores(terms(state["standalone"])) evidence = [all_evidence[i] for i in sorted(range(len(scores)), key=lambda i: scores[i], reverse=True)[:4]] found = {e.id.split(":")[0] for e in evidence} evidence += [evidence_for(d)[-1] for d in state["docs"] if d.id in found] return {"evidence": list({e.id: e for e in evidence}.values())} def answer(state): result = engine.generate(state["standalone"], state["evidence"]) return {"result": result} def collection(state): claims, sources, notes = [], [], [] for doc in state["docs"]: part = engine.generate(state["standalone"] + " Answer only for this document; do not infer a total " "across documents. Report only supported matches.", evidence_for(doc)) claims.extend(part.claims) sources.extend(part.sources) if part.status != "checked": notes.append(f"{doc.name}: no checked answer ({'; '.join(part.notes)}).") notes.append("Each selected document was examined. These are supported findings, not a guaranteed exhaustive list or count.") return {"result": Result(status="checked" if claims else "abstained", claims=claims, sources=list({e.id: e for e in sources}.values()), notes=notes, displayed=[c.text for c in claims])} def remember(state): result = state["result"] if result.status == "checked": engine.history = (engine.history + [{"question": state["standalone"], "answer": " ".join(c.text for c in result.claims)}])[-4:] return {} graph = StateGraph(QuestionState) for name, node in [("conversation", conversation), ("check_scope", scope), ("rewrite_followup", rewrite), ("retrieve", retrieve), ("answer", answer), ("examine_documents", collection), ("remember", remember)]: graph.add_node(name, node) graph.add_edge(START, "conversation") graph.add_conditional_edges("conversation", lambda s: "done" if s.get("result") is not None else "document", {"done": END, "document": "check_scope"}) graph.add_conditional_edges("check_scope", lambda s: "done" if s.get("result") is not None else "rewrite", {"done": END, "rewrite": "rewrite_followup"}) graph.add_conditional_edges("rewrite_followup", route, {"done": END, "collection": "examine_documents", "retrieve": "retrieve"}) graph.add_edge("retrieve", "answer") graph.add_edge("answer", "remember") graph.add_edge("examine_documents", "remember") graph.add_edge("remember", END) return graph.compile(name="document_qa") def build_ingestion_graph(engine): def extract(state): return {"document": read_document(state["name"], state["data"], engine.models)} def store(state): doc = state["document"] engine.documents[doc.id] = doc engine.invalidate() return {} graph = StateGraph(IngestState) graph.add_node("gemini_pdf_extraction", extract) graph.add_node("store_document", store) graph.add_edge(START, "gemini_pdf_extraction") graph.add_edge("gemini_pdf_extraction", "store_document") graph.add_edge("store_document", END) return graph.compile(name="document_ingestion") def build_report_graph(engine): def check(state): doc = engine.documents[state["doc_id"]] if not doc.ready: return {"report": {"fields": {}, "summary": Result(status="needs_review", notes=["This document has no usable text. Upload a clearer PDF."]).model_dump()}} return {"document": doc, "evidence": evidence_for(doc)} def fields(state): fields, notes = engine._extract_fields(state["evidence"]) return {"fields": fields, "notes": notes} def summary(state): return {"summary": engine.generate("Summarize the complete document: " + state["document"].name, state["evidence"], summary=True)} def finish(state): report = {"fields": state["fields"], "field_notes": state["notes"], "summary": state["summary"].model_dump()} if state["summary"].status == "checked" and not state["notes"]: engine.reports[state["doc_id"]] = report return {"report": report} graph = StateGraph(ReportState) for name, node in [("check_ready", check), ("extract_and_check_fields", fields), ("independent_summary", summary), ("finish", finish)]: graph.add_node(name, node) graph.add_edge(START, "check_ready") graph.add_conditional_edges("check_ready", lambda s: "done" if "report" in s else "fields", {"done": END, "fields": "extract_and_check_fields"}) graph.add_edge("extract_and_check_fields", "independent_summary") graph.add_edge("independent_summary", "finish") graph.add_edge("finish", END) return graph.compile(name="independent_document_report")