graph/langgraph_pipeline.py

from typing import Any, TypedDict

from langgraph.graph import END, StateGraph

from agents.claim_extractor import ClaimExtractor
from agents.evidence_retriever import EvidenceRetriever
from agents.explainer import Explainer
from graph.jury_subgraph import run_jury_subgraph


class ClaimResult(TypedDict):
    claim: str
    evidence: list[dict[str, Any]]
    verdict: dict[str, Any]
    explanation: dict[str, Any]


class PipelineState(TypedDict):
    input_text: str
    claims: list[str]
    claim_results: list[ClaimResult]
    explanation: dict[str, Any] | None


claim_extractor = ClaimExtractor()
evidence_retriever = EvidenceRetriever()
explainer = Explainer()


VERDICT_LABELS = {
    "SUPPORTED": "مدعوم بالأدلة",
    "REFUTED": "مدحوض",
    "PARTIALLY_TRUE": "صحيح جزئيا",
    "UNVERIFIABLE": "لا يمكن التحقق منه",
}


def _to_plain_dict(value: Any) -> dict[str, Any]:
    if hasattr(value, "model_dump"):
        return value.model_dump()
    if isinstance(value, dict):
        return value
    return {}


def _overall_verdict(claim_results: list[ClaimResult]) -> str:
    if not claim_results:
        return "UNVERIFIABLE"

    verdicts = {
        result.get("verdict", {}).get("verdict", "UNVERIFIABLE")
        for result in claim_results
    }
    verdicts.discard("")

    if len(verdicts) == 1:
        return next(iter(verdicts))

    return "PARTIALLY_TRUE"


def _average_confidence(claim_results: list[ClaimResult]) -> float:
    if not claim_results:
        return 0.0

    confidences = []
    for result in claim_results:
        try:
            confidences.append(
                float(result.get("verdict", {}).get("confidence", 0.0)))
        except (TypeError, ValueError):
            confidences.append(0.0)

    return round(sum(confidences) / len(confidences), 3) if confidences else 0.0


def _dedupe_citations(claim_results: list[ClaimResult]) -> list[dict[str, str]]:
    citations: list[dict[str, str]] = []
    seen: set[tuple[str, str]] = set()

    for result in claim_results:
        for citation in result.get("explanation", {}).get("citations", []):
            if not isinstance(citation, dict):
                continue
            title = str(citation.get("title", "")).strip()
            url = str(citation.get("url", "")).strip()
            key = (title, url)
            if key in seen or (not title and not url):
                continue
            seen.add(key)
            citations.append({"title": title, "url": url})

    return citations


def _build_article_explanation(

    claims: list[str],

    claim_results: list[ClaimResult],

) -> dict[str, Any]:
    if not claims:
        return {
            "arabic_explanation": "لم يتم استخراج أي ادعاءات قابلة للتحقق من النص المرسل.",
            "verdict": "UNVERIFIABLE",
            "confidence": 0.0,
            "citations": [],
        }

    if len(claim_results) == 1:
        explanation = dict(claim_results[0]["explanation"])
        explanation.setdefault("citations", [])
        return explanation

    overall_verdict = _overall_verdict(claim_results)
    overall_confidence = _average_confidence(claim_results)
    overview = (
        f"تم استخراج {len(claim_results)} ادعاءات قابلة للتحقق من النص. "
        f"الحكم الإجمالي على مستوى المقال هو: {VERDICT_LABELS.get(overall_verdict, overall_verdict)}."
    )

    details = []
    for index, result in enumerate(claim_results, start=1):
        verdict = result.get("verdict", {}).get("verdict", "UNVERIFIABLE")
        explanation = str(result.get("explanation", {}).get(
            "arabic_explanation", "")).strip()
        details.append(
            f"{index}. الادعاء: \"{result['claim']}\"\n"
            f"الحكم: {VERDICT_LABELS.get(verdict, verdict)}.\n"
            f"{explanation}"
        )

    return {
        "arabic_explanation": overview + "\n\n" + "\n\n".join(details),
        "verdict": overall_verdict,
        "confidence": overall_confidence,
        "citations": _dedupe_citations(claim_results),
    }


async def run_claim_extractor(state: PipelineState) -> PipelineState:
    claims = await claim_extractor.run(state["input_text"])
    return {**state, "claims": claims}


async def run_evidence_retriever(state: PipelineState) -> PipelineState:
    claim_results: list[ClaimResult] = []

    for claim in state["claims"]:
        evidence = await evidence_retriever.run(claim)
        claim_results.append(
            {
                "claim": claim,
                "evidence": evidence,
                "verdict": {},
                "explanation": {},
            }
        )

    return {**state, "claim_results": claim_results}


async def run_jury(state: PipelineState) -> PipelineState:
    updated_results: list[ClaimResult] = []

    for result in state["claim_results"]:
        verdict = await run_jury_subgraph(result["claim"], result["evidence"])
        updated_results.append({**result, "verdict": verdict})

    return {**state, "claim_results": updated_results}


