diff --git a/sentry_sdk/consts.py b/sentry_sdk/consts.py index c6db94d780..ef468d9602 100644 --- a/sentry_sdk/consts.py +++ b/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. diff --git a/sentry_sdk/integrations/langchain.py b/sentry_sdk/integrations/langchain.py index 61b3cda772..e53851a676 100644 --- a/sentry_sdk/integrations/langchain.py +++ b/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,14 +728,13 @@ 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: @@ -735,14 +742,44 @@ def _extract_tokens_from_generations( 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" diff --git a/tests/conftest.py b/tests/conftest.py index f3ae302057..599c075224 100644 --- a/tests/conftest.py +++ b/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( diff --git a/tests/integrations/langchain/test_langchain.py b/tests/integrations/langchain/test_langchain.py index 892491bec4..e0ca232302 100644 --- a/tests/integrations/langchain/test_langchain.py +++ b/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):