Files
crewAI/src/crewai/agents/agent_builder/utilities/base_token_process.py
Brandon Hancock (bhancock_ai) 627bb3f5f6 Brandon/new release cleanup (#1918)
* WIP

* fixes to match enterprise changes
2025-01-18 15:46:41 -03:00

34 lines
1.1 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:
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,
)