Skip to content

Commit 03659a4

Browse files
antonpirkerantonpirker
committed
Merge branch 'master' into potel-base
2 parents 0f9f51c + 6f71a1b commit 03659a4

File tree

6 files changed

+67
-63
lines changed

6 files changed

+67
-63
lines changed

‎sentry_sdk/ai/monitoring.py‎

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -103,15 +103,19 @@ def record_token_usage(
103103
ai_pipeline_name = get_ai_pipeline_name()
104104
if ai_pipeline_name:
105105
span.set_attribute(SPANDATA.AI_PIPELINE_NAME, ai_pipeline_name)
106+
106107
if prompt_tokens is not None:
107-
span.set_attribute(SPANDATA.AI_PROMPT_TOKENS_USED, prompt_tokens)
108+
span.set_attribute(SPANDATA.GEN_AI_USAGE_INPUT_TOKENS, prompt_tokens)
109+
108110
if completion_tokens is not None:
109-
span.set_attribute(SPANDATA.AI_COMPLETION_TOKENS_USED, completion_tokens)
111+
span.set_attribute(SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS, completion_tokens)
112+
110113
if (
111114
total_tokens is None
112115
and prompt_tokens is not None
113116
and completion_tokens is not None
114117
):
115118
total_tokens = prompt_tokens + completion_tokens
119+
116120
if total_tokens is not None:
117-
span.set_attribute(SPANDATA.AI_TOTAL_TOKENS_USED, total_tokens)
121+
span.set_attribute(SPANDATA.GEN_AI_USAGE_TOTAL_TOKENS, total_tokens)

‎tests/integrations/anthropic/test_anthropic.py‎

Lines changed: 29 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -127,10 +127,10 @@ 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"][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
133-
assert span["data"][SPANDATA.AI_STREAMING] is False
130+
assert span["data"]["gen_ai.usage.input_tokens"] == 10
131+
assert span["data"]["gen_ai.usage.output_tokens"] == 20
132+
assert span["data"]["gen_ai.usage.total_tokens"] == 30
133+
assert span["data"]["ai.streaming"] is False
134134

135135

136136
@pytest.mark.asyncio
@@ -197,10 +197,10 @@ 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"][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
203-
assert span["data"][SPANDATA.AI_STREAMING] is False
200+
assert span["data"]["gen_ai.usage.input_tokens"] == 10
201+
assert span["data"]["gen_ai.usage.output_tokens"] == 20
202+
assert span["data"]["gen_ai.usage.total_tokens"] == 30
203+
assert span["data"]["ai.streaming"] is False
204204

205205

206206
@pytest.mark.parametrize(
@@ -299,10 +299,10 @@ 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"][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
305-
assert span["data"][SPANDATA.AI_STREAMING] is True
302+
assert span["data"]["gen_ai.usage.input_tokens"] == 10
303+
assert span["data"]["gen_ai.usage.output_tokens"] == 30
304+
assert span["data"]["gen_ai.usage.total_tokens"] == 40
305+
assert span["data"]["ai.streaming"] is True
306306

307307

308308
@pytest.mark.asyncio
@@ -404,10 +404,10 @@ 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"][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
410-
assert span["data"][SPANDATA.AI_STREAMING] is True
407+
assert span["data"]["gen_ai.usage.input_tokens"] == 10
408+
assert span["data"]["gen_ai.usage.output_tokens"] == 30
409+
assert span["data"]["gen_ai.usage.total_tokens"] == 40
410+
assert span["data"]["ai.streaming"] is True
411411

412412

413413
@pytest.mark.skipif(
@@ -536,10 +536,10 @@ 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"][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
542-
assert span["data"][SPANDATA.AI_STREAMING] is True
539+
assert span["data"]["gen_ai.usage.input_tokens"] == 366
540+
assert span["data"]["gen_ai.usage.output_tokens"] == 51
541+
assert span["data"]["gen_ai.usage.total_tokens"] == 417
542+
assert span["data"]["ai.streaming"] is True
543543

544544

545545
@pytest.mark.asyncio
@@ -675,10 +675,10 @@ 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"][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
681-
assert span["data"][SPANDATA.AI_STREAMING] is True
678+
assert span["data"]["gen_ai.usage.input_tokens"] == 366
679+
assert span["data"]["gen_ai.usage.output_tokens"] == 51
680+
assert span["data"]["gen_ai.usage.total_tokens"] == 417
681+
assert span["data"]["ai.streaming"] is True
682682

683683

684684
@pytest.mark.forked
@@ -827,10 +827,10 @@ def test_add_ai_data_to_span_with_input_json_delta(sentry_init, capture_events):
827827
assert len(event["spans"]) == 1
828828
(span,) = event["spans"]
829829

830-
assert span["data"][SPANDATA.AI_RESPONSES] == _serialize_span_attribute(
830+
assert span["data"]["ai.responses"] == _serialize_span_attribute(
831831
[{"type": "text", "text": "{'test': 'data','more': 'json'}"}]
832832
)
833-
assert span["data"][SPANDATA.AI_STREAMING] is True
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
833+
assert span["data"]["ai.streaming"] is True
834+
assert span["data"]["gen_ai.usage.input_tokens"] == 10
835+
assert span["data"]["gen_ai.usage.output_tokens"] == 20
836+
assert span["data"]["gen_ai.usage.total_tokens"] == 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"][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
68+
assert span["data"]["gen_ai.usage.output_tokens"] == 10
69+
assert span["data"]["gen_ai.usage.input_tokens"] == 20
70+
assert span["data"]["gen_ai.usage.total_tokens"] == 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"][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
140+
assert span["data"]["gen_ai.usage.output_tokens"] == 10
141+
assert span["data"]["gen_ai.usage.input_tokens"] == 20
142+
assert span["data"]["gen_ai.usage.total_tokens"] == 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"][SPANDATA.AI_PROMPT_TOKENS_USED] == 10
205-
assert span["data"][SPANDATA.AI_TOTAL_TOKENS_USED] == 10
204+
assert span["data"]["gen_ai.usage.input_tokens"] == 10
205+
assert span["data"]["gen_ai.usage.total_tokens"] == 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"][SPANDATA.AI_TOTAL_TOKENS_USED] == 10
78+
assert span["data"]["gen_ai.usage.total_tokens"] == 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"][SPANDATA.AI_TOTAL_TOKENS_USED] == 10
137+
assert span["data"]["gen_ai.usage.total_tokens"] == 10
138138

139139

140140
def test_bad_chat_completion(sentry_init, capture_events):

‎tests/integrations/langchain/test_langchain.py‎

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

188188
if use_unknown_llm_type:
189-
assert SPANDATA.AI_PROMPT_TOKENS_USED in chat_spans[0]["data"]
190-
assert SPANDATA.AI_TOTAL_TOKENS_USED in chat_spans[0]["data"]
189+
assert "gen_ai.usage.input_tokens" in chat_spans[0]["data"]
190+
assert "gen_ai.usage.total_tokens" in chat_spans[0]["data"]
191191
else:
192-
# important: to avoid double counting, we do *not* measure
192+
# important: to avoid double counting, we do *not* count
193193
# tokens used if we have an explicit integration (e.g. OpenAI)
194-
assert SPANDATA.AI_PROMPT_TOKENS_USED not in chat_spans[0]["data"]
195-
assert SPANDATA.AI_TOTAL_TOKENS_USED not in chat_spans[0]["data"]
194+
assert "gen_ai.usage.input_tokens" not in chat_spans[0]["data"]
195+
assert "gen_ai.usage.total_tokens" not in chat_spans[0]["data"]
196196

197197
if send_default_pii and include_prompts:
198198
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"][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
93+
assert span["data"]["gen_ai.usage.output_tokens"] == 10
94+
assert span["data"]["gen_ai.usage.input_tokens"] == 20
95+
assert span["data"]["gen_ai.usage.total_tokens"] == 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"][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
135+
assert span["data"]["gen_ai.usage.output_tokens"] == 10
136+
assert span["data"]["gen_ai.usage.input_tokens"] == 20
137+
assert span["data"]["gen_ai.usage.total_tokens"] == 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"][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
231+
assert span["data"]["gen_ai.usage.output_tokens"] == 2
232+
assert span["data"]["gen_ai.usage.input_tokens"] == 1
233+
assert span["data"]["gen_ai.usage.total_tokens"] == 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"][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
327+
assert span["data"]["gen_ai.usage.output_tokens"] == 2
328+
assert span["data"]["gen_ai.usage.input_tokens"] == 1
329+
assert span["data"]["gen_ai.usage.total_tokens"] == 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"][SPANDATA.AI_PROMPT_TOKENS_USED] == 20
414-
assert span["data"][SPANDATA.AI_TOTAL_TOKENS_USED] == 30
413+
assert span["data"]["gen_ai.usage.input_tokens"] == 20
414+
assert span["data"]["gen_ai.usage.total_tokens"] == 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"][SPANDATA.AI_PROMPT_TOKENS_USED] == 20
462-
assert span["data"][SPANDATA.AI_TOTAL_TOKENS_USED] == 30
461+
assert span["data"]["gen_ai.usage.input_tokens"] == 20
462+
assert span["data"]["gen_ai.usage.total_tokens"] == 30
463463

464464

465465
@pytest.mark.forked

0 commit comments

Comments
 (0)