portfolio-optimiser-claude/src/portfolio_optimiser_claude/loop.py

"""The agentic loop, steps 2–5 (method-spec §3): generate, debate, gate, refine.

The model is reached ONLY through the ``ModelClient`` protocol — an injected
client (the Claude Agent SDK on the run path, S10; a scripted stand-in in the
offline suite, honesty rule §1). Step 2: a reply that fails to parse into the
typed IR is retried BLIND, never silently accepted or repaired downstream,
bounded by the budget meter (§8). Step 3: the round-capped maker-checker
debate, fresh state per run, the checker instructed to end with exactly one
verdict line. Step 4 (second falsifier): the checker gate — opt-in-reject
(fail-open), parsed from the checker's LAST surfaced output; an explicit
REJECT overrides an otherwise-validated outcome; a validator rejection stands
regardless. Step 5: informed refinement — the most recent rejection REASON
verbatim in the next prompt (never history, never the prior proposal JSON),
under the existing ``max_attempts`` and meter caps. The two falsifiers are
recorded separately (§9): ``validator_decision`` mirrors the deterministic
validator only, stamped BEFORE any checker override.
"""

from __future__ import annotations

import json
import re
from dataclasses import dataclass
from typing import Literal, Protocol

from pydantic import ValidationError

from portfolio_optimiser_claude.budget import BudgetMeter
from portfolio_optimiser_claude.ir import SavingsProposal
from portfolio_optimiser_claude.provenance import stamp_validator_decision
from portfolio_optimiser_claude.validator import Rejection, ValidatedProposal, validate_proposal

_PROPOSER_ROLE = "proposer"
_CHECKER_ROLE = "checker"
# §3 Step 3: the checker's mandated closing line, both marker forms spelled out.
_CHECKER_INSTRUCTION = (
    "End your reply with exactly one verdict line: `VERDICT: APPROVE` if the "
    "reasoning holds, or `VERDICT: REJECT - <short reason>` if not."
)
# §3 Step 4: parsed case-insensitively; the reject marker takes precedence and
# its trailing text is the reason.
_REJECT_PATTERN = re.compile(r"VERDICT:\s*REJECT(?:\s*[-–—:]\s*(.*))?", re.IGNORECASE)
_APPROVE_PATTERN = re.compile(r"VERDICT:\s*APPROVE", re.IGNORECASE)

CheckerDecision = Literal["approve", "reject", "absent"]


@dataclass(frozen=True)
class ModelReply:
    """One model reply: text + provider-reported usage (+ real model id, §9)."""

    text: str
    usage_tokens: int | None = None
    model: str | None = None


class ModelClient(Protocol):
    """The injected model seam — SDK client on the run path, stand-in in tests."""

    def complete(self, prompt: str, *, role: str) -> ModelReply: ...


def _guarded_complete(
    client: ModelClient, prompt: str, *, role: str, meter: BudgetMeter
) -> ModelReply:
    """One model call, pre-guarded by the run-total USD belt (C3.5, §8).

    Every model call in the loop goes through here. The pre-call guard reads the
    client's accumulated ``total_cost_usd`` (0.0 for scripted clients that carry
    no cost) and refuses to make the call once the run-total USD cap is crossed
    — the belt that stops the loop from spending past its run budget, on TOP of
    the per-call SDK cap and the post-charge token/round meter.
    """
    spent_usd: float = getattr(client, "total_cost_usd", 0.0)
    meter.guard_before_call(spent_usd)
    return client.complete(prompt, role=role)


# --- Step 2: hypothesise (structured candidate generation) ---------------------------------


def _parse_candidate(text: str, default_project_id: str | None) -> SavingsProposal:
    raw = json.loads(text)
    if not isinstance(raw, dict):
        raise ValueError("candidate reply is not a JSON object")
    if default_project_id is not None:
        # §3 Step 2: project_id MAY be defaulted from the project when omitted.
        raw.setdefault("project_id", default_project_id)
    return SavingsProposal.model_validate(raw)


def generate_candidate(
    client: ModelClient,
    prompt: str,
    *,
    meter: BudgetMeter,
    default_project_id: str | None = None,
) -> SavingsProposal:
    """Ask for exactly one candidate as JSON for the IR — blind parse-retry (§3 Step 2).

    A reply that fails to parse into the typed IR is retried with the SAME
    prompt, never silently accepted or repaired downstream. The retry loop is
    bounded by the budget meter: a round tick is charged between attempts (§8).
    """
    while True:
        model_reply = _guarded_complete(client, prompt, role=_PROPOSER_ROLE, meter=meter)
        meter.charge_tokens(model_reply.usage_tokens)
        try:
            # JSONDecodeError and pydantic's ValidationError are ValueErrors.
            return _parse_candidate(model_reply.text, default_project_id)
        except (ValueError, ValidationError):
            meter.charge_round()


