app/services/orchestrator.py

"""Top-level pipeline orchestration — 5-agent pipeline.

Flow:
  parse → NLP entity extraction → segment → extractor agent
  → risk analyst (parallel, concurrency=2)
  → devil's advocate (sequential, medium+)
  → legal context agent (sequential, critical only)
  → final RiskReport

Emits AgentMessage + ClauseUpdate events through Job queue for SSE live streaming.
"""

from __future__ import annotations

import asyncio
import logging
import time
from typing import Any

from app.agents.devil_advocate import DevilAdvocateAgent
from app.agents.extractor import ExtractorAgent
from app.agents.legal_context import LegalContextAgent
from app.agents.risk_analyst import RiskAnalystAgent
from app.config import get_settings
from app.schemas import AgentMessage, Clause, DocumentEntities, RiskReport, SEVERITY_RANK
from app.services import retrieval
from app.services.entities import extract_entities
from app.services.jobs import Job
from app.services.parser import parse
from app.services.segmenter import segment

logger = logging.getLogger(__name__)

_TECH_STACK = [
    "Multi-Agent AI (5 agents)",
    "Groq Llama-3.3-70B",
    "RAG · ChromaDB",
    "BGE-small Embeddings",
    "NLP Entity Extraction",
    "Explainable AI",
    "OCR · pytesseract",
    "Vector Semantic Search",
]


async def run_analysis(job: Job, filename: str, raw: bytes) -> None:
    settings = get_settings()
    started = time.perf_counter()
    trace: list[AgentMessage] = []

    def push(message: AgentMessage) -> None:
        trace.append(message)
        job.emit_agent_message(message)

    try:
        # ── 1. Parse ────────────────────────────────────────────────────────
        parsed = parse(filename, raw)
        ocr_note = " (OCR)" if parsed.ocr_used else ""
        push(AgentMessage(
            agent="extractor",
            content=(
                f"Parsed {parsed.source_format.upper()}{ocr_note} — "
                f"{parsed.char_count:,} chars across {parsed.page_count} page(s)."
            ),
            timestamp=time.time(),
        ))

        # ── 2. NLP entity extraction (zero LLM cost) ─────────────────────
        doc_entities: DocumentEntities = extract_entities(parsed.text)
        entity_summary_parts = []
        if doc_entities.parties:
            entity_summary_parts.append(f"Parties: {', '.join(doc_entities.parties[:4])}")
        if doc_entities.jurisdictions:
            entity_summary_parts.append(f"Jurisdiction: {', '.join(doc_entities.jurisdictions[:3])}")
        if doc_entities.monetary_amounts:
            entity_summary_parts.append(f"Amounts: {', '.join(doc_entities.monetary_amounts[:4])}")
        if doc_entities.durations:
            entity_summary_parts.append(f"Durations: {', '.join(doc_entities.durations[:4])}")
        if doc_entities.key_obligations:
            entity_summary_parts.append(f"Key patterns: {', '.join(doc_entities.key_obligations)}")
        if entity_summary_parts:
            push(AgentMessage(
                agent="extractor",
                content="NLP entity extraction complete. " + " · ".join(entity_summary_parts),
                timestamp=time.time(),
            ))

        # ── 3. Segment ───────────────────────────────────────────────────
        candidates = segment(parsed.text)
        candidates = candidates[: settings.max_clauses_per_doc]
        push(AgentMessage(
            agent="extractor",
            content=f"Segmented into {len(candidates)} candidate clauses for analysis.",
            timestamp=time.time(),
        ))

        # ── 4. Extractor agent ───────────────────────────────────────────
        extractor = ExtractorAgent(document=parsed.text, emit=push)
        extraction = await extractor.run(candidates)
        clauses = extraction.clauses
        doc_type = extraction.doc_type

        if not clauses:
            report = RiskReport(
                job_id=job.id,
                doc_type=doc_type,
                overall_risk=0.0,
                clauses=[],
                summary=(
                    "No substantive clauses were detected. "
                    "The document may be a template, table of contents, or otherwise lack binding obligations."
                ),
                agent_trace=trace,
                document_entities=doc_entities,
                tech_stack=_TECH_STACK,
            )
            job.complete(report)
            return

        # ── 5. Risk Analyst (parallel, concurrency=2) ─────────────────────
        analyst = RiskAnalystAgent(document=parsed.text, emit=push)

        async def analyze_one(clause: Clause) -> Clause:
            matches = retrieval.fetch_benchmark_matches(
                clause.text, clause.type, k=settings.rag_top_k
            )
            scored = await analyst.run(clause, matches)
            job.emit_clause_update(scored)
            return scored

        scored_clauses = await _gather_with_concurrency(2, [analyze_one(c) for c in clauses])

        # ── 6. Devil's Advocate (sequential, medium+) ─────────────────────
        min_rank = SEVERITY_RANK[settings.devil_advocate_min_severity]  # type: ignore[index]
        devil_targets = [c for c in scored_clauses if SEVERITY_RANK[c.severity] >= min_rank]
        if devil_targets:
            devil = DevilAdvocateAgent(document=parsed.text, emit=push)

            async def amplify(clause: Clause) -> Clause:
                try:
                    worst = await devil.run(clause)
                    clause.worst_case = worst
                    job.emit_clause_update(clause)
                except Exception as exc:
                    logger.warning("devil_advocate skipped clause %s: %s", clause.id, exc)
                return clause

            await _gather_with_concurrency(1, [amplify(c) for c in devil_targets])

        # ── 7. Legal Context agent (sequential, critical only) ─────────────
        critical_clauses = [c for c in scored_clauses if c.severity == "critical"]
        if critical_clauses:
            legal = LegalContextAgent(document=parsed.text, emit=push)

            async def contextualize(clause: Clause) -> Clause:
                try:
                    updated = await legal.run(clause)
                    job.emit_clause_update(updated)
                    return updated
                except Exception as exc:
                    logger.warning("legal_context skipped clause %s: %s", clause.id, exc)
                    return clause

            await _gather_with_concurrency(1, [contextualize(c) for c in critical_clauses])

        # ── 8. Assemble final report ───────────────────────────────────────
        overall = _aggregate_risk(scored_clauses)
        summary = _build_summary(doc_type, scored_clauses, overall)

        report = RiskReport(
            job_id=job.id,
            doc_type=doc_type,
            overall_risk=overall,
            clauses=scored_clauses,
            summary=summary,
            agent_trace=trace,
            document_entities=doc_entities,
            tech_stack=_TECH_STACK,
        )
        job.complete(report)
        logger.info(
            "job %s complete in %.1fs — %d clauses, overall_risk=%.1f",
            job.id,
            time.perf_counter() - started,
            len(scored_clauses),
            overall,
        )

    except Exception as exc:
        logger.exception("orchestrator failure on job %s", job.id)
        job.fail(str(exc) or exc.__class__.__name__)


