diff --git a/langfuse/callback/langchain.py b/langfuse/callback/langchain.py index 674ab0b21..1986ed87b 100644 --- a/langfuse/callback/langchain.py +++ b/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