andre/crewAI
1
0
Fork 0
mirror of https://github.com/crewAIInc/crewAI.git synced 2026-01-10 00:28:31 +00:00
Code Issues Actions 1 Packages Projects Releases Wiki Activity

Cached prompt tokens on usage metrics

Browse Source
This commit is contained in:
Thiago Moretto
2024-11-13 10:16:30 -03:00
parent bcfcf88e78
commit 36aa4bcb46
5 changed files with 137 additions and 95 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
src/crewai/agents/agent_builder/utilities/base_token_process.py
View File

@@ -4,6 +4,7 @@ from crewai.types.usage_metrics import UsageMetrics
class TokenProcess:
total_tokens: int = 0
prompt_tokens: int = 0
cached_prompt_tokens: int = 0
completion_tokens: int = 0
successful_requests: int = 0
@@ -15,6 +16,10 @@ class TokenProcess:
self.completion_tokens = self.completion_tokens + tokens
self.total_tokens = self.total_tokens + tokens
def sum_cached_prompt_tokens(self, tokens: int):
self.cached_prompt_tokens = self.cached_prompt_tokens + tokens
self.total_tokens = self.total_tokens + tokens
def sum_successful_requests(self, requests: int):
self.successful_requests = self.successful_requests + requests
@@ -22,6 +27,7 @@ class TokenProcess:
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,
)

5
src/crewai/types/usage_metrics.py
View File

@@ -8,6 +8,7 @@ class UsageMetrics(BaseModel):
Attributes:
total_tokens: Total number of tokens used.
prompt_tokens: Number of tokens used in prompts.
cached_prompt_tokens: Number of cached prompt tokens used.
completion_tokens: Number of tokens used in completions.
successful_requests: Number of successful requests made.
"""
@@ -16,6 +17,9 @@ class UsageMetrics(BaseModel):
prompt_tokens: int = Field(
default=0, description="Number of tokens used in prompts."
)
cached_prompt_tokens: int = Field(
default=0, description="Number of cached prompt tokens used."
)
completion_tokens: int = Field(
default=0, description="Number of tokens used in completions."
)
@@ -32,5 +36,6 @@ class UsageMetrics(BaseModel):
"""
self.total_tokens += usage_metrics.total_tokens
self.prompt_tokens += usage_metrics.prompt_tokens
self.cached_prompt_tokens += usage_metrics.cached_prompt_tokens
self.completion_tokens += usage_metrics.completion_tokens
self.successful_requests += usage_metrics.successful_requests

13
src/crewai/utilities/token_counter_callback.py
View File

@@ -1,5 +1,5 @@
from litellm.integrations.custom_logger import CustomLogger
from litellm.types.utils import Usage
from crewai.agents.agent_builder.utilities.base_token_process import TokenProcess
@@ -11,8 +11,11 @@ class TokenCalcHandler(CustomLogger):
if self.token_cost_process is None:
return
usage : Usage = response_obj["usage"]
self.token_cost_process.sum_successful_requests(1)
self.token_cost_process.sum_prompt_tokens(response_obj["usage"].prompt_tokens)
self.token_cost_process.sum_completion_tokens(
response_obj["usage"].completion_tokens
)
self.token_cost_process.sum_prompt_tokens(usage.prompt_tokens)
self.token_cost_process.sum_completion_tokens(usage.completion_tokens)
if usage.prompt_tokens_details:
self.token_cost_process.sum_cached_prompt_tokens(
usage.prompt_tokens_details.cached_tokens
)