async def run_explainer(state: PipelineState) -> PipelineState:
    updated_results: list[ClaimResult] = []

    for result in state["claim_results"]:
        verdict_block = result.get("verdict", {})

        explainer_verdict = {
            "verdict":            verdict_block.get("verdict", "UNVERIFIABLE"),
            "confidence":         verdict_block.get("confidence", 0.0),
            "reasoning":          verdict_block.get("reasoning", ""),
            "jury_outputs":       verdict_block.get("jury_outputs", []),
            "debate_log":         verdict_block.get("debate_log", []),
            "needs_human_review": verdict_block.get("needs_human_review", False),
        }

        explanation = await explainer.run(
            result["claim"],
            result["evidence"],
            explainer_verdict,
        )
        updated_results.append(
            {
                **result,
                "explanation": _to_plain_dict(explanation),
            }
        )

    article_explanation = _build_article_explanation(
        state["claims"], updated_results)
    return {
        **state,
        "claim_results": updated_results,
        "explanation": article_explanation,
    }


def build_pipeline() -> StateGraph:
    graph = StateGraph(PipelineState)
    graph.add_node("claim_extractor",    run_claim_extractor)
    graph.add_node("evidence_retriever", run_evidence_retriever)
    graph.add_node("jury",               run_jury)
    graph.add_node("explainer",          run_explainer)

    graph.set_entry_point("claim_extractor")
    graph.add_edge("claim_extractor",    "evidence_retriever")
    graph.add_edge("evidence_retriever", "jury")
    graph.add_edge("jury",               "explainer")
    graph.add_edge("explainer",          END)

    return graph.compile()


def build_jury_explainer_pipeline() -> StateGraph:
    """jury -> explainer only. Used when claims+evidence are supplied

    externally (e.g. test harness holding retrieval fixed to isolate

    reasoning-only variance). Not used by the main endpoint."""
    graph = StateGraph(PipelineState)
    graph.add_node("jury",      run_jury)
    graph.add_node("explainer", run_explainer)

    graph.set_entry_point("jury")
    graph.add_edge("jury",      "explainer")
    graph.add_edge("explainer", END)

    return graph.compile()


pipeline = build_pipeline()
jury_explainer_pipeline = build_jury_explainer_pipeline()


def _serialize_claim_results(claim_results: list[ClaimResult]) -> list[dict[str, Any]]:
    out = []
    for result in claim_results:
        exp = result.get("explanation", {})
        verdict_block = result.get("verdict", {})
        out.append({
            "claim":       result["claim"],
            "verdict":     verdict_block.get("verdict", "UNVERIFIABLE"),
            "confidence":  round(float(verdict_block.get("confidence", 0.0)), 3),
            "explanation": exp.get("arabic_explanation", ""),
            "citations": [
                {"title": c.get("title", ""), "url": c.get("url", "")}
                for c in exp.get("citations", [])
                if isinstance(c, dict)
            ],
            "needs_human_review": verdict_block.get("needs_human_review", False),
            "jury_outputs":       verdict_block.get("jury_outputs", []),
            "debate_log":         verdict_block.get("debate_log", []),
        })
    return out


def _finalize(final_state: PipelineState) -> dict[str, Any]:
    """Shared by run_pipeline and run_pipeline_from_evidence. Pure

    function of final_state -> identical output shape either way."""
    article = final_state.get("explanation") or {
        "arabic_explanation": "لم يتمكن النظام من توليد نتيجة نهائية.",
        "verdict": "UNVERIFIABLE",
        "confidence": 0.0,
        "citations": [],
    }

    claim_results = final_state.get("claim_results", [])
    needs_human_review = any(
        r.get("verdict", {}).get("needs_human_review", False)
        for r in claim_results
    )

    return {
        "verdict":            article.get("verdict", "UNVERIFIABLE"),
        "confidence":         round(float(article.get("confidence", 0.0)), 3),
        "arabic_explanation": article.get("arabic_explanation", ""),
        "citations":          article.get("citations", []),
        "needs_human_review": needs_human_review,
        "claims":             _serialize_claim_results(claim_results),
    }


async def run_pipeline(text: str) -> dict[str, Any]:
    """Main endpoint entry point — behavior unchanged."""
    initial_state: PipelineState = {
        "input_text": text,
        "claims": [],
        "claim_results": [],
        "explanation": None,
    }
    final_state = await pipeline.ainvoke(initial_state)
    return _finalize(final_state)


async def run_pipeline_from_evidence(

    claims: list[str],

    evidence_map: dict[str, list[dict[str, Any]]],

) -> dict[str, Any]:
    """TEST-ONLY entry point. Skips claim_extractor/evidence_retriever;

    consumes pre-fetched claims+evidence so retrieval is pinned and only

    jury/explainer reasoning varies across repeated runs."""
    claim_results: list[ClaimResult] = [
        {
            "claim": claim,
            "evidence": evidence_map.get(claim, []),
            "verdict": {},
            "explanation": {},
        }
        for claim in claims
    ]
    initial_state: PipelineState = {
        "input_text": "",
        "claims": claims,
        "claim_results": claim_results,
        "explanation": None,
    }
    final_state = await jury_explainer_pipeline.ainvoke(initial_state)
    return _finalize(final_state)