chore(closes OPEN-8441): tokens not captured for Bedrock traced via the LangChain callback handler

Author: gustavocidornelas

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