diff --git a/nanobot/api/server.py b/nanobot/api/server.py index 3262e39b..0fd35a97 100644 --- a/nanobot/api/server.py +++ b/nanobot/api/server.py @@ -54,7 +54,14 @@ def _error_json(status: int, message: str, err_type: str = "invalid_request_erro ) -def _chat_completion_response(content: str, model: str) -> dict[str, Any]: +def _chat_completion_response( + content: str, + model: str, + usage: dict[str, int] | None = None, +) -> dict[str, Any]: + prompt = (usage or {}).get("prompt_tokens", 0) + completion = (usage or {}).get("completion_tokens", 0) + total = (usage or {}).get("total_tokens", 0) or prompt + completion return { "id": f"chatcmpl-{uuid.uuid4().hex[:12]}", "object": "chat.completion", @@ -67,7 +74,11 @@ def _chat_completion_response(content: str, model: str) -> dict[str, Any]: "finish_reason": "stop", } ], - "usage": {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0}, + "usage": { + "prompt_tokens": prompt, + "completion_tokens": completion, + "total_tokens": total, + }, } @@ -346,7 +357,9 @@ async def handle_chat_completions(request: web.Request) -> web.Response: logger.exception("Unexpected API lock error for session {}", session_key) return _error_json(500, "Internal server error", err_type="server_error") - return web.json_response(_chat_completion_response(response_text, model_name)) + return web.json_response( + _chat_completion_response(response_text, model_name, getattr(agent_loop, "_last_usage", None)) + ) async def handle_models(request: web.Request) -> web.Response: diff --git a/tests/test_api_attachment.py b/tests/test_api_attachment.py index 92e09ef8..3cdbc847 100644 --- a/tests/test_api_attachment.py +++ b/tests/test_api_attachment.py @@ -32,6 +32,7 @@ def _make_mock_agent(response_text: str = "mock response") -> MagicMock: agent.process_direct = AsyncMock(return_value=response_text) agent._connect_mcp = AsyncMock() agent.close_mcp = AsyncMock() + agent._last_usage = {} return agent diff --git a/tests/test_api_stream.py b/tests/test_api_stream.py index b98e5b8d..f23339bf 100644 --- a/tests/test_api_stream.py +++ b/tests/test_api_stream.py @@ -75,6 +75,7 @@ def _make_streaming_agent(tokens: list[str]) -> MagicMock: return " ".join(tokens) agent.process_direct = fake_process_direct + agent._last_usage = {} return agent @@ -133,6 +134,7 @@ async def test_stream_false_returns_json(aiohttp_client) -> None: agent.process_direct = AsyncMock(return_value="normal reply") agent._connect_mcp = AsyncMock() agent.close_mcp = AsyncMock() + agent._last_usage = {} app = create_app(agent, model_name="m") client = await aiohttp_client(app) @@ -155,6 +157,7 @@ async def test_stream_default_is_false(aiohttp_client) -> None: agent.process_direct = AsyncMock(return_value="default reply") agent._connect_mcp = AsyncMock() agent.close_mcp = AsyncMock() + agent._last_usage = {} app = create_app(agent, model_name="m") client = await aiohttp_client(app) @@ -209,6 +212,7 @@ async def test_stream_passes_on_stream_callbacks(aiohttp_client) -> None: agent.process_direct = fake_process_direct agent._connect_mcp = AsyncMock() agent.close_mcp = AsyncMock() + agent._last_usage = {} app = create_app(agent, model_name="m") client = await aiohttp_client(app) @@ -241,6 +245,7 @@ async def test_stream_segment_end_does_not_close_sse(aiohttp_client) -> None: agent.process_direct = fake_process_direct agent._connect_mcp = AsyncMock() agent.close_mcp = AsyncMock() + agent._last_usage = {} app = create_app(agent, model_name="m") client = await aiohttp_client(app) @@ -279,6 +284,7 @@ async def test_stream_uses_final_response_when_no_deltas(aiohttp_client) -> None agent.process_direct = fake_process_direct agent._connect_mcp = AsyncMock() agent.close_mcp = AsyncMock() + agent._last_usage = {} app = create_app(agent, model_name="m") client = await aiohttp_client(app) @@ -320,6 +326,7 @@ async def test_stream_with_session_id(aiohttp_client) -> None: agent.process_direct = fake_process_direct agent._connect_mcp = AsyncMock() agent.close_mcp = AsyncMock() + agent._last_usage = {} app = create_app(agent, model_name="m") client = await aiohttp_client(app) @@ -348,6 +355,7 @@ async def test_streaming_backend_failure_does_not_emit_success_terminator(aiohtt agent.process_direct = boom agent._connect_mcp = AsyncMock() agent.close_mcp = AsyncMock() + agent._last_usage = {} app = create_app(agent, model_name="m") client = await aiohttp_client(app) diff --git a/tests/test_openai_api.py b/tests/test_openai_api.py index 839c1705..ba75435a 100644 --- a/tests/test_openai_api.py +++ b/tests/test_openai_api.py @@ -33,6 +33,7 @@ def _make_mock_agent(response_text: str = "mock response") -> MagicMock: agent.process_direct = AsyncMock(return_value=response_text) agent._connect_mcp = AsyncMock() agent.close_mcp = AsyncMock() + agent._last_usage = {"prompt_tokens": 100, "completion_tokens": 50} return agent @@ -78,6 +79,25 @@ def test_chat_completion_response() -> None: assert result["choices"][0]["message"]["content"] == "hello world" assert result["choices"][0]["finish_reason"] == "stop" assert result["id"].startswith("chatcmpl-") + assert result["usage"]["prompt_tokens"] == 0 + assert result["usage"]["completion_tokens"] == 0 + assert result["usage"]["total_tokens"] == 0 + + +def test_chat_completion_response_with_usage() -> None: + usage = {"prompt_tokens": 150, "completion_tokens": 42} + result = _chat_completion_response("hello world", "test-model", usage) + assert result["usage"]["prompt_tokens"] == 150 + assert result["usage"]["completion_tokens"] == 42 + assert result["usage"]["total_tokens"] == 192 + + +def test_chat_completion_response_preserves_provider_total_usage() -> None: + usage = {"total_tokens": 77} + result = _chat_completion_response("hello world", "test-model", usage) + assert result["usage"]["prompt_tokens"] == 0 + assert result["usage"]["completion_tokens"] == 0 + assert result["usage"]["total_tokens"] == 77 @pytest.mark.skipif(not HAS_AIOHTTP, reason="aiohttp not installed") @@ -213,6 +233,7 @@ async def test_followup_requests_share_same_session_key(aiohttp_client) -> None: agent.process_direct = fake_process agent._connect_mcp = AsyncMock() agent.close_mcp = AsyncMock() + agent._last_usage = {} app = create_app(agent, model_name="m") client = await aiohttp_client(app) @@ -250,6 +271,7 @@ async def test_fixed_session_requests_are_serialized(aiohttp_client) -> None: agent.process_direct = slow_process agent._connect_mcp = AsyncMock() agent.close_mcp = AsyncMock() + agent._last_usage = {} app = create_app(agent, model_name="m") client = await aiohttp_client(app) @@ -364,6 +386,7 @@ async def test_empty_response_retry_then_success(aiohttp_client) -> None: agent.process_direct = sometimes_empty agent._connect_mcp = AsyncMock() agent.close_mcp = AsyncMock() + agent._last_usage = {} app = create_app(agent, model_name="m") client = await aiohttp_client(app) @@ -393,6 +416,7 @@ async def test_empty_response_falls_back(aiohttp_client) -> None: agent.process_direct = always_empty agent._connect_mcp = AsyncMock() agent.close_mcp = AsyncMock() + agent._last_usage = {} app = create_app(agent, model_name="m") client = await aiohttp_client(app)