mirror of
https://github.com/crewAIInc/crewAI.git
synced 2026-01-27 00:58:13 +00:00
chore: add typing and docstrings to base_token_process module (#3486)
Co-authored-by: Lucas Gomide <lucaslg200@gmail.com>
This commit is contained in:
@@ -1,8 +1,25 @@
|
|||||||
|
"""Token usage tracking utilities.
|
||||||
|
|
||||||
|
This module provides utilities for tracking token consumption and request
|
||||||
|
metrics during agent execution.
|
||||||
|
"""
|
||||||
|
|
||||||
from crewai.types.usage_metrics import UsageMetrics
|
from crewai.types.usage_metrics import UsageMetrics
|
||||||
|
|
||||||
|
|
||||||
class TokenProcess:
|
class TokenProcess:
|
||||||
|
"""Track token usage during agent processing.
|
||||||
|
|
||||||
|
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.
|
||||||
|
"""
|
||||||
|
|
||||||
def __init__(self) -> None:
|
def __init__(self) -> None:
|
||||||
|
"""Initialize token tracking with zero values."""
|
||||||
self.total_tokens: int = 0
|
self.total_tokens: int = 0
|
||||||
self.prompt_tokens: int = 0
|
self.prompt_tokens: int = 0
|
||||||
self.cached_prompt_tokens: int = 0
|
self.cached_prompt_tokens: int = 0
|
||||||
@@ -10,20 +27,45 @@ class TokenProcess:
|
|||||||
self.successful_requests: int = 0
|
self.successful_requests: int = 0
|
||||||
|
|
||||||
def sum_prompt_tokens(self, tokens: int) -> None:
|
def sum_prompt_tokens(self, tokens: int) -> None:
|
||||||
|
"""Add prompt tokens to the running totals.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
tokens: Number of prompt tokens to add.
|
||||||
|
"""
|
||||||
self.prompt_tokens += tokens
|
self.prompt_tokens += tokens
|
||||||
self.total_tokens += tokens
|
self.total_tokens += tokens
|
||||||
|
|
||||||
def sum_completion_tokens(self, tokens: int) -> None:
|
def sum_completion_tokens(self, tokens: int) -> None:
|
||||||
|
"""Add completion tokens to the running totals.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
tokens: Number of completion tokens to add.
|
||||||
|
"""
|
||||||
self.completion_tokens += tokens
|
self.completion_tokens += tokens
|
||||||
self.total_tokens += tokens
|
self.total_tokens += tokens
|
||||||
|
|
||||||
def sum_cached_prompt_tokens(self, tokens: int) -> None:
|
def sum_cached_prompt_tokens(self, tokens: int) -> None:
|
||||||
|
"""Add cached prompt tokens to the running total.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
tokens: Number of cached prompt tokens to add.
|
||||||
|
"""
|
||||||
self.cached_prompt_tokens += tokens
|
self.cached_prompt_tokens += tokens
|
||||||
|
|
||||||
def sum_successful_requests(self, requests: int) -> None:
|
def sum_successful_requests(self, requests: int) -> None:
|
||||||
|
"""Add successful requests to the running total.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
requests: Number of successful requests to add.
|
||||||
|
"""
|
||||||
self.successful_requests += requests
|
self.successful_requests += requests
|
||||||
|
|
||||||
def get_summary(self) -> UsageMetrics:
|
def get_summary(self) -> UsageMetrics:
|
||||||
|
"""Get a summary of all tracked metrics.
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
UsageMetrics object with current totals.
|
||||||
|
"""
|
||||||
return UsageMetrics(
|
return UsageMetrics(
|
||||||
total_tokens=self.total_tokens,
|
total_tokens=self.total_tokens,
|
||||||
prompt_tokens=self.prompt_tokens,
|
prompt_tokens=self.prompt_tokens,
|
||||||
|
|||||||
Reference in New Issue
Block a user