Skip to content

Commit a3cdae1

Browse files
sentrivanasentrivana
committed
Merge branch 'master' into ivana/move-http-crumbs-2
2 parents 70ad7a9 + 9ee8af5 commit a3cdae1

22 files changed

Lines changed: 2854 additions & 428 deletions

File tree

‎sentry_sdk/_span_batcher.py‎

Lines changed: 9 additions & 38 deletions
Original file line numberDiff line numberDiff line change
@@ -21,16 +21,12 @@ class SpanBatcher(Batcher["SpanJSON"]):
2121
# MAX_BEFORE_FLUSH should be lower than MAX_BEFORE_DROP, so that there is
2222
# a bit of a buffer for spans that appear between the trigger to flush
2323
# and actually flushing the buffer.
24+
#
25+
# The max limits are all per trace (per bucket).
2426
MAX_ENVELOPE_SIZE = 1000 # spans
25-
2627
MAX_BEFORE_FLUSH = 1000
27-
GLOBAL_MAX_BEFORE_FLUSH = 5_000
28-
2928
MAX_BEFORE_DROP = 2000
30-
GLOBAL_MAX_BEFORE_DROP = 10_000
31-
3229
MAX_BYTES_BEFORE_FLUSH = 5 * 1024 * 1024 # 5 MB
33-
GLOBAL_MAX_BYTES_BEFORE_FLUSH = 25 * 1024 * 1024 # 25 MB
3430

3531
FLUSH_WAIT_TIME = 5.0
3632

