portfolio-optimiser/src/portfolio_optimiser/run.py

"""Vertical-slice orchestrator + single-command entry (two-layer HITL wiring).

``run_project`` composes the whole method for ONE synthetic project on real (or injected)
chat clients:

1. ``load_contracts`` — fail-fast: every config (incl. the verdict-feedback shape) is
   validated BEFORE any chat client is built.
2. load the project + retrieve cited chunks via the Step-7 data source -> first-class
   ``provenance.Citation`` list.
3. a FRESH maker-checker ``GroupChat`` debate (``fresh_workflow``), round-capped
   (``with_max_rounds``); **Layer-1 HITL** = the optional in-run ``with_request_info`` gate.
4. ``generate_via_llm`` -> blocking ``validate_proposal`` -> ``ValidatedProposal | Rejection``
   (the token bound is the ``meter`` checked in the generate loop).
5. attach a first-class ``ProvenanceStamp``.
6. **Layer-2 (out-of-band)**: ``capture_verdict`` mints a stable id from the proposal's
   features; the decision/rationale come from ``verdict_input`` (function arg / CLI / fixture).
   The B11 expert notification is a STUB (``notify``) in Fase 2.
7. persist to the ``VerdictStore`` -> the next run's ``ExpeLContextProvider`` retrieval (the
   learning loop; this run also exercises the two-arg ``extend_instructions`` injection).

The two HITL layers are deliberately distinct: **Layer-1** is the optional synchronous in-run
review gate (no checkpoint — research 01: durable resume is fragile); **Layer-2** is the
durable learned verdict captured out-of-band in the VerdictStore (D7-portable).
"""

from __future__ import annotations

from collections.abc import Callable
from dataclasses import dataclass

from agent_framework import BaseChatClient, SessionContext

from portfolio_optimiser.backends import Profile, get_backend, resolve_model
from portfolio_optimiser.budget import Budget, BudgetMiddleware, TokenMeter
from portfolio_optimiser.contracts import load_contracts
from portfolio_optimiser.datasource import (
    chunk_dict_to_citation,
    make_retrieval_tool,
    retrieve_chunks,
)
from portfolio_optimiser.generate import generate_via_llm
from portfolio_optimiser.ir import SavingsProposal
from portfolio_optimiser.provenance import ProvenanceStamp
from portfolio_optimiser.reference_domain import Project, load_reference_projects
from portfolio_optimiser.validator import Rejection, ValidatedProposal
from portfolio_optimiser.verdicts import (
    ExpeLContextProvider,
    ProposalFeatures,
    Verdict,
    VerdictStore,
    capture_verdict,
)
from portfolio_optimiser.workflow import fresh_workflow


@dataclass(frozen=True)
class RunResult:
    """The outcome of one project run: the validated/rejected proposal, its first-class
    provenance, the captured (Layer-2) verdict, the ExpeL hits surfaced for it, and the store."""

    outcome: ValidatedProposal | Rejection
    provenance: ProvenanceStamp
    verdict: Verdict
    retrieved: list[Verdict]
    store: VerdictStore


def _project_by_id(project_id: str) -> Project:
    for project in load_reference_projects():
        if project.id == project_id:
            return project
    raise ValueError(f"unknown project_id: {project_id!r}")


def _features_of(proposal: SavingsProposal) -> ProposalFeatures:
    return ProposalFeatures(
        affected_codes=frozenset(item.code for item in proposal.affected_items),
        measure_type=proposal.measure,
        claimed_saving_nok=proposal.claimed_saving_nok,
        description=proposal.measure,
    )


def _default_factory(profile: Profile | str) -> Callable[[str], BaseChatClient]:
    def factory(role: str) -> BaseChatClient:
        return get_backend(profile).create_chat_client(model=resolve_model(profile, role))

    return factory


