diff --git a/sentry_sdk/integrations/langchain.py b/sentry_sdk/integrations/langchain.py index 91bb70f550..c80c7d7278 100644 --- a/sentry_sdk/integrations/langchain.py +++ b/sentry_sdk/integrations/langchain.py @@ -3,12 +3,14 @@ import sys import warnings from collections import OrderedDict +from dataclasses import dataclass from functools import wraps from typing import TYPE_CHECKING import sentry_sdk from sentry_sdk.ai.utils import ( GEN_AI_ALLOWED_MESSAGE_ROLES, + _set_span_data_attribute, get_start_span_function, normalize_message_roles, set_data_normalized, @@ -716,47 +718,116 @@ def on_tool_error( self._handle_error(run_id, error) +@dataclass +class _TokenUsage: + """ + Normalized token usage, matching UsageMetadata from Langchain + """ + + input_tokens: "Optional[int|float]" = None + input_tokens_audio: "Optional[int|float]" = None + input_tokens_cache_read: "Optional[int|float]" = None + input_tokens_cache_creation: "Optional[int|float]" = None + output_tokens: "Optional[int|float]" = None + output_tokens_audio: "Optional[int|float]" = None + output_tokens_reasoning: "Optional[int|float]" = None + total_tokens: "Optional[int|float]" = None + + def _extract_tokens( token_usage: "Any", -) -> "tuple[Optional[int], Optional[int], Optional[int]]": +) -> _TokenUsage: + usage = _TokenUsage() if not token_usage: - return None, None, None + return usage - input_tokens = _get_value(token_usage, "prompt_tokens") or _get_value( + usage.input_tokens = _get_value(token_usage, "prompt_tokens") or _get_value( token_usage, "input_tokens" ) - output_tokens = _get_value(token_usage, "completion_tokens") or _get_value( + usage.output_tokens = _get_value(token_usage, "completion_tokens") or _get_value( token_usage, "output_tokens" ) - total_tokens = _get_value(token_usage, "total_tokens") + usage.total_tokens = _get_value(token_usage, "total_tokens") + + input_token_details = _get_value(token_usage, "input_token_details") + if input_token_details is not None: + usage.input_tokens_audio = _get_value(input_token_details, "audio") + usage.input_tokens_cache_read = _get_value(input_token_details, "cache_read") + usage.input_tokens_cache_creation = _get_value( + input_token_details, "cache_creation" + ) - return input_tokens, output_tokens, total_tokens + output_token_details = _get_value(token_usage, "output_token_details") + if output_token_details is not None: + usage.output_tokens_audio = _get_value(output_token_details, "audio") + usage.output_tokens_reasoning = _get_value(output_token_details, "reasoning") + return usage def _extract_tokens_from_generations( generations: "Any", -) -> "tuple[Optional[int], Optional[int], Optional[int]]": +) -> _TokenUsage: """Extract token usage from response.generations structure.""" + total = _TokenUsage() + if not generations: - return None, None, None + return total + + def _is_number(value: "Any") -> bool: + return isinstance(value, (int, float)) - total_input = 0 - total_output = 0 - total_total = 0 + def _add(left: "Any", right: "Any") -> "Union[int, float]": + left = left if _is_number(left) else 0 + right = right if _is_number(right) else 0 + return left + right + + def _none_or_zero(value: "Optional[int|float]") -> bool: + return value is None or value <= 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) - 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 - - 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, - ) + tokens = _extract_tokens(token_usage) + total.input_tokens = _add(total.input_tokens, tokens.input_tokens) + total.output_tokens = _add(total.output_tokens, tokens.output_tokens) + total.total_tokens = _add(total.total_tokens, tokens.total_tokens) + + # We don't want to default zero here as it oculd be a valid value that is given by the + # provider and can be a common value for caches + if _is_number(tokens.input_tokens_cache_read): + total.input_tokens_cache_read = _add( + total.input_tokens_cache_read, tokens.input_tokens_cache_read + ) + + if _is_number(tokens.input_tokens_cache_creation): + total.input_tokens_cache_creation = _add( + total.input_tokens_cache_creation, + tokens.input_tokens_cache_creation, + ) + + if _is_number(tokens.input_tokens_audio): + total.input_tokens_audio = _add( + total.input_tokens_audio, tokens.input_tokens_audio + ) + + if _is_number(tokens.output_tokens_audio): + total.output_tokens_audio = _add( + total.output_tokens_audio, tokens.output_tokens_audio + ) + + if _is_number(tokens.output_tokens_reasoning): + total.output_tokens_reasoning = _add( + total.output_tokens_reasoning, tokens.output_tokens_reasoning + ) + + if _none_or_zero(total.input_tokens): + total.input_tokens = None + if _none_or_zero(total.output_tokens): + total.output_tokens = None + if _none_or_zero(total.total_tokens): + total.total_tokens = None + # we keep cached/reasoning token counts as is + return total def _get_token_usage(obj: "Any") -> "Optional[Dict[str, Any]]": @@ -788,26 +859,50 @@ 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) - else: - input_tokens, output_tokens, total_tokens = _extract_tokens_from_generations( - response.generations + # Prefer usage_metadata: UsageMetadata from "generations" as it's Langchain's provider-agnostic + # shape. Legacy usages rely on response.llm_output["token_usage"] + tokens = _extract_tokens_from_generations(response.generations) + if ( + tokens.input_tokens is None + and tokens.output_tokens is None + and tokens.total_tokens is None + ): + token_usage = _get_token_usage(response) + if token_usage: + tokens = _extract_tokens(token_usage) + + if tokens.input_tokens is not None: + _set_span_data_attribute( + span, SPANDATA.GEN_AI_USAGE_INPUT_TOKENS, tokens.input_tokens + ) + if tokens.output_tokens is not None: + _set_span_data_attribute( + span, SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS, tokens.output_tokens + ) + if tokens.total_tokens is not None: + _set_span_data_attribute( + span, SPANDATA.GEN_AI_USAGE_TOTAL_TOKENS, tokens.total_tokens ) - set_on_span = ( - span.set_attribute if isinstance(span, StreamedSpan) else span.set_data - ) - - if input_tokens is not None: - set_on_span(SPANDATA.GEN_AI_USAGE_INPUT_TOKENS, input_tokens) - - if output_tokens is not None: - set_on_span(SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS, output_tokens) - - if total_tokens is not None: - set_on_span(SPANDATA.GEN_AI_USAGE_TOTAL_TOKENS, total_tokens) + # TODO: add input/output audio tokens when Sentry supports them + if tokens.input_tokens_cache_read is not None: + _set_span_data_attribute( + span, + SPANDATA.GEN_AI_USAGE_INPUT_TOKENS_CACHED, + tokens.input_tokens_cache_read, + ) + if tokens.input_tokens_cache_creation is not None: + _set_span_data_attribute( + span, + SPANDATA.GEN_AI_USAGE_INPUT_TOKENS_CACHE_WRITE, + tokens.input_tokens_cache_creation, + ) + if tokens.output_tokens_reasoning is not None: + _set_span_data_attribute( + span, + SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS_REASONING, + tokens.output_tokens_reasoning, + ) def _get_request_data( diff --git a/tests/conftest.py b/tests/conftest.py index 6b406d6a06..ffc69f3326 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1358,6 +1358,12 @@ def streaming_chat_completions_model_response(): prompt_tokens=10, completion_tokens=20, total_tokens=30, + completion_tokens_details=openai.types.CompletionTokensDetails( + reasoning_tokens=12 + ), + prompt_tokens_details=openai.types.PromptTokensDetails( + cached_tokens=8, cache_write_tokens=2 + ), ), ), ] @@ -1439,8 +1445,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( diff --git a/tests/integrations/langchain/test_langchain.py b/tests/integrations/langchain/test_langchain.py index 0983aa2094..738ef1f459 100644 --- a/tests/integrations/langchain/test_langchain.py +++ b/tests/integrations/langchain/test_langchain.py @@ -744,7 +744,9 @@ def test_langchain_create_agent( assert chat_spans[0]["attributes"]["gen_ai.agent.name"] == "word_length_agent" assert chat_spans[0]["attributes"]["gen_ai.usage.input_tokens"] == 10 + assert chat_spans[0]["attributes"]["gen_ai.usage.input_tokens.cached"] == 4 assert chat_spans[0]["attributes"]["gen_ai.usage.output_tokens"] == 20 + assert chat_spans[0]["attributes"]["gen_ai.usage.output_tokens.reasoning"] == 5 assert chat_spans[0]["attributes"]["gen_ai.usage.total_tokens"] == 30 if LANGCHAIN_OPENAI_VERSION >= (0, 3, 13): @@ -813,7 +815,9 @@ def test_langchain_create_agent( assert chat_spans[0]["data"]["gen_ai.agent.name"] == "word_length_agent" assert chat_spans[0]["data"]["gen_ai.usage.input_tokens"] == 10 + assert chat_spans[0]["data"]["gen_ai.usage.input_tokens.cached"] == 4 assert chat_spans[0]["data"]["gen_ai.usage.output_tokens"] == 20 + assert chat_spans[0]["data"]["gen_ai.usage.output_tokens.reasoning"] == 5 assert chat_spans[0]["data"]["gen_ai.usage.total_tokens"] == 30 if LANGCHAIN_OPENAI_VERSION >= (0, 3, 13): @@ -886,24 +890,24 @@ 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( - reasoning_tokens=0, + reasoning_tokens=10, ), total_tokens=192, ), 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( - reasoning_tokens=0, + reasoning_tokens=11, ), total_tokens=117, ), @@ -1091,12 +1095,16 @@ def test_tool_execution_span( assert tool_exec_span["attributes"]["gen_ai.agent.name"] == "word_length_agent" assert chat_spans[0]["attributes"]["gen_ai.usage.input_tokens"] == 142 + assert chat_spans[0]["attributes"]["gen_ai.usage.input_tokens.cached"] == 69 assert chat_spans[0]["attributes"]["gen_ai.usage.output_tokens"] == 50 + assert chat_spans[0]["attributes"]["gen_ai.usage.output_tokens.reasoning"] == 10 assert chat_spans[0]["attributes"]["gen_ai.usage.total_tokens"] == 192 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.input_tokens.cached"] == 69 assert chat_spans[1]["attributes"]["gen_ai.usage.output_tokens"] == 28 + assert chat_spans[1]["attributes"]["gen_ai.usage.output_tokens.reasoning"] == 11 assert chat_spans[1]["attributes"]["gen_ai.usage.total_tokens"] == 117 assert chat_spans[1]["attributes"]["gen_ai.system"] == "openai-chat" @@ -1204,12 +1212,16 @@ def test_tool_execution_span( assert tool_exec_span["data"]["gen_ai.agent.name"] == "word_length_agent" assert chat_spans[0]["data"]["gen_ai.usage.input_tokens"] == 142 + assert chat_spans[0]["data"]["gen_ai.usage.input_tokens.cached"] == 69 assert chat_spans[0]["data"]["gen_ai.usage.output_tokens"] == 50 + assert chat_spans[0]["data"]["gen_ai.usage.output_tokens.reasoning"] == 10 assert chat_spans[0]["data"]["gen_ai.usage.total_tokens"] == 192 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.input_tokens.cached"] == 69 assert chat_spans[1]["data"]["gen_ai.usage.output_tokens"] == 28 + assert chat_spans[1]["data"]["gen_ai.usage.output_tokens.reasoning"] == 11 assert chat_spans[1]["data"]["gen_ai.usage.total_tokens"] == 117 assert chat_spans[1]["data"]["gen_ai.system"] == "openai-chat"