Skip to content
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
49 changes: 0 additions & 49 deletions sentry_sdk/ai/monitoring.py

This file was deleted.

29 changes: 21 additions & 8 deletions sentry_sdk/integrations/anthropic.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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":
Expand Down
67 changes: 51 additions & 16 deletions sentry_sdk/integrations/cohere.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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
51 changes: 31 additions & 20 deletions sentry_sdk/integrations/huggingface_hub.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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):
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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)
Expand Down
26 changes: 19 additions & 7 deletions sentry_sdk/integrations/litellm.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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")
Expand Down
59 changes: 42 additions & 17 deletions sentry_sdk/integrations/openai.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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(
Expand Down Expand Up @@ -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(
Expand Down
Loading