[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

Author: ishaan-jaff

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")
+        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),
             )
-
-            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,
+        else:
+            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,
-            )
-        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(
+                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,
-                batch_id=batch_id,
-                api_base=api_base,
-                vertex_project=vertex_ai_project,
-                vertex_location=vertex_ai_location,
-                vertex_credentials=vertex_credentials,
+                client=client
+                if client is not None
+                and isinstance(client, (HTTPHandler, AsyncHTTPHandler))
+                else None,
                 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
-                ),
+                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
 

litellm/llms/base_llm/batches/transformation.py

@@ -158,6 +158,48 @@ def transform_create_batch_response(
         """
         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]