push

Author: hassiebp

langfuse/langchain/CallbackHandler.py

@@ -875,22 +875,28 @@ def _parse_usage_model(usage: typing.Union[pydantic.BaseModel, dict]) -> Any:
     usage_model = cast(Dict, usage.copy())  # Copy all existing key-value pairs
 
     # Skip OpenAI usage types as they are handled server side
-    if not all(
+    if all(
         openai_key in usage_model
-        for openai_key in ["prompt_tokens", "completion_tokens", "total_tokens"]
+        for openai_key in [
+            "prompt_tokens",
+            "completion_tokens",
+            "total_tokens",
+            "prompt_tokens_details",
+            "completion_tokens_details",
+        ]
     ):
-        for model_key, langfuse_key in conversion_list:
-            if model_key in usage_model:
-                captured_count = usage_model.pop(model_key)
-                final_count = (
-                    sum(captured_count)
-                    if isinstance(captured_count, list)
-                    else captured_count
-                )  # For Bedrock, the token count is a list when streamed
-
-                usage_model[langfuse_key] = (
-                    final_count  # Translate key and keep the value
-                )
+        return usage_model
+
+    for model_key, langfuse_key in conversion_list:
+        if model_key in usage_model:
+            captured_count = usage_model.pop(model_key)
+            final_count = (
+                sum(captured_count)
+                if isinstance(captured_count, list)
+                else captured_count
+            )  # For Bedrock, the token count is a list when streamed
+
+            usage_model[langfuse_key] = final_count  # Translate key and keep the value
 
     if isinstance(usage_model, dict):
         if "input_token_details" in usage_model: