mirror of
https://github.com/crewAIInc/crewAI.git
synced 2026-01-12 09:38:31 +00:00
4e84b98ac2
- Fix import sorting in tests - Add docstrings and type validation - Add comprehensive test cases - Add validation for negative token counts Addresses review feedback on #2198 Co-Authored-By: Joe Moura <joao@crewai.com>
46 lines
1.5 KiB
Python
46 lines
1.5 KiB
Python
from crewai.types.usage_metrics import UsageMetrics
|
|
|
|
|
|
class TokenProcess:
|
|
def __init__(self) -> None:
|
|
self.total_tokens: int = 0
|
|
self.prompt_tokens: int = 0
|
|
self.cached_prompt_tokens: int = 0
|
|
self.completion_tokens: int = 0
|
|
self.successful_requests: int = 0
|
|
|
|
def sum_prompt_tokens(self, tokens: int) -> None:
|
|
self.prompt_tokens += tokens
|
|
self.total_tokens += tokens
|
|
|
|
def sum_completion_tokens(self, tokens: int) -> None:
|
|
self.completion_tokens += tokens
|
|
self.total_tokens += tokens
|
|
|
|
def sum_cached_prompt_tokens(self, tokens: int | None) -> None:
|
|
"""
|
|
Adds the given token count to cached prompt tokens.
|
|
|
|
Args:
|
|
tokens (int | None): Number of tokens to add. None values are ignored.
|
|
|
|
Raises:
|
|
ValueError: If tokens is negative.
|
|
"""
|
|
if tokens is not None:
|
|
if tokens < 0:
|
|
raise ValueError("Token count cannot be negative")
|
|
self.cached_prompt_tokens += tokens
|
|
|
|
def sum_successful_requests(self, requests: int) -> None:
|
|
self.successful_requests += requests
|
|
|
|
def get_summary(self) -> UsageMetrics:
|
|
return UsageMetrics(
|
|
total_tokens=self.total_tokens,
|
|
prompt_tokens=self.prompt_tokens,
|
|
cached_prompt_tokens=self.cached_prompt_tokens,
|
|
completion_tokens=self.completion_tokens,
|
|
successful_requests=self.successful_requests,
|
|
)
|