(Openinference Migration: Langchain) - Add support for cache and reasoning token counts (PR #1)

instrumentation/opentelemetry-instrumentation-genai-langchain/.changelog/272.added

@@ -0,0 +1 @@
+(Openinference Migration: Langchain) - Add support for cache and reasoning token counts

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

@@ -21,6 +21,7 @@
 )
 from opentelemetry.instrumentation.genai.langchain.utils import (
     _normalize_role,
+    extract_token_details,
     make_input_message,
     make_last_output_message,
     normalize_provider,
@@ -362,18 +363,50 @@ def on_llm_end(
 
                     # Get token usage if available
                     if chat_generation.message.usage_metadata:
-                        input_tokens = (
-                            chat_generation.message.usage_metadata.get(
-                                "input_tokens", 0
-                            )
-                        )
+                        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 = (
-                            chat_generation.message.usage_metadata.get(
-                                "output_tokens", 0
+                        output_tokens = usage_metadata.get("output_tokens", 0)
+                        if not isinstance(output_tokens, int) or isinstance(
+                            output_tokens, bool
+                        ):
+                            output_tokens = 0
+
+                        # Cache/reasoning break-downs (Anthropic, OpenAI
+                        # reasoning models, Bedrock). Audio tokens are dropped
+                        # (no GenAI semconv attribute).
+                        token_details = extract_token_details(
+                            cast(dict[str, Any], usage_metadata)
+                        )
+                        cache_creation = token_details.get(
+                            "cache_creation_input_tokens"
+                        )
+                        if cache_creation is not None:
+                            llm_invocation.cache_creation_input_tokens = (
+                                cache_creation
                             )
+                        cache_read = token_details.get(
+                            "cache_read_input_tokens"
+                        )
+                        if cache_read is not None:
+                            llm_invocation.cache_read_input_tokens = cache_read
+                        reasoning_tokens = token_details.get(
+                            "reasoning_tokens"
                         )
+                        if reasoning_tokens is not None:
+                            llm_invocation.thinking_tokens = reasoning_tokens
+                            # LangChain folds reasoning into ``output_tokens``;
+                            # util-genai re-sums ``output + thinking`` so
+                            # subtract here to keep the provider total intact.
+                            output_tokens = max(
+                                output_tokens - reasoning_tokens, 0
+                            )
+
                         llm_invocation.output_tokens = output_tokens
 
         llm_invocation.output_messages = output_messages

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

@@ -319,3 +319,35 @@ def serialize(obj: Any) -> Optional[str]:
         return json.dumps(obj, ensure_ascii=False, default=str)
     except (TypeError, ValueError):
         return None
+
+
+def extract_token_details(usage_metadata: dict[str, Any]) -> dict[str, int]:
+    """Extract cache/reasoning token break-downs from LangChain usage metadata."""
+
+    token_details: dict[str, int] = {}
+    raw_input_details = usage_metadata.get("input_token_details")
+    input_details: dict[str, Any] = (
+        cast(dict[str, Any], raw_input_details)
+        if isinstance(raw_input_details, dict)
+        else {}
+    )
+    raw_output_details = usage_metadata.get("output_token_details")
+    output_details: dict[str, Any] = (
+        cast(dict[str, Any], raw_output_details)
+        if isinstance(raw_output_details, dict)
+        else {}
+    )
+
+    cache_creation = input_details.get("cache_creation")
+    if isinstance(cache_creation, int) and cache_creation:
+        token_details["cache_creation_input_tokens"] = cache_creation
+
+    cache_read = input_details.get("cache_read")
+    if isinstance(cache_read, int) and cache_read:
+        token_details["cache_read_input_tokens"] = cache_read
+
+    reasoning = output_details.get("reasoning")
+    if isinstance(reasoning, int) and reasoning:
+        token_details["reasoning_tokens"] = reasoning
+
+    return token_details

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

@@ -19,6 +19,7 @@
     OpenTelemetryLangChainCallbackHandler,
 )
 from opentelemetry.instrumentation.genai.langchain.utils import (
+    extract_token_details,
     make_input_message,
     make_last_output_message,
     make_output_message,
@@ -1206,3 +1207,108 @@ def test_bedrock_tool_use_finish_reason_produces_tool_call_request(self):
         assert part.name == "get_weather"
         assert part.id == "tooluse_abc"
         assert part.arguments == {"location": "London"}
+
+
+# ---------------------------------------------------------------------------
+# on_llm_end – token usage break-downs
+# ---------------------------------------------------------------------------
+
+
+class TestOnLlmEndTokenDetails:
+    def test_cache_and_reasoning_tokens_set_on_invocation(self):
+        run_id = _run_id()
+        handler, _, llm_inv = _make_handler_with_llm_invocation(run_id)
+
+        ai_msg = AIMessage(
+            content="hi there",
+            usage_metadata={
+                "input_tokens": 10,
+                "output_tokens": 20,
+                "total_tokens": 30,
+                "input_token_details": {
+                    "cache_creation": 3,
+                    "cache_read": 2,
+                },
+                "output_token_details": {"reasoning": 5},
+            },
+        )
+        gen = ChatGeneration(
+            message=ai_msg, generation_info={"finish_reason": "stop"}
+        )
+        response = LLMResult(generations=[[gen]])
+
+        handler.on_llm_end(response=response, run_id=run_id)
+
+        assert llm_inv.input_tokens == 10
+        assert llm_inv.cache_creation_input_tokens == 3
+        assert llm_inv.cache_read_input_tokens == 2
+        assert llm_inv.thinking_tokens == 5
+        assert llm_inv.output_tokens == 15
+
+    def test_audio_tokens_ignored(self):
+        run_id = _run_id()
+        handler, _, llm_inv = _make_handler_with_llm_invocation(run_id)
+
+        ai_msg = AIMessage(
+            content="hi",
+            usage_metadata={
+                "input_tokens": 10,
+                "output_tokens": 20,
+                "total_tokens": 30,
+                "input_token_details": {"audio": 5},
+                "output_token_details": {"audio": 4},
+            },
+        )
+        gen = ChatGeneration(
+            message=ai_msg, generation_info={"finish_reason": "stop"}
+        )
+        response = LLMResult(generations=[[gen]])
+
+        handler.on_llm_end(response=response, run_id=run_id)
+
+        assert llm_inv.input_tokens == 10
+        assert llm_inv.output_tokens == 20
+
+
+# ---------------------------------------------------------------------------
+# utils.extract_token_details
+# ---------------------------------------------------------------------------
+
+
+def test_extract_token_details_cache_and_reasoning():
+    usage = {
+        "input_tokens": 10,
+        "output_tokens": 20,
+        "total_tokens": 30,
+        "input_token_details": {"cache_creation": 3, "cache_read": 2},
+        "output_token_details": {"reasoning": 5},
+    }
+    details = extract_token_details(usage)
+    assert details == {
+        "cache_creation_input_tokens": 3,
+        "cache_read_input_tokens": 2,
+        "reasoning_tokens": 5,
+    }
+
+
+def test_extract_token_details_ignores_audio_tokens():
+    usage = {
+        "input_tokens": 10,
+        "output_tokens": 20,
+        "input_token_details": {"audio": 5},
+        "output_token_details": {"audio": 4},
+    }
+    assert extract_token_details(usage) == {}
+
+
+def test_extract_token_details_zero_values_omitted():
+    usage = {
+        "input_tokens": 10,
+        "output_tokens": 20,
+        "input_token_details": {"cache_creation": 0, "cache_read": 0},
+    }
+    assert extract_token_details(usage) == {}
+
+
+def test_extract_token_details_no_details_key():
+    assert extract_token_details({"input_tokens": 1, "output_tokens": 2}) == {}

util/opentelemetry-util-genai/src/opentelemetry/util/genai/_inference_invocation.py

@@ -180,10 +180,10 @@ def _get_metric_token_counts(self) -> dict[str, int]:
         counts: dict[str, int] = {}
         if self.input_tokens is not None:
             counts[GenAI.GenAiTokenTypeValues.INPUT.value] = self.input_tokens
-        if self.output_tokens is not None:
+        if self.output_tokens is not None or self.thinking_tokens is not None:
             counts[GenAI.GenAiTokenTypeValues.OUTPUT.value] = (
-                self.output_tokens
-            )
+                self.output_tokens or 0
+            ) + (self.thinking_tokens or 0)
         return counts
 
     def _apply_finish(self, error: Error | None = None) -> None: