fix: support usage as part of generation

Author: k11kirky

posthog/ai/langchain/callbacks.py

@@ -597,9 +597,13 @@ def _parse_usage_model(
         # Bedrock: https://docs.aws.amazon.com/bedrock/latest/userguide/monitoring-cw.html#runtime-cloudwatch-metrics
         ("inputTokenCount", "input"),
         ("outputTokenCount", "output"),
+        # Bedrock Anthropic
+        ("prompt_tokens", "input"),
+        ("completion_tokens", "output"),
         # langchain-ibm https://pypi.org/project/langchain-ibm/
         ("input_token_count", "input"),
         ("generated_token_count", "output"),
+
     ]
 
     parsed_usage = {}
@@ -627,6 +631,10 @@ def _parse_usage(response: LLMResult):
 
     if hasattr(response, "generations"):
         for generation in response.generations:
+            if "usage" in generation:
+                llm_usage = _parse_usage_model(generation["usage"])
+                break
+            
             for generation_chunk in generation:
                 if generation_chunk.generation_info and ("usage_metadata" in generation_chunk.generation_info):
                     llm_usage = _parse_usage_model(generation_chunk.generation_info["usage_metadata"])
@@ -652,6 +660,7 @@ def _parse_usage(response: LLMResult):
                     llm_usage = _parse_usage_model(chunk_usage)
                     break
 
+
     return llm_usage