import os import time from typing import Any, Dict, List, Literal, Optional, Union, cast from httpx import Headers, Response from litellm.llms.base_llm.batches.transformation import BaseBatchesConfig from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.secret_managers.main import get_secret_str from litellm.types.llms.bedrock import ( BedrockCreateBatchRequest, BedrockCreateBatchResponse, BedrockInputDataConfig, BedrockOutputDataConfig, BedrockS3InputDataConfig, BedrockS3OutputDataConfig, ) from litellm.types.llms.openai import ( AllMessageValues, CreateBatchRequest, ) from litellm.types.utils import LiteLLMBatch, LlmProviders from ..base_aws_llm import BaseAWSLLM from ..common_utils import CommonBatchFilesUtils class BedrockBatchesConfig(BaseAWSLLM, BaseBatchesConfig): """ Config for Bedrock Batches - handles batch job creation and management for Bedrock """ def __init__(self): super().__init__() self.common_utils = CommonBatchFilesUtils() @property def custom_llm_provider(self) -> LlmProviders: return LlmProviders.BEDROCK def validate_environment( self, headers: dict, model: str, messages: List[AllMessageValues], optional_params: dict, litellm_params: dict, api_key: Optional[str] = None, api_base: Optional[str] = None, ) -> dict: """ Validate and prepare environment for Bedrock batch requests. AWS credentials are handled by BaseAWSLLM. """ # Add any Bedrock-specific headers if needed return headers def get_complete_batch_url( self, api_base: Optional[str], api_key: Optional[str], model: str, optional_params: Dict, litellm_params: Dict, data: CreateBatchRequest, ) -> str: """ Get the complete URL for Bedrock batch creation. Bedrock batch jobs are created via the model invocation job API. """ aws_region_name = self._get_aws_region_name(optional_params, model) # Bedrock model invocation job endpoint # Format: https://bedrock.{region}.amazonaws.com/model-invocation-job bedrock_endpoint = f"https://bedrock.{aws_region_name}.amazonaws.com/model-invocation-job" return bedrock_endpoint def transform_create_batch_request( self, model: str, create_batch_data: CreateBatchRequest, optional_params: dict, litellm_params: dict, ) -> Dict[str, Any]: """ Transform the batch creation request to Bedrock format. Bedrock batch inference requires: - modelId: The Bedrock model ID - jobName: Unique name for the batch job - inputDataConfig: Configuration for input data (S3 location) - outputDataConfig: Configuration for output data (S3 location) - roleArn: IAM role ARN for the batch job """ # Get required parameters input_file_id = create_batch_data.get("input_file_id") if not input_file_id: raise ValueError("input_file_id is required for Bedrock batch creation") # Extract S3 information from file ID using common utility input_bucket, input_key = self.common_utils.parse_s3_uri(input_file_id) # Get output S3 configuration output_bucket = litellm_params.get("s3_output_bucket_name") or os.getenv("AWS_S3_OUTPUT_BUCKET_NAME") if not output_bucket: # Use same bucket as input if no output bucket specified output_bucket = input_bucket # Get IAM role ARN role_arn = ( litellm_params.get("aws_batch_role_arn") or optional_params.get("aws_batch_role_arn") or os.getenv("AWS_BATCH_ROLE_ARN") ) if not role_arn: raise ValueError( "AWS IAM role ARN is required for Bedrock batch jobs. " "Set 'aws_batch_role_arn' in litellm_params or AWS_BATCH_ROLE_ARN env var" ) if not model: raise ValueError("Could not determine Bedrock model ID. Please pass `model` in your request body.") # Generate job name with the correct model ID using common utility job_name = self.common_utils.generate_unique_job_name(model, prefix="litellm") output_key = f"litellm-batch-outputs/{job_name}/" # Build input data config input_data_config: BedrockInputDataConfig = { "s3InputDataConfig": BedrockS3InputDataConfig( s3Uri=f"s3://{input_bucket}/{input_key}" ) } # Build output data config s3_output_config: BedrockS3OutputDataConfig = BedrockS3OutputDataConfig( s3Uri=f"s3://{output_bucket}/{output_key}" ) # Add optional KMS encryption key ID if provided s3_encryption_key_id = ( litellm_params.get("s3_encryption_key_id") or get_secret_str("AWS_S3_ENCRYPTION_KEY_ID") ) if s3_encryption_key_id: s3_output_config["s3EncryptionKeyId"] = s3_encryption_key_id output_data_config: BedrockOutputDataConfig = { "s3OutputDataConfig": s3_output_config } # Create Bedrock batch request with proper typing bedrock_request: BedrockCreateBatchRequest = { "modelId": model, "jobName": job_name, "inputDataConfig": input_data_config, "outputDataConfig": output_data_config, "roleArn": role_arn } # Add optional parameters if provided completion_window = create_batch_data.get("completion_window") if completion_window: # Map OpenAI completion window to Bedrock timeout # OpenAI uses "24h", Bedrock expects timeout in hours if completion_window == "24h": bedrock_request["timeoutDurationInHours"] = 24 # For Bedrock, we need to return a pre-signed request with AWS auth headers # Use common utility for AWS signing endpoint_url = f"https://bedrock.{self._get_aws_region_name(optional_params, model)}.amazonaws.com/model-invocation-job" signed_headers, signed_data = self.common_utils.sign_aws_request( service_name="bedrock", data=bedrock_request, endpoint_url=endpoint_url, optional_params=optional_params, method="POST" ) # Return a pre-signed request format that the HTTP handler can use return { "method": "POST", "url": endpoint_url, "headers": signed_headers, "data": signed_data.decode('utf-8') } def transform_create_batch_response( self, model: Optional[str], raw_response: Response, logging_obj: Any, litellm_params: dict, ) -> LiteLLMBatch: """ Transform Bedrock batch creation response to LiteLLM format. """ try: response_data: BedrockCreateBatchResponse = raw_response.json() except Exception as e: raise ValueError(f"Failed to parse Bedrock batch response: {e}") # Extract information from typed Bedrock response job_arn = response_data.get("jobArn", "") status_str: str = str(response_data.get("status", "Submitted")) # Map Bedrock status to OpenAI-compatible status status_mapping: Dict[str, str] = { "Submitted": "validating", "Validating": "validating", "Scheduled": "in_progress", "InProgress": "in_progress", "PartiallyCompleted": "completed", "Completed": "completed", "Failed": "failed", "Stopping": "cancelling", "Stopped": "cancelled", "Expired": "expired", } openai_status = cast(Literal["validating", "failed", "in_progress", "finalizing", "completed", "expired", "cancelling", "cancelled"], status_mapping.get(status_str, "validating")) # Get original request data from litellm_params if available original_request = litellm_params.get("original_batch_request", {}) # Create LiteLLM batch object return LiteLLMBatch( id=job_arn, # Use ARN as the batch ID object="batch", endpoint=original_request.get("endpoint", "/v1/chat/completions"), errors=None, input_file_id=original_request.get("input_file_id", ""), completion_window=original_request.get("completion_window", "24h"), status=openai_status, output_file_id=None, # Will be populated when job completes error_file_id=None, created_at=int(time.time()), in_progress_at=int(time.time()) if status_str == "InProgress" else None, expires_at=None, finalizing_at=None, completed_at=None, failed_at=None, expired_at=None, cancelling_at=None, cancelled_at=None, request_counts=None, metadata=original_request.get("metadata", {}), ) def transform_retrieve_batch_request( self, batch_id: str, optional_params: dict, litellm_params: dict, ) -> Dict[str, Any]: """ Transform batch retrieval request for Bedrock. Args: batch_id: Bedrock job ARN optional_params: Optional parameters litellm_params: LiteLLM parameters Returns: Transformed request data for Bedrock GetModelInvocationJob API """ # For Bedrock, batch_id should be the full job ARN # The GetModelInvocationJob API expects the full ARN as the identifier if not batch_id.startswith("arn:aws:bedrock:"): raise ValueError(f"Invalid batch_id format. Expected ARN, got: {batch_id}") # Extract the job identifier from the ARN - use the full ARN path part # ARN format: arn:aws:bedrock:region:account:model-invocation-job/job-name arn_parts = batch_id.split(":") if len(arn_parts) < 6: raise ValueError(f"Invalid ARN format: {batch_id}") region = arn_parts[3] # arn_parts[5] contains "model-invocation-job/{jobId}" # Build the endpoint URL for GetModelInvocationJob # AWS API format: GET /model-invocation-job/{jobIdentifier} # Use the FULL ARN as jobIdentifier and URL-encode it (includes ':' and '/') import urllib.parse as _ul encoded_arn = _ul.quote(batch_id, safe="") endpoint_url = f"https://bedrock.{region}.amazonaws.com/model-invocation-job/{encoded_arn}" # Use common utility for AWS signing signed_headers, _ = self.common_utils.sign_aws_request( service_name="bedrock", data={}, # GET request has no body endpoint_url=endpoint_url, optional_params=optional_params, method="GET" ) # Return pre-signed request format return { "method": "GET", "url": endpoint_url, "headers": signed_headers, "data": None } def _parse_timestamps_and_status(self, response_data, status_str: str): """Helper to parse timestamps based on status.""" import datetime def parse_timestamp(ts_str: Optional[str]) -> Optional[int]: if not ts_str: return None try: dt = datetime.datetime.fromisoformat(ts_str.replace('Z', '+00:00')) return int(dt.timestamp()) except Exception: return None created_at = parse_timestamp(str(response_data.get("submitTime")) if response_data.get("submitTime") is not None else None) in_progress_states = {"InProgress", "Validating", "Scheduled"} in_progress_at = ( parse_timestamp(str(response_data.get("lastModifiedTime")) if response_data.get("lastModifiedTime") is not None else None) if status_str in in_progress_states else None ) completed_at = parse_timestamp(str(response_data.get("endTime")) if response_data.get("endTime") is not None else None) if status_str in {"Completed", "PartiallyCompleted"} else None failed_at = parse_timestamp(str(response_data.get("endTime")) if response_data.get("endTime") is not None else None) if status_str == "Failed" else None cancelled_at = parse_timestamp(str(response_data.get("endTime")) if response_data.get("endTime") is not None else None) if status_str == "Stopped" else None expires_at = parse_timestamp(str(response_data.get("jobExpirationTime")) if response_data.get("jobExpirationTime") is not None else None) return created_at, in_progress_at, completed_at, failed_at, cancelled_at, expires_at def _extract_file_configs(self, response_data): """Helper to extract input and output file configurations.""" # Extract input file ID input_file_id = "" input_data_config = response_data.get("inputDataConfig", {}) if isinstance(input_data_config, dict): s3_input_config = input_data_config.get("s3InputDataConfig", {}) if isinstance(s3_input_config, dict): input_file_id = s3_input_config.get("s3Uri", "") # Extract output file ID output_file_id = None output_data_config = response_data.get("outputDataConfig", {}) if isinstance(output_data_config, dict): s3_output_config = output_data_config.get("s3OutputDataConfig", {}) if isinstance(s3_output_config, dict): output_file_id = s3_output_config.get("s3Uri", "") return input_file_id, output_file_id def _extract_errors_and_metadata(self, response_data, raw_response): """Helper to extract errors and enriched metadata.""" # Extract errors message = response_data.get("message") errors = None if message: from openai.types.batch import Errors from openai.types.batch_error import BatchError errors = Errors( data=[BatchError(message=message, code=str(raw_response.status_code))], object="list" ) # Enrich metadata with useful Bedrock fields enriched_metadata_raw: Dict[str, Any] = { "jobName": response_data.get("jobName"), "clientRequestToken": response_data.get("clientRequestToken"), "modelId": response_data.get("modelId"), "roleArn": response_data.get("roleArn"), "timeoutDurationInHours": response_data.get("timeoutDurationInHours"), "vpcConfig": response_data.get("vpcConfig"), } import json as _json enriched_metadata: Dict[str, str] = {} for _k, _v in enriched_metadata_raw.items(): if _v is None: continue if isinstance(_v, (dict, list)): try: enriched_metadata[_k] = _json.dumps(_v) except Exception: enriched_metadata[_k] = str(_v) else: enriched_metadata[_k] = str(_v) return errors, enriched_metadata def transform_retrieve_batch_response( self, model: Optional[str], raw_response: Response, logging_obj: Any, litellm_params: dict, ) -> LiteLLMBatch: """ Transform Bedrock batch retrieval response to LiteLLM format. """ from litellm.types.llms.bedrock import BedrockGetBatchResponse try: response_data: BedrockGetBatchResponse = raw_response.json() except Exception as e: raise ValueError(f"Failed to parse Bedrock batch response: {e}") job_arn = response_data.get("jobArn", "") status_str: str = str(response_data.get("status", "Submitted")) # Map Bedrock status to OpenAI-compatible status status_mapping: Dict[str, str] = { "Submitted": "validating", "Validating": "validating", "Scheduled": "in_progress", "InProgress": "in_progress", "PartiallyCompleted": "completed", "Completed": "completed", "Failed": "failed", "Stopping": "cancelling", "Stopped": "cancelled", "Expired": "expired" } openai_status = cast(Literal["validating", "failed", "in_progress", "finalizing", "completed", "expired", "cancelling", "cancelled"], status_mapping.get(status_str, "validating")) # Parse timestamps created_at, in_progress_at, completed_at, failed_at, cancelled_at, expires_at = self._parse_timestamps_and_status(response_data, status_str) # Extract file configurations input_file_id, output_file_id = self._extract_file_configs(response_data) # Extract errors and metadata errors, enriched_metadata = self._extract_errors_and_metadata(response_data, raw_response) return LiteLLMBatch( id=job_arn, object="batch", endpoint="/v1/chat/completions", errors=errors, input_file_id=input_file_id, completion_window="24h", status=openai_status, output_file_id=output_file_id, error_file_id=None, created_at=created_at or int(time.time()), in_progress_at=in_progress_at, expires_at=expires_at, finalizing_at=None, completed_at=completed_at, failed_at=failed_at, expired_at=None, cancelling_at=None, cancelled_at=cancelled_at, request_counts=None, metadata=enriched_metadata, ) def get_error_class( self, error_message: str, status_code: int, headers: Union[Dict, Headers] ) -> BaseLLMException: """ Get Bedrock-specific error class using common utility. """ return self.common_utils.get_error_class(error_message, status_code, headers)