From 7573c4439f4980be181b14532dade101e16a5bed Mon Sep 17 00:00:00 2001 From: Kjell Tore Guttormsen Date: Fri, 26 Jun 2026 00:37:45 +0200 Subject: [PATCH] feat(fase2): wire BudgetMiddleware + retrieval tool onto the debate in run_project --- src/portfolio_optimiser/run.py | 21 ++++++-- tests/test_vertical_slice_e2e.py | 90 ++++++++++++++++++++++++++------ 2 files changed, 92 insertions(+), 19 deletions(-) diff --git a/src/portfolio_optimiser/run.py b/src/portfolio_optimiser/run.py index 3ecb6ed..f3f75f5 100644 --- a/src/portfolio_optimiser/run.py +++ b/src/portfolio_optimiser/run.py @@ -31,9 +31,13 @@ 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, TokenMeter +from portfolio_optimiser.budget import Budget, BudgetMiddleware, TokenMeter from portfolio_optimiser.contracts import load_contracts -from portfolio_optimiser.datasource import chunk_dict_to_citation, retrieve_chunks +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 @@ -118,9 +122,20 @@ async def run_project( 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) - debate = fresh_workflow(factory, max_rounds=max_rounds, enable_layer1_hitl=enable_layer1_hitl) + 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. diff --git a/tests/test_vertical_slice_e2e.py b/tests/test_vertical_slice_e2e.py index bc74d05..d60b1ae 100644 --- a/tests/test_vertical_slice_e2e.py +++ b/tests/test_vertical_slice_e2e.py @@ -27,8 +27,12 @@ _VI = {"decision": "approved", "rationale": "feasible within range"} async def test_a_valid_proposal_end_to_end(docs_dir, make_client_factory, fresh_store) -> None: result = await run_project( - "FV42-GSV-E1", "local", docs_dir=docs_dir, verdict_input=_VI, - client_factory=make_client_factory(_VALID), store=fresh_store, + "FV42-GSV-E1", + "local", + docs_dir=docs_dir, + verdict_input=_VI, + client_factory=make_client_factory(_VALID), + store=fresh_store, ) assert isinstance(result.outcome, ValidatedProposal) stamp = result.provenance @@ -39,9 +43,12 @@ async def test_a_valid_proposal_end_to_end(docs_dir, make_client_factory, fresh_ async def test_b_out_of_range_is_rejected(docs_dir, make_client_factory, fresh_store) -> None: result = await run_project( - "FV42-GSV-E1", "local", docs_dir=docs_dir, + "FV42-GSV-E1", + "local", + docs_dir=docs_dir, verdict_input={"decision": "rejected", "rationale": "claim too high"}, - client_factory=make_client_factory(_OUT_OF_RANGE), store=fresh_store, + client_factory=make_client_factory(_OUT_OF_RANGE), + store=fresh_store, ) assert isinstance(result.outcome, Rejection) assert result.outcome.reason @@ -49,8 +56,12 @@ async def test_b_out_of_range_is_rejected(docs_dir, make_client_factory, fresh_s async def test_c_layer2_verdict_is_persisted(docs_dir, make_client_factory, fresh_store) -> None: result = await run_project( - "FV42-GSV-E1", "local", docs_dir=docs_dir, verdict_input=_VI, - client_factory=make_client_factory(_VALID), store=fresh_store, + "FV42-GSV-E1", + "local", + docs_dir=docs_dir, + verdict_input=_VI, + client_factory=make_client_factory(_VALID), + store=fresh_store, ) assert any(v.id == result.verdict.id for v in result.store.verdicts) @@ -60,12 +71,20 @@ async def test_d_second_run_retrieves_prior_verdict( ) -> None: factory = make_client_factory(_VALID) r1 = await run_project( - "FV42-GSV-E1", "local", docs_dir=docs_dir, verdict_input=_VI, - client_factory=factory, store=fresh_store, + "FV42-GSV-E1", + "local", + docs_dir=docs_dir, + verdict_input=_VI, + client_factory=factory, + store=fresh_store, ) r2 = await run_project( - "FV42-GSV-E1", "local", docs_dir=docs_dir, verdict_input=_VI, - client_factory=factory, store=fresh_store, + "FV42-GSV-E1", + "local", + docs_dir=docs_dir, + verdict_input=_VI, + client_factory=factory, + store=fresh_store, ) assert r2.retrieved # ExpeL surfaced prior verdicts on the second run assert r2.retrieved[0].id == r1.verdict.id # structurally identical -> stable id @@ -76,15 +95,50 @@ async def test_e_tiny_budget_halts_without_exceeding( ) -> None: with pytest.raises(BudgetExceeded): await run_project( - "FV42-GSV-E1", "local", docs_dir=docs_dir, verdict_input=_VI, + "FV42-GSV-E1", + "local", + docs_dir=docs_dir, + verdict_input=_VI, client_factory=make_client_factory(_VALID, tokens=100), - store=fresh_store, max_tokens=10, # synthetic 100 tokens > cap 10 -> halt + store=fresh_store, + max_tokens=10, # synthetic 100 tokens > cap 10 -> halt ) -async def test_f_malformed_contract_raises_before_any_chat( - docs_dir, make_client_factory +async def test_wiring_budget_middleware_and_retrieval_tool( + docs_dir, make_client_factory, fresh_store, monkeypatch ) -> None: + """Step 3 (F2/F5/F7 wiring): run_project constructs a BudgetMiddleware + a retrieval + FunctionTool and passes BOTH into fresh_workflow. Spy the factory call; build a real + workflow so the rest of the run completes.""" + from agent_framework import FunctionTool + + from portfolio_optimiser import run as run_mod + from portfolio_optimiser.budget import BudgetMiddleware + + captured: dict[str, object] = {} + real_fresh = run_mod.fresh_workflow + + def spy_fresh(factory, **kwargs): + captured.update(kwargs) + return real_fresh(factory, **kwargs) + + monkeypatch.setattr(run_mod, "fresh_workflow", spy_fresh) + await run_project( + "FV42-GSV-E1", + "local", + docs_dir=docs_dir, + verdict_input=_VI, + client_factory=make_client_factory(_VALID), + store=fresh_store, + ) + middleware = captured.get("middleware") or [] + tools = captured.get("tools") or [] + assert any(isinstance(m, BudgetMiddleware) for m in middleware) + assert any(isinstance(t, FunctionTool) for t in tools) + + +async def test_f_malformed_contract_raises_before_any_chat(docs_dir, make_client_factory) -> None: calls = {"n": 0} base = make_client_factory(_VALID) @@ -94,7 +148,11 @@ async def test_f_malformed_contract_raises_before_any_chat( with pytest.raises(ValidationError): await run_project( - "FV42-GSV-E1", "local", docs_dir=docs_dir, verdict_input=_VI, - client_factory=spy, top_k=0, # malformed data-source contract + "FV42-GSV-E1", + "local", + docs_dir=docs_dir, + verdict_input=_VI, + client_factory=spy, + top_k=0, # malformed data-source contract ) assert calls["n"] == 0 # validation fails before any chat client is touched