#3621 - Pass s3:// file URLs directly to API in BedrockConverseModel

Author: Motta Kin

pydantic_ai_slim/pydantic_ai/models/bedrock.py

@@ -643,20 +643,25 @@ async def _map_user_prompt(part: UserPromptPart, document_count: Iterator[int])
                     else:
                         raise NotImplementedError('Binary content is not supported yet.')
                 elif isinstance(item, ImageUrl | DocumentUrl | VideoUrl):
-                    downloaded_item = await download_item(item, data_format='bytes', type_format='extension')
-                    format = downloaded_item['data_type']
+                    source: dict[str, Any]
+                    if item.url.startswith('s3://'):
+                        source = {'s3Location': {'uri': item.url}}
+                    else:
+                        downloaded_item = await download_item(item, data_format='bytes', type_format='extension')
+                        source = {'bytes': downloaded_item['data']}
+
                     if item.kind == 'image-url':
                         format = item.media_type.split('/')[1]
                         assert format in ('jpeg', 'png', 'gif', 'webp'), f'Unsupported image format: {format}'
-                        image: ImageBlockTypeDef = {'format': format, 'source': {'bytes': downloaded_item['data']}}
+                        image: ImageBlockTypeDef = {'format': format, 'source': cast(Any, source)}
                         content.append({'image': image})
 
                     elif item.kind == 'document-url':
                         name = f'Document {next(document_count)}'
                         document: DocumentBlockTypeDef = {
                             'name': name,
                             'format': item.format,
-                            'source': {'bytes': downloaded_item['data']},
+                            'source': cast(Any, source),
                         }
                         content.append({'document': document})
 
@@ -673,7 +678,7 @@ async def _map_user_prompt(part: UserPromptPart, document_count: Iterator[int])
                             'wmv',
                             'three_gp',
                         ), f'Unsupported video format: {format}'
-                        video: VideoBlockTypeDef = {'format': format, 'source': {'bytes': downloaded_item['data']}}
+                        video: VideoBlockTypeDef = {'format': format, 'source': cast(Any, source)}
                         content.append({'video': video})
                 elif isinstance(item, AudioUrl):  # pragma: no cover
                     raise NotImplementedError('Audio is not supported yet.')

tests/models/cassettes/test_bedrock/test_s3_document_url_input.yaml

@@ -0,0 +1,42 @@
+interactions:
+- request:
+    body: '{"messages": [{"role": "user", "content": [{"text": "What is the main content on this document?"}, {"document": {"format": "pdf", "name": "test-doc.pdf", "source": {"s3Location": {"uri": "s3://my-bucket/documents/test-doc.pdf"}}}}]}], "system": [{"text": "You are a helpful chatbot."}], "inferenceConfig": {}}'
+    headers:
+      amz-sdk-invocation-id:
+      - !!binary |
+        ZGQxNWI1ODItMTk4Yy00NWZhLTllZjYtODFlY2IzZmUxNWM2
+      amz-sdk-request:
+      - !!binary |
+        YXR0ZW1wdD0x
+      content-length:
+      - '280'
+      content-type:
+      - !!binary |
+        YXBwbGljYXRpb24vanNvbg==
+    method: POST
+    uri: https://bedrock-runtime.us-east-1.amazonaws.com/model/anthropic.claude-v2/converse
+  response:
+    headers:
+      connection:
+      - keep-alive
+      content-length:
+      - '420'
+      content-type:
+      - application/json
+    parsed_body:
+      metrics:
+        latencyMs: 600
+      output:
+        message:
+          content:
+          - text: Based on the provided document, the main content discusses best practices for cloud storage and data management.
+          role: assistant
+      stopReason: end_turn
+      usage:
+        inputTokens: 35
+        outputTokens: 18
+        totalTokens: 53
+    status:
+      code: 200
+      message: OK
+version: 1

tests/models/cassettes/test_bedrock/test_s3_image_url_input.yaml

@@ -0,0 +1,42 @@
+interactions:
+- request:
+    body: '{"messages": [{"role": "user", "content": [{"text": "What is in this image?"}, {"image": {"format": "jpeg", "source": {"s3Location": {"uri": "s3://my-bucket/images/test-image.jpg"}}}}]}], "system": [{"text": "You are a helpful chatbot."}], "inferenceConfig": {}}'
+    headers:
+      amz-sdk-invocation-id:
+      - !!binary |
+        ZGQxNWI1ODItMTk4Yy00NWZhLTllZjYtODFlY2IzZmUxNWM2
+      amz-sdk-request:
+      - !!binary |
+        YXR0ZW1wdD0x
+      content-length:
+      - '250'
+      content-type:
+      - !!binary |
+        YXBwbGljYXRpb24vanNvbg==
+    method: POST
+    uri: https://bedrock-runtime.us-east-1.amazonaws.com/model/us.amazon.nova-pro-v1%3A0/converse
+  response:
+    headers:
+      connection:
+      - keep-alive
+      content-length:
+      - '400'
+      content-type:
+      - application/json
+    parsed_body:
+      metrics:
+        latencyMs: 450
+      output:
+        message:
+          content:
+          - text: The image shows a scenic landscape with mountains in the background and a clear blue sky above.
+          role: assistant
+      stopReason: end_turn
+      usage:
+        inputTokens: 25
+        outputTokens: 20
+        totalTokens: 45
+    status:
+      code: 200
+      message: OK
+version: 1

tests/models/cassettes/test_bedrock/test_s3_video_url_input.yaml

@@ -0,0 +1,42 @@
+interactions:
+- request:
+    body: '{"messages": [{"role": "user", "content": [{"text": "Describe this video"}, {"video": {"format": "mp4", "source": {"s3Location": {"uri": "s3://my-bucket/videos/test-video.mp4"}}}}]}], "system": [{"text": "You are a helpful chatbot."}], "inferenceConfig": {}}'
+    headers:
+      amz-sdk-invocation-id:
+      - !!binary |
+        ZGQxNWI1ODItMTk4Yy00NWZhLTllZjYtODFlY2IzZmUxNWM2
+      amz-sdk-request:
+      - !!binary |
+        YXR0ZW1wdD0x
+      content-length:
+      - '250'
+      content-type:
+      - !!binary |
+        YXBwbGljYXRpb24vanNvbg==
+    method: POST
+    uri: https://bedrock-runtime.us-east-1.amazonaws.com/model/us.amazon.nova-pro-v1%3A0/converse
+  response:
+    headers:
+      connection:
+      - keep-alive
+      content-length:
+      - '400'
+      content-type:
+      - application/json
+    parsed_body:
+      metrics:
+        latencyMs: 550
+      output:
+        message:
+          content:
+          - text: The video shows a time-lapse of a sunset over the ocean with waves gently rolling onto the shore.
+          role: assistant
+      stopReason: end_turn
+      usage:
+        inputTokens: 30
+        outputTokens: 22
+        totalTokens: 52
+    status:
+      code: 200
+      message: OK
+version: 1

tests/models/test_bedrock.py

@@ -729,6 +729,45 @@ async def test_text_document_url_input(allow_model_requests: None, bedrock_provi
 """)
 
 
+@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.'
+    )
+
+
 @pytest.mark.vcr()
 async def test_text_as_binary_content_input(allow_model_requests: None, bedrock_provider: BedrockProvider):
     m = BedrockConverseModel('us.amazon.nova-pro-v1:0', provider=bedrock_provider)