feat(langchain): Record cached token usage attributes

sentry_sdk/consts.py

@@ -782,6 +782,20 @@ class SPANDATA:
     Example: "rainy, 57°F"
     """
 
+    GEN_AI_USAGE_CACHE_READ_INPUT_TOKENS = "gen_ai.usage.cache_read.input_tokens"
+    """
+    The number of cached tokens used to process the AI input (prompt).
+    Example: 50
+    """
+
+    GEN_AI_USAGE_CACHE_CREATION_INPUT_TOKENS = (
+        "gen_ai.usage.cache_creation.input_tokens"
+    )
+    """
+    The number of tokens written to the cache when processing the AI input (prompt).
+    Example: 100
+    """
+
     GEN_AI_USAGE_INPUT_TOKENS = "gen_ai.usage.input_tokens"
     """
     The number of tokens in the input.

sentry_sdk/integrations/langchain.py

@@ -4,7 +4,7 @@
 import warnings
 from collections import OrderedDict
 from functools import wraps
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, NamedTuple
 
 import sentry_sdk
 from sentry_sdk.ai.utils import (
@@ -50,7 +50,7 @@
         Callbacks,
         manager,
     )
-    from langchain_core.messages import BaseMessage
+    from langchain_core.messages import AIMessage, BaseMessage
     from langchain_core.outputs import (
         ChatGeneration,
         ChatGenerationChunk,
@@ -63,6 +63,14 @@
     raise DidNotEnable("langchain not installed")
 
 
+class TokenUsage(NamedTuple):
+    input_tokens: "Optional[int]"
+    output_tokens: "Optional[int]"
+    total_tokens: "Optional[int]"
+    cache_read: "Optional[int]"
+    cache_creation: "Optional[int]"
+
+
 try:
     # >=v1
     from langchain_classic.agents import AgentExecutor  # type: ignore[import-not-found]
@@ -720,29 +728,58 @@ def _extract_tokens(
 
 def _extract_tokens_from_generations(
     generations: "list[list[Generation | ChatGeneration | GenerationChunk | ChatGenerationChunk]]",
-) -> "tuple[Optional[int], Optional[int], Optional[int]]":
+) -> "TokenUsage":
     """Extract token usage from response.generations structure."""
-    if not generations:
-        return None, None, None
-
     total_input = 0
     total_output = 0
     total_total = 0
+    total_cache_read = None
+    total_cache_creation = None
 
     for gen_list in generations:
         if not gen_list:
             continue
 
         token_usage = _get_token_usage(gen_list[0])
         input_tokens, output_tokens, total_tokens = _extract_tokens(token_usage)
-        total_input += input_tokens if input_tokens is not None else 0
-        total_output += output_tokens if output_tokens is not None else 0
-        total_total += total_tokens if total_tokens is not None else 0
+        total_input += input_tokens if isinstance(input_tokens, int) else 0
+        total_output += output_tokens if isinstance(output_tokens, int) else 0
+        total_total += total_tokens if isinstance(total_tokens, int) else 0
+
+        if not isinstance(gen_list[0], ChatGeneration):
+            continue
+
+        message = gen_list[0].message
+
+        # The property was added in https://github.com/langchain-ai/langchain/commit/fbfed65fb1ccff3eb8477c4f114450537a0510b2
+        if not isinstance(message, AIMessage) or not hasattr(message, "usage_metadata"):
+            continue
+
+        usage_metadata = message.usage_metadata
 
-    return (
+        if not isinstance(usage_metadata, dict):
+            continue
+
+        input_token_details = usage_metadata.get("input_token_details")
+        if not isinstance(input_token_details, dict):
+            continue
+
+        if isinstance(input_token_details.get("cache_read"), int):
+            total_cache_read = (total_cache_read or 0) + input_token_details[
+                "cache_read"
+            ]
+
+        if isinstance(input_token_details.get("cache_creation"), int):
+            total_cache_creation = (total_cache_creation or 0) + input_token_details[
+                "cache_creation"
+            ]
+
+    return TokenUsage(
         total_input if total_input > 0 else None,
         total_output if total_output > 0 else None,
         total_total if total_total > 0 else None,
+        total_cache_read,
+        total_cache_creation,
     )
 
 
@@ -777,13 +814,30 @@ def _get_token_usage(obj: "Any") -> "Optional[Dict[str, Any]]":
 def _record_token_usage(
     span: "Union[Span, StreamedSpan]", response: "LLMResult"
 ) -> None:
+    input_tokens = None
+    output_tokens = None
+    total_tokens = None
+    cache_read_tokens = None
+    cache_creation_tokens = None
+
+    # Legacy that reads provider-specific token information.
     token_usage = _get_token_usage(response)
     if token_usage:
         input_tokens, output_tokens, total_tokens = _extract_tokens(token_usage)
-    else:
-        input_tokens, output_tokens, total_tokens = _extract_tokens_from_generations(
-            response.generations
-        )
+
+    # Prefer provider-agnostic UsageMetadata if available.
+    if response.generations is not None:
+        token_usage = _extract_tokens_from_generations(response.generations)
+        if token_usage.input_tokens is not None:
+            input_tokens = token_usage.input_tokens
+        if token_usage.output_tokens is not None:
+            output_tokens = token_usage.output_tokens
+        if token_usage.total_tokens is not None:
+            total_tokens = token_usage.total_tokens
+        if token_usage.cache_read is not None:
+            cache_read_tokens = token_usage.cache_read
+        if token_usage.cache_creation is not None:
+            cache_creation_tokens = token_usage.cache_creation
 
     set_on_span = (
         span.set_attribute if isinstance(span, StreamedSpan) else span.set_data
@@ -798,6 +852,14 @@ def _record_token_usage(
     if total_tokens is not None:
         set_on_span(SPANDATA.GEN_AI_USAGE_TOTAL_TOKENS, total_tokens)
 
+    if cache_read_tokens is not None:
+        set_on_span(SPANDATA.GEN_AI_USAGE_CACHE_READ_INPUT_TOKENS, cache_read_tokens)
+
+    if cache_creation_tokens is not None:
+        set_on_span(
+            SPANDATA.GEN_AI_USAGE_CACHE_CREATION_INPUT_TOKENS, cache_creation_tokens
+        )
+
 
 def _get_request_data(
     obj: "Any", args: "Any", kwargs: "Any"

tests/conftest.py

@@ -1439,8 +1439,8 @@ def nonstreaming_responses_model_response():
         usage=openai.types.responses.ResponseUsage(
             input_tokens=10,
             input_tokens_details=openai.types.responses.response_usage.InputTokensDetails(
-                cached_tokens=0,
-                cache_write_tokens=0,
+                cached_tokens=4,
+                cache_write_tokens=6,
             ),
             output_tokens=20,
             output_tokens_details=openai.types.responses.response_usage.OutputTokensDetails(

tests/integrations/langchain/test_langchain.py

@@ -268,6 +268,7 @@ def nonstreaming_multi_candidate_google_genai_model_response():
         ],
         model_version="gemini/gemini-pro",
         usage_metadata=google.genai.types.GenerateContentResponseUsageMetadata(
+            cached_content_token_count=4,
             prompt_token_count=10,
             candidates_token_count=20,
             total_token_count=30,
@@ -592,6 +593,11 @@ def test_langchain_multi_choice_response(
         assert chat_spans[0]["attributes"]["gen_ai.usage.input_tokens"] == 10
         assert chat_spans[0]["attributes"]["gen_ai.usage.output_tokens"] == 20
         assert chat_spans[0]["attributes"]["gen_ai.usage.total_tokens"] == 30
+
+        assert (
+            chat_spans[0]["attributes"][SPANDATA.GEN_AI_USAGE_CACHE_READ_INPUT_TOKENS]
+            == 4
+        )
     else:
         events = capture_events()
 
@@ -613,6 +619,8 @@ def test_langchain_multi_choice_response(
         assert chat_spans[0]["data"]["gen_ai.usage.output_tokens"] == 20
         assert chat_spans[0]["data"]["gen_ai.usage.total_tokens"] == 30
 
+        assert chat_spans[0]["data"][SPANDATA.GEN_AI_USAGE_CACHE_READ_INPUT_TOKENS] == 4
+
 
 @pytest.mark.parametrize("span_streaming", [True, False])
 @pytest.mark.parametrize("stream_gen_ai_spans", [True, False])
@@ -798,6 +806,17 @@ def test_langchain_create_agent(
         assert chat_spans[0]["attributes"]["gen_ai.usage.output_tokens"] == 20
         assert chat_spans[0]["attributes"]["gen_ai.usage.total_tokens"] == 30
 
+        assert (
+            chat_spans[0]["attributes"][SPANDATA.GEN_AI_USAGE_CACHE_READ_INPUT_TOKENS]
+            == 4
+        )
+        assert (
+            chat_spans[0]["attributes"][
+                SPANDATA.GEN_AI_USAGE_CACHE_CREATION_INPUT_TOKENS
+            ]
+            == 6
+        )
+
         if LANGCHAIN_OPENAI_VERSION >= (0, 3, 13):
             assert (
                 chat_spans[0]["attributes"][SPANDATA.GEN_AI_RESPONSE_MODEL] == "gpt-4"
@@ -874,6 +893,17 @@ def test_langchain_create_agent(
         assert chat_spans[0]["attributes"]["gen_ai.usage.output_tokens"] == 20
         assert chat_spans[0]["attributes"]["gen_ai.usage.total_tokens"] == 30
 
+        assert (
+            chat_spans[0]["attributes"][SPANDATA.GEN_AI_USAGE_CACHE_READ_INPUT_TOKENS]
+            == 4
+        )
+        assert (
+            chat_spans[0]["attributes"][
+                SPANDATA.GEN_AI_USAGE_CACHE_CREATION_INPUT_TOKENS
+            ]
+            == 6
+        )
+
         if LANGCHAIN_OPENAI_VERSION >= (0, 3, 13):
             assert (
                 chat_spans[0]["attributes"][SPANDATA.GEN_AI_RESPONSE_MODEL] == "gpt-4"
@@ -943,6 +973,12 @@ def test_langchain_create_agent(
         assert chat_spans[0]["data"]["gen_ai.usage.output_tokens"] == 20
         assert chat_spans[0]["data"]["gen_ai.usage.total_tokens"] == 30
 
+        assert chat_spans[0]["data"][SPANDATA.GEN_AI_USAGE_CACHE_READ_INPUT_TOKENS] == 4
+        assert (
+            chat_spans[0]["data"][SPANDATA.GEN_AI_USAGE_CACHE_CREATION_INPUT_TOKENS]
+            == 6
+        )
+
         if LANGCHAIN_OPENAI_VERSION >= (0, 3, 13):
             assert chat_spans[0]["data"][SPANDATA.GEN_AI_RESPONSE_MODEL] == "gpt-4"
 
@@ -1013,8 +1049,8 @@ def test_tool_execution_span(
                 ResponseUsage(
                     input_tokens=142,
                     input_tokens_details=InputTokensDetails(
-                        cached_tokens=0,
-                        cache_write_tokens=0,
+                        cached_tokens=69,
+                        cache_write_tokens=31,
                     ),
                     output_tokens=50,
                     output_tokens_details=OutputTokensDetails(
@@ -1025,8 +1061,8 @@ def test_tool_execution_span(
                 ResponseUsage(
                     input_tokens=89,
                     input_tokens_details=InputTokensDetails(
-                        cached_tokens=0,
-                        cache_write_tokens=0,
+                        cached_tokens=69,
+                        cache_write_tokens=10,
                     ),
                     output_tokens=28,
                     output_tokens_details=OutputTokensDetails(
@@ -1107,11 +1143,31 @@ def test_tool_execution_span(
         assert chat_spans[0]["attributes"]["gen_ai.usage.input_tokens"] == 142
         assert chat_spans[0]["attributes"]["gen_ai.usage.output_tokens"] == 50
         assert chat_spans[0]["attributes"]["gen_ai.usage.total_tokens"] == 192
+        assert (
+            chat_spans[0]["attributes"][SPANDATA.GEN_AI_USAGE_CACHE_READ_INPUT_TOKENS]
+            == 69
+        )
+        assert (
+            chat_spans[0]["attributes"][
+                SPANDATA.GEN_AI_USAGE_CACHE_CREATION_INPUT_TOKENS
+            ]
+            == 31
+        )
         assert chat_spans[0]["attributes"]["gen_ai.system"] == "openai-chat"
 
         assert chat_spans[1]["attributes"]["gen_ai.usage.input_tokens"] == 89
         assert chat_spans[1]["attributes"]["gen_ai.usage.output_tokens"] == 28
         assert chat_spans[1]["attributes"]["gen_ai.usage.total_tokens"] == 117
+        assert (
+            chat_spans[1]["attributes"][SPANDATA.GEN_AI_USAGE_CACHE_READ_INPUT_TOKENS]
+            == 69
+        )
+        assert (
+            chat_spans[1]["attributes"][
+                SPANDATA.GEN_AI_USAGE_CACHE_CREATION_INPUT_TOKENS
+            ]
+            == 10
+        )
         assert chat_spans[1]["attributes"]["gen_ai.system"] == "openai-chat"
 
         if LANGCHAIN_OPENAI_VERSION >= (0, 3, 13):
@@ -1220,11 +1276,31 @@ def test_tool_execution_span(
         assert chat_spans[0]["attributes"]["gen_ai.usage.input_tokens"] == 142
         assert chat_spans[0]["attributes"]["gen_ai.usage.output_tokens"] == 50
         assert chat_spans[0]["attributes"]["gen_ai.usage.total_tokens"] == 192
+        assert (
+            chat_spans[0]["attributes"][SPANDATA.GEN_AI_USAGE_CACHE_READ_INPUT_TOKENS]
+            == 69
+        )
+        assert (
+            chat_spans[0]["attributes"][
+                SPANDATA.GEN_AI_USAGE_CACHE_CREATION_INPUT_TOKENS
+            ]
+            == 31
+        )
         assert chat_spans[0]["attributes"]["gen_ai.system"] == "openai-chat"
 
         assert chat_spans[1]["attributes"]["gen_ai.usage.input_tokens"] == 89
         assert chat_spans[1]["attributes"]["gen_ai.usage.output_tokens"] == 28
         assert chat_spans[1]["attributes"]["gen_ai.usage.total_tokens"] == 117
+        assert (
+            chat_spans[1]["attributes"][SPANDATA.GEN_AI_USAGE_CACHE_READ_INPUT_TOKENS]
+            == 69
+        )
+        assert (
+            chat_spans[1]["attributes"][
+                SPANDATA.GEN_AI_USAGE_CACHE_CREATION_INPUT_TOKENS
+            ]
+            == 10
+        )
         assert chat_spans[1]["attributes"]["gen_ai.system"] == "openai-chat"
 
         if LANGCHAIN_OPENAI_VERSION >= (0, 3, 13):
@@ -1333,11 +1409,25 @@ def test_tool_execution_span(
         assert chat_spans[0]["data"]["gen_ai.usage.input_tokens"] == 142
         assert chat_spans[0]["data"]["gen_ai.usage.output_tokens"] == 50
         assert chat_spans[0]["data"]["gen_ai.usage.total_tokens"] == 192
+        assert (
+            chat_spans[0]["data"][SPANDATA.GEN_AI_USAGE_CACHE_READ_INPUT_TOKENS] == 69
+        )
+        assert (
+            chat_spans[0]["data"][SPANDATA.GEN_AI_USAGE_CACHE_CREATION_INPUT_TOKENS]
+            == 31
+        )
         assert chat_spans[0]["data"]["gen_ai.system"] == "openai-chat"
 
         assert chat_spans[1]["data"]["gen_ai.usage.input_tokens"] == 89
         assert chat_spans[1]["data"]["gen_ai.usage.output_tokens"] == 28
         assert chat_spans[1]["data"]["gen_ai.usage.total_tokens"] == 117
+        assert (
+            chat_spans[1]["data"][SPANDATA.GEN_AI_USAGE_CACHE_READ_INPUT_TOKENS] == 69
+        )
+        assert (
+            chat_spans[1]["data"][SPANDATA.GEN_AI_USAGE_CACHE_CREATION_INPUT_TOKENS]
+            == 10
+        )
         assert chat_spans[1]["data"]["gen_ai.system"] == "openai-chat"
 
         if LANGCHAIN_OPENAI_VERSION >= (0, 3, 13):