Skip to content

Commit 7cb8094

Browse files
alexander-alderman-webbsentrivana
alexander-alderman-webb
and
sentrivana
authored
chore: Deprecate ai.monitoring.record_token_usage() (#7918)
Deprecate the function, pending removal in the upcoming major. Inline uses of the function in integrations. Co-authored-by: Ivana Kellyer <ivana.kellyer@sentry.io>
1 parent 1f4ba26 commit 7cb8094

6 files changed

Lines changed: 206 additions & 71 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_attribute() (in streaming mode) or 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: 58 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,18 +133,52 @@ def collect_chat_response_fields(
134133
set_data_normalized(span, "ai." + attr, getattr(res, attr))
135134

136135
if hasattr(res, "meta"):
136+
set_on_span = (
137+
span.set_attribute if isinstance(span, StreamedSpan) else span.set_data
138+
)
139+
137140
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,
142-
)
141+
if res.meta.billed_units.input_tokens is not None:
142+
set_on_span(
143+
SPANDATA.GEN_AI_USAGE_INPUT_TOKENS,
144+
res.meta.billed_units.input_tokens,
145+
)
146+
147+
if res.meta.billed_units.output_tokens is not None:
148+
set_on_span(
149+
SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS,
150+
res.meta.billed_units.output_tokens,
151+
)
152+
153+
if (
154+
res.meta.billed_units.input_tokens is not None
155+
and res.meta.billed_units.output_tokens is not None
156+
):
157+
set_on_span(
158+
SPANDATA.GEN_AI_USAGE_TOTAL_TOKENS,
159+
res.meta.billed_units.input_tokens
160+
+ res.meta.billed_units.output_tokens,
161+
)
143162
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,
148-
)
163+
if res.meta.tokens.input_tokens is not None:
164+
set_on_span(
165+
SPANDATA.GEN_AI_USAGE_INPUT_TOKENS, res.meta.tokens.input_tokens
166+
)
167+
168+
if res.meta.tokens.output_tokens is not None:
169+
set_on_span(
170+
SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS,
171+
res.meta.tokens.output_tokens,
172+
)
173+
174+
if (
175+
res.meta.tokens.input_tokens is not None
176+
and res.meta.tokens.output_tokens is not None
177+
):
178+
set_on_span(
179+
SPANDATA.GEN_AI_USAGE_TOTAL_TOKENS,
180+
res.meta.tokens.input_tokens + res.meta.tokens.output_tokens,
181+
)
149182

150183
if hasattr(res.meta, "warnings"):
151184
set_data_normalized(span, SPANDATA.AI_WARNINGS, res.meta.warnings)
@@ -268,13 +301,17 @@ def new_embed(*args: "Any", **kwargs: "Any") -> "Any":
268301
"sentry.origin": CohereIntegration.origin,
269302
},
270303
)
304+
305+
set_on_span = span_ctx.set_attribute
271306
else:
272307
span_ctx = get_start_span_function()(
273308
op=consts.OP.COHERE_EMBEDDINGS_CREATE,
274309
name="Cohere Embedding Creation",
275310
origin=CohereIntegration.origin,
276311
)
277312

313+
set_on_span = span_ctx.set_data
314+
278315
with span_ctx as span:
279316
if "texts" in kwargs and _should_record(integration, "inputs"):
280317
if isinstance(kwargs["texts"], str):
@@ -302,11 +339,16 @@ def new_embed(*args: "Any", **kwargs: "Any") -> "Any":
302339
and hasattr(res.meta, "billed_units")
303340
and hasattr(res.meta.billed_units, "input_tokens")
304341
):
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-
)
342+
if res.meta.billed_units.input_tokens is not None:
343+
set_on_span(
344+
SPANDATA.GEN_AI_USAGE_INPUT_TOKENS,
345+
res.meta.billed_units.input_tokens,
346+
)
347+
348+
set_on_span(
349+
SPANDATA.GEN_AI_USAGE_TOTAL_TOKENS,
350+
res.meta.billed_units.input_tokens,
351+
)
310352
return res
311353

312354
return new_embed

‎sentry_sdk/integrations/huggingface_hub.py‎

Lines changed: 55 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,6 @@
44
from typing import TYPE_CHECKING, cast
55

