fix: capture LangChain token usage fallbacks

instrumentation/opentelemetry-instrumentation-genai-langchain/.changelog/640.fixed

@@ -0,0 +1 @@
+Capture LangChain chat token usage from generation metadata and LLM output fallbacks.

instrumentation/opentelemetry-instrumentation-genai-langchain/src/opentelemetry/instrumentation/genai/langchain/callback_handler.py

@@ -33,6 +33,7 @@
     _message_name,
     _normalize_role,
     extract_token_details,
+    extract_usage_tokens,
     is_stream_end_marker,
     make_input_message,
     make_last_output_message,
@@ -80,6 +81,40 @@ def _conversation_id(metadata: dict[str, Any] | None) -> str | None:
     return None
 
 
+def _usage_metadata_candidates(
+    chat_generation: Any,
+    llm_output: Any,
+) -> list[Mapping[str, Any]]:
+    candidates: list[Mapping[str, Any]] = []
+    message = getattr(chat_generation, "message", None)
+    usage_metadata = _usage_mapping(getattr(message, "usage_metadata", None))
+    if usage_metadata is not None:
+        candidates.append(usage_metadata)
+
+    generation_info = _usage_mapping(
+        getattr(chat_generation, "generation_info", None)
+    )
+    if generation_info is not None:
+        usage_metadata = _usage_mapping(generation_info.get("usage_metadata"))
+        if usage_metadata is not None:
+            candidates.append(usage_metadata)
+
+    llm_output_mapping = _usage_mapping(llm_output)
+    if llm_output_mapping is not None:
+        for key in ("token_usage", "usage"):
+            usage_metadata = _usage_mapping(llm_output_mapping.get(key))
+            if usage_metadata is not None:
+                candidates.append(usage_metadata)
+
+    return candidates
+
+
+def _usage_mapping(value: Any) -> Mapping[str, Any] | None:
+    if isinstance(value, Mapping):
+        return cast(Mapping[str, Any], value)
+    return None
+
+
 class OpenTelemetryLangChainCallbackHandler(BaseCallbackHandler):
     """
     A callback handler for LangChain that uses OpenTelemetry to create spans for LLM calls and chains, tools etc,. in future.
@@ -399,6 +434,7 @@ def on_llm_end(
         served_model: str | None = None
         generation_model: str | None = None
         generation_response_id: str | None = None
+        llm_output = getattr(response, "llm_output", None)
         for generation in getattr(response, "generations", []):
             for chat_generation in generation:
                 message = chat_generation.message
@@ -521,25 +557,24 @@ def on_llm_end(
                     output_messages.append(output_message)
 
                     # Get token usage if available
-                    if chat_generation.message.usage_metadata:
-                        usage_metadata = chat_generation.message.usage_metadata
-                        input_tokens = usage_metadata.get("input_tokens", 0)
-                        if not isinstance(input_tokens, int) or isinstance(
-                            input_tokens, bool
-                        ):
-                            input_tokens = 0
-                        llm_invocation.input_tokens = input_tokens
-
-                        output_tokens = usage_metadata.get("output_tokens", 0)
-                        if not isinstance(output_tokens, int) or isinstance(
-                            output_tokens, bool
-                        ):
-                            output_tokens = 0
+                    has_input_tokens = False
+                    has_output_tokens = False
+                    for usage_metadata in _usage_metadata_candidates(
+                        chat_generation,
+                        llm_output,
+                    ):
+                        input_tokens, output_tokens = extract_usage_tokens(
+                            usage_metadata
+                        )
+                        if input_tokens is not None and not has_input_tokens:
+                            llm_invocation.input_tokens = input_tokens
+                            has_input_tokens = True
+                        if output_tokens is not None and not has_output_tokens:
+                            llm_invocation.output_tokens = output_tokens
+                            has_output_tokens = True
 
                         # Cache, reasoning, and modality token break-downs
-                        token_details = extract_token_details(
-                            cast(dict[str, Any], usage_metadata)
-                        )
+                        token_details = extract_token_details(usage_metadata)
                         if (
                             cache_write := token_details.get(
                                 "cache_write_input_tokens"
@@ -591,12 +626,10 @@ def on_llm_end(
                         ) is not None:
                             llm_invocation.audio_output_tokens = audio_out
 
-                        llm_invocation.output_tokens = output_tokens
-
         llm_invocation.output_messages = output_messages
 
         response_model, response_id = resolve_response_model_and_id(
-            llm_output=getattr(response, "llm_output", None),
+            llm_output=llm_output,
             served_model=served_model,
             generation_model=generation_model,
             generation_response_id=generation_response_id,

instrumentation/opentelemetry-instrumentation-genai-langchain/src/opentelemetry/instrumentation/genai/langchain/utils.py

@@ -616,7 +616,14 @@ def resolve_response_model_and_id(
     return response_model, response_id
 
 
-def extract_token_details(usage_metadata: dict[str, Any]) -> dict[str, int]:
+def _get_positive_int(values: Mapping[str, Any], key: str) -> int | None:
+    value = values.get(key)
+    if isinstance(value, int) and not isinstance(value, bool) and value > 0:
+        return value
+    return None
+
+
+def extract_token_details(usage_metadata: Mapping[str, Any]) -> dict[str, int]:
     """Extract cache, reasoning, and modality token break-downs from LangChain usage metadata."""
 
     token_details: dict[str, int] = {}
@@ -633,21 +640,22 @@ def extract_token_details(usage_metadata: dict[str, Any]) -> dict[str, int]:
         else {}
     )
 
-    def _get_positive_int(d: dict[str, Any], key: str) -> int | None:
-        val = d.get(key)
-        if isinstance(val, int) and not isinstance(val, bool) and val > 0:
-            return val
-        return None
-
     cache_write = _get_positive_int(input_details, "cache_write")
     if cache_write is None:
         cache_write = _get_positive_int(input_details, "cache_creation")
+    if cache_write is None:
+        cache_write = _get_positive_int(
+            usage_metadata, "cache_creation_input_tokens"
+        )
     if cache_write is not None:
         token_details["cache_write_input_tokens"] = cache_write
 
-    if (
-        cache_read := _get_positive_int(input_details, "cache_read")
-    ) is not None:
+    cache_read = _get_positive_int(input_details, "cache_read")
+    if cache_read is None:
+        cache_read = _get_positive_int(
+            usage_metadata, "cache_read_input_tokens"
+        )
+    if cache_read is not None:
         token_details["cache_read_input_tokens"] = cache_read
 
     if (
@@ -672,3 +680,41 @@ def _get_positive_int(d: dict[str, Any], key: str) -> int | None:
         token_details["audio_output_tokens"] = audio_out
 
     return token_details
+
+
+def extract_usage_tokens(
+    usage_metadata: Mapping[str, Any],
+) -> tuple[
+    int | None,
+    int | None,
+]:
+    """Extract input and output token counts from LangChain usage mappings."""
+
+    input_tokens = _first_int_value(
+        usage_metadata.get("input_tokens"),
+        usage_metadata.get("prompt_tokens"),
+        usage_metadata.get("prompt_token_count"),
+    )
+    if input_tokens is not None:
+        input_tokens += sum(
+            cache_tokens
+            for key in (
+                "cache_creation_input_tokens",
+                "cache_read_input_tokens",
+            )
+            if (cache_tokens := _get_positive_int(usage_metadata, key))
+            is not None
+        )
+    output_tokens = _first_int_value(
+        usage_metadata.get("output_tokens"),
+        usage_metadata.get("completion_tokens"),
+        usage_metadata.get("candidates_token_count"),
+    )
+    return input_tokens, output_tokens
+
+
+def _first_int_value(*values: Any) -> int | None:
+    for value in values:
+        if isinstance(value, int) and not isinstance(value, bool):
+            return value
+    return None

instrumentation/opentelemetry-instrumentation-genai-langchain/tests/test_callback_handler.py

@@ -37,12 +37,14 @@
 
 from opentelemetry.instrumentation.genai.langchain.callback_handler import (
     OpenTelemetryLangChainCallbackHandler,
+    _usage_metadata_candidates,
 )
 from opentelemetry.instrumentation.genai.langchain.utils import (
     _legacy_function_call_request,
     _media_part,
     _normalize_role,
     extract_token_details,
+    extract_usage_tokens,
     make_input_message,
     make_last_output_message,
     make_output_message,
@@ -1672,6 +1674,137 @@ def test_modality_tokens_set_on_invocation(self):
         assert llm_inv.image_output_tokens == 10
         assert llm_inv.audio_output_tokens == 5
 
+    @pytest.mark.parametrize(
+        ("generation_info", "llm_output", "input_tokens", "output_tokens"),
+        [
+            (
+                {"finish_reason": "stop"},
+                {
+                    "token_usage": {
+                        "prompt_tokens": 13,
+                        "completion_tokens": 7,
+                    }
+                },
+                13,
+                7,
+            ),
+            (
+                {"finish_reason": "stop"},
+                {"usage": {"input_tokens": 10, "output_tokens": 4}},
+                10,
+                4,
+            ),
+            (
+                {
+                    "finish_reason": "stop",
+                    "usage_metadata": {
+                        "input_tokens": 11,
+                        "output_tokens": 5,
+                    },
+                },
+                None,
+                11,
+                5,
+            ),
+        ],
+    )
+    def test_token_usage_falls_back_to_alternate_metadata(
+        self,
+        generation_info,
+        llm_output,
+        input_tokens,
+        output_tokens,
+    ):
+        run_id = _run_id()
+        handler, _, llm_inv = _make_handler_with_llm_invocation(run_id)
+
+        ai_msg = AIMessage(content="hi")
+        gen = ChatGeneration(message=ai_msg, generation_info=generation_info)
+        response = LLMResult(generations=[[gen]], llm_output=llm_output)
+
+        handler.on_llm_end(response=response, run_id=run_id)
+
+        assert llm_inv.input_tokens == input_tokens
+        assert llm_inv.output_tokens == output_tokens
+
+    def test_empty_generation_info_usage_does_not_mask_llm_output(self):
+        run_id = _run_id()
+        handler, _, llm_inv = _make_handler_with_llm_invocation(run_id)
+
+        ai_msg = AIMessage(content="hi")
+        gen = ChatGeneration(
+            message=ai_msg,
+            generation_info={
+                "finish_reason": "stop",
+                "usage_metadata": {},
+            },
+        )
+        response = LLMResult(
+            generations=[[gen]],
+            llm_output={
+                "token_usage": {
+                    "prompt_tokens": 13,
+                    "completion_tokens": 7,
+                },
+            },
+        )
+
+        handler.on_llm_end(response=response, run_id=run_id)
+
+        assert llm_inv.input_tokens == 13
+        assert llm_inv.output_tokens == 7
+
+    def test_vertexai_generation_info_usage_metadata(self):
+        run_id = _run_id()
+        handler, _, llm_inv = _make_handler_with_llm_invocation(run_id)
+
+        ai_msg = AIMessage(content="hi")
+        gen = ChatGeneration(
+            message=ai_msg,
+            generation_info={
+                "finish_reason": "stop",
+                "usage_metadata": {
+                    "prompt_token_count": 13,
+                    "candidates_token_count": 7,
+                    "total_token_count": 20,
+                },
+            },
+        )
+        response = LLMResult(generations=[[gen]])
+
+        handler.on_llm_end(response=response, run_id=run_id)
+
+        assert llm_inv.input_tokens == 13
+        assert llm_inv.output_tokens == 7
+
+    def test_anthropic_llm_output_usage_includes_cached_input_tokens(self):
+        run_id = _run_id()
+        handler, _, llm_inv = _make_handler_with_llm_invocation(run_id)
+
+        ai_msg = AIMessage(content="hi")
+        gen = ChatGeneration(
+            message=ai_msg,
+            generation_info={"finish_reason": "stop"},
+        )
+        response = LLMResult(
+            generations=[[gen]],
+            llm_output={
+                "usage": {
+                    "input_tokens": 10,
+                    "cache_creation_input_tokens": 3,
+                    "cache_read_input_tokens": 2,
+                    "output_tokens": 4,
+                }
+            },
+        )
+
+        handler.on_llm_end(response=response, run_id=run_id)
+
+        assert llm_inv.input_tokens == 15
+        assert llm_inv.cache_write_input_tokens == 3
+        assert llm_inv.cache_read_input_tokens == 2
+        assert llm_inv.output_tokens == 4
+
 
 # ---------------------------------------------------------------------------
 # utils.extract_token_details
@@ -1768,6 +1901,37 @@ def test_extract_token_details_no_details_key():
     assert extract_token_details({"input_tokens": 1, "output_tokens": 2}) == {}
 
 
+def test_extract_usage_tokens_langchain_keys():
+    assert extract_usage_tokens({"input_tokens": 10, "output_tokens": 20}) == (
+        10,
+        20,
+    )
+
+
+def test_extract_usage_tokens_openai_keys():
+    assert extract_usage_tokens(
+        {"prompt_tokens": 10, "completion_tokens": 20}
+    ) == (10, 20)
+
+
+def test_extract_usage_tokens_ignores_bool_and_unknown_values():
+    assert extract_usage_tokens(
+        {"input_tokens": True, "output_tokens": "20"}
+    ) == (None, None)
+
+
+def test_usage_metadata_candidates_handles_missing_message_usage_metadata():
+    chat_generation = mock.Mock(
+        message=object(),
+        generation_info={},
+    )
+    llm_output = {"token_usage": {"prompt_tokens": 10}}
+
+    assert _usage_metadata_candidates(chat_generation, llm_output) == [
+        llm_output["token_usage"]
+    ]
+
+
 def test_extract_token_details_boolean_and_non_int_omitted():
     usage = {
         "input_token_details": {