chore: initial public snapshot for github upload
This commit is contained in:
Your Name
@@ -0,0 +1,549 @@
|
||||
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)
|
||||
Reference in New Issue
Block a user