async def _gather_with_concurrency(limit: int, coros: list) -> list[Any]:
    semaphore = asyncio.Semaphore(limit)

    async def guarded(coro):
        async with semaphore:
            return await coro

    return await asyncio.gather(*(guarded(c) for c in coros))


def _aggregate_risk(clauses: list[Clause]) -> float:
    if not clauses:
        return 0.0
    weights = {"low": 0.4, "medium": 1.0, "high": 2.0, "critical": 3.0}
    weighted_sum = sum(c.risk_score * weights[c.severity] for c in clauses)
    weight_total = sum(weights[c.severity] for c in clauses)
    if weight_total == 0:
        return 0.0
    return round(weighted_sum / weight_total, 1)


def _build_summary(doc_type: str, clauses: list[Clause], overall: float) -> str:
    by_sev: dict[str, int] = {"low": 0, "medium": 0, "high": 0, "critical": 0}
    for c in clauses:
        by_sev[c.severity] = by_sev.get(c.severity, 0) + 1

    parts = [
        f"Detected a `{doc_type}` document with {len(clauses)} substantive clauses; "
        f"overall risk score {overall:.0f}/100.",
    ]
    if by_sev["critical"]:
        parts.append(
            f"{by_sev['critical']} clause(s) flagged CRITICAL — do not sign without amendment."
        )
    if by_sev["high"]:
        parts.append(f"{by_sev['high']} clause(s) flagged HIGH — push to renegotiate.")
    if by_sev["medium"]:
        parts.append(f"{by_sev['medium']} clause(s) flagged MEDIUM — be aware before signing.")
    if by_sev["low"] and not (by_sev["critical"] or by_sev["high"]):
        parts.append("Remaining clauses are within standard-of-market ranges.")
    return " ".join(parts)