diff --git a/sentry_sdk/integrations/langchain.py b/sentry_sdk/integrations/langchain.py index 9dcbb189ce..8fcb010dff 100644 --- a/sentry_sdk/integrations/langchain.py +++ b/sentry_sdk/integrations/langchain.py @@ -730,9 +730,9 @@ def on_tool_error( 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" @@ -742,32 +742,64 @@ def _extract_tokens( ) total_tokens = _get_value(token_usage, "total_tokens") - return input_tokens, output_tokens, total_tokens + # LangChain's usage_metadata nests these under input/output_token_details; + # OpenAI-style dicts use prompt/completion_tokens_details. + input_details = _get_value(token_usage, "input_token_details") or _get_value( + token_usage, "prompt_tokens_details" + ) + cached_tokens = None + if input_details is not None: + cached_tokens = _get_value(input_details, "cache_read") or _get_value( + input_details, "cached_tokens" + ) + + output_details = _get_value(token_usage, "output_token_details") or _get_value( + token_usage, "completion_tokens_details" + ) + reasoning_tokens = None + if output_details is not None: + reasoning_tokens = _get_value(output_details, "reasoning") or _get_value( + output_details, "reasoning_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 +834,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 +863,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..74935cf000 100644 --- a/sentry_sdk/integrations/langgraph.py +++ b/sentry_sdk/integrations/langgraph.py @@ -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") @@ -446,6 +448,12 @@ def _set_usage_data(span: "sentry_sdk.tracing.Span", messages: "Any") -> None: output_tokens += int(token_usage.get("completion_tokens", 0)) total_tokens += int(token_usage.get("total_tokens", 0)) + input_details = token_usage.get("prompt_tokens_details") or {} + cached_tokens += int(input_details.get("cached_tokens") or 0) + + output_details = token_usage.get("completion_tokens_details") or {} + reasoning_tokens += int(output_details.get("reasoning_tokens") or 0) + set_on_span = ( span.set_attribute if isinstance(span, StreamedSpan) else span.set_data ) @@ -462,6 +470,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..4f215ddfe4 100644 --- a/tests/integrations/langchain/test_langchain.py +++ b/tests/integrations/langchain/test_langchain.py @@ -4977,3 +4977,58 @@ 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) diff --git a/tests/integrations/langgraph/test_langgraph.py b/tests/integrations/langgraph/test_langgraph.py index fa90a89efe..0eae1dca75 100644 --- a/tests/integrations/langgraph/test_langgraph.py +++ b/tests/integrations/langgraph/test_langgraph.py @@ -2132,3 +2132,40 @@ 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)