Skip to content

fix: CitationLocation is UnionType, and correctly joining citation chunks when streaming is being used #1341

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
dbschmigelski merged 6 commits into from
Dec 18, 2025
Merged

fix: CitationLocation is UnionType, and correctly joining citation chunks when streaming is being used #1341

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
9413c0b
Fixes CitationLocation UnionType, and when streaming is being used
ericfzhu Dec 16, 2025
f5dfbe7
Modification to citations.py for linter
ericfzhu Dec 16, 2025
88ac264
Ignore dynamic dict construction in bedrock.py for mypy
ericfzhu Dec 16, 2025
937e8f4
feat: align citation content block pattern with ToolChoice
dbschmigelski Dec 16, 2025
35e70da
fix(citations): remove unnecessary key from CitationStreamEvent
dbschmigelski Dec 16, 2025
9383622
Update test__events.py
dbschmigelski Dec 17, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 4 additions & 3 deletions src/strands/event_loop/streaming.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -289,12 +289,13 @@ def handle_content_block_stop(state: dict[str, Any]) -> dict[str, Any]:
state["current_tool_use"] = {}

elif text:
content.append({"text": text})
state["text"] = ""
if citations_content:
citations_block: CitationsContentBlock = {"citations": citations_content}
citations_block: CitationsContentBlock = {"citations": citations_content, "content": [{"text": text}]}
content.append({"citationsContent": citations_block})
state["citationsContent"] = []
else:
content.append({"text": text})
state["text"] = ""

