From 1c5dae9b6c40e3fb237849955353f2943f96a5dc Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Andr=C3=A9=20Ferreira?= Date: Tue, 18 Jun 2024 16:22:42 +0200 Subject: [PATCH] fix: :bug: Skip updating estimates if event isn't a prompt nor message --- tokencost/callbacks/llama_index.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tokencost/callbacks/llama_index.py b/tokencost/callbacks/llama_index.py index ee4ef8e..d11c1ab 100644 --- a/tokencost/callbacks/llama_index.py +++ b/tokencost/callbacks/llama_index.py @@ -39,6 +39,8 @@ class TokenCostHandler(BaseCallbackHandler): estimates = calculate_all_costs_and_tokens( messages_str, response, self.model ) + else: + return self.prompt_cost += estimates["prompt_cost"] self.completion_cost += estimates["completion_cost"]