# --- Step 3: debate (maker-checker) ---------------------------------------------------------


@dataclass(frozen=True)
class DebateResult:
    """The debate's converged output + the checker's LAST surfaced reply (§3)."""

    proposer_output: str
    checker_last: str
    rounds: int


def check_turn_safety_net(turns: int, max_rounds: int) -> None:
    """The turn-count termination safety net ABOVE the round cap (§3 Step 3, §8).

    Under the ``range(max_rounds)``-bounded debate loop, ``turns`` never exceeds
    ``2 * max_rounds``, so this net is structurally unreachable — it is a
    deliberate belt that fires only if a refactor breaks the loop's own bound.
    """
    if turns > 2 * max_rounds + 2:
        raise RuntimeError(
            f"debate turn-count safety net tripped: {turns} turns with max_rounds={max_rounds}"
        )


def _proposer_debate_prompt(context: str, critique: str | None) -> str:
    prompt = (
        f"{context}\n\n"
        "Propose the reasoning for exactly one cost-saving candidate measure "
        "for this project."
    )
    if critique is not None:
        prompt += f"\n\nThe checker challenged your reasoning:\n{critique}\n\nAddress it."
    return prompt


def _checker_prompt(proposer_output: str) -> str:
    return (
        f"Check the proposer's reasoning for flaws:\n\n{proposer_output}\n\n{_CHECKER_INSTRUCTION}"
    )


def run_debate(
    client: ModelClient, context: str, *, max_rounds: int, meter: BudgetMeter
) -> DebateResult:
    """Alternate proposer/checker turns — round-capped, fresh state per run (§3 Step 3).

    Converges when the checker approves; otherwise the checker's critique
    feeds the next proposer turn until the round cap. All state is local to
    this call — nothing survives from one project run into the next.
    """
    if max_rounds <= 0:
        raise ValueError(f"max_rounds must be positive, got {max_rounds}")
    proposer_output = ""
    checker_last = ""
    critique: str | None = None
    rounds = 0
    turns = 0
    for _ in range(max_rounds):
        turns += 1
        check_turn_safety_net(turns, max_rounds)
        proposer_reply = _guarded_complete(
            client, _proposer_debate_prompt(context, critique), role=_PROPOSER_ROLE, meter=meter
        )
        meter.charge_tokens(proposer_reply.usage_tokens)
        proposer_output = proposer_reply.text

        turns += 1
        check_turn_safety_net(turns, max_rounds)
        checker_reply = _guarded_complete(
            client, _checker_prompt(proposer_output), role=_CHECKER_ROLE, meter=meter
        )
        meter.charge_tokens(checker_reply.usage_tokens)
        checker_last = checker_reply.text

        rounds += 1
        meter.charge_round()
        if parse_checker_verdict(checker_last).decision == "approve":
            break
        critique = checker_last
    return DebateResult(proposer_output=proposer_output, checker_last=checker_last, rounds=rounds)


# --- Step 4 (second falsifier): the checker gate --------------------------------------------


@dataclass(frozen=True)
class CheckerVerdict:
    """The parsed checker decision: approve / reject(+reason) / absent (§3 Step 4)."""

    decision: CheckerDecision
    reason: str | None


def parse_checker_verdict(text: str) -> CheckerVerdict:
    """Parse the verdict marker case-insensitively; REJECT takes precedence (§3 Step 4)."""
    reject = _REJECT_PATTERN.search(text)
    if reject:
        reason = (reject.group(1) or "").strip()
        return CheckerVerdict(decision="reject", reason=reason)
    if _APPROVE_PATTERN.search(text):
        return CheckerVerdict(decision="approve", reason=None)
    return CheckerVerdict(decision="absent", reason=None)


def apply_checker_gate(
    outcome: ValidatedProposal | Rejection, verdict: CheckerVerdict
) -> ValidatedProposal | Rejection:
    """Opt-in-reject gate (fail-open): only an explicit REJECT overrides (§3 Step 4).

    A validator rejection stands regardless of the checker; APPROVE or a
    missing/unparseable marker never blocks. An explicit REJECT turns an
    otherwise-validated outcome into a rejection whose reason is prefixed
    with the checker's reason.
    """
    if isinstance(outcome, Rejection):
        return outcome
    if verdict.decision != "reject":
        return outcome
    reason = verdict.reason or "checker rejected the reasoning"
    return Rejection(reason=f"{reason} (checker REJECT overrode a validated outcome)")


