Skip to content

Commit 1792ddb

Browse files
ericfzhudbschmigelski
ericfzhu
and
dbschmigelski
authored
fix(bedrock): CitationLocation is UnionType, and correctly joining citation chunks when streaming is being used (#1341)
--------- Co-authored-by: Eric Zhu <[email protected]> Co-authored-by: Dean Schmigelski <[email protected]>
1 parent 82f5bcf commit 1792ddb

File tree

8 files changed

+323
-22
lines changed

8 files changed

+323
-22
lines changed

‎src/strands/event_loop/streaming.py‎

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -289,12 +289,13 @@ def handle_content_block_stop(state: dict[str, Any]) -> dict[str, Any]:
289289
state["current_tool_use"] = {}
290290

291291
elif text:
292-
content.append({"text": text})
293-
state["text"] = ""
294292
if citations_content:
295-
citations_block: CitationsContentBlock = {"citations": citations_content}
293+
citations_block: CitationsContentBlock = {"citations": citations_content, "content": [{"text": text}]}
296294
content.append({"citationsContent": citations_block})
297295
state["citationsContent"] = []
296+
else:
297+
content.append({"text": text})
298+
state["text"] = ""
298299

299300
elif reasoning_text:
300301
content_block: ContentBlock = {

‎src/strands/models/bedrock.py‎

Lines changed: 1 addition & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -500,16 +500,7 @@ def _format_request_message_content(self, content: ContentBlock) -> dict[str, An
500500
for citation in citations["citations"]:
501501
filtered_citation: dict[str, Any] = {}
502502
if "location" in citation:
503-
location = citation["location"]
504-
filtered_location = {}
505-
# Filter location fields to only include Bedrock-supported ones
506-
if "documentIndex" in location:
507-
filtered_location["documentIndex"] = location["documentIndex"]
508-
if "start" in location:
509-
filtered_location["start"] = location["start"]
510-
if "end" in location:
511-
filtered_location["end"] = location["end"]
512-
filtered_citation["location"] = filtered_location
503+
filtered_citation["location"] = citation["location"]
513504
if "sourceContent" in citation:
514505
filtered_source_content: list[dict[str, Any]] = []
515506
for source_content in citation["sourceContent"]:

‎src/strands/types/_events.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -161,7 +161,7 @@ class CitationStreamEvent(ModelStreamEvent):
161161

162162
def __init__(self, delta: ContentBlockDelta, citation: Citation) -> None:
163163
"""Initialize with delta and citation content."""
164-
super().__init__({"callback": {"citation": citation, "delta": delta}})
164+
super().__init__({"citation": citation, "delta": delta})
165165

166166

167167
class ReasoningTextStreamEvent(ModelStreamEvent):

‎src/strands/types/citations.py‎

Lines changed: 12 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@
33
These types are modeled after the Bedrock API.
44
"""
55

6-
from typing import List, Union
6+
from typing import List, Literal, Union
77

88
from typing_extensions import TypedDict
99

@@ -77,8 +77,17 @@ class DocumentPageLocation(TypedDict, total=False):
7777
end: int
7878

7979

80-
# Union type for citation locations
81-
CitationLocation = Union[DocumentCharLocation, DocumentChunkLocation, DocumentPageLocation]
80+
# Tagged union type aliases following the ToolChoice pattern
81+
DocumentCharLocationDict = dict[Literal["documentChar"], DocumentCharLocation]
82+
DocumentPageLocationDict = dict[Literal["documentPage"], DocumentPageLocation]
83+
DocumentChunkLocationDict = dict[Literal["documentChunk"], DocumentChunkLocation]
84+
85+
# Union type for citation locations - tagged union format matching AWS Bedrock API
86+
CitationLocation = Union[
87+
DocumentCharLocationDict,
88+
DocumentPageLocationDict,
89+
DocumentChunkLocationDict,
90+
]
8291

8392

8493
class CitationSourceContent(TypedDict, total=False):

‎tests/strands/event_loop/test_streaming.py‎

Lines changed: 224 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -215,6 +215,59 @@ def test_handle_content_block_start(chunk: ContentBlockStartEvent, exp_tool_use)
215215
{},
216216
{},
217217
),
218+
# Citation - New
219+
(
220+
{
221+
"delta": {
222+
"citation": {
223+
"location": {"documentChar": {"documentIndex": 0, "start": 10, "end": 20}},
224+
"title": "Test Doc",
225+
}
226+
}
227+
},
228+
{},
229+
{},
230+
{
231+
"citationsContent": [
232+
{"location": {"documentChar": {"documentIndex": 0, "start": 10, "end": 20}}, "title": "Test Doc"}
233+
]
234+
},
235+
{
236+
"citation": {
237+
"location": {"documentChar": {"documentIndex": 0, "start": 10, "end": 20}},
238+
"title": "Test Doc",
239+
}
240+
},
241+
),
242+
# Citation - Existing
243+
(
244+
{
245+
"delta": {
246+
"citation": {
247+
"location": {"documentPage": {"documentIndex": 1, "start": 5, "end": 6}},
248+
"title": "Another Doc",
249+
}
250+
}
251+
},
252+
{},
253+
{
254+
"citationsContent": [
255+
{"location": {"documentChar": {"documentIndex": 0, "start": 10, "end": 20}}, "title": "Test Doc"}
256+
]
257+
},
258+
{
259+
"citationsContent": [
260+
{"location": {"documentChar": {"documentIndex": 0, "start": 10, "end": 20}}, "title": "Test Doc"},
261+
{"location": {"documentPage": {"documentIndex": 1, "start": 5, "end": 6}}, "title": "Another Doc"},
262+
]
263+
},
264+
{
265+
"citation": {
266+
"location": {"documentPage": {"documentIndex": 1, "start": 5, "end": 6}},
267+
"title": "Another Doc",
268+
}
269+
},
270+
),
218271
# Empty
219272
(
220273
{"delta": {}},
@@ -294,22 +347,59 @@ def test_handle_content_block_delta(event: ContentBlockDeltaEvent, event_type, s
294347
"redactedContent": b"",
295348
},
296349
),
297-
# Citations
350+
# Text with Citations
351+
(
352+
{
353+
"content": [],
354+
"current_tool_use": {},
355+
"text": "This is cited text",
356+
"reasoningText": "",
357+
"citationsContent": [
358+
{"location": {"documentChar": {"documentIndex": 0, "start": 10, "end": 20}}, "title": "Test Doc"}
359+
],
360+
"redactedContent": b"",
361+
},
362+
{
363+
"content": [
364+
{
365+
"citationsContent": {
366+
"citations": [
367+
{
368+
"location": {"documentChar": {"documentIndex": 0, "start": 10, "end": 20}},
369+
"title": "Test Doc",
370+
}
371+
],
372+
"content": [{"text": "This is cited text"}],
373+
}
374+
}
375+
],
376+
"current_tool_use": {},
377+
"text": "",
378+
"reasoningText": "",
379+
"citationsContent": [],
380+
"redactedContent": b"",
381+
},
382+
),
383+
# Citations without text (should not create content block)
298384
(
299385
{
300386
"content": [],
301387
"current_tool_use": {},
302388
"text": "",
303389
"reasoningText": "",
304-
"citationsContent": [{"citations": [{"text": "test", "source": "test"}]}],
390+
"citationsContent": [
391+
{"location": {"documentChar": {"documentIndex": 0, "start": 10, "end": 20}}, "title": "Test Doc"}
392+
],
305393
"redactedContent": b"",
306394
},
307395
{
308396
"content": [],
309397
"current_tool_use": {},
310398
"text": "",
311399
"reasoningText": "",
312-
"citationsContent": [{"citations": [{"text": "test", "source": "test"}]}],
400+
"citationsContent": [
401+
{"location": {"documentChar": {"documentIndex": 0, "start": 10, "end": 20}}, "title": "Test Doc"}
402+
],
313403
"redactedContent": b"",
314404
},
315405
),
@@ -578,6 +668,137 @@ def test_extract_usage_metrics_empty_metadata():
578668
},
579669
],
580670
),
671+
# Message with Citations
672+
(
673+
[
674+
{"messageStart": {"role": "assistant"}},
675+
{"contentBlockStart": {"start": {}}},
676+
{"contentBlockDelta": {"delta": {"text": "This is cited text"}}},
677+
{
678+
"contentBlockDelta": {
679+
"delta": {
680+
"citation": {
681+
"location": {"documentChar": {"documentIndex": 0, "start": 10, "end": 20}},
682+
"title": "Test Doc",
683+
}
684+
}
685+
}
686+
},
687+
{
688+
"contentBlockDelta": {
689+
"delta": {
690+
"citation": {
691+
"location": {"documentPage": {"documentIndex": 1, "start": 5, "end": 6}},
692+
"title": "Another Doc",
693+
}
694+
}
695+
}
696+
},
697+
{"contentBlockStop": {}},
698+
{"messageStop": {"stopReason": "end_turn"}},
699+
{
700+
"metadata": {
701+
"usage": {"inputTokens": 5, "outputTokens": 10, "totalTokens": 15},
702+
"metrics": {"latencyMs": 100},
703+
}
704+
},
705+
],
706+
[
707+
{"event": {"messageStart": {"role": "assistant"}}},
708+
{"event": {"contentBlockStart": {"start": {}}}},
709+
{"event": {"contentBlockDelta": {"delta": {"text": "This is cited text"}}}},
710+
{"data": "This is cited text", "delta": {"text": "This is cited text"}},
711+
{
712+
"event": {
713+
"contentBlockDelta": {
714+
"delta": {
715+
"citation": {
716+
"location": {"documentChar": {"documentIndex": 0, "start": 10, "end": 20}},
717+
"title": "Test Doc",
718+
}
719+
}
720+
}
721+
}
722+
},
723+
{
724+
"citation": {
725+
"location": {"documentChar": {"documentIndex": 0, "start": 10, "end": 20}},
726+
"title": "Test Doc",
727+
},
728+
"delta": {
729+
"citation": {
730+
"location": {"documentChar": {"documentIndex": 0, "start": 10, "end": 20}},
731+
"title": "Test Doc",
732+
}
733+
},
734+
},
735+
{
736+
"event": {
737+
"contentBlockDelta": {
738+
"delta": {
739+
"citation": {
740+
"location": {"documentPage": {"documentIndex": 1, "start": 5, "end": 6}},
741+
"title": "Another Doc",
742+
}
743+
}
744+
}
745+
}
746+
},
747+
{
748+
"citation": {
749+
"location": {"documentPage": {"documentIndex": 1, "start": 5, "end": 6}},
750+
"title": "Another Doc",
751+
},
752+
"delta": {
753+
"citation": {
754+
"location": {"documentPage": {"documentIndex": 1, "start": 5, "end": 6}},
755+
"title": "Another Doc",
756+
}
757+
},
758+
},
759+
{"event": {"contentBlockStop": {}}},
760+
{"event": {"messageStop": {"stopReason": "end_turn"}}},
761+
{
762+
"event": {
763+
"metadata": {
764+
"usage": {"inputTokens": 5, "outputTokens": 10, "totalTokens": 15},
765+
"metrics": {"latencyMs": 100},
766+
}
767+
}
768+
},
769+
{
770+
"stop": (
771+
"end_turn",
772+
{
773+
"role": "assistant",
774+
"content": [
775+
{
776+
"citationsContent": {
777+
"citations": [
778+
{
779+
"location": {
780+
"documentChar": {"documentIndex": 0, "start": 10, "end": 20}
781+
},
782+
"title": "Test Doc",
783+
},
784+
{
785+
"location": {
786+
"documentPage": {"documentIndex": 1, "start": 5, "end": 6}
787+
},
788+
"title": "Another Doc",
789+
},
790+
],
791+
"content": [{"text": "This is cited text"}],
792+
}
793+
}
794+
],
795+
},
796+
{"inputTokens": 5, "outputTokens": 10, "totalTokens": 15},
797+
{"latencyMs": 100},
798+
)
799+
},
800+
],
801+
),
581802
# Empty Message
582803
(
583804
[{}],

0 commit comments

Comments
 (0)