Skip to content

Commit 6f26527

Browse files
alexander-alderman-webbalexander-alderman-webb
committed
chore: Deprecate ai.monitoring.record_token_usage()
1 parent 4b9d329 commit 6f26527

5 files changed

Lines changed: 150 additions & 47 deletions

File tree

‎sentry_sdk/ai/monitoring.py‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -192,6 +192,13 @@ def record_token_usage(
192192
output_tokens_reasoning: "Optional[int]" = None,
193193
total_tokens: "Optional[int]" = None,
194194
) -> None:
195+
warnings.warn(
196+
"record_token_usage() is deprecated and will be removed in version 3.0 of sentry-sdk. "
197+
"Use the manual span API instead, e.g. Span.set_data().",
198+
DeprecationWarning,
199+
stacklevel=2,
200+
)
201+
195202
# TODO: move pipeline name elsewhere
196203
ai_pipeline_name = get_ai_pipeline_name()
197204
if ai_pipeline_name:

‎sentry_sdk/integrations/anthropic.py‎

Lines changed: 17 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,6 @@
55
from typing import TYPE_CHECKING, cast
66

77
import sentry_sdk
8-
from sentry_sdk.ai.monitoring import record_token_usage
98
from sentry_sdk.ai.utils import (
109
GEN_AI_ALLOWED_MESSAGE_ROLES,
1110
get_start_span_function,
@@ -714,13 +713,23 @@ def _set_output_data(
714713
span, SPANDATA.GEN_AI_RESPONSE_TEXT, output_messages["response"]
715714
)
716715

717-
record_token_usage(
718-
span,
719-
input_tokens=input_tokens,
720-
output_tokens=output_tokens,
721-
input_tokens_cached=cache_read_input_tokens,
722-
input_tokens_cache_write=cache_write_input_tokens,
723-
)
716+
if input_tokens is not None:
717+
set_on_span(SPANDATA.GEN_AI_USAGE_INPUT_TOKENS, input_tokens)
718+
719+
if cache_read_input_tokens is not None:
720+
set_on_span(SPANDATA.GEN_AI_USAGE_INPUT_TOKENS_CACHED, cache_read_input_tokens)
721+
722+
if cache_write_input_tokens is not None:
723+
set_on_span(
724+
SPANDATA.GEN_AI_USAGE_INPUT_TOKENS_CACHE_WRITE,
725+
cache_write_input_tokens,
726+
)
727+
728+
if output_tokens is not None:
729+
set_on_span(SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS, output_tokens)
730+
731+
if input_tokens is not None and output_tokens is not None:
732+
set_on_span(SPANDATA.GEN_AI_USAGE_TOTAL_TOKENS, input_tokens + output_tokens)
724733

725734

726735
def _sentry_patched_create_sync(f: "Any", *args: "Any", **kwargs: "Any") -> "Any":

‎sentry_sdk/integrations/cohere.py‎

Lines changed: 57 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,6 @@
33
from typing import TYPE_CHECKING
44

55
from sentry_sdk import consts
6-
from sentry_sdk.ai.monitoring import record_token_usage
76
from sentry_sdk.ai.utils import get_start_span_function, set_data_normalized
87
from sentry_sdk.consts import SPANDATA
98
from sentry_sdk.traces import StreamedSpan
@@ -134,19 +133,52 @@ def collect_chat_response_fields(
134133
set_data_normalized(span, "ai." + attr, getattr(res, attr))
135134

136135
if hasattr(res, "meta"):
137-
if hasattr(res.meta, "billed_units"):
138-
record_token_usage(
139-
span,
140-
input_tokens=res.meta.billed_units.input_tokens,
141-
output_tokens=res.meta.billed_units.output_tokens,
136+
set_on_span = (
137+
span.set_attribute if isinstance(span, StreamedSpan) else span.set_data
138+
)
139+
140+
if res.meta.billed_units.input_tokens is not None:
141+
set_on_span(
142+
SPANDATA.GEN_AI_USAGE_INPUT_TOKENS,
143+
res.meta.billed_units.input_tokens,
142144
)
143-
elif hasattr(res.meta, "tokens"):
144-
record_token_usage(
145-
span,
146-
input_tokens=res.meta.tokens.input_tokens,
147-
output_tokens=res.meta.tokens.output_tokens,
145+
146+
if res.meta.billed_units.output_tokens is not None:
147+
set_on_span(
148+
SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS,
149+
res.meta.billed_units.output_tokens,
148150
)
149151

152+
if (
153+
res.meta.billed_units.input_tokens is not None
154+
and res.meta.billed_units.output_tokens is not None
155+
):
156+
set_on_span(
157+
SPANDATA.GEN_AI_USAGE_TOTAL_TOKENS,
158+
res.meta.billed_units.input_tokens
159+
+ res.meta.billed_units.output_tokens,
160+
)
161+
elif hasattr(res.meta, "tokens"):
162+
if res.meta.tokens.input_tokens is not None:
163+
set_on_span(
164+
SPANDATA.GEN_AI_USAGE_INPUT_TOKENS, res.meta.tokens.input_tokens
165+
)
166+
167+
if res.meta.tokens.output_tokens is not None:
168+
set_on_span(
169+
SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS,
170+
res.meta.tokens.output_tokens,
171+
)
172+
173+
if (
174+
res.meta.tokens.input_tokens is not None
175+
and res.meta.tokens.output_tokens is not None
176+
):
177+
set_on_span(
178+
SPANDATA.GEN_AI_USAGE_TOTAL_TOKENS,
179+
res.meta.tokens.input_tokens + res.meta.tokens.output_tokens,
180+
)
181+
150182
if hasattr(res.meta, "warnings"):
151183
set_data_normalized(span, SPANDATA.AI_WARNINGS, res.meta.warnings)
152184

@@ -268,13 +300,17 @@ def new_embed(*args: "Any", **kwargs: "Any") -> "Any":
268300
"sentry.origin": CohereIntegration.origin,
269301
},
270302
)
303+
304+
set_on_span = span_ctx.set_attribute
271305
else:
272306
span_ctx = get_start_span_function()(
273307
op=consts.OP.COHERE_EMBEDDINGS_CREATE,
274308
name="Cohere Embedding Creation",
275309
origin=CohereIntegration.origin,
276310
)
277311

312+
set_on_span = span_ctx.set_data
313+
278314
with span_ctx as span:
279315
if "texts" in kwargs and _should_record(integration, "inputs"):
280316
if isinstance(kwargs["texts"], str):
@@ -302,11 +338,16 @@ def new_embed(*args: "Any", **kwargs: "Any") -> "Any":
302338
and hasattr(res.meta, "billed_units")
303339
and hasattr(res.meta.billed_units, "input_tokens")
304340
):
305-
record_token_usage(
306-
span,
307-
input_tokens=res.meta.billed_units.input_tokens,
308-
total_tokens=res.meta.billed_units.input_tokens,
309-
)
341+
if res.meta.billed_units.input_tokens is not None:
342+
set_on_span(
343+
SPANDATA.GEN_AI_USAGE_INPUT_TOKENS,
344+
res.meta.billed_units.input_tokens,
345+
)
346+
347+
set_on_span(
348+
SPANDATA.GEN_AI_USAGE_TOTAL_TOKENS,
349+
res.meta.billed_units.input_tokens,
350+
)
310351
return res
311352

312353
return new_embed

‎sentry_sdk/integrations/litellm.py‎

Lines changed: 23 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,6 @@
33

44
import sentry_sdk
55
from sentry_sdk import consts
6-
from sentry_sdk.ai.monitoring import record_token_usage
76
from sentry_sdk.ai.utils import (
87
get_start_span_function,
98
set_data_normalized,
@@ -14,6 +13,7 @@
1413
from sentry_sdk.consts import SPANDATA
1514
from sentry_sdk.integrations import DidNotEnable, Integration
1615
from sentry_sdk.scope import should_send_default_pii
16+
from sentry_sdk.traces import StreamedSpan
1717
from sentry_sdk.tracing_utils import (
1818
has_span_streaming_enabled,
1919
should_truncate_gen_ai_input,
@@ -257,13 +257,30 @@ def _success_callback(
257257
# Record token usage
258258
if hasattr(completion_response, "usage"):
259259
usage = completion_response.usage
260-
record_token_usage(
261-
span,
262-
input_tokens=getattr(usage, "prompt_tokens", None),
263-
output_tokens=getattr(usage, "completion_tokens", None),
264-
total_tokens=getattr(usage, "total_tokens", None),
260+
261+
set_on_span = (
262+
span.set_attribute if isinstance(span, StreamedSpan) else span.set_data
265263
)
266264

265+
input_tokens = getattr(usage, "prompt_tokens", None)
266+
if input_tokens is not None:
267+
set_on_span(SPANDATA.GEN_AI_USAGE_INPUT_TOKENS, input_tokens)
268+
269+
output_tokens = getattr(usage, "completion_tokens", None)
270+
if output_tokens is not None:
271+
set_on_span(SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS, output_tokens)
272+
273+
total_tokens = getattr(usage, "total_tokens", None)
274+
if (
275+
total_tokens is None
276+
and input_tokens is not None
277+
and output_tokens is not None
278+
):
279+
total_tokens = input_tokens + output_tokens
280+
281+
if total_tokens is not None:
282+
set_on_span(SPANDATA.GEN_AI_USAGE_TOTAL_TOKENS, total_tokens)
283+
267284
finally:
268285
is_streaming = kwargs.get("stream")
269286
# Callback is fired multiple times when streaming a response.

‎sentry_sdk/integrations/openai.py‎

Lines changed: 46 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -34,7 +34,6 @@
3434
from sentry_sdk.ai._openai_responses_api import (
3535
_transform_tool_definitions as _transform_tool_definitions_responses,
3636
)
37-
from sentry_sdk.ai.monitoring import record_token_usage
3837
from sentry_sdk.ai.utils import (
3938
get_start_span_function,
4039
normalize_message_roles,
@@ -239,21 +238,36 @@ def _calculate_completions_token_usage(
239238
if hasattr(choice, "message") and hasattr(choice.message, "content"):
240239
output_tokens += count_tokens(choice.message.content)
241240

241+
set_on_span = (
242+
span.set_attribute if isinstance(span, StreamedSpan) else span.set_data
243+
)
244+
242245
# Do not set token data if it is 0
243246
input_tokens = input_tokens or None
247+
if input_tokens is not None:
248+
set_on_span(SPANDATA.GEN_AI_USAGE_INPUT_TOKENS, input_tokens)
249+
244250
input_tokens_cached = input_tokens_cached or None
251+
if input_tokens_cached is not None:
252+
set_on_span(SPANDATA.GEN_AI_USAGE_INPUT_TOKENS_CACHED, input_tokens_cached)
253+
245254
output_tokens = output_tokens or None
255+
if output_tokens is not None:
256+
set_on_span(SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS, output_tokens)
257+
246258
output_tokens_reasoning = output_tokens_reasoning or None
259+
if output_tokens_reasoning is not None:
260+
set_on_span(
261+
SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS_REASONING,
262+
output_tokens_reasoning,
263+
)
264+
247265
total_tokens = total_tokens or None
266+
if total_tokens is None and input_tokens is not None and output_tokens is not None:
267+
total_tokens = input_tokens + output_tokens
248268

249-
record_token_usage(
250-
span,
251-
input_tokens=input_tokens,
252-
input_tokens_cached=input_tokens_cached,
253-
output_tokens=output_tokens,
254-
output_tokens_reasoning=output_tokens_reasoning,
255-
total_tokens=total_tokens,
256-
)
269+
if total_tokens is not None:
270+
set_on_span(SPANDATA.GEN_AI_USAGE_TOTAL_TOKENS, total_tokens)
257271

258272

259273
def _calculate_responses_token_usage(
@@ -317,21 +331,36 @@ def _calculate_responses_token_usage(
317331
if hasattr(content_item, "text"):
318332
output_tokens += count_tokens(content_item.text)
319333

334+
set_on_span = (
335+
span.set_attribute if isinstance(span, StreamedSpan) else span.set_data
336+
)
337+
320338
# Do not set token data if it is 0
321339
input_tokens = input_tokens or None
340+
if input_tokens is not None:
341+
set_on_span(SPANDATA.GEN_AI_USAGE_INPUT_TOKENS, input_tokens)
342+
322343
input_tokens_cached = input_tokens_cached or None
344+
if input_tokens_cached is not None:
345+
set_on_span(SPANDATA.GEN_AI_USAGE_INPUT_TOKENS_CACHED, input_tokens_cached)
346+
323347
output_tokens = output_tokens or None
348+
if output_tokens is not None:
349+
set_on_span(SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS, output_tokens)
350+
324351
output_tokens_reasoning = output_tokens_reasoning or None
352+
if output_tokens_reasoning is not None:
353+
set_on_span(
354+
SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS_REASONING,
355+
output_tokens_reasoning,
356+
)
357+
325358
total_tokens = total_tokens or None
359+
if total_tokens is None and input_tokens is not None and output_tokens is not None:
360+
total_tokens = input_tokens + output_tokens
326361

327-
record_token_usage(
328-
span,
329-
input_tokens=input_tokens,
330-
input_tokens_cached=input_tokens_cached,
331-
output_tokens=output_tokens,
332-
output_tokens_reasoning=output_tokens_reasoning,
333-
total_tokens=total_tokens,
334-
)
362+
if total_tokens is not None:
363+
set_on_span(SPANDATA.GEN_AI_USAGE_TOTAL_TOKENS, total_tokens)
335364

336365

337366
def _set_responses_api_input_data(

0 commit comments

Comments
 (0)