Fix Langchain ChatVertexAI

langfuse/callback/langchain.py

@@ -1121,6 +1121,48 @@ def _parse_usage(response: LLMResult):
     # langchain-anthropic uses the usage field
     llm_usage_keys = ["token_usage", "usage"]
     llm_usage = None
+    model_name = _parse_model(response)
+    if (
+        model_name
+        and "gemini" in model_name.lower()
+        and hasattr(response, "generations")
+    ):
+        for generation in response.generations:
+            for generation_chunk in generation:
+                if (
+                    generation_chunk.generation_info
+                    and "usage_metadata" in generation_chunk.generation_info
+                ):
+                    vertex_ai_usage = generation_chunk.generation_info["usage_metadata"]
+
+                    # Extract the information.
+                    llm_usage = {}
+                    llm_usage["prompt_tokens"] = vertex_ai_usage.get(
+                        "prompt_token_count"
+                    )
+                    llm_usage["completion_tokens"] = vertex_ai_usage.get(
+                        "candidates_token_count"
+                    )
+                    llm_usage["total_tokens"] = vertex_ai_usage.get("total_token_count")
+
+                    # Ensure the extracted values are integers.
+                    for key in ["prompt_tokens", "completion_tokens", "total_tokens"]:
+                        if llm_usage.get(key) is not None:
+                            if isinstance(llm_usage.get(key), float):
+                                llm_usage[key] = int(llm_usage[key])  # Convert if float
+                            if not isinstance(llm_usage.get(key), int):
+                                try:
+                                    llm_usage[key] = int(
+                                        llm_usage[key]
+                                    )  # Try conversion
+                                except (ValueError, TypeError):
+                                    log.warning(
+                                        f"Could not convert {key} to integer: {vertex_ai_usage.get(key)}"
+                                    )
+                                    llm_usage[key] = 0  # Default to 0
+
+                    return llm_usage  # Stop after Vertex AI usage is found
+
     if response.llm_output is not None:
         for key in llm_usage_keys:
             if key in response.llm_output and response.llm_output[key]:
@@ -1164,15 +1206,26 @@ def _parse_usage(response: LLMResult):
     return llm_usage
 
 
-def _parse_model(response: LLMResult):
-    # langchain-anthropic uses the usage field
-    llm_model_keys = ["model_name"]
+def _parse_model(response: LLMResult) -> str | None:
+    """Extract the model name from the LLMResult, handling different providers."""
     llm_model = None
-    if response.llm_output is not None:
-        for key in llm_model_keys:
-            if key in response.llm_output and response.llm_output[key]:
-                llm_model = response.llm_output[key]
-                break
+
+    # 1. Check llm_output first (e.g., for OpenAI)
+    if response.llm_output and "model_name" in response.llm_output:
+        llm_model = response.llm_output["model_name"]
+        return llm_model
+
+    # 2. If llm_output is None, check generations for Vertex AI
+    if hasattr(response, "generations") and response.generations:
+        first_generation = response.generations[0]
+        if first_generation:
+            first_generation_chunk = first_generation[0]
+            message = getattr(first_generation_chunk, "message", None)
+            if message and hasattr(message, "response_metadata"):
+                response_metadata = getattr(message, "response_metadata", None)
+                if response_metadata and "model_name" in response_metadata:
+                    llm_model = response_metadata["model_name"]
+                    return llm_model
 
     return llm_model