diff --git a/sentry_sdk/ai/monitoring.py b/sentry_sdk/ai/monitoring.py deleted file mode 100644 index d0664adac4..0000000000 --- a/sentry_sdk/ai/monitoring.py +++ /dev/null @@ -1,49 +0,0 @@ -from typing import TYPE_CHECKING - -from sentry_sdk.consts import SPANDATA - -if TYPE_CHECKING: - from typing import Any, Awaitable, Callable, Optional, TypeVar, Union - - from sentry_sdk.traces import Span - - F = TypeVar("F", bound=Union[Callable[..., Any], Callable[..., Awaitable[Any]]]) - - -def record_token_usage( - span: "Span", - input_tokens: "Optional[int]" = None, - input_tokens_cached: "Optional[int]" = None, - input_tokens_cache_write: "Optional[int]" = None, - output_tokens: "Optional[int]" = None, - output_tokens_reasoning: "Optional[int]" = None, - total_tokens: "Optional[int]" = None, -) -> None: - if input_tokens is not None: - span.set_attribute(SPANDATA.GEN_AI_USAGE_INPUT_TOKENS, input_tokens) - - if input_tokens_cached is not None: - span.set_attribute( - SPANDATA.GEN_AI_USAGE_CACHE_READ_INPUT_TOKENS, input_tokens_cached - ) - - if input_tokens_cache_write is not None: - span.set_attribute( - SPANDATA.GEN_AI_USAGE_CACHE_CREATION_INPUT_TOKENS, - input_tokens_cache_write, - ) - - if output_tokens is not None: - span.set_attribute(SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS, output_tokens) - - if output_tokens_reasoning is not None: - span.set_attribute( - SPANDATA.GEN_AI_USAGE_REASONING_OUTPUT_TOKENS, - output_tokens_reasoning, - ) - - if total_tokens is None and input_tokens is not None and output_tokens is not None: - total_tokens = input_tokens + output_tokens - - if total_tokens is not None: - span.set_attribute(SPANDATA.GEN_AI_USAGE_TOTAL_TOKENS, total_tokens) diff --git a/sentry_sdk/integrations/anthropic.py b/sentry_sdk/integrations/anthropic.py index 070fbbb0cc..39965eccbe 100644 --- a/sentry_sdk/integrations/anthropic.py +++ b/sentry_sdk/integrations/anthropic.py @@ -5,7 +5,6 @@ from typing import TYPE_CHECKING, cast import sentry_sdk -from sentry_sdk.ai.monitoring import record_token_usage from sentry_sdk.ai.utils import ( GEN_AI_ALLOWED_MESSAGE_ROLES, normalize_message_roles, @@ -692,13 +691,27 @@ def _set_output_data( span, SPANDATA.GEN_AI_RESPONSE_TEXT, output_messages["response"] ) - record_token_usage( - span, - input_tokens=input_tokens, - output_tokens=output_tokens, - input_tokens_cached=cache_read_input_tokens, - input_tokens_cache_write=cache_write_input_tokens, - ) + if input_tokens is not None: + span.set_attribute(SPANDATA.GEN_AI_USAGE_INPUT_TOKENS, input_tokens) + + if cache_read_input_tokens is not None: + span.set_attribute( + SPANDATA.GEN_AI_USAGE_CACHE_READ_INPUT_TOKENS, cache_read_input_tokens + ) + + if cache_write_input_tokens is not None: + span.set_attribute( + SPANDATA.GEN_AI_USAGE_CACHE_CREATION_INPUT_TOKENS, + cache_write_input_tokens, + ) + + if output_tokens is not None: + span.set_attribute(SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS, output_tokens) + + if input_tokens is not None and output_tokens is not None: + span.set_attribute( + SPANDATA.GEN_AI_USAGE_TOTAL_TOKENS, input_tokens + output_tokens + ) def _sentry_patched_create_sync(f: "Any", *args: "Any", **kwargs: "Any") -> "Any": diff --git a/sentry_sdk/integrations/cohere.py b/sentry_sdk/integrations/cohere.py index 5aa25acf4a..740eb5ea6b 100644 --- a/sentry_sdk/integrations/cohere.py +++ b/sentry_sdk/integrations/cohere.py @@ -3,7 +3,6 @@ from typing import TYPE_CHECKING from sentry_sdk import consts -from sentry_sdk.ai.monitoring import record_token_usage from sentry_sdk.ai.utils import set_data_normalized from sentry_sdk.consts import SPANDATA from sentry_sdk.traces import Span @@ -133,17 +132,47 @@ def collect_chat_response_fields( if hasattr(res, "meta"): if hasattr(res.meta, "billed_units"): - record_token_usage( - span, - input_tokens=res.meta.billed_units.input_tokens, - output_tokens=res.meta.billed_units.output_tokens, - ) + if res.meta.billed_units.input_tokens is not None: + span.set_attribute( + SPANDATA.GEN_AI_USAGE_INPUT_TOKENS, + res.meta.billed_units.input_tokens, + ) + + if res.meta.billed_units.output_tokens is not None: + span.set_attribute( + SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS, + res.meta.billed_units.output_tokens, + ) + + if ( + res.meta.billed_units.input_tokens is not None + and res.meta.billed_units.output_tokens is not None + ): + span.set_attribute( + SPANDATA.GEN_AI_USAGE_TOTAL_TOKENS, + res.meta.billed_units.input_tokens + + res.meta.billed_units.output_tokens, + ) elif hasattr(res.meta, "tokens"): - record_token_usage( - span, - input_tokens=res.meta.tokens.input_tokens, - output_tokens=res.meta.tokens.output_tokens, - ) + if res.meta.tokens.input_tokens is not None: + span.set_attribute( + SPANDATA.GEN_AI_USAGE_INPUT_TOKENS, res.meta.tokens.input_tokens + ) + + if res.meta.tokens.output_tokens is not None: + span.set_attribute( + SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS, + res.meta.tokens.output_tokens, + ) + + if ( + res.meta.tokens.input_tokens is not None + and res.meta.tokens.output_tokens is not None + ): + span.set_attribute( + SPANDATA.GEN_AI_USAGE_TOTAL_TOKENS, + res.meta.tokens.input_tokens + res.meta.tokens.output_tokens, + ) if hasattr(res.meta, "warnings"): set_data_normalized(span, SPANDATA.AI_WARNINGS, res.meta.warnings) @@ -277,11 +306,17 @@ def new_embed(*args: "Any", **kwargs: "Any") -> "Any": and hasattr(res.meta, "billed_units") and hasattr(res.meta.billed_units, "input_tokens") ): - record_token_usage( - span, - input_tokens=res.meta.billed_units.input_tokens, - total_tokens=res.meta.billed_units.input_tokens, - ) + if res.meta.billed_units.input_tokens is not None: + span.set_attribute( + SPANDATA.GEN_AI_USAGE_INPUT_TOKENS, + res.meta.billed_units.input_tokens, + ) + + span.set_attribute( + SPANDATA.GEN_AI_USAGE_TOTAL_TOKENS, + res.meta.billed_units.input_tokens, + ) + return res return new_embed diff --git a/sentry_sdk/integrations/huggingface_hub.py b/sentry_sdk/integrations/huggingface_hub.py index 1fb951c5dc..c8660ce24f 100644 --- a/sentry_sdk/integrations/huggingface_hub.py +++ b/sentry_sdk/integrations/huggingface_hub.py @@ -4,7 +4,6 @@ from typing import TYPE_CHECKING, cast import sentry_sdk -from sentry_sdk.ai.monitoring import record_token_usage from sentry_sdk.ai.utils import set_data_normalized from sentry_sdk.consts import OP, SPANDATA from sentry_sdk.integrations import DidNotEnable, Integration, _check_minimum_version @@ -244,19 +243,24 @@ def new_huggingface_task(*args: "Any", **kwargs: "Any") -> "Any": text_response, ) - if usage is not None: - record_token_usage( - span, - input_tokens=usage.prompt_tokens, - output_tokens=usage.completion_tokens, - total_tokens=usage.total_tokens, + if usage is not None and usage.prompt_tokens is not None: + span.set_attribute( + SPANDATA.GEN_AI_USAGE_INPUT_TOKENS, usage.prompt_tokens ) - elif tokens_used > 0: - record_token_usage( - span, - total_tokens=tokens_used, + + if usage is not None and usage.completion_tokens is not None: + span.set_attribute( + SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS, usage.completion_tokens + ) + + if usage is not None and usage.total_tokens is not None: + span.set_attribute( + SPANDATA.GEN_AI_USAGE_TOTAL_TOKENS, usage.total_tokens ) + elif tokens_used > 0 and tokens_used is not None: + span.set_attribute(SPANDATA.GEN_AI_USAGE_TOTAL_TOKENS, tokens_used) + # If the response is not a generator (meaning a streaming response) # we are done and can return the response if not inspect.isgenerator(res): @@ -318,9 +322,8 @@ def new_details_iterator() -> "Iterable[Any]": ) if tokens_used > 0: - record_token_usage( - span, - total_tokens=tokens_used, + span.set_attribute( + SPANDATA.GEN_AI_USAGE_TOTAL_TOKENS, tokens_used ) span.__exit__(None, None, None) @@ -430,12 +433,20 @@ def new_iterator() -> "Iterable[ChatCompletionStreamOutput]": text_response, ) - if usage is not None: - record_token_usage( - span, - input_tokens=usage.prompt_tokens, - output_tokens=usage.completion_tokens, - total_tokens=usage.total_tokens, + if usage is not None and usage.prompt_tokens is not None: + span.set_attribute( + SPANDATA.GEN_AI_USAGE_INPUT_TOKENS, usage.prompt_tokens + ) + + if usage is not None and usage.completion_tokens is not None: + span.set_attribute( + SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS, + usage.completion_tokens, + ) + + if usage is not None and usage.total_tokens is not None: + span.set_attribute( + SPANDATA.GEN_AI_USAGE_TOTAL_TOKENS, usage.total_tokens ) span.__exit__(None, None, None) diff --git a/sentry_sdk/integrations/litellm.py b/sentry_sdk/integrations/litellm.py index 7ca8894aed..9e87e7f04d 100644 --- a/sentry_sdk/integrations/litellm.py +++ b/sentry_sdk/integrations/litellm.py @@ -3,7 +3,6 @@ import sentry_sdk from sentry_sdk import consts -from sentry_sdk.ai.monitoring import record_token_usage from sentry_sdk.ai.utils import ( set_data_normalized, transform_openai_content_part, @@ -229,12 +228,25 @@ def _success_callback( # Record token usage if hasattr(completion_response, "usage"): usage = completion_response.usage - record_token_usage( - span, - input_tokens=getattr(usage, "prompt_tokens", None), - output_tokens=getattr(usage, "completion_tokens", None), - total_tokens=getattr(usage, "total_tokens", None), - ) + + input_tokens = getattr(usage, "prompt_tokens", None) + if input_tokens is not None: + span.set_attribute(SPANDATA.GEN_AI_USAGE_INPUT_TOKENS, input_tokens) + + output_tokens = getattr(usage, "completion_tokens", None) + if output_tokens is not None: + span.set_attribute(SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS, output_tokens) + + total_tokens = getattr(usage, "total_tokens", None) + if ( + total_tokens is None + and input_tokens is not None + and output_tokens is not None + ): + total_tokens = input_tokens + output_tokens + + if total_tokens is not None: + span.set_attribute(SPANDATA.GEN_AI_USAGE_TOTAL_TOKENS, total_tokens) finally: is_streaming = kwargs.get("stream") diff --git a/sentry_sdk/integrations/openai.py b/sentry_sdk/integrations/openai.py index 41a5a07be7..1bc2c2cc2d 100644 --- a/sentry_sdk/integrations/openai.py +++ b/sentry_sdk/integrations/openai.py @@ -34,7 +34,6 @@ from sentry_sdk.ai._openai_responses_api import ( _transform_tool_definitions as _transform_tool_definitions_responses, ) -from sentry_sdk.ai.monitoring import record_token_usage from sentry_sdk.ai.utils import ( normalize_message_roles, set_data_normalized, @@ -238,19 +237,32 @@ def _calculate_completions_token_usage( # Do not set token data if it is 0 input_tokens = input_tokens or None + if input_tokens is not None: + span.set_attribute(SPANDATA.GEN_AI_USAGE_INPUT_TOKENS, input_tokens) + input_tokens_cached = input_tokens_cached or None + if input_tokens_cached is not None: + span.set_attribute( + SPANDATA.GEN_AI_USAGE_CACHE_READ_INPUT_TOKENS, input_tokens_cached + ) + output_tokens = output_tokens or None + if output_tokens is not None: + span.set_attribute(SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS, output_tokens) + output_tokens_reasoning = output_tokens_reasoning or None + if output_tokens_reasoning is not None: + span.set_attribute( + SPANDATA.GEN_AI_USAGE_REASONING_OUTPUT_TOKENS, + output_tokens_reasoning, + ) + total_tokens = total_tokens or None + if total_tokens is None and input_tokens is not None and output_tokens is not None: + total_tokens = input_tokens + output_tokens - record_token_usage( - span, - input_tokens=input_tokens, - input_tokens_cached=input_tokens_cached, - output_tokens=output_tokens, - output_tokens_reasoning=output_tokens_reasoning, - total_tokens=total_tokens, - ) + if total_tokens is not None: + span.set_attribute(SPANDATA.GEN_AI_USAGE_TOTAL_TOKENS, total_tokens) def _calculate_responses_token_usage( @@ -316,19 +328,32 @@ def _calculate_responses_token_usage( # Do not set token data if it is 0 input_tokens = input_tokens or None + if input_tokens is not None: + span.set_attribute(SPANDATA.GEN_AI_USAGE_INPUT_TOKENS, input_tokens) + input_tokens_cached = input_tokens_cached or None + if input_tokens_cached is not None: + span.set_attribute( + SPANDATA.GEN_AI_USAGE_CACHE_READ_INPUT_TOKENS, input_tokens_cached + ) + output_tokens = output_tokens or None + if output_tokens is not None: + span.set_attribute(SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS, output_tokens) + output_tokens_reasoning = output_tokens_reasoning or None + if output_tokens_reasoning is not None: + span.set_attribute( + SPANDATA.GEN_AI_USAGE_REASONING_OUTPUT_TOKENS, + output_tokens_reasoning, + ) + total_tokens = total_tokens or None + if total_tokens is None and input_tokens is not None and output_tokens is not None: + total_tokens = input_tokens + output_tokens - record_token_usage( - span, - input_tokens=input_tokens, - input_tokens_cached=input_tokens_cached, - output_tokens=output_tokens, - output_tokens_reasoning=output_tokens_reasoning, - total_tokens=total_tokens, - ) + if total_tokens is not None: + span.set_attribute(SPANDATA.GEN_AI_USAGE_TOTAL_TOKENS, total_tokens) def _set_responses_api_input_data(