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 a ``GenerationResult`` (the outcome PLUS the falsifications that informed it); never a
  malformed proposal; raises ``BudgetExceeded`` when the meter cap is crossed.
"""

from __future__ import annotations

import json
from collections.abc import Callable, Mapping
from dataclasses import dataclass, field
from typing import Any

from agent_framework import BaseChatClient, Message
from pydantic import BaseModel, ValidationError

from portfolio_optimiser.budget import TokenMeter
from portfolio_optimiser.ir import CostBaseline, SavingsProposal
from portfolio_optimiser.mandate import Approach
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."""


class StructuredOutputUnsupported(TypeError):
    """A schema node cannot be expressed in the provider's strict structured-output subset.

    Fail-closed, and deliberately so (mirrors ``write_concept_file`` / ``promote_verdict``:
    validation, never repair). The alternative — silently dropping what cannot be expressed — would
    stop commissioning a field without saying so, and the field it would have dropped first is
    ``assumptions``, whose absence makes the Monte Carlo falsifier inert while it still reports
    percentiles. A schema this module cannot express is a decision for a human, not a default.
    """


#: Type-specific JSON Schema keywords the provider's structured-output subset does NOT support,
#: transcribed from Azure's published table (Structured outputs -> "Unsupported type-specific
#: keywords", https://learn.microsoft.com/azure/foundry/openai/how-to/structured-outputs), which
#: states it is the same subset OpenAI accepts.
#:
#: ``exclusiveMinimum``/``exclusiveMaximum`` are NOT literally in that table — it names
#: ``minimum maximum multipleOf`` — but they are the same family, and pydantic emits them for
#: ``Field(gt=...)``/``Field(lt=...)``, which is exactly how this repo's IR spells its bounds. Being
#: stricter than the table costs nothing here: every constraint stripped is re-applied by pydantic in
#: ``_parse_ir`` and by ``validate_proposal``. The schema's job is SHAPE; the validator's job is
#: VALUES. ``default`` is stripped for a different reason — strict mode requires every property to be
#: required, so a default can never apply.
UNSUPPORTED_SCHEMA_KEYWORDS = frozenset(
    {
        # String
        "minLength",
        "maxLength",
        "pattern",
        "format",
        # Number
        "minimum",
        "maximum",
        "multipleOf",
        "exclusiveMinimum",
        "exclusiveMaximum",
        # Objects
        "patternProperties",
        "unevaluatedProperties",
        "propertyNames",
        "minProperties",
        "maxProperties",
        # Arrays
        "unevaluatedItems",
        "contains",
        "minContains",
        "maxContains",
        "minItems",
        "maxItems",
        "uniqueItems",
        # Meaningless once every property is required
        "default",
    }
)

#: The strict-legal stand-in for ``SavingsProposal.assumptions``.
#:
#: The IR spells the uncertainty bands as ``dict[str, tuple[float, float]]`` — a free-form map whose
#: values are tuples. Neither half is expressible: strict mode requires ``additionalProperties:
#: false`` in every object (so a map with arbitrary keys cannot be described), and tuples arrive as
#: ``prefixItems``, which is outside the supported type list. Dropping the field instead would be
#: silent damage: ``validator._monte_carlo`` falls back to the item's stated ``unit_cost`` for every
#: code with no band, so with no bands at all the samples are identical and P10 == P50 == P90 — the
#: stochastic falsifier goes inert while still reporting percentiles.
#:
#: So the WIRE carries an array of named entries and ``_parse_ir`` folds it back into the IR's map.
#: The IR itself is untouched; the entry names spell out what the tuple positions mean, which the
#: model would otherwise have to guess.
_ASSUMPTIONS_WIRE_NODE: dict[str, Any] = {
    "type": "array",
    "description": (
        "Uncertainty band per affected cost line: the low and high unit cost the true price is "
        "expected to fall between. The band MUST enclose that item's own unit_cost. Omit an entry "
        "for a line whose unit cost is certain; an empty list means no uncertainty is claimed."
    ),
    "items": {
        "type": "object",
        "properties": {
            "code": {"type": "string"},
            "low_unit_cost": {"type": "number"},
            "high_unit_cost": {"type": "number"},
        },
    },
}

#: Dotted paths (from the root model's own properties) whose node is replaced before sanitising.
_PROPOSAL_SCHEMA_OVERRIDES: Mapping[str, dict[str, Any]] = {"assumptions": _ASSUMPTIONS_WIRE_NODE}


