lexquery/workflows.py

"""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")