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

"""Decision-lesson synthesis (#158 / INV-LRN8).

The learning channels (panel, curator, chair) accumulate overlapping ``decision_lessons``
on the same style dimension. The writer only consumes the 15 most-recent APPROVED ones
per practice_area (a silent cap, #157), so beyond that lessons pile up unused. This pass
**clusters** near-duplicate lessons within a (practice_area, category) shard and **merges**
each cluster into ONE richer, generalised "super-lesson" — so the writer-fed set shrinks
to a small, high-quality profile (Authorial Style Profiling, ספ §0.1) and the cap stops
biting.

Mirrors the canonical-halacha synthesis (V41 / INV-LRN6) on the SAME table (no parallel
store, G2). Invariants:
  • INV-AH  — the super-lesson is GROUNDED in the source lessons only; the model abstains
              rather than invent, and a re-embedding DRIFT guard rejects a rewrite that
              drifts from the cluster centroid.
  • INV-LRN1/G10 — graduated gate (chair decision 2026-06-28): since every source is
              already ``approved``, the super-lesson flows as ``approved`` (reversible —
              chair veto in /training restores the sources). A non-approved source ⇒ proposed.
  • G2  — single synthesis path; the MCP tool and the backfill script both call
          :func:`run_shard` / :func:`synthesize_cluster` here.
  • G9  — every outcome (accepted / abstained / drift_rejected / merged-duplicate) returned.

LLM calls go through ``claude_session`` (local ``claude -p`` CLI) only — never from the
FastAPI container (see claude_session docstring).
"""

from __future__ import annotations

import logging
import math
from uuid import UUID

from legal_mcp import config
from legal_mcp.services import claude_session, db, embeddings

logger = logging.getLogger(__name__)

_SYSTEM = (
    "אתה עורך-דין בכיר המזקק כללי-סגנון-וכתיבה לבסיס-ידע של ועדת ערר לתכנון ובנייה. "
    "תפקידך למזג כמה לקחי-סגנון חופפים לכלל אחד, כללי ומדויק, על *איך* כותבים — לא להמציא "
    "כלל חדש ולא להוסיף מהות משפטית."
)


def _cosine(a, b) -> float:
    dot = sum(x * y for x, y in zip(a, b))
    na = math.sqrt(sum(x * x for x in a))
    nb = math.sqrt(sum(y * y for y in b))
    if na == 0 or nb == 0:
        return 0.0
    return dot / (na * nb)


def _centroid(vecs: list[list[float]]) -> list[float]:
    n = len(vecs)
    dim = len(vecs[0])
    return [sum(v[i] for v in vecs) / n for i in range(dim)]


def _build_prompt(members: list[dict]) -> str:
    blocks = []
    for i, m in enumerate(members, 1):
        blocks.append(f"### לקח {i}\n{m['lesson_text']}")
    evidence = "\n\n".join(blocks)
    return f"""{_SYSTEM}

לקחי-המקור (כולם מאותו תחום וקטגוריה, חופפים בנושא):
{evidence}

## המשימה
מזג את לקחי-המקור לכלל-סגנון **אחד** עשיר ומוכלל המשותף לכולם. שמר כל ניואנס מובחן שמופיע
באחד הלקחים, אך נסח אותו פעם אחת, נקי וכללי.

## כללים מחייבים (INV-AH — עיגון, ללא הזיה)
1. **עיגון-מקור בלבד.** הכלל חייב לנבוע מלקחי-המקור שלמעלה. אסור להוסיף כלל, חריג או דוגמה שאינם עולים מהם.
2. **סגנון/שיטה בלבד, לא מהות.** אל תכניס הלכה, עובדה, מספר-תיק או תקדים ספציפי — רק *איך* דפנה כותבת.
3. **כללי ובלתי-תלוי-תיק.** הסר פרטים קונקרטיים; נסח כלל רב-תחולה.
4. **רגיסטר נקי** בעברית, משפט אחד עד שלושה, בלי מילות-מסגרת ("יש לזכור ש...") — רק הכלל עצמו.
5. **הימנעות עדיפה על המצאה.** אם הלקחים אינם באמת מתמזגים לכלל אחד מעוגן — החזר grounded=false.

## פלט — JSON בלבד, ללא markdown וללא הסבר:
{{
  "lesson_text": "<כלל-הסגנון הממוזג>",
  "grounded": true,
  "reason": "<משפט קצר: מה אוחד>"
}}"""


def _greedy_clusters(candidates: list[dict], threshold: float) -> list[list[dict]]:
    """Greedy single-link clustering by cosine over candidate embeddings. Each candidate
    has an 'embedding' (python list). Returns clusters of size ≥2 only (singletons are
    nothing to merge)."""
    remaining = [c for c in candidates if c.get("embedding") is not None]
    clusters: list[list[dict]] = []
    used: set = set()
    for i, seed in enumerate(remaining):
        if seed["id"] in used:
            continue
        cluster = [seed]
        used.add(seed["id"])
        for other in remaining[i + 1:]:
            if other["id"] in used:
                continue
            if _cosine(seed["embedding"], other["embedding"]) >= threshold:
                cluster.append(other)
                used.add(other["id"])
        if len(cluster) >= 2:
            clusters.append(cluster)
    return clusters


async def _ensure_embeddings(candidates: list[dict]) -> list[dict]:
    """Lazy-backfill: embed any candidate whose stored embedding is NULL, persist it,
    and return the candidates with embeddings populated (skips ones that still fail)."""
    missing = [c for c in candidates if c.get("embedding") is None]
    if missing:
        vecs = await embeddings.embed_texts([c["lesson_text"] for c in missing])
        for c, v in zip(missing, vecs):
            c["embedding"] = list(v)
            await db.set_lesson_embedding(c["id"], c["embedding"])
    return [c for c in candidates if c.get("embedding") is not None]