elif reasoning_text:
content_block: ContentBlock = {
Expand Down
11 changes: 1 addition & 10 deletions src/strands/models/bedrock.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -500,16 +500,7 @@ def _format_request_message_content(self, content: ContentBlock) -> dict[str, An
for citation in citations["citations"]:
filtered_citation: dict[str, Any] = {}
if "location" in citation:
location = citation["location"]
filtered_location = {}
# Filter location fields to only include Bedrock-supported ones
if "documentIndex" in location:
filtered_location["documentIndex"] = location["documentIndex"]
if "start" in location:
filtered_location["start"] = location["start"]
if "end" in location:
filtered_location["end"] = location["end"]
filtered_citation["location"] = filtered_location
filtered_citation["location"] = citation["location"]
if "sourceContent" in citation:
filtered_source_content: list[dict[str, Any]] = []
for source_content in citation["sourceContent"]:
Expand Down
2 changes: 1 addition & 1 deletion src/strands/types/_events.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -161,7 +161,7 @@ class CitationStreamEvent(ModelStreamEvent):

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


class ReasoningTextStreamEvent(ModelStreamEvent):
Expand Down
15 changes: 12 additions & 3 deletions src/strands/types/citations.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@
These types are modeled after the Bedrock API.
"""

from typing import List, Union
from typing import List, Literal, Union

from typing_extensions import TypedDict

Expand Down Expand Up @@ -77,8 +77,17 @@ class DocumentPageLocation(TypedDict, total=False):
end: int


# Union type for citation locations
CitationLocation = Union[DocumentCharLocation, DocumentChunkLocation, DocumentPageLocation]
# Tagged union type aliases following the ToolChoice pattern
DocumentCharLocationDict = dict[Literal["documentChar"], DocumentCharLocation]
DocumentPageLocationDict = dict[Literal["documentPage"], DocumentPageLocation]
DocumentChunkLocationDict = dict[Literal["documentChunk"], DocumentChunkLocation]

# Union type for citation locations - tagged union format matching AWS Bedrock API
CitationLocation = Union[
DocumentCharLocationDict,
DocumentPageLocationDict,
DocumentChunkLocationDict,
]


class CitationSourceContent(TypedDict, total=False):
Expand Down
227 changes: 224 additions & 3 deletions tests/strands/event_loop/test_streaming.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -215,6 +215,59 @@ def test_handle_content_block_start(chunk: ContentBlockStartEvent, exp_tool_use)
{},
{},
),
# Citation - New
(
{
"delta": {
"citation": {
"location": {"documentChar": {"documentIndex": 0, "start": 10, "end": 20}},
"title": "Test Doc",
}
}
},
{},
{},
{
"citationsContent": [
{"location": {"documentChar": {"documentIndex": 0, "start": 10, "end": 20}}, "title": "Test Doc"}
]
},
{
"citation": {
"location": {"documentChar": {"documentIndex": 0, "start": 10, "end": 20}},
"title": "Test Doc",
}
},
),
# Citation - Existing
(
{
"delta": {
"citation": {
"location": {"documentPage": {"documentIndex": 1, "start": 5, "end": 6}},
"title": "Another Doc",
}
}
},
{},
{
"citationsContent": [
{"location": {"documentChar": {"documentIndex": 0, "start": 10, "end": 20}}, "title": "Test Doc"}
]
},
{
"citationsContent": [
{"location": {"documentChar": {"documentIndex": 0, "start": 10, "end": 20}}, "title": "Test Doc"},
{"location": {"documentPage": {"documentIndex": 1, "start": 5, "end": 6}}, "title": "Another Doc"},
]
},
{
"citation": {
"location": {"documentPage": {"documentIndex": 1, "start": 5, "end": 6}},
"title": "Another Doc",
}
},
),
# Empty
(
{"delta": {}},
Expand Down Expand Up @@ -294,22 +347,59 @@ def test_handle_content_block_delta(event: ContentBlockDeltaEvent, event_type, s
"redactedContent": b"",
},
),
# Citations
# Text with Citations
(
{
"content": [],
"current_tool_use": {},
"text": "This is cited text",
"reasoningText": "",
"citationsContent": [
{"location": {"documentChar": {"documentIndex": 0, "start": 10, "end": 20}}, "title": "Test Doc"}
],
"redactedContent": b"",
},
{
"content": [
{
"citationsContent": {
"citations": [
{
"location": {"documentChar": {"documentIndex": 0, "start": 10, "end": 20}},
"title": "Test Doc",
}
],
"content": [{"text": "This is cited text"}],
}
}
],
"current_tool_use": {},
"text": "",
"reasoningText": "",
"citationsContent": [],
"redactedContent": b"",
},
),
# Citations without text (should not create content block)
(
{
"content": [],
"current_tool_use": {},
"text": "",
"reasoningText": "",
"citationsContent": [{"citations": [{"text": "test", "source": "test"}]}],
"citationsContent": [
{"location": {"documentChar": {"documentIndex": 0, "start": 10, "end": 20}}, "title": "Test Doc"}
],
"redactedContent": b"",
},
{
"content": [],
"current_tool_use": {},
"text": "",
"reasoningText": "",
"citationsContent": [{"citations": [{"text": "test", "source": "test"}]}],
"citationsContent": [
{"location": {"documentChar": {"documentIndex": 0, "start": 10, "end": 20}}, "title": "Test Doc"}
],
"redactedContent": b"",
},
),
Expand Down Expand Up @@ -578,6 +668,137 @@ def test_extract_usage_metrics_empty_metadata():
},
],
),
# Message with Citations
(
[
{"messageStart": {"role": "assistant"}},
{"contentBlockStart": {"start": {}}},
{"contentBlockDelta": {"delta": {"text": "This is cited text"}}},
{
"contentBlockDelta": {
"delta": {
"citation": {
"location": {"documentChar": {"documentIndex": 0, "start": 10, "end": 20}},
"title": "Test Doc",
}
}
}
},
{
"contentBlockDelta": {
"delta": {
"citation": {
"location": {"documentPage": {"documentIndex": 1, "start": 5, "end": 6}},
"title": "Another Doc",
}
}
}
},
{"contentBlockStop": {}},
{"messageStop": {"stopReason": "end_turn"}},
{
"metadata": {
"usage": {"inputTokens": 5, "outputTokens": 10, "totalTokens": 15},
"metrics": {"latencyMs": 100},
}
},
],
[
{"event": {"messageStart": {"role": "assistant"}}},
{"event": {"contentBlockStart": {"start": {}}}},
{"event": {"contentBlockDelta": {"delta": {"text": "This is cited text"}}}},
{"data": "This is cited text", "delta": {"text": "This is cited text"}},
{
"event": {
"contentBlockDelta": {
"delta": {
"citation": {
"location": {"documentChar": {"documentIndex": 0, "start": 10, "end": 20}},
"title": "Test Doc",
}
}
}
}
},
{
"citation": {
"location": {"documentChar": {"documentIndex": 0, "start": 10, "end": 20}},
"title": "Test Doc",
},
"delta": {
"citation": {
"location": {"documentChar": {"documentIndex": 0, "start": 10, "end": 20}},
"title": "Test Doc",
}
},
},
{
"event": {
"contentBlockDelta": {
"delta": {
"citation": {
"location": {"documentPage": {"documentIndex": 1, "start": 5, "end": 6}},
"title": "Another Doc",
}
}
}
}
},
{
"citation": {
"location": {"documentPage": {"documentIndex": 1, "start": 5, "end": 6}},
"title": "Another Doc",
},
"delta": {
"citation": {
"location": {"documentPage": {"documentIndex": 1, "start": 5, "end": 6}},
"title": "Another Doc",
}
},
},
{"event": {"contentBlockStop": {}}},
{"event": {"messageStop": {"stopReason": "end_turn"}}},
{
"event": {
"metadata": {
"usage": {"inputTokens": 5, "outputTokens": 10, "totalTokens": 15},
"metrics": {"latencyMs": 100},
}
}
},
{
"stop": (
"end_turn",
{
"role": "assistant",
"content": [
{
"citationsContent": {
"citations": [
{
"location": {
"documentChar": {"documentIndex": 0, "start": 10, "end": 20}
},
"title": "Test Doc",
},
{
"location": {
"documentPage": {"documentIndex": 1, "start": 5, "end": 6}
},
"title": "Another Doc",
},
],
"content": [{"text": "This is cited text"}],
}
}
],
},
{"inputTokens": 5, "outputTokens": 10, "totalTokens": 15},
{"latencyMs": 100},
)
},
],
),
# Empty Message
(
[{}],
Expand Down
Loading
Loading