def _sanitise_schema_node(node: Any, *, path: str, overrides: Mapping[str, dict[str, Any]]) -> Any:
    """Rewrite one JSON Schema node into the strict subset, or raise ``StructuredOutputUnsupported``.

    An override is applied FIRST, so a declared replacement is what gets checked and emitted — that
    is how the one inexpressible node in this repo's IR (``assumptions``) is expressed rather than
    excused. The replacement is then sanitised by the same code as everything else, so an override
    cannot smuggle in an illegal node.
    """
    if not isinstance(node, Mapping):
        return node
    if path in overrides:
        node = overrides[path]

    if "prefixItems" in node:
        raise StructuredOutputUnsupported(
            f"{path or '<root>'}: tuple types (prefixItems) are outside the strict subset"
        )
    for combinator in ("oneOf", "allOf"):
        if combinator in node:
            raise StructuredOutputUnsupported(
                f"{path or '<root>'}: {combinator} is outside the strict subset (anyOf is the "
                "only supported combinator)"
            )
    if isinstance(node.get("additionalProperties"), Mapping):
        raise StructuredOutputUnsupported(
            f"{path or '<root>'}: a free-form map cannot be expressed — strict mode requires "
            "additionalProperties: false in every object. Declare an override that spells the "
            "entries out as an array."
        )

    out: dict[str, Any] = {}
    for key, value in node.items():
        if key in UNSUPPORTED_SCHEMA_KEYWORDS:
            continue
        if key == "properties" and isinstance(value, Mapping):
            out[key] = {
                name: _sanitise_schema_node(
                    sub, path=f"{path}.{name}" if path else name, overrides=overrides
                )
                for name, sub in value.items()
            }
        elif key == "$defs" and isinstance(value, Mapping):
            out[key] = {
                name: _sanitise_schema_node(sub, path=f"$defs.{name}", overrides=overrides)
                for name, sub in value.items()
            }
        elif key == "items":
            out[key] = _sanitise_schema_node(value, path=f"{path}[]", overrides=overrides)
        elif key == "anyOf" and isinstance(value, list):
            out[key] = [_sanitise_schema_node(sub, path=path, overrides=overrides) for sub in value]
        else:
            out[key] = value

    if "properties" in out:
        # Strict mode's two structural demands, applied to EVERY object rather than the root only:
        # no undeclared keys, and every declared key required.
        out["additionalProperties"] = False
        out["required"] = sorted(out["properties"])
    return out


def strict_json_schema(
    model: type[BaseModel], *, overrides: Mapping[str, dict[str, Any]] | None = None
) -> dict[str, Any]:
    """Derive a strict-structured-output schema from ``model``'s own pydantic schema.

    DERIVED rather than hand-written on purpose: a hand-written copy of a shape that already exists
    in ``ir.py`` is the second copy that drifts (kø-(p)), and it drifts silently — the model would
    keep being commissioned for the old shape. ``$defs``/``$ref`` are kept (the published subset
    supports definitions), so nested models need no inlining.
    """
    schema = _sanitise_schema_node(model.model_json_schema(), path="", overrides=overrides or {})
    assert isinstance(schema, dict)  # a model's root schema is always an object
    return schema


def proposal_response_format() -> dict[str, Any]:
    """The ``response_format`` mapping commissioning a ``SavingsProposal`` from the proposer.

    A MAPPING, not the ``type[BaseModel]`` the option also accepts, and the reason is measured: given
    a class, the client converts it with ``type_to_response_format_param``, which emits ``minimum`` /
    ``exclusiveMinimum`` / ``minItems`` / ``prefixItems`` and an ``assumptions`` node whose
    ``additionalProperties`` is a schema — four things the published subset rules out. Our own
    mapping is the only way to control what reaches the wire.

    ONE mapping serves both wired profiles (measured against agent-framework-openai 1.8.2 /
    agent-framework-foundry 1.8.2): the Chat Completions client passes it through verbatim, and the
    Responses client — which ``FoundryChatClient`` delegates to — converts this exact envelope into
    ``text.format``.
    """
    return {
        "type": "json_schema",
        "json_schema": {
            "name": SavingsProposal.__name__,
            "strict": True,
            "schema": strict_json_schema(SavingsProposal, overrides=_PROPOSAL_SCHEMA_OVERRIDES),
        },
    }


