Skip to content

Commit 07ba0f0

Browse files
antonpirkerantonpirker
authored
Make use of SPANDATA consistent in AI integrations.` (#4376)
Not using plain strings, but always use `SPANDATA`. Follow up to this PR: #4373
1 parent bf0393b commit 07ba0f0

File tree

7 files changed

+72
-54
lines changed

7 files changed

+72
-54
lines changed

‎sentry_sdk/ai/monitoring.py

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -107,14 +107,14 @@ def record_token_usage(
107107
if ai_pipeline_name:
108108
span.set_attribute(SPANDATA.AI_PIPELINE_NAME, ai_pipeline_name)
109109
if prompt_tokens is not None:
110-
span.set_attribute("ai.prompt_tokens.used", prompt_tokens)
110+
span.set_attribute(SPANDATA.AI_PROMPT_TOKENS_USED, prompt_tokens)
111111
if completion_tokens is not None:
112-
span.set_attribute("ai.completion_tokens.used", completion_tokens)
112+
span.set_attribute(SPANDATA.AI_COMPLETION_TOKENS_USED, completion_tokens)
113113
if (
114114
total_tokens is None
115115
and prompt_tokens is not None
116116
and completion_tokens is not None
117117
):
118118
total_tokens = prompt_tokens + completion_tokens
119119
if total_tokens is not None:
120-
span.set_attribute("ai.total_tokens.used", total_tokens)
120+
span.set_attribute(SPANDATA.AI_TOTAL_TOKENS_USED, total_tokens)

‎sentry_sdk/consts.py

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -251,6 +251,24 @@ class SPANDATA:
251251
Example: "qa-pipeline"
252252
"""
253253

254+
AI_PROMPT_TOKENS_USED = "ai.prompt_tokens.used"
255+
"""
256+
The number of input prompt tokens used by the model.
257+
Example: 10
258+
"""
259+
260+
AI_COMPLETION_TOKENS_USED = "ai.completion_tokens.used"
261+
"""
262+
The number of output completion tokens used by the model.
263+
Example: 10
264+
"""
265+
266+
AI_TOTAL_TOKENS_USED = "ai.total_tokens.used"
267+
"""
268+
The total number of tokens (input + output) used by the request to the model.
269+
Example: 20
270+
"""
271+
254272
AI_TEXTS = "ai.texts"
255273
"""
256274
Raw text inputs provided to the model.

‎tests/integrations/anthropic/test_anthropic.py

Lines changed: 21 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -127,9 +127,9 @@ def test_nonstreaming_create_message(
127127
assert SPANDATA.AI_INPUT_MESSAGES not in span["data"]
128128
assert SPANDATA.AI_RESPONSES not in span["data"]
129129

130-
assert span["data"]["ai.prompt_tokens.used"] == 10
131-
assert span["data"]["ai.completion_tokens.used"] == 20
132-
assert span["data"]["ai.total_tokens.used"] == 30
130+
assert span["data"][SPANDATA.AI_PROMPT_TOKENS_USED] == 10
131+
assert span["data"][SPANDATA.AI_COMPLETION_TOKENS_USED] == 20
132+
assert span["data"][SPANDATA.AI_TOTAL_TOKENS_USED] == 30
133133
assert span["data"][SPANDATA.AI_STREAMING] is False
134134

135135

@@ -197,9 +197,9 @@ async def test_nonstreaming_create_message_async(
197197
assert SPANDATA.AI_INPUT_MESSAGES not in span["data"]
198198
assert SPANDATA.AI_RESPONSES not in span["data"]
199199

200-
assert span["data"]["ai.prompt_tokens.used"] == 10
201-
assert span["data"]["ai.completion_tokens.used"] == 20
202-
assert span["data"]["ai.total_tokens.used"] == 30
200+
assert span["data"][SPANDATA.AI_PROMPT_TOKENS_USED] == 10
201+
assert span["data"][SPANDATA.AI_COMPLETION_TOKENS_USED] == 20
202+
assert span["data"][SPANDATA.AI_TOTAL_TOKENS_USED] == 30
203203
assert span["data"][SPANDATA.AI_STREAMING] is False
204204

205205

@@ -299,9 +299,9 @@ def test_streaming_create_message(
299299
assert SPANDATA.AI_INPUT_MESSAGES not in span["data"]
300300
assert SPANDATA.AI_RESPONSES not in span["data"]
301301

302-
assert span["data"]["ai.prompt_tokens.used"] == 10
303-
assert span["data"]["ai.completion_tokens.used"] == 30
304-
assert span["data"]["ai.total_tokens.used"] == 40
302+
assert span["data"][SPANDATA.AI_PROMPT_TOKENS_USED] == 10
303+
assert span["data"][SPANDATA.AI_COMPLETION_TOKENS_USED] == 30
304+
assert span["data"][SPANDATA.AI_TOTAL_TOKENS_USED] == 40
305305
assert span["data"][SPANDATA.AI_STREAMING] is True
306306

307307

@@ -404,9 +404,9 @@ async def test_streaming_create_message_async(
404404
assert SPANDATA.AI_INPUT_MESSAGES not in span["data"]
405405
assert SPANDATA.AI_RESPONSES not in span["data"]
406406

407-
assert span["data"]["ai.prompt_tokens.used"] == 10
408-
assert span["data"]["ai.completion_tokens.used"] == 30
409-
assert span["data"]["ai.total_tokens.used"] == 40
407+
assert span["data"][SPANDATA.AI_PROMPT_TOKENS_USED] == 10
408+
assert span["data"][SPANDATA.AI_COMPLETION_TOKENS_USED] == 30
409+
assert span["data"][SPANDATA.AI_TOTAL_TOKENS_USED] == 40
410410
assert span["data"][SPANDATA.AI_STREAMING] is True
411411

412412

@@ -536,9 +536,9 @@ def test_streaming_create_message_with_input_json_delta(
536536
assert SPANDATA.AI_INPUT_MESSAGES not in span["data"]
537537
assert SPANDATA.AI_RESPONSES not in span["data"]
538538

539-
assert span["data"]["ai.prompt_tokens.used"] == 366
540-
assert span["data"]["ai.completion_tokens.used"] == 51
541-
assert span["data"]["ai.total_tokens.used"] == 417
539+
assert span["data"][SPANDATA.AI_PROMPT_TOKENS_USED] == 366
540+
assert span["data"][SPANDATA.AI_COMPLETION_TOKENS_USED] == 51
541+
assert span["data"][SPANDATA.AI_TOTAL_TOKENS_USED] == 417
542542
assert span["data"][SPANDATA.AI_STREAMING] is True
543543

544544

@@ -675,9 +675,9 @@ async def test_streaming_create_message_with_input_json_delta_async(
675675
assert SPANDATA.AI_INPUT_MESSAGES not in span["data"]
676676
assert SPANDATA.AI_RESPONSES not in span["data"]
677677

678-
assert span["data"]["ai.prompt_tokens.used"] == 366
679-
assert span["data"]["ai.completion_tokens.used"] == 51
680-
assert span["data"]["ai.total_tokens.used"] == 417
678+
assert span["data"][SPANDATA.AI_PROMPT_TOKENS_USED] == 366
679+
assert span["data"][SPANDATA.AI_COMPLETION_TOKENS_USED] == 51
680+
assert span["data"][SPANDATA.AI_TOTAL_TOKENS_USED] == 417
681681
assert span["data"][SPANDATA.AI_STREAMING] is True
682682

683683

@@ -831,6 +831,6 @@ def test_add_ai_data_to_span_with_input_json_delta(sentry_init, capture_events):
831831
[{"type": "text", "text": "{'test': 'data','more': 'json'}"}]
832832
)
833833
assert span["data"][SPANDATA.AI_STREAMING] is True
834-
assert span["data"]["ai.prompt_tokens.used"] == 10
835-
assert span["data"]["ai.completion_tokens.used"] == 20
836-
assert span["data"]["ai.total_tokens.used"] == 30
834+
assert span["data"][SPANDATA.AI_PROMPT_TOKENS_USED] == 10
835+
assert span["data"][SPANDATA.AI_COMPLETION_TOKENS_USED] == 20
836+
assert span["data"][SPANDATA.AI_TOTAL_TOKENS_USED] == 30

‎tests/integrations/cohere/test_cohere.py

Lines changed: 8 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -65,9 +65,9 @@ def test_nonstreaming_chat(
6565
assert SPANDATA.AI_INPUT_MESSAGES not in span["data"]
6666
assert SPANDATA.AI_RESPONSES not in span["data"]
6767

68-
assert span["data"]["ai.completion_tokens.used"] == 10
69-
assert span["data"]["ai.prompt_tokens.used"] == 20
70-
assert span["data"]["ai.total_tokens.used"] == 30
68+
assert span["data"][SPANDATA.AI_COMPLETION_TOKENS_USED] == 10
69+
assert span["data"][SPANDATA.AI_PROMPT_TOKENS_USED] == 20
70+
assert span["data"][SPANDATA.AI_TOTAL_TOKENS_USED] == 30
7171

7272

7373
# noinspection PyTypeChecker
@@ -137,9 +137,9 @@ def test_streaming_chat(sentry_init, capture_events, send_default_pii, include_p
137137
assert SPANDATA.AI_INPUT_MESSAGES not in span["data"]
138138
assert SPANDATA.AI_RESPONSES not in span["data"]
139139

140-
assert span["data"]["ai.completion_tokens.used"] == 10
141-
assert span["data"]["ai.prompt_tokens.used"] == 20
142-
assert span["data"]["ai.total_tokens.used"] == 30
140+
assert span["data"][SPANDATA.AI_COMPLETION_TOKENS_USED] == 10
141+
assert span["data"][SPANDATA.AI_PROMPT_TOKENS_USED] == 20
142+
assert span["data"][SPANDATA.AI_TOTAL_TOKENS_USED] == 30
143143

144144

145145
def test_bad_chat(sentry_init, capture_events):
@@ -201,8 +201,8 @@ def test_embed(sentry_init, capture_events, send_default_pii, include_prompts):
201201
else:
202202
assert SPANDATA.AI_INPUT_MESSAGES not in span["data"]
203203

204-
assert span["data"]["ai.prompt_tokens.used"] == 10
205-
assert span["data"]["ai.total_tokens.used"] == 10
204+
assert span["data"][SPANDATA.AI_PROMPT_TOKENS_USED] == 10
205+
assert span["data"][SPANDATA.AI_TOTAL_TOKENS_USED] == 10
206206

207207

208208
def test_span_origin_chat(sentry_init, capture_events):

‎tests/integrations/huggingface_hub/test_huggingface_hub.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -75,7 +75,7 @@ def test_nonstreaming_chat_completion(
7575
assert SPANDATA.AI_RESPONSES not in span["data"]
7676

7777
if details_arg:
78-
assert span["data"]["ai.total_tokens.used"] == 10
78+
assert span["data"][SPANDATA.AI_TOTAL_TOKENS_USED] == 10
7979

8080

8181
@pytest.mark.parametrize(
@@ -134,7 +134,7 @@ def test_streaming_chat_completion(
134134
assert SPANDATA.AI_RESPONSES not in span["data"]
135135

136136
if details_arg:
137-
assert span["data"]["ai.total_tokens.used"] == 10
137+
assert span["data"][SPANDATA.AI_TOTAL_TOKENS_USED] == 10
138138

139139

140140
def test_bad_chat_completion(sentry_init, capture_events):

‎tests/integrations/langchain/test_langchain.py

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -181,13 +181,13 @@ def test_langchain_agent(
181181
assert len(list(x for x in tx["spans"] if x["op"] == "ai.run.langchain")) > 0
182182

183183
if use_unknown_llm_type:
184-
assert "ai.prompt_tokens.used" in chat_spans[0]["data"]
185-
assert "ai.total_tokens.used" in chat_spans[0]["data"]
184+
assert SPANDATA.AI_PROMPT_TOKENS_USED in chat_spans[0]["data"]
185+
assert SPANDATA.AI_TOTAL_TOKENS_USED in chat_spans[0]["data"]
186186
else:
187187
# important: to avoid double counting, we do *not* measure
188188
# tokens used if we have an explicit integration (e.g. OpenAI)
189-
assert "ai.prompt_tokens.used" not in chat_spans[0]["data"]
190-
assert "ai.total_tokens.used" not in chat_spans[0]["data"]
189+
assert SPANDATA.AI_PROMPT_TOKENS_USED not in chat_spans[0]["data"]
190+
assert SPANDATA.AI_TOTAL_TOKENS_USED not in chat_spans[0]["data"]
191191

192192
if send_default_pii and include_prompts:
193193
assert (

‎tests/integrations/openai/test_openai.py

Lines changed: 16 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -90,9 +90,9 @@ def test_nonstreaming_chat_completion(
9090
assert SPANDATA.AI_INPUT_MESSAGES not in span["data"]
9191
assert SPANDATA.AI_RESPONSES not in span["data"]
9292

93-
assert span["data"]["ai.completion_tokens.used"] == 10
94-
assert span["data"]["ai.prompt_tokens.used"] == 20
95-
assert span["data"]["ai.total_tokens.used"] == 30
93+
assert span["data"][SPANDATA.AI_COMPLETION_TOKENS_USED] == 10
94+
assert span["data"][SPANDATA.AI_PROMPT_TOKENS_USED] == 20
95+
assert span["data"][SPANDATA.AI_TOTAL_TOKENS_USED] == 30
9696

9797

9898
@pytest.mark.asyncio
@@ -132,9 +132,9 @@ async def test_nonstreaming_chat_completion_async(
132132
assert SPANDATA.AI_INPUT_MESSAGES not in span["data"]
133133
assert SPANDATA.AI_RESPONSES not in span["data"]
134134

135-
assert span["data"]["ai.completion_tokens.used"] == 10
136-
assert span["data"]["ai.prompt_tokens.used"] == 20
137-
assert span["data"]["ai.total_tokens.used"] == 30
135+
assert span["data"][SPANDATA.AI_COMPLETION_TOKENS_USED] == 10
136+
assert span["data"][SPANDATA.AI_PROMPT_TOKENS_USED] == 20
137+
assert span["data"][SPANDATA.AI_TOTAL_TOKENS_USED] == 30
138138

139139

140140
def tiktoken_encoding_if_installed():
@@ -228,9 +228,9 @@ def test_streaming_chat_completion(
228228
try:
229229
import tiktoken # type: ignore # noqa # pylint: disable=unused-import
230230

231-
assert span["data"]["ai.completion_tokens.used"] == 2
232-
assert span["data"]["ai.prompt_tokens.used"] == 1
233-
assert span["data"]["ai.total_tokens.used"] == 3
231+
assert span["data"][SPANDATA.AI_COMPLETION_TOKENS_USED] == 2
232+
assert span["data"][SPANDATA.AI_PROMPT_TOKENS_USED] == 1
233+
assert span["data"][SPANDATA.AI_TOTAL_TOKENS_USED] == 3
234234
except ImportError:
235235
pass # if tiktoken is not installed, we can't guarantee token usage will be calculated properly
236236

@@ -324,9 +324,9 @@ async def test_streaming_chat_completion_async(
324324
try:
325325
import tiktoken # type: ignore # noqa # pylint: disable=unused-import
326326

327-
assert span["data"]["ai.completion_tokens.used"] == 2
328-
assert span["data"]["ai.prompt_tokens.used"] == 1
329-
assert span["data"]["ai.total_tokens.used"] == 3
327+
assert span["data"][SPANDATA.AI_COMPLETION_TOKENS_USED] == 2
328+
assert span["data"][SPANDATA.AI_PROMPT_TOKENS_USED] == 1
329+
assert span["data"][SPANDATA.AI_TOTAL_TOKENS_USED] == 3
330330
except ImportError:
331331
pass # if tiktoken is not installed, we can't guarantee token usage will be calculated properly
332332

@@ -410,8 +410,8 @@ def test_embeddings_create(
410410
else:
411411
assert SPANDATA.AI_INPUT_MESSAGES not in span["data"]
412412

413-
assert span["data"]["ai.prompt_tokens.used"] == 20
414-
assert span["data"]["ai.total_tokens.used"] == 30
413+
assert span["data"][SPANDATA.AI_PROMPT_TOKENS_USED] == 20
414+
assert span["data"][SPANDATA.AI_TOTAL_TOKENS_USED] == 30
415415

416416

417417
@pytest.mark.asyncio
@@ -458,8 +458,8 @@ async def test_embeddings_create_async(
458458
else:
459459
assert SPANDATA.AI_INPUT_MESSAGES not in span["data"]
460460

461-
assert span["data"]["ai.prompt_tokens.used"] == 20
462-
assert span["data"]["ai.total_tokens.used"] == 30
461+
assert span["data"][SPANDATA.AI_PROMPT_TOKENS_USED] == 20
462+
assert span["data"][SPANDATA.AI_TOTAL_TOKENS_USED] == 30
463463

464464

465465
@pytest.mark.forked

0 commit comments

Comments
 (0)