async def run_project(
    project_id: str,
    profile: Profile | str = Profile.LOCAL,
    *,
    docs_dir: str,
    verdict_input: dict[str, str],
    store: VerdictStore | None = None,
    client_factory: Callable[[str], BaseChatClient] | None = None,
    max_rounds: int = 3,
    max_tokens: int = 100_000,
    top_k: int = 3,
    enable_layer1_hitl: bool = False,
    notify: Callable[[Verdict], None] | None = None,
) -> RunResult:
    """Run the vertical slice for ONE project. ``client_factory`` is the test-injection seam
    (defaults to the real backend). ``verdict_input`` carries the expert decision/rationale
    (Layer-2). Raises ``pydantic.ValidationError`` on a bad contract and ``BudgetExceeded``
    when the token/round cap is crossed."""
    # 1. Fail-fast: validate ALL contracts (incl. the verdict-feedback shape) before any client.
    load_contracts(
        {"docs_dir": docs_dir, "top_k": top_k},
        {"max_rounds": max_rounds, "max_tokens": max_tokens},
        verdict_input,
    )

    # 2-3. Project + cited chunks (first-class provenance citations).
    project = _project_by_id(project_id)
    chunks = retrieve_chunks("cost saving measure", docs_dir, top_k)
    citations = [chunk_dict_to_citation(c) for c in chunks]
    if not citations:
        raise ValueError(f"no citable content in docs_dir: {docs_dir!r}")
    context = "\n".join(c["snippet"] for c in chunks)

    # 4. Budget + maker-checker debate (round-capped; Layer-1 HITL optional).
    # The shared meter is driven on the debate's chat calls by the BudgetMiddleware
    # (the brief's named short-circuit mechanism); the retrieval tool exposes the
    # citation-bearing data source to the agents (no longer string-stuffed out-of-band).
    meter = TokenMeter(Budget(max_tokens=max_tokens, max_rounds=max(max_rounds * 4, 4)))
    factory = client_factory if client_factory is not None else _default_factory(profile)
    budget_mw = BudgetMiddleware(meter)
    retrieval_tool = make_retrieval_tool(docs_dir, top_k=top_k)
    debate = fresh_workflow(
        factory,
        max_rounds=max_rounds,
        enable_layer1_hitl=enable_layer1_hitl,
        tools=[retrieval_tool],
        middleware=[budget_mw],
    )
    await debate.run(f"Find a cost-saving measure for {project.id}.\nContext:\n{context}")

    # 5. Structured candidate -> blocking validation; token bound = the meter in this loop.
    outcome = await generate_via_llm(factory("proposer"), project, context, meter)
    proposal = outcome.proposal

    # 6. First-class provenance stamp (authoritative; independent of MAF Annotation).
    model = "fake-model" if client_factory is not None else resolve_model(profile, "proposer")
    stamp = ProvenanceStamp(
        citations=citations,
        model=model,
        role="proposer",
        validator_decision="validated" if isinstance(outcome, ValidatedProposal) else "rejected",
        token_usage=meter.tokens,
    )

    # 7. ExpeL: surface prior verdicts for this proposal (exercises the two-arg
    #    extend_instructions injection on a real SessionContext — the learning loop).
    store = store if store is not None else VerdictStore(verdicts=[])
    features = _features_of(proposal)
    provider = ExpeLContextProvider(store, features, k=top_k)
    sctx = SessionContext(input_messages=[], instructions=[])
    await provider.before_run(agent=None, session=None, context=sctx, state={})
    retrieved = store.retrieve(features, k=top_k) if store.verdicts else []

    # 8. Layer-2 (out-of-band): capture the durable verdict + persist; B11 notify is a stub.
    verdict = capture_verdict(features, verdict_input["decision"], verdict_input["rationale"])
    store.add(verdict)
    if notify is not None:
        notify(verdict)

    return RunResult(
        outcome=outcome, provenance=stamp, verdict=verdict, retrieved=retrieved, store=store
    )


def main(argv: list[str] | None = None) -> int:
    """Single-command console entry: run the slice for one project against a docs folder."""
    import argparse
    import asyncio

    parser = argparse.ArgumentParser(description="portfolio-optimiser vertical slice")
    parser.add_argument("project_id")
    parser.add_argument("--profile", default="local")
    parser.add_argument("--docs-dir", required=True)
    parser.add_argument("--decision", default="approved", choices=["approved", "rejected"])
    parser.add_argument("--rationale", default="reviewed by expert")
    args = parser.parse_args(argv)

    result = asyncio.run(
        run_project(
            args.project_id,
            args.profile,
            docs_dir=args.docs_dir,
            verdict_input={"decision": args.decision, "rationale": args.rationale},
        )
    )
    kind = type(result.outcome).__name__
    print(f"{args.project_id}: {kind} (verdict id={result.verdict.id}, decision={args.decision})")
    return 0


if __name__ == "__main__":  # pragma: no cover - console entry
    raise SystemExit(main())