Repository navigation
Python: [Feature]: Non-streaming mode token usage event/hook #4671
Description
Description
For streaming mode, I could fetch ["usage"] from update.contents to get real time per tool call actual token usage for the inner chat/agent run.
But there is no clean way to get per tool call usage event in non-streaming mode. The only way to do it right now is via a monkey patch hook like below.
Any plan to support such feature in a public way?
Code Sample
def attach_per_call_usage_hook(
client: Any,
tracker: UsageTrackingMiddleware,
) -> None:
"""Monkey-patch a chat client to fire usage callbacks after each model call.
This wraps ``client._inner_get_response()`` so that for **non-streaming**
calls, after the underlying API returns a ``ChatResponse``, the tracker's
``_handle_usage`` is called immediately with ``dict(response.usage_details)``.
Also sets ``tracker._per_call_hook_active = True`` so the middleware's
``process()`` avoids firing duplicate callbacks for the final response.
In the agent-framework MRO::
ChatMiddlewareLayer [once]
→ FunctionInvocationLayer [tool loop]
→ ChatTelemetryLayer [per-call] ← same level
→ BaseChatClient.get_response()
→ _inner_get_response() ← we wrap here
This gives us the same per-iteration accuracy as OpenTelemetry spans
without requiring otel to be enabled.
Streaming calls are left unmodified — use ``ResponseStream`` hooks instead.
"""
tracker._per_call_hook_active = True
callback = tracker._handle_usage
original = client._inner_get_response
def _wrapped(*, messages: Any, stream: bool, options: Any, **kwargs: Any) -> Any:
result = original(messages=messages, stream=stream, options=options, **kwargs)
if stream:
return result # Streaming handled by ResponseStream hooks
# Wrap the awaitable to intercept the response
async def _intercept() -> Any:
response = await result
if response and hasattr(response, "usage_details") and response.usage_details:
callback(dict(response.usage_details))
return response
return _intercept()
client._inner_get_response = _wrappedLanguage/SDK
Python
Metadata
Metadata
Type
Projects
- Status
Activity
[-][Feature]: Non-streaming mode token usage event[/-][+]Python: [Feature]: Non-streaming mode token usage event[/+]moonbox3 commentedon Mar 13, 2026
Thanks for raising this. In non-streaming mode,
AgentResponsealready providesusage_detailswith aggregate token counts, and the fullmessageslist includes all intermediate steps (tool calls, tool results, etc.). Could you share more about what specific information you need that isn't available from the response today?0x7c13 commentedon Mar 13, 2026
Yes, the final response has everything but it will only be ready and available at the end of a single run() meaning there is no usage info can be obtained mid-way. Saying agent.run() invokes 100 tool calls and there is no way to view current token usage (context usage) in between tool call 13->14 in non-streaming mode. I need this information to update context display or do something in a customized way say my own "_compaction.py"
[-]Python: [Feature]: Non-streaming mode token usage event[/-][+]Python: [Feature]: Non-streaming mode token usage event/hook[/+]moonbox3 commentedon Mar 13, 2026
Thanks for clarifying. Mid-run usage visibility during
agent.run()isn't available in non-streaming mode by design — it's a single awaitable that returns when complete. The aggregateusage_detailson the response covers all LLM calls made during the run.For mid-run observability, a couple of options:
Content.from_usage()items are emitted per LLM call, so you can track usage as it happensmarkwallace-microsoft commentedon Mar 20, 2026
Jiaqi Liu (@0x7c13), friendly reminder — this issue is waiting on your response. Please share any updates when you get a chance. (This is an automated message.)
0x7c13 commentedon Mar 20, 2026
Fixed in #4746, closing the issue now.