legal-ai/mcp-server/src/legal_mcp/services/case_citation_verification.py

"""Citation-verification view (X11 Phase 2 / #154) — the chair's "אימות פסיקה" tab.

Assembles, per legal ARGUMENT of a case, the supporting precedents the chair should
verify before the writer cites them:

  • in-corpus suggestions — per-issue semantic retrieval over the authoritative
    precedent library (``search_library``), each carrying the cumulative authority
    signal (``cited_by``: followed/distinguished — db.citation_authority, X11).
  • attached/verified state — any ``case_precedents`` row already attached to the
    argument (verified flag + chair_note), merged onto the matching suggestion.
  • radar — UNLINKED digests relevant to the same issue (rulings we don't hold yet),
    from ``case_digest_radar`` grouped by matched issue.

Pure read/assembly — never writes, never cites (INV-DIG1/INV-AH). The chair verifies
through ``db.set_case_precedent_verified`` / attach; the writer consumes only verified
rows. Reuses the one corpus search + the one authority query + the one radar (G2).
"""

from __future__ import annotations

import asyncio
import logging
from uuid import UUID

from legal_mcp.services import (
    argument_aggregator,
    db,
    digest_library,
    precedent_library,
)

logger = logging.getLogger(__name__)

_SUGGEST_PER_ISSUE = 4
_SUGGEST_FLOOR = 0.45

#: Concurrent per-argument retrievals. The fan-out used to be unbounded — one
#: task per legal argument — which is self-defeating, not merely risky: each
#: task opens a Voyage embed call, and past ~8 in flight Voyage throttles, so
#: every request stalls together. Measured on this corpus (32 arguments):
#:
#:     one search alone ......  1.2s
#:     32 unbounded .......... 30.4s   ← 25× a single search
#:     32 at 8 concurrent .... 22.1s   ← bounded is FASTER
#:
#: A 69-argument case (1069-04-26) therefore blew past the 30s proxy timeout and
#: returned 500, and overlapping page reloads piled contention into Postgres
#: deadlocks. Bounding the fan-out both fixes the failure and speeds it up.
_MAX_CONCURRENT_LOOKUPS = 8

#: Wall-clock ceiling for the whole retrieval phase. The proxy gives up at 30s;
#: cutting ourselves off earlier lets us return the suggestions that DID land
#: instead of a 500 that shows the chair nothing. Partial results are labelled
#: (``retrieval_complete: false``) rather than passed off as the full picture.
_RETRIEVAL_BUDGET_S = 22.0


