mirror of
https://github.com/crewAIInc/crewAI.git
synced 2026-01-09 08:08:32 +00:00
0b0f2d30ab
* sort imports * update --------- Co-authored-by: Brandon Hancock (bhancock_ai) <109994880+bhancockio@users.noreply.github.com> Co-authored-by: Eduardo Chiarotti <dudumelgaco@hotmail.com>
23 lines
902 B
Python
23 lines
902 B
Python
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):
|
|
def __init__(self, token_cost_process: TokenProcess):
|
|
self.token_cost_process = token_cost_process
|
|
|
|
def log_success_event(self, kwargs, response_obj, start_time, end_time):
|
|
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(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
|
|
)
|