mirror of
https://github.com/tiennm99/litellm.git
synced 2026-08-13 22:23:04 +00:00
Merge pull request #17470 from BerriAI/litellm_batches_bedrock_content
Add support for file content download for bedrock batches
This commit is contained in:
Sameer Kankute
@@ -301,6 +301,17 @@ content = await litellm.afile_content(
|
||||
print("file content=", content)
|
||||
```
|
||||
|
||||
**Get File Content (Bedrock)**
|
||||
```python
|
||||
# For Bedrock batch output files stored in S3
|
||||
content = await litellm.afile_content(
|
||||
file_id="s3://bucket-name/path/to/file.jsonl", # S3 URI or unified file ID
|
||||
custom_llm_provider="bedrock",
|
||||
aws_region_name="us-west-2"
|
||||
)
|
||||
print("file content=", content.text)
|
||||
```
|
||||
|
||||
</TabItem>
|
||||
</Tabs>
|
||||
|
||||
@@ -313,4 +324,6 @@ print("file content=", content)
|
||||
|
||||
### [Vertex AI](./providers/vertex#batch-apis)
|
||||
|
||||
### [Bedrock](./providers/bedrock_batches#4-retrieve-batch-results)
|
||||
|
||||
## [Swagger API Reference](https://litellm-api.up.railway.app/#/files)
|
||||
|
||||
@@ -172,6 +172,97 @@ curl http://localhost:4000/v1/batches \
|
||||
</TabItem>
|
||||
</Tabs>
|
||||
|
||||
### 4. Retrieve batch results
|
||||
|
||||
Once the batch job is completed, download the results from S3:
|
||||
|
||||
<Tabs>
|
||||
<TabItem value="python" label="Python">
|
||||
|
||||
```python showLineNumbers title="bedrock_batch.py"
|
||||
...
|
||||
# Wait for batch completion (check status periodically)
|
||||
batch_status = client.batches.retrieve(batch_id=batch.id)
|
||||
|
||||
if batch_status.status == "completed":
|
||||
# Download the output file
|
||||
result = client.files.content(
|
||||
file_id=batch_status.output_file_id,
|
||||
extra_headers={"custom-llm-provider": "bedrock"}
|
||||
)
|
||||
|
||||
# Save or process the results
|
||||
with open("batch_output.jsonl", "wb") as f:
|
||||
f.write(result.content)
|
||||
|
||||
# Parse JSONL results
|
||||
for line in result.text.strip().split('\n'):
|
||||
record = json.loads(line)
|
||||
print(f"Record ID: {record['recordId']}")
|
||||
print(f"Output: {record.get('modelOutput', {})}")
|
||||
```
|
||||
|
||||
</TabItem>
|
||||
<TabItem value="curl" label="Curl">
|
||||
|
||||
```bash showLineNumbers title="Download Batch Results"
|
||||
# First retrieve batch to get output_file_id
|
||||
curl http://localhost:4000/v1/batches/batch_abc123 \
|
||||
-H "Authorization: Bearer sk-1234"
|
||||
|
||||
# Then download the output file
|
||||
curl http://localhost:4000/v1/files/{output_file_id}/content \
|
||||
-H "Authorization: Bearer sk-1234" \
|
||||
-H "custom-llm-provider: bedrock" \
|
||||
-o batch_output.jsonl
|
||||
```
|
||||
|
||||
</TabItem>
|
||||
<TabItem value="litellm-direct" label="LiteLLM Direct">
|
||||
|
||||
```python showLineNumbers title="bedrock_batch.py"
|
||||
import litellm
|
||||
from litellm import file_content
|
||||
|
||||
# Download using litellm directly (bypasses proxy managed files)
|
||||
result = file_content(
|
||||
file_id=batch_status.output_file_id, # Can be S3 URI or unified file ID
|
||||
custom_llm_provider="bedrock",
|
||||
aws_region_name="us-west-2",
|
||||
)
|
||||
|
||||
# Process results
|
||||
print(result.text)
|
||||
```
|
||||
|
||||
</TabItem>
|
||||
</Tabs>
|
||||
|
||||
**Output Format:**
|
||||
|
||||
The batch output file is in JSONL format with each line containing:
|
||||
|
||||
```json
|
||||
{
|
||||
"recordId": "request-1",
|
||||
"modelInput": {
|
||||
"messages": [...],
|
||||
"max_tokens": 1000
|
||||
},
|
||||
"modelOutput": {
|
||||
"content": [...],
|
||||
"id": "msg_abc123",
|
||||
"model": "claude-3-5-sonnet-20240620-v1:0",
|
||||
"role": "assistant",
|
||||
"stop_reason": "end_turn",
|
||||
"usage": {
|
||||
"input_tokens": 15,
|
||||
"output_tokens": 10
|
||||
}
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
## FAQ
|
||||
|
||||
### Where are my files written?
|
||||
|
||||
+14
-3
@@ -18,6 +18,7 @@ from litellm import get_secret_str
|
||||
from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
from litellm.llms.azure.files.handler import AzureOpenAIFilesAPI
|
||||
from litellm.llms.bedrock.files.handler import BedrockFilesHandler
|
||||
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
|
||||
from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler
|
||||
from litellm.llms.openai.openai import FileDeleted, FileObject, OpenAIFilesAPI
|
||||
@@ -47,6 +48,7 @@ base_llm_http_handler = BaseLLMHTTPHandler()
|
||||
openai_files_instance = OpenAIFilesAPI()
|
||||
azure_files_instance = AzureOpenAIFilesAPI()
|
||||
vertex_ai_files_instance = VertexAIFilesHandler()
|
||||
bedrock_files_instance = BedrockFilesHandler()
|
||||
#################################################
|
||||
|
||||
|
||||
@@ -755,7 +757,7 @@ def file_list(
|
||||
@client
|
||||
async def afile_content(
|
||||
file_id: str,
|
||||
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm"] = "openai",
|
||||
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "bedrock", "hosted_vllm"] = "openai",
|
||||
extra_headers: Optional[Dict[str, str]] = None,
|
||||
extra_body: Optional[Dict[str, str]] = None,
|
||||
**kwargs,
|
||||
@@ -800,7 +802,7 @@ def file_content(
|
||||
file_id: str,
|
||||
model: Optional[str] = None,
|
||||
custom_llm_provider: Optional[
|
||||
Union[Literal["openai", "azure", "vertex_ai", "hosted_vllm"], str]
|
||||
Union[Literal["openai", "azure", "vertex_ai", "bedrock", "hosted_vllm"], str]
|
||||
] = None,
|
||||
extra_headers: Optional[Dict[str, str]] = None,
|
||||
extra_body: Optional[Dict[str, str]] = None,
|
||||
@@ -938,9 +940,18 @@ def file_content(
|
||||
timeout=timeout,
|
||||
max_retries=optional_params.max_retries,
|
||||
)
|
||||
elif custom_llm_provider == "bedrock":
|
||||
response = bedrock_files_instance.file_content(
|
||||
_is_async=_is_async,
|
||||
file_content_request=_file_content_request,
|
||||
api_base=optional_params.api_base,
|
||||
optional_params=litellm_params_dict,
|
||||
timeout=timeout,
|
||||
max_retries=optional_params.max_retries,
|
||||
)
|
||||
else:
|
||||
raise litellm.exceptions.BadRequestError(
|
||||
message="LiteLLM doesn't support {} for 'custom_llm_provider'. Supported providers are 'openai', 'azure', 'vertex_ai'.".format(
|
||||
message="LiteLLM doesn't support {} for 'custom_llm_provider'. Supported providers are 'openai', 'azure', 'vertex_ai', 'bedrock'.".format(
|
||||
custom_llm_provider
|
||||
),
|
||||
model="n/a",
|
||||
|
||||
@@ -0,0 +1,206 @@
|
||||
import asyncio
|
||||
import base64
|
||||
from typing import Any, Coroutine, Optional, Tuple, Union
|
||||
|
||||
import httpx
|
||||
|
||||
from litellm import LlmProviders
|
||||
from litellm.llms.custom_httpx.http_handler import get_async_httpx_client
|
||||
from litellm.types.llms.openai import (
|
||||
FileContentRequest,
|
||||
HttpxBinaryResponseContent,
|
||||
)
|
||||
from litellm.types.utils import SpecialEnums
|
||||
|
||||
from ..base_aws_llm import BaseAWSLLM
|
||||
|
||||
|
||||
class BedrockFilesHandler(BaseAWSLLM):
|
||||
"""
|
||||
Handles downloading files from S3 for Bedrock batch processing.
|
||||
|
||||
This implementation downloads files from S3 buckets where Bedrock
|
||||
stores batch output files.
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.async_httpx_client = get_async_httpx_client(
|
||||
llm_provider=LlmProviders.BEDROCK,
|
||||
)
|
||||
|
||||
def _extract_s3_uri_from_file_id(self, file_id: str) -> str:
|
||||
"""
|
||||
Extract S3 URI from encoded file ID.
|
||||
|
||||
The file ID can be in two formats:
|
||||
1. Base64-encoded unified file ID containing: llm_output_file_id,s3://bucket/path
|
||||
2. Direct S3 URI: s3://bucket/path
|
||||
|
||||
Args:
|
||||
file_id: Encoded file ID or direct S3 URI
|
||||
|
||||
Returns:
|
||||
S3 URI (e.g., "s3://bucket-name/path/to/file")
|
||||
"""
|
||||
# First, try to decode if it's a base64-encoded unified file ID
|
||||
try:
|
||||
# Add padding if needed
|
||||
padded = file_id + "=" * (-len(file_id) % 4)
|
||||
decoded = base64.urlsafe_b64decode(padded).decode()
|
||||
|
||||
# Check if it's a unified file ID format
|
||||
if decoded.startswith(SpecialEnums.LITELM_MANAGED_FILE_ID_PREFIX.value):
|
||||
# Extract llm_output_file_id from the decoded string
|
||||
if "llm_output_file_id," in decoded:
|
||||
s3_uri = decoded.split("llm_output_file_id,")[1].split(";")[0]
|
||||
return s3_uri
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# If not base64 encoded or doesn't contain llm_output_file_id, assume it's already an S3 URI
|
||||
if file_id.startswith("s3://"):
|
||||
return file_id
|
||||
|
||||
# If it doesn't start with s3://, assume it's a direct S3 URI and add the prefix
|
||||
return f"s3://{file_id}"
|
||||
|
||||
def _parse_s3_uri(self, s3_uri: str) -> Tuple[str, str]:
|
||||
"""
|
||||
Parse S3 URI to extract bucket name and object key.
|
||||
|
||||
Args:
|
||||
s3_uri: S3 URI (e.g., "s3://bucket-name/path/to/file")
|
||||
|
||||
Returns:
|
||||
Tuple of (bucket_name, object_key)
|
||||
"""
|
||||
if not s3_uri.startswith("s3://"):
|
||||
raise ValueError(f"Invalid S3 URI format: {s3_uri}. Expected format: s3://bucket-name/path/to/file")
|
||||
|
||||
# Remove 's3://' prefix
|
||||
path = s3_uri[5:]
|
||||
|
||||
if "/" in path:
|
||||
bucket_name, object_key = path.split("/", 1)
|
||||
else:
|
||||
bucket_name = path
|
||||
object_key = ""
|
||||
|
||||
return bucket_name, object_key
|
||||
|
||||
async def afile_content(
|
||||
self,
|
||||
file_content_request: FileContentRequest,
|
||||
optional_params: dict,
|
||||
timeout: Union[float, httpx.Timeout],
|
||||
max_retries: Optional[int],
|
||||
) -> HttpxBinaryResponseContent:
|
||||
"""
|
||||
Download file content from S3 bucket for Bedrock files.
|
||||
|
||||
Args:
|
||||
file_content_request: Contains file_id (encoded or S3 URI)
|
||||
optional_params: Optional parameters containing AWS credentials
|
||||
timeout: Request timeout
|
||||
max_retries: Max retry attempts
|
||||
|
||||
Returns:
|
||||
HttpxBinaryResponseContent: Binary content wrapped in compatible response format
|
||||
"""
|
||||
import boto3
|
||||
from botocore.credentials import Credentials
|
||||
|
||||
file_id = file_content_request.get("file_id")
|
||||
if not file_id:
|
||||
raise ValueError("file_id is required in file_content_request")
|
||||
|
||||
# Extract S3 URI from file ID
|
||||
s3_uri = self._extract_s3_uri_from_file_id(file_id)
|
||||
bucket_name, object_key = self._parse_s3_uri(s3_uri)
|
||||
|
||||
# Get AWS credentials
|
||||
aws_region_name = self._get_aws_region_name(
|
||||
optional_params=optional_params, model=""
|
||||
)
|
||||
credentials: Credentials = self.get_credentials(
|
||||
aws_access_key_id=optional_params.get("aws_access_key_id"),
|
||||
aws_secret_access_key=optional_params.get("aws_secret_access_key"),
|
||||
aws_session_token=optional_params.get("aws_session_token"),
|
||||
aws_region_name=aws_region_name,
|
||||
aws_session_name=optional_params.get("aws_session_name"),
|
||||
aws_profile_name=optional_params.get("aws_profile_name"),
|
||||
aws_role_name=optional_params.get("aws_role_name"),
|
||||
aws_web_identity_token=optional_params.get("aws_web_identity_token"),
|
||||
aws_sts_endpoint=optional_params.get("aws_sts_endpoint"),
|
||||
)
|
||||
|
||||
# Create S3 client
|
||||
s3_client = boto3.client(
|
||||
"s3",
|
||||
aws_access_key_id=credentials.access_key,
|
||||
aws_secret_access_key=credentials.secret_key,
|
||||
aws_session_token=credentials.token,
|
||||
region_name=aws_region_name,
|
||||
)
|
||||
|
||||
# Download file from S3
|
||||
try:
|
||||
response = s3_client.get_object(Bucket=bucket_name, Key=object_key)
|
||||
file_content = response["Body"].read()
|
||||
except Exception as e:
|
||||
raise ValueError(f"Failed to download file from S3: {s3_uri}. Error: {str(e)}")
|
||||
|
||||
# Create mock HTTP response
|
||||
mock_response = httpx.Response(
|
||||
status_code=200,
|
||||
content=file_content,
|
||||
headers={"content-type": "application/octet-stream"},
|
||||
request=httpx.Request(method="GET", url=s3_uri),
|
||||
)
|
||||
|
||||
return HttpxBinaryResponseContent(response=mock_response)
|
||||
|
||||
def file_content(
|
||||
self,
|
||||
_is_async: bool,
|
||||
file_content_request: FileContentRequest,
|
||||
api_base: Optional[str],
|
||||
optional_params: dict,
|
||||
timeout: Union[float, httpx.Timeout],
|
||||
max_retries: Optional[int],
|
||||
) -> Union[
|
||||
HttpxBinaryResponseContent, Coroutine[Any, Any, HttpxBinaryResponseContent]
|
||||
]:
|
||||
"""
|
||||
Download file content from S3 bucket for Bedrock files.
|
||||
Supports both sync and async operations.
|
||||
|
||||
Args:
|
||||
_is_async: Whether to run asynchronously
|
||||
file_content_request: Contains file_id (encoded or S3 URI)
|
||||
api_base: API base (unused for S3 operations)
|
||||
optional_params: Optional parameters containing AWS credentials
|
||||
timeout: Request timeout
|
||||
max_retries: Max retry attempts
|
||||
|
||||
Returns:
|
||||
HttpxBinaryResponseContent or Coroutine: Binary content wrapped in compatible response format
|
||||
"""
|
||||
if _is_async:
|
||||
return self.afile_content(
|
||||
file_content_request=file_content_request,
|
||||
optional_params=optional_params,
|
||||
timeout=timeout,
|
||||
max_retries=max_retries,
|
||||
)
|
||||
else:
|
||||
return asyncio.run(
|
||||
self.afile_content(
|
||||
file_content_request=file_content_request,
|
||||
optional_params=optional_params,
|
||||
timeout=timeout,
|
||||
max_retries=max_retries,
|
||||
)
|
||||
)
|
||||
|
||||
@@ -0,0 +1,110 @@
|
||||
"""
|
||||
Test Bedrock files integration with main files API
|
||||
"""
|
||||
|
||||
import base64
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
import litellm
|
||||
from litellm.types.llms.openai import HttpxBinaryResponseContent
|
||||
from litellm.types.utils import SpecialEnums
|
||||
|
||||
|
||||
class TestBedrockFilesIntegration:
|
||||
"""Test integration of Bedrock files with main litellm API"""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_litellm_afile_content_bedrock_provider_with_s3_uri(self):
|
||||
"""Test litellm.afile_content with bedrock provider using direct S3 URI"""
|
||||
file_id = "s3://test-bucket/test-file.jsonl"
|
||||
expected_content = b'{"recordId": "request-1", "modelInput": {}, "modelOutput": {}}'
|
||||
|
||||
# Mock the bedrock_files_instance.file_content method
|
||||
with patch(
|
||||
"litellm.files.main.bedrock_files_instance.file_content",
|
||||
new_callable=AsyncMock,
|
||||
) as mock_file_content:
|
||||
# Create a mock HttpxBinaryResponseContent response
|
||||
import httpx
|
||||
|
||||
mock_response = httpx.Response(
|
||||
status_code=200,
|
||||
content=expected_content,
|
||||
headers={"content-type": "application/octet-stream"},
|
||||
request=httpx.Request(
|
||||
method="GET", url="s3://test-bucket/test-file.jsonl"
|
||||
),
|
||||
)
|
||||
mock_file_content.return_value = HttpxBinaryResponseContent(
|
||||
response=mock_response
|
||||
)
|
||||
|
||||
# Call litellm.afile_content
|
||||
result = await litellm.afile_content(
|
||||
file_id=file_id,
|
||||
custom_llm_provider="bedrock",
|
||||
aws_region_name="us-west-2",
|
||||
)
|
||||
|
||||
# Verify the result
|
||||
assert isinstance(result, HttpxBinaryResponseContent)
|
||||
assert result.response.content == expected_content
|
||||
assert result.response.status_code == 200
|
||||
|
||||
# Verify the mock was called with correct parameters
|
||||
mock_file_content.assert_called_once()
|
||||
call_kwargs = mock_file_content.call_args.kwargs
|
||||
assert call_kwargs["_is_async"] is True
|
||||
assert call_kwargs["file_content_request"]["file_id"] == file_id
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_litellm_afile_content_bedrock_provider_with_unified_file_id(self):
|
||||
"""Test litellm.afile_content with bedrock provider using unified file ID"""
|
||||
# Create a unified file ID
|
||||
s3_uri = "s3://test-bucket/batch-outputs/output.jsonl"
|
||||
unified_id = "test-unified-id-123"
|
||||
model_id = "test-model-id-456"
|
||||
|
||||
unified_file_id_str = f"litellm_proxy:application/json;unified_id,{unified_id};target_model_names,;llm_output_file_id,{s3_uri};llm_output_file_model_id,{model_id}"
|
||||
encoded_file_id = base64.urlsafe_b64encode(unified_file_id_str.encode()).decode().rstrip("=")
|
||||
|
||||
expected_content = b'{"recordId": "request-1", "modelInput": {}, "modelOutput": {}}'
|
||||
|
||||
# Mock the bedrock_files_instance.file_content method
|
||||
with patch(
|
||||
"litellm.files.main.bedrock_files_instance.file_content",
|
||||
new_callable=AsyncMock,
|
||||
) as mock_file_content:
|
||||
# Create a mock HttpxBinaryResponseContent response
|
||||
import httpx
|
||||
|
||||
mock_response = httpx.Response(
|
||||
status_code=200,
|
||||
content=expected_content,
|
||||
headers={"content-type": "application/octet-stream"},
|
||||
request=httpx.Request(method="GET", url=s3_uri),
|
||||
)
|
||||
mock_file_content.return_value = HttpxBinaryResponseContent(
|
||||
response=mock_response
|
||||
)
|
||||
|
||||
# Call litellm.afile_content with unified file ID
|
||||
result = await litellm.afile_content(
|
||||
file_id=encoded_file_id,
|
||||
custom_llm_provider="bedrock",
|
||||
aws_region_name="us-west-2",
|
||||
)
|
||||
|
||||
# Verify the result
|
||||
assert isinstance(result, HttpxBinaryResponseContent)
|
||||
assert result.response.content == expected_content
|
||||
assert result.response.status_code == 200
|
||||
|
||||
# Verify the mock was called - the handler should extract S3 URI from unified file ID
|
||||
mock_file_content.assert_called_once()
|
||||
call_kwargs = mock_file_content.call_args.kwargs
|
||||
assert call_kwargs["_is_async"] is True
|
||||
# The handler extracts S3 URI from the unified file ID
|
||||
assert call_kwargs["file_content_request"]["file_id"] == encoded_file_id
|
||||
Reference in New Issue
Block a user