|
5 | 5 | from cohere import Client, ChatMessage |
6 | 6 |
|
7 | 7 | from sentry_sdk import start_span |
| 8 | +from sentry_sdk.consts import SPANDATA |
8 | 9 | from sentry_sdk.integrations.cohere import CohereIntegration |
9 | 10 |
|
10 | 11 | from unittest import mock # python 3.3 and above |
@@ -53,16 +54,16 @@ def test_nonstreaming_chat( |
53 | 54 | assert tx["type"] == "transaction" |
54 | 55 | span = tx["spans"][0] |
55 | 56 | assert span["op"] == "ai.chat_completions.create.cohere" |
56 | | - assert span["data"]["ai.model_id"] == "some-model" |
| 57 | + assert span["data"][SPANDATA.AI_MODEL_ID] == "some-model" |
57 | 58 |
|
58 | 59 | if send_default_pii and include_prompts: |
59 | | - input_messages = json.loads(span["data"]["ai.input_messages"]) |
| 60 | + input_messages = json.loads(span["data"][SPANDATA.AI_INPUT_MESSAGES]) |
60 | 61 | assert "some context" in input_messages[0]["content"] |
61 | 62 | assert "hello" in input_messages[1]["content"] |
62 | | - assert "the model response" in span["data"]["ai.responses"] |
| 63 | + assert "the model response" in span["data"][SPANDATA.AI_RESPONSES] |
63 | 64 | else: |
64 | | - assert "ai.input_messages" not in span["data"] |
65 | | - assert "ai.responses" not in span["data"] |
| 65 | + assert SPANDATA.AI_INPUT_MESSAGES not in span["data"] |
| 66 | + assert SPANDATA.AI_RESPONSES not in span["data"] |
66 | 67 |
|
67 | 68 | assert span["data"]["ai.completion_tokens.used"] == 10 |
68 | 69 | assert span["data"]["ai.prompt_tokens.used"] == 20 |
@@ -125,16 +126,16 @@ def test_streaming_chat(sentry_init, capture_events, send_default_pii, include_p |
125 | 126 | assert tx["type"] == "transaction" |
126 | 127 | span = tx["spans"][0] |
127 | 128 | assert span["op"] == "ai.chat_completions.create.cohere" |
128 | | - assert span["data"]["ai.model_id"] == "some-model" |
| 129 | + assert span["data"][SPANDATA.AI_MODEL_ID] == "some-model" |
129 | 130 |
|
130 | 131 | if send_default_pii and include_prompts: |
131 | | - input_messages = json.loads(span["data"]["ai.input_messages"]) |
| 132 | + input_messages = json.loads(span["data"][SPANDATA.AI_INPUT_MESSAGES]) |
132 | 133 | assert "some context" in input_messages[0]["content"] |
133 | 134 | assert "hello" in input_messages[1]["content"] |
134 | | - assert "the model response" in span["data"]["ai.responses"] |
| 135 | + assert "the model response" in span["data"][SPANDATA.AI_RESPONSES] |
135 | 136 | else: |
136 | | - assert "ai.input_messages" not in span["data"] |
137 | | - assert "ai.responses" not in span["data"] |
| 137 | + assert SPANDATA.AI_INPUT_MESSAGES not in span["data"] |
| 138 | + assert SPANDATA.AI_RESPONSES not in span["data"] |
138 | 139 |
|
139 | 140 | assert span["data"]["ai.completion_tokens.used"] == 10 |
140 | 141 | assert span["data"]["ai.prompt_tokens.used"] == 20 |
@@ -196,9 +197,9 @@ def test_embed(sentry_init, capture_events, send_default_pii, include_prompts): |
196 | 197 | span = tx["spans"][0] |
197 | 198 | assert span["op"] == "ai.embeddings.create.cohere" |
198 | 199 | if send_default_pii and include_prompts: |
199 | | - assert "hello" in span["data"]["ai.input_messages"] |
| 200 | + assert "hello" in span["data"][SPANDATA.AI_INPUT_MESSAGES] |
200 | 201 | else: |
201 | | - assert "ai.input_messages" not in span["data"] |
| 202 | + assert SPANDATA.AI_INPUT_MESSAGES not in span["data"] |
202 | 203 |
|
203 | 204 | assert span["data"]["ai.prompt_tokens.used"] == 10 |
204 | 205 | assert span["data"]["ai.total_tokens.used"] == 10 |
|
0 commit comments