Merge branch 'master' into antonpirker/better-aws-tests

Author: antonpirker

sentry_sdk/integrations/anthropic.py

@@ -101,6 +101,8 @@ def _collect_ai_data(event, input_tokens, output_tokens, content_blocks):
             elif event.type == "content_block_delta":
                 if hasattr(event.delta, "text"):
                     content_blocks.append(event.delta.text)
+                elif hasattr(event.delta, "partial_json"):
+                    content_blocks.append(event.delta.partial_json)
             elif event.type == "content_block_stop":
                 pass
             elif event.type == "message_delta":

sentry_sdk/integrations/aws_lambda.py

@@ -61,7 +61,10 @@ def sentry_init_error(*args, **kwargs):
 
             else:
                 # Fall back to AWS lambdas JSON representation of the error
-                sentry_event = _event_from_error_json(json.loads(args[1]))
+                error_info = args[1]
+                if isinstance(error_info, str):
+                    error_info = json.loads(error_info)
+                sentry_event = _event_from_error_json(error_info)
                 sentry_sdk.capture_event(sentry_event)
 
         return init_error(*args, **kwargs)

sentry_sdk/profiler/__init__.py

@@ -1,4 +1,9 @@
-from sentry_sdk.profiler.continuous_profiler import start_profiler, stop_profiler
+from sentry_sdk.profiler.continuous_profiler import (
+    start_profile_session,
+    start_profiler,
+    stop_profile_session,
+    stop_profiler,
+)
 from sentry_sdk.profiler.transaction_profiler import (
     MAX_PROFILE_DURATION_NS,
     PROFILE_MINIMUM_SAMPLES,
@@ -20,8 +25,10 @@
 )
 
 __all__ = [
-    "start_profiler",
-    "stop_profiler",
+    "start_profile_session",
+    "start_profiler",  # TODO: Deprecate this in favor of `start_profile_session`
+    "stop_profile_session",
+    "stop_profiler",  # TODO: Deprecate this in favor of `stop_profile_session`
     # DEPRECATED: The following was re-exported for backwards compatibility. It
     # will be removed from sentry_sdk.profiler in a future release.
     "MAX_PROFILE_DURATION_NS",

sentry_sdk/scope.py

@@ -1568,7 +1568,7 @@ def update_from_kwargs(
         user=None,  # type: Optional[Any]
         level=None,  # type: Optional[LogLevelStr]
         extras=None,  # type: Optional[Dict[str, Any]]
-        contexts=None,  # type: Optional[Dict[str, Any]]
+        contexts=None,  # type: Optional[Dict[str, Dict[str, Any]]]
         tags=None,  # type: Optional[Dict[str, str]]
         fingerprint=None,  # type: Optional[List[str]]
     ):

tests/integrations/anthropic/test_anthropic.py

@@ -1,5 +1,6 @@
 from unittest import mock
 
+
 try:
     from unittest.mock import AsyncMock
 except ImportError:
@@ -10,7 +11,7 @@ async def __call__(self, *args, **kwargs):
 
 
 import pytest
-from anthropic import AsyncAnthropic, Anthropic, AnthropicError, AsyncStream, Stream
+from anthropic import Anthropic, AnthropicError, AsyncAnthropic, AsyncStream, Stream
 from anthropic.types import MessageDeltaUsage, TextDelta, Usage
 from anthropic.types.content_block_delta_event import ContentBlockDeltaEvent
 from anthropic.types.content_block_start_event import ContentBlockStartEvent
@@ -19,6 +20,7 @@ async def __call__(self, *args, **kwargs):
 from anthropic.types.message_delta_event import MessageDeltaEvent
 from anthropic.types.message_start_event import MessageStartEvent
 
+from sentry_sdk.integrations.anthropic import _add_ai_data_to_span, _collect_ai_data
 from sentry_sdk.utils import package_version
 
 try:
@@ -42,7 +44,7 @@ async def __call__(self, *args, **kwargs):
 except ImportError:
     from anthropic.types.content_block import ContentBlock as TextBlock
 
-from sentry_sdk import start_transaction
+from sentry_sdk import start_transaction, start_span
 from sentry_sdk.consts import OP, SPANDATA
 from sentry_sdk.integrations.anthropic import AnthropicIntegration
 
@@ -517,9 +519,8 @@ def test_streaming_create_message_with_input_json_delta(
     if send_default_pii and include_prompts:
         assert span["data"][SPANDATA.AI_INPUT_MESSAGES] == messages
         assert span["data"][SPANDATA.AI_RESPONSES] == [
-            {"text": "", "type": "text"}
-        ]  # we do not record InputJSONDelta because it could contain PII
-
+            {"text": "{'location': 'San Francisco, CA'}", "type": "text"}
+        ]
     else:
         assert SPANDATA.AI_INPUT_MESSAGES not in span["data"]
         assert SPANDATA.AI_RESPONSES not in span["data"]
@@ -654,8 +655,8 @@ async def test_streaming_create_message_with_input_json_delta_async(
     if send_default_pii and include_prompts:
         assert span["data"][SPANDATA.AI_INPUT_MESSAGES] == messages
         assert span["data"][SPANDATA.AI_RESPONSES] == [
-            {"text": "", "type": "text"}
-        ]  # we do not record InputJSONDelta because it could contain PII
+            {"text": "{'location': 'San Francisco, CA'}", "type": "text"}
+        ]
 
     else:
         assert SPANDATA.AI_INPUT_MESSAGES not in span["data"]
@@ -757,3 +758,59 @@ async def test_span_origin_async(sentry_init, capture_events):
 
     assert event["contexts"]["trace"]["origin"] == "manual"
     assert event["spans"][0]["origin"] == "auto.ai.anthropic"
+
+
+@pytest.mark.skipif(
+    ANTHROPIC_VERSION < (0, 27),
+    reason="Versions <0.27.0 do not include InputJSONDelta.",
+)
+def test_collect_ai_data_with_input_json_delta():
+    event = ContentBlockDeltaEvent(
+        delta=InputJSONDelta(partial_json="test", type="input_json_delta"),
+        index=0,
+        type="content_block_delta",
+    )
+
+    input_tokens = 10
+    output_tokens = 20
+    content_blocks = []
+
+    new_input_tokens, new_output_tokens, new_content_blocks = _collect_ai_data(
+        event, input_tokens, output_tokens, content_blocks
+    )
+
+    assert new_input_tokens == input_tokens
+    assert new_output_tokens == output_tokens
+    assert new_content_blocks == ["test"]
+
+
+@pytest.mark.skipif(
+    ANTHROPIC_VERSION < (0, 27),
+    reason="Versions <0.27.0 do not include InputJSONDelta.",
+)
+def test_add_ai_data_to_span_with_input_json_delta(sentry_init):
+    sentry_init(
+        integrations=[AnthropicIntegration(include_prompts=True)],
+        traces_sample_rate=1.0,
+        send_default_pii=True,
+    )
+
+    with start_transaction(name="test"):
+        span = start_span()
+        integration = AnthropicIntegration()
+
+        _add_ai_data_to_span(
+            span,
+            integration,
+            input_tokens=10,
+            output_tokens=20,
+            content_blocks=["{'test': 'data',", "'more': 'json'}"],
+        )
+
+        assert span._data.get(SPANDATA.AI_RESPONSES) == [
+            {"type": "text", "text": "{'test': 'data','more': 'json'}"}
+        ]
+        assert span._data.get("ai.streaming") is True
+        assert span._measurements.get("ai_prompt_tokens_used")["value"] == 10
+        assert span._measurements.get("ai_completion_tokens_used")["value"] == 20
+        assert span._measurements.get("ai_total_tokens_used")["value"] == 30

tests/test_feature_flags.py

@@ -170,6 +170,25 @@ def test_flag_tracking():
         {"flag": "f", "result": False},
     ]
 
+    # Test updates
+    buffer.set("e", True)
+    buffer.set("e", False)
+    buffer.set("e", True)
+    flags = buffer.get()
+    assert flags == [
+        {"flag": "d", "result": False},
+        {"flag": "f", "result": False},
+        {"flag": "e", "result": True},
+    ]
+
+    buffer.set("d", True)
+    flags = buffer.get()
+    assert flags == [
+        {"flag": "f", "result": False},
+        {"flag": "e", "result": True},
+        {"flag": "d", "result": True},
+    ]
+
 
 def test_flag_buffer_concurrent_access():
     buffer = FlagBuffer(capacity=100)