Merge pull request #138 from rowboatlabs/tokens_used

fixed token_used error for other providers
This commit is contained in:
arkml 2025-05-20 18:47:28 +07:00 committed by GitHub
commit df3e06a08d
No known key found for this signature in database
GPG key ID: B5690EEEBB952194

View file

@ -196,12 +196,15 @@ async def run_turn_streamed(
if event.type == "raw_response_event": if event.type == "raw_response_event":
# Handle token usage counting # Handle token usage counting
if hasattr(event.data, 'type') and event.data.type == "response.completed" and hasattr(event.data.response, 'usage'): if hasattr(event.data, 'type') and event.data.type == "response.completed" and hasattr(event.data.response, 'usage'):
try:
tokens_used["total"] += event.data.response.usage.total_tokens tokens_used["total"] += event.data.response.usage.total_tokens
tokens_used["prompt"] += event.data.response.usage.input_tokens tokens_used["prompt"] += event.data.response.usage.input_tokens
tokens_used["completion"] += event.data.response.usage.output_tokens tokens_used["completion"] += event.data.response.usage.output_tokens
print('-'*50) print('-'*50)
print(f"Found usage information. Updated cumulative tokens: {tokens_used}") print(f"Found usage information. Updated cumulative tokens: {tokens_used}")
print('-'*50) print('-'*50)
except Exception as e:
print(f"Warning: Tokens used is likely not available for your chosen model: {e}")
web_search_messages = handle_web_search_event(event, current_agent) web_search_messages = handle_web_search_event(event, current_agent)
for message in web_search_messages: for message in web_search_messages: