mirror of
https://github.com/crewAIInc/crewAI.git
synced 2026-01-09 16:18:30 +00:00
42311d9c7a
* Fix SQLite log handling issue causing ValueError: Logs cannot be None in tests - Add proper error handling in SQLite storage operations - Set up isolated test environment with temporary storage directory - Ensure consistent error messages across all database operations Co-Authored-By: Joe Moura <joao@crewai.com> * fix: Sort imports in conftest.py Co-Authored-By: Joe Moura <joao@crewai.com> * fix: Convert TokenProcess counters to instance variables to fix callback tracking Co-Authored-By: Joe Moura <joao@crewai.com> * refactor: Replace print statements with logging and improve error handling - Add proper logging setup in kickoff_task_outputs_storage.py - Replace self._printer.print() with logger calls - Use appropriate log levels (error/warning) - Add directory validation in test environment setup - Maintain consistent error messages with DatabaseError format Co-Authored-By: Joe Moura <joao@crewai.com> * fix: Comprehensive improvements to database and token handling - Fix SQLite database path handling in storage classes - Add proper directory creation and error handling - Improve token tracking with robust type checking - Convert TokenProcess counters to instance variables - Add standardized database error handling - Set up isolated test environment with temporary storage Resolves test failures in PR #1899 Co-Authored-By: Joe Moura <joao@crewai.com> --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Co-authored-by: Joe Moura <joao@crewai.com> Co-authored-by: João Moura <joaomdmoura@gmail.com>
38 lines
1.5 KiB
Python
38 lines
1.5 KiB
Python
import warnings
|
|
from typing import Any, Dict, Optional
|
|
|
|
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: Optional[TokenProcess]):
|
|
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:
|
|
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:
|
|
self.token_cost_process.sum_cached_prompt_tokens(
|
|
usage.prompt_tokens_details.cached_tokens
|
|
)
|