mirror of
https://github.com/tiennm99/litellm.git
synced 2026-07-25 20:19:51 +00:00
[Feat] Batches - Add bedrock retrieve endpoint support (#14618)
* feat: add bedrock retrieve endpoint * feat: feat: add bedrock retrieve endpoint * test: batches mocked transform * ruff fix * refactor * fix transform * fix: parse_timestamp
This commit is contained in:
Ishaan Jaff
+169
-107
@@ -340,7 +340,7 @@ def create_batch(
|
||||
@client
|
||||
async def aretrieve_batch(
|
||||
batch_id: str,
|
||||
custom_llm_provider: Literal["openai", "azure", "vertex_ai"] = "openai",
|
||||
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "bedrock"] = "openai",
|
||||
metadata: Optional[Dict[str, str]] = None,
|
||||
extra_headers: Optional[Dict[str, str]] = None,
|
||||
extra_body: Optional[Dict[str, str]] = None,
|
||||
@@ -378,11 +378,129 @@ async def aretrieve_batch(
|
||||
except Exception as e:
|
||||
raise e
|
||||
|
||||
def _handle_retrieve_batch_providers_without_provider_config(
|
||||
batch_id: str,
|
||||
optional_params: GenericLiteLLMParams,
|
||||
timeout: Union[float, httpx.Timeout],
|
||||
litellm_params: dict,
|
||||
_retrieve_batch_request: RetrieveBatchRequest,
|
||||
_is_async: bool,
|
||||
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "bedrock"] = "openai",
|
||||
):
|
||||
api_base: Optional[str] = None
|
||||
if custom_llm_provider == "openai":
|
||||
# for deepinfra/perplexity/anyscale/groq we check in get_llm_provider and pass in the api base from there
|
||||
api_base = (
|
||||
optional_params.api_base
|
||||
or litellm.api_base
|
||||
or os.getenv("OPENAI_BASE_URL")
|
||||
or os.getenv("OPENAI_API_BASE")
|
||||
or "https://api.openai.com/v1"
|
||||
)
|
||||
organization = (
|
||||
optional_params.organization
|
||||
or litellm.organization
|
||||
or os.getenv("OPENAI_ORGANIZATION", None)
|
||||
or None # default - https://github.com/openai/openai-python/blob/284c1799070c723c6a553337134148a7ab088dd8/openai/util.py#L105
|
||||
)
|
||||
# set API KEY
|
||||
api_key = (
|
||||
optional_params.api_key
|
||||
or litellm.api_key # for deepinfra/perplexity/anyscale we check in get_llm_provider and pass in the api key from there
|
||||
or litellm.openai_key
|
||||
or os.getenv("OPENAI_API_KEY")
|
||||
)
|
||||
|
||||
response = openai_batches_instance.retrieve_batch(
|
||||
_is_async=_is_async,
|
||||
retrieve_batch_data=_retrieve_batch_request,
|
||||
api_base=api_base,
|
||||
api_key=api_key,
|
||||
organization=organization,
|
||||
timeout=timeout,
|
||||
max_retries=optional_params.max_retries,
|
||||
)
|
||||
elif custom_llm_provider == "azure":
|
||||
api_base = (
|
||||
optional_params.api_base
|
||||
or litellm.api_base
|
||||
or get_secret_str("AZURE_API_BASE")
|
||||
)
|
||||
api_version = (
|
||||
optional_params.api_version
|
||||
or litellm.api_version
|
||||
or get_secret_str("AZURE_API_VERSION")
|
||||
)
|
||||
|
||||
api_key = (
|
||||
optional_params.api_key
|
||||
or litellm.api_key
|
||||
or litellm.azure_key
|
||||
or get_secret_str("AZURE_OPENAI_API_KEY")
|
||||
or get_secret_str("AZURE_API_KEY")
|
||||
)
|
||||
|
||||
extra_body = optional_params.get("extra_body", {})
|
||||
if extra_body is not None:
|
||||
extra_body.pop("azure_ad_token", None)
|
||||
else:
|
||||
get_secret_str("AZURE_AD_TOKEN") # type: ignore
|
||||
|
||||
response = azure_batches_instance.retrieve_batch(
|
||||
_is_async=_is_async,
|
||||
api_base=api_base,
|
||||
api_key=api_key,
|
||||
api_version=api_version,
|
||||
timeout=timeout,
|
||||
max_retries=optional_params.max_retries,
|
||||
retrieve_batch_data=_retrieve_batch_request,
|
||||
litellm_params=litellm_params,
|
||||
)
|
||||
elif custom_llm_provider == "vertex_ai":
|
||||
api_base = optional_params.api_base or ""
|
||||
vertex_ai_project = (
|
||||
optional_params.vertex_project
|
||||
or litellm.vertex_project
|
||||
or get_secret_str("VERTEXAI_PROJECT")
|
||||
)
|
||||
vertex_ai_location = (
|
||||
optional_params.vertex_location
|
||||
or litellm.vertex_location
|
||||
or get_secret_str("VERTEXAI_LOCATION")
|
||||
)
|
||||
vertex_credentials = optional_params.vertex_credentials or get_secret_str(
|
||||
"VERTEXAI_CREDENTIALS"
|
||||
)
|
||||
|
||||
response = vertex_ai_batches_instance.retrieve_batch(
|
||||
_is_async=_is_async,
|
||||
batch_id=batch_id,
|
||||
api_base=api_base,
|
||||
vertex_project=vertex_ai_project,
|
||||
vertex_location=vertex_ai_location,
|
||||
vertex_credentials=vertex_credentials,
|
||||
timeout=timeout,
|
||||
max_retries=optional_params.max_retries,
|
||||
)
|
||||
else:
|
||||
raise litellm.exceptions.BadRequestError(
|
||||
message="LiteLLM doesn't support {} for 'create_batch'. Only 'openai' is supported.".format(
|
||||
custom_llm_provider
|
||||
),
|
||||
model="n/a",
|
||||
llm_provider=custom_llm_provider,
|
||||
response=httpx.Response(
|
||||
status_code=400,
|
||||
content="Unsupported provider",
|
||||
request=httpx.Request(method="create_thread", url="https://github.com/BerriAI/litellm"), # type: ignore
|
||||
),
|
||||
)
|
||||
return response
|
||||
|
||||
@client
|
||||
def retrieve_batch(
|
||||
batch_id: str,
|
||||
custom_llm_provider: Literal["openai", "azure", "vertex_ai"] = "openai",
|
||||
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "bedrock"] = "openai",
|
||||
metadata: Optional[Dict[str, str]] = None,
|
||||
extra_headers: Optional[Dict[str, str]] = None,
|
||||
extra_body: Optional[Dict[str, str]] = None,
|
||||
@@ -430,115 +548,59 @@ def retrieve_batch(
|
||||
)
|
||||
|
||||
_is_async = kwargs.pop("aretrieve_batch", False) is True
|
||||
api_base: Optional[str] = None
|
||||
if custom_llm_provider == "openai":
|
||||
# for deepinfra/perplexity/anyscale/groq we check in get_llm_provider and pass in the api base from there
|
||||
api_base = (
|
||||
optional_params.api_base
|
||||
or litellm.api_base
|
||||
or os.getenv("OPENAI_BASE_URL")
|
||||
or os.getenv("OPENAI_API_BASE")
|
||||
or "https://api.openai.com/v1"
|
||||
)
|
||||
organization = (
|
||||
optional_params.organization
|
||||
or litellm.organization
|
||||
or os.getenv("OPENAI_ORGANIZATION", None)
|
||||
or None # default - https://github.com/openai/openai-python/blob/284c1799070c723c6a553337134148a7ab088dd8/openai/util.py#L105
|
||||
)
|
||||
# set API KEY
|
||||
api_key = (
|
||||
optional_params.api_key
|
||||
or litellm.api_key # for deepinfra/perplexity/anyscale we check in get_llm_provider and pass in the api key from there
|
||||
or litellm.openai_key
|
||||
or os.getenv("OPENAI_API_KEY")
|
||||
)
|
||||
|
||||
response = openai_batches_instance.retrieve_batch(
|
||||
_is_async=_is_async,
|
||||
retrieve_batch_data=_retrieve_batch_request,
|
||||
api_base=api_base,
|
||||
api_key=api_key,
|
||||
organization=organization,
|
||||
timeout=timeout,
|
||||
max_retries=optional_params.max_retries,
|
||||
)
|
||||
elif custom_llm_provider == "azure":
|
||||
api_base = (
|
||||
optional_params.api_base
|
||||
or litellm.api_base
|
||||
or get_secret_str("AZURE_API_BASE")
|
||||
)
|
||||
api_version = (
|
||||
optional_params.api_version
|
||||
or litellm.api_version
|
||||
or get_secret_str("AZURE_API_VERSION")
|
||||
)
|
||||
|
||||
api_key = (
|
||||
optional_params.api_key
|
||||
or litellm.api_key
|
||||
or litellm.azure_key
|
||||
or get_secret_str("AZURE_OPENAI_API_KEY")
|
||||
or get_secret_str("AZURE_API_KEY")
|
||||
)
|
||||
|
||||
extra_body = optional_params.get("extra_body", {})
|
||||
if extra_body is not None:
|
||||
extra_body.pop("azure_ad_token", None)
|
||||
else:
|
||||
get_secret_str("AZURE_AD_TOKEN") # type: ignore
|
||||
|
||||
response = azure_batches_instance.retrieve_batch(
|
||||
_is_async=_is_async,
|
||||
api_base=api_base,
|
||||
api_key=api_key,
|
||||
api_version=api_version,
|
||||
timeout=timeout,
|
||||
max_retries=optional_params.max_retries,
|
||||
retrieve_batch_data=_retrieve_batch_request,
|
||||
litellm_params=litellm_params,
|
||||
)
|
||||
elif custom_llm_provider == "vertex_ai":
|
||||
api_base = optional_params.api_base or ""
|
||||
vertex_ai_project = (
|
||||
optional_params.vertex_project
|
||||
or litellm.vertex_project
|
||||
or get_secret_str("VERTEXAI_PROJECT")
|
||||
)
|
||||
vertex_ai_location = (
|
||||
optional_params.vertex_location
|
||||
or litellm.vertex_location
|
||||
or get_secret_str("VERTEXAI_LOCATION")
|
||||
)
|
||||
vertex_credentials = optional_params.vertex_credentials or get_secret_str(
|
||||
"VERTEXAI_CREDENTIALS"
|
||||
)
|
||||
|
||||
response = vertex_ai_batches_instance.retrieve_batch(
|
||||
_is_async=_is_async,
|
||||
batch_id=batch_id,
|
||||
api_base=api_base,
|
||||
vertex_project=vertex_ai_project,
|
||||
vertex_location=vertex_ai_location,
|
||||
vertex_credentials=vertex_credentials,
|
||||
timeout=timeout,
|
||||
max_retries=optional_params.max_retries,
|
||||
client = kwargs.get("client", None)
|
||||
|
||||
# Try to use provider config first (for providers like bedrock)
|
||||
model: Optional[str] = kwargs.get("model", None)
|
||||
if model is not None:
|
||||
provider_config = ProviderConfigManager.get_provider_batches_config(
|
||||
model=model,
|
||||
provider=LlmProviders(custom_llm_provider),
|
||||
)
|
||||
else:
|
||||
raise litellm.exceptions.BadRequestError(
|
||||
message="LiteLLM doesn't support {} for 'create_batch'. Only 'openai' is supported.".format(
|
||||
custom_llm_provider
|
||||
),
|
||||
model="n/a",
|
||||
llm_provider=custom_llm_provider,
|
||||
response=httpx.Response(
|
||||
status_code=400,
|
||||
content="Unsupported provider",
|
||||
request=httpx.Request(method="create_thread", url="https://github.com/BerriAI/litellm"), # type: ignore
|
||||
provider_config = None
|
||||
|
||||
if provider_config is not None:
|
||||
response = base_llm_http_handler.retrieve_batch(
|
||||
batch_id=batch_id,
|
||||
provider_config=provider_config,
|
||||
litellm_params=litellm_params,
|
||||
headers=extra_headers or {},
|
||||
api_base=optional_params.api_base,
|
||||
api_key=optional_params.api_key,
|
||||
logging_obj=litellm_logging_obj or LiteLLMLoggingObj(
|
||||
model=model or "bedrock/unknown",
|
||||
messages=[],
|
||||
stream=False,
|
||||
call_type="batch_retrieve",
|
||||
start_time=None,
|
||||
litellm_call_id="batch_retrieve_" + batch_id,
|
||||
function_id="batch_retrieve",
|
||||
),
|
||||
_is_async=_is_async,
|
||||
client=client
|
||||
if client is not None
|
||||
and isinstance(client, (HTTPHandler, AsyncHTTPHandler))
|
||||
else None,
|
||||
timeout=timeout,
|
||||
model=model,
|
||||
)
|
||||
return response
|
||||
return response
|
||||
|
||||
|
||||
#########################################################
|
||||
# Handle providers without provider config
|
||||
#########################################################
|
||||
return _handle_retrieve_batch_providers_without_provider_config(
|
||||
batch_id=batch_id,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
optional_params=optional_params,
|
||||
litellm_params=litellm_params,
|
||||
_retrieve_batch_request=_retrieve_batch_request,
|
||||
_is_async=_is_async,
|
||||
timeout=timeout,
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
raise e
|
||||
|
||||
|
||||
@@ -158,6 +158,48 @@ class BaseBatchesConfig(ABC):
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def transform_retrieve_batch_request(
|
||||
self,
|
||||
batch_id: str,
|
||||
optional_params: dict,
|
||||
litellm_params: dict,
|
||||
) -> Union[bytes, str, Dict[str, Any]]:
|
||||
"""
|
||||
Transform the batch retrieval request to provider-specific format.
|
||||
|
||||
Args:
|
||||
batch_id: Batch ID to retrieve
|
||||
optional_params: Optional parameters
|
||||
litellm_params: LiteLLM parameters
|
||||
|
||||
Returns:
|
||||
Transformed request data
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def transform_retrieve_batch_response(
|
||||
self,
|
||||
model: Optional[str],
|
||||
raw_response: httpx.Response,
|
||||
logging_obj: LiteLLMLoggingObj,
|
||||
litellm_params: dict,
|
||||
) -> LiteLLMBatch:
|
||||
"""
|
||||
Transform provider-specific batch retrieval response to LiteLLM format.
|
||||
|
||||
Args:
|
||||
model: Model name
|
||||
raw_response: Raw HTTP response
|
||||
logging_obj: Logging object
|
||||
litellm_params: LiteLLM parameters
|
||||
|
||||
Returns:
|
||||
LiteLLM batch object
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def get_error_class(
|
||||
self, error_message: str, status_code: int, headers: Union[Dict, Headers]
|
||||
|
||||
@@ -7,7 +7,6 @@ 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.types.llms.bedrock import (
|
||||
BedrockBatchJobStatus,
|
||||
BedrockCreateBatchRequest,
|
||||
BedrockCreateBatchResponse,
|
||||
BedrockInputDataConfig,
|
||||
@@ -200,19 +199,23 @@ class BedrockBatchesConfig(BaseAWSLLM, BaseBatchesConfig):
|
||||
|
||||
# Extract information from typed Bedrock response
|
||||
job_arn = response_data.get("jobArn", "")
|
||||
status: BedrockBatchJobStatus = response_data.get("status", "Submitted")
|
||||
status_str: str = str(response_data.get("status", "Submitted"))
|
||||
|
||||
# Map Bedrock status to OpenAI-compatible status
|
||||
status_mapping: Dict[BedrockBatchJobStatus, str] = {
|
||||
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"
|
||||
"Stopped": "cancelled",
|
||||
"Expired": "expired",
|
||||
}
|
||||
|
||||
openai_status = cast(Literal["validating", "failed", "in_progress", "finalizing", "completed", "expired", "cancelling", "cancelled"], status_mapping.get(status, "validating"))
|
||||
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", {})
|
||||
@@ -229,7 +232,7 @@ class BedrockBatchesConfig(BaseAWSLLM, BaseBatchesConfig):
|
||||
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 == "InProgress" else None,
|
||||
in_progress_at=int(time.time()) if status_str == "InProgress" else None,
|
||||
expires_at=None,
|
||||
finalizing_at=None,
|
||||
completed_at=None,
|
||||
@@ -241,6 +244,203 @@ class BedrockBatchesConfig(BaseAWSLLM, BaseBatchesConfig):
|
||||
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:
|
||||
|
||||
@@ -738,19 +738,24 @@ class CommonBatchFilesUtils:
|
||||
)
|
||||
|
||||
# Prepare the request data
|
||||
if isinstance(data, dict):
|
||||
import json
|
||||
request_data = json.dumps(data)
|
||||
method_upper = method.upper()
|
||||
if method_upper == "GET":
|
||||
# GET requests should be signed with an empty payload
|
||||
request_data = ""
|
||||
headers = {}
|
||||
else:
|
||||
request_data = data
|
||||
|
||||
# Prepare headers
|
||||
headers = {"Content-Type": "application/json"}
|
||||
if isinstance(data, dict):
|
||||
import json
|
||||
request_data = json.dumps(data)
|
||||
else:
|
||||
request_data = data
|
||||
# Prepare headers for non-GET requests
|
||||
headers = {"Content-Type": "application/json"}
|
||||
|
||||
# Create AWS request and sign it
|
||||
sigv4 = SigV4Auth(credentials, service_name, aws_region_name)
|
||||
request = AWSRequest(
|
||||
method=method.upper(), url=endpoint_url, data=request_data, headers=headers
|
||||
method=method_upper, url=endpoint_url, data=request_data, headers=headers
|
||||
)
|
||||
sigv4.add_auth(request)
|
||||
prepped = request.prepare()
|
||||
|
||||
@@ -2520,6 +2520,95 @@ class BaseLLMHTTPHandler:
|
||||
litellm_params=litellm_params_with_request,
|
||||
)
|
||||
|
||||
def retrieve_batch(
|
||||
self,
|
||||
batch_id: str,
|
||||
litellm_params: dict,
|
||||
provider_config: "BaseBatchesConfig",
|
||||
headers: dict,
|
||||
api_base: Optional[str],
|
||||
api_key: Optional[str],
|
||||
logging_obj: "LiteLLMLoggingObj",
|
||||
_is_async: bool = False,
|
||||
client: Optional[Union["HTTPHandler", "AsyncHTTPHandler"]] = None,
|
||||
timeout: Optional[Union[float, httpx.Timeout]] = None,
|
||||
model: Optional[str] = None,
|
||||
) -> Union["LiteLLMBatch", Coroutine[Any, Any, "LiteLLMBatch"]]:
|
||||
"""
|
||||
Retrieve a batch using provider-specific configuration.
|
||||
"""
|
||||
# Transform the request using provider config
|
||||
transformed_request = provider_config.transform_retrieve_batch_request(
|
||||
batch_id=batch_id,
|
||||
optional_params=litellm_params,
|
||||
litellm_params=litellm_params,
|
||||
)
|
||||
|
||||
if _is_async:
|
||||
return self.async_retrieve_batch(
|
||||
transformed_request=transformed_request,
|
||||
litellm_params=litellm_params,
|
||||
provider_config=provider_config,
|
||||
headers=headers,
|
||||
api_base=api_base,
|
||||
logging_obj=logging_obj,
|
||||
client=client,
|
||||
timeout=timeout,
|
||||
batch_id=batch_id,
|
||||
model=model,
|
||||
)
|
||||
|
||||
if client is None or not isinstance(client, HTTPHandler):
|
||||
sync_httpx_client = _get_httpx_client()
|
||||
else:
|
||||
sync_httpx_client = client
|
||||
|
||||
try:
|
||||
if (
|
||||
isinstance(transformed_request, dict)
|
||||
and "method" in transformed_request
|
||||
):
|
||||
# Handle pre-signed requests (e.g., from Bedrock with AWS auth)
|
||||
method = transformed_request["method"].lower()
|
||||
request_kwargs = {
|
||||
"url": transformed_request["url"],
|
||||
"headers": transformed_request["headers"],
|
||||
}
|
||||
|
||||
# Only add data for non-GET requests
|
||||
if method != "get" and transformed_request.get("data") is not None:
|
||||
request_kwargs["data"] = transformed_request["data"]
|
||||
|
||||
batch_response = getattr(sync_httpx_client, method)(**request_kwargs)
|
||||
elif isinstance(transformed_request, dict) and api_base:
|
||||
# For other providers that use JSON requests
|
||||
batch_response = sync_httpx_client.get(
|
||||
url=api_base,
|
||||
headers={**headers, "Content-Type": "application/json"},
|
||||
params=transformed_request,
|
||||
)
|
||||
else:
|
||||
# Handle other request types if needed
|
||||
if not api_base:
|
||||
raise ValueError("api_base is required for non-pre-signed requests")
|
||||
batch_response = sync_httpx_client.get(
|
||||
url=api_base,
|
||||
headers=headers,
|
||||
)
|
||||
except Exception as e:
|
||||
verbose_logger.exception(f"Error retrieving batch: {e}")
|
||||
raise self._handle_error(
|
||||
e=e,
|
||||
provider_config=provider_config,
|
||||
)
|
||||
|
||||
return provider_config.transform_retrieve_batch_response(
|
||||
model=model,
|
||||
raw_response=batch_response,
|
||||
logging_obj=logging_obj,
|
||||
litellm_params=litellm_params,
|
||||
)
|
||||
|
||||
async def async_create_batch(
|
||||
self,
|
||||
transformed_request: Union[bytes, str, dict],
|
||||
@@ -2606,6 +2695,89 @@ class BaseLLMHTTPHandler:
|
||||
litellm_params=litellm_params_with_request,
|
||||
)
|
||||
|
||||
async def async_retrieve_batch(
|
||||
self,
|
||||
transformed_request: Union[bytes, str, dict],
|
||||
litellm_params: dict,
|
||||
provider_config: "BaseBatchesConfig",
|
||||
headers: dict,
|
||||
api_base: Optional[str],
|
||||
logging_obj: "LiteLLMLoggingObj",
|
||||
client: Optional[Union["HTTPHandler", "AsyncHTTPHandler"]] = None,
|
||||
timeout: Optional[Union[float, httpx.Timeout]] = None,
|
||||
batch_id: Optional[str] = None,
|
||||
model: Optional[str] = None,
|
||||
):
|
||||
"""
|
||||
Async version of retrieve_batch
|
||||
"""
|
||||
if client is None or not isinstance(client, AsyncHTTPHandler):
|
||||
async_httpx_client = get_async_httpx_client(
|
||||
llm_provider=provider_config.custom_llm_provider
|
||||
)
|
||||
else:
|
||||
async_httpx_client = client
|
||||
|
||||
#########################################################
|
||||
# Debug Logging
|
||||
#########################################################
|
||||
logging_obj.pre_call(
|
||||
input="",
|
||||
api_key="",
|
||||
additional_args={
|
||||
"complete_input_dict": transformed_request,
|
||||
"api_base": api_base,
|
||||
"headers": headers,
|
||||
"batch_id": batch_id,
|
||||
},
|
||||
)
|
||||
|
||||
try:
|
||||
if (
|
||||
isinstance(transformed_request, dict)
|
||||
and "method" in transformed_request
|
||||
):
|
||||
# Handle pre-signed requests (e.g., from Bedrock with AWS auth)
|
||||
method = transformed_request["method"].lower()
|
||||
request_kwargs = {
|
||||
"url": transformed_request["url"],
|
||||
"headers": transformed_request["headers"],
|
||||
}
|
||||
|
||||
# Only add data for non-GET requests
|
||||
if method != "get" and transformed_request.get("data") is not None:
|
||||
request_kwargs["data"] = transformed_request["data"]
|
||||
|
||||
batch_response = await getattr(async_httpx_client, method)(**request_kwargs)
|
||||
elif isinstance(transformed_request, dict) and api_base:
|
||||
# For other providers that use JSON requests
|
||||
batch_response = await async_httpx_client.get(
|
||||
url=api_base,
|
||||
headers={**headers, "Content-Type": "application/json"},
|
||||
params=transformed_request,
|
||||
)
|
||||
else:
|
||||
# Handle other request types if needed
|
||||
if not api_base:
|
||||
raise ValueError("api_base is required for non-pre-signed requests")
|
||||
batch_response = await async_httpx_client.get(
|
||||
url=api_base,
|
||||
headers=headers,
|
||||
)
|
||||
except Exception as e:
|
||||
verbose_logger.exception(f"Error retrieving batch: {e}")
|
||||
raise self._handle_error(
|
||||
e=e,
|
||||
provider_config=provider_config,
|
||||
)
|
||||
|
||||
return provider_config.transform_retrieve_batch_response(
|
||||
model=model,
|
||||
raw_response=batch_response,
|
||||
logging_obj=logging_obj,
|
||||
litellm_params=litellm_params,
|
||||
)
|
||||
|
||||
def cancel_response_api_handler(
|
||||
self,
|
||||
response_id: str,
|
||||
|
||||
@@ -75,6 +75,19 @@ async def test_async_file_and_batch():
|
||||
)
|
||||
print("CREATED BATCH RESPONSE=", create_batch_response)
|
||||
|
||||
# retrieve batch
|
||||
retrieve_batch_response = await litellm.aretrieve_batch(
|
||||
batch_id=create_batch_response.id,
|
||||
custom_llm_provider="bedrock",
|
||||
model="us.anthropic.claude-3-5-sonnet-20240620-v1:0",
|
||||
)
|
||||
print("RETRIEVED BATCH RESPONSE=", retrieve_batch_response)
|
||||
|
||||
# Validate the response
|
||||
assert retrieve_batch_response.id == create_batch_response.id
|
||||
assert retrieve_batch_response.object == "batch"
|
||||
assert retrieve_batch_response.status in ["validating", "in_progress", "completed", "failed", "cancelled"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio()
|
||||
async def test_mock_bedrock_file_url_mapping():
|
||||
@@ -118,3 +131,65 @@ async def test_mock_bedrock_file_url_mapping():
|
||||
expected_s3_uri, _ = bedrock_config._convert_https_url_to_s3_uri(captured_put_url)
|
||||
assert file_obj.id == expected_s3_uri
|
||||
|
||||
|
||||
@pytest.mark.asyncio()
|
||||
async def test_bedrock_retrieve_batch():
|
||||
"""
|
||||
Test bedrock batch retrieval functionality, validating that input and output file IDs
|
||||
are correctly extracted from the Bedrock response and included in the final transformed response.
|
||||
"""
|
||||
print("Testing bedrock batch retrieval")
|
||||
|
||||
# Mock bedrock batch response
|
||||
mock_bedrock_response = {
|
||||
"jobArn": "arn:aws:bedrock:us-west-2:123456789012:model-invocation-job/test-job-123",
|
||||
"jobName": "test-job-123",
|
||||
"modelId": "us.anthropic.claude-3-5-sonnet-20240620-v1:0",
|
||||
"roleArn": "arn:aws:iam::123456789012:role/service-role/AmazonBedrockExecutionRoleForAgents_TEST",
|
||||
"status": "InProgress",
|
||||
"message": "Job is in progress",
|
||||
"submitTime": "2024-01-01T12:00:00Z",
|
||||
"lastModifiedTime": "2024-01-01T12:30:00Z",
|
||||
"inputDataConfig": {
|
||||
"s3InputDataConfig": {
|
||||
"s3Uri": "s3://test-bucket/input/test-input.jsonl"
|
||||
}
|
||||
},
|
||||
"outputDataConfig": {
|
||||
"s3OutputDataConfig": {
|
||||
"s3Uri": "s3://test-bucket/output/"
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
# Mock the HTTP response
|
||||
mock_response = MagicMock()
|
||||
mock_response.json.return_value = mock_bedrock_response
|
||||
mock_response.status_code = 200
|
||||
|
||||
# Print the mock response to debug
|
||||
print("MOCK RESPONSE DATA:", mock_bedrock_response)
|
||||
|
||||
with patch("litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.get") as mock_get:
|
||||
mock_response.raise_for_status.return_value = None
|
||||
mock_get.return_value = mock_response
|
||||
|
||||
# Test retrieve batch
|
||||
batch_response = await litellm.aretrieve_batch(
|
||||
batch_id="arn:aws:bedrock:us-west-2:123456789012:model-invocation-job/test-job-123",
|
||||
custom_llm_provider="bedrock",
|
||||
model="us.anthropic.claude-3-5-sonnet-20240620-v1:0",
|
||||
)
|
||||
|
||||
print("MOCKED BATCH RESPONSE=", batch_response)
|
||||
|
||||
# Validate the response
|
||||
assert batch_response.id == "arn:aws:bedrock:us-west-2:123456789012:model-invocation-job/test-job-123"
|
||||
assert batch_response.object == "batch"
|
||||
assert batch_response.status == "in_progress" # Bedrock "InProgress" maps to "in_progress"
|
||||
assert batch_response.endpoint == "/v1/chat/completions"
|
||||
|
||||
# Validate input and output file IDs in the final transformed response
|
||||
assert batch_response.input_file_id == "s3://test-bucket/input/test-input.jsonl"
|
||||
assert batch_response.output_file_id == "s3://test-bucket/output/"
|
||||
|
||||
|
||||
Reference in New Issue
Block a user