AIRegulation-DocAnalysis/backend/app/infrastructure/perception/llm_pipeline.py

"""LLM-driven pipeline for regulatory event enrichment."""

from __future__ import annotations

import json
from typing import Any

from loguru import logger

from app.config.settings import settings
from app.infrastructure.perception.regulation_differ import (
    ParagraphChange,
    RegulationDiffer,
)
from app.services.llm.llm_factory import get_llm_client

_EXTRACT_SYSTEM = (
    "You are a regulatory compliance expert specialising in automotive standards "
    "(GB, UN-ECE, ISO, EU). Extract structured information from regulation text. "
    "Return valid JSON only — no markdown fences, no extra keys."
)

_ASSESS_SYSTEM = (
    "You are an automotive compliance analyst. Given a regulation and related document excerpts, "
    "identify which documents are affected and what actions are required. "
    "Return a JSON array only."
)

_DIFF_SYSTEM = (
    "You are a regulatory change analyst. You are given the OLD and NEW version of "
    "one regulation paragraph, with the exact edits marked <DEL>removed</DEL> and "
    "<INS>added</INS>. Classify the legal effect of the change. "
    "Return JSON only: {\"change_type\": \"tightened|relaxed|numeric|clarified|scope\", "
    "\"legal_effect\": \"one sentence on what this means for compliance\"}"
)


def _marked_diff(change: ParagraphChange) -> str:
    """Render a paragraph change with the exact edits marked for the model.

    The model is shown where the edit is rather than being asked to find it,
    and is never asked to reproduce the changed text — the differ already
    computed those spans exactly, so there is nothing for the model to
    hallucinate.
    """
    marked = "".join(
        text if op == 0 else (f"<DEL>{text}</DEL>" if op < 0 else f"<INS>{text}</INS>")
        for op, text in change.diff_spans
    )
    return (
        f"OLD: {change.old_text[:500]}\n"
        f"NEW: {change.new_text[:500]}\n"
        f"MARKED: {marked[:800]}"
    )


def _llm_json(client: Any, messages: list[dict]) -> Any:
    """Call LLM and parse JSON response; return None on failure."""
    try:
        resp = client.chat(messages)
        text = (resp.content or "").strip()
        if text.startswith("```"):
            text = text.split("```")[1]
            if text.startswith("json"):
                text = text[4:]
        return json.loads(text)
    except Exception as exc:
        logger.warning("LLM JSON parse failed: {}", exc)
        return None


