diff --git a/docs/my-website/docs/files_endpoints.md b/docs/my-website/docs/files_endpoints.md index fc0484e921..30677c748a 100644 --- a/docs/my-website/docs/files_endpoints.md +++ b/docs/my-website/docs/files_endpoints.md @@ -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) +``` + @@ -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) diff --git a/docs/my-website/docs/providers/bedrock_batches.md b/docs/my-website/docs/providers/bedrock_batches.md index a1116f4107..19446fda83 100644 --- a/docs/my-website/docs/providers/bedrock_batches.md +++ b/docs/my-website/docs/providers/bedrock_batches.md @@ -172,6 +172,97 @@ curl http://localhost:4000/v1/batches \ +### 4. Retrieve batch results + +Once the batch job is completed, download the results from S3: + + + + +```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', {})}") +``` + + + + +```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 +``` + + + + +```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) +``` + + + + +**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? diff --git a/litellm/files/main.py b/litellm/files/main.py index 71139001e5..9378715a47 100644 --- a/litellm/files/main.py +++ b/litellm/files/main.py @@ -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", diff --git a/litellm/llms/bedrock/files/handler.py b/litellm/llms/bedrock/files/handler.py new file mode 100644 index 0000000000..d6177e090d --- /dev/null +++ b/litellm/llms/bedrock/files/handler.py @@ -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, + ) + ) + diff --git a/tests/test_litellm/llms/bedrock/files/test_bedrock_files_integration.py b/tests/test_litellm/llms/bedrock/files/test_bedrock_files_integration.py new file mode 100644 index 0000000000..37a0daa1d5 --- /dev/null +++ b/tests/test_litellm/llms/bedrock/files/test_bedrock_files_integration.py @@ -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