klbr/khoj
1
0
Fork 0
mirror of https://github.com/khoaliber/khoj.git synced 2026-03-02 21:19:12 +00:00
Code Issues Packages Projects Releases Wiki Activity

Log ai model usage on every call to get_chat_usage_metrics in debug mode

Browse Source
This commit is contained in:
Debanjum
2025-05-30 16:02:45 -07:00
parent 4cb900658d
commit 3797f03625
2 changed files with 4 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
src/khoj/processor/operator/operator_agent_base.py
View File

@@ -98,7 +98,6 @@ class OperatorAgent(ABC):
self.tracer["usage"] = get_chat_usage_metrics( self.tracer["usage"] = get_chat_usage_metrics(
self.vision_model.name, input_tokens, output_tokens, cache_read, cache_write, usage=self.tracer.get("usage") 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): def _commit_trace(self):
self.tracer["chat_model"] = self.vision_model.name self.tracer["chat_model"] = self.vision_model.name

5
src/khoj/utils/helpers.py
View File

@@ -46,6 +46,7 @@ if TYPE_CHECKING:
from khoj.utils.models import BaseEncoder from khoj.utils.models import BaseEncoder
from khoj.utils.rawconfig import AppConfig from khoj.utils.rawconfig import AppConfig
logger = logging.getLogger(__name__)
# Initialize Magika for file type identification # Initialize Magika for file type identification
magika = Magika() magika = Magika()
@@ -680,7 +681,7 @@ def get_chat_usage_metrics(
"cache_write_tokens": 0, "cache_write_tokens": 0,
"cost": 0.0, "cost": 0.0,
} }
return { current_usage = {
"input_tokens": prev_usage["input_tokens"] + input_tokens, "input_tokens": prev_usage["input_tokens"] + input_tokens,
"output_tokens": prev_usage["output_tokens"] + output_tokens, "output_tokens": prev_usage["output_tokens"] + output_tokens,
"thought_tokens": prev_usage.get("thought_tokens", 0) + thought_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"], prev_cost=prev_usage["cost"],
), ),
} }
logger.debug(f"AI API usage by {model_name}: {current_usage}")
return current_usage
class AiApiInfo(NamedTuple): class AiApiInfo(NamedTuple):