fix(proxy): resolve managed batch type-check diagnostics

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

enterprise/litellm_enterprise/proxy/hooks/managed_files.py

@@ -63,7 +63,7 @@ from litellm.proxy._types import (
 )
 from litellm.proxy.litellm_pre_call_utils import (
     LiteLLMProxyRequestSetup,
-    _sanitize_for_log,
+    sanitize_for_log,
 )
 from litellm.proxy.openai_files_endpoints.common_utils import (
     BATCH_CREATE_HIDDEN_PARAM,
@@ -212,7 +212,7 @@ def _provider_file_retrieve_credentials(
     except Exception as error:
         verbose_logger.warning(
             "Failed to retrieve credentials for provider file "
-            f"model_id={_sanitize_for_log(model_id)}: {_sanitize_for_log(error)}"
+            f"model_id={sanitize_for_log(model_id)}: {sanitize_for_log(error)}"
         )
         return None
     return cast(Mapping[str, object], credentials) if credentials else None
@@ -819,8 +819,8 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
         except Exception as error:
             verbose_logger.warning(
                 "Failed to retrieve batch file object for "
-                f"provider_file_id={_sanitize_for_log(provider_file_id)}: "
-                f"{type(error).__name__} {_sanitize_for_log(error)}"
+                f"provider_file_id={sanitize_for_log(provider_file_id)}: "
+                f"{type(error).__name__} {sanitize_for_log(error)}"
             )
             if raise_on_failure:
                 if isinstance(error, TimeoutError) and not str(error):
@@ -1946,7 +1946,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
                 except Exception as error:
                     verbose_logger.warning(
                         "Failed to refresh batch file object for "
-                        f"file_id={_sanitize_for_log(file_id)}: {_sanitize_for_log(error)}"
+                        f"file_id={sanitize_for_log(file_id)}: {sanitize_for_log(error)}"
                     )
             return _public_file_object(file_object, file_id)
 

litellm/proxy/batches_endpoints/endpoints.py

@@ -75,7 +75,7 @@ from litellm.repositories.managed_batch_repository import ManagedBatchRepository
 from litellm.repositories.managed_file_repository import ManagedFileRepository
 from litellm.router import Router
 from litellm.types.llms.openai import LiteLLMBatchCreateRequest
-from litellm.types.utils import LiteLLMBatch
+from litellm.types.utils import LiteLLMBatch, LLMResponseTypes
 
 if TYPE_CHECKING:
     from prisma.models import LiteLLM_ManagedObjectTable
@@ -96,6 +96,12 @@ def _request_tags(data: Mapping[str, object]) -> tuple[str, ...] | None:
     return request_tags_from_metadata(_METADATA_ADAPTER.validate_python(metadata))
 
 
+def _require_batch_response(response: LLMResponseTypes) -> LiteLLMBatch:
+    if not isinstance(response, LiteLLMBatch):
+        raise TypeError("Batch endpoint received a non-batch response")
+    return response
+
+
 def _litellm_executed_batch_runner(llm_router: Router, proxy_logging_obj: ProxyLogging) -> LiteLLMExecutedBatchRunner:
     from litellm.proxy.proxy_server import general_settings, prisma_client
 
@@ -652,8 +658,9 @@ async def retrieve_batch(
             # The DB may store raw provider file IDs (before hooks translate them).
             # Register any missing managed-file rows and return unified IDs.
             if unified_batch_id:
+                batch_response: Final = _require_batch_response(response)
                 await ensure_batch_response_managed_file_ids(
-                    response=response,
+                    response=batch_response,
                     managed_files_obj=managed_files_obj,
                     prisma_client=prisma_client,
                     verbose_proxy_logger=verbose_proxy_logger,
@@ -799,8 +806,9 @@ async def retrieve_batch(
         # Fix: bug_feb14_batch_retrieve_returns_raw_input_file_id
         # Register any missing managed-file rows and return unified IDs.
         if unified_batch_id:
+            batch_response: Final = _require_batch_response(response)
             await ensure_batch_response_managed_file_ids(
-                response=response,
+                response=batch_response,
                 managed_files_obj=managed_files_obj,
                 prisma_client=prisma_client,
                 verbose_proxy_logger=verbose_proxy_logger,

litellm/proxy/litellm_pre_call_utils.py

@@ -209,6 +209,10 @@ def _sanitize_for_log(value: object) -> str:
     return text.replace("\r", "").replace("\n", "")
 
 
+def sanitize_for_log(value: object) -> str:
+    return _sanitize_for_log(value)
+
+
 from litellm.router import Router
 from litellm.secret_managers.main import get_secret_bool
 from litellm.types.llms.anthropic import ANTHROPIC_API_HEADERS

tests/unit/enterprise/proxy/hooks/test_managed_files.py

@@ -2452,9 +2452,9 @@ async def test_store_batch_output_file_marks_fallback_after_four_transient_failu
     stored_object = managed_file_table.rows["unified-output"].file_object
     assert stored_object is not None
     assert stored_object.litellm_details_fallback is True
-    assert json.loads(managed_file_table.upsert_calls[0][1]["create"]["file_object"])[
-        "litellm_details_fallback"
-    ] is True
+    stored_file_object_json: Final = managed_file_table.upsert_calls[0][1]["create"]["file_object"]
+    assert isinstance(stored_file_object_json, str)
+    assert json.loads(stored_file_object_json)["litellm_details_fallback"] is True
 
 
 @pytest.mark.asyncio