diff --git a/src/crewai/agents/agent_builder/utilities/base_token_process.py b/src/crewai/agents/agent_builder/utilities/base_token_process.py index 50a02023f..ced4dd1a7 100644 --- a/src/crewai/agents/agent_builder/utilities/base_token_process.py +++ b/src/crewai/agents/agent_builder/utilities/base_token_process.py @@ -1,4 +1,8 @@ from crewai.types.usage_metrics import UsageMetrics +import logging + + +logger = logging.getLogger(__name__) class TokenProcess: @@ -27,10 +31,12 @@ class TokenProcess: Raises: ValueError: If tokens is negative. """ - if tokens is not None: - if tokens < 0: - raise ValueError("Token count cannot be negative") - self.cached_prompt_tokens += tokens + if tokens is None: + logger.debug("Received None value for token count") + return + if tokens < 0: + raise ValueError("Token count cannot be negative") + self.cached_prompt_tokens += tokens def sum_successful_requests(self, requests: int) -> None: self.successful_requests += requests diff --git a/tests/utilities/test_token_process.py b/tests/utilities/test_token_process.py index 7fb89dab7..98ae5e69b 100644 --- a/tests/utilities/test_token_process.py +++ b/tests/utilities/test_token_process.py @@ -4,7 +4,10 @@ from crewai.agents.agent_builder.utilities.base_token_process import TokenProces class TestTokenProcess(unittest.TestCase): + """Test suite for TokenProcess class token counting functionality.""" + def setUp(self): + """Initialize a fresh TokenProcess instance before each test.""" self.token_process = TokenProcess() def test_sum_cached_prompt_tokens_with_none(self):