diff --git a/litellm/batches/main.py b/litellm/batches/main.py index e5c7a2e4b6..37b9aff4ef 100644 --- a/litellm/batches/main.py +++ b/litellm/batches/main.py @@ -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 diff --git a/litellm/llms/base_llm/batches/transformation.py b/litellm/llms/base_llm/batches/transformation.py index 1d3e54fae6..9e67689fcd 100644 --- a/litellm/llms/base_llm/batches/transformation.py +++ b/litellm/llms/base_llm/batches/transformation.py @@ -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] diff --git a/litellm/llms/bedrock/batches/transformation.py b/litellm/llms/bedrock/batches/transformation.py index 3338450f79..2f3d00dddd 100644 --- a/litellm/llms/bedrock/batches/transformation.py +++ b/litellm/llms/bedrock/batches/transformation.py @@ -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: diff --git a/litellm/llms/bedrock/common_utils.py b/litellm/llms/bedrock/common_utils.py index 831a6da93b..c7f7acf331 100644 --- a/litellm/llms/bedrock/common_utils.py +++ b/litellm/llms/bedrock/common_utils.py @@ -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() diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 8b925a375a..dc64fea1a3 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -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, diff --git a/tests/batches_tests/test_bedrock_files_and_batches.py b/tests/batches_tests/test_bedrock_files_and_batches.py index 3f90411f62..d082ed41ea 100644 --- a/tests/batches_tests/test_bedrock_files_and_batches.py +++ b/tests/batches_tests/test_bedrock_files_and_batches.py @@ -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/" +