class LlmPipeline:
    """Three-step enrichment pipeline for crawled regulatory events."""

    def __init__(self) -> None:
        self._client = get_llm_client(
            provider=settings.llm_provider,
            model=settings.llm_model,
        )
        # Change detection is deterministic; the differ needs no model and no
        # network, so the pipeline no longer constructs an embedding provider.
        self._differ = RegulationDiffer()

    # ------------------------------------------------------------------
    # Step 1: Structure extraction
    # ------------------------------------------------------------------

    def extract_structure(self, event: dict) -> dict:
        """Extract obligations, deadlines, scope, penalties, impact_level from event text."""
        prompt = f"""Extract structured compliance information from this regulation:

Standard: {event.get('standard_code', '')}
Title: {event.get('title', '')}
Source: {event.get('source_label', '')}
Summary: {event.get('summary', '')}
Tags: {', '.join(event.get('tags') or [])}

Return JSON with exactly these keys:
{{
  "obligations": [{{"text": "...", "deontic": "must|shall|may|prohibited", "subject": "...", "object": "...", "condition": ""}}],
  "deadlines": [{{"date": "YYYY-MM-DD or null", "description": "..."}}],
  "scope": "one sentence describing who/what this applies to",
  "penalties": "one sentence on consequences of non-compliance, or null",
  "impact_level": "high|medium|low"
}}"""

        messages = [
            {"role": "system", "content": _EXTRACT_SYSTEM},
            {"role": "user", "content": prompt},
        ]
        result = _llm_json(self._client, messages)
        if not isinstance(result, dict):
            return {
                "obligations": [],
                "deadlines": [],
                "scope": "",
                "penalties": "",
                "impact_level": "medium",
            }
        return result

    # ------------------------------------------------------------------
    # Step 2: Impact assessment
    # ------------------------------------------------------------------

    def assess_impact(self, event: dict, retrieval_service: Any) -> list[dict]:
        """Use RAG to find affected documents and generate recommendations."""
        obligations = event.get("obligations") or []
        obligation_texts = " ".join(o.get("text", "") for o in obligations[:3])
        query = f"{event.get('standard_code', '')} {event.get('title', '')} {obligation_texts}"

        try:
            chunks = retrieval_service.retrieve(query=query, top_k=5)
        except Exception as exc:
            logger.warning("RAG retrieval failed: {}", exc)
            return []

        if not chunks:
            return []

        seen: set[str] = set()
        doc_excerpts: list[dict] = []
        for chunk in chunks:
            if chunk.doc_id not in seen:
                seen.add(chunk.doc_id)
                doc_excerpts.append({
                    "doc_id": chunk.doc_id,
                    "doc_name": chunk.doc_title,
                    "score": round(float(chunk.score if chunk.score is not None else 0), 4),
                    "snippet": (chunk.text or "")[:300],
                    "clause": getattr(chunk, "section_title", "") or "",
                })

        context = "\n".join(
            f"[{d['doc_name']} {d['clause']}] score={d['score']}: {d['snippet']}"
            for d in doc_excerpts
        )
        prompt = f"""Regulation: {event.get('standard_code')} — {event.get('title')}
Obligations: {obligation_texts or event.get('summary', '')}

Affected documents found in knowledge base:
{context}

For each document, assess impact and recommend action. Return JSON array:
[{{"doc_id":"...","doc_name":"...","score":0.0,"key_clauses":"...","recommendation":"one sentence action"}}]"""

        messages = [
            {"role": "system", "content": _ASSESS_SYSTEM},
            {"role": "user", "content": prompt},
        ]
        result = _llm_json(self._client, messages)
        if isinstance(result, list):
            score_map = {d["doc_id"]: d["score"] for d in doc_excerpts}
            for item in result:
                if isinstance(item, dict) and item.get("doc_id") in score_map:
                    item["score"] = score_map[item["doc_id"]]
            return result
        return doc_excerpts

    # ------------------------------------------------------------------
    # Step 3: Deterministic diff with gated LLM classification
    # ------------------------------------------------------------------

    def compute_diff(self, old_text: str, new_text: str) -> dict:
        """Compare old and new regulation text; return changed sections and summary.

        Detection is deterministic — see regulation_differ for why embedding
        similarity was removed. The LLM is called only for paragraphs the
        differ marked significant, and only to explain the legal effect of a
        change that has already been located exactly.
        """
        changes = self._differ.diff(old_text, new_text)
        if not changes:
            return {
                "changed_sections": [],
                "change_summary": "No substantive changes detected between versions.",
            }

        changed_sections = [self._describe(change) for change in changes]

        types = sorted({section["change_type"] for section in changed_sections})
        gated = sum(1 for change in changes if change.needs_llm)
        change_summary = (
            f"{len(changed_sections)} paragraph(s) changed ({', '.join(types)}); "
            f"{gated} significant. "
            + (changed_sections[0].get("summary") or "")
        ).strip()

        return {"changed_sections": changed_sections, "change_summary": change_summary}

    def _describe(self, change: ParagraphChange) -> dict:
        """Turn one detected change into the API payload, classifying if warranted."""
        section = {
            "old_text": change.old_text[:300],
            "new_text": change.new_text[:300],
            "change_type": change.change_type,
            "change_ratio": round(change.change_ratio, 3),
            "numeric_changed": change.numeric_changed,
            "deontic_changed": change.deontic_changed,
            "summary": "",
        }

        if not change.needs_llm:
            return section

        classification = _llm_json(
            self._client,
            [
                {"role": "system", "content": _DIFF_SYSTEM},
                {"role": "user", "content": _marked_diff(change)},
            ],
        )
        if isinstance(classification, dict):
            section["change_type"] = classification.get("change_type") or change.change_type
            section["summary"] = classification.get("legal_effect") or ""
        # A failed or malformed model response must not discard a change that
        # deterministic analysis already proved real; the section keeps its
        # spans, flags, and alignment-derived type with an empty summary.

        if change.numeric_changed:
            # Models routinely label a changed threshold as "clarified". The
            # deterministic pass already knows a number moved, so it wins.
            section["change_type"] = "numeric"

        return section