@@ -48,11 +44,7 @@ def __init__(
4844
# envelope.
4945
# trace_id -> span buffer
5046
self._span_buffer: dict[str, list["SpanJSON"]] = defaultdict(list)
51-
self._span_number: int = 0
52-
5347
self._running_size: dict[str, int] = defaultdict(lambda: 0)
54-
self._total_running_size: int = 0
55-
5648
self._capture_func = capture_func
5749
self._record_lost_func = record_lost_func
5850
self._running = True
@@ -79,11 +71,7 @@ def _reset_in_child() -> None:
7971

8072
def _reset_thread_state(self) -> None:
8173
self._span_buffer = defaultdict(list)
82-
self._span_number = 0
83-
8474
self._running_size = defaultdict(lambda: 0)
85-
self._total_running_size = 0
86-
8775
self._running = True
8876

8977
self._lock = threading.Lock()
@@ -106,12 +94,8 @@ def _flush_loop(self) -> None:
10694
self._flush(only_pending=True)
10795

10896
if (
109-
self._span_number >= self.GLOBAL_MAX_BEFORE_FLUSH
110-
or self._total_running_size >= self.GLOBAL_MAX_BYTES_BEFORE_FLUSH
111-
or (
112-
time.monotonic() - self._last_full_flush
113-
>= self.FLUSH_WAIT_TIME + jitter
114-
)
97+
time.monotonic() - self._last_full_flush
98+
>= self.FLUSH_WAIT_TIME + jitter
11599
):
116100
self._flush()
117101
self._last_full_flush = time.monotonic()
@@ -132,10 +116,8 @@ def add(self, span: "SpanJSON") -> None:
132116
return None
133117

134118
with self._lock:
135-
if (
136-
self._span_number >= self.GLOBAL_MAX_BEFORE_DROP
137-
or len(self._span_buffer[span["trace_id"]]) >= self.MAX_BEFORE_DROP
138-
):
119+
size = len(self._span_buffer[span["trace_id"]])
120+
if size >= self.MAX_BEFORE_DROP:
139121
self._record_lost_func(
140122
reason="queue_overflow",
141123
data_category="span",
@@ -144,25 +126,17 @@ def add(self, span: "SpanJSON") -> None:
144126
return None
145127

146128
self._span_buffer[span["trace_id"]].append(span)
147-
self._span_number += 1
148-
149-
estimated_size = self._estimate_size(span)
150-
self._running_size[span["trace_id"]] += estimated_size
151-
self._total_running_size += estimated_size
129+
self._running_size[span["trace_id"]] += self._estimate_size(span)
152130

153131
if (
154-
len(self._span_buffer[span["trace_id"]]) >= self.MAX_BEFORE_FLUSH
132+
size + 1 >= self.MAX_BEFORE_FLUSH
155133
or self._running_size[span["trace_id"]]
156134
>= self.MAX_BYTES_BEFORE_FLUSH
157135
):
158136
self._pending_flush.add(span["trace_id"])
159137
notify = True
160138
else:
161-
notify = (
162-
self._span_number >= self.GLOBAL_MAX_BEFORE_FLUSH
163-
or self._total_running_size
164-
>= self.GLOBAL_MAX_BYTES_BEFORE_FLUSH
165-
)
139+
notify = False
166140

167141
if notify:
168142
self._flush_event.set()
@@ -253,10 +227,7 @@ def _flush(self, only_pending: bool = False) -> None:
253227

254228
envelopes.append(envelope)
255229

256-
self._span_number -= len(self._span_buffer[bucket_id])
257230
del self._span_buffer[bucket_id]
258-
259-
self._total_running_size -= self._running_size[bucket_id]
260231
del self._running_size[bucket_id]
261232

262233
for envelope in envelopes:

‎sentry_sdk/integrations/anthropic.py‎

Lines changed: 7 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -608,14 +608,16 @@ def _set_output_data(
608608
set_on_span(SPANDATA.GEN_AI_RESPONSE_FINISH_REASONS, [finish_reason])
609609

610610
client = sentry_sdk.get_client()
611+
record_inputs = False
611612
record_outputs = False
612613
if has_data_collection_enabled(client.options):
613-
if client.options["data_collection"]["gen_ai"]["outputs"]:
614-
record_outputs = True
614+
record_inputs = client.options["data_collection"]["gen_ai"]["inputs"]
615+
record_outputs = client.options["data_collection"]["gen_ai"]["outputs"]
615616
elif should_send_default_pii() and integration.include_prompts:
617+
record_inputs = True
616618
record_outputs = True
617619

618-
if record_outputs:
620+
if record_inputs or record_outputs:
619621
output_messages: "dict[str, list[Any]]" = {
620622
"response": [],
621623
"tool": [],
@@ -627,15 +629,15 @@ def _set_output_data(
627629
elif output["type"] == "tool_use":
628630
output_messages["tool"].append(output)
629631

630-
if len(output_messages["tool"]) > 0:
632+
if record_inputs and len(output_messages["tool"]) > 0:
631633
set_data_normalized(
632634
span,
633635
SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS,
634636
output_messages["tool"],
635637
unpack=False,
636638
)
637639

638-
if len(output_messages["response"]) > 0:
640+
if record_outputs and len(output_messages["response"]) > 0:
639641
set_data_normalized(
640642
span, SPANDATA.GEN_AI_RESPONSE_TEXT, output_messages["response"]
641643
)

‎sentry_sdk/integrations/google_genai/streaming.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -160,7 +160,7 @@ def set_span_data_for_streaming_response(
160160

161161
if accumulated_response.get("tool_calls"):
162162
if has_data_collection_enabled(client.options):
163-
if client.options["data_collection"]["gen_ai"]["outputs"]:
163+
if client.options["data_collection"]["gen_ai"]["inputs"]:
164164
set_on_span(
165165
SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS,
166166
safe_serialize(accumulated_response["tool_calls"]),

‎sentry_sdk/integrations/google_genai/utils.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1034,7 +1034,7 @@ def set_span_data_for_response(
10341034
tool_calls = extract_tool_calls(response)
10351035
if tool_calls:
10361036
if has_data_collection_enabled(client.options):
1037-
if client.options["data_collection"]["gen_ai"]["outputs"]:
1037+
if client.options["data_collection"]["gen_ai"]["inputs"]:
10381038
set_on_span(
10391039
SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS, safe_serialize(tool_calls)
10401040
)

‎sentry_sdk/integrations/huggingface_hub.py‎

Lines changed: 84 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@
1818
from sentry_sdk.utils import (
1919
capture_internal_exceptions,
2020
event_from_exception,
21+
has_data_collection_enabled,
2122
reraise,
2223
)
2324

@@ -74,7 +75,8 @@ def _capture_exception(exc: "Any") -> None:
7475
def _wrap_huggingface_task(f: "Callable[..., Any]", op: str) -> "Callable[..., Any]":
7576
@wraps(f)
7677
def new_huggingface_task(*args: "Any", **kwargs: "Any") -> "Any":
77-
integration = sentry_sdk.get_client().get_integration(HuggingfaceHubIntegration)
78+
client = sentry_sdk.get_client()
79+
integration = client.get_integration(HuggingfaceHubIntegration)
7880
if integration is None:
7981
return f(*args, **kwargs)
8082

@@ -91,12 +93,12 @@ def new_huggingface_task(*args: "Any", **kwargs: "Any") -> "Any":
9193
# invalid call, dont instrument, let it return error
9294
return f(*args, **kwargs)
9395

94-
client = args[0]
95-
model = client.model or kwargs.get("model") or ""
96+
hf_client = args[0]
97+
model = hf_client.model or kwargs.get("model") or ""
9698
operation_name = op.split(".")[-1]
9799

98100
span: "Union[Span, StreamedSpan]"
99-
if has_span_streaming_enabled(sentry_sdk.get_client().options):
101+
if has_span_streaming_enabled(client.options):
100102
span = sentry_sdk.traces.start_span(
101103
name=f"{operation_name} {model}",
102104
attributes={
@@ -117,14 +119,7 @@ def new_huggingface_task(*args: "Any", **kwargs: "Any") -> "Any":
117119
if model:
118120
_set_span_data_attribute(span, SPANDATA.GEN_AI_REQUEST_MODEL, model)
119121

120-
# Input attributes
121-
if should_send_default_pii() and integration.include_prompts:
122-
set_data_normalized(
123-
span, SPANDATA.GEN_AI_REQUEST_MESSAGES, prompt, unpack=False
124-
)
125-
126122
attribute_mapping = {
127-
"tools": SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS,
128123
"frequency_penalty": SPANDATA.GEN_AI_REQUEST_FREQUENCY_PENALTY,
129124
"max_tokens": SPANDATA.GEN_AI_REQUEST_MAX_TOKENS,
130125
"presence_penalty": SPANDATA.GEN_AI_REQUEST_PRESENCE_PENALTY,
@@ -134,6 +129,24 @@ def new_huggingface_task(*args: "Any", **kwargs: "Any") -> "Any":
134129
"stream": SPANDATA.GEN_AI_RESPONSE_STREAMING,
135130
}
136131

132+
if has_data_collection_enabled(client.options):
133+
if client.options["data_collection"]["gen_ai"]["inputs"]:
134+
attribute_mapping["tools"] = SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS
135+
else:
136+
# Legacy behaviour where we unconditionally set this. Remove when data collection is fully rolled out
137+
attribute_mapping["tools"] = SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS
138+
139+
# Input attributes
140+
if has_data_collection_enabled(client.options):
141+
if client.options["data_collection"]["gen_ai"]["inputs"]:
142+
set_data_normalized(
143+
span, SPANDATA.GEN_AI_REQUEST_MESSAGES, prompt, unpack=False
144+
)
145+
elif should_send_default_pii() and integration.include_prompts:
146+
set_data_normalized(
147+
span, SPANDATA.GEN_AI_REQUEST_MESSAGES, prompt, unpack=False
148+
)
149+
137150
for attribute, span_attribute in attribute_mapping.items():
138151
value = kwargs.get(attribute, None)
139152
if value is not None:
@@ -210,18 +223,34 @@ def new_huggingface_task(*args: "Any", **kwargs: "Any") -> "Any":
210223
finish_reason,
211224
)
212225

213-
if should_send_default_pii() and integration.include_prompts:
214-
if tool_calls is not None and len(tool_calls) > 0:
226+
if tool_calls is not None and len(tool_calls) > 0:
227+
if has_data_collection_enabled(client.options):
228+
if client.options["data_collection"]["gen_ai"]["inputs"]:
229+
set_data_normalized(
230+
span,
231+
SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS,
232+
tool_calls,
233+
unpack=False,
234+
)
235+
elif should_send_default_pii() and integration.include_prompts:
215236
set_data_normalized(
216237
span,
217238
SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS,
218239
tool_calls,
219240
unpack=False,
220241
)
221242

222-
if len(response_text_buffer) > 0:
223-
text_response = "".join(response_text_buffer)
224-
if text_response:
243+
if len(response_text_buffer) > 0:
244+
text_response = "".join(response_text_buffer)
245+
if text_response:
246+
if has_data_collection_enabled(client.options):
247+
if client.options["data_collection"]["gen_ai"]["outputs"]:
248+
set_data_normalized(
249+
span,
250+
SPANDATA.GEN_AI_RESPONSE_TEXT,
251+
text_response,
252+
)
253+
elif should_send_default_pii() and integration.include_prompts:
225254
set_data_normalized(
226255
span,
227256
SPANDATA.GEN_AI_RESPONSE_TEXT,
@@ -284,7 +313,14 @@ def new_details_iterator() -> "Iterable[Any]":
284313
finish_reason,
285314
)
286315

287-
if should_send_default_pii() and integration.include_prompts:
316+
should_set_response_text = False
317+
if has_data_collection_enabled(client.options):
318+
if client.options["data_collection"]["gen_ai"]["outputs"]:
319+
should_set_response_text = True
320+
elif should_send_default_pii() and integration.include_prompts:
321+
should_set_response_text = True
322+
323+
if should_set_response_text:
288324
if len(response_text_buffer) > 0:
289325
text_response = "".join(response_text_buffer)
290326
if text_response:
@@ -363,18 +399,44 @@ def new_iterator() -> "Iterable[ChatCompletionStreamOutput]":
363399
finish_reason,
364400
)
365401

366-
if should_send_default_pii() and integration.include_prompts:
367-
if tool_calls is not None and len(tool_calls) > 0:
402+
if tool_calls is not None and len(tool_calls) > 0:
403+
if has_data_collection_enabled(client.options):
404+
if client.options["data_collection"]["gen_ai"][
405+
"inputs"
406+
]:
407+
set_data_normalized(
408+
span,
409+
SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS,
410+
tool_calls,
411+
unpack=False,
412+
)
413+
elif (
414+
should_send_default_pii()
415+
and integration.include_prompts
416+
):
368417
set_data_normalized(
369418
span,
370419
SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS,
371420
tool_calls,
372421
unpack=False,
373422
)
374423

375-
if len(response_text_buffer) > 0:
376-
text_response = "".join(response_text_buffer)
377-
if text_response:
424+
if len(response_text_buffer) > 0:
425+
text_response = "".join(response_text_buffer)
426+
if text_response:
427+
if has_data_collection_enabled(client.options):
428+
if client.options["data_collection"]["gen_ai"][
429+
"outputs"
430+
]:
431+
set_data_normalized(
432+
span,
433+
SPANDATA.GEN_AI_RESPONSE_TEXT,
434+
text_response,
435+
)
436+
elif (
437+
should_send_default_pii()
438+
and integration.include_prompts
439+
):
378440
set_data_normalized(
379441
span,
380442
SPANDATA.GEN_AI_RESPONSE_TEXT,

0 commit comments

Comments
 (0)