Skip to content

Commit 115af46

Browse files
constantiniusconstantinius
committed
test(anthropic) update tests to use GEN_AI attributes
1 parent 1af3203 commit 115af46

File tree

1 file changed

+58
-58
lines changed

1 file changed

+58
-58
lines changed

‎tests/integrations/anthropic/test_anthropic.py‎

Lines changed: 58 additions & 58 deletions
Original file line numberDiff line numberDiff line change
@@ -114,21 +114,21 @@ def test_nonstreaming_create_message(
114114

115115
assert span["op"] == OP.ANTHROPIC_MESSAGES_CREATE
116116
assert span["description"] == "Anthropic messages create"
117-
assert span["data"][SPANDATA.AI_MODEL_ID] == "model"
117+
assert span["data"][SPANDATA.GEN_AI_REQUEST_MODEL] == "model"
118118

119119
if send_default_pii and include_prompts:
120-
assert span["data"][SPANDATA.AI_INPUT_MESSAGES] == messages
121-
assert span["data"][SPANDATA.AI_RESPONSES] == [
120+
assert span["data"][SPANDATA.GEN_AI_REQUEST_MESSAGES] == messages
121+
assert span["data"][SPANDATA.GEN_AI_RESPONSE_TEXT] == [
122122
{"type": "text", "text": "Hi, I'm Claude."}
123123
]
124124
else:
125-
assert SPANDATA.AI_INPUT_MESSAGES not in span["data"]
126-
assert SPANDATA.AI_RESPONSES not in span["data"]
125+
assert SPANDATA.GEN_AI_REQUEST_MESSAGES not in span["data"]
126+
assert SPANDATA.GEN_AI_RESPONSE_TEXT not in span["data"]
127127

128-
assert span["data"]["gen_ai.usage.input_tokens"] == 10
129-
assert span["data"]["gen_ai.usage.output_tokens"] == 20
130-
assert span["data"]["gen_ai.usage.total_tokens"] == 30
131-
assert span["data"][SPANDATA.AI_STREAMING] is False
128+
assert span["data"][SPANDATA.GEN_AI_USAGE_INPUT_TOKENS] == 10
129+
assert span["data"][SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS] == 20
130+
assert span["data"][SPANDATA.GEN_AI_USAGE_TOTAL_TOKENS] == 30
131+
assert span["data"][SPANDATA.GEN_AI_RESPONSE_STREAMING] is False
132132

133133

134134
@pytest.mark.asyncio
@@ -182,21 +182,21 @@ async def test_nonstreaming_create_message_async(
182182

183183
assert span["op"] == OP.ANTHROPIC_MESSAGES_CREATE
184184
assert span["description"] == "Anthropic messages create"
185-
assert span["data"][SPANDATA.AI_MODEL_ID] == "model"
185+
assert span["data"][SPANDATA.GEN_AI_REQUEST_MODEL] == "model"
186186

187187
if send_default_pii and include_prompts:
188-
assert span["data"][SPANDATA.AI_INPUT_MESSAGES] == messages
189-
assert span["data"][SPANDATA.AI_RESPONSES] == [
188+
assert span["data"][SPANDATA.GEN_AI_REQUEST_MESSAGES] == messages
189+
assert span["data"][SPANDATA.GEN_AI_RESPONSE_TEXT] == [
190190
{"type": "text", "text": "Hi, I'm Claude."}
191191
]
192192
else:
193-
assert SPANDATA.AI_INPUT_MESSAGES not in span["data"]
194-
assert SPANDATA.AI_RESPONSES not in span["data"]
193+
assert SPANDATA.GEN_AI_REQUEST_MESSAGES not in span["data"]
194+
assert SPANDATA.GEN_AI_RESPONSE_TEXT not in span["data"]
195195

196-
assert span["data"]["gen_ai.usage.input_tokens"] == 10
197-
assert span["data"]["gen_ai.usage.output_tokens"] == 20
198-
assert span["data"]["gen_ai.usage.total_tokens"] == 30
199-
assert span["data"][SPANDATA.AI_STREAMING] is False
196+
assert span["data"][SPANDATA.GEN_AI_USAGE_INPUT_TOKENS] == 10
197+
assert span["data"][SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS] == 20
198+
assert span["data"][SPANDATA.GEN_AI_USAGE_TOTAL_TOKENS] == 30
199+
assert span["data"][SPANDATA.GEN_AI_RESPONSE_STREAMING] is False
200200

201201

202202
@pytest.mark.parametrize(
@@ -281,22 +281,22 @@ def test_streaming_create_message(
281281

282282
assert span["op"] == OP.ANTHROPIC_MESSAGES_CREATE
283283
assert span["description"] == "Anthropic messages create"
284-
assert span["data"][SPANDATA.AI_MODEL_ID] == "model"
284+
assert span["data"][SPANDATA.GEN_AI_REQUEST_MODEL] == "model"
285285

286286
if send_default_pii and include_prompts:
287-
assert span["data"][SPANDATA.AI_INPUT_MESSAGES] == messages
288-
assert span["data"][SPANDATA.AI_RESPONSES] == [
287+
assert span["data"][SPANDATA.GEN_AI_REQUEST_MESSAGES] == messages
288+
assert span["data"][SPANDATA.GEN_AI_RESPONSE_TEXT] == [
289289
{"type": "text", "text": "Hi! I'm Claude!"}
290290
]
291291

292292
else:
293-
assert SPANDATA.AI_INPUT_MESSAGES not in span["data"]
294-
assert SPANDATA.AI_RESPONSES not in span["data"]
293+
assert SPANDATA.GEN_AI_REQUEST_MESSAGES not in span["data"]
294+
assert SPANDATA.GEN_AI_RESPONSE_TEXT not in span["data"]
295295

296-
assert span["data"]["gen_ai.usage.input_tokens"] == 10
297-
assert span["data"]["gen_ai.usage.output_tokens"] == 30
298-
assert span["data"]["gen_ai.usage.total_tokens"] == 40
299-
assert span["data"][SPANDATA.AI_STREAMING] is True
296+
assert span["data"][SPANDATA.GEN_AI_USAGE_INPUT_TOKENS] == 10
297+
assert span["data"][SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS] == 30
298+
assert span["data"][SPANDATA.GEN_AI_USAGE_TOTAL_TOKENS] == 40
299+
assert span["data"][SPANDATA.GEN_AI_RESPONSE_STREAMING] is True
300300

301301

302302
@pytest.mark.asyncio
@@ -384,22 +384,22 @@ async def test_streaming_create_message_async(
384384

385385
assert span["op"] == OP.ANTHROPIC_MESSAGES_CREATE
386386
assert span["description"] == "Anthropic messages create"
387-
assert span["data"][SPANDATA.AI_MODEL_ID] == "model"
387+
assert span["data"][SPANDATA.GEN_AI_REQUEST_MODEL] == "model"
388388

389389
if send_default_pii and include_prompts:
390-
assert span["data"][SPANDATA.AI_INPUT_MESSAGES] == messages
391-
assert span["data"][SPANDATA.AI_RESPONSES] == [
390+
assert span["data"][SPANDATA.GEN_AI_REQUEST_MESSAGES] == messages
391+
assert span["data"][SPANDATA.GEN_AI_RESPONSE_TEXT] == [
392392
{"type": "text", "text": "Hi! I'm Claude!"}
393393
]
394394

395395
else:
396-
assert SPANDATA.AI_INPUT_MESSAGES not in span["data"]
397-
assert SPANDATA.AI_RESPONSES not in span["data"]
396+
assert SPANDATA.GEN_AI_REQUEST_MESSAGES not in span["data"]
397+
assert SPANDATA.GEN_AI_RESPONSE_TEXT not in span["data"]
398398

399-
assert span["data"]["gen_ai.usage.input_tokens"] == 10
400-
assert span["data"]["gen_ai.usage.output_tokens"] == 30
401-
assert span["data"]["gen_ai.usage.total_tokens"] == 40
402-
assert span["data"][SPANDATA.AI_STREAMING] is True
399+
assert span["data"][SPANDATA.GEN_AI_USAGE_INPUT_TOKENS] == 10
400+
assert span["data"][SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS] == 30
401+
assert span["data"][SPANDATA.GEN_AI_USAGE_TOTAL_TOKENS] == 40
402+
assert span["data"][SPANDATA.GEN_AI_RESPONSE_STREAMING] is True
403403

404404

405405
@pytest.mark.skipif(
@@ -514,21 +514,21 @@ def test_streaming_create_message_with_input_json_delta(
514514

515515
assert span["op"] == OP.ANTHROPIC_MESSAGES_CREATE
516516
assert span["description"] == "Anthropic messages create"
517-
assert span["data"][SPANDATA.AI_MODEL_ID] == "model"
517+
assert span["data"][SPANDATA.GEN_AI_REQUEST_MODEL] == "model"
518518

519519
if send_default_pii and include_prompts:
520-
assert span["data"][SPANDATA.AI_INPUT_MESSAGES] == messages
521-
assert span["data"][SPANDATA.AI_RESPONSES] == [
520+
assert span["data"][SPANDATA.GEN_AI_REQUEST_MESSAGES] == messages
521+
assert span["data"][SPANDATA.GEN_AI_RESPONSE_TEXT] == [
522522
{"text": "{'location': 'San Francisco, CA'}", "type": "text"}
523523
]
524524
else:
525-
assert SPANDATA.AI_INPUT_MESSAGES not in span["data"]
526-
assert SPANDATA.AI_RESPONSES not in span["data"]
525+
assert SPANDATA.GEN_AI_REQUEST_MESSAGES not in span["data"]
526+
assert SPANDATA.GEN_AI_RESPONSE_TEXT not in span["data"]
527527

528-
assert span["data"]["gen_ai.usage.input_tokens"] == 366
529-
assert span["data"]["gen_ai.usage.output_tokens"] == 51
530-
assert span["data"]["gen_ai.usage.total_tokens"] == 417
531-
assert span["data"][SPANDATA.AI_STREAMING] is True
528+
assert span["data"][SPANDATA.GEN_AI_USAGE_INPUT_TOKENS] == 366
529+
assert span["data"][SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS] == 51
530+
assert span["data"][SPANDATA.GEN_AI_USAGE_TOTAL_TOKENS] == 417
531+
assert span["data"][SPANDATA.GEN_AI_RESPONSE_STREAMING] is True
532532

533533

534534
@pytest.mark.asyncio
@@ -650,22 +650,22 @@ async def test_streaming_create_message_with_input_json_delta_async(
650650

651651
assert span["op"] == OP.ANTHROPIC_MESSAGES_CREATE
652652
assert span["description"] == "Anthropic messages create"
653-
assert span["data"][SPANDATA.AI_MODEL_ID] == "model"
653+
assert span["data"][SPANDATA.GEN_AI_REQUEST_MODEL] == "model"
654654

655655
if send_default_pii and include_prompts:
656-
assert span["data"][SPANDATA.AI_INPUT_MESSAGES] == messages
657-
assert span["data"][SPANDATA.AI_RESPONSES] == [
656+
assert span["data"][SPANDATA.GEN_AI_REQUEST_MESSAGES] == messages
657+
assert span["data"][SPANDATA.GEN_AI_RESPONSE_TEXT] == [
658658
{"text": "{'location': 'San Francisco, CA'}", "type": "text"}
659659
]
660660

661661
else:
662-
assert SPANDATA.AI_INPUT_MESSAGES not in span["data"]
663-
assert SPANDATA.AI_RESPONSES not in span["data"]
662+
assert SPANDATA.GEN_AI_REQUEST_MESSAGES not in span["data"]
663+
assert SPANDATA.GEN_AI_RESPONSE_TEXT not in span["data"]
664664

665-
assert span["data"]["gen_ai.usage.input_tokens"] == 366
666-
assert span["data"]["gen_ai.usage.output_tokens"] == 51
667-
assert span["data"]["gen_ai.usage.total_tokens"] == 417
668-
assert span["data"][SPANDATA.AI_STREAMING] is True
665+
assert span["data"][SPANDATA.GEN_AI_USAGE_INPUT_TOKENS] == 366
666+
assert span["data"][SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS] == 51
667+
assert span["data"][SPANDATA.GEN_AI_USAGE_TOTAL_TOKENS] == 417
668+
assert span["data"][SPANDATA.GEN_AI_RESPONSE_STREAMING] is True
669669

670670

671671
def test_exception_message_create(sentry_init, capture_events):
@@ -810,7 +810,7 @@ def test_add_ai_data_to_span_with_input_json_delta(sentry_init):
810810
assert span._data.get("ai.responses") == [
811811
{"type": "text", "text": "{'test': 'data','more': 'json'}"}
812812
]
813-
assert span._data.get("ai.streaming") is True
814-
assert span._data.get("gen_ai.usage.input_tokens") == 10
815-
assert span._data.get("gen_ai.usage.output_tokens") == 20
816-
assert span._data.get("gen_ai.usage.total_tokens") == 30
813+
assert span._data.get(SPANDATA.GEN_AI_RESPONSE_STREAMING) is True
814+
assert span._data.get(SPANDATA.GEN_AI_USAGE_INPUT_TOKENS) == 10
815+
assert span._data.get(SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS) == 20
816+
assert span._data.get(SPANDATA.GEN_AI_USAGE_TOTAL_TOKENS) == 30

0 commit comments

Comments
 (0)