diff --git a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py index 45a9fdf2af3..4c06f0b9802 100644 --- a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py +++ b/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) diff --git a/litellm/proxy/batches_endpoints/endpoints.py b/litellm/proxy/batches_endpoints/endpoints.py index 9071ef163ea..bb606633572 100644 --- a/litellm/proxy/batches_endpoints/endpoints.py +++ b/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, diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 66dca57e782..5130ab35da5 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/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 diff --git a/tests/unit/enterprise/proxy/hooks/test_managed_files.py b/tests/unit/enterprise/proxy/hooks/test_managed_files.py index df05b82c924..ed80bd9b928 100644 --- a/tests/unit/enterprise/proxy/hooks/test_managed_files.py +++ b/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