fix: implement token tracking in RouterAgent for cost observability#1190
fix: implement token tracking in RouterAgent for cost observability#1190MervinPraison merged 2 commits intomainfrom
Conversation
- Add TokenUsage dataclass with structured token consumption data - Enhance LLM.get_response() with optional return_token_usage parameter - Update RouterAgent to track tokens and costs per model selection - Emit token usage via existing trace infrastructure for observability - Maintain full backward compatibility with zero performance regression Resolves issue #1138 - RouterAgent now provides complete cost visibility for multi-agent routing scenarios, addressing the TODO at line 234. 🤖 Generated with [Claude Code](https://claude.ai/code) Co-Authored-By: Claude <noreply@anthropic.com>
|
@coderabbitai review |
|
/review |
✅ Actions performedReview triggered.
|
Code Review by Qodo
1.
|
|
Claude encountered an error —— View job Analyzing Issue and Implementing Fix
|
|
Claude encountered an error —— View job Analyzing PR and Issue #1138
|
📝 WalkthroughWalkthroughThe changes implement token usage tracking across the LLM and routing layers: a new Changes
Sequence DiagramsequenceDiagram
participant Client as Client/RouterAgent
participant LLM as LLM.get_response()
participant LiteLLM as LiteLLM Response
participant ModelRouter as model_router
participant Emitter as Telemetry Emitter
Client->>LLM: Call with return_token_usage=True
LLM->>LiteLLM: Request response
LiteLLM-->>LLM: Return response + usage metadata
LLM->>LLM: _extract_token_usage() from metadata
LLM-->>Client: Return (response_text, TokenUsage)
Client->>ModelRouter: get_model_info() for cost estimation
ModelRouter-->>Client: Model info + pricing data
Client->>Client: estimate_cost() + aggregate per-model stats
Client->>Emitter: output() with token_usage dict, cost, calls
Emitter-->>Client: Emit trace event (or log failure at debug level)
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes Possibly related issues
Possibly related PRs
Suggested labels
Poem
🚥 Pre-merge checks | ✅ 3✅ Passed checks (3 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![qodo-code-review[bot]](https://avatars.githubusercontent.com/in/484649?s=60&v=4){kind=small}
| # Execute with the selected model, requesting token usage tracking | ||
| result = llm_instance.get_response( | ||
| prompt=full_prompt, | ||
| system_prompt=self._build_system_prompt(), | ||
| tools=tools, |
There was a problem hiding this comment.
2. Routeragent token_usage not persisted 📎 Requirement gap ✧ Quality
RouterAgent computes token_usage and an estimated_cost but only stores them in in-memory model_usage_stats and emits them to trace metadata. It does not persist token/cost data into chat_history or session metadata for later attribution.
Agent Prompt
## Issue description
RouterAgent tracks per-call `token_usage` and `estimated_cost` but does not persist these values into `chat_history` or session metadata, so later attribution/analysis is not possible.
## Issue Context
The project includes a SessionStore that supports per-message `metadata`, and compliance requires storing routing token/cost tracking in chat history or session metadata after routed interactions.
## Fix Focus Areas
- src/praisonai-agents/praisonaiagents/agent/router_agent.py[217-269]
ⓘ Copy this prompt and use it to remediate the issue with your preferred AI generation tools
|
@copilot Do a thorough review of this PR. Read ALL existing reviewer comments above first. Review areas:
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 5
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@src/praisonai-agents/praisonaiagents/agent/router_agent.py`:
- Around line 245-249: The trace currently emits a cumulative estimated cost
(model_usage_stats[model_name]['cost']) after you add the current call's cost,
causing later events to include prior spend; change the trace payload to include
the per-call cost variable you compute (cost =
self.model_router.estimate_cost(...)) instead of the accumulated
model_usage_stats value, or if you intentionally want cumulative, rename the
emitted field to cumulative_estimated_cost; update references around
model_router.get_model_info, estimate_cost, model_usage_stats and where the
trace event is built so the event-level key holds the single-call cost (or the
renamed cumulative field) accordingly.
In `@src/praisonai-agents/praisonaiagents/llm/llm.py`:
- Around line 1606-1607: The code only sets _final_llm_response in non-streaming
Chat Completions branches so Responses API and successful streaming flows lose
usage data; modify the Responses API branches to capture the raw resp (the
OpenAI Responses API return) into _final_llm_response, and update the streaming
helper(s) that currently synthesize final_response to return/propagate the
terminal raw response or at least its usage object back to the caller (instead
of rebuilding a usage-less dict); ensure callers that honor return_token_usage
read usage from _final_llm_response (or the value threaded out from the
streaming helpers) so return_token_usage=True yields correct metrics for models
observed by RouterAgent.
- Around line 1600-1602: The async method get_response_async currently lacks the
return_token_usage parameter and still declares a return type of str; update
get_response_async to mirror the sync variant by adding return_token_usage: bool
= False to its signature, change its declared return type to Union[str,
tuple[str, TokenUsage]], and ensure the implementation collects TokenUsage (same
structure/type used by the sync get_response) and returns (response_text,
token_usage) when return_token_usage is True, otherwise just response_text;
locate the implementation inside get_response_async and propagate the flag
through any helper calls so token accounting is computed in the async path.
- Around line 3058-3068: Move the nested helper into a private instance method
named _prepare_return_value(self, response_text: str) at the class level (not
inside get_response), update all call sites (e.g., get_response,
_process_stream_delta and other methods that call it) to use
self._prepare_return_value(...), and remove the infinite recursion by making the
method simply return response_text when return_token_usage is False and return
(response_text, token_usage) when True—compute token_usage by calling
self._extract_token_usage(_final_llm_response) (falling back to TokenUsage() if
None) and do not call _prepare_return_value recursively.
- Around line 4245-4286: The _extract_token_usage method fails to handle
Responses API field names (input_tokens/output_tokens), so update TokenUsage
extraction in both dict and object branches of _extract_token_usage to: if
prompt_tokens/completion_tokens are zero or missing, fall back to input_tokens
and output_tokens respectively; if total_tokens is missing or zero, compute it
as the sum of prompt/completion (or input/output) tokens; preserve other fields
(cached_tokens, reasoning_tokens, audio_*). Modify the dict branch (where
usage.get(...) is used) and the object branch (where getattr(usage, '...', 0) is
used) to implement these fallbacks so RouterAgent cost calculations
(estimate_cost) see correct token counts.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: edfa84b7-0830-49f3-8cc1-79b98d9cfece
📒 Files selected for processing (3)
src/praisonai-agents/praisonaiagents/agent/router_agent.pysrc/praisonai-agents/praisonaiagents/llm/__init__.pysrc/praisonai-agents/praisonaiagents/llm/llm.py
| # Calculate and store cost estimate | ||
| model_info = self.model_router.get_model_info(model_name) | ||
| if model_info and token_usage.total_tokens > 0: | ||
| cost = self.model_router.estimate_cost(model_name, token_usage.total_tokens) | ||
| self.model_usage_stats[model_name]['cost'] += cost |
There was a problem hiding this comment.
Emit per-decision cost in the trace event.
estimated_cost is populated with the running model total after accumulation. From the second call onward, each trace event re-includes earlier spend, so any event-level aggregation will overcount. Emit the current call's cost here, or rename the field to cumulative_estimated_cost.
Suggested patch
- model_info = self.model_router.get_model_info(model_name)
- if model_info and token_usage.total_tokens > 0:
- cost = self.model_router.estimate_cost(model_name, token_usage.total_tokens)
+ cost = 0.0
+ model_info = self.model_router.get_model_info(model_name)
+ if model_info and token_usage.total_tokens > 0:
+ cost = self.model_router.estimate_cost(model_name, token_usage.total_tokens)
self.model_usage_stats[model_name]['cost'] += cost- 'estimated_cost': self.model_usage_stats[model_name]['cost'],
+ 'estimated_cost': cost,
+ 'cumulative_estimated_cost': self.model_usage_stats[model_name]['cost'],Also applies to: 251-263
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@src/praisonai-agents/praisonaiagents/agent/router_agent.py` around lines 245
- 249, The trace currently emits a cumulative estimated cost
(model_usage_stats[model_name]['cost']) after you add the current call's cost,
causing later events to include prior spend; change the trace payload to include
the per-call cost variable you compute (cost =
self.model_router.estimate_cost(...)) instead of the accumulated
model_usage_stats value, or if you intentionally want cumulative, rename the
emitted field to cumulative_estimated_cost; update references around
model_router.get_model_info, estimate_cost, model_usage_stats and where the
trace event is built so the event-level key holds the single-call cost (or the
renamed cumulative field) accordingly.
| return_token_usage: bool = False, | ||
| **kwargs | ||
| ) -> str: | ||
| ) -> Union[str, tuple[str, TokenUsage]]: |
There was a problem hiding this comment.
Mirror return_token_usage in get_response_async().
The sync API now exposes token usage, but the async counterpart still advertises -> str and has no matching flag/tuple contract. That makes cost observability depend on whether the caller used the sync or async path. As per coding guidelines, "All I/O operations must have both sync and async variants; never block the event loop with sync I/O in async context; use asyncio primitives for coordination, not threading".
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@src/praisonai-agents/praisonaiagents/llm/llm.py` around lines 1600 - 1602,
The async method get_response_async currently lacks the return_token_usage
parameter and still declares a return type of str; update get_response_async to
mirror the sync variant by adding return_token_usage: bool = False to its
signature, change its declared return type to Union[str, tuple[str,
TokenUsage]], and ensure the implementation collects TokenUsage (same
structure/type used by the sync get_response) and returns (response_text,
token_usage) when return_token_usage is True, otherwise just response_text;
locate the implementation inside get_response_async and propagate the flag
through any helper calls so token accounting is computed in the async path.
| # Variable to store final response for token usage extraction | ||
| _final_llm_response = None |
There was a problem hiding this comment.
Most execution paths still drop the raw usage payload.
_final_llm_response is only populated in the non-streaming Chat Completions branches added here. The OpenAI Responses API flow and the successful streaming flow still synthesize final_response objects without any usage, so return_token_usage=True falls back to empty metrics for exactly the models RouterAgent now observes.
Please capture resp in the Responses API branches as well, and thread the terminal raw response or usage object out of the streaming helpers instead of rebuilding a usage-less dict.
Also applies to: 1903-1903, 2134-2134, 2326-2326
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@src/praisonai-agents/praisonaiagents/llm/llm.py` around lines 1606 - 1607,
The code only sets _final_llm_response in non-streaming Chat Completions
branches so Responses API and successful streaming flows lose usage data; modify
the Responses API branches to capture the raw resp (the OpenAI Responses API
return) into _final_llm_response, and update the streaming helper(s) that
currently synthesize final_response to return/propagate the terminal raw
response or at least its usage object back to the caller (instead of rebuilding
a usage-less dict); ensure callers that honor return_token_usage read usage from
_final_llm_response (or the value threaded out from the streaming helpers) so
return_token_usage=True yields correct metrics for models observed by
RouterAgent.
| # Helper function to return appropriate format based on return_token_usage | ||
| def _prepare_return_value(response_text: str) -> Union[str, tuple[str, TokenUsage]]: | ||
| if not return_token_usage: | ||
| return _prepare_return_value(response_text) | ||
|
|
||
| token_usage = self._extract_token_usage(_final_llm_response) if _final_llm_response else None | ||
| if token_usage is None: | ||
| token_usage = TokenUsage() # Return empty TokenUsage if extraction fails | ||
|
|
||
| return _prepare_return_value(response_text), token_usage | ||
|
|
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
# First, let's look at the target file and lines
wc -l src/praisonai-agents/praisonaiagents/llm/llm.pyRepository: MervinPraison/PraisonAI
Length of output: 119
🏁 Script executed:
# Get the code around lines 3058-3068 with broader context
sed -n '3000,3100p' src/praisonai-agents/praisonaiagents/llm/llm.py | cat -nRepository: MervinPraison/PraisonAI
Length of output: 6057
🏁 Script executed:
# Search for _prepare_return_value calls in the file to identify usage order
rg '_prepare_return_value' src/praisonai-agents/praisonaiagents/llm/llm.py -nRepository: MervinPraison/PraisonAI
Length of output: 1783
🏁 Script executed:
# Also search for _process_stream_delta, _extract_from_responses_output, response, aresponse to understand context
rg '(_process_stream_delta|_extract_from_responses_output|def response|def aresponse)' src/praisonai-agents/praisonaiagents/llm/llm.py -n | head -20Repository: MervinPraison/PraisonAI
Length of output: 1292
🏁 Script executed:
# Get broader context around the function definition
sed -n '3040,3080p' src/praisonai-agents/praisonaiagents/llm/llm.py | cat -nRepository: MervinPraison/PraisonAI
Length of output: 2364
🏁 Script executed:
# Check if this is inside a method (search for def get_response context)
sed -n '2990,3080p' src/praisonai-agents/praisonaiagents/llm/llm.py | cat -nRepository: MervinPraison/PraisonAI
Length of output: 5778
🏁 Script executed:
# Look for what the function should actually return
sed -n '3050,3085p' src/praisonai-agents/praisonaiagents/llm/llm.pyRepository: MervinPraison/PraisonAI
Length of output: 1774
🏁 Script executed:
# Check if _prepare_return_value is defined elsewhere
rg 'def _prepare_return_value' src/praisonai-agents/praisonaiagents/llm/llm.py -nRepository: MervinPraison/PraisonAI
Length of output: 164
🏁 Script executed:
# Look at the method that contains this - find where get_response method starts
rg 'def get_response\(' src/praisonai-agents/praisonaiagents/llm/llm.py -n -B 2 -A 5Repository: MervinPraison/PraisonAI
Length of output: 335
🏁 Script executed:
# Check the context around line 1143 to see what method that first call is in
sed -n '1100,1160p' src/praisonai-agents/praisonaiagents/llm/llm.py | cat -nRepository: MervinPraison/PraisonAI
Length of output: 3427
🏁 Script executed:
# Check which method contains the call at line 1143
sed -n '1100,1150p' src/praisonai-agents/praisonaiagents/llm/llm.py | cat -nRepository: MervinPraison/PraisonAI
Length of output: 2813
🏁 Script executed:
# Verify line numbers - which method does line 1143 belong to
sed -n '1000,1200p' src/praisonai-agents/praisonaiagents/llm/llm.py | grep -E '^\s*(def |class )' | head -5Repository: MervinPraison/PraisonAI
Length of output: 641
🏁 Script executed:
# Get the actual line 1143 in context - what method is it in?
sed -n '1140,1150p' src/praisonai-agents/praisonaiagents/llm/llm.pyRepository: MervinPraison/PraisonAI
Length of output: 532
🏁 Script executed:
# Verify that _prepare_return_value is indeed nested inside get_response by checking indentation
sed -n '2800,2880p' src/praisonai-agents/praisonaiagents/llm/llm.py | grep -E '(def get_response|def _extract_from_responses_output|def _prepare_return_value)' -A 2Repository: MervinPraison/PraisonAI
Length of output: 49
🏁 Script executed:
# Check for any other definition of _prepare_return_value at class level (not nested)
rg '^\s{0,4}def _prepare_return_value' src/praisonai-agents/praisonaiagents/llm/llm.py -nRepository: MervinPraison/PraisonAI
Length of output: 49
🏁 Script executed:
# Let's check the structure - is _prepare_return_value nested inside get_response?
# Get the indentation of the function definition at line 3059
sed -n '3059p' src/praisonai-agents/praisonaiagents/llm/llm.py | od -c | head -5Repository: MervinPraison/PraisonAI
Length of output: 107
🏁 Script executed:
# Compare with the indentation of get_response method definition
sed -n '1575p' src/praisonai-agents/praisonaiagents/llm/llm.py | od -c | head -5Repository: MervinPraison/PraisonAI
Length of output: 107
🏁 Script executed:
# Better approach - see the structure around line 3059
sed -n '3050,3075p' src/praisonai-agents/praisonaiagents/llm/llm.pyRepository: MervinPraison/PraisonAI
Length of output: 1315
🏁 Script executed:
# Check what method contains line 3059 by looking at def keywords before it
sed -n '3000,3059p' src/praisonai-agents/praisonaiagents/llm/llm.py | grep -E 'def ' | tail -2Repository: MervinPraison/PraisonAI
Length of output: 159
🏁 Script executed:
# Verify scope - is _process_stream_delta a class method (not nested)?
sed -n '1112,1150p' src/praisonai-agents/praisonaiagents/llm/llm.py | cat -nRepository: MervinPraison/PraisonAI
Length of output: 2146
🏁 Script executed:
# Confirm _process_stream_delta is a class method and verify it's trying to call _prepare_return_value
sed -n '1100,1115p' src/praisonai-agents/praisonaiagents/llm/llm.pyRepository: MervinPraison/PraisonAI
Length of output: 882
🏁 Script executed:
# Check the structure more clearly - count spaces for indentation
sed -n '1112p' src/praisonai-agents/praisonaiagents/llm/llm.py | sed 's/[^ ].*//' | wc -cRepository: MervinPraison/PraisonAI
Length of output: 68
🏁 Script executed:
# Check indentation of get_response (should be same as _process_stream_delta if both class methods)
sed -n '1575p' src/praisonai-agents/praisonaiagents/llm/llm.py | sed 's/[^ ].*//' | wc -cRepository: MervinPraison/PraisonAI
Length of output: 68
🏁 Script executed:
# Check indentation of _prepare_return_value definition
sed -n '3059p' src/praisonai-agents/praisonaiagents/llm/llm.py | sed 's/[^ ].*//' | wc -cRepository: MervinPraison/PraisonAI
Length of output: 68
🏁 Script executed:
# Let's see the actual spaces more clearly
sed -n '3059p' src/praisonai-agents/praisonaiagents/llm/llm.py | cat -ARepository: MervinPraison/PraisonAI
Length of output: 160
🏁 Script executed:
# Verify the scope issue - check what method line 3059 is nested in
# Find the last 'def' before line 3059 at lower indentation
sed -n '1575,3059p' src/praisonai-agents/praisonaiagents/llm/llm.py | grep -E '^ def ' | tail -1Repository: MervinPraison/PraisonAI
Length of output: 88
🏁 Script executed:
# Confirm the recursive call issue - line 3061 calls itself
sed -n '3059,3068p' src/praisonai-agents/praisonaiagents/llm/llm.pyRepository: MervinPraison/PraisonAI
Length of output: 586
Move _prepare_return_value() to a private instance method and fix infinite recursion.
This nested function is defined at line 3059 inside get_response(), but it's called from class-level methods like _process_stream_delta() (line 1143) and other methods starting at lines 2740, 2829, 2858, 2859, 2975, 2984, 3041, 3052—all before the definition. This causes a NameError at runtime since class methods cannot access local functions nested in other methods.
Additionally, both branches recursively call themselves (lines 3061 and 3067), causing infinite recursion. Line 3061 should return response_text directly when return_token_usage=False, and line 3067 should return the tuple without calling itself again.
Convert _prepare_return_value() to a private instance method (def _prepare_return_value(self, response_text: str)), place it at the class level, and apply the fix at final public return sites only.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@src/praisonai-agents/praisonaiagents/llm/llm.py` around lines 3058 - 3068,
Move the nested helper into a private instance method named
_prepare_return_value(self, response_text: str) at the class level (not inside
get_response), update all call sites (e.g., get_response, _process_stream_delta
and other methods that call it) to use self._prepare_return_value(...), and
remove the infinite recursion by making the method simply return response_text
when return_token_usage is False and return (response_text, token_usage) when
True—compute token_usage by calling
self._extract_token_usage(_final_llm_response) (falling back to TokenUsage() if
None) and do not call _prepare_return_value recursively.
| def _extract_token_usage(self, response: Union[Dict[str, Any], Any]) -> Optional[TokenUsage]: | ||
| """Extract token usage from LiteLLM response for public API.""" | ||
| try: | ||
| usage = None | ||
|
|
||
| # Handle both dict and ModelResponse object formats | ||
| if isinstance(response, dict): | ||
| usage = response.get("usage", {}) | ||
| else: | ||
| # ModelResponse object | ||
| usage = getattr(response, 'usage', None) | ||
|
|
||
| if not usage: | ||
| return None | ||
|
|
||
| # Extract token counts with support for both dict and object access | ||
| if isinstance(usage, dict): | ||
| return TokenUsage( | ||
| prompt_tokens=usage.get("prompt_tokens", 0), | ||
| completion_tokens=usage.get("completion_tokens", 0), | ||
| total_tokens=usage.get("total_tokens", 0), | ||
| cached_tokens=usage.get("cached_tokens", 0), | ||
| reasoning_tokens=usage.get("reasoning_tokens", 0), | ||
| audio_input_tokens=usage.get("audio_input_tokens", 0), | ||
| audio_output_tokens=usage.get("audio_output_tokens", 0), | ||
| ) | ||
| else: | ||
| # Object-style access | ||
| return TokenUsage( | ||
| prompt_tokens=getattr(usage, 'prompt_tokens', 0) or 0, | ||
| completion_tokens=getattr(usage, 'completion_tokens', 0) or 0, | ||
| total_tokens=getattr(usage, 'total_tokens', 0) or 0, | ||
| cached_tokens=getattr(usage, 'cached_tokens', 0) or 0, | ||
| reasoning_tokens=getattr(usage, 'reasoning_tokens', 0) or 0, | ||
| audio_input_tokens=getattr(usage, 'audio_input_tokens', 0) or 0, | ||
| audio_output_tokens=getattr(usage, 'audio_output_tokens', 0) or 0, | ||
| ) | ||
|
|
||
| except Exception as e: | ||
| if self.verbose: | ||
| logging.warning(f"Failed to extract token usage: {e}") | ||
| return None |
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
find . -name "llm.py" -path "*/praisonai-agents/*" | head -5Repository: MervinPraison/PraisonAI
Length of output: 116
🏁 Script executed:
wc -l ./src/praisonai-agents/praisonaiagents/llm/llm.pyRepository: MervinPraison/PraisonAI
Length of output: 121
🏁 Script executed:
sed -n '4245,4286p' ./src/praisonai-agents/praisonaiagents/llm/llm.pyRepository: MervinPraison/PraisonAI
Length of output: 2140
🏁 Script executed:
sed -n '1750,1760p' ./src/praisonai-agents/praisonaiagents/llm/llm.pyRepository: MervinPraison/PraisonAI
Length of output: 794
🏁 Script executed:
rg "RouterAgent" ./src/praisonai-agents/praisonaiagents/ -lRepository: MervinPraison/PraisonAI
Length of output: 184
🏁 Script executed:
cat -n ./src/praisonai-agents/praisonaiagents/agent/router_agent.py | head -100Repository: MervinPraison/PraisonAI
Length of output: 4645
🏁 Script executed:
rg "cost" ./src/praisonai-agents/praisonaiagents/agent/router_agent.py -A 3 -B 3Repository: MervinPraison/PraisonAI
Length of output: 3576
🏁 Script executed:
rg "class TokenUsage" ./src/praisonai-agents/praisonaiagents/ -A 20Repository: MervinPraison/PraisonAI
Length of output: 1811
🏁 Script executed:
rg "_extract_token_usage" ./src/praisonai-agents/praisonaiagents/ -B 5 -A 5Repository: MervinPraison/PraisonAI
Length of output: 2215
🏁 Script executed:
rg "estimate_cost" ./src/praisonai-agents/praisonaiagents/llm/model_router.py -B 2 -A 10Repository: MervinPraison/PraisonAI
Length of output: 541
🏁 Script executed:
sed -n '170,185p' ./src/praisonai-agents/praisonaiagents/agent/router_agent.pyRepository: MervinPraison/PraisonAI
Length of output: 573
🏁 Script executed:
sed -n '190,240p' ./src/praisonai-agents/praisonaiagents/agent/router_agent.pyRepository: MervinPraison/PraisonAI
Length of output: 2043
🏁 Script executed:
sed -n '240,280p' ./src/praisonai-agents/praisonaiagents/agent/router_agent.pyRepository: MervinPraison/PraisonAI
Length of output: 2072
🏁 Script executed:
rg "total_tokens" ./src/praisonai-agents/praisonaiagents/llm/model_router.py -B 2 -A 2Repository: MervinPraison/PraisonAI
Length of output: 49
🏁 Script executed:
sed -n '1750,1770p' ./src/praisonai-agents/praisonaiagents/llm/llm.pyRepository: MervinPraison/PraisonAI
Length of output: 1319
🏁 Script executed:
rg "input_tokens|output_tokens" ./src/praisonai-agents/praisonaiagents/llm/llm.py | head -20Repository: MervinPraison/PraisonAI
Length of output: 1069
Add fallback to Responses API token field names in _extract_token_usage().
The method at lines 4245-4286 only looks for prompt_tokens and completion_tokens, but the codebase already handles the Responses API format with input_tokens and output_tokens at line 1754. When responses use only the Responses API names, the method returns zero tokens, which causes RouterAgent's cost calculation to report 0.0 cost (since estimate_cost() multiplies by token count).
Add fallback logic to check input_tokens/output_tokens if the standard names are not present, and calculate total_tokens as their sum when not explicitly provided in the response.
Suggested patch
if isinstance(usage, dict):
+ prompt_tokens = usage.get("prompt_tokens", usage.get("input_tokens", 0))
+ completion_tokens = usage.get("completion_tokens", usage.get("output_tokens", 0))
return TokenUsage(
- prompt_tokens=usage.get("prompt_tokens", 0),
- completion_tokens=usage.get("completion_tokens", 0),
- total_tokens=usage.get("total_tokens", 0),
+ prompt_tokens=prompt_tokens,
+ completion_tokens=completion_tokens,
+ total_tokens=usage.get("total_tokens", prompt_tokens + completion_tokens),
cached_tokens=usage.get("cached_tokens", 0),
reasoning_tokens=usage.get("reasoning_tokens", 0),
audio_input_tokens=usage.get("audio_input_tokens", 0),
audio_output_tokens=usage.get("audio_output_tokens", 0),
)
else:
+ prompt_tokens = getattr(usage, "prompt_tokens", None)
+ if prompt_tokens is None:
+ prompt_tokens = getattr(usage, "input_tokens", 0) or 0
+ completion_tokens = getattr(usage, "completion_tokens", None)
+ if completion_tokens is None:
+ completion_tokens = getattr(usage, "output_tokens", 0) or 0
return TokenUsage(
- prompt_tokens=getattr(usage, 'prompt_tokens', 0) or 0,
- completion_tokens=getattr(usage, 'completion_tokens', 0) or 0,
- total_tokens=getattr(usage, 'total_tokens', 0) or 0,
+ prompt_tokens=prompt_tokens,
+ completion_tokens=completion_tokens,
+ total_tokens=getattr(usage, 'total_tokens', prompt_tokens + completion_tokens) or (prompt_tokens + completion_tokens),
cached_tokens=getattr(usage, 'cached_tokens', 0) or 0,
reasoning_tokens=getattr(usage, 'reasoning_tokens', 0) or 0,
audio_input_tokens=getattr(usage, 'audio_input_tokens', 0) or 0,
audio_output_tokens=getattr(usage, 'audio_output_tokens', 0) or 0,
)🧰 Tools
🪛 Ruff (0.15.7)
[warning] 4283-4283: Do not catch blind exception: Exception
(BLE001)
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@src/praisonai-agents/praisonaiagents/llm/llm.py` around lines 4245 - 4286,
The _extract_token_usage method fails to handle Responses API field names
(input_tokens/output_tokens), so update TokenUsage extraction in both dict and
object branches of _extract_token_usage to: if prompt_tokens/completion_tokens
are zero or missing, fall back to input_tokens and output_tokens respectively;
if total_tokens is missing or zero, compute it as the sum of prompt/completion
(or input/output) tokens; preserve other fields (cached_tokens,
reasoning_tokens, audio_*). Modify the dict branch (where usage.get(...) is
used) and the object branch (where getattr(usage, '...', 0) is used) to
implement these fallbacks so RouterAgent cost calculations (estimate_cost) see
correct token counts.
|
@copilot Review this PR thoroughly. Check code quality, security, performance, backward compatibility, and naming conventions. |
|
Claude encountered an error —— View job I'll analyze this and get back to you. |
- Move _prepare_return_value closure to before the try block in get_response() so it's defined before use - Fix infinite recursion: both branches called themselves instead of returning the text value - Revert _prepare_return_value usage from 7 out-of-scope methods (_process_stream_delta, get_response_async, _extract_from_responses_output, _stream_responses_api, _stream_responses_api_async, response, aresponse) where it was a NameError and would corrupt return tuple structure Agent-Logs-Url: https://github.com/MervinPraison/PraisonAI/sessions/1b069906-c514-4c5f-8c93-34c345db5fd0 Co-authored-by: MervinPraison <454862+MervinPraison@users.noreply.github.com>
|
@claude Final architectural review. Check: 1) Protocol-driven (no heavy impl in core) 2) Lazy imports 3) Thread/async safety 4) Backward compat 5) No scope creep beyond the issue. Approve or request changes. |
|
Claude encountered an error —— View job I'll analyze this and get back to you. |
_prepare_return_valueplacement (moved before try block — was unreachable dead code afterraise)_prepare_return_value_prepare_return_valuecalls in 7 out-of-scope methods:_process_stream_delta,get_response_async(3 returns),_extract_from_responses_output,_stream_responses_api,_stream_responses_api_async,response,aresponse