Repository: pydantic/pydantic-ai
File: tests/models/test_bedrock.py (Python)

Code before:
    agent = Agent(m, system_prompt='You are a helpful chatbot.')

    text_document_url = DocumentUrl(url='https://example-files.online-convert.com/document/txt/example.txt')

    result = await agent.run(['What is the main content on this document?', text_document_url])
    assert result.output == snapshot("""\
Based on the text in the <document_content> tag, the main content of this document appears to be:

An example text describing the use of "John Doe" as a placeholder name in legal cases, hospitals, and other contexts where a party's real identity is unknown or needs to be withheld. It provides background on how "John Doe" and "Jane Doe" are commonly used in the United States and Canada for this purpose, in contrast to other English speaking countries that use names like "Joe Bloggs". The text gives examples of using John/Jane Doe for legal cases, unidentified corpses, and as generic names on forms. It also mentions how "Baby Doe" and "Precious Doe" are used for unidentified children.\
""")


@pytest.mark.vcr()
async def test_s3_image_url_input(allow_model_requests: None, bedrock_provider: BedrockProvider):
    """Test that s3:// image URLs are passed directly to Bedrock API without downloading."""
    m = BedrockConverseModel('us.amazon.nova-pro-v1:0', provider=bedrock_provider)
    agent = Agent(m, system_prompt='You are a helpful chatbot.')
    image_url = ImageUrl(url='s3://my-bucket/images/test-image.jpg', media_type='image/jpeg')

    result = await agent.run(['What is in this image?', image_url])
    assert result.output == snapshot(
        'The image shows a scenic landscape with mountains in the background and a clear blue sky above.'
    )


@pytest.mark.vcr()
async def test_s3_video_url_input(allow_model_requests: None, bedrock_provider: BedrockProvider):
    """Test that s3:// video URLs are passed directly to Bedrock API."""
    m = BedrockConverseModel('us.amazon.nova-pro-v1:0', provider=bedrock_provider)
    agent = Agent(m, system_prompt='You are a helpful chatbot.')
    video_url = VideoUrl(url='s3://my-bucket/videos/test-video.mp4', media_type='video/mp4')

    result = await agent.run(['Describe this video', video_url])
    assert result.output == snapshot(
        'The video shows a time-lapse of a sunset over the ocean with waves gently rolling onto the shore.'
    )


@pytest.mark.vcr()
async def test_s3_document_url_input(allow_model_requests: None, bedrock_provider: BedrockProvider):
    """Test that s3:// document URLs are passed directly to Bedrock API."""
    m = BedrockConverseModel('anthropic.claude-v2', provider=bedrock_provider)
    agent = Agent(m, system_prompt='You are a helpful chatbot.')
    document_url = DocumentUrl(url='s3://my-bucket/documents/test-doc.pdf', media_type='application/pdf')

    result = await agent.run(['What is the main content on this document?', document_url])
    assert result.output == snapshot(
        'Based on the provided document, the main content discusses best practices for cloud storage and data management.'
    )



Review comment (suggestion):
I'm guessing these are not real files in S3 and the cassettes were manually modified? That'll get us into trouble if we re-record the cassettes :) So I'd rather either use a _real_ S3 URL (may be hard), _or_ directly test the result of `_map_messages` as we do in other tests already

Code after:
    assert result.output == snapshot(
        'Based on the provided XML data, the main content of the document is "Dummy PDF file". This is contained in the <document_content> tag for the document with index="1".'
    )


@pytest.mark.vcr()
async def test_text_document_url_input(allow_model_requests: None, bedrock_provider: BedrockProvider):
    m = BedrockConverseModel('anthropic.claude-v2', provider=bedrock_provider)
    agent = Agent(m, system_prompt='You are a helpful chatbot.')

    text_document_url = DocumentUrl(url='https://example-files.online-convert.com/document/txt/example.txt')

    result = await agent.run(['What is the main content on this document?', text_document_url])
    assert result.output == snapshot(
        """\
Based on the text in the <document_content> tag, the main content of this document appears to be:

An example text describing the use of "John Doe" as a placeholder name in legal cases, hospitals, and other contexts where a party's real identity is unknown or needs to be withheld. It provides background on how "John Doe" and "Jane Doe" are commonly used in the United States and Canada for this purpose, in contrast to other English speaking countries that use names like "Joe Bloggs". The text gives examples of using John/Jane Doe for legal cases, unidentified corpses, and as generic names on forms. It also mentions how "Baby Doe" and "Precious Doe" are used for unidentified children.\
"""
    )


async def test_s3_image_url_input(bedrock_provider: BedrockProvider):
    """Test that s3:// image URLs are passed directly to Bedrock API without downloading."""
    model = BedrockConverseModel('us.amazon.nova-pro-v1:0', provider=bedrock_provider)
    image_url = ImageUrl(url='s3://my-bucket/images/test-image.jpg', media_type='image/jpeg')

    req = [
        ModelRequest(parts=[UserPromptPart(content=['What is in this image?', image_url])]),
    ]

    _, bedrock_messages = await model._map_messages(req, ModelRequestParameters(), None)  # type: ignore[reportPrivateUsage]

    assert bedrock_messages == snapshot(
        [
            {
                'role': 'user',
                'content': [
                    {'text': 'What is in this image?'},
                    {
                        'image': {
                            'format': 'jpeg',
                            'source': {'s3Location': {'uri': 's3://my-bucket/images/test-image.jpg'}},
                        }
                    },
                ],
            }
        ]
    )