async def build_view(case_number: str) -> dict:
    case = await db.get_case_by_number(case_number)
    if not case:
        return {"status": "case_not_found", "case_number": case_number, "arguments": []}
    case_id = case["id"]
    if isinstance(case_id, str):
        case_id = UUID(case_id)
    ctx = " ".join(x for x in [case.get("title") or "", case.get("appeal_subtype") or ""] if x).strip()

    args = await argument_aggregator.get_legal_arguments(case_id)

    # Attached precedents already on the case → grouped by argument_id, keyed by
    # the resolved corpus ruling so we can merge verify-state onto a suggestion.
    attached = await db.list_case_precedents(case_id)
    attached_by_arg: dict[str, dict[str, dict]] = {}
    for p in attached:
        aid = str(p.get("argument_id") or "")
        clid = str(p.get("case_law_id") or "")
        if aid and clid:
            attached_by_arg.setdefault(aid, {})[clid] = p

    # Radar (unlinked digests) once, grouped by the issue label it matched.
    radar_by_issue: dict[str, list[dict]] = {}
    try:
        radar = await digest_library.case_digest_radar(case_number, limit=12, min_score=0.42)
        for lead in radar.get("leads", []):
            for label in (lead.get("matched_issues") or ["—"]):
                radar_by_issue.setdefault(label, []).append(lead)
    except Exception as e:  # noqa: BLE001 — radar is best-effort
        logger.warning("citation_verification radar failed for %s: %s", case_number, e)

    async def _fetch(a: dict) -> tuple[list[dict], dict]:
        """Per-argument corpus search + batched authority — run concurrently.

        Each call is one Voyage embed + one vector search (+ one batched authority
        query); independent across arguments, so they fan out rather than waterfall.
        """
        title = (a.get("argument_title") or "").strip()
        topic = (a.get("legal_topic") or "").strip()
        query = f"{ctx} {title}. {topic}".strip()
        hits: list[dict] = []
        try:
            hits = await precedent_library.search_library(
                query=query, limit=_SUGGEST_PER_ISSUE, include_halachot=True)
        except Exception as e:  # noqa: BLE001
            logger.warning("citation_verification search failed (%s): %s", title[:30], e)
        # Resolve the authority breakdown for the hit set in one batched query.
        clids = [UUID(str(h["case_law_id"])) for h in hits
                 if h.get("case_law_id") and float(h.get("score", 0) or 0) >= _SUGGEST_FLOOR]
        authority = await db.citation_authority(clids) if clids else {}
        return hits, authority

    # Fan out the expensive per-argument retrieval, but BOUNDED — see
    # _MAX_CONCURRENT_LOOKUPS. gather preserves order, so the zip below still
    # pairs each argument with its own result.
    _sem = asyncio.Semaphore(_MAX_CONCURRENT_LOOKUPS)

    async def _fetch_bounded(a: dict) -> tuple[list[dict], dict]:
        async with _sem:
            return await _fetch(a)

    retrieval_complete = True
    fetched: list[tuple[list[dict], dict]] = []
    if args:
        # Harvest whatever finished inside the budget, per argument. Deliberately
        # NOT wait_for(gather(...)): that cancels every task on timeout, so one
        # slow lookup would throw away the 30 that already succeeded and the page
        # would show nothing at all.
        tasks = [asyncio.ensure_future(_fetch_bounded(a)) for a in args]
        done, pending = await asyncio.wait(tasks, timeout=_RETRIEVAL_BUDGET_S)
        for t in pending:
            t.cancel()
        if pending:
            # Let the cancellations settle before the caller's DB pool unwinds —
            # a task cancelled mid-query otherwise surfaces as a stray
            # "connection_lost" future with no owner.
            await asyncio.gather(*pending, return_exceptions=True)
            retrieval_complete = False
            logger.warning(
                "citation_verification: retrieval budget of %.0fs exhausted for %s "
                "— %d of %d arguments returned suggestions, the rest are empty",
                _RETRIEVAL_BUDGET_S, case_number, len(done), len(args),
            )
        for t in tasks:
            if t in done and not t.cancelled() and t.exception() is None:
                fetched.append(t.result())
            else:
                fetched.append(([], {}))

    out_args: list[dict] = []
    n_verified = 0
    for a, (hits, authority) in zip(args, fetched):
        aid = str(a["id"])
        title = (a.get("argument_title") or "").strip()
        topic = (a.get("legal_topic") or "").strip()

        seen: set[str] = set()
        supporting: list[dict] = []
        for h in hits:
            clid = str(h.get("case_law_id") or "")
            if not clid or clid in seen:
                continue
            if float(h.get("score", 0) or 0) < _SUGGEST_FLOOR:
                continue
            seen.add(clid)
            att = attached_by_arg.get(aid, {}).get(clid)
            if att and att.get("verified"):
                n_verified += 1
            supporting.append({
                "case_law_id": clid,
                "case_number": h.get("case_number") or "",
                "case_name": h.get("case_name") or "",
                "quote": h.get("supporting_quote") or h.get("rule_statement") or "",
                "score": round(float(h.get("score", 0) or 0), 3),
                "cited_by": authority.get(clid, {"total": 0, "positive": 0,
                                                 "negative": 0, "unclassified": 0,
                                                 "by_treatment": {}}),
                "attached_id": str(att["id"]) if att else None,
                "verified": bool(att.get("verified")) if att else False,
                "chair_note": (att.get("chair_note") or "") if att else "",
            })

        out_args.append({
            "argument_id": aid,
            "title": title,
            "legal_topic": topic,
            "priority": a.get("priority") or "",
            "party": a.get("party") or "",
            "supporting": supporting,
            "radar": radar_by_issue.get(title, []),
        })

    return {
        "status": "ok",
        "case_number": case_number,
        # False when the retrieval budget ran out: the attached/verified rows and
        # the radar are complete, but the corpus SUGGESTIONS are missing. The UI
        # must say so — an empty suggestion list otherwise reads as "no precedent
        # in the corpus supports this argument", which is a different claim.
        "retrieval_complete": retrieval_complete,
        "arguments": out_args,
        "summary": {
            "arguments_total": len(out_args),
            "arguments_with_support": sum(1 for x in out_args if x["supporting"]),
            "verified": n_verified,
            "radar_leads": sum(len(x["radar"]) for x in out_args),
        },
    }