Skip to content

Commit aaa51a2

Browse files
alexander-alderman-webbalexander-alderman-webb
committed
merge
2 parents 5253e29 + e72921f commit aaa51a2

2 files changed

Lines changed: 306 additions & 115 deletions

File tree

‎sentry_sdk/integrations/boto3/_services/bedrock.py‎

Lines changed: 170 additions & 61 deletions
Original file line numberDiff line numberDiff line change
@@ -94,7 +94,7 @@ def sentry_patched_iter(
9494
)
9595

9696
role: "Optional[str]" = None
97-
content_blocks: "dict[int, dict[str, Any]]" = {}
97+
parts: "dict[int, Union[TextPart, ReasoningPart, ToolCallPart, ToolCallResponsePart, UriPart, BlobPart]]" = {}
9898

9999
try:
100100
for event in original_iter(self):
@@ -132,6 +132,10 @@ def sentry_patched_iter(
132132
yield event
133133
continue
134134

135+
if "messageStop" in event:
136+
finish_reasons: "list[str]" = [event["messageStop"]["stopReason"]]
137+
set_on_span(SPANDATA.GEN_AI_RESPONSE_FINISH_REASONS, finish_reasons)
138+
135139
client = sentry_sdk.get_client()
136140
if (
137141
not has_data_collection_enabled(client.options)
@@ -149,77 +153,182 @@ def sentry_patched_iter(
149153
yield event
150154
continue
151155

152-
if "contentBlockStart" in event:
153-
block_start = event["contentBlockStart"]
154-
tool_use = block_start["start"].get("toolUse")
155-
if tool_use is not None:
156-
tool_use_block: "dict[str, Any]" = {"input": ""}
157-
if "name" in tool_use:
158-
tool_use_block["name"] = tool_use["name"]
159-
content_blocks[block_start["contentBlockIndex"]] = {
160-
"toolUse": tool_use_block
161-
}
156+
if "messageStop" in event and role is not None:
157+
set_on_span(
158+
SPANDATA.GEN_AI_OUTPUT_MESSAGES,
159+
json.dumps(
160+
[
161+
{
162+
"role": role,
163+
"parts": [parts[index] for index in sorted(parts)],
164+
}
165+
]
166+
),
167+
)
162168

163169
yield event
164170
continue
165171

166-
if "contentBlockDelta" in event:
167-
index = event["contentBlockDelta"]["contentBlockIndex"]
168-
delta = event["contentBlockDelta"]["delta"]
172+
if (
173+
"contentBlockStart" in event
174+
and "toolUse" in event["contentBlockStart"]["start"]
175+
):
176+
tool_call_part: "ToolCallPart" = {
177+
"type": "tool_call",
178+
"name": event["contentBlockStart"]["start"]["toolUse"]["name"],
179+
}
180+
parts[event["contentBlockStart"]["contentBlockIndex"]] = (
181+
tool_call_part
182+
)
183+
184+
yield event
185+
continue
169186

170-
if "text" in delta:
171-
block = content_blocks.setdefault(index, {"text": ""})
172-
block["text"] += delta["text"]
173-
elif "toolUse" in delta:
174-
block = content_blocks.setdefault(
175-
index, {"toolUse": {"input": ""}}
176-
)
177-
block["toolUse"]["input"] += delta["toolUse"]["input"]
178-
elif "reasoningContent" in delta:
179-
block = content_blocks.setdefault(
180-
index, {"reasoningContent": {}}
181-
)
182-
reasoning_text = delta["reasoningContent"].get("text")
183-
if reasoning_text is not None:
184-
block["reasoningContent"].setdefault(
185-
"reasoningText", {"text": ""}
186-
)["text"] += reasoning_text
187+
if (
188+
"contentBlockStart" in event
189+
and "toolResult" in event["contentBlockStart"]["start"]
190+
):
191+
tool_call_response_part: "ToolCallResponsePart" = {
192+
"type": "tool_call_response",
193+
}
194+
parts[event["contentBlockStart"]["contentBlockIndex"]] = (
195+
tool_call_response_part
196+
)
187197

188198
yield event
189199
continue
190200

191-
if "messageStop" in event:
192-
finish_reasons: "list[str]" = [event["messageStop"]["stopReason"]]
193-
set_on_span(SPANDATA.GEN_AI_RESPONSE_FINISH_REASONS, finish_reasons)
201+
if "contentBlockDelta" not in event:
202+
yield event
203+
continue
194204

195-
if "messageStop" in event and role is not None:
196-
for block in content_blocks.values():
197-
if "toolUse" not in block:
198-
continue
199-
# Tool input is streamed as a JSON string; non-streaming responses contain the parsed object.
200-
try:
201-
block["toolUse"]["input"] = json.loads(
202-
block["toolUse"]["input"]
203-
)
204-
except ValueError:
205-
pass
205+
index = event["contentBlockDelta"]["contentBlockIndex"]
206+
delta = event["contentBlockDelta"]["delta"]
206207

207-
set_on_span(
208-
SPANDATA.GEN_AI_OUTPUT_MESSAGES,
209-
json.dumps(
210-
_transform_messages(
211-
[
212-
{
213-
"role": role,
214-
"content": [
215-
content_blocks[index]
216-
for index in sorted(content_blocks)
217-
],
218-
}
219-
]
220-
)
221-
),
222-
)
208+
if "text" in delta and index not in parts:
209+
parts[index] = {"type": "text", "content": delta["text"]}
210+
211+
yield event
212+
continue
213+
214+
if "text" in delta:
215+
cast("TextPart", parts[index])["content"] += delta["text"]
216+
217+
yield event
218+
continue
219+
220+
if "toolUse" in delta and index not in parts:
221+
parts[index] = {
222+
"type": "tool_call",
223+
"arguments": delta["toolUse"]["input"],
224+
}
225+
226+
yield event
227+
continue
228+
229+
if "toolUse" in delta and "arguments" in parts[index]:
230+
cast("ToolCallPart", parts[index])["arguments"] += delta["toolUse"][
231+
"input"
232+
]
233+
234+
yield event
235+
continue
236+
237+
if "toolUse" in delta:
238+
cast("ToolCallPart", parts[index])["arguments"] = delta["toolUse"][
239+
"input"
240+
]
241+
242+
yield event
243+
continue
244+
245+
if "toolResult" in delta and index not in parts:
246+
parts[index] = {
247+
"type": "tool_call_response",
248+
"response": [
249+
item
250+
for item in delta["toolResult"]
251+
if "json" in item or "text" in item
252+
],
253+
}
254+
255+
yield event
256+
continue
257+
258+
if "toolResult" in delta and "response" in parts[index]:
259+
cast("ToolCallResponsePart", parts[index])["response"] += [
260+
item
261+
for item in delta["toolResult"]
262+
if "json" in item or "text" in item
263+
]
264+
265+
yield event
266+
continue
267+
268+
if "toolResult" in delta:
269+
cast("ToolCallResponsePart", parts[index])["response"] = [
270+
item
271+
for item in delta["toolResult"]
272+
if "json" in item or "text" in item
273+
]
274+
275+
yield event
276+
continue
277+
278+
if "reasoningContent" in delta and index not in parts:
279+
parts[index] = {"type": "reasoning"}
280+
281+
if "text" in delta["reasoningContent"]:
282+
cast("ReasoningPart", parts[index])["content"] = delta[
283+
"reasoningContent"
284+
]["text"]
285+
286+
yield event
287+
continue
288+
289+
if (
290+
"reasoningContent" in delta
291+
and "text" in delta["reasoningContent"]
292+
and "content" in parts[index]
293+
):
294+
cast("ReasoningPart", parts[index])["content"] += delta[
295+
"reasoningContent"
296+
]["text"]
297+
298+
yield event
299+
continue
300+
301+
if "reasoningContent" in delta and "text" in delta["reasoningContent"]:
302+
cast("ReasoningPart", parts[index])["content"] = delta[
303+
"reasoningContent"
304+
]["text"]
305+
306+
yield event
307+
continue
308+
309+
if (
310+
"image" in delta
311+
and "source" in delta["image"]
312+
and "s3Location" in delta["image"]["source"]
313+
):
314+
parts[index] = {
315+
"type": "uri",
316+
"modality": "image",
317+
"uri": delta["image"]["source"]["s3Location"]["uri"],
318+
}
319+
320+
yield event
321+
continue
322+
323+
if (
324+
"image" in delta
325+
and "source" in delta["image"]
326+
and "bytes" in delta["image"]["source"]
327+
):
328+
parts[index] = {
329+
"type": "blob",
330+
"modality": "image",
331+
}
223332

224333
yield event
225334
continue

0 commit comments

Comments
 (0)