"""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)