portfolio-optimiser/src/portfolio_optimiser/generate.py

"""LLM->IR generation wired to validator-as-retry (B1 + research 03 Dim 3).

A NON-STREAMING chat call asks the model for a structured ``SavingsProposal``; the reply is
parsed into the typed IR. Small local models leak text or emit wrong-typed tool calls
(research 03 Dim 3), so a malformed reply is RETRIED — never silently accepted. The
deterministic validator is the reliability mechanism.

The token bound lives HERE, in the generate loop (``meter`` checked between attempts) — so
``validator.py`` stays the verbatim Step-2 module (it is in this step's ``forbidden_paths``).

Two entry points, because the LLM call is async while ``validator.self_repair`` is sync:

* ``generate_with_validation`` — the SYNC validator-as-retry primitive: it drives
  ``validator.self_repair`` over a sync candidate source and adds the meter bound between
  attempts. Used for deterministic candidate sources.
* ``generate_via_llm`` — the ASYNC LLM path: an async mirror of the same bounded retry that
  awaits the chat call (parse-retry inside the meter budget, then ``validate_proposal``).
  Returns ``ValidatedProposal | Rejection``; never a malformed proposal; raises
  ``BudgetExceeded`` when the meter cap is crossed.
"""

from __future__ import annotations

import json
from collections.abc import Callable

from agent_framework import BaseChatClient, Message
from pydantic import ValidationError

from portfolio_optimiser.budget import TokenMeter
from portfolio_optimiser.ir import SavingsProposal
from portfolio_optimiser.reference_domain import Project
from portfolio_optimiser.validator import (
    Rejection,
    ValidatedProposal,
    self_repair,
    validate_proposal,
)


class GenerationError(RuntimeError):
    """No parseable proposal could be produced within the attempt budget."""


def _build_messages(
    project: Project, context: str, prior_rejection: Rejection | None = None
) -> list[Message]:
    """Build the hypothesis prompt. When ``prior_rejection`` is set (Step 5, målbilde §5/§7),
    append a revision block carrying ONLY the falsification *reason* verbatim — never the prior
    proposal JSON (minimal honest payload: the model must address the falsification, not parrot
    the rejected candidate back). ``None`` -> the byte-identical base prompt, so attempt 1 is
    unchanged. The reason carries only the rejected claim/feasible figures, which deliberately
    do not collide with other load-bearing prompt markers."""
    prompt = (
        "Propose ONE concrete cost-saving measure for this project.\n"
        f"Project: {project.id} - {project.name}\n"
        f"Context (prior verdicts / cited cost docs):\n{context}\n\n"
        "Respond with ONLY a JSON object for a SavingsProposal with keys: project_id, "
        "measure, affected_items (list of {code, quantity, unit_cost}), claimed_saving_nok, "
        "and optional assumptions."
    )
    if prior_rejection is not None:
        prompt += (
            "\n\nYour previous proposal was REJECTED by the deterministic validator.\n"
            f"Reason: {prior_rejection.reason}\n"
            "Produce a REVISED SavingsProposal that resolves this."
        )
    return [Message(role="user", contents=[prompt])]


def _parse_ir(text: str, project: Project) -> SavingsProposal:
    """Parse the model's structured reply into the typed IR. Raises on malformed/text-leaked
    output (JSON error or Pydantic ``ValidationError``)."""
    data = json.loads(text)
    if not isinstance(data, dict):
        raise ValueError("reply is not a JSON object")
    data.setdefault("project_id", project.id)
    return SavingsProposal(**data)


def _charge_usage(meter: TokenMeter, reply: object) -> None:
    usage = getattr(reply, "usage_details", None)
    total = usage.get("total_token_count") if usage else None
    if total:
        meter.charge(int(total))  # raises BudgetExceeded over cap


def generate_with_validation(
    make_proposal: Callable[[int], SavingsProposal],
    meter: TokenMeter,
    *,
    max_attempts: int = 3,
) -> ValidatedProposal | Rejection:
    """Sync validator-as-retry: drive ``validator.self_repair`` over a sync candidate source,
    checking the token meter between attempts (the token bound lives HERE, never in
    ``validator.py``). Returns ``ValidatedProposal | Rejection``; raises ``BudgetExceeded`` on
    a meter cap."""

    def _attempt(attempt: int) -> SavingsProposal:
        meter.tick_round()  # between-attempt iteration bound (BudgetExceeded over cap)
        return make_proposal(attempt)

    return self_repair(_attempt, max_attempts=max_attempts)


async def generate_via_llm(
    chat_client: BaseChatClient,
    project: Project,
    context: str,
    meter: TokenMeter,
    *,
    max_attempts: int = 3,
) -> ValidatedProposal | Rejection:
    """Async LLM path: non-streaming chat -> parse -> validate, with TWO bounded retry kinds,
    the meter checked in this loop:

    * malformed/text-leaked reply -> BLIND parse-retry (inner loop): re-fetch until the reply
      parses; never silently accepted.
    * validator rejection -> INFORMED refinement (outer ``max_attempts`` loop; Step 5,
      målbilde §5/§7): the previous attempt's ``Rejection.reason`` is fed into the next
      attempt's prompt (``_build_messages(prior_rejection=...)``) so the proposer can correct
      rather than re-answer blindly. Bounded by ``max_attempts`` + the meter (no new loop; §6).

    The only per-attempt falsifier here is the deterministic validator (the numbers). The
    checker is a run-level, one-shot signal (run.py, before generation); seeding generation
    with the checker critique is separately scoped and NOT done here. Returns
    ``ValidatedProposal | Rejection``; never a malformed proposal; raises ``BudgetExceeded``
    when the meter cap is crossed."""

    async def _fetch_parsed(messages: list[Message]) -> SavingsProposal:
        # Parse-robust: a malformed/text-leaked reply is retried; the meter caps total work.
        while True:
            meter.tick_round()  # between-attempt bound (BudgetExceeded over cap)
            reply = await chat_client.get_response(messages)  # non-streaming
            _charge_usage(meter, reply)
            try:
                return _parse_ir(reply.text, project)
            except (ValidationError, ValueError, TypeError):
                continue

    last: Rejection | None = None
    for _ in range(max_attempts):
        # Informed refinement: feed the PREVIOUS attempt's validator rejection into this
        # attempt's prompt. ``last`` is None on attempt 1 -> the unchanged base prompt; it is
        # overwritten each round -> only the most-recent falsification ("forrige"), never an
        # accumulated history (bounded prompt growth).
        messages = _build_messages(project, context, prior_rejection=last)
        candidate = await _fetch_parsed(messages)
        result = validate_proposal(candidate)
        if isinstance(result, ValidatedProposal):
            return result
        last = result
    assert last is not None  # max_attempts >= 1, so at least one validation ran
    return last  # validation never passed within the attempt budget -> typed Rejection