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)}