async def synthesize_cluster(
    members: list[dict],
    *,
    model: str | None = None,
    effort: str | None = None,
    drift_floor: float | None = None,
) -> dict:
    """Merge one cluster of lessons. PURE — no DB writes. Returns:
      {status, proposed, embedding, members:[ids], drift_cosine, reason}
    status ∈ {accepted, abstained, drift_rejected, llm_error, too_small}.
    """
    model = model or config.LESSON_SYNTH_MODEL
    effort = effort or config.LESSON_SYNTH_EFFORT
    drift_floor = config.LESSON_SYNTH_DRIFT_FLOOR if drift_floor is None else drift_floor
    ids = [str(m["id"]) for m in members]
    base = {"members": ids, "proposed": "", "embedding": None,
            "drift_cosine": None, "reason": ""}
    if len(members) < 2:
        return {**base, "status": "too_small", "reason": "cluster < 2"}

    try:
        result = await claude_session.query_json(
            _build_prompt(members), model=model, effort=effort, tools="",
        )
    except Exception as e:
        logger.warning("synthesize_cluster %s: LLM error: %s", ids, e)
        return {**base, "status": "llm_error", "reason": str(e)}

    if not isinstance(result, dict) or not result.get("lesson_text"):
        return {**base, "status": "llm_error", "reason": "malformed LLM output"}
    if not result.get("grounded", True):
        return {**base, "status": "abstained",
                "reason": result.get("reason") or "model abstained (not grounded)"}

    proposed = str(result["lesson_text"]).strip()
    if not proposed:
        return {**base, "status": "abstained", "reason": "empty proposal"}

    # Drift guard: the merged lesson must stay near the cluster centroid.
    new_emb = list((await embeddings.embed_texts([proposed]))[0])
    centroid = _centroid([m["embedding"] for m in members])
    drift = _cosine(new_emb, centroid)
    if drift < drift_floor:
        return {**base, "status": "drift_rejected", "proposed": proposed,
                "drift_cosine": round(drift, 4),
                "reason": f"drift {drift:.3f} < floor {drift_floor}"}

    return {**base, "status": "accepted", "proposed": proposed, "embedding": new_emb,
            "drift_cosine": round(drift, 4), "reason": result.get("reason") or "merged"}


async def run_shard(
    practice_area: str,
    category: str,
    *,
    apply: bool,
    model: str | None = None,
    effort: str | None = None,
    drift_floor: float | None = None,
    cluster_threshold: float | None = None,
) -> dict:
    """Synthesize all clusters in one (practice_area, category) shard.

    Returns {practice_area, category, candidates, clusters:[result...]}. Each result is
    a synthesize_cluster outcome augmented with ``applied`` and (when applied) ``new_id``.
    With apply=False this is a pure dry-run (no writes beyond lazy embedding backfill).
    """
    cluster_threshold = (config.LESSON_SYNTH_CLUSTER_THRESHOLD
                         if cluster_threshold is None else cluster_threshold)
    candidates = await db.fetch_synthesis_candidates(practice_area, category)
    candidates = await _ensure_embeddings(candidates)
    clusters = _greedy_clusters(candidates, cluster_threshold)

    results = []
    for members in clusters:
        res = await synthesize_cluster(
            members, model=model, effort=effort, drift_floor=drift_floor,
        )
        res["applied"] = False
        if apply and res["status"] == "accepted":
            # Idempotency: skip if a near-identical synthesis already exists (re-run safe).
            dup = await db.nearest_synthesis_lesson(
                res["embedding"], category, config.HALACHA_CANONICAL_THRESHOLD,
            )
            if dup:
                res["status"] = "duplicate_skipped"
                res["reason"] = f"near existing synthesis {dup[0]} (sim {dup[1]:.3f})"
            else:
                # graduated gate: all sources are approved (fetch filter) → approved.
                row = await db.apply_lesson_synthesis(
                    corpus_id=members[0]["style_corpus_id"],
                    lesson_text=res["proposed"],
                    category=category,
                    embedding=res["embedding"],
                    source_ids=[m["id"] for m in members],
                    review_status="approved",
                )
                res["applied"] = True
                res["new_id"] = str(row.get("id", ""))
        results.append(res)

    return {"practice_area": practice_area or "*", "category": category,
            "candidates": len(candidates), "clusters": results}


async def run_pending(
    practice_area: str = "",
    category: str = "",
    *,
    apply: bool,
    model: str | None = None,
    effort: str | None = None,
    drift_floor: float | None = None,
    cluster_threshold: float | None = None,
) -> list[dict]:
    """Run synthesis across shards. If practice_area+category are given, one shard;
    otherwise iterate every shard with ≥2 live approved lessons. Single entry point (G2)."""
    if practice_area and category:
        shards = [{"practice_area": practice_area, "category": category}]
    else:
        shards = await db.synthesis_shards(min_size=2)
        if practice_area:
            shards = [s for s in shards if s["practice_area"] == practice_area]
        if category:
            shards = [s for s in shards if s["category"] == category]
    out = []
    for s in shards:
        out.append(await run_shard(
            s["practice_area"], s["category"], apply=apply,
            model=model, effort=effort, drift_floor=drift_floor,
            cluster_threshold=cluster_threshold,
        ))
    return out