# --- Step 5: refine, informed and bounded ----------------------------------------------------


@dataclass(frozen=True)
class CandidateRun:
    """The refinement loop's result: last proposal, its outcome, attempts used."""

    proposal: SavingsProposal
    outcome: ValidatedProposal | Rejection
    attempts: int


def _informed_prompt(base_prompt: str, reason: str) -> str:
    # Only the most recent rejection REASON crosses attempts — never an
    # accumulated history, never the prior proposal JSON (§3 Step 5).
    return (
        f"{base_prompt}\n\n"
        "The previous attempt was rejected by the deterministic validator. "
        "Revise the candidate to address this falsification:\n"
        f"{reason}"
    )


def run_candidate_loop(
    client: ModelClient,
    base_prompt: str,
    *,
    meter: BudgetMeter,
    max_attempts: int = 3,
    default_project_id: str | None = None,
) -> CandidateRun:
    """Generate → validate, informed by the last rejection reason (§3 Steps 2+5).

    Attempt 1 uses the unchanged base prompt; each later attempt appends the
    previous rejection reason VERBATIM as a revision instruction. Bounded by
    ``max_attempts`` and the meter (a round tick between attempts, §8). The
    only per-attempt falsifier is the deterministic validator.
    """
    if max_attempts <= 0:
        raise ValueError(f"max_attempts must be positive, got {max_attempts}")
    reason: str | None = None
    for attempt in range(1, max_attempts + 1):
        prompt = base_prompt if reason is None else _informed_prompt(base_prompt, reason)
        proposal = generate_candidate(
            client, prompt, meter=meter, default_project_id=default_project_id
        )
        outcome = validate_proposal(proposal)
        if isinstance(outcome, ValidatedProposal):
            return CandidateRun(proposal=proposal, outcome=outcome, attempts=attempt)
        reason = outcome.reason
        if attempt < max_attempts:
            meter.charge_round()
    return CandidateRun(proposal=proposal, outcome=outcome, attempts=max_attempts)


# --- The run: debate → generation → both falsifiers, recorded separately --------------------


@dataclass(frozen=True)
class RunResult:
    """One project run (§3 Step 6): a validated proposal or a TYPED rejection.

    ``validator_decision`` mirrors the DETERMINISTIC VALIDATOR only, stamped
    before any checker override (§9); ``checker_decision`` is the second
    falsifier's own result field — the two are never conflated.
    """

    outcome: ValidatedProposal | Rejection
    validator_decision: Literal["validated", "rejected"]
    checker_decision: CheckerDecision
    attempts: int
    proposal: SavingsProposal


def _generation_prompt(context: str, converged_reasoning: str) -> str:
    # §3 Step 3: the debate's converged proposer output feeds generation.
    return (
        f"{context}\n\n"
        f"Converged reasoning from the maker-checker debate:\n{converged_reasoning}\n\n"
        "Reply with exactly one candidate measure as a JSON object with the "
        "fields: project_id, measure, affected_items (list of {code, quantity, "
        "unit_cost}), claimed_saving_nok, and optionally assumptions. "
        "Output ONLY the raw JSON object — no markdown fences, no commentary."
    )


def run_project(
    client: ModelClient,
    context: str,
    *,
    meter: BudgetMeter,
    max_debate_rounds: int,
    max_attempts: int = 3,
    default_project_id: str | None = None,
) -> RunResult:
    """Run steps 2–5 for one project: debate, generate, validate, gate (§3)."""
    debate = run_debate(client, context, max_rounds=max_debate_rounds, meter=meter)
    run = run_candidate_loop(
        client,
        _generation_prompt(context, debate.proposer_output),
        meter=meter,
        max_attempts=max_attempts,
        default_project_id=default_project_id,
    )
    verdict = parse_checker_verdict(debate.checker_last)
    # §9: stamped from the validator's outcome BEFORE the checker override.
    validator_decision = stamp_validator_decision(run.outcome)
    return RunResult(
        outcome=apply_checker_gate(run.outcome, verdict),
        validator_decision=validator_decision,
        checker_decision=verdict.decision,
        attempts=run.attempts,
        proposal=run.proposal,
    )