refactor(providers): define typed usage contract

This commit is contained in:
chengyongru
2026-08-25 01:04:25 +08:00
committed by chengyongru
parent 89c94d8744
commit 9895c23cb5
84 changed files with 1643 additions and 726 deletions
+1 -1
View File
@@ -128,7 +128,7 @@ async def test_pending_document_attachment_keeps_body_out_of_prompt(
nonlocal call_count
call_count += 1
captured_messages.append([dict(message) for message in messages])
return LLMResponse(content=f"answer-{call_count}", tool_calls=[], usage={})
return LLMResponse(content=f"answer-{call_count}", tool_calls=[], usage=None)
loop = _make_loop(workspace)
loop.provider.chat_with_retry = chat_with_retry
+4 -4
View File
@@ -412,7 +412,7 @@ class TestEphemeralDirect:
provider.supports_tools = True
provider.generation = MagicMock(max_tokens=4096)
provider.chat_with_retry = AsyncMock(
return_value=LLMResponse(content="done", tool_calls=[], finish_reason="stop", usage={})
return_value=LLMResponse(content="done", tool_calls=[], finish_reason="stop", usage=None)
)
with (
@@ -556,9 +556,9 @@ class TestEphemeralDirect:
"new_text": "replacement",
},
)],
usage={},
usage=None,
),
LLMResponse(content="done", finish_reason="stop", tool_calls=[], usage={}),
LLMResponse(content="done", finish_reason="stop", tool_calls=[], usage=None),
])
resp = await loop.process_direct(
@@ -646,7 +646,7 @@ class TestEphemeralHooks:
provider.generation = MagicMock(max_tokens=4096)
provider.chat_with_retry = AsyncMock(
return_value=LLMResponse(
content="done", finish_reason="stop", tool_calls=[], usage={},
content="done", finish_reason="stop", tool_calls=[], usage=None,
)
)
+1 -1
View File
@@ -12,7 +12,7 @@ from nanobot.utils.evaluator import (
class DummyProvider(LLMProvider):
def __init__(self, responses: list[LLMResponse]):
super().__init__()
super().__init__(provider_name="dummy")
self._responses = list(responses)
async def chat(self, *args, **kwargs) -> LLMResponse:
+2 -2
View File
@@ -69,7 +69,7 @@ def test_explicit_message_limit_still_starts_at_user_turn() -> None:
async def test_process_message_replays_with_token_budget_only(tmp_path: Path) -> None:
loop = _make_loop(tmp_path, context_window_tokens=32_768)
loop.provider.chat_with_retry = AsyncMock(
return_value=LLMResponse(content="ok", tool_calls=[], usage={})
return_value=LLMResponse(content="ok", tool_calls=[], usage=None)
)
loop.tools.get_definitions = MagicMock(return_value=[])
loop.consolidator.maybe_consolidate_by_tokens = AsyncMock(return_value=False) # type: ignore[method-assign]
@@ -91,7 +91,7 @@ async def test_process_message_replays_with_token_budget_only(tmp_path: Path) ->
async def test_token_budget_keeps_current_user_as_replay_boundary(tmp_path: Path) -> None:
loop = _make_loop(tmp_path, context_window_tokens=8_000)
loop.provider.chat_with_retry = AsyncMock(
return_value=LLMResponse(content="ok", tool_calls=[], usage={})
return_value=LLMResponse(content="ok", tool_calls=[], usage=None)
)
loop.tools.get_definitions = MagicMock(return_value=[])
loop.consolidator.maybe_consolidate_by_tokens = AsyncMock(return_value=False) # type: ignore[method-assign]
+4 -4
View File
@@ -453,7 +453,7 @@ async def test_agent_loop_extra_hook_receives_calls(tmp_path):
loop = _make_loop(tmp_path, hooks=[TrackingHook()])
loop.provider.chat_with_retry = AsyncMock(
return_value=LLMResponse(content="done", tool_calls=[], usage={})
return_value=LLMResponse(content="done", tool_calls=[], usage=None)
)
loop.tools.get_definitions = MagicMock(return_value=[])
@@ -494,7 +494,7 @@ async def test_agent_loop_turn_hook_factories_receive_context(tmp_path):
loop = _make_loop(tmp_path, hook_factories=[factory("registered")])
loop.provider.chat_with_retry = AsyncMock(
return_value=LLMResponse(content="done", tool_calls=[], usage={})
return_value=LLMResponse(content="done", tool_calls=[], usage=None)
)
loop.tools.get_definitions = MagicMock(return_value=[])
@@ -541,7 +541,7 @@ async def test_agent_loop_extra_hook_error_isolation(tmp_path):
loop = _make_loop(tmp_path, hooks=[BadHook()])
loop.provider.chat_with_retry = AsyncMock(
return_value=LLMResponse(content="still works", tool_calls=[], usage={})
return_value=LLMResponse(content="still works", tool_calls=[], usage=None)
)
loop.tools.get_definitions = MagicMock(return_value=[])
@@ -562,7 +562,7 @@ async def test_agent_loop_extra_hooks_do_not_swallow_loop_hook_errors(tmp_path):
loop.provider.chat_with_retry = AsyncMock(return_value=LLMResponse(
content="working",
tool_calls=[ToolCallRequest(id="c1", name="list_dir", arguments={"path": "."})],
usage={},
usage=None,
))
loop.tools.get_definitions = MagicMock(return_value=[])
loop.tools.execute = AsyncMock(return_value="ok")
+2 -2
View File
@@ -393,9 +393,9 @@ class TestToolEventProgress:
},
)
],
usage={},
usage=None,
)
return LLMResponse(content="Done", tool_calls=[], usage={})
return LLMResponse(content="Done", tool_calls=[], usage=None)
provider.chat_stream_with_retry = chat_stream_with_retry
provider.chat_with_retry = AsyncMock()
+22 -22
View File
@@ -48,7 +48,7 @@ async def test_ephemeral_runner_enters_and_restores_turn_scopes(tmp_path):
async def chat_with_retry(**_kwargs):
assert goal_mutation_allowed() is True
return LLMResponse(content="done", tool_calls=[], usage={})
return LLMResponse(content="done", tool_calls=[], usage=None)
loop.provider.chat_with_retry = AsyncMock(side_effect=chat_with_retry)
loop.tools.get_definitions = MagicMock(return_value=[])
@@ -83,7 +83,7 @@ async def test_goal_command_can_implement_plan_from_prior_discussion(tmp_path):
},
)
],
usage={},
usage=None,
),
LLMResponse(
content="closing goal",
@@ -94,7 +94,7 @@ async def test_goal_command_can_implement_plan_from_prior_discussion(tmp_path):
arguments={"action": "complete", "recap": "Implemented and tested."},
)
],
usage={},
usage=None,
),
LLMResponse(
content="trying to start another goal",
@@ -105,9 +105,9 @@ async def test_goal_command_can_implement_plan_from_prior_discussion(tmp_path):
arguments={"objective": "Start an unrelated follow-up."},
)
],
usage={},
usage=None,
),
LLMResponse(content="done", tool_calls=[], usage={}),
LLMResponse(content="done", tool_calls=[], usage=None),
])
loop = AgentLoop(bus=MessageBus(), provider=provider, workspace=tmp_path, model="test-model")
loop.consolidator.maybe_consolidate_by_tokens = AsyncMock(return_value=None)
@@ -160,8 +160,8 @@ async def test_runtime_context_is_persisted_as_next_turn_prompt_prefix(tmp_path)
provider.get_default_model.return_value = "test-model"
provider.generation = GenerationSettings()
provider.chat_with_retry = AsyncMock(side_effect=[
LLMResponse(content="first answer", usage={}),
LLMResponse(content="second answer", usage={}),
LLMResponse(content="first answer", usage=None),
LLMResponse(content="second answer", usage=None),
])
loop = AgentLoop(bus=MessageBus(), provider=provider, workspace=tmp_path, model="test-model")
loop.consolidator.maybe_consolidate_by_tokens = AsyncMock(return_value=None)
@@ -216,7 +216,7 @@ async def test_webui_quote_reaches_model_without_leaking_into_public_history(tmp
provider = MagicMock()
provider.get_default_model.return_value = "test-model"
provider.generation = GenerationSettings()
provider.chat_with_retry = AsyncMock(return_value=LLMResponse(content="answer", usage={}))
provider.chat_with_retry = AsyncMock(return_value=LLMResponse(content="answer", usage=None))
loop = AgentLoop(bus=MessageBus(), provider=provider, workspace=tmp_path, model="test-model")
loop.consolidator.maybe_consolidate_by_tokens = AsyncMock(return_value=None)
session = loop.sessions.get_or_create("websocket:chat")
@@ -258,9 +258,9 @@ async def test_runtime_context_provider_runs_once_across_tool_iterations(tmp_pat
name="read_file",
arguments={"path": "note.txt"},
)],
usage={},
usage=None,
),
LLMResponse(content="done", usage={}),
LLMResponse(content="done", usage=None),
])
loop = AgentLoop(bus=MessageBus(), provider=provider, workspace=tmp_path, model="test-model")
loop.consolidator.maybe_consolidate_by_tokens = AsyncMock(return_value=None)
@@ -303,9 +303,9 @@ async def test_non_goal_direct_turn_cannot_reuse_prior_goal_command(tmp_path):
arguments={"objective": "Unauthorized persistent objective."},
)
],
usage={},
usage=None,
),
LLMResponse(content="handled as a one-time task", tool_calls=[], usage={}),
LLMResponse(content="handled as a one-time task", tool_calls=[], usage=None),
])
loop = AgentLoop(bus=MessageBus(), provider=provider, workspace=tmp_path, model="test-model")
loop.consolidator.maybe_consolidate_by_tokens = AsyncMock(return_value=None)
@@ -383,7 +383,7 @@ async def test_loop_stream_filter_handles_think_only_prefix_without_crashing(tmp
async def chat_stream_with_retry(*, on_content_delta, **kwargs):
await on_content_delta("<think>hidden")
await on_content_delta("</think>Hello")
return LLMResponse(content="<think>hidden</think>Hello", tool_calls=[], usage={})
return LLMResponse(content="<think>hidden</think>Hello", tool_calls=[], usage=None)
loop.provider.chat_stream_with_retry = chat_stream_with_retry
@@ -413,7 +413,7 @@ async def test_loop_stream_filter_hides_partial_trailing_think_prefix(tmp_path):
async def chat_stream_with_retry(*, on_content_delta, **kwargs):
await on_content_delta("Hello <thin")
await on_content_delta("k>hidden</think>World")
return LLMResponse(content="Hello <think>hidden</think>World", tool_calls=[], usage={})
return LLMResponse(content="Hello <think>hidden</think>World", tool_calls=[], usage=None)
loop.provider.chat_stream_with_retry = chat_stream_with_retry
@@ -436,7 +436,7 @@ async def test_loop_stream_filter_hides_complete_trailing_think_tag(tmp_path):
async def chat_stream_with_retry(*, on_content_delta, **kwargs):
await on_content_delta("Hello <think>")
await on_content_delta("hidden</think>World")
return LLMResponse(content="Hello <think>hidden</think>World", tool_calls=[], usage={})
return LLMResponse(content="Hello <think>hidden</think>World", tool_calls=[], usage=None)
loop.provider.chat_stream_with_retry = chat_stream_with_retry
@@ -459,8 +459,8 @@ async def test_loop_retries_think_only_final_response(tmp_path):
async def chat_with_retry(**kwargs):
call_count["n"] += 1
if call_count["n"] == 1:
return LLMResponse(content="<think>hidden</think>", tool_calls=[], usage={})
return LLMResponse(content="Recovered answer", tool_calls=[], usage={})
return LLMResponse(content="<think>hidden</think>", tool_calls=[], usage=None)
return LLMResponse(content="Recovered answer", tool_calls=[], usage=None)
loop.provider.chat_with_retry = chat_with_retry
@@ -485,7 +485,7 @@ async def test_streamed_flag_not_set_on_llm_error(tmp_path):
provider.get_default_model.return_value = "test-model"
loop = AgentLoop(bus=bus, provider=provider, workspace=tmp_path, model="test-model")
error_resp = LLMResponse(
content="503 service unavailable", finish_reason="error", tool_calls=[], usage={},
content="503 service unavailable", finish_reason="error", tool_calls=[], usage=None,
)
loop.provider.chat_with_retry = AsyncMock(return_value=error_resp)
loop.provider.chat_stream_with_retry = AsyncMock(return_value=error_resp)
@@ -523,14 +523,14 @@ async def test_ssrf_soft_block_can_finalize_after_streamed_tool_call(tmp_path):
name="exec",
arguments={"command": "curl http://169.254.169.254/latest/meta-data/"},
)],
usage={},
usage=None,
)
responses = iter([
tool_call_resp,
LLMResponse(
content="I cannot access private URLs. Please share the local file.",
tool_calls=[],
usage={},
usage=None,
),
])
@@ -569,8 +569,8 @@ async def test_next_turn_after_llm_error_keeps_turn_boundary(tmp_path):
provider = MagicMock()
provider.get_default_model.return_value = "test-model"
provider.chat_with_retry = AsyncMock(side_effect=[
LLMResponse(content="429 rate limit exceeded", finish_reason="error", tool_calls=[], usage={}),
LLMResponse(content="Recovered answer", tool_calls=[], usage={}),
LLMResponse(content="429 rate limit exceeded", finish_reason="error", tool_calls=[], usage=None),
LLMResponse(content="Recovered answer", tool_calls=[], usage=None),
])
loop = AgentLoop(bus=MessageBus(), provider=provider, workspace=tmp_path, model="test-model")
+5 -6
View File
@@ -20,7 +20,7 @@ from nanobot.bus.outbound_events import (
)
from nanobot.bus.queue import MessageBus
from nanobot.cron.session_turns import CRON_HISTORY_META, CRON_TRIGGER_META
from nanobot.providers.base import LLMProvider, LLMResponse, ProviderConversationState
from nanobot.providers.base import LLMProvider, LLMResponse, LLMUsage, ProviderConversationState
from nanobot.providers.factory import ProviderSnapshot
from nanobot.runtime_context import (
RUNTIME_CONTEXT_HISTORY_META,
@@ -1883,7 +1883,7 @@ async def test_turn_usage_is_persisted_with_the_saved_session(tmp_path: Path) ->
loop.consolidator.maybe_consolidate_by_tokens = AsyncMock(return_value=False) # type: ignore[method-assign]
async def fake_run_agent_loop(initial_messages, **_kwargs):
loop._last_usage = {"prompt_tokens": 64, "completion_tokens": 9}
loop._last_usage = LLMUsage.reported(input_tokens=64, output_tokens=9)
return (
"done",
[],
@@ -1898,10 +1898,9 @@ async def test_turn_usage_is_persisted_with_the_saved_session(tmp_path: Path) ->
)
loop.sessions.invalidate("cli:usage")
assert loop.sessions.get_or_create("cli:usage").metadata["_last_usage"] == {
"prompt_tokens": 64,
"completion_tokens": 9,
}
assert loop.sessions.get_or_create("cli:usage").metadata["_last_usage"] == (
LLMUsage.reported(input_tokens=64, output_tokens=9).to_dict()
)
@pytest.mark.asyncio
+1 -1
View File
@@ -30,7 +30,7 @@ def _loop(tmp_path, responses: list[str], **kwargs) -> AgentLoop:
provider.get_default_model.return_value = "test-model"
provider.generation = GenerationSettings()
provider.chat_with_retry = AsyncMock(
side_effect=[LLMResponse(content=response, usage={}) for response in responses]
side_effect=[LLMResponse(content=response, usage=None) for response in responses]
)
return AgentLoop(
bus=MessageBus(),
+190 -31
View File
@@ -14,6 +14,7 @@ from nanobot.config.schema import AgentDefaults
from nanobot.providers.base import (
LLMProvider,
LLMResponse,
LLMUsage,
ProviderCallContext,
ProviderConversationState,
ToolCallRequest,
@@ -22,6 +23,163 @@ from nanobot.providers.base import (
_MAX_TOOL_RESULT_CHARS = AgentDefaults().max_tool_result_chars
def _make_usage_spec(provider, tools):
return make_run_spec(
provider,
initial_messages=[{"role": "user", "content": "hello"}],
tools=tools,
model="test-model",
max_iterations=1,
max_tool_result_chars=_MAX_TOOL_RESULT_CHARS,
)
def test_usage_or_estimate_replaces_reported_zero_for_content(monkeypatch) -> None:
from nanobot.agent.runner import AgentRunner
provider = MagicMock(spec=LLMProvider)
tools = MagicMock()
tools.get_definitions.return_value = []
monkeypatch.setattr(
"nanobot.agent.runner.estimate_prompt_tokens_chain",
lambda provider, model, messages, definitions: (12, "test"),
)
monkeypatch.setattr("nanobot.agent.runner.estimate_message_tokens", lambda message: 7)
response = LLMResponse(
content="answer",
usage=LLMUsage.reported(input_tokens=0, output_tokens=0),
generation_ms=25,
ttft_ms=5,
)
usage = AgentRunner()._usage_or_estimate(
_make_usage_spec(provider, tools),
[{"role": "user", "content": "hello"}],
response,
)
assert usage == LLMUsage.estimated(input_tokens=12, output_tokens=7).with_timing(
generation_ms=25,
ttft_ms=5,
)
assert usage.source == "estimated"
assert usage.total_tokens == 19
def test_usage_or_estimate_counts_tool_call_output_for_reported_zero(monkeypatch) -> None:
from nanobot.agent.runner import AgentRunner
provider = MagicMock(spec=LLMProvider)
tools = MagicMock()
tools.get_definitions.return_value = []
captured_message: dict = {}
monkeypatch.setattr(
"nanobot.agent.runner.estimate_prompt_tokens_chain",
lambda provider, model, messages, definitions: (13, "test"),
)
def estimate_output(message):
captured_message.update(message)
return 9
monkeypatch.setattr("nanobot.agent.runner.estimate_message_tokens", estimate_output)
response = LLMResponse(
content=None,
tool_calls=[
ToolCallRequest(
id="call_1",
name="lookup",
arguments={"query": "nanobot"},
)
],
finish_reason="tool_calls",
usage=LLMUsage.reported(input_tokens=0, output_tokens=0),
)
usage = AgentRunner()._usage_or_estimate(
_make_usage_spec(provider, tools),
[{"role": "user", "content": "hello"}],
response,
)
assert usage == LLMUsage.estimated(input_tokens=13, output_tokens=9)
assert usage.total_tokens == 22
assert captured_message["tool_calls"][0]["function"]["name"] == "lookup"
@pytest.mark.parametrize(
"provider_usage",
[None, LLMUsage.reported(input_tokens=0, output_tokens=0)],
)
def test_usage_or_estimate_counts_error_without_estimating_tokens(
monkeypatch,
provider_usage: LLMUsage | None,
) -> None:
from nanobot.agent.runner import AgentRunner
provider = MagicMock(spec=LLMProvider)
tools = MagicMock()
estimate = MagicMock()
runner = AgentRunner()
monkeypatch.setattr(runner, "_estimate_response_usage", estimate)
response = LLMResponse(
content="upstream failed",
finish_reason="error",
usage=provider_usage,
)
usage = runner._usage_or_estimate(
_make_usage_spec(provider, tools),
[{"role": "user", "content": "hello"}],
response,
)
assert usage is not None
assert usage.total_tokens == 0
assert usage.request_count == 1
assert usage.context_tokens is None
aggregate = LLMUsage.reported(input_tokens=12, output_tokens=3) + usage
assert aggregate.context_tokens == 12
assert aggregate.request_count == 2
estimate.assert_not_called()
def test_usage_or_estimate_trusts_positive_reported_total(monkeypatch) -> None:
from nanobot.agent.runner import AgentRunner
provider = MagicMock(spec=LLMProvider)
tools = MagicMock()
estimate = MagicMock()
runner = AgentRunner()
monkeypatch.setattr(runner, "_estimate_response_usage", estimate)
response = LLMResponse(
content="answer",
usage=LLMUsage.reported(
input_tokens=15,
output_tokens=18,
total_tokens=175,
),
generation_ms=30,
ttft_ms=6,
)
usage = runner._usage_or_estimate(
_make_usage_spec(provider, tools),
[{"role": "user", "content": "hello"}],
response,
)
assert usage is not None
assert usage.source == "reported"
assert usage.input_tokens == 15
assert usage.output_tokens == 18
assert usage.total_tokens == 175
assert usage.reported_tokens == 175
assert usage.generation_ms == 30
assert usage.ttft_ms == 6
estimate.assert_not_called()
@pytest.mark.asyncio
async def test_runner_preserves_reasoning_fields_and_tool_results():
from nanobot.agent.runner import AgentRunner
@@ -38,10 +196,10 @@ async def test_runner_preserves_reasoning_fields_and_tool_results():
tool_calls=[ToolCallRequest(id="call_1", name="list_dir", arguments={"path": "."})],
reasoning_content="hidden reasoning",
thinking_blocks=[{"type": "thinking", "thinking": "step"}],
usage={"prompt_tokens": 5, "completion_tokens": 3},
usage=LLMUsage.reported(input_tokens=5, output_tokens=3),
)
captured_second_call[:] = messages
return LLMResponse(content="done", tool_calls=[], usage={})
return LLMResponse(content="done", tool_calls=[], usage=None)
provider.chat_with_retry = chat_with_retry
tools = MagicMock()
@@ -441,7 +599,7 @@ async def test_runner_uses_no_tools_finalization_after_max_iterations():
return LLMResponse(
content="Read the directory twice. More investigation remains.",
tool_calls=[],
usage={"prompt_tokens": 10, "completion_tokens": 7},
usage=LLMUsage.reported(input_tokens=10, output_tokens=7),
)
provider.chat_with_retry = chat_with_retry
@@ -713,10 +871,10 @@ async def test_runner_replaces_empty_tool_result_with_marker():
return LLMResponse(
content="working",
tool_calls=[ToolCallRequest(id="call_1", name="noop", arguments={})],
usage={},
usage=None,
)
captured_second_call[:] = messages
return LLMResponse(content="done", tool_calls=[], usage={})
return LLMResponse(content="done", tool_calls=[], usage=None)
provider.chat_with_retry = chat_with_retry
tools = MagicMock()
@@ -751,12 +909,12 @@ async def test_runner_retries_empty_final_response_with_summary_prompt():
return LLMResponse(
content=None,
tool_calls=[],
usage={"prompt_tokens": 5, "completion_tokens": 1},
usage=LLMUsage.reported(input_tokens=5, output_tokens=1),
)
return LLMResponse(
content="final answer",
tool_calls=[],
usage={"prompt_tokens": 3, "completion_tokens": 7},
usage=LLMUsage.reported(input_tokens=3, output_tokens=7),
)
provider.chat_with_retry = chat_with_retry
@@ -778,8 +936,9 @@ async def test_runner_retries_empty_final_response_with_summary_prompt():
assert calls[0]["tools"] is not None
assert calls[1]["tools"] is not None
assert calls[2]["tools"] is None
assert result.usage["prompt_tokens"] == 13
assert result.usage["completion_tokens"] == 9
assert result.usage is not None
assert result.usage.input_tokens == 13
assert result.usage.output_tokens == 9
@pytest.mark.asyncio
@@ -851,7 +1010,7 @@ async def test_runner_uses_specific_message_after_empty_finalization_retry():
provider = MagicMock(spec=LLMProvider)
async def chat_with_retry(*, messages, **kwargs):
return LLMResponse(content=None, tool_calls=[], usage={})
return LLMResponse(content=None, tool_calls=[], usage=None)
provider.chat_with_retry = chat_with_retry
tools = MagicMock()
@@ -891,14 +1050,14 @@ async def test_empty_finalization_retry_discards_candidate_provider_state():
provider = MagicMock(spec=LLMProvider)
provider.can_resume_conversation_state.return_value = True
provider.chat_with_retry = AsyncMock(side_effect=[
LLMResponse(content=None, tool_calls=[], usage={}),
LLMResponse(content=None, tool_calls=[], usage={}),
LLMResponse(content=None, tool_calls=[], usage=None),
LLMResponse(content=None, tool_calls=[], usage=None),
LLMResponse(
content="finalized without tools",
tool_calls=[ToolCallRequest(id="call_1", name="exec", arguments={})],
finish_reason="stop",
provider_state=candidate,
usage={},
usage=None,
),
])
tools = MagicMock()
@@ -1037,20 +1196,20 @@ async def test_runner_empty_response_does_not_break_tool_chain():
return LLMResponse(
content=None,
tool_calls=[ToolCallRequest(id="tc1", name="read_file", arguments={"path": "a.txt"})],
usage={"prompt_tokens": 10, "completion_tokens": 5},
usage=LLMUsage.reported(input_tokens=10, output_tokens=5),
)
if call_count == 2:
return LLMResponse(content=None, tool_calls=[], usage={"prompt_tokens": 10, "completion_tokens": 1})
return LLMResponse(content=None, tool_calls=[], usage=LLMUsage.reported(input_tokens=10, output_tokens=1))
if call_count == 3:
return LLMResponse(
content=None,
tool_calls=[ToolCallRequest(id="tc2", name="read_file", arguments={"path": "b.txt"})],
usage={"prompt_tokens": 10, "completion_tokens": 5},
usage=LLMUsage.reported(input_tokens=10, output_tokens=5),
)
return LLMResponse(
content="Here are the results.",
tool_calls=[],
usage={"prompt_tokens": 10, "completion_tokens": 10},
usage=LLMUsage.reported(input_tokens=10, output_tokens=10),
)
provider.chat_with_retry = chat_with_retry
@@ -1079,9 +1238,8 @@ async def test_runner_empty_response_does_not_break_tool_chain():
@pytest.mark.asyncio
async def test_runner_accumulates_usage_and_preserves_cached_tokens():
"""Runner should accumulate prompt/completion tokens across iterations
and preserve cached_tokens from provider responses."""
async def test_runner_accumulates_usage_and_preserves_cache_reads():
"""Runner accumulates usage across iterations, including cache reads."""
from nanobot.agent.runner import AgentRunner
provider = MagicMock(spec=LLMProvider)
@@ -1093,12 +1251,12 @@ async def test_runner_accumulates_usage_and_preserves_cached_tokens():
return LLMResponse(
content="thinking",
tool_calls=[ToolCallRequest(id="call_1", name="read_file", arguments={"path": "x"})],
usage={"prompt_tokens": 100, "completion_tokens": 10, "cached_tokens": 80},
usage=LLMUsage.reported(input_tokens=100, output_tokens=10, cache_read_tokens=80),
)
return LLMResponse(
content="done",
tool_calls=[],
usage={"prompt_tokens": 200, "completion_tokens": 20, "cached_tokens": 150},
usage=LLMUsage.reported(input_tokens=200, output_tokens=20, cache_read_tokens=150),
)
provider.chat_with_retry = chat_with_retry
@@ -1116,11 +1274,12 @@ async def test_runner_accumulates_usage_and_preserves_cached_tokens():
))
# Usage should be accumulated across iterations
assert result.usage["prompt_tokens"] == 300 # 100 + 200
assert result.usage["completion_tokens"] == 30 # 10 + 20
assert result.usage["cached_tokens"] == 230 # 80 + 150
assert result.usage["context_tokens"] == 200
assert result.usage["request_count"] == 2
assert result.usage is not None
assert result.usage.input_tokens == 300 # 100 + 200
assert result.usage.output_tokens == 30 # 10 + 20
assert result.usage.cache_read_tokens == 230 # 80 + 150
assert result.usage.context_tokens == 200
assert result.usage.request_count == 2
@pytest.mark.asyncio
@@ -1137,7 +1296,7 @@ async def test_runner_binds_on_retry_wait_to_retry_callback_not_progress():
async def chat_with_retry(**kwargs):
captured.update(kwargs)
return LLMResponse(content="done", tool_calls=[], usage={})
return LLMResponse(content="done", tool_calls=[], usage=None)
provider = MagicMock(spec=LLMProvider)
provider.chat_with_retry = chat_with_retry
@@ -1179,7 +1338,7 @@ async def test_runner_passes_temperature_to_provider():
async def chat_with_retry(**kwargs):
captured.update(kwargs)
return LLMResponse(content="done", tool_calls=[], usage={})
return LLMResponse(content="done", tool_calls=[], usage=None)
provider = MagicMock(spec=LLMProvider)
provider.chat_with_retry = chat_with_retry
@@ -1208,7 +1367,7 @@ async def test_runner_passes_max_tokens_to_provider():
async def chat_with_retry(**kwargs):
captured.update(kwargs)
return LLMResponse(content="done", tool_calls=[], usage={})
return LLMResponse(content="done", tool_calls=[], usage=None)
provider = MagicMock(spec=LLMProvider)
provider.chat_with_retry = chat_with_retry
@@ -1237,7 +1396,7 @@ async def test_runner_passes_reasoning_effort_to_provider():
async def chat_with_retry(**kwargs):
captured.update(kwargs)
return LLMResponse(content="done", tool_calls=[], usage={})
return LLMResponse(content="done", tool_calls=[], usage=None)
provider = MagicMock(spec=LLMProvider)
provider.chat_with_retry = chat_with_retry
+8 -8
View File
@@ -90,7 +90,7 @@ async def test_llm_error_not_appended_to_session_messages():
provider = MagicMock(spec=LLMProvider)
provider.chat_with_retry = AsyncMock(return_value=LLMResponse(
content="429 rate limit exceeded", finish_reason="error", tool_calls=[], usage={},
content="429 rate limit exceeded", finish_reason="error", tool_calls=[], usage=None,
))
tools = MagicMock()
tools.get_definitions.return_value = []
@@ -158,7 +158,7 @@ async def test_runner_ignores_tool_calls_when_finish_reason_blocks_execution(
content="Request blocked by provider policy.",
finish_reason=finish_reason,
tool_calls=[ToolCallRequest(id="call_1", name="exec", arguments={"command": "echo nope"})],
usage={},
usage=None,
))
tools = MagicMock()
tools.get_definitions.return_value = []
@@ -189,7 +189,7 @@ async def test_runner_tool_error_sets_final_content():
return LLMResponse(
content="working",
tool_calls=[ToolCallRequest(id="call_1", name="read_file", arguments={"path": "x"})],
usage={},
usage=None,
)
provider.chat_with_retry = chat_with_retry
@@ -224,9 +224,9 @@ async def test_runner_preserves_successful_exec_output_that_starts_with_error():
tool_calls=[
ToolCallRequest(id="call_1", name="exec", arguments={"command": "report"})
],
usage={},
usage=None,
)
return LLMResponse(content="done", usage={})
return LLMResponse(content="done", usage=None)
provider.chat_with_retry = chat_with_retry
output = "Error: generated report successfully\n\nExit code: 0"
@@ -266,7 +266,7 @@ async def test_runner_tool_error_preserves_tool_results_in_messages():
ToolCallRequest(id="tc1", name="read_file", arguments={"path": "a"}),
ToolCallRequest(id="tc2", name="exec", arguments={"cmd": "bad"}),
],
usage={},
usage=None,
)
provider.chat_with_retry = chat_with_retry
@@ -333,9 +333,9 @@ async def test_length_finish_with_blank_content_routes_to_length_recovery():
content="",
finish_reason="length",
tool_calls=[ToolCallRequest(id="call_1", name="exec", arguments={})],
usage={},
usage=None,
),
LLMResponse(content="done", finish_reason="stop", tool_calls=[], usage={}),
LLMResponse(content="done", finish_reason="stop", tool_calls=[], usage=None),
])
tools = MagicMock()
tools.get_definitions.return_value = []
+36 -1
View File
@@ -78,7 +78,7 @@ class _FakeProvider(LLMProvider):
*,
responses: list[LLMResponse] | None = None,
):
super().__init__()
super().__init__(provider_name=name)
self.name = name
self._response = response or _make_response()
self._responses = iter(responses) if responses is not None else None
@@ -260,6 +260,41 @@ def test_provider_snapshot_uses_smallest_fallback_context_window() -> None:
assert snapshot.provider._primary_context_window_tokens == 128000
def test_factory_injects_configured_identity_into_primary_and_fallback_leaves() -> None:
from nanobot.config.schema import Config
from nanobot.providers.factory import build_provider_snapshot
config = Config.model_validate({
"agents": {
"defaults": {
"modelPreset": "primary",
"fallbackModels": ["backup"],
}
},
"modelPresets": {
"primary": {"model": "primary-model", "provider": "primary_edge"},
"backup": {"model": "backup-model", "provider": "backup_edge"},
},
"providers": {
"primary_edge": {
"apiKey": "primary-key",
"apiBase": "https://primary.example/v1",
},
"backup_edge": {
"apiKey": "backup-key",
"apiBase": "https://backup.example/v1",
},
},
})
snapshot = build_provider_snapshot(config)
assert isinstance(snapshot.provider, FallbackProvider)
assert snapshot.provider._primary.provider_name == "primary_edge"
fallback = snapshot.provider._provider_factory(snapshot.provider._fallback_presets[0])
assert fallback.provider_name == "backup_edge"
def test_inline_fallback_reasoning_effort_does_not_inherit_primary() -> None:
from nanobot.config.schema import Config
from nanobot.providers.factory import provider_signature
+8 -8
View File
@@ -25,7 +25,7 @@ async def test_runner_exits_normally_without_predicate():
provider = MagicMock(spec=LLMProvider)
provider.chat_with_retry = AsyncMock(return_value=LLMResponse(
content="all done", tool_calls=[], usage={},
content="all done", tool_calls=[], usage=None,
))
tools = MagicMock()
tools.get_definitions.return_value = []
@@ -50,7 +50,7 @@ async def test_runner_exits_normally_with_inactive_goal():
provider = MagicMock(spec=LLMProvider)
provider.chat_with_retry = AsyncMock(return_value=LLMResponse(
content="all done", tool_calls=[], usage={},
content="all done", tool_calls=[], usage=None,
))
tools = MagicMock()
tools.get_definitions.return_value = []
@@ -82,7 +82,7 @@ async def test_runner_forces_continue_when_goal_active():
provider = MagicMock(spec=LLMProvider)
provider.chat_with_retry = AsyncMock(return_value=LLMResponse(
content="still working", tool_calls=[], usage={},
content="still working", tool_calls=[], usage=None,
))
tools = MagicMock()
tools.get_definitions.return_value = []
@@ -112,7 +112,7 @@ async def test_runner_respects_max_iterations_even_with_active_goal():
provider = MagicMock(spec=LLMProvider)
provider.chat_with_retry = AsyncMock(return_value=LLMResponse(
content="still working", tool_calls=[], usage={},
content="still working", tool_calls=[], usage=None,
))
tools = MagicMock()
tools.get_definitions.return_value = []
@@ -137,7 +137,7 @@ async def test_runner_goal_continue_not_limited_by_injection_cycle_cap():
provider = MagicMock(spec=LLMProvider)
provider.chat_with_retry = AsyncMock(return_value=LLMResponse(
content="still working", tool_calls=[], usage={},
content="still working", tool_calls=[], usage=None,
))
tools = MagicMock()
tools.get_definitions.return_value = []
@@ -165,7 +165,7 @@ async def test_runner_does_not_force_continue_on_error():
provider = MagicMock(spec=LLMProvider)
provider.chat_with_retry = AsyncMock(return_value=LLMResponse(
content=None, tool_calls=[], usage={},
content=None, tool_calls=[], usage=None,
finish_reason="error",
))
tools = MagicMock()
@@ -191,7 +191,7 @@ async def test_runner_uses_custom_goal_continue_message():
provider = MagicMock(spec=LLMProvider)
provider.chat_with_retry = AsyncMock(return_value=LLMResponse(
content="still working", tool_calls=[], usage={},
content="still working", tool_calls=[], usage=None,
))
tools = MagicMock()
tools.get_definitions.return_value = []
@@ -220,7 +220,7 @@ async def test_runner_resolves_goal_continue_message_lazily():
provider = MagicMock(spec=LLMProvider)
provider.chat_with_retry = AsyncMock(return_value=LLMResponse(
content="still working", tool_calls=[], usage={},
content="still working", tool_calls=[], usage=None,
))
tools = MagicMock()
tools.get_definitions.return_value = []
+3 -3
View File
@@ -273,7 +273,7 @@ async def test_runner_drops_orphan_tool_results_before_model_request():
async def chat_with_retry(*, messages, **kwargs):
captured_messages[:] = messages
return LLMResponse(content="done", tool_calls=[], usage={})
return LLMResponse(content="done", tool_calls=[], usage=None)
provider.chat_with_retry = chat_with_retry
tools = MagicMock()
@@ -312,7 +312,7 @@ async def test_backfill_repairs_model_context_without_shifting_save_turn_boundar
provider = MagicMock()
provider.get_default_model.return_value = "test-model"
response = LLMResponse(content="new answer", tool_calls=[], usage={})
response = LLMResponse(content="new answer", tool_calls=[], usage=None)
provider.chat_with_retry = AsyncMock(return_value=response)
provider.chat_stream_with_retry = AsyncMock(return_value=response)
@@ -397,7 +397,7 @@ async def test_runner_backfill_only_mutates_model_context_not_returned_messages(
async def chat_with_retry(*, messages, **kwargs):
captured_messages[:] = messages
return LLMResponse(content="done", tool_calls=[], usage={})
return LLMResponse(content="done", tool_calls=[], usage=None)
provider.chat_with_retry = chat_with_retry
tools = MagicMock()
+27 -34
View File
@@ -9,7 +9,7 @@ import pytest
from agent.runner_helpers import make_run_spec
from nanobot.config.schema import AgentDefaults
from nanobot.providers.base import LLMProvider, LLMResponse, ToolCallRequest
from nanobot.providers.base import LLMProvider, LLMResponse, LLMUsage, ToolCallRequest
_MAX_TOOL_RESULT_CHARS = AgentDefaults().max_tool_result_chars
@@ -30,7 +30,7 @@ async def test_runner_calls_hooks_in_order():
content="thinking",
tool_calls=[ToolCallRequest(id="call_1", name="list_dir", arguments={"path": "."})],
)
return LLMResponse(content="done", tool_calls=[], usage={})
return LLMResponse(content="done", tool_calls=[], usage=None)
provider.chat_with_retry = chat_with_retry
tools = MagicMock()
@@ -110,7 +110,7 @@ async def test_runner_streaming_hook_receives_deltas_and_end_signal():
async def chat_stream_with_retry(*, on_content_delta, **kwargs):
await on_content_delta("he")
await on_content_delta("llo")
return LLMResponse(content="hello", tool_calls=[], usage={})
return LLMResponse(content="hello", tool_calls=[], usage=None)
provider.chat_stream_with_retry = chat_stream_with_retry
provider.chat_with_retry = AsyncMock()
@@ -155,7 +155,7 @@ async def test_runner_measures_stream_generation_without_time_to_first_token():
await on_content_delta("llo")
return LLMResponse(
content="hello",
usage={"prompt_tokens": 100, "completion_tokens": 12},
usage=LLMUsage.reported(input_tokens=100, output_tokens=12),
)
provider.chat_stream_with_retry = chat_stream_with_retry
@@ -181,10 +181,11 @@ async def test_runner_measures_stream_generation_without_time_to_first_token():
hook=StreamingHook(),
))
assert result.usage["generation_ms"] == 600
assert result.usage["measured_completion_tokens"] == 12
assert result.usage["ttft_ms"] == 200
assert result.usage["timed_requests"] == 1
assert result.usage is not None
assert result.usage.generation_ms == 600
assert result.usage.measured_output_tokens == 12
assert result.usage.ttft_ms == 200
assert result.usage.timed_requests == 1
@pytest.mark.asyncio
@@ -240,23 +241,24 @@ async def test_runner_length_recovery_streams_segments_once_and_returns_all_cont
@pytest.mark.asyncio
async def test_runner_passes_cached_tokens_to_hook_context():
"""Hook context.usage should contain cached_tokens."""
async def test_runner_passes_cache_read_tokens_to_hook_context():
"""Hook context usage preserves a reported cache-read count."""
from nanobot.agent.hook import AgentHook, AgentHookContext
from nanobot.agent.runner import AgentRunner
provider = MagicMock(spec=LLMProvider)
captured_usage: list[dict] = []
captured_usage: list[LLMUsage] = []
class UsageHook(AgentHook):
async def after_iteration(self, context: AgentHookContext) -> None:
captured_usage.append(dict(context.usage))
assert context.usage is not None
captured_usage.append(context.usage)
async def chat_with_retry(**kwargs):
return LLMResponse(
content="done",
tool_calls=[],
usage={"prompt_tokens": 200, "completion_tokens": 20, "cached_tokens": 150},
usage=LLMUsage.reported(input_tokens=200, output_tokens=20, cache_read_tokens=150),
)
provider.chat_with_retry = chat_with_retry
@@ -274,8 +276,8 @@ async def test_runner_passes_cached_tokens_to_hook_context():
))
assert len(captured_usage) == 1
assert captured_usage[0]["cached_tokens"] == 150
assert captured_usage[0]["provider_tokens"] == 220
assert captured_usage[0].cache_read_tokens == 150
assert captured_usage[0].reported_tokens == 220
@pytest.mark.asyncio
@@ -284,14 +286,15 @@ async def test_runner_estimates_usage_when_provider_omits_usage(monkeypatch):
from nanobot.agent.runner import AgentRunner
provider = MagicMock(spec=LLMProvider)
captured_usage: list[dict] = []
captured_usage: list[LLMUsage] = []
class UsageHook(AgentHook):
async def after_iteration(self, context: AgentHookContext) -> None:
captured_usage.append(dict(context.usage))
assert context.usage is not None
captured_usage.append(context.usage)
async def chat_with_retry(**kwargs):
return LLMResponse(content="done", tool_calls=[], usage={})
return LLMResponse(content="done", tool_calls=[], usage=None)
provider.chat_with_retry = chat_with_retry
tools = MagicMock()
@@ -312,11 +315,8 @@ async def test_runner_estimates_usage_when_provider_omits_usage(monkeypatch):
hook=UsageHook(),
))
assert result.usage["prompt_tokens"] == 123
assert result.usage["completion_tokens"] == 7
assert result.usage["total_tokens"] == 130
assert result.usage["estimated_tokens"] == 130
assert captured_usage[0]["estimated_tokens"] == 130
assert result.usage == LLMUsage.estimated(input_tokens=123, output_tokens=7)
assert captured_usage[0].estimated_tokens == 130
@pytest.mark.asyncio
@@ -332,7 +332,7 @@ async def test_runner_calls_run_level_hooks_on_success():
return LLMResponse(
content="done",
tool_calls=[],
usage={"prompt_tokens": 3, "completion_tokens": 2},
usage=LLMUsage.reported(input_tokens=3, output_tokens=2),
)
provider.chat_with_retry = chat_with_retry
@@ -350,7 +350,7 @@ async def test_runner_calls_run_level_hooks_on_success():
context.final_content,
context.stop_reason,
context.error,
dict(context.usage),
context.usage,
[msg["role"] for msg in context.messages],
))
@@ -379,14 +379,7 @@ async def test_runner_calls_run_level_hooks_on_success():
"done",
"completed",
None,
{
"prompt_tokens": 3,
"completion_tokens": 2,
"total_tokens": 5,
"provider_tokens": 5,
"request_count": 1,
"context_tokens": 3,
},
LLMUsage.reported(input_tokens=3, output_tokens=2),
["user", "assistant"],
),
("on_finally", "completed", None),
@@ -410,7 +403,7 @@ async def test_runner_run_level_context_is_detached_snapshot():
content="thinking",
tool_calls=[ToolCallRequest(id="call_1", name="list_dir", arguments={"path": "."})],
)
return LLMResponse(content="done", tool_calls=[], usage={})
return LLMResponse(content="done", tool_calls=[], usage=None)
provider.chat_with_retry = chat_with_retry
tools = MagicMock()
+28 -28
View File
@@ -263,9 +263,9 @@ async def test_checkpoint1_injects_after_tool_execution():
return LLMResponse(
content="using tool",
tool_calls=[ToolCallRequest(id="c1", name="read_file", arguments={"path": "x"})],
usage={},
usage=None,
)
return LLMResponse(content="final answer", tool_calls=[], usage={})
return LLMResponse(content="final answer", tool_calls=[], usage=None)
provider.chat_with_retry = chat_with_retry
tools = MagicMock()
@@ -323,8 +323,8 @@ async def test_checkpoint2_injects_after_final_response_with_resuming_stream():
async def chat_stream_with_retry(*, messages, on_content_delta=None, **kwargs):
call_count["n"] += 1
if call_count["n"] == 1:
return LLMResponse(content="first answer", tool_calls=[], usage={})
return LLMResponse(content="second answer", tool_calls=[], usage={})
return LLMResponse(content="first answer", tool_calls=[], usage=None)
return LLMResponse(content="second answer", tool_calls=[], usage=None)
provider.chat_stream_with_retry = chat_stream_with_retry
tools = MagicMock()
@@ -411,8 +411,8 @@ async def test_checkpoint2_preserves_final_response_in_history_before_followup()
call_count["n"] += 1
captured_messages.append([dict(message) for message in messages])
if call_count["n"] == 1:
return LLMResponse(content="first answer", tool_calls=[], usage={})
return LLMResponse(content="second answer", tool_calls=[], usage={})
return LLMResponse(content="first answer", tool_calls=[], usage=None)
return LLMResponse(content="second answer", tool_calls=[], usage=None)
provider.chat_with_retry = chat_with_retry
tools = MagicMock()
@@ -474,8 +474,8 @@ async def test_loop_injected_followup_preserves_image_media(tmp_path):
call_count["n"] += 1
captured_messages.append(list(messages))
if call_count["n"] == 1:
return LLMResponse(content="first answer", tool_calls=[], usage={})
return LLMResponse(content="second answer", tool_calls=[], usage={})
return LLMResponse(content="first answer", tool_calls=[], usage=None)
return LLMResponse(content="second answer", tool_calls=[], usage=None)
provider.chat_with_retry = chat_with_retry
loop = AgentLoop(bus=bus, provider=provider, workspace=tmp_path, model="test-model")
@@ -528,8 +528,8 @@ async def test_pending_injection_resolves_its_own_runtime_context(tmp_path):
provider = MagicMock()
provider.get_default_model.return_value = "test-model"
provider.chat_with_retry = AsyncMock(side_effect=[
LLMResponse(content="first answer", tool_calls=[], usage={}),
LLMResponse(content="second answer", tool_calls=[], usage={}),
LLMResponse(content="first answer", tool_calls=[], usage=None),
LLMResponse(content="second answer", tool_calls=[], usage=None),
])
loop = AgentLoop(
bus=MessageBus(),
@@ -653,8 +653,8 @@ async def test_subagent_pending_injection_is_hidden_history_and_not_merged(tmp_p
async def chat_with_retry(*, messages, **kwargs):
call_count["n"] += 1
if call_count["n"] == 1:
return LLMResponse(content="first answer", tool_calls=[], usage={})
return LLMResponse(content="second answer", tool_calls=[], usage={})
return LLMResponse(content="first answer", tool_calls=[], usage=None)
return LLMResponse(content="second answer", tool_calls=[], usage=None)
provider.chat_with_retry = chat_with_retry
loop = AgentLoop(bus=bus, provider=provider, workspace=tmp_path, model="test-model")
@@ -714,8 +714,8 @@ async def test_runner_merges_multiple_injected_user_messages_without_losing_medi
call_count["n"] += 1
captured_messages.append([dict(message) for message in messages])
if call_count["n"] == 1:
return LLMResponse(content="first answer", tool_calls=[], usage={})
return LLMResponse(content="second answer", tool_calls=[], usage={})
return LLMResponse(content="first answer", tool_calls=[], usage=None)
return LLMResponse(content="second answer", tool_calls=[], usage=None)
provider.chat_with_retry = chat_with_retry
tools = MagicMock()
@@ -841,7 +841,7 @@ async def test_injection_cycles_capped_at_max():
async def chat_with_retry(*, messages, **kwargs):
call_count["n"] += 1
return LLMResponse(content=f"answer-{call_count['n']}", tool_calls=[], usage={})
return LLMResponse(content=f"answer-{call_count['n']}", tool_calls=[], usage=None)
provider.chat_with_retry = chat_with_retry
tools = MagicMock()
@@ -879,7 +879,7 @@ async def test_no_injections_flag_is_false_by_default():
provider = MagicMock()
async def chat_with_retry(**kwargs):
return LLMResponse(content="done", tool_calls=[], usage={})
return LLMResponse(content="done", tool_calls=[], usage=None)
provider.chat_with_retry = chat_with_retry
tools = MagicMock()
@@ -903,7 +903,7 @@ async def test_pending_queue_cleanup_on_dispatch(tmp_path):
loop = _make_loop(tmp_path)
async def chat_with_retry(**kwargs):
return LLMResponse(content="done", tool_calls=[], usage={})
return LLMResponse(content="done", tool_calls=[], usage=None)
loop.provider.chat_with_retry = chat_with_retry
@@ -1329,7 +1329,7 @@ async def test_pending_queue_preserves_overflow_for_next_injection_cycle(tmp_pat
async def chat_with_retry(*, messages, **kwargs):
call_count["n"] += 1
captured_messages.append([dict(message) for message in messages])
return LLMResponse(content=f"answer-{call_count['n']}", tool_calls=[], usage={})
return LLMResponse(content=f"answer-{call_count['n']}", tool_calls=[], usage=None)
provider.chat_with_retry = chat_with_retry
loop = AgentLoop(bus=bus, provider=provider, workspace=tmp_path, model="test-model")
@@ -1502,16 +1502,16 @@ async def test_drain_injections_on_fatal_tool_error():
return LLMResponse(
content="stale prefix ",
finish_reason="length",
usage={},
usage=None,
)
if call_count["n"] == 2:
return LLMResponse(
content="",
tool_calls=[ToolCallRequest(id="c1", name="exec", arguments={"cmd": "bad"})],
usage={},
usage=None,
)
# Third call: respond normally to the injected follow-up.
return LLMResponse(content="reply to follow-up", tool_calls=[], usage={})
return LLMResponse(content="reply to follow-up", tool_calls=[], usage=None)
provider.chat_with_retry = chat_with_retry
tools = MagicMock()
@@ -1563,10 +1563,10 @@ async def test_drain_injections_on_llm_error():
content=None,
tool_calls=[],
finish_reason="error",
usage={},
usage=None,
)
# Second call: respond normally to the injected follow-up
return LLMResponse(content="recovered answer", tool_calls=[], usage={})
return LLMResponse(content="recovered answer", tool_calls=[], usage=None)
provider.chat_with_retry = chat_with_retry
tools = MagicMock()
@@ -1614,9 +1614,9 @@ async def test_drain_injections_on_empty_final_response():
async def chat_with_retry(*, messages, **kwargs):
call_count["n"] += 1
if call_count["n"] <= _MAX_EMPTY_RETRIES + 1:
return LLMResponse(content="", tool_calls=[], usage={})
return LLMResponse(content="", tool_calls=[], usage=None)
# After retries exhausted + injection drain, respond normally
return LLMResponse(content="answer after empty", tool_calls=[], usage={})
return LLMResponse(content="answer after empty", tool_calls=[], usage=None)
provider.chat_with_retry = chat_with_retry
tools = MagicMock()
@@ -1671,7 +1671,7 @@ async def test_drain_injections_on_max_iterations():
return LLMResponse(
content="",
tool_calls=[ToolCallRequest(id=f"c{call_count['n']}", name="read_file", arguments={"path": "x"})],
usage={},
usage=None,
)
provider.chat_with_retry = chat_with_retry
@@ -1723,7 +1723,7 @@ async def test_drain_injections_set_flag_when_followup_arrives_after_last_iterat
return LLMResponse(
content="",
tool_calls=[ToolCallRequest(id=f"c{call_count['n']}", name="read_file", arguments={"path": "x"})],
usage={},
usage=None,
)
provider.chat_with_retry = chat_with_retry
@@ -1786,7 +1786,7 @@ async def test_injection_cycle_cap_on_error_path():
content=None,
tool_calls=[],
finish_reason="error",
usage={},
usage=None,
)
provider.chat_with_retry = chat_with_retry
+7 -7
View File
@@ -8,7 +8,7 @@ from unittest.mock import AsyncMock, MagicMock, patch
from agent.runner_helpers import make_run_spec
from nanobot.config.schema import AgentDefaults
from nanobot.providers.base import LLMResponse, ToolCallRequest
from nanobot.providers.base import LLMResponse, LLMUsage, ToolCallRequest
_MAX_TOOL_RESULT_CHARS = AgentDefaults().max_tool_result_chars
@@ -25,10 +25,10 @@ async def test_runner_persists_large_tool_results_for_follow_up_calls(tmp_path):
return LLMResponse(
content="working",
tool_calls=[ToolCallRequest(id="call_big", name="list_dir", arguments={"path": "."})],
usage={"prompt_tokens": 5, "completion_tokens": 3},
usage=LLMUsage.reported(input_tokens=5, output_tokens=3),
)
captured_second_call[:] = messages
return LLMResponse(content="done", tool_calls=[], usage={})
return LLMResponse(content="done", tool_calls=[], usage=None)
provider.chat_with_retry = chat_with_retry
tools = MagicMock()
@@ -138,10 +138,10 @@ async def test_read_file_result_is_not_offloaded(tmp_path):
return LLMResponse(
content="reading",
tool_calls=[ToolCallRequest(id="call_rf", name="read_file", arguments={"path": "big.txt"})],
usage={"prompt_tokens": 5, "completion_tokens": 3},
usage=LLMUsage.reported(input_tokens=5, output_tokens=3),
)
captured_second_call[:] = messages
return LLMResponse(content="done", tool_calls=[], usage={})
return LLMResponse(content="done", tool_calls=[], usage=None)
provider.chat_with_retry = chat_with_retry
tools = MagicMock()
@@ -183,10 +183,10 @@ async def test_runner_keeps_going_when_tool_result_persistence_fails():
return LLMResponse(
content="working",
tool_calls=[ToolCallRequest(id="call_1", name="list_dir", arguments={"path": "."})],
usage={"prompt_tokens": 5, "completion_tokens": 3},
usage=LLMUsage.reported(input_tokens=5, output_tokens=3),
)
captured_second_call[:] = messages
return LLMResponse(content="done", tool_calls=[], usage={})
return LLMResponse(content="done", tool_calls=[], usage=None)
provider.chat_with_retry = chat_with_retry
tools = MagicMock()
+10 -10
View File
@@ -23,7 +23,7 @@ async def test_runner_can_disable_provider_progress_delta_streaming():
provider = MagicMock()
provider.supports_progress_deltas = True
provider.chat_with_retry = AsyncMock(
return_value=LLMResponse(content="done", tool_calls=[], usage={})
return_value=LLMResponse(content="done", tool_calls=[], usage=None)
)
provider.chat_stream_with_retry = AsyncMock()
tools = MagicMock()
@@ -59,7 +59,7 @@ async def test_runner_streams_provider_progress_deltas_by_default():
async def chat_stream_with_retry(*, on_content_delta, **kwargs):
await on_content_delta("he")
await on_content_delta("llo")
return LLMResponse(content="hello", tool_calls=[], usage={})
return LLMResponse(content="hello", tool_calls=[], usage=None)
provider.chat_stream_with_retry = chat_stream_with_retry
provider.chat_with_retry = AsyncMock()
@@ -113,7 +113,7 @@ async def test_runner_routes_hosted_tool_events_to_structured_progress():
"result": {"name": "x_semantic_search"},
})
await on_content_delta("done")
return LLMResponse(content="done", tool_calls=[], usage={})
return LLMResponse(content="done", tool_calls=[], usage=None)
provider.chat_stream_with_retry = chat_stream_with_retry
provider.chat_with_retry = AsyncMock()
@@ -264,9 +264,9 @@ async def test_runner_emits_write_file_diff_from_tool_execution_snapshots(tmp_pa
arguments={"path": "big.txt", "content": "line\n" * 24},
)
],
usage={},
usage=None,
)
return LLMResponse(content="done", tool_calls=[], usage={})
return LLMResponse(content="done", tool_calls=[], usage=None)
provider.chat_stream_with_retry = chat_stream_with_retry
provider.chat_with_retry = AsyncMock()
@@ -338,9 +338,9 @@ async def test_runner_emits_edit_file_diff_from_tool_execution_snapshots(tmp_pat
},
)
],
usage={},
usage=None,
)
return LLMResponse(content="done", tool_calls=[], usage={})
return LLMResponse(content="done", tool_calls=[], usage=None)
provider.chat_stream_with_retry = chat_stream_with_retry
provider.chat_with_retry = AsyncMock()
@@ -404,9 +404,9 @@ async def test_runner_marks_file_edit_activity_failed_when_tool_errors(tmp_path)
arguments={"path": "aborted.txt"},
)
],
usage={},
usage=None,
)
return LLMResponse(content="done", tool_calls=[], usage={})
return LLMResponse(content="done", tool_calls=[], usage=None)
provider.chat_stream_with_retry = chat_stream_with_retry
provider.chat_with_retry = AsyncMock()
@@ -469,7 +469,7 @@ async def test_runner_marks_file_edit_activity_failed_when_cancelled(tmp_path):
arguments={"path": "cancelled.txt", "content": "new\n"},
)
],
usage={},
usage=None,
)
provider.chat_stream_with_retry = chat_stream_with_retry
+12 -12
View File
@@ -16,7 +16,7 @@ import pytest
from agent.runner_helpers import make_run_spec
from nanobot.agent.hook import AgentHook, AgentHookContext
from nanobot.config.schema import AgentDefaults
from nanobot.providers.base import LLMResponse, ToolCallRequest
from nanobot.providers.base import LLMResponse, LLMUsage, ToolCallRequest
_MAX_TOOL_RESULT_CHARS = AgentDefaults().max_tool_result_chars
@@ -53,10 +53,10 @@ async def test_runner_preserves_reasoning_fields_in_assistant_history():
tool_calls=[ToolCallRequest(id="call_1", name="list_dir", arguments={"path": "."})],
reasoning_content="hidden reasoning",
thinking_blocks=[{"type": "thinking", "thinking": "step"}],
usage={"prompt_tokens": 5, "completion_tokens": 3},
usage=LLMUsage.reported(input_tokens=5, output_tokens=3),
)
captured_second_call[:] = messages
return LLMResponse(content="done", tool_calls=[], usage={})
return LLMResponse(content="done", tool_calls=[], usage=None)
provider.chat_with_retry = chat_with_retry
tools = MagicMock()
@@ -99,7 +99,7 @@ async def test_runner_emits_anthropic_thinking_blocks():
{"type": "thinking", "thinking": "After careful consideration.", "signature": "sig2"},
],
tool_calls=[],
usage={"prompt_tokens": 5, "completion_tokens": 3},
usage=LLMUsage.reported(input_tokens=5, output_tokens=3),
)
provider.chat_with_retry = chat_with_retry
@@ -135,7 +135,7 @@ async def test_runner_emits_inline_think_content_as_reasoning():
return LLMResponse(
content="<think>Let me think about this...\nThe answer is 42.</think>The answer is 42.",
tool_calls=[],
usage={"prompt_tokens": 5, "completion_tokens": 3},
usage=LLMUsage.reported(input_tokens=5, output_tokens=3),
)
provider.chat_with_retry = chat_with_retry
@@ -171,7 +171,7 @@ async def test_runner_prefers_reasoning_content_over_inline_think():
content="<think>inline thinking</think>The answer.",
reasoning_content="dedicated reasoning field",
tool_calls=[],
usage={"prompt_tokens": 5, "completion_tokens": 3},
usage=LLMUsage.reported(input_tokens=5, output_tokens=3),
)
provider.chat_with_retry = chat_with_retry
@@ -211,7 +211,7 @@ async def test_runner_emits_reasoning_content_even_when_answer_was_streamed():
content="The answer.",
reasoning_content="step-by-step deduction",
tool_calls=[],
usage={"prompt_tokens": 5, "completion_tokens": 3},
usage=LLMUsage.reported(input_tokens=5, output_tokens=3),
)
provider.chat_stream_with_retry = chat_stream_with_retry
@@ -257,7 +257,7 @@ async def test_runner_does_not_double_emit_when_inline_think_already_streamed():
return LLMResponse(
content="<think>working...</think>The answer.",
tool_calls=[],
usage={"prompt_tokens": 5, "completion_tokens": 3},
usage=LLMUsage.reported(input_tokens=5, output_tokens=3),
)
provider.chat_stream_with_retry = chat_stream_with_retry
@@ -299,7 +299,7 @@ async def test_runner_closes_reasoning_stream_after_one_shot_response():
content="answer",
reasoning_content="hidden thought",
tool_calls=[],
usage={"prompt_tokens": 5, "completion_tokens": 3},
usage=LLMUsage.reported(input_tokens=5, output_tokens=3),
)
provider.chat_with_retry = chat_with_retry
@@ -350,7 +350,7 @@ async def test_runner_streams_native_thinking_deltas_without_post_hoc_dup():
content="done",
tool_calls=[],
thinking_blocks=[{"type": "thinking", "thinking": "part1part2"}],
usage={"prompt_tokens": 1, "completion_tokens": 2},
usage=LLMUsage.reported(input_tokens=1, output_tokens=2),
)
provider.chat_stream_with_retry = chat_stream_with_retry
@@ -387,7 +387,7 @@ async def test_runner_strips_thinking_tags_from_native_thinking_deltas():
await on_thinking_delta("</thinking>")
if on_content_delta:
await on_content_delta("done")
return LLMResponse(content="done", tool_calls=[], usage={})
return LLMResponse(content="done", tool_calls=[], usage=None)
provider.chat_stream_with_retry = chat_stream_with_retry
tools = MagicMock()
@@ -425,7 +425,7 @@ async def test_runner_ignores_empty_thinking_marker_before_final_reasoning():
content="done",
reasoning_content="Preparing final response",
tool_calls=[],
usage={},
usage=None,
)
provider.chat_stream_with_retry = chat_stream_with_retry
+8 -8
View File
@@ -106,7 +106,7 @@ async def _run_optional_tool_response(response: LLMResponse):
calls["n"] += 1
if calls["n"] == 1:
return response
return LLMResponse(content="done", tool_calls=[], usage={})
return LLMResponse(content="done", tool_calls=[], usage=None)
provider.chat_with_retry = chat_with_retry
tools = ToolRegistry()
@@ -272,10 +272,10 @@ async def test_runner_rejects_near_miss_tool_name_without_executing():
)
],
finish_reason="tool_calls",
usage={},
usage=None,
)
captured_second_call[:] = messages
return LLMResponse(content="done", tool_calls=[], usage={})
return LLMResponse(content="done", tool_calls=[], usage=None)
provider.chat_with_retry = chat_with_retry
tools = ToolRegistry()
@@ -403,7 +403,7 @@ async def test_runner_treats_legacy_entry_point_error_prefix_as_tool_error(tmp_p
provider.chat_with_retry = AsyncMock(return_value=LLMResponse(
content="working",
tool_calls=[ToolCallRequest(id="call_1", name="legacy_plugin", arguments={})],
usage={},
usage=None,
))
result = await AgentRunner().run(make_run_spec(provider,
@@ -430,9 +430,9 @@ async def test_runner_preserves_structured_plugin_success_that_starts_with_error
tool_calls=[
ToolCallRequest(id="call_1", name="structured_success_plugin", arguments={})
],
usage={},
usage=None,
),
LLMResponse(content="done", tool_calls=[], usage={}),
LLMResponse(content="done", tool_calls=[], usage=None),
])
result = await AgentRunner().run(make_run_spec(provider,
@@ -466,10 +466,10 @@ async def test_runner_blocks_repeated_external_fetches():
return LLMResponse(
content="working",
tool_calls=[ToolCallRequest(id=f"call_{call_count['n']}", name="web_fetch", arguments={"url": "https://example.com"})],
usage={},
usage=None,
)
captured_final_call[:] = messages
return LLMResponse(content="done", tool_calls=[], usage={})
return LLMResponse(content="done", tool_calls=[], usage=None)
provider.chat_with_retry = chat_with_retry
tools = MagicMock()
+1 -1
View File
@@ -18,7 +18,7 @@ def _loop(tmp_path: Path) -> AgentLoop:
provider.get_default_model.return_value = "test-model"
provider.generation = SimpleNamespace(max_tokens=4096)
provider.chat_with_retry = AsyncMock(
return_value=LLMResponse(content="Reviewed", tool_calls=[], usage={})
return_value=LLMResponse(content="Reviewed", tool_calls=[], usage=None)
)
return AgentLoop(
bus=MessageBus(),
+1 -1
View File
@@ -18,7 +18,7 @@ from nanobot.utils.llm_runtime import LLMRuntime
class RecordingProvider(LLMProvider):
def __init__(self, name: str) -> None:
super().__init__()
super().__init__(provider_name=name)
self.name = name
self.generation = GenerationSettings(max_tokens=256, temperature=0.1)
self.calls: list[str | None] = []
+5 -5
View File
@@ -16,7 +16,7 @@ from nanobot.agent.subagent import (
)
from nanobot.agent.tools.context import current_request_context
from nanobot.bus.queue import MessageBus
from nanobot.providers.base import GenerationSettings, LLMProvider
from nanobot.providers.base import GenerationSettings, LLMProvider, LLMUsage
from nanobot.utils.llm_runtime import LLMRuntime
# ---------------------------------------------------------------------------
@@ -49,7 +49,7 @@ def _make_hook_context(**overrides) -> AgentHookContext:
tool_calls=[],
tool_events=[],
messages=[],
usage={},
usage=None,
error=None,
stop_reason="completed",
final_content="ok",
@@ -97,7 +97,7 @@ class TestSubagentStatus:
assert s.phase == "initializing"
assert s.iteration == 0
assert s.tool_events == []
assert s.usage == {}
assert s.usage is None
assert s.stop_reason is None
assert s.error is None
@@ -677,12 +677,12 @@ class TestSubagentHook:
ctx = _make_hook_context(
iteration=3,
tool_events=[{"name": "read_file", "status": "ok", "detail": ""}],
usage={"prompt_tokens": 100},
usage=LLMUsage.reported(input_tokens=100, output_tokens=0),
)
await hook.after_iteration(ctx)
assert status.iteration == 3
assert len(status.tool_events) == 1
assert status.usage == {"prompt_tokens": 100}
assert status.usage == LLMUsage.reported(input_tokens=100, output_tokens=0)
@pytest.mark.asyncio
async def test_after_iteration_no_status_noop(self):
+10 -9
View File
@@ -15,6 +15,7 @@ from nanobot.agent.tools.self import MyTool
from nanobot.agent.tools.shell import ExecToolConfig
from nanobot.agent.tools.web import WebSearchConfig, WebToolsConfig
from nanobot.config.schema import ModelPresetConfig
from nanobot.providers.base import LLMUsage
# ---------------------------------------------------------------------------
# Helpers
@@ -31,7 +32,7 @@ def _make_mock_loop(**overrides):
loop._start_time = 1000.0
loop.exec_config = ExecToolConfig()
loop.channels_config = MagicMock()
loop._last_usage = {"prompt_tokens": 100, "completion_tokens": 50}
loop._last_usage = LLMUsage.reported(input_tokens=100, output_tokens=50)
loop.last_usage = loop._last_usage
loop._current_iteration = 0
loop.current_iteration = loop._current_iteration
@@ -163,9 +164,9 @@ class TestInspectPathNavigation:
@pytest.mark.asyncio
async def test_inspect_dict_key_via_dotpath(self):
loop = _make_mock_loop()
loop._last_usage = {"prompt_tokens": 100, "completion_tokens": 50}
loop._last_usage = LLMUsage.reported(input_tokens=100, output_tokens=50)
tool = _make_tool(loop=loop)
result = await tool.execute(action="check", key="_last_usage.prompt_tokens")
result = await tool.execute(action="check", key="_last_usage.input_tokens")
assert "100" in result
@pytest.mark.asyncio
@@ -624,7 +625,7 @@ class TestSubagentStatusFormatting:
{"name": "grep", "status": "ok", "detail": "searched ERROR"},
{"name": "exec", "status": "error", "detail": "timeout"},
],
usage={"prompt_tokens": 4500, "completion_tokens": 1200},
usage=LLMUsage.reported(input_tokens=4500, output_tokens=1200),
)
result = MyTool._format_value(status)
assert "abc12345" in result
@@ -698,14 +699,14 @@ class TestSubagentHookStatus:
iteration=5,
messages=[],
tool_events=[{"name": "read_file", "status": "ok", "detail": "ok"}],
usage={"prompt_tokens": 100, "completion_tokens": 50},
usage=LLMUsage.reported(input_tokens=100, output_tokens=50),
)
await hook.after_iteration(context)
assert status.iteration == 5
assert len(status.tool_events) == 1
assert status.tool_events[0]["name"] == "read_file"
assert status.usage == {"prompt_tokens": 100, "completion_tokens": 50}
assert status.usage == LLMUsage.reported(input_tokens=100, output_tokens=50)
@pytest.mark.asyncio
async def test_after_iteration_with_error(self):
@@ -821,7 +822,7 @@ class TestInspectTaskStatuses:
phase="awaiting_tools",
iteration=2,
tool_events=[{"name": "read_file", "status": "ok", "detail": "ok"}],
usage={"prompt_tokens": 500, "completion_tokens": 100},
usage=LLMUsage.reported(input_tokens=500, output_tokens=100),
),
}
tool = _make_tool(loop=loop)
@@ -1127,12 +1128,12 @@ class TestLastUsageInSummary:
tool = _make_tool()
result = await tool.execute(action="check")
assert "_last_usage" in result
assert "prompt_tokens" in result
assert "input_tokens" in result
@pytest.mark.asyncio
async def test_last_usage_not_shown_when_empty(self):
loop = _make_mock_loop()
loop._last_usage = {}
loop._last_usage = None
loop.last_usage = loop._last_usage
tool = _make_tool(loop=loop)
result = await tool.execute(action="check")
+4 -4
View File
@@ -456,7 +456,7 @@ async def test_agent_loop_syncs_updated_max_iterations_before_run(tmp_path):
error=None,
tool_events=[],
messages=[],
usage={},
usage=None,
had_injections=False,
tools_used=[],
)
@@ -501,7 +501,7 @@ async def test_drain_pending_blocks_while_subagents_running(tmp_path):
error=None,
tool_events=[],
messages=[],
usage={},
usage=None,
had_injections=False,
tools_used=[],
provider_state=None,
@@ -587,7 +587,7 @@ async def test_drain_pending_no_block_when_no_subagents(tmp_path):
error=None,
tool_events=[],
messages=[],
usage={},
usage=None,
had_injections=False,
tools_used=[],
provider_state=None,
@@ -637,7 +637,7 @@ async def test_drain_pending_timeout(tmp_path):
error=None,
tool_events=[],
messages=[],
usage={},
usage=None,
had_injections=False,
tools_used=[],
provider_state=None,