From 67e4433dc282644ba7e8146f0ab3e676b8cbd2a7 Mon Sep 17 00:00:00 2001 From: Thiago Moretto Date: Mon, 5 Aug 2024 13:29:39 -0300 Subject: [PATCH 1/2] Fix flaky test due to suppressed error on on_llm_start callback --- src/crewai/utilities/token_counter_callback.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/src/crewai/utilities/token_counter_callback.py b/src/crewai/utilities/token_counter_callback.py index 51200f37e4..c5f632d5f5 100644 --- a/src/crewai/utilities/token_counter_callback.py +++ b/src/crewai/utilities/token_counter_callback.py @@ -10,24 +10,24 @@ class TokenCalcHandler(BaseCallbackHandler): model_name: str = "" token_cost_process: TokenProcess + encoding: tiktoken.Encoding def __init__(self, model_name, token_cost_process): self.model_name = model_name self.token_cost_process = token_cost_process + try: + self.encoding = tiktoken.encoding_for_model(self.model_name) + except KeyError as e: + self.encoding = tiktoken.get_encoding("cl100k_base") def on_llm_start( self, serialized: Dict[str, Any], prompts: List[str], **kwargs: Any ) -> None: - try: - encoding = tiktoken.encoding_for_model(self.model_name) - except KeyError: - encoding = tiktoken.get_encoding("cl100k_base") - if self.token_cost_process is None: return for prompt in prompts: - self.token_cost_process.sum_prompt_tokens(len(encoding.encode(prompt))) + self.token_cost_process.sum_prompt_tokens(len(self.encoding.encode(prompt))) async def on_llm_new_token(self, token: str, **kwargs) -> None: self.token_cost_process.sum_completion_tokens(1) From f3b3d321e5799153efbf9ba81c73b97f9ed04d3a Mon Sep 17 00:00:00 2001 From: Thiago Moretto Date: Mon, 5 Aug 2024 13:34:03 -0300 Subject: [PATCH 2/2] Fix lint issue --- src/crewai/utilities/token_counter_callback.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/crewai/utilities/token_counter_callback.py b/src/crewai/utilities/token_counter_callback.py index c5f632d5f5..dc0cef71df 100644 --- a/src/crewai/utilities/token_counter_callback.py +++ b/src/crewai/utilities/token_counter_callback.py @@ -17,7 +17,7 @@ def __init__(self, model_name, token_cost_process): self.token_cost_process = token_cost_process try: self.encoding = tiktoken.encoding_for_model(self.model_name) - except KeyError as e: + except KeyError: self.encoding = tiktoken.get_encoding("cl100k_base") def on_llm_start(