@dataclass(frozen=True)
class ParseFailure:
    """One model reply that did NOT parse into the typed IR, kept VERBATIM (Fase 1b, funn 1).

    ``text`` is the reply exactly as the model produced it — never truncated, stripped or
    summarised. It is the thing the run PAID for and the only evidence of *why* the reply did not
    parse; a paraphrase would make the next paid run a guess again, which is the defect this type
    exists to close. ``error`` names the parse error itself (``json.JSONDecodeError`` vs a pydantic
    ``ValidationError`` are very different diagnoses: leaked prose vs a wrong-shaped object).

    Collected into a CALLER-OWNED sink rather than returned — see ``generate_via_llm``.
    """

    text: str
    error: str


@dataclass(frozen=True)
class GenerationResult:
    """What one ``generate_via_llm`` call produced: the outcome, and the falsification history that
    informed it (Step 5, målbilde §5/§7).

    A TYPED RETURN VALUE rather than an out-parameter or a callback, deliberately: the informed
    refinement loop already computed this history internally and then dropped it, so Step 5 was the
    one step of the eight with no observable output. A returned value cannot be silently lost by a
    caller that forgets to pass a collector, and it forces every call site to acknowledge the seam.

    ``refinements`` holds ONLY the rejections that were actually fed back into a later attempt's
    prompt — the honest reading of "informed refinement". When the attempt budget runs out, the
    final rejection IS ``outcome``: it informed nothing and is not repeated here. So the total
    number of validator falsifications this call produced is ``len(refinements)`` plus one when
    ``outcome`` is itself a ``Rejection``. It is empty on the common single-attempt path, which is
    honest rather than merely convenient: nothing was falsified, so there is nothing to show.
    """

    outcome: ValidatedProposal | Rejection
    refinements: tuple[Rejection, ...] = field(default=())


