loop.py

"""The bounded, code-driven ReAct loop.

reason -> act -> observe over a fixed whitelist (parse, chunk, select, extract,
validate)
ending at a hard interrupt. The model is called only inside extract; every other
step is deterministic code. The loop cannot exceed STEP_CAP and cannot call
outside the whitelist — those invariants are what the eval harness asserts.
"""

from __future__ import annotations

from collections.abc import Callable
from dataclasses import dataclass, field

from hipo_contracts import DegradationState, Rejection, ToolName, TraceStep

from hipo_api.agent.extract import EXTRACT_MODEL, ExtractParseError, extract_claims
from hipo_api.agent.normalize import normalize_claims
from hipo_api.agent.select import select_chunks
from hipo_api.build import build_candidate
from hipo_api.chunk import chunk_pages
from hipo_api.parse import parse_pdf
from hipo_api.validate import validate_factsheet

STEP_CAP = 8


@dataclass
class AgentResult:
    steps: list[TraceStep] = field(default_factory=list)
    state: DegradationState = DegradationState.REJECTED
    payload: dict = field(default_factory=dict)
    terminated_at_interrupt: bool = False


def run_agent(
    pdf_bytes: bytes,
    entry: dict,
    client,
    *,
    model: str = EXTRACT_MODEL,
    on_step: Callable[[TraceStep], None] | None = None,
    on_token: Callable[[str], None] | None = None,
) -> AgentResult:
    result = AgentResult()
    n = 0

    def record(thought, tool, args_digest, observation, model_name=None):
        nonlocal n
        step = TraceStep(
            step=n,
            thought=thought,
            tool=tool,
            args_digest=args_digest,
            observation=observation,
            model=model_name,
        )
        result.steps.append(step)
        if on_step is not None:
            on_step(step)
        n += 1

    doc = parse_pdf(pdf_bytes)
    record(
        "read the filing",
        ToolName.PARSE,
        {"pages": doc.page_count},
        f"{doc.page_count} pages",
    )

    chunks = chunk_pages(doc)
    record(
        "select candidate pages",
        ToolName.CHUNK,
        {"chunks": len(chunks)},
        f"{len(chunks)} non-empty chunks",
    )

    selected, dropped_pages = select_chunks(chunks)
    record(
        "narrow to the most relevant pages within budget",
        ToolName.SELECT,
        {"kept": len(selected), "dropped": dropped_pages},
        f"selected {len(selected)} of {len(chunks)} pages",
    )

    try:
        claims = extract_claims(selected, client, model=model, on_token=on_token)
    except ExtractParseError as exc:
        record(
            "extract claims from the filing",
            ToolName.EXTRACT,
            {},
            f"extraction failed: {exc}",
            model_name=model,
        )
        rej = Rejection.single("extract", "model extraction failed")
        result.payload = rej.model_dump()
        result.state = DegradationState.REJECTED
        record(
            "stop — nothing to surface",
            ToolName.INTERRUPT,
            {},
            "surfaced rejection",
        )
        result.terminated_at_interrupt = True
        return result

    n_claims = len(claims.get("financial_highlights", [])) + len(
        claims.get("risks", [])
    )
    record(
        "extract claims from the filing",
        ToolName.EXTRACT,
        {"claims": n_claims},
        f"model returned {n_claims} claims",
        model_name=model,
    )

    base = build_candidate(doc, entry)
    candidate = normalize_claims(base, claims)
    kept = len(candidate["financial_highlights"]) + len(candidate["risks"])
    dropped = n_claims - kept

    outcome = validate_factsheet(candidate, doc)
    if isinstance(outcome, Rejection):
        result.payload = outcome.model_dump()
        result.state = DegradationState.REJECTED
        record(
            "validate schema + citations",
            ToolName.VALIDATE,
            {"dropped": dropped},
            "validation rejected the candidate",
        )
    else:
        result.payload = outcome.model_dump()
        result.state = (
            DegradationState.FULL if dropped == 0 else DegradationState.PARTIAL
        )
        record(
            "validate schema + citations",
            ToolName.VALIDATE,
            {"dropped": dropped},
            f"validated; {dropped} claim(s) dropped",
        )

    record(
        "draft ready — never auto-act",
        ToolName.INTERRUPT,
        {},
        "surfaced draft factsheet",
    )
    result.terminated_at_interrupt = True
    return result