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

[ericfzhu]

Description

When using a Bedrock model with citations enabled, the returned CitationLocation object inside of Citation is a UNION, which means that currently it must be one of documentChar, documentChunk, documentPage, searchResultLocation, and web.

Changing _format_request_message_content would cause the following error to surface:

ValidationException: An error occurred (ValidationException) when calling the ConverseStream operation: The model returned the following errors: messages.1.content.2.text.text: Field required

As handle_content_block_stop from streaming.py is not formatting the chunks correctly when citations are used

Related Issues

[BUG] Bedrock Runtime API CitationLocation Tagged Union Structure Error #1323

Documentation PR

https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_ContentBlockDelta.html

Type of Change

Bug fix

Testing

How have you tested the change? Verify that the changes do not break functionality or introduce warnings in consuming repositories: agents-docs, agents-tools, agents-cli

• I ran hatch run prepare

Checklist

• I have read the CONTRIBUTING document
• I have added any necessary tests that prove my fix is effective or my feature works
• I have updated the documentation accordingly
• I have added an appropriate example to the documentation to outline the feature, or no new docs are needed
• My changes generate no new warnings
• Any dependent changes have been merged and published

---

By submitting this pull request, I confirm that you can use, modify, copy, and redistribute this contribution, under the terms of your choice.

[dbschmigelski]

Adding a bit more light for future reviewers.

Multi-turn conversations which received CitationContentBlocks do not work. The call fails with

botocore.errorfactory.ValidationException: An error occurred (ValidationException) when calling the ConverseStream operation: The model returned the following errors: messages.1.content.2.text.text: Field required

The error message is a little ambiguous but the problem is that in a message array looking like

{
    'role': 'user',
    'content': [{...}]
},
 {
    'role': 'assistant',
    'content': [{
        'text': "Based on the document, Amazon views artificial intelligence as transformative & is investing heavily in it. Here are the key points about AI from the letter:\n\n## AI's Transformative Potential\n\n"
    }, {
        'text': "Generative AI is going to reinvent virtually every customer experience we know, and enable altogether new ones about which we've only fantasized."
    }, {
        'citationsContent': {
            'citations': [{
                'title': 'letter to shareholders',
                'sourceContent': [{
                    'text': 'Generative AI is going to reinvent virtually every customer experience we know, and enable altogether new\r\nones about which we’ve only fantasized. '
                }],
                'location': {
                    'documentPage': {
                        'documentIndex': 0,
                        'start': 5,
                        'end': 6
                    }
                }
            }]
        }
    }, 
    {
        'text': ' The CEO emphasizes that '
    }, 
    ...,
  {
    'role': 'user',
    'content': [{...}]
},

the citationsContent message needs to contain both a citations and a content key if we want to include it.

This gives us two options

1. We can filter the citationsContent from the message before sending it back, since the generated content is already reflected in the text. block immediately above it
2. We can include the content and limit manipulation by returning the content block as it was intended.

I do not want to continue with option 1 for forward looking reasons. Right now, citationsContent only contains text, but that can change in the future according to https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_CitationGeneratedContent.html. I do not want to continue the practice of extracting the cited content blocks with additional logic when we can simply return the blocks as they were received.

Option 2 requires that we pass the content back. We can do this by in streaming.py conditioning on the presence of citations then not writing the separate text block. This prevents us from double writing: once as a text block and once as a generated citation block.

This implementation will break, or rather not automatically work, when new shapes are added to the GeneratedContent Union. This is acceptable for now.

---

Multi turn tests have been added for both the streaming and non streaming paths.

---

This change also proposes a breaking change to the CitationStreamEvent to remove an unnecessary nesting with the "callback" key. Given that citations were broken for multi-turn. This breaking change seems like it can be made freely.

---

The key decision points here are

• How to handle the disambiguation of the CitationLocation. The proposal here is to use the ToolChoice pattern
• CitationStreamEvent breaking change
• To drop citations in the request or keep them
• If we want a long term solution regarding non-text citations

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!

[dbschmigelski]

note, this is a breaking change as well. However, as noted citations is a broken feature at the moment.