66
import sentry_sdk
7-
from sentry_sdk.ai.monitoring import record_token_usage
87
from sentry_sdk.ai.utils import (
98
_set_span_data_attribute,
109
get_start_span_function,
@@ -13,7 +12,6 @@
1312
from sentry_sdk.consts import OP, SPANDATA
1413
from sentry_sdk.integrations import DidNotEnable, Integration
1514
from sentry_sdk.scope import should_send_default_pii
16-
from sentry_sdk.traces import StreamedSpan
1715
from sentry_sdk.tracing_utils import has_span_streaming_enabled
1816
from sentry_sdk.utils import (
1917
capture_internal_exceptions,
@@ -23,13 +21,12 @@
2321
)
2422

2523
if TYPE_CHECKING:
26-
from typing import Any, Callable, Iterable, Union
24+
from typing import Any, Callable, Iterable
2725

2826
from huggingface_hub import (
2927
ChatCompletionStreamOutput,
3028
)
3129

32-
from sentry_sdk.tracing import Span
3330

3431
try:
3532
import huggingface_hub.inference._client
@@ -97,7 +94,6 @@ def new_huggingface_task(*args: "Any", **kwargs: "Any") -> "Any":
9794
model = hf_client.model or kwargs.get("model") or ""
9895
operation_name = op.split(".")[-1]
9996

100-
span: "Union[Span, StreamedSpan]"
10197
if has_span_streaming_enabled(client.options):
10298
span = sentry_sdk.traces.start_span(
10399
name=f"{operation_name} {model}",
@@ -106,12 +102,16 @@ def new_huggingface_task(*args: "Any", **kwargs: "Any") -> "Any":
106102
"sentry.origin": HuggingfaceHubIntegration.origin,
107103
},
108104
)
105+
106+
set_on_span = span.set_attribute
109107
else:
110108
span = get_start_span_function()(
111109
op=op,
112110
name=f"{operation_name} {model}",
113111
origin=HuggingfaceHubIntegration.origin,
114112
)
113+
114+
set_on_span = span.set_data
115115
span.__enter__()
116116

117117
_set_span_data_attribute(span, SPANDATA.GEN_AI_OPERATION_NAME, operation_name)
@@ -258,16 +258,30 @@ def new_huggingface_task(*args: "Any", **kwargs: "Any") -> "Any":
258258
)
259259

260260
if usage is not None:
261-
record_token_usage(
262-
span,
263-
input_tokens=usage.prompt_tokens,
264-
output_tokens=usage.completion_tokens,
265-
total_tokens=usage.total_tokens,
266-
)
261+
if usage is not None and usage.prompt_tokens is not None:
262+
set_on_span(SPANDATA.GEN_AI_USAGE_INPUT_TOKENS, usage.prompt_tokens)
263+
264+
if usage is not None and usage.completion_tokens is not None:
265+
set_on_span(
266+
SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS, usage.completion_tokens
267+
)
268+
269+
if usage is not None and usage.total_tokens is not None:
270+
set_on_span(SPANDATA.GEN_AI_USAGE_TOTAL_TOKENS, usage.total_tokens)
271+
272+
elif (
273+
usage is not None
274+
and usage.prompt_tokens is not None
275+
and usage.completion_tokens is not None
276+
):
277+
set_on_span(
278+
SPANDATA.GEN_AI_USAGE_TOTAL_TOKENS,
279+
usage.prompt_tokens + usage.completion_tokens,
280+
)
267281
elif tokens_used > 0:
268-
record_token_usage(
269-
span,
270-
total_tokens=tokens_used,
282+
set_on_span(
283+
SPANDATA.GEN_AI_USAGE_TOTAL_TOKENS,
284+
tokens_used,
271285
)
272286

273287
# If the response is not a generator (meaning a streaming response)
@@ -331,10 +345,7 @@ def new_details_iterator() -> "Iterable[Any]":
331345
)
332346

333347
if tokens_used > 0:
334-
record_token_usage(
335-
span,
336-
total_tokens=tokens_used,
337-
)
348+
set_on_span(SPANDATA.GEN_AI_USAGE_TOTAL_TOKENS, tokens_used)
338349

339350
span.__exit__(None, None, None)
340351

@@ -443,12 +454,32 @@ def new_iterator() -> "Iterable[ChatCompletionStreamOutput]":
443454
text_response,
444455
)
445456

446-
if usage is not None:
447-
record_token_usage(
448-
span,
449-
input_tokens=usage.prompt_tokens,
450-
output_tokens=usage.completion_tokens,
451-
total_tokens=usage.total_tokens,
457+
if usage is None:
458+
span.__exit__(None, None, None)
459+
return
460+
461+
if usage.prompt_tokens is not None:
462+
set_on_span(
463+
SPANDATA.GEN_AI_USAGE_INPUT_TOKENS, usage.prompt_tokens
464+
)
465+
466+
if usage.completion_tokens is not None:
467+
set_on_span(
468+
SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS,
469+
usage.completion_tokens,
470+
)
471+
472+
if usage.total_tokens is not None:
473+
set_on_span(
474+
SPANDATA.GEN_AI_USAGE_TOTAL_TOKENS, usage.total_tokens
475+
)
476+
elif (
477+
usage.prompt_tokens is not None
478+
and usage.completion_tokens is not None
479+
):
480+
set_on_span(
481+
SPANDATA.GEN_AI_USAGE_TOTAL_TOKENS,
482+
usage.prompt_tokens + usage.completion_tokens,
452483
)
453484

454485
span.__exit__(None, None, None)

‎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.

0 commit comments

Comments
 (0)