From 3797f0362590c862dc638e40e65e8ff8aac86222 Mon Sep 17 00:00:00 2001 From: Debanjum Date: Fri, 30 May 2025 16:02:45 -0700 Subject: [PATCH] Log ai model usage on every call to get_chat_usage_metrics in debug mode --- src/khoj/processor/operator/operator_agent_base.py | 1 - src/khoj/utils/helpers.py | 5 ++++- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/src/khoj/processor/operator/operator_agent_base.py b/src/khoj/processor/operator/operator_agent_base.py index 8686b708..a04b44d4 100644 --- a/src/khoj/processor/operator/operator_agent_base.py +++ b/src/khoj/processor/operator/operator_agent_base.py @@ -98,7 +98,6 @@ class OperatorAgent(ABC): self.tracer["usage"] = get_chat_usage_metrics( self.vision_model.name, input_tokens, output_tokens, cache_read, cache_write, usage=self.tracer.get("usage") ) - logger.debug(f"Operator usage by {self.vision_model.model_type}: {self.tracer['usage']}") def _commit_trace(self): self.tracer["chat_model"] = self.vision_model.name diff --git a/src/khoj/utils/helpers.py b/src/khoj/utils/helpers.py index 1ddda3fb..4e877c11 100644 --- a/src/khoj/utils/helpers.py +++ b/src/khoj/utils/helpers.py @@ -46,6 +46,7 @@ if TYPE_CHECKING: from khoj.utils.models import BaseEncoder from khoj.utils.rawconfig import AppConfig +logger = logging.getLogger(__name__) # Initialize Magika for file type identification magika = Magika() @@ -680,7 +681,7 @@ def get_chat_usage_metrics( "cache_write_tokens": 0, "cost": 0.0, } - return { + current_usage = { "input_tokens": prev_usage["input_tokens"] + input_tokens, "output_tokens": prev_usage["output_tokens"] + output_tokens, "thought_tokens": prev_usage.get("thought_tokens", 0) + thought_tokens, @@ -697,6 +698,8 @@ def get_chat_usage_metrics( prev_cost=prev_usage["cost"], ), } + logger.debug(f"AI API usage by {model_name}: {current_usage}") + return current_usage class AiApiInfo(NamedTuple):