diff --git a/sentry_sdk/integrations/langchain.py b/sentry_sdk/integrations/langchain.py index 9dcbb189ce..ce01b222c0 100644 --- a/sentry_sdk/integrations/langchain.py +++ b/sentry_sdk/integrations/langchain.py @@ -728,46 +728,89 @@ def on_tool_error( self._handle_error(run_id, error) +# LangChain's usage_metadata is a leaky abstraction: the token-detail dicts pass +# provider-formatted keys through unchanged. Instead of sniffing ad-hoc keys at +# each call site, the known spellings per field are listed once here, in +# precedence order (LangChain's own normalized names first, then provider ones). +# Adding support for another provider's spelling means extending one tuple. +_INPUT_TOKEN_KEYS = ("prompt_tokens", "input_tokens") +_OUTPUT_TOKEN_KEYS = ("completion_tokens", "output_tokens") +_INPUT_DETAILS_KEYS = ("input_token_details", "prompt_tokens_details") +_OUTPUT_DETAILS_KEYS = ("output_token_details", "completion_tokens_details") +_CACHED_TOKEN_KEYS = ("cache_read", "cached_tokens") +_REASONING_TOKEN_KEYS = ("reasoning", "reasoning_tokens") + + +def _first_value(obj: "Any", keys: "tuple[str, ...]") -> "Optional[Any]": + """Return the first non-None value among `keys`. + + Deliberately not an `or`-chain: a legitimate count of 0 must be kept + rather than falling through to the next provider's key. + """ + if obj is None: + return None + for key in keys: + value = _get_value(obj, key) + if value is not None: + return value + return None + + def _extract_tokens( token_usage: "Any", -) -> "tuple[Optional[int], Optional[int], Optional[int]]": +) -> "tuple[Optional[int], Optional[int], Optional[int], Optional[int], Optional[int]]": if not token_usage: - return None, None, None + return None, None, None, None, None - input_tokens = _get_value(token_usage, "prompt_tokens") or _get_value( - token_usage, "input_tokens" + input_tokens = _first_value(token_usage, _INPUT_TOKEN_KEYS) + output_tokens = _first_value(token_usage, _OUTPUT_TOKEN_KEYS) + total_tokens = _get_value(token_usage, "total_tokens") + + cached_tokens = _first_value( + _first_value(token_usage, _INPUT_DETAILS_KEYS), _CACHED_TOKEN_KEYS ) - output_tokens = _get_value(token_usage, "completion_tokens") or _get_value( - token_usage, "output_tokens" + reasoning_tokens = _first_value( + _first_value(token_usage, _OUTPUT_DETAILS_KEYS), _REASONING_TOKEN_KEYS ) - total_tokens = _get_value(token_usage, "total_tokens") - return input_tokens, output_tokens, total_tokens + return input_tokens, output_tokens, total_tokens, cached_tokens, reasoning_tokens def _extract_tokens_from_generations( generations: "Any", -) -> "tuple[Optional[int], Optional[int], Optional[int]]": +) -> "tuple[Optional[int], Optional[int], Optional[int], Optional[int], Optional[int]]": """Extract token usage from response.generations structure.""" if not generations: - return None, None, None + return None, None, None, None, None total_input = 0 total_output = 0 total_total = 0 + total_cached = 0 + total_reasoning = 0 for gen_list in generations: for gen in gen_list: token_usage = _get_token_usage(gen) - input_tokens, output_tokens, total_tokens = _extract_tokens(token_usage) + ( + input_tokens, + output_tokens, + total_tokens, + cached_tokens, + reasoning_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_cached += cached_tokens if cached_tokens is not None else 0 + total_reasoning += reasoning_tokens if reasoning_tokens is not None else 0 return ( 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_cached if total_cached > 0 else None, + total_reasoning if total_reasoning > 0 else None, ) @@ -802,11 +845,21 @@ def _get_token_usage(obj: "Any") -> "Optional[Dict[str, Any]]": def _record_token_usage(span: "Union[Span, StreamedSpan]", response: "Any") -> None: token_usage = _get_token_usage(response) if token_usage: - input_tokens, output_tokens, total_tokens = _extract_tokens(token_usage) + ( + input_tokens, + output_tokens, + total_tokens, + cached_tokens, + reasoning_tokens, + ) = _extract_tokens(token_usage) else: - input_tokens, output_tokens, total_tokens = _extract_tokens_from_generations( - response.generations - ) + ( + input_tokens, + output_tokens, + total_tokens, + cached_tokens, + reasoning_tokens, + ) = _extract_tokens_from_generations(response.generations) set_on_span = ( span.set_attribute if isinstance(span, StreamedSpan) else span.set_data @@ -821,6 +874,12 @@ def _record_token_usage(span: "Union[Span, StreamedSpan]", response: "Any") -> N if total_tokens is not None: set_on_span(SPANDATA.GEN_AI_USAGE_TOTAL_TOKENS, total_tokens) + if cached_tokens is not None: + set_on_span(SPANDATA.GEN_AI_USAGE_INPUT_TOKENS_CACHED, cached_tokens) + + if reasoning_tokens is not None: + set_on_span(SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS_REASONING, reasoning_tokens) + def _get_request_data( obj: "Any", args: "Any", kwargs: "Any" diff --git a/sentry_sdk/integrations/langgraph.py b/sentry_sdk/integrations/langgraph.py index 3d3856a913..575d74abc6 100644 --- a/sentry_sdk/integrations/langgraph.py +++ b/sentry_sdk/integrations/langgraph.py @@ -12,7 +12,7 @@ from sentry_sdk.integrations import DidNotEnable, Integration # This is fine because langgraph depends on langchain-base, and LangchainIntegration only imports from langchain-base. -from sentry_sdk.integrations.langchain import LangchainIntegration +from sentry_sdk.integrations.langchain import LangchainIntegration, _extract_tokens from sentry_sdk.scope import should_send_default_pii from sentry_sdk.traces import StreamedSpan from sentry_sdk.tracing_utils import ( @@ -432,6 +432,8 @@ def _set_usage_data(span: "sentry_sdk.tracing.Span", messages: "Any") -> None: input_tokens = 0 output_tokens = 0 total_tokens = 0 + cached_tokens = 0 + reasoning_tokens = 0 for message in messages: response_metadata = message.get("response_metadata") @@ -442,9 +444,16 @@ def _set_usage_data(span: "sentry_sdk.tracing.Span", messages: "Any") -> None: if not token_usage: continue - input_tokens += int(token_usage.get("prompt_tokens", 0)) - output_tokens += int(token_usage.get("completion_tokens", 0)) - total_tokens += int(token_usage.get("total_tokens", 0)) + # Single extraction path shared with the Langchain integration, so both + # integrations understand the same provider token-usage shapes. + message_input, message_output, message_total, message_cached, message_reasoning = ( + _extract_tokens(token_usage) + ) + input_tokens += int(message_input or 0) + output_tokens += int(message_output or 0) + total_tokens += int(message_total or 0) + cached_tokens += int(message_cached or 0) + reasoning_tokens += int(message_reasoning or 0) set_on_span = ( span.set_attribute if isinstance(span, StreamedSpan) else span.set_data @@ -462,6 +471,12 @@ def _set_usage_data(span: "sentry_sdk.tracing.Span", messages: "Any") -> None: total_tokens, ) + if cached_tokens > 0: + set_on_span(SPANDATA.GEN_AI_USAGE_INPUT_TOKENS_CACHED, cached_tokens) + + if reasoning_tokens > 0: + set_on_span(SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS_REASONING, reasoning_tokens) + def _set_response_model_name(span: "sentry_sdk.tracing.Span", messages: "Any") -> None: if len(messages) == 0: diff --git a/tests/integrations/langchain/test_langchain.py b/tests/integrations/langchain/test_langchain.py index f46cec9652..ac28fed0aa 100644 --- a/tests/integrations/langchain/test_langchain.py +++ b/tests/integrations/langchain/test_langchain.py @@ -4977,3 +4977,72 @@ def test_transform_list_with_legacy_image_url(self): "mime_type": "image/jpeg", "content": "/9j/4AAQ...", } + + +def test_extract_tokens_includes_cached_and_reasoning_details(): + from sentry_sdk.integrations.langchain import _extract_tokens + + # LangChain usage_metadata shape + usage = { + "input_tokens": 100, + "output_tokens": 50, + "total_tokens": 150, + "input_token_details": {"cache_read": 40}, + "output_token_details": {"reasoning": 10}, + } + assert _extract_tokens(usage) == (100, 50, 150, 40, 10) + + # OpenAI-style details shape + usage = { + "prompt_tokens": 100, + "completion_tokens": 50, + "total_tokens": 150, + "prompt_tokens_details": {"cached_tokens": 30}, + "completion_tokens_details": {"reasoning_tokens": 5}, + } + assert _extract_tokens(usage) == (100, 50, 150, 30, 5) + + # No details present + usage = {"input_tokens": 1, "output_tokens": 2, "total_tokens": 3} + assert _extract_tokens(usage) == (1, 2, 3, None, None) + + +def test_record_token_usage_sets_cached_and_reasoning_span_data(): + from unittest.mock import MagicMock + + from sentry_sdk.consts import SPANDATA + from sentry_sdk.integrations.langchain import _record_token_usage + + span = MagicMock(spec=["set_data"]) + response = MagicMock() + response.llm_output = None + response.generations = [] + response.usage = None + response.token_usage = None + response.message = None + response.usage_metadata = { + "input_tokens": 100, + "output_tokens": 50, + "total_tokens": 150, + "input_token_details": {"cache_read": 40}, + "output_token_details": {"reasoning": 10}, + } + + _record_token_usage(span, response) + + span.set_data.assert_any_call(SPANDATA.GEN_AI_USAGE_INPUT_TOKENS_CACHED, 40) + span.set_data.assert_any_call(SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS_REASONING, 10) + + +def test_extract_tokens_preserves_zero_counts(): + from sentry_sdk.integrations.langchain import _extract_tokens + + # A legitimate 0 must not fall through to the other provider's key. + usage = { + "input_tokens": 0, + "output_tokens": 0, + "total_tokens": 0, + "input_token_details": {"cache_read": 0, "cached_tokens": 7}, + "output_token_details": {"reasoning": 0, "reasoning_tokens": 9}, + } + assert _extract_tokens(usage) == (0, 0, 0, 0, 0) diff --git a/tests/integrations/langgraph/test_langgraph.py b/tests/integrations/langgraph/test_langgraph.py index fa90a89efe..b9fec3d2be 100644 --- a/tests/integrations/langgraph/test_langgraph.py +++ b/tests/integrations/langgraph/test_langgraph.py @@ -2132,3 +2132,70 @@ def test_graph_bubble_up_ignored(sentry_init, capture_items): model.invoke([HumanMessage(content="hi")]) assert len(events) == 0 + + +def test_set_usage_data_includes_cached_and_reasoning_tokens(): + from unittest.mock import MagicMock + + from sentry_sdk.consts import SPANDATA + from sentry_sdk.integrations.langgraph import _set_usage_data + + span = MagicMock(spec=["set_data"]) + messages = [ + { + "response_metadata": { + "token_usage": { + "prompt_tokens": 100, + "completion_tokens": 50, + "total_tokens": 150, + "prompt_tokens_details": {"cached_tokens": 40}, + "completion_tokens_details": {"reasoning_tokens": 10}, + } + } + }, + { + "response_metadata": { + "token_usage": { + "prompt_tokens": 10, + "completion_tokens": 5, + "total_tokens": 15, + } + } + }, + ] + + _set_usage_data(span, messages) + + span.set_data.assert_any_call(SPANDATA.GEN_AI_USAGE_INPUT_TOKENS, 110) + span.set_data.assert_any_call(SPANDATA.GEN_AI_USAGE_INPUT_TOKENS_CACHED, 40) + span.set_data.assert_any_call(SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS_REASONING, 10) + + +def test_set_usage_data_reads_langchain_detail_shapes(): + from unittest.mock import MagicMock + + from sentry_sdk.consts import SPANDATA + from sentry_sdk.integrations.langgraph import _set_usage_data + + # LangChain-style usage_metadata detail keys (not OpenAI-style) must also be + # understood, via the extraction shared with the Langchain integration. + span = MagicMock(spec=["set_data"]) + messages = [ + { + "response_metadata": { + "token_usage": { + "input_tokens": 80, + "output_tokens": 20, + "total_tokens": 100, + "input_token_details": {"cache_read": 30}, + "output_token_details": {"reasoning": 5}, + } + } + }, + ] + + _set_usage_data(span, messages) + + span.set_data.assert_any_call(SPANDATA.GEN_AI_USAGE_INPUT_TOKENS, 80) + span.set_data.assert_any_call(SPANDATA.GEN_AI_USAGE_INPUT_TOKENS_CACHED, 30) + span.set_data.assert_any_call(SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS_REASONING, 5)