fix(api): forward real LLM usage in /v1/chat/completions response (#4310)
* fix(api): forward real LLM usage in /v1/chat/completions response _chat_completion_response() hardcoded prompt_tokens/completion_tokens to zero. Now reads agent_loop._last_usage (set by process_direct after every LLM call) and forwards the actual prompt/completion counts. Streaming path is unchanged; usage is only surfaced in non-streaming responses for now. Fixes #4309 * fix: use defensive getattr for _last_usage and add it to all test mock agents - Use getattr(agent_loop, '_last_usage', None) in server.py for safety - Add _last_usage = {} to mock agents in test_api_attachment.py and test_api_stream.py - Prevents AttributeError/500 when mock agents don't have the attribute * fix(api): preserve provider total usage --------- Co-authored-by: michaelxer <michaelxer@users.noreply.github.com> Co-authored-by: Xubin Ren <52506698+Re-bin@users.noreply.github.com>
This commit is contained in:
co-authored by
michaelxer
Xubin Ren
parent
f85101f017
commit
9814a3b9fe
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user