Merge pull request #22081 from BerriAI/litellm_fix_bedrock_batch_health_check

Fix: Test connect failing for bedrock batches mode

litellm/batches/main.py

@@ -37,7 +37,9 @@ from litellm.types.llms.openai import (
 )
 from litellm.types.router import GenericLiteLLMParams
 from litellm.types.utils import (
+    LIST_BATCHES_SUPPORTED_PROVIDERS,
     OPENAI_COMPATIBLE_BATCH_AND_FILES_PROVIDERS,
+    ListBatchesSupportedProvider,
     LiteLLMBatch,
     LlmProviders,
 )
@@ -674,7 +676,7 @@ def retrieve_batch(
 async def alist_batches(
     after: Optional[str] = None,
     limit: Optional[int] = None,
-    custom_llm_provider: Literal["openai", "azure", "hosted_vllm", "vertex_ai"] = "openai",
+    custom_llm_provider: ListBatchesSupportedProvider = "openai",
     metadata: Optional[Dict[str, str]] = None,
     extra_headers: Optional[Dict[str, str]] = None,
     extra_body: Optional[Dict[str, str]] = None,
@@ -717,7 +719,7 @@ async def alist_batches(
 def list_batches(
     after: Optional[str] = None,
     limit: Optional[int] = None,
-    custom_llm_provider: Literal["openai", "azure", "hosted_vllm", "vertex_ai"] = "openai",
+    custom_llm_provider: ListBatchesSupportedProvider = "openai",
     extra_headers: Optional[Dict[str, str]] = None,
     extra_body: Optional[Dict[str, str]] = None,
     **kwargs,
@@ -843,8 +845,9 @@ def list_batches(
             )
         else:
             raise litellm.exceptions.BadRequestError(
-                message="LiteLLM doesn't support {} for 'list_batch'. Supported providers: openai, azure, vertex_ai.".format(
-                    custom_llm_provider
+                message="LiteLLM doesn't support {} for 'list_batch'. Supported providers: {}.".format(
+                    custom_llm_provider,
+                    ", ".join(sorted(LIST_BATCHES_SUPPORTED_PROVIDERS)),
                 ),
                 model="n/a",
                 llm_provider=custom_llm_provider,

litellm/litellm_core_utils/health_check_helpers.py

@@ -4,6 +4,8 @@ Helper functions for health check calls.
 
 from typing import TYPE_CHECKING, Callable, Dict, Literal, Optional
 
+from litellm.types.utils import LIST_BATCHES_SUPPORTED_PROVIDERS
+
 if TYPE_CHECKING:
     from litellm.litellm_core_utils.litellm_logging import Logging
 
@@ -82,6 +84,27 @@ class HealthCheckHelpers:
             "tags": [LITTELM_INTERNAL_HEALTH_SERVICE_ACCOUNT_NAME],
         }
 
+    @staticmethod
+    async def _batch_health_check(
+        custom_llm_provider: str,
+        model_params: dict,
+        filtered_model_params: dict,
+    ) -> dict:
+        """
+        Health check for batch mode.
+
+        Calls list_batches for providers that support it (openai, hosted_vllm, azure,
+        vertex_ai). For all other providers (e.g. bedrock) the batch API surface doesn't
+        include list_batches, so we fall back to acompletion to verify connectivity and
+        credential validity instead.
+        """
+        import litellm
+
+        if custom_llm_provider in LIST_BATCHES_SUPPORTED_PROVIDERS:
+            return await litellm.alist_batches(**filtered_model_params)
+        else:
+            return await litellm.acompletion(**model_params)
+
     @staticmethod
     def get_mode_handlers(
         model: str,
@@ -176,8 +199,10 @@ class HealthCheckHelpers:
                 api_key=model_params.get("api_key", None),
                 api_version=model_params.get("api_version", None),
             ),
-            "batch": lambda: litellm.alist_batches(
-                **_filter_model_params(model_params=model_params),
+            "batch": lambda: HealthCheckHelpers._batch_health_check(
+                custom_llm_provider=custom_llm_provider,
+                model_params=model_params,
+                filtered_model_params=_filter_model_params(model_params=model_params),
             ),
             "responses": lambda: litellm.aresponses(
                 **_filter_model_params(model_params=model_params),

litellm/types/utils.py

@@ -1,7 +1,17 @@
 import json
 import time
 from enum import Enum
-from typing import TYPE_CHECKING, Any, Dict, List, Literal, Mapping, Optional, Union
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Dict,
+    List,
+    Literal,
+    Mapping,
+    Optional,
+    Union,
+    get_args,
+)
 
 from openai._models import BaseModel as OpenAIObject
 from openai.types.audio.transcription_create_params import (
@@ -3186,6 +3196,12 @@ OPENAI_COMPATIBLE_BATCH_AND_FILES_PROVIDERS: set[str] = {
     LlmProviders.HOSTED_VLLM.value,
 }
 
+ListBatchesSupportedProvider = Literal["openai", "azure", "hosted_vllm", "vertex_ai"]
+
+LIST_BATCHES_SUPPORTED_PROVIDERS: frozenset[str] = frozenset(
+    get_args(ListBatchesSupportedProvider)
+)
+
 
 class SearchProviders(str, Enum):
     """