feat(fase2): wire BudgetMiddleware + retrieval tool onto the debate in run_project

This commit is contained in:
Kjell Tore Guttormsen 2026-06-26 00:37:45 +02:00
commit 7573c4439f
2 changed files with 92 additions and 19 deletions

View file

@ -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.

View file

@ -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