app/agents/extractor.py

"""Extractor agent: classifies pre-segmented candidates and filters boilerplate.

Operates in batches (≤10 candidates per LLM call) to stay well within token limits and
give incremental output for SSE streaming.
"""

from __future__ import annotations

import logging
import uuid
from dataclasses import dataclass
from typing import ClassVar

from app.agents.base import BaseAgent
from app.schemas import AgentName, Clause, ClauseType, DocType
from app.services.segmenter import CandidateClause

logger = logging.getLogger(__name__)

BATCH_SIZE = 10
_ALLOWED_TYPES: set[str] = {
    "non_compete",
    "arbitration",
    "ip_assignment",
    "auto_renewal",
    "termination",
    "liability",
    "data_collection",
    "payment",
    "confidentiality",
    "indemnification",
    "warranty",
    "governing_law",
    "other",
}
_ALLOWED_DOC_TYPES: set[str] = {
    "employment",
    "freelance",
    "saas_tos",
    "privacy_policy",
    "rental",
    "vendor",
    "nda",
    "other",
}


@dataclass
class ExtractionResult:
    clauses: list[Clause]
    doc_type: DocType


class ExtractorAgent(BaseAgent):
    name: ClassVar[AgentName] = "extractor"
    prompt_file: ClassVar[str] = "extractor"

    async def run(self, candidates: list[CandidateClause]) -> ExtractionResult:
        if not candidates:
            return ExtractionResult(clauses=[], doc_type="other")

        all_clauses: list[Clause] = []
        doc_type_votes: dict[str, int] = {}

        for batch_start in range(0, len(candidates), BATCH_SIZE):
            batch = candidates[batch_start : batch_start + BATCH_SIZE]
            user_prompt = _format_batch(batch)
            response = await self._call(
                user_prompt,
                max_tokens=1500,
                temperature=0.0,
                log_label=f"batch={batch_start // BATCH_SIZE}",
            )
            try:
                items = response.extract_json()
            except ValueError as exc:
                logger.warning("extractor batch parse failure: %s", exc)
                continue

            for item in items:
                if not isinstance(item, dict):
                    continue
                idx = item.get("index")
                if not isinstance(idx, int) or idx < 0 or idx >= len(batch):
                    continue
                if not item.get("keep"):
                    continue
                ctype_raw = item.get("type")
                ctype: ClauseType = (
                    ctype_raw if isinstance(ctype_raw, str) and ctype_raw in _ALLOWED_TYPES else "other"
                )
                candidate = batch[idx]
                clause = Clause(
                    id=f"c_{uuid.uuid4().hex[:8]}",
                    type=ctype,
                    text=candidate.text,
                    span=candidate.span,
                    heading=candidate.heading,
                )
                all_clauses.append(clause)

                signal = item.get("doc_type_signal")
                if isinstance(signal, str) and signal in _ALLOWED_DOC_TYPES:
                    doc_type_votes[signal] = doc_type_votes.get(signal, 0) + 1

        doc_type: DocType = (
            max(doc_type_votes.items(), key=lambda kv: kv[1])[0]  # type: ignore[assignment]
            if doc_type_votes
            else "other"
        )

        summary = (
            f"Extractor kept {len(all_clauses)}/{len(candidates)} candidates; "
            f"document classified as `{doc_type}`."
        )
        self._emit_message(summary)
        logger.info(summary)

        return ExtractionResult(clauses=all_clauses, doc_type=doc_type)


def _format_batch(batch: list[CandidateClause]) -> str:
    lines = [
        "Classify the following candidate clauses. Respond with the JSON array described in your role instructions.",
        "",
    ]
    for i, c in enumerate(batch):
        heading = f"[heading: {c.heading}]" if c.heading else "[no heading]"
        lines.append(f"--- Candidate {i} {heading} ---")
        lines.append(c.text)
        lines.append("")
    return "\n".join(lines)