diff --git a/app/desktop/studio_server/eval_api.py b/app/desktop/studio_server/eval_api.py index 3c42b3a485..89c236a9e1 100644 --- a/app/desktop/studio_server/eval_api.py +++ b/app/desktop/studio_server/eval_api.py @@ -735,11 +735,26 @@ def scored_trace_usage(trace: TaskRun) -> Usage | None: Latency still reads from `usage` — `cumulative_usage` deliberately carries none, since per-message latencies don't aggregate meaningfully. - An eval-driven conversation stores the synthetic-user driver model's spend - in `synthetic_user_usage`, beside the assistant-only `usage`; the honest - total sums the two null-tolerantly. Today that field carries cost only, so - token counts pass through from the assistant side. Migrated legacy traces - have the blend fused inside `usage` with `synthetic_user_usage` None, so - the same sum reads both record generations correctly. + in `synthetic_user_usage`, beside the assistant-only `usage`. Only its + **cost** is blended in here, deliberately, even though the field carries + the driver's tokens and latency too: + + * Cost is total-spend semantics — what this trace cost to produce, both + models included. That is what a run-config summary should report, and it + matches migrated legacy traces, which have the blend fused inside + `usage` with `synthetic_user_usage` None. + * Tokens are not. The synthetic user is usually a different model on a + different provider from the agent under test, so folding its ~3.5k input + tokens per conversation into this figure would attribute them to the + agent and make `cost / total_tokens` meaningless — the exact conflation + `synthetic_user_usage` exists to undo. + * Latency is not. This summary reports how responsive the agent is; + driver-side wall clock is not the agent's, and summing it would make + every driven run config look slower than it is. + + Blending everything would also make this field mean two different + quantities depending on record age: legacy records blend cost only, since + that is all the old field carried. None when the record has nothing to report, so it contributes nothing to an average instead of counting as a zero. @@ -760,7 +775,9 @@ def scored_trace_usage(trace: TaskRun) -> Usage | None: base = trace.usage if trace.synthetic_user_usage is not None: - base = (base or Usage()) + trace.synthetic_user_usage + # Cost only — see the docstring. Adding the whole object would fold the + # driver's tokens and latency into the agent's figures. + base = (base or Usage()) + Usage(cost=trace.synthetic_user_usage.cost) if base is None or all(v is None for v in base.model_dump().values()): return None return base diff --git a/app/desktop/studio_server/test_eval_api.py b/app/desktop/studio_server/test_eval_api.py index bf18d3c6d6..c1b50be8ee 100644 --- a/app/desktop/studio_server/test_eval_api.py +++ b/app/desktop/studio_server/test_eval_api.py @@ -2374,6 +2374,66 @@ def test_migrated_legacy_trace_reads_unchanged(self, mock_task, data_source): assert trace.synthetic_user_usage is None assert scored_trace_usage(trace) == blended + def test_synthetic_user_blends_cost_only(self, mock_task, data_source): + """The driver's cost joins the total; its tokens and latency do not. + + The synthetic user is normally a different model on a different provider + from the agent under test. Folding its tokens in would attribute them to + the agent (~3.5k input per conversation) and make cost/token meaningless, + and folding its latency in would make every driven run config look slower + than it is. Cost alone is total-spend semantics, and it is what migrated + legacy records already blend — so all three quantities keep one meaning + across record generations. + """ + trace = TaskRun( + parent=mock_task, + input="in", + input_source=data_source, + output=TaskOutput(output="out", source=data_source), + usage=Usage( + input_tokens=100, + output_tokens=50, + total_tokens=150, + cost=1.0, + total_llm_latency_ms=4000, + ), + synthetic_user_usage=Usage( + input_tokens=3548, + output_tokens=61, + total_tokens=3609, + cost=0.25, + total_llm_latency_ms=9000, + ), + ) + + usage = scored_trace_usage(trace) + + assert usage is not None + # Cost blends: both models' spend produced this trace. + assert usage.cost == pytest.approx(1.25) + # Tokens and latency stay the agent's alone. + assert usage.input_tokens == 100 + assert usage.output_tokens == 50 + assert usage.total_tokens == 150 + assert usage.total_llm_latency_ms == 4000 + + def test_synthetic_user_without_cost_leaves_the_total_alone( + self, mock_task, data_source + ): + """A driver that reported tokens but no cost must not perturb anything — + including not turning an all-agent figure into a different object.""" + agent = Usage(input_tokens=100, total_tokens=150, cost=1.0) + trace = TaskRun( + parent=mock_task, + input="in", + input_source=data_source, + output=TaskOutput(output="out", source=data_source), + usage=agent, + synthetic_user_usage=Usage(input_tokens=3548, total_tokens=3609), + ) + + assert scored_trace_usage(trace) == agent + def test_nothing_to_report_reads_as_none(self, mock_task, data_source): trace = TaskRun( parent=mock_task, diff --git a/libs/core/kiln_ai/adapters/eval/eval_runner.py b/libs/core/kiln_ai/adapters/eval/eval_runner.py index 7a706cc8d1..57865bc97d 100644 --- a/libs/core/kiln_ai/adapters/eval/eval_runner.py +++ b/libs/core/kiln_ai/adapters/eval/eval_runner.py @@ -1048,10 +1048,12 @@ async def _persist_driven_conversation( trace=leaf.trace, usage=_conversation_usage(drive_result.chain), # The synthetic-user driver's spend rides its own field so `usage` - # stays honestly assistant-only; a zero-cost drive records None. - synthetic_user_usage=Usage(cost=drive_result.su_total_cost) - if drive_result.su_total_cost > 0 - else None, + # stays honestly assistant-only. The whole Usage, not a cost-only + # stub: the SU is usually a different model on a different provider + # from the agent, so its tokens are the half that makes the figure + # reconcilable. `drive_case` already returns None for a drive whose + # provider reported nothing. + synthetic_user_usage=drive_result.su_usage, cumulative_usage=leaf.cumulative_usage or MessageUsage.from_trace(leaf.trace), eval_source=EvalItemSource(source_type=source_type, source_id=source_id), diff --git a/libs/core/kiln_ai/adapters/eval/test_eval_runner.py b/libs/core/kiln_ai/adapters/eval/test_eval_runner.py index c2f1d6e880..132d872839 100644 --- a/libs/core/kiln_ai/adapters/eval/test_eval_runner.py +++ b/libs/core/kiln_ai/adapters/eval/test_eval_runner.py @@ -3518,12 +3518,16 @@ def multi_turn_eval_input(mock_task): def _fresh_leaf( task: Task, data_source: DataSource, - su_total_cost: float = 0.0, + su_usage: Usage | None = None, cumulative_usage: MessageUsage | None = None, trace: list[ChatCompletionMessageParam] | None = MULTI_TURN_TRACE, ) -> DriveCaseResult: """The in-memory DriveCaseResult drive_case_for_eval would return: - an id-less, trace-carrying, never-saved leaf plus the SU-side spend.""" + an id-less, trace-carrying, never-saved leaf plus the SU-side spend. + + `su_usage` defaults to None — the shape a drive whose provider reported + nothing produces, which is what the tests that don't care about SU spend + want.""" leaf = TaskRun( input="opening message", input_source=data_source, @@ -3533,7 +3537,7 @@ def _fresh_leaf( parent=task, ) leaf.id = None - return DriveCaseResult(chain=[leaf], su_total_cost=su_total_cost) + return DriveCaseResult(chain=[leaf], su_usage=su_usage) class TestRunV2MultiTurnRedrive: @@ -4274,7 +4278,7 @@ async def test_unsaved_fresh_run_does_not_crash_record_keys_on_item( class TestEvalRunUsageRecording: @pytest.mark.asyncio - async def test_drive_records_agent_usage_plus_su_cost( + async def test_drive_records_agent_usage_and_full_su_usage( self, mock_task, mock_run_config, @@ -4285,7 +4289,9 @@ async def test_drive_records_agent_usage_plus_su_cost( drive_result = _fresh_leaf( mock_task, data_source, - su_total_cost=0.25, + su_usage=Usage( + input_tokens=3548, output_tokens=61, total_tokens=3609, cost=0.25 + ), cumulative_usage=MessageUsage( input_tokens=100, output_tokens=50, total_tokens=150, cost=1.0 ), @@ -4327,9 +4333,15 @@ async def test_drive_records_agent_usage_plus_su_cost( assert trace.usage.cost == pytest.approx(1.0) assert trace.usage.input_tokens == 100 assert trace.usage.total_tokens == 150 - # The synthetic-user driver's spend rides its own field. + # The synthetic-user driver's spend rides its own field, with its own + # tokens — the agent's counts above are a different model on a different + # provider, so a cost-only SU record could be reconciled against neither + # invoice and split per model not at all. assert trace.synthetic_user_usage is not None assert trace.synthetic_user_usage.cost == pytest.approx(0.25) + assert trace.synthetic_user_usage.input_tokens == 3548 + assert trace.synthetic_user_usage.output_tokens == 61 + assert trace.synthetic_user_usage.total_tokens == 3609 assert trace.cumulative_usage == MessageUsage( input_tokens=100, output_tokens=50, total_tokens=150, cost=1.0 ) @@ -4345,7 +4357,7 @@ async def test_zero_cost_drive_records_no_synthetic_user_usage( ): """None, not a zero-cost Usage: the rollup's null-tolerant blend would read a zero object as a real 0.0 cost and count it in averages.""" - drive_result = _fresh_leaf(mock_task, data_source, su_total_cost=0.0) + drive_result = _fresh_leaf(mock_task, data_source, su_usage=None) runner = EvalRunner( eval_configs=[mock_v2_redrive_config], run_configs=[mock_run_config], diff --git a/libs/core/kiln_ai/synthetic_user/drive_loop.py b/libs/core/kiln_ai/synthetic_user/drive_loop.py index 3424ae3409..b8bccaaf8f 100644 --- a/libs/core/kiln_ai/synthetic_user/drive_loop.py +++ b/libs/core/kiln_ai/synthetic_user/drive_loop.py @@ -16,6 +16,7 @@ from typing import Protocol from kiln_ai.datamodel.task_run import TaskRun +from kiln_ai.datamodel.usage import Usage from kiln_ai.synthetic_user.driver import SyntheticUserDriver from kiln_ai.utils.open_ai_types import ChatCompletionMessageParam @@ -63,14 +64,32 @@ class DriveCaseResult: no stop_reason field — every case ends after exactly `turns` iterations by design. - `su_total_cost` sums the SU driver's per-turn LLM cost across the - case. SU turns aren't persisted as TaskRuns, so this is the only - place that spend surfaces — the runner adds it to the target's - `cumulative_usage.cost` to produce an honest total. + `su_usage` sums the SU driver model's usage across the case's turns. + SU turns aren't persisted as TaskRuns, so this is the only place that + spend surfaces at all — and it carries the driver's tokens, not just + its cost, because the SU is usually a different model on a different + provider from the agent under test. A cost alone can be reconciled + against no invoice and split per model not at all. + + None when no turn reported usage, rather than a zeroed Usage: an + unmeasured drive must not read as a genuinely free one. """ chain: list[TaskRun] - su_total_cost: float + su_usage: Usage | None + + @property + def su_total_cost(self) -> float: + """The case's SU cost, or 0.0 when nothing was reported. + + Kept as a derived property so the interactive runner's + `CaseCompletedEvent.total_cost` is unchanged in behaviour — it wants a + float it can add, and "no usage reported" has always summed as zero + there. + """ + if self.su_usage is None or self.su_usage.cost is None: + return 0.0 + return float(self.su_usage.cost) async def drive_case( @@ -111,7 +130,9 @@ async def drive_case( prev_run: TaskRun | None = None prev_trace: list[ChatCompletionMessageParam] | None = None chain: list[TaskRun] = [] - su_total_cost: float = 0.0 + # Stays None until some turn actually reports usage, so a drive whose + # provider reported nothing is distinguishable from a free one. + su_usage: Usage | None = None for turn in range(1, turns + 1): new_run = await target_invoker( @@ -128,8 +149,13 @@ async def drive_case( su_message: str | None = None su_cost = 0.0 if turn < turns: - su_message, su_cost = await su_driver.respond(new_run.trace or []) - su_total_cost += su_cost + su_message, turn_usage = await su_driver.respond(new_run.trace or []) + if turn_usage is not None: + # Usage.__add__ is None-graceful per field, so a turn that + # reports only cost doesn't erase another turn's token counts. + su_usage = turn_usage if su_usage is None else su_usage + turn_usage + if turn_usage.cost is not None: + su_cost = float(turn_usage.cost) if on_turn is not None: await on_turn(run=new_run, su_message=su_message, su_cost=su_cost) @@ -139,4 +165,4 @@ async def drive_case( prev_run = new_run prev_trace = new_run.trace - return DriveCaseResult(chain=chain, su_total_cost=su_total_cost) + return DriveCaseResult(chain=chain, su_usage=su_usage) diff --git a/libs/core/kiln_ai/synthetic_user/driver.py b/libs/core/kiln_ai/synthetic_user/driver.py index 6c0780d792..383d75467e 100644 --- a/libs/core/kiln_ai/synthetic_user/driver.py +++ b/libs/core/kiln_ai/synthetic_user/driver.py @@ -14,6 +14,7 @@ from kiln_ai.datamodel.datamodel_enums import StructuredOutputMode from kiln_ai.datamodel.run_config import KilnAgentRunConfigProperties, ToolsRunConfig from kiln_ai.datamodel.task import Task +from kiln_ai.datamodel.usage import Usage from kiln_ai.synthetic_user.models import ( SyntheticUserDriverConfig, SyntheticUserInfo, @@ -84,11 +85,20 @@ def __init__( async def respond( self, conversation: list[ChatCompletionMessageParam] - ) -> tuple[str, float]: - """Return the SU's next message and the per-call cost. + ) -> tuple[str, Usage | None]: + """Return the SU's next message and the driver model's usage for the call. `conversation` is in the eval frame and must end on an `assistant` (target) turn. Drive-loop termination is the caller's concern. + + The whole `Usage` rather than just its cost: the SU's TaskRun is never + persisted, so this in-memory value is the only place the driver model's + tokens ever exist. A cost alone can neither be split per model against an + invoice nor recomputed at a different price, and `cost / total_tokens` + over a figure whose tokens are the agent's is meaningless. + + None when the provider reported nothing — distinct from a zeroed Usage, + which would read as a genuinely free call rather than an unmeasured one. """ # 1) Filter to visible roles (drop system/tool if present). visible = [ @@ -133,12 +143,4 @@ async def respond( if not isinstance(raw, str): raise RuntimeError("synthetic user returned non-string output") - # Per-call cost; defaults to 0.0 when the provider doesn't - # surface pricing. - cost = ( - float(task_run.usage.cost) - if task_run.usage is not None and task_run.usage.cost is not None - else 0.0 - ) - - return raw, cost + return raw, task_run.usage diff --git a/libs/core/kiln_ai/synthetic_user/eval_drive.py b/libs/core/kiln_ai/synthetic_user/eval_drive.py index fe941a8821..aec7e3197d 100644 --- a/libs/core/kiln_ai/synthetic_user/eval_drive.py +++ b/libs/core/kiln_ai/synthetic_user/eval_drive.py @@ -46,7 +46,7 @@ async def drive_case_for_eval( The result's leaf (`chain[-1]`) has `.trace` holding the full cumulative conversation and its id is None (nothing touches disk). The result also - carries `su_total_cost` — the synthetic user's LLM spend, which surfaces + carries `su_usage` — the synthetic user model's spend, which surfaces nowhere else since SU turns aren't persisted. `skills` must be preloaded by the caller — the adapter raises on skill tools with no injected dict. """ diff --git a/libs/core/kiln_ai/synthetic_user/test_drive_loop.py b/libs/core/kiln_ai/synthetic_user/test_drive_loop.py index 6b9b3cfbb4..fad922e07f 100644 --- a/libs/core/kiln_ai/synthetic_user/test_drive_loop.py +++ b/libs/core/kiln_ai/synthetic_user/test_drive_loop.py @@ -11,6 +11,7 @@ import pytest from kiln_ai.datamodel.task_run import TaskRun +from kiln_ai.datamodel.usage import Usage from kiln_ai.synthetic_user.drive_loop import DriveCaseResult, drive_case from kiln_ai.synthetic_user.driver import SyntheticUserDriver @@ -73,13 +74,27 @@ async def __call__( return _fake_run(new_trace, run_id=f"run-turn-{len(self.calls)}") -def _su_driver_with_replies(replies: list[str], cost_per_reply: float = 0.0) -> Mock: +def _su_driver_with_replies( + replies: list[str], + cost_per_reply: float = 0.0, + usage_per_reply: Usage | None = None, +) -> Mock: """Mock(spec=SyntheticUserDriver) with respond() returning canned - (message, cost) tuples. `cost_per_reply` lets cost-aware tests inject - a non-zero per-call cost; defaults to 0.0 for the legacy tests. + (message, Usage | None) tuples. + + `usage_per_reply` is the direct control, for tests that care about the + driver model's tokens. `cost_per_reply` is kept for the cost-only tests: + a non-zero value becomes a cost-only Usage, and the 0.0 default becomes + None — the shape a provider that reported nothing produces. """ + if usage_per_reply is not None: + usage: Usage | None = usage_per_reply + elif cost_per_reply: + usage = Usage(cost=cost_per_reply) + else: + usage = None drv = Mock(spec=SyntheticUserDriver) - drv.respond = AsyncMock(side_effect=[(r, cost_per_reply) for r in replies]) + drv.respond = AsyncMock(side_effect=[(r, usage) for r in replies]) return drv @@ -192,6 +207,87 @@ async def test_drive_case_skips_su_reply_after_final_turn() -> None: assert result.su_total_cost == pytest.approx(0.01) +@pytest.mark.asyncio +async def test_drive_case_sums_su_usage_across_turns() -> None: + """The driver model's tokens are summed, not just its cost. + + SU turns are never persisted as TaskRuns, so anything the loop drops here + exists nowhere on disk afterwards. The SU is normally a different model on a + different provider from the agent, so its token counts are what make the + figure reconcilable against an invoice at all. + """ + invoker = _FakeInvoker(assistant_replies=["a1", "a2", "a3"]) + su = _su_driver_with_replies( + ["u2", "u3"], + usage_per_reply=Usage( + input_tokens=1000, output_tokens=20, total_tokens=1020, cost=0.01 + ), + ) + + result = await drive_case( + seed_prompt="u1", + target_invoker=invoker, + su_driver=su, + turns=3, + ) + + # Two SU calls for three turns — the loop skips the SU after the final turn. + assert su.respond.await_count == 2 + assert result.su_usage is not None + assert result.su_usage.input_tokens == 2000 + assert result.su_usage.output_tokens == 40 + assert result.su_usage.total_tokens == 2040 + assert result.su_usage.cost == pytest.approx(0.02) + # The derived cost keeps the interactive runner's total unchanged. + assert result.su_total_cost == pytest.approx(0.02) + + +@pytest.mark.asyncio +async def test_drive_case_su_usage_is_none_when_no_turn_reports() -> None: + """None, not a zeroed Usage: an unmeasured drive must not read as a free + one. `su_total_cost` still answers 0.0, so cost sums stay well-defined.""" + invoker = _FakeInvoker(assistant_replies=["a1", "a2"]) + su = _su_driver_with_replies(["u2"]) + + result = await drive_case( + seed_prompt="u1", + target_invoker=invoker, + su_driver=su, + turns=2, + ) + + assert su.respond.await_count == 1 + assert result.su_usage is None + assert result.su_total_cost == 0.0 + + +@pytest.mark.asyncio +async def test_drive_case_su_usage_keeps_tokens_when_a_turn_reports_only_cost() -> None: + """A provider that surfaces pricing but no token counts on one turn must not + wipe out the counts another turn did report — Usage.__add__ is None-graceful + per field, and the loop relies on exactly that.""" + invoker = _FakeInvoker(assistant_replies=["a1", "a2", "a3"]) + su = Mock(spec=SyntheticUserDriver) + su.respond = AsyncMock( + side_effect=[ + ("u2", Usage(input_tokens=800, output_tokens=10, cost=0.01)), + ("u3", Usage(cost=0.02)), + ] + ) + + result = await drive_case( + seed_prompt="u1", + target_invoker=invoker, + su_driver=su, + turns=3, + ) + + assert result.su_usage is not None + assert result.su_usage.input_tokens == 800 + assert result.su_usage.output_tokens == 10 + assert result.su_usage.cost == pytest.approx(0.03) + + @pytest.mark.asyncio async def test_drive_case_single_turn_never_calls_su_driver() -> None: """turns=1 is seed → target → done: the SU driver is never invoked.""" diff --git a/libs/core/kiln_ai/synthetic_user/test_driver.py b/libs/core/kiln_ai/synthetic_user/test_driver.py index 8bfcc73e3b..c49eb0f37a 100644 --- a/libs/core/kiln_ai/synthetic_user/test_driver.py +++ b/libs/core/kiln_ai/synthetic_user/test_driver.py @@ -42,20 +42,25 @@ def _patch_adapter( monkeypatch: pytest.MonkeyPatch, return_value: RunOutput, cost: float | None = None, + usage: Usage | None = None, ) -> Mock: """Replace adapter_for_task with a stub returning a mock adapter whose invoke_returning_run_output yields (Mock(spec=TaskRun), return_value). Returns the mock adapter so tests can assert call args. - Pass `cost` to populate the in-memory TaskRun's `.usage.cost`; when - omitted, `.usage` is None and `respond()` should report cost=0.0. + Pass `usage` to set the in-memory TaskRun's `.usage` outright, or `cost` + for the cost-only shorthand. With neither, `.usage` is None and + `respond()` reports None. """ task_run = Mock(spec=TaskRun) - task_run.usage = ( - Usage(input_tokens=0, output_tokens=0, total_tokens=0, cost=cost) - if cost is not None - else None - ) + if usage is not None: + task_run.usage = usage + elif cost is not None: + task_run.usage = Usage( + input_tokens=0, output_tokens=0, total_tokens=0, cost=cost + ) + else: + task_run.usage = None adapter = Mock() adapter.invoke_returning_run_output = AsyncMock( return_value=(task_run, return_value) @@ -97,8 +102,11 @@ def test_construction_renders_system_prompt_once( async def test_respond_returns_adapter_output_and_zero_cost_when_unset( monkeypatch: pytest.MonkeyPatch, ) -> None: - """When the provider doesn't surface pricing, cost defaults to 0.0 - so downstream sums stay well-defined. + """A provider that reports nothing yields None, not a zeroed Usage. + + The distinction is the point of returning Usage at all: a zeroed record + would read downstream as a genuinely free call rather than an unmeasured + one, and the drive's total would claim a precision it doesn't have. """ adapter = _patch_adapter(monkeypatch, _fake_run_output("the SU's reply")) drv = SyntheticUserDriver(_INFO, _DRIVER_CONFIG) @@ -107,28 +115,42 @@ async def test_respond_returns_adapter_output_and_zero_cost_when_unset( {"role": "assistant", "content": "a1"}, ] - message, cost = await drv.respond(conversation) + message, usage = await drv.respond(conversation) assert message == "the SU's reply" - assert cost == 0.0 + assert usage is None adapter.invoke_returning_run_output.assert_awaited_once() @pytest.mark.asyncio -async def test_respond_returns_cost_from_task_run_usage( +async def test_respond_returns_usage_from_task_run( monkeypatch: pytest.MonkeyPatch, ) -> None: - """Per-call cost is read from the in-memory TaskRun's `usage.cost`.""" - _patch_adapter(monkeypatch, _fake_run_output("hi"), cost=0.0123) + """The whole Usage comes back, tokens included — not just the cost. + + The SU's TaskRun is never persisted, so anything dropped here is gone for + good: these token counts exist nowhere else on disk. + """ + _patch_adapter( + monkeypatch, + _fake_run_output("hi"), + usage=Usage( + input_tokens=3548, output_tokens=61, total_tokens=3609, cost=0.0123 + ), + ) drv = SyntheticUserDriver(_INFO, _DRIVER_CONFIG) conversation: list[ChatCompletionMessageParam] = [ {"role": "user", "content": "u1"}, {"role": "assistant", "content": "a1"}, ] - _, cost = await drv.respond(conversation) + _, usage = await drv.respond(conversation) - assert cost == pytest.approx(0.0123) + assert usage is not None + assert usage.input_tokens == 3548 + assert usage.output_tokens == 61 + assert usage.total_tokens == 3609 + assert usage.cost == pytest.approx(0.0123) @pytest.mark.asyncio diff --git a/libs/core/kiln_ai/synthetic_user/test_eval_drive.py b/libs/core/kiln_ai/synthetic_user/test_eval_drive.py index 2ea373a331..0abb1f8981 100644 --- a/libs/core/kiln_ai/synthetic_user/test_eval_drive.py +++ b/libs/core/kiln_ai/synthetic_user/test_eval_drive.py @@ -96,7 +96,7 @@ def _factory(task, run_config, base_adapter_config=None): def fake_su_driver(monkeypatch: pytest.MonkeyPatch) -> Mock: instance = Mock(spec=SyntheticUserDriver) instance.respond = AsyncMock( - side_effect=[(f"follow-up-{i}", 0.0) for i in range(1, 10)] + side_effect=[(f"follow-up-{i}", None) for i in range(1, 10)] ) captured_ctor: dict[str, Any] = {} diff --git a/libs/core/kiln_ai/synthetic_user/test_runner.py b/libs/core/kiln_ai/synthetic_user/test_runner.py index 39d8e80ade..87731ede46 100644 --- a/libs/core/kiln_ai/synthetic_user/test_runner.py +++ b/libs/core/kiln_ai/synthetic_user/test_runner.py @@ -21,6 +21,7 @@ from kiln_ai.datamodel.run_config import KilnAgentRunConfigProperties, ToolsRunConfig from kiln_ai.datamodel.task import Task from kiln_ai.datamodel.task_run import TaskRun +from kiln_ai.datamodel.usage import Usage from kiln_ai.synthetic_user import runner as runner_mod from kiln_ai.synthetic_user.case import SyntheticUserCase from kiln_ai.synthetic_user.driver import SyntheticUserDriver @@ -120,9 +121,10 @@ def _ctor(info, config): else list(replies_per_case) ) instance = Mock(spec=SyntheticUserDriver) - # respond() returns (message, cost). Tests that don't care about - # cost get 0.0 — the runner adds it to total_cost regardless. - instance.respond = AsyncMock(side_effect=[(r, 0.0) for r in replies]) + # respond() returns (message, Usage | None). Tests that don't care about + # the driver's spend hand back None — the shape a provider that reported + # nothing produces, which the runner totals as zero. + instance.respond = AsyncMock(side_effect=[(r, None) for r in replies]) return instance monkeypatch.setattr(runner_mod, "SyntheticUserDriver", _ctor) @@ -254,7 +256,7 @@ async def test_total_cost_sums_target_and_su_driver_spend( # at $0.01 → $0.01 SU per case. def _ctor(info, config): instance = Mock(spec=SyntheticUserDriver) - instance.respond = AsyncMock(side_effect=[("u2", 0.01)]) + instance.respond = AsyncMock(side_effect=[("u2", Usage(cost=0.01))]) return instance monkeypatch.setattr(runner_mod, "SyntheticUserDriver", _ctor) @@ -450,7 +452,7 @@ async def test_malformed_blob_surfaces_as_case_failed( def _ctor(info, config): instance = Mock(spec=SyntheticUserDriver) - instance.respond = AsyncMock(return_value=("ok", 0.0)) + instance.respond = AsyncMock(return_value=("ok", None)) return instance _patch_su_driver_factory(monkeypatch, _ctor) @@ -574,7 +576,7 @@ def _ctor(info, config): instance = Mock(spec=SyntheticUserDriver) # Turn 1's SU reply succeeds; turn 2's SU call dies mid-case. instance.respond = AsyncMock( - side_effect=[("u2", 0.0), ValueError("su blew up")] + side_effect=[("u2", None), ValueError("su blew up")] ) return instance @@ -774,7 +776,7 @@ async def test_retried_case_batch_total_includes_both_attempts_costs( # One SU call per attempt (turns=2), at $0.01. def _ctor(info, config): instance = Mock(spec=SyntheticUserDriver) - instance.respond = AsyncMock(side_effect=[("u2", 0.01)]) + instance.respond = AsyncMock(side_effect=[("u2", Usage(cost=0.01))]) return instance monkeypatch.setattr(runner_mod, "SyntheticUserDriver", _ctor) @@ -824,7 +826,7 @@ async def test_dead_case_failed_event_reports_all_attempts_spend( def _ctor(info, config): instance = Mock(spec=SyntheticUserDriver) - instance.respond = AsyncMock(side_effect=[("u2", 0.01)]) + instance.respond = AsyncMock(side_effect=[("u2", Usage(cost=0.01))]) return instance monkeypatch.setattr(runner_mod, "SyntheticUserDriver", _ctor)