From 76ac43dba1e4ef4c83b1dd0abf1d8a6f80c4d501 Mon Sep 17 00:00:00 2001 From: Gustavo Cid Date: Fri, 12 Dec 2025 15:42:19 -0300 Subject: [PATCH] chore(closes OPEN-8441): tokens not captured for Bedrock traced via the LangChain callback handler --- .../lib/integrations/langchain_callback.py | 13 ++++++++++++- 1 file changed, 12 insertions(+), 1 deletion(-) diff --git a/src/openlayer/lib/integrations/langchain_callback.py b/src/openlayer/lib/integrations/langchain_callback.py index 6d3da477..33d79edc 100644 --- a/src/openlayer/lib/integrations/langchain_callback.py +++ b/src/openlayer/lib/integrations/langchain_callback.py @@ -29,6 +29,7 @@ "openai-chat": "OpenAI", "chat-ollama": "Ollama", "vertexai": "Google", + "amazon_bedrock_converse_chat": "Bedrock", } @@ -398,7 +399,8 @@ def _extract_token_info( # Fallback to generation info for providers like Ollama/Google if not token_usage and response.generations: - generation_info = response.generations[0][0].generation_info or {} + gen = response.generations[0][0] + generation_info = gen.generation_info or {} # Ollama style if "prompt_eval_count" in generation_info: @@ -417,6 +419,15 @@ def _extract_token_info( "completion_tokens": usage.get("candidates_token_count", 0), "total_tokens": usage.get("total_token_count", 0), } + # AWS Bedrock / newer LangChain style - usage_metadata on the message + elif hasattr(gen, "message") and hasattr(gen.message, "usage_metadata"): + usage = gen.message.usage_metadata + if usage: + token_usage = { + "prompt_tokens": usage.get("input_tokens", 0), + "completion_tokens": usage.get("output_tokens", 0), + "total_tokens": usage.get("total_tokens", 0), + } return { "prompt_tokens": token_usage.get("prompt_tokens", 0),