fix: handle None values in token counter

- Update sum_cached_prompt_tokens to handle None values gracefully
- Add unit tests for token counting with None values
- Fixes #2197

Co-Authored-By: Joe Moura <joao@crewai.com>
This commit is contained in:
Devin AI
2025-02-22 01:29:47 +00:00
parent b50772a38b
commit 3f02d10626
2 changed files with 20 additions and 2 deletions

View File

@@ -17,7 +17,8 @@ class TokenProcess:
self.completion_tokens += tokens
self.total_tokens += tokens
def sum_cached_prompt_tokens(self, tokens: int) -> None:
def sum_cached_prompt_tokens(self, tokens: int | None) -> None:
if tokens is not None:
self.cached_prompt_tokens += tokens
def sum_successful_requests(self, requests: int) -> None:

View File

@@ -0,0 +1,17 @@
import unittest
from crewai.agents.agent_builder.utilities.base_token_process import TokenProcess
class TestTokenProcess(unittest.TestCase):
def setUp(self):
self.token_process = TokenProcess()
def test_sum_cached_prompt_tokens_with_none(self):
initial_tokens = self.token_process.cached_prompt_tokens
self.token_process.sum_cached_prompt_tokens(None)
self.assertEqual(self.token_process.cached_prompt_tokens, initial_tokens)
def test_sum_cached_prompt_tokens_with_int(self):
initial_tokens = self.token_process.cached_prompt_tokens
self.token_process.sum_cached_prompt_tokens(5)
self.assertEqual(self.token_process.cached_prompt_tokens, initial_tokens + 5)