From 855c49d0ef01a09161f8d2ec195be01447669f1a Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 8 Aug 2026 01:44:56 -0700 Subject: [PATCH] fix(proxy): skip prisma-dependent hooks when no database is attached --- .../storage_backend_service.py | 10 ++ litellm/proxy/utils.py | 5 + .../test_storage_backend_service.py | 127 ++++++++++++++++++ .../utils/proxy_logging/test_lifecycle.py | 94 +++++++++++-- 4 files changed, 228 insertions(+), 8 deletions(-) create mode 100644 tests/test_litellm/proxy/openai_files_endpoint/test_storage_backend_service.py diff --git a/litellm/proxy/openai_files_endpoints/storage_backend_service.py b/litellm/proxy/openai_files_endpoints/storage_backend_service.py index 4c301c96f30..e766f335071 100644 --- a/litellm/proxy/openai_files_endpoints/storage_backend_service.py +++ b/litellm/proxy/openai_files_endpoints/storage_backend_service.py @@ -68,6 +68,16 @@ class StorageBackendFileService: code=400, ) + if target_model_names: + managed_files_hook: Final = proxy_logging_obj.get_proxy_hook("managed_files") + if not isinstance(managed_files_hook, BaseFileEndpoints): + raise ProxyException( + message="Uploading with target_model_names requires a database-connected proxy, and this proxy has no database configured", + type="invalid_request_error", + param="target_model_names", + code=400, + ) + # Extract file information file_content: Final = file_data["content"] filename: Final = file_data.get("filename", "file") diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index e59c6adaf22..5f22ca021ac 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -544,6 +544,11 @@ class ProxyLogging: for hook in PROXY_HOOKS: proxy_hook = get_proxy_hook(hook) expected_args = inspect.getfullargspec(proxy_hook).args + if "prisma_client" in expected_args and prisma_client is None: + verbose_proxy_logger.debug( + "Skipping proxy hook %s: it requires a database and no prisma client is configured", hook + ) + continue passed_in_args: dict[str, Any] = {} if "internal_usage_cache" in expected_args: passed_in_args["internal_usage_cache"] = self.internal_usage_cache diff --git a/tests/test_litellm/proxy/openai_files_endpoint/test_storage_backend_service.py b/tests/test_litellm/proxy/openai_files_endpoint/test_storage_backend_service.py new file mode 100644 index 00000000000..07a85a70815 --- /dev/null +++ b/tests/test_litellm/proxy/openai_files_endpoint/test_storage_backend_service.py @@ -0,0 +1,127 @@ +import pytest + +from litellm.llms.base_llm.files.transformation import BaseFileEndpoints +from litellm.proxy._types import ProxyException, UserAPIKeyAuth +from litellm.proxy.openai_files_endpoints import storage_backend_service +from litellm.proxy.openai_files_endpoints.storage_backend_service import ( + StorageBackendFileService, +) + + +class _RecordingStorageBackend: + def __init__(self): + self.upload_calls = [] + + async def upload_file(self, **kwargs): + self.upload_calls.append(kwargs) + return "https://storage.example/blob-1" + + +class _FakeManagedFilesHook(BaseFileEndpoints): + def __init__(self): + self.stored = [] + + async def acreate_file( + self, create_file_request, llm_router, target_model_names_list, litellm_parent_otel_span, user_api_key_dict + ): + raise NotImplementedError + + async def afile_retrieve(self, file_id, litellm_parent_otel_span, llm_router=None): + raise NotImplementedError + + async def afile_list(self, purpose, litellm_parent_otel_span, **data): + raise NotImplementedError + + async def afile_delete(self, file_id, litellm_parent_otel_span, llm_router, **data): + raise NotImplementedError + + async def afile_content(self, file_id, litellm_parent_otel_span, llm_router, **data): + raise NotImplementedError + + async def store_unified_file_id(self, **kwargs): + self.stored.append(kwargs) + + +class _FakeProxyLogging: + def __init__(self, hook): + self._hook = hook + + def get_proxy_hook(self, hook_name): + return self._hook if hook_name == "managed_files" else None + + +def _file_data(): + return {"content": b"x", "filename": "input.jsonl", "content_type": "application/jsonl"} + + +@pytest.mark.asyncio +async def test_upload_with_target_model_names_but_no_hook_raises_before_uploading(monkeypatch): + backend = _RecordingStorageBackend() + monkeypatch.setattr(storage_backend_service, "get_storage_backend", lambda name: backend) + + with pytest.raises(ProxyException) as exc_info: + await StorageBackendFileService.upload_file_to_storage_backend( + file_data=_file_data(), + target_storage="azure_storage", + target_model_names=["gpt-x"], + purpose="batch", + proxy_logging_obj=_FakeProxyLogging(hook=None), + user_api_key_dict=UserAPIKeyAuth(api_key="sk-test"), + ) + + snapshot = { + "code": exc_info.value.code, + "message_names_requirement": "requires a database-connected proxy" in exc_info.value.message, + "upload_calls": backend.upload_calls, + } + assert snapshot == {"code": "400", "message_names_requirement": True, "upload_calls": []} + + +@pytest.mark.asyncio +async def test_upload_without_target_model_names_skips_hook_requirement(monkeypatch): + backend = _RecordingStorageBackend() + monkeypatch.setattr(storage_backend_service, "get_storage_backend", lambda name: backend) + + file_object = await StorageBackendFileService.upload_file_to_storage_backend( + file_data=_file_data(), + target_storage="azure_storage", + target_model_names=[], + purpose="batch", + proxy_logging_obj=_FakeProxyLogging(hook=None), + user_api_key_dict=UserAPIKeyAuth(api_key="sk-test"), + ) + + snapshot = { + "upload_count": len(backend.upload_calls), + "id_prefix": file_object.id.split("-")[0], + } + assert snapshot == {"upload_count": 1, "id_prefix": "file"} + + +@pytest.mark.asyncio +async def test_upload_with_target_model_names_and_hook_stores_unified_id(monkeypatch): + backend = _RecordingStorageBackend() + monkeypatch.setattr(storage_backend_service, "get_storage_backend", lambda name: backend) + hook = _FakeManagedFilesHook() + + file_object = await StorageBackendFileService.upload_file_to_storage_backend( + file_data=_file_data(), + target_storage="azure_storage", + target_model_names=["gpt-x"], + purpose="batch", + proxy_logging_obj=_FakeProxyLogging(hook=hook), + user_api_key_dict=UserAPIKeyAuth(api_key="sk-test"), + ) + + snapshot = { + "upload_count": len(backend.upload_calls), + "store_count": len(hook.stored), + "stored_id_matches_response": hook.stored[0]["file_id"] == file_object.id, + "model_mappings": hook.stored[0]["model_mappings"], + } + assert snapshot == { + "upload_count": 1, + "store_count": 1, + "stored_id_matches_response": True, + "model_mappings": {"gpt-x": "https://storage.example/blob-1"}, + } diff --git a/tests/test_litellm/proxy/utils/proxy_logging/test_lifecycle.py b/tests/test_litellm/proxy/utils/proxy_logging/test_lifecycle.py index e33da672599..cf906259246 100644 --- a/tests/test_litellm/proxy/utils/proxy_logging/test_lifecycle.py +++ b/tests/test_litellm/proxy/utils/proxy_logging/test_lifecycle.py @@ -8,7 +8,6 @@ because they are direct dependents on the lifecycle state. from __future__ import annotations -import asyncio from typing import Any, Dict, List from unittest.mock import AsyncMock, MagicMock, patch @@ -17,7 +16,6 @@ import pytest import litellm from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache from litellm.proxy.utils import ( - InternalUsageCache, ProxyLogging, ) @@ -102,9 +100,7 @@ def test_update_values_with_no_args_is_noop(proxy_logging): def test_update_values_invalid_type_for_alerting_raises(proxy_logging): - proxy_logging.slack_alerting_instance = MagicMock( - update_values=MagicMock(side_effect=TypeError("bad type")) - ) + proxy_logging.slack_alerting_instance = MagicMock(update_values=MagicMock(side_effect=TypeError("bad type"))) with pytest.raises(TypeError): proxy_logging.update_values(alerting={"not": "a list"}) # type: ignore[arg-type] @@ -190,6 +186,90 @@ def test_add_proxy_hooks_registers_callbacks(proxy_logging, monkeypatch): } +def _stub_hook_classes(): + class _PrismaFreeHook: + def __init__(self, internal_usage_cache): + self.internal_usage_cache = internal_usage_cache + + class _PrismaRequiringHook: + def __init__(self, internal_usage_cache, prisma_client): + self.internal_usage_cache = internal_usage_cache + self.prisma_client = prisma_client + + class _PrismaOnlyHook: + def __init__(self, prisma_client): + self.prisma_client = prisma_client + + return { + "cache_control_check": _PrismaFreeHook, + "needs_db_hook": _PrismaRequiringHook, + "db_only_hook": _PrismaOnlyHook, + } + + +def test_add_proxy_hooks_skips_prisma_requiring_hook_when_no_db(proxy_logging, monkeypatch): + hook_classes = _stub_hook_classes() + registered: List[Any] = [] + + from litellm.proxy import utils as utils_mod + + monkeypatch.setattr(utils_mod, "PROXY_HOOKS", list(hook_classes.keys())) + monkeypatch.setattr(utils_mod, "get_proxy_hook", hook_classes.__getitem__) + monkeypatch.setattr( + litellm.logging_callback_manager, + "add_litellm_callback", + lambda cb: registered.append(cb), + ) + + with patch("litellm.proxy.proxy_server.prisma_client", None): + proxy_logging._add_proxy_hooks(llm_router=None) + + snapshot = { + "mapping_keys": list(proxy_logging.proxy_hook_mapping.keys()), + "registered_types": [type(r).__name__ for r in registered], + "needs_db_hook_lookup": proxy_logging.get_proxy_hook("needs_db_hook"), + "db_only_hook_lookup": proxy_logging.get_proxy_hook("db_only_hook"), + } + assert snapshot == { + "mapping_keys": ["cache_control_check"], + "registered_types": ["_PrismaFreeHook"], + "needs_db_hook_lookup": None, + "db_only_hook_lookup": None, + } + + +def test_add_proxy_hooks_registers_prisma_requiring_hook_with_db(proxy_logging, monkeypatch): + hook_classes = _stub_hook_classes() + registered: List[Any] = [] + fake_prisma = MagicMock() + + from litellm.proxy import utils as utils_mod + + monkeypatch.setattr(utils_mod, "PROXY_HOOKS", list(hook_classes.keys())) + monkeypatch.setattr(utils_mod, "get_proxy_hook", hook_classes.__getitem__) + monkeypatch.setattr( + litellm.logging_callback_manager, + "add_litellm_callback", + lambda cb: registered.append(cb), + ) + + with patch("litellm.proxy.proxy_server.prisma_client", fake_prisma): + proxy_logging._add_proxy_hooks(llm_router=None) + + snapshot = { + "mapping_keys": list(proxy_logging.proxy_hook_mapping.keys()), + "registered_count": len(registered), + "needs_db_hook_got_prisma": proxy_logging.proxy_hook_mapping["needs_db_hook"].prisma_client is fake_prisma, + "db_only_hook_got_prisma": proxy_logging.proxy_hook_mapping["db_only_hook"].prisma_client is fake_prisma, + } + assert snapshot == { + "mapping_keys": ["cache_control_check", "needs_db_hook", "db_only_hook"], + "registered_count": 3, + "needs_db_hook_got_prisma": True, + "db_only_hook_got_prisma": True, + } + + def test_add_proxy_hooks_unknown_hook_raises(proxy_logging, monkeypatch): from litellm.proxy import utils as utils_mod @@ -267,9 +347,7 @@ def test_init_litellm_callbacks_replaces_string_with_instance(proxy_logging, mon snapshot = { "replaced_first_item": litellm.callbacks[0] is sentinel_instance, "callbacks_grew_with_service": len(litellm.callbacks) >= 2, - "service_logging_appended": any( - "ServiceLogging" in type(c).__name__ for c in litellm.callbacks - ), + "service_logging_appended": any("ServiceLogging" in type(c).__name__ for c in litellm.callbacks), } assert snapshot == { "replaced_first_item": True,