mirror of
https://github.com/crewAIInc/crewAI.git
synced 2026-01-09 08:08:32 +00:00
8ec2eb7d72
- Add proper debug, info, warning, and error methods to Logger class - Ensure warnings and errors are always shown regardless of verbose mode - Fix token process initialization and tracking in Agent class - Update TokenProcess import to use correct class from agent_builder utilities Co-Authored-By: Joe Moura <joao@crewai.com>
75 lines
2.3 KiB
Python
75 lines
2.3 KiB
Python
"""Token processing utility for tracking and managing token usage."""
|
|
from __future__ import annotations
|
|
|
|
from typing import Any, Dict, List, Optional, Union
|
|
|
|
from crewai.types.usage_metrics import UsageMetrics
|
|
|
|
|
|
class TokenProcess:
|
|
"""Handles token processing and tracking for agents."""
|
|
|
|
def __init__(self):
|
|
"""Initialize the token processor."""
|
|
self._total_tokens = 0
|
|
self._prompt_tokens = 0
|
|
self._completion_tokens = 0
|
|
self._cached_prompt_tokens = 0
|
|
self._successful_requests = 0
|
|
|
|
def sum_prompt_tokens(self, count: int) -> None:
|
|
"""Add to prompt token count.
|
|
|
|
Args:
|
|
count (int): Number of prompt tokens to add
|
|
"""
|
|
self._prompt_tokens += count
|
|
self._total_tokens += count
|
|
|
|
def sum_completion_tokens(self, count: int) -> None:
|
|
"""Add to completion token count.
|
|
|
|
Args:
|
|
count (int): Number of completion tokens to add
|
|
"""
|
|
self._completion_tokens += count
|
|
self._total_tokens += count
|
|
|
|
def sum_cached_prompt_tokens(self, count: int) -> None:
|
|
"""Add to cached prompt token count.
|
|
|
|
Args:
|
|
count (int): Number of cached prompt tokens to add
|
|
"""
|
|
self._cached_prompt_tokens += count
|
|
|
|
def sum_successful_requests(self, count: int) -> None:
|
|
"""Add to successful requests count.
|
|
|
|
Args:
|
|
count (int): Number of successful requests to add
|
|
"""
|
|
self._successful_requests += count
|
|
|
|
def reset(self) -> None:
|
|
"""Reset all token counts to zero."""
|
|
self._total_tokens = 0
|
|
self._prompt_tokens = 0
|
|
self._completion_tokens = 0
|
|
self._cached_prompt_tokens = 0
|
|
self._successful_requests = 0
|
|
|
|
def get_summary(self) -> UsageMetrics:
|
|
"""Get a summary of token usage.
|
|
|
|
Returns:
|
|
UsageMetrics: Object containing token usage metrics
|
|
"""
|
|
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
|
|
)
|