From 33e09bccea38ae18fc2a7b342732eb29520ef031 Mon Sep 17 00:00:00 2001 From: Hassieb Pakzad <68423100+hassiebp@users.noreply.github.com> Date: Tue, 27 May 2025 11:31:51 +0200 Subject: [PATCH] fix(langchain): token counts for vertex --- langfuse/langchain/CallbackHandler.py | 24 +++++++++++++++--------- 1 file changed, 15 insertions(+), 9 deletions(-) diff --git a/langfuse/langchain/CallbackHandler.py b/langfuse/langchain/CallbackHandler.py index fe96516e2..576bfdebd 100644 --- a/langfuse/langchain/CallbackHandler.py +++ b/langfuse/langchain/CallbackHandler.py @@ -844,9 +844,11 @@ def _parse_usage_model(usage: typing.Union[pydantic.BaseModel, dict]): and "modality" in item and "token_count" in item ): - usage_model[f"input_modality_{item['modality']}"] = item[ - "token_count" - ] + value = item["token_count"] + usage_model[f"input_modality_{item['modality']}"] = value + + if "input" in usage_model: + usage_model["input"] = max(0, usage_model["input"] - value) # Vertex AI if "candidates_tokens_details" in usage_model and isinstance( @@ -860,9 +862,11 @@ def _parse_usage_model(usage: typing.Union[pydantic.BaseModel, dict]): and "modality" in item and "token_count" in item ): - usage_model[f"output_modality_{item['modality']}"] = item[ - "token_count" - ] + value = item["token_count"] + usage_model[f"output_modality_{item['modality']}"] = value + + if "output" in usage_model: + usage_model["output"] = max(0, usage_model["output"] - value) # Vertex AI if "cache_tokens_details" in usage_model and isinstance( @@ -876,9 +880,11 @@ def _parse_usage_model(usage: typing.Union[pydantic.BaseModel, dict]): and "modality" in item and "token_count" in item ): - usage_model[f"cached_modality_{item['modality']}"] = item[ - "token_count" - ] + value = item["token_count"] + usage_model[f"cached_modality_{item['modality']}"] = value + + if "input" in usage_model: + usage_model["input"] = max(0, usage_model["input"] - value) usage_model = {k: v for k, v in usage_model.items() if not isinstance(v, str)}