def _build_messages(
    project: Project,
    context: str,
    prior_rejection: Rejection | None = None,
    *,
    approach: Approach | 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.

    When ``approach`` is set (Trekk A3, krav 1) the opening instruction switches from *find one*
    to *quantify THIS one*: the domain expert has already decided what shall be evaluated, and the
    model's job is the numbers, not the direction. The expert's ``label`` and ``description`` are
    carried VERBATIM — the description is the reason the approach is worth trying, which is
    exactly the part the model cannot infer from the cost data. ``None`` -> the byte-identical
    base prompt, so an un-commissioned run is untouched (mirrors ``prior_rejection``).

    The two are composable: a commissioned approach that the validator rejects is refined through
    the SAME informed-refinement block, still bound to that approach.
    """
    if approach is None:
        head = "Propose ONE concrete cost-saving measure for this project.\n"
    else:
        head = (
            "A domain expert has commissioned ONE specific approach for this project. "
            "Quantify THAT approach as a concrete cost-saving measure — do not substitute a "
            "different measure. If it does not apply to this project, say so through the "
            "numbers rather than proposing something else.\n"
            f"Approach: {approach.label}\n"
        )
        if approach.description:
            head += f"Why the expert wants it evaluated: {approach.description}\n"
    prompt = (
        f"{head}"
        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 _normalise_assumptions(data: dict[str, Any]) -> None:
    """Fold the WIRE's array-of-entries assumption bands back into the IR's ``code -> (low, high)``
    map, in place.

    ADDITIVE, never a replacement: a reply that already uses the IR's map form (every scripted reply
    in the suite, and any model that answers without honouring the schema) is left untouched. A
    malformed entry is raised as ``ValueError`` rather than ``KeyError`` on purpose — ``ValueError``
    is what ``_fetch_parsed`` catches, so a bad band is captured as the parse failure it is instead
    of escaping the loop and killing the run.
    """
    entries = data.get("assumptions")
    if not isinstance(entries, list):
        return
    bands: dict[str, tuple[Any, Any]] = {}
    for entry in entries:
        if (
            not isinstance(entry, Mapping)
            or not {
                "code",
                "low_unit_cost",
                "high_unit_cost",
            }
            <= entry.keys()
        ):
            raise ValueError(
                f"each assumption entry needs code, low_unit_cost and high_unit_cost; got {entry!r}"
            )
        bands[entry["code"]] = (entry["low_unit_cost"], entry["high_unit_cost"])
    data["assumptions"] = bands


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)
    _normalise_assumptions(data)
    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,
    baseline: CostBaseline | None = None,
    approach: Approach | None = None,
    parse_failures: list[ParseFailure] | None = None,
) -> GenerationResult:
    """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.

    ``approach`` (Trekk A3, krav 1) binds every attempt of this call to ONE expert-commissioned
    approach. It changes only the prompt: ``validate_proposal`` is called exactly as before, so a
    commissioned approach gets **no discount at the deterministic gate** — the expert directs what
    is evaluated, never what is approved. A commissioned proposal that fails is refined through the
    same informed-refinement path, still bound to that approach, and returns a typed ``Rejection``
    when the attempt budget runs out.

    ``baseline`` (S4.0) is handed straight to ``validate_proposal``, so a fabricated cost line is
    falsified per ATTEMPT like any other rejection — and its reason feeds the next attempt's prompt
    through the SAME informed-refinement path (Step 5), which is why no new loop appears here.

    ``parse_failures`` (Fase 1b, funn 1) is a CALLER-OWNED sink: every reply that fails to parse is
    appended to it VERBATIM, at the moment it fails. It is an out-parameter and not part of the
    return value ON PURPOSE, and the reason is measured rather than stylistic. ``meter.tick_round``
    raises ``BudgetExceeded`` inside the inner fetch loop, so on the path this capture exists for —
    a model whose replies never parse, which burns the round ledger — this function raises and
    returns NOTHING. That is exactly the live Fase-1b failure. A field on ``GenerationResult`` (the
    Step-5 ``refinements`` shape) would be blind to it, as would any artefact written by the caller
    *after* a successful return. The sink mirrors ``meter`` instead: a caller-owned accumulator this
    loop mutates, whose contents the caller still holds however the loop ended. Step 5's "a returned
    value cannot be silently lost by a caller that forgets to pass a collector" governs a value that
    REACHES the caller; here it does not, so the rule is cited and departed from deliberately. That
    a caller can forget is answered by a test on the wiring, not by a shape that cannot work.

    Returns a ``GenerationResult``: the ``ValidatedProposal | Rejection`` outcome plus every
    rejection that was fed back into a later attempt's prompt. Surfacing that history changes
    nothing about the loop's BOUND — ``max_attempts`` and ``meter.tick_round`` are exactly as
    before ("refine until good enough" without a cap stays forbidden, §6); it only stops the loop
    from discarding what it already knew. 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)
            # Fase 1b, funn 1b: hand the model a GRAMMAR, not a prose request. The prompt's
            # "Respond with ONLY a JSON object" line stays — a provider that ignores
            # ``response_format`` (or a local model that does not implement it) must still be told
            # what is wanted, and the parse-retry below remains the backstop either way.
            reply = await chat_client.get_response(  # non-streaming
                messages, options={"response_format": proposal_response_format()}
            )
            _charge_usage(meter, reply)
            try:
                return _parse_ir(reply.text, project)
            except (ValidationError, ValueError, TypeError) as exc:
                # Capture BEFORE the retry: this reply was paid for, and once ``continue`` runs the
                # only record of what the model actually said is gone (Fase 1b, funn 1). Verbatim —
                # the operator is diagnosing a format failure, so any shortening removes evidence.
                if parse_failures is not None:
                    parse_failures.append(
                        ParseFailure(text=reply.text, error=f"{type(exc).__name__}: {exc}")
                    )
                continue

    last: Rejection | None = None
    # The falsifications that were FED BACK, in attempt order. ``last`` still drives the PROMPT and
    # is still overwritten each round -- only the most-recent falsification reaches the model, so
    # prompt growth is unchanged. This list is a record for the CALLER, appended to only once a
    # rejection is about to inform a further attempt; it is never read back into a prompt.
    fed_back: list[Rejection] = []
    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).
        if last is not None:
            fed_back.append(last)
        messages = _build_messages(project, context, prior_rejection=last, approach=approach)
        candidate = await _fetch_parsed(messages)
        result = validate_proposal(candidate, baseline=baseline)
        if isinstance(result, ValidatedProposal):
            return GenerationResult(outcome=result, refinements=tuple(fed_back))
        last = result
    assert last is not None  # max_attempts >= 1, so at least one validation ran
    # Validation never passed within the attempt budget -> typed Rejection. ``last`` is the outcome
    # and was never fed back, so it is deliberately absent from ``refinements``.
    return GenerationResult(outcome=last, refinements=tuple(fed_back))