|
1 |
| -import pytest |
2 | 1 | from unittest import mock
|
3 |
| -from anthropic import Anthropic, Stream, AnthropicError |
4 |
| -from anthropic.types import Usage, MessageDeltaUsage, TextDelta |
| 2 | + |
| 3 | +import pytest |
| 4 | +from anthropic import Anthropic, AnthropicError, Stream |
| 5 | +from anthropic.types import MessageDeltaUsage, TextDelta, Usage |
| 6 | +from anthropic.types.content_block_delta_event import ContentBlockDeltaEvent |
| 7 | +from anthropic.types.content_block_start_event import ContentBlockStartEvent |
| 8 | +from anthropic.types.content_block_stop_event import ContentBlockStopEvent |
5 | 9 | from anthropic.types.message import Message
|
6 | 10 | from anthropic.types.message_delta_event import MessageDeltaEvent
|
7 | 11 | from anthropic.types.message_start_event import MessageStartEvent
|
8 |
| -from anthropic.types.content_block_start_event import ContentBlockStartEvent |
9 |
| -from anthropic.types.content_block_delta_event import ContentBlockDeltaEvent |
10 |
| -from anthropic.types.content_block_stop_event import ContentBlockStopEvent |
| 12 | + |
| 13 | +from sentry_sdk.utils import package_version |
| 14 | + |
| 15 | +try: |
| 16 | + from anthropic.types import InputJSONDelta |
| 17 | +except ImportError: |
| 18 | + try: |
| 19 | + from anthropic.types import InputJsonDelta as InputJSONDelta |
| 20 | + except ImportError: |
| 21 | + pass |
11 | 22 |
|
12 | 23 | try:
|
13 | 24 | # 0.27+
|
14 | 25 | from anthropic.types.raw_message_delta_event import Delta
|
| 26 | + from anthropic.types.tool_use_block import ToolUseBlock |
15 | 27 | except ImportError:
|
16 | 28 | # pre 0.27
|
17 | 29 | from anthropic.types.message_delta_event import Delta
|
|
25 | 37 | from sentry_sdk.consts import OP, SPANDATA
|
26 | 38 | from sentry_sdk.integrations.anthropic import AnthropicIntegration
|
27 | 39 |
|
28 |
| - |
| 40 | +ANTHROPIC_VERSION = package_version("anthropic") |
29 | 41 | EXAMPLE_MESSAGE = Message(
|
30 | 42 | id="id",
|
31 | 43 | model="model",
|
@@ -203,6 +215,136 @@ def test_streaming_create_message(
|
203 | 215 | assert span["data"]["ai.streaming"] is True
|
204 | 216 |
|
205 | 217 |
|
| 218 | +@pytest.mark.skipif( |
| 219 | + ANTHROPIC_VERSION < (0, 27), |
| 220 | + reason="Versions <0.27.0 do not include InputJSONDelta, which was introduced in >=0.27.0 along with a new message delta type for tool calling.", |
| 221 | +) |
| 222 | +@pytest.mark.parametrize( |
| 223 | + "send_default_pii, include_prompts", |
| 224 | + [ |
| 225 | + (True, True), |
| 226 | + (True, False), |
| 227 | + (False, True), |
| 228 | + (False, False), |
| 229 | + ], |
| 230 | +) |
| 231 | +def test_streaming_create_message_with_input_json_delta( |
| 232 | + sentry_init, capture_events, send_default_pii, include_prompts |
| 233 | +): |
| 234 | + client = Anthropic(api_key="z") |
| 235 | + returned_stream = Stream(cast_to=None, response=None, client=client) |
| 236 | + returned_stream._iterator = [ |
| 237 | + MessageStartEvent( |
| 238 | + message=Message( |
| 239 | + id="msg_0", |
| 240 | + content=[], |
| 241 | + model="claude-3-5-sonnet-20240620", |
| 242 | + role="assistant", |
| 243 | + stop_reason=None, |
| 244 | + stop_sequence=None, |
| 245 | + type="message", |
| 246 | + usage=Usage(input_tokens=366, output_tokens=10), |
| 247 | + ), |
| 248 | + type="message_start", |
| 249 | + ), |
| 250 | + ContentBlockStartEvent( |
| 251 | + type="content_block_start", |
| 252 | + index=0, |
| 253 | + content_block=ToolUseBlock( |
| 254 | + id="toolu_0", input={}, name="get_weather", type="tool_use" |
| 255 | + ), |
| 256 | + ), |
| 257 | + ContentBlockDeltaEvent( |
| 258 | + delta=InputJSONDelta(partial_json="", type="input_json_delta"), |
| 259 | + index=0, |
| 260 | + type="content_block_delta", |
| 261 | + ), |
| 262 | + ContentBlockDeltaEvent( |
| 263 | + delta=InputJSONDelta(partial_json="{'location':", type="input_json_delta"), |
| 264 | + index=0, |
| 265 | + type="content_block_delta", |
| 266 | + ), |
| 267 | + ContentBlockDeltaEvent( |
| 268 | + delta=InputJSONDelta(partial_json=" 'S", type="input_json_delta"), |
| 269 | + index=0, |
| 270 | + type="content_block_delta", |
| 271 | + ), |
| 272 | + ContentBlockDeltaEvent( |
| 273 | + delta=InputJSONDelta(partial_json="an ", type="input_json_delta"), |
| 274 | + index=0, |
| 275 | + type="content_block_delta", |
| 276 | + ), |
| 277 | + ContentBlockDeltaEvent( |
| 278 | + delta=InputJSONDelta(partial_json="Francisco, C", type="input_json_delta"), |
| 279 | + index=0, |
| 280 | + type="content_block_delta", |
| 281 | + ), |
| 282 | + ContentBlockDeltaEvent( |
| 283 | + delta=InputJSONDelta(partial_json="A'}", type="input_json_delta"), |
| 284 | + index=0, |
| 285 | + type="content_block_delta", |
| 286 | + ), |
| 287 | + ContentBlockStopEvent(type="content_block_stop", index=0), |
| 288 | + MessageDeltaEvent( |
| 289 | + delta=Delta(stop_reason="tool_use", stop_sequence=None), |
| 290 | + usage=MessageDeltaUsage(output_tokens=41), |
| 291 | + type="message_delta", |
| 292 | + ), |
| 293 | + ] |
| 294 | + |
| 295 | + sentry_init( |
| 296 | + integrations=[AnthropicIntegration(include_prompts=include_prompts)], |
| 297 | + traces_sample_rate=1.0, |
| 298 | + send_default_pii=send_default_pii, |
| 299 | + ) |
| 300 | + events = capture_events() |
| 301 | + client.messages._post = mock.Mock(return_value=returned_stream) |
| 302 | + |
| 303 | + messages = [ |
| 304 | + { |
| 305 | + "role": "user", |
| 306 | + "content": "What is the weather like in San Francisco?", |
| 307 | + } |
| 308 | + ] |
| 309 | + |
| 310 | + with start_transaction(name="anthropic"): |
| 311 | + message = client.messages.create( |
| 312 | + max_tokens=1024, messages=messages, model="model", stream=True |
| 313 | + ) |
| 314 | + |
| 315 | + for _ in message: |
| 316 | + pass |
| 317 | + |
| 318 | + assert message == returned_stream |
| 319 | + assert len(events) == 1 |
| 320 | + (event,) = events |
| 321 | + |
| 322 | + assert event["type"] == "transaction" |
| 323 | + assert event["transaction"] == "anthropic" |
| 324 | + |
| 325 | + assert len(event["spans"]) == 1 |
| 326 | + (span,) = event["spans"] |
| 327 | + |
| 328 | + assert span["op"] == OP.ANTHROPIC_MESSAGES_CREATE |
| 329 | + assert span["description"] == "Anthropic messages create" |
| 330 | + assert span["data"][SPANDATA.AI_MODEL_ID] == "model" |
| 331 | + |
| 332 | + if send_default_pii and include_prompts: |
| 333 | + assert span["data"][SPANDATA.AI_INPUT_MESSAGES] == messages |
| 334 | + assert span["data"][SPANDATA.AI_RESPONSES] == [ |
| 335 | + {"text": "", "type": "text"} |
| 336 | + ] # we do not record InputJSONDelta because it could contain PII |
| 337 | + |
| 338 | + else: |
| 339 | + assert SPANDATA.AI_INPUT_MESSAGES not in span["data"] |
| 340 | + assert SPANDATA.AI_RESPONSES not in span["data"] |
| 341 | + |
| 342 | + assert span["measurements"]["ai_prompt_tokens_used"]["value"] == 366 |
| 343 | + assert span["measurements"]["ai_completion_tokens_used"]["value"] == 51 |
| 344 | + assert span["measurements"]["ai_total_tokens_used"]["value"] == 417 |
| 345 | + assert span["data"]["ai.streaming"] is True |
| 346 | + |
| 347 | + |
206 | 348 | def test_exception_message_create(sentry_init, capture_events):
|
207 | 349 | sentry_init(integrations=[AnthropicIntegration()], traces_sample_rate=1.0)
|
208 | 350 | events = capture_events()
|
|
0 commit comments