feat: add token usage tracking for tool calls #3327
Conversation
|
Important Review skippedAuto reviews are disabled on this repository. Please check the settings in the CodeRabbit UI or the You can disable this status message by setting the ✨ Finishing touches🧪 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 |
There was a problem hiding this comment.
thanks @lllllyh01 for the PR, the calculation logic needs to be fixed
| output_tokens = self._count_tokens(output_text) | ||
|
|
||
| # Calculate total tokens | ||
| total_tokens = profile.base_tokens + input_tokens + output_tokens |
There was a problem hiding this comment.
It seems the token calculation might be unreliable. Could you clarify the reasoning for setting the base token count to 100?
| self._initialize_default_profiles() | ||
| ) | ||
|
|
||
| def _initialize_default_profiles(self) -> Dict[str, ToolCostProfile]: |
There was a problem hiding this comment.
i think this part is not necessary and hard to maintain
| # Extract text from common result fields | ||
| text_fields = ["content", "text", "message", "result", "output"] |
There was a problem hiding this comment.
would this be able to extract all text information?
| try: | ||
| return len(self.token_counter.encode(text)) | ||
| except Exception as e: | ||
| logger.error(f"Error counting tokens: {e}") | ||
| pass | ||
|
|
||
| return len(text.split()) |
There was a problem hiding this comment.
why not use token_counter.count_tokens_from_messages
| class ToolCategory(Enum): | ||
| """Categories of tools based on their typical token consumption.""" | ||
|
|
||
| # Low-cost tools (typically < 100 tokens) | ||
| SEARCH_API = "search_api" # Google Search, Bing Search, etc. | ||
| SIMPLE_UTILITY = ( | ||
| "simple_utility" # Basic file operations, simple calculations | ||
| ) | ||
|
|
||
| # Medium-cost tools (100-1000 tokens) | ||
| CODE_EXECUTION = "code_execution" # Python execution, shell commands | ||
| DOCUMENT_PROCESSING = "document_processing" # PDF parsing, text analysis | ||
| API_CALLS = "api_calls" # REST API calls, webhooks | ||
|
|
||
| # High-cost tools (1000+ tokens) | ||
| BROWSER_AUTOMATION = ( | ||
| "browser_automation" # Browser interactions, screenshots | ||
| ) | ||
| MULTIMODAL_PROCESSING = ( | ||
| "multimodal_processing" # Image analysis, audio processing | ||
| ) | ||
| LLM_CALLS = "llm_calls" # Sub-agent calls, complex reasoning | ||
|
|
||
|
|
||
| @dataclass | ||
| class ToolCostProfile: | ||
| """Cost profile for a specific tool type.""" | ||
|
|
||
| category: ToolCategory | ||
| base_tokens: int # Base token consumption | ||
|
|
||
|
|
||
| class ToolCostInfo(TypedDict): |
There was a problem hiding this comment.
i think these classes may not necessary
Wendong-Fan
added
Waiting for Update
2 participants
Description
Describe your changes in detail (optional if the linked issue already contains a detailed description of the changes).
Fixes #3219Currently,
ChatAgentonly records token usage on step level. This PR adds finer-grained token usage tracking for each tool calls in one step.Checklist
Go over all the following points, and put an
xin all the boxes that apply.Fixes #issue-numberin the PR description (required)pyproject.tomlanduv lockIf you are unsure about any of these, don't hesitate to ask. We are here to help!