diff --git a/src/strands/types/citations.py b/src/strands/types/citations.py index 41f2fa4e0..623f6ddc7 100644 --- a/src/strands/types/citations.py +++ b/src/strands/types/citations.py @@ -77,16 +77,55 @@ class DocumentPageLocation(TypedDict, total=False): end: int +class SearchResultLocation(TypedDict, total=False): + """Specifies a search result location within the content array. + + Provides positioning information for cited content using search result + index and block positions. + + Attributes: + searchResultIndex: The index of the search result content block where + the cited content is found. Minimum value of 0. + start: The starting position in the content array where the cited + content begins. Minimum value of 0. + end: The ending position in the content array where the cited + content ends. Minimum value of 0. + """ + + searchResultIndex: int + start: int + end: int + + +class WebLocation(TypedDict, total=False): + """Provides the URL and domain information for a cited website. + + Contains information about the website that was cited when performing + a web search. + + Attributes: + url: The URL that was cited when performing a web search. + domain: The domain that was cited when performing a web search. + """ + + url: str + domain: str + + # Tagged union type aliases following the ToolChoice pattern DocumentCharLocationDict = dict[Literal["documentChar"], DocumentCharLocation] DocumentPageLocationDict = dict[Literal["documentPage"], DocumentPageLocation] DocumentChunkLocationDict = dict[Literal["documentChunk"], DocumentChunkLocation] +SearchResultLocationDict = dict[Literal["searchResultLocation"], SearchResultLocation] +WebLocationDict = dict[Literal["web"], WebLocation] # Union type for citation locations - tagged union format matching AWS Bedrock API CitationLocation = Union[ DocumentCharLocationDict, DocumentPageLocationDict, DocumentChunkLocationDict, + SearchResultLocationDict, + WebLocationDict, ] diff --git a/tests/strands/models/test_bedrock.py b/tests/strands/models/test_bedrock.py index 5ec5a7072..33be44b1b 100644 --- a/tests/strands/models/test_bedrock.py +++ b/tests/strands/models/test_bedrock.py @@ -2078,14 +2078,15 @@ async def test_citations_content_preserves_tagged_union_structure(bedrock_client This test verifies that when messages contain citationsContent with tagged union CitationLocation objects, the structure is preserved when sent to AWS Bedrock API. AWS Bedrock expects CitationLocation to be a - tagged union with exactly one wrapper key (documentChar, documentPage, etc.) containing the location fields. + tagged union with exactly one wrapper key (documentChar, documentPage, documentChunk, searchResultLocation, web) + containing the location fields. """ # Mock the Bedrock response bedrock_client.converse_stream.return_value = {"stream": []} - # Messages with citationsContent using tagged union CitationLocation structure + # Messages with citationsContent using all tagged union CitationLocation types messages = [ - {"role": "user", "content": [{"text": "Analyze this document"}]}, + {"role": "user", "content": [{"text": "Analyze multiple sources"}]}, { "role": "assistant", "content": [ @@ -2104,8 +2105,34 @@ async def test_citations_content_preserves_tagged_union_structure(bedrock_client "sourceContent": [{"text": "Vacation policy allows 15 days per year"}], "title": "Vacation Policy", }, + { + "location": {"documentChunk": {"documentIndex": 1, "start": 5, "end": 8}}, + "sourceContent": [{"text": "Company culture emphasizes work-life balance"}], + "title": "Culture Section", + }, + { + "location": { + "searchResultLocation": { + "searchResultIndex": 0, + "start": 25, + "end": 150, + } + }, + "sourceContent": [{"text": "Search results show industry best practices"}], + "title": "Search Results", + }, + { + "location": { + "web": { + "url": "https://example.com/hr-policies", + "domain": "example.com", + } + }, + "sourceContent": [{"text": "External HR policy guidelines"}], + "title": "External Reference", + }, ], - "content": [{"text": "Based on the document, employees receive comprehensive benefits."}], + "content": [{"text": "Based on multiple sources, the company offers comprehensive benefits."}], } } ], @@ -2123,7 +2150,7 @@ async def test_citations_content_preserves_tagged_union_structure(bedrock_client formatted_messages = call_args["messages"] citations_content = formatted_messages[1]["content"][0]["citationsContent"] - # Verify the tagged union structure is preserved + # Verify the tagged union structure is preserved for all location types expected_citations = [ { "location": {"documentChar": {"documentIndex": 0, "start": 150, "end": 300}}, @@ -2135,6 +2162,32 @@ async def test_citations_content_preserves_tagged_union_structure(bedrock_client "sourceContent": [{"text": "Vacation policy allows 15 days per year"}], "title": "Vacation Policy", }, + { + "location": {"documentChunk": {"documentIndex": 1, "start": 5, "end": 8}}, + "sourceContent": [{"text": "Company culture emphasizes work-life balance"}], + "title": "Culture Section", + }, + { + "location": { + "searchResultLocation": { + "searchResultIndex": 0, + "start": 25, + "end": 150, + } + }, + "sourceContent": [{"text": "Search results show industry best practices"}], + "title": "Search Results", + }, + { + "location": { + "web": { + "url": "https://example.com/hr-policies", + "domain": "example.com", + } + }, + "sourceContent": [{"text": "External HR policy guidelines"}], + "title": "External Reference", + }, ] assert citations_content["citations"] == expected_citations, (