mirror of
https://github.com/crewAIInc/crewAI.git
synced 2026-05-04 08:42:38 +00:00
c3ad5887ef
- Update to Python 3.10+ typing across LLM, callbacks, storage, and errors - Complete typing updates for crew_chat and hitl - Add stop attr to mock LLM, suppress test warnings - Add type-ignore for aisuite import
72 lines
2.6 KiB
Python
72 lines
2.6 KiB
Python
"""Token counting callback handler for LLM interactions.
|
|
|
|
This module provides a callback handler that tracks token usage
|
|
for LLM API calls through the litellm library.
|
|
"""
|
|
|
|
import warnings
|
|
from typing import Any
|
|
|
|
from litellm.integrations.custom_logger import CustomLogger
|
|
from litellm.types.utils import Usage
|
|
|
|
from crewai.agents.agent_builder.utilities.base_token_process import TokenProcess
|
|
|
|
|
|
class TokenCalcHandler(CustomLogger):
|
|
"""Handler for calculating and tracking token usage in LLM calls.
|
|
|
|
This handler integrates with litellm's logging system to track
|
|
prompt tokens, completion tokens, and cached tokens across requests.
|
|
|
|
Attributes:
|
|
token_cost_process: The token process tracker to accumulate usage metrics.
|
|
"""
|
|
|
|
def __init__(self, token_cost_process: TokenProcess | None) -> None:
|
|
"""Initialize the token calculation handler.
|
|
|
|
Args:
|
|
token_cost_process: Optional token process tracker for accumulating metrics.
|
|
"""
|
|
self.token_cost_process = token_cost_process
|
|
|
|
def log_success_event(
|
|
self,
|
|
kwargs: dict[str, Any],
|
|
response_obj: dict[str, Any],
|
|
start_time: float,
|
|
end_time: float,
|
|
) -> None:
|
|
"""Log successful LLM API call and track token usage.
|
|
|
|
Args:
|
|
kwargs: The arguments passed to the LLM call.
|
|
response_obj: The response object from the LLM API.
|
|
start_time: The timestamp when the call started.
|
|
end_time: The timestamp when the call completed.
|
|
"""
|
|
if self.token_cost_process is None:
|
|
return
|
|
|
|
with warnings.catch_warnings():
|
|
warnings.simplefilter("ignore", UserWarning)
|
|
if isinstance(response_obj, dict) and "usage" in response_obj:
|
|
usage: Usage = response_obj["usage"]
|
|
if usage:
|
|
self.token_cost_process.sum_successful_requests(1)
|
|
if hasattr(usage, "prompt_tokens"):
|
|
self.token_cost_process.sum_prompt_tokens(usage.prompt_tokens)
|
|
if hasattr(usage, "completion_tokens"):
|
|
self.token_cost_process.sum_completion_tokens(
|
|
usage.completion_tokens
|
|
)
|
|
if (
|
|
hasattr(usage, "prompt_tokens_details")
|
|
and usage.prompt_tokens_details
|
|
and usage.prompt_tokens_details.cached_tokens
|
|
):
|
|
self.token_cost_process.sum_cached_prompt_tokens(
|
|
usage.prompt_tokens_details.cached_tokens
|
|
)
|