From c9125626726f93172ea1cbe007e91b6977e60d62 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 20 Oct 2026 19:23:31 -0700 Subject: [PATCH 01/11] test(proxy): cover managed files hook registration without the repo-only enterprise package --- .../unit/proxy/hooks/test_proxy_hooks_init.py | 22 +++++++++++++++++++ 1 file changed, 22 insertions(+) diff --git a/tests/unit/proxy/hooks/test_proxy_hooks_init.py b/tests/unit/proxy/hooks/test_proxy_hooks_init.py index a6edd3db944..009d69912e0 100644 --- a/tests/unit/proxy/hooks/test_proxy_hooks_init.py +++ b/tests/unit/proxy/hooks/test_proxy_hooks_init.py @@ -14,8 +14,12 @@ silently swallowed the ImportError in `hooks/__init__.py`, leaving ``managed_files`` unregistered and the /files endpoint returning 500. """ +import importlib +import sys + import pytest +import litellm.proxy from litellm.proxy.hooks import PROXY_HOOKS, get_proxy_hook @@ -49,3 +53,21 @@ def test_isolation_module_does_not_pull_in_proxy_utils(): importlib.import_module("litellm.llms.base_llm.managed_resources.isolation") assert "litellm.proxy.utils" not in sys.modules assert "litellm.proxy.management_endpoints.common_utils" not in sys.modules + + +def test_managed_hooks_register_from_litellm_enterprise_without_the_repo_only_enterprise_package( + monkeypatch: pytest.MonkeyPatch, +): + pytest.importorskip("litellm_enterprise") + from litellm_enterprise.proxy.hooks.managed_files import PROXY_LiteLLMManagedFiles + from litellm_enterprise.proxy.hooks.managed_vector_stores import PROXY_LiteLLMManagedVectorStores + + monkeypatch.setitem(sys.modules, "enterprise", None) + monkeypatch.setitem(sys.modules, "enterprise.enterprise_hooks", None) + monkeypatch.delitem(sys.modules, "litellm.proxy.hooks") + monkeypatch.delattr(litellm.proxy, "hooks") + + wheel_install_hooks = importlib.import_module("litellm.proxy.hooks") + + assert wheel_install_hooks.get_proxy_hook("managed_files") is PROXY_LiteLLMManagedFiles + assert wheel_install_hooks.get_proxy_hook("managed_vector_stores") is PROXY_LiteLLMManagedVectorStores From 35d30d6428181a11bc1ad62ce30e7ec029126889 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 20 Oct 2026 19:25:25 -0700 Subject: [PATCH 02/11] fix(proxy): register managed files hooks from the litellm_enterprise package A pip-installed litellm[proxy] could not import the repo-only enterprise folder, so the managed_files and managed_vector_stores hooks never registered and POST /v1/files with target_model_names answered 500. Register both straight from litellm_enterprise, which the proxy extra installs, and log a warning naming both features when that import fails --- litellm/proxy/hooks/__init__.py | 29 ++++++++++--------- .../proxy/hooks/autorouter_baseline_cache.py | 6 ++-- 2 files changed, 20 insertions(+), 15 deletions(-) diff --git a/litellm/proxy/hooks/__init__.py b/litellm/proxy/hooks/__init__.py index 138050ccc9a..444bc086d35 100644 --- a/litellm/proxy/hooks/__init__.py +++ b/litellm/proxy/hooks/__init__.py @@ -1,6 +1,9 @@ import os from typing import Final, Literal +from litellm._logging import verbose_proxy_logger +from litellm.integrations.custom_logger import CustomLogger + from . import * from .autorouter_baseline_cache import AutoRouterBaselineCache from .cache_control_check import ( # noqa: F401 # backwards-compatible package export @@ -32,10 +35,10 @@ from .sensitive_data_routing import ( # noqa: F401 # backwards-compatible pack ) # List of all available hooks that can be enabled. -# Defined before the enterprise import below so that any module re-imported -# transitively through `enterprise.enterprise_hooks` can resolve `PROXY_HOOKS` +# Defined before the litellm_enterprise hook imports below so that any module re-imported +# transitively through them can resolve `PROXY_HOOKS` # and `get_proxy_hook` from this partially-initialized module without circling. -PROXY_HOOKS: Final = { +PROXY_HOOKS: Final[dict[str, type[CustomLogger]]] = { # mutable-ok: shared with cycle importers, extended below "parallel_request_limiter": PROXY_MaxParallelRequestsHandler_v3, "cache_control_check": PROXY_CacheControlCheck, "responses_id_security": ResponsesIDSecurity, @@ -63,14 +66,14 @@ def get_proxy_hook( return PROXY_HOOKS[hook_name] -### CHECK IF ENTERPRISE HOOKS ARE AVAILABLE #### - try: - from enterprise.enterprise_hooks import ENTERPRISE_PROXY_HOOKS -except ImportError: - ENTERPRISE_PROXY_HOOKS = {} - - -### update PROXY_HOOKS with ENTERPRISE_PROXY_HOOKS ### - -PROXY_HOOKS.update(ENTERPRISE_PROXY_HOOKS) + from litellm_enterprise.proxy.hooks.managed_files import PROXY_LiteLLMManagedFiles + from litellm_enterprise.proxy.hooks.managed_vector_stores import PROXY_LiteLLMManagedVectorStores +except ImportError as litellm_enterprise_import_error: + verbose_proxy_logger.warning( + "Managed files and managed vector stores are disabled because litellm_enterprise failed to import: %s", + litellm_enterprise_import_error, + ) +else: + PROXY_HOOKS["managed_files"] = PROXY_LiteLLMManagedFiles + PROXY_HOOKS["managed_vector_stores"] = PROXY_LiteLLMManagedVectorStores diff --git a/litellm/proxy/hooks/autorouter_baseline_cache.py b/litellm/proxy/hooks/autorouter_baseline_cache.py index c4232609124..dc0f58d48b5 100644 --- a/litellm/proxy/hooks/autorouter_baseline_cache.py +++ b/litellm/proxy/hooks/autorouter_baseline_cache.py @@ -137,7 +137,9 @@ class AutoRouterBaselineCache(CustomLogger): return try: raw_metadata: Final = kwargs.get(get_metadata_variable_name_from_kwargs(kwargs)) - metadata: Final = _METADATA.validate_python(raw_metadata) if isinstance(raw_metadata, Mapping) else {} + metadata: Final[Mapping[str, object]] = ( + _METADATA.validate_python(raw_metadata) if isinstance(raw_metadata, Mapping) else {} + ) if metadata.get(INTERNAL_CALL_ORIGIN_METADATA_KEY): return if logging_obj.baseline_cache_context is not None: @@ -157,7 +159,7 @@ class AutoRouterBaselineCache(CustomLogger): prices: Final = _PRICES.validate_python( _effective_model_info(router, request.route.baseline_deployment_id, request.route.baseline_model) ) - params: Final = ( + params: Final[Mapping[str, object]] = ( _METADATA.validate_python(deployment.litellm_params.model_dump(mode="json")) if deployment else {} ) projected: Final = ( From dc488f4b10820163b010c69a1afeda4381652bdc Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 20 Oct 2026 20:09:46 -0700 Subject: [PATCH 03/11] test(proxy): cover the warning when litellm_enterprise cannot load --- .../unit/proxy/hooks/test_proxy_hooks_init.py | 24 +++++++++++++++++++ 1 file changed, 24 insertions(+) diff --git a/tests/unit/proxy/hooks/test_proxy_hooks_init.py b/tests/unit/proxy/hooks/test_proxy_hooks_init.py index 009d69912e0..6aeafda75a3 100644 --- a/tests/unit/proxy/hooks/test_proxy_hooks_init.py +++ b/tests/unit/proxy/hooks/test_proxy_hooks_init.py @@ -15,11 +15,13 @@ silently swallowed the ImportError in `hooks/__init__.py`, leaving """ import importlib +import logging import sys import pytest import litellm.proxy +from litellm._logging import verbose_proxy_logger from litellm.proxy.hooks import PROXY_HOOKS, get_proxy_hook @@ -71,3 +73,25 @@ def test_managed_hooks_register_from_litellm_enterprise_without_the_repo_only_en assert wheel_install_hooks.get_proxy_hook("managed_files") is PROXY_LiteLLMManagedFiles assert wheel_install_hooks.get_proxy_hook("managed_vector_stores") is PROXY_LiteLLMManagedVectorStores + + +def test_proxy_hooks_still_import_and_warn_when_litellm_enterprise_cannot_load( + monkeypatch: pytest.MonkeyPatch, + caplog: pytest.LogCaptureFixture, +): + monkeypatch.setitem(sys.modules, "litellm_enterprise.proxy.hooks.managed_files", None) + monkeypatch.delitem(sys.modules, "litellm.proxy.hooks") + monkeypatch.delattr(litellm.proxy, "hooks") + verbose_proxy_logger.addHandler(caplog.handler) + try: + with caplog.at_level(logging.WARNING, logger=verbose_proxy_logger.name): + hooks_without_enterprise = importlib.import_module("litellm.proxy.hooks") + finally: + verbose_proxy_logger.removeHandler(caplog.handler) + + assert "managed_files" not in hooks_without_enterprise.PROXY_HOOKS + assert "managed_vector_stores" not in hooks_without_enterprise.PROXY_HOOKS + assert "parallel_request_limiter" in hooks_without_enterprise.PROXY_HOOKS + assert any( + "Managed files and managed vector stores are disabled" in record.getMessage() for record in caplog.records + ) From f57972b10123220c38e108abc5400cd83263353b Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 20 Oct 2026 20:32:19 -0700 Subject: [PATCH 04/11] refactor(enterprise): drop the unused enterprise hook registry and update the Dockerfile notes --- Dockerfile | 4 +-- docker/Dockerfile.database | 4 +-- docker/Dockerfile.non_root | 4 +-- enterprise/enterprise_hooks/__init__.py | 33 ------------------------- 4 files changed, 6 insertions(+), 39 deletions(-) diff --git a/Dockerfile b/Dockerfile index 890c90fe449..ba914bd44e3 100644 --- a/Dockerfile +++ b/Dockerfile @@ -161,8 +161,8 @@ COPY --from=builder /app/docker /app/docker COPY --from=builder /app/schema.prisma /app/schema.prisma COPY --from=builder /app/litellm/proxy/prisma_migration.py /app/litellm/proxy/prisma_migration.py # enterprise/ is imported by source path at runtime (proxy_cli puts the -# working directory on sys.path; litellm/proxy/hooks resolves -# enterprise.enterprise_hooks from it) +# working directory on sys.path; the moderation and blocked user list +# callbacks import enterprise.enterprise_hooks from it) COPY --from=builder /app/enterprise /app/enterprise COPY --from=builder /app/litellm-proxy-extras /app/litellm-proxy-extras # Prisma CLI + engines are baked under /opt/prisma, a fixed path every diff --git a/docker/Dockerfile.database b/docker/Dockerfile.database index 56210458532..f1fe4a0a757 100644 --- a/docker/Dockerfile.database +++ b/docker/Dockerfile.database @@ -141,8 +141,8 @@ COPY --from=builder /app/docker /app/docker COPY --from=builder /app/schema.prisma /app/schema.prisma COPY --from=builder /app/litellm/proxy/prisma_migration.py /app/litellm/proxy/prisma_migration.py # enterprise/ is imported by source path at runtime (proxy_cli puts the -# working directory on sys.path; litellm/proxy/hooks resolves -# enterprise.enterprise_hooks from it) +# working directory on sys.path; the moderation and blocked user list +# callbacks import enterprise.enterprise_hooks from it) COPY --from=builder /app/enterprise /app/enterprise COPY --from=builder /app/litellm-proxy-extras /app/litellm-proxy-extras # Prisma CLI + engines are baked under /opt/prisma, a fixed path every diff --git a/docker/Dockerfile.non_root b/docker/Dockerfile.non_root index 086c38e0d60..0115597517a 100644 --- a/docker/Dockerfile.non_root +++ b/docker/Dockerfile.non_root @@ -147,8 +147,8 @@ COPY --from=builder /app/docker /app/docker COPY --from=builder /app/schema.prisma /app/schema.prisma COPY --from=builder /app/litellm/proxy/prisma_migration.py /app/litellm/proxy/prisma_migration.py # enterprise/ is imported by source path at runtime (proxy_cli puts the -# working directory on sys.path; litellm/proxy/hooks resolves -# enterprise.enterprise_hooks from it) +# working directory on sys.path; the moderation and blocked user list +# callbacks import enterprise.enterprise_hooks from it) COPY --from=builder /app/enterprise /app/enterprise COPY --from=builder /app/litellm-proxy-extras /app/litellm-proxy-extras # Prisma CLI + engines are baked under /opt/prisma, a fixed path every runtime diff --git a/enterprise/enterprise_hooks/__init__.py b/enterprise/enterprise_hooks/__init__.py index 6de3325f113..e69de29bb2d 100644 --- a/enterprise/enterprise_hooks/__init__.py +++ b/enterprise/enterprise_hooks/__init__.py @@ -1,33 +0,0 @@ -from typing import Dict, Literal, Type, Union - -from litellm_enterprise.proxy.hooks.managed_files import PROXY_LiteLLMManagedFiles -from litellm_enterprise.proxy.hooks.managed_vector_stores import ( - PROXY_LiteLLMManagedVectorStores, -) - -from litellm.integrations.custom_logger import CustomLogger - -ENTERPRISE_PROXY_HOOKS: Dict[str, Type[CustomLogger]] = { - "managed_files": PROXY_LiteLLMManagedFiles, - "managed_vector_stores": PROXY_LiteLLMManagedVectorStores, -} - - -def get_enterprise_proxy_hook( - hook_name: Union[ - Literal[ - "managed_files", - "managed_vector_stores", - "max_parallel_requests", - ], - str, - ], -): - """ - Factory method to get a enterprise hook instance by name - """ - if hook_name not in ENTERPRISE_PROXY_HOOKS: - raise ValueError( - f"Unknown hook: {hook_name}. Available hooks: {list(ENTERPRISE_PROXY_HOOKS.keys())}" - ) - return ENTERPRISE_PROXY_HOOKS[hook_name] From 6336e9732834c89e98b2aa99f06d92d2dfa03961 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 20 Oct 2026 21:07:52 -0700 Subject: [PATCH 05/11] test(proxy): annotate the hook registry test bindings as Final --- tests/unit/proxy/hooks/test_proxy_hooks_init.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/tests/unit/proxy/hooks/test_proxy_hooks_init.py b/tests/unit/proxy/hooks/test_proxy_hooks_init.py index 6aeafda75a3..82c2a24f4c1 100644 --- a/tests/unit/proxy/hooks/test_proxy_hooks_init.py +++ b/tests/unit/proxy/hooks/test_proxy_hooks_init.py @@ -17,6 +17,7 @@ silently swallowed the ImportError in `hooks/__init__.py`, leaving import importlib import logging import sys +from typing import Final import pytest @@ -69,7 +70,7 @@ def test_managed_hooks_register_from_litellm_enterprise_without_the_repo_only_en monkeypatch.delitem(sys.modules, "litellm.proxy.hooks") monkeypatch.delattr(litellm.proxy, "hooks") - wheel_install_hooks = importlib.import_module("litellm.proxy.hooks") + wheel_install_hooks: Final = importlib.import_module("litellm.proxy.hooks") assert wheel_install_hooks.get_proxy_hook("managed_files") is PROXY_LiteLLMManagedFiles assert wheel_install_hooks.get_proxy_hook("managed_vector_stores") is PROXY_LiteLLMManagedVectorStores @@ -85,7 +86,7 @@ def test_proxy_hooks_still_import_and_warn_when_litellm_enterprise_cannot_load( verbose_proxy_logger.addHandler(caplog.handler) try: with caplog.at_level(logging.WARNING, logger=verbose_proxy_logger.name): - hooks_without_enterprise = importlib.import_module("litellm.proxy.hooks") + hooks_without_enterprise: Final = importlib.import_module("litellm.proxy.hooks") finally: verbose_proxy_logger.removeHandler(caplog.handler) From 53ced6a7d72c79459265c9dd619676b3d3b11594 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 28 Oct 2026 16:44:27 -0700 Subject: [PATCH 06/11] fix(proxy): route managed vector store retrieve and delete to the provider store GET and DELETE /v1/vector_stores/{id} with the managed id the proxy returned on create sent that id to the provider, which answered 400. The managed vector stores hook now swaps in the provider store id and routes to the store's model for retrieve and delete, as it already did for search. --- Dockerfile | 4 +- docker/Dockerfile.database | 4 +- docker/Dockerfile.non_root | 4 +- .../proxy/hooks/managed_vector_stores.py | 100 +++++------------ .../proxy/hooks/test_managed_vector_stores.py | 105 ++++++++++++++++++ 5 files changed, 139 insertions(+), 78 deletions(-) create mode 100644 tests/unit/enterprise/proxy/hooks/test_managed_vector_stores.py diff --git a/Dockerfile b/Dockerfile index d8ec3a654ee..46a69f5d910 100644 --- a/Dockerfile +++ b/Dockerfile @@ -162,8 +162,8 @@ COPY --from=builder /app/docker /app/docker COPY --from=builder /app/schema.prisma /app/schema.prisma COPY --from=builder /app/litellm/proxy/prisma_migration.py /app/litellm/proxy/prisma_migration.py # enterprise/ is imported by source path at runtime (proxy_cli puts the -# working directory on sys.path; the moderation and blocked user list -# callbacks import enterprise.enterprise_hooks from it) +# working directory on sys.path; the moderation, blocked user list, and +# banned keywords callbacks import enterprise.enterprise_hooks from it) COPY --from=builder /app/enterprise /app/enterprise COPY --from=builder /app/litellm-proxy-extras /app/litellm-proxy-extras # Prisma CLI + engines are baked under /opt/prisma, a fixed path every diff --git a/docker/Dockerfile.database b/docker/Dockerfile.database index 198abc0a210..e2726b347b1 100644 --- a/docker/Dockerfile.database +++ b/docker/Dockerfile.database @@ -142,8 +142,8 @@ COPY --from=builder /app/docker /app/docker COPY --from=builder /app/schema.prisma /app/schema.prisma COPY --from=builder /app/litellm/proxy/prisma_migration.py /app/litellm/proxy/prisma_migration.py # enterprise/ is imported by source path at runtime (proxy_cli puts the -# working directory on sys.path; the moderation and blocked user list -# callbacks import enterprise.enterprise_hooks from it) +# working directory on sys.path; the moderation, blocked user list, and +# banned keywords callbacks import enterprise.enterprise_hooks from it) COPY --from=builder /app/enterprise /app/enterprise COPY --from=builder /app/litellm-proxy-extras /app/litellm-proxy-extras # Prisma CLI + engines are baked under /opt/prisma, a fixed path every diff --git a/docker/Dockerfile.non_root b/docker/Dockerfile.non_root index e19765ed74c..183c5141744 100644 --- a/docker/Dockerfile.non_root +++ b/docker/Dockerfile.non_root @@ -148,8 +148,8 @@ COPY --from=builder /app/docker /app/docker COPY --from=builder /app/schema.prisma /app/schema.prisma COPY --from=builder /app/litellm/proxy/prisma_migration.py /app/litellm/proxy/prisma_migration.py # enterprise/ is imported by source path at runtime (proxy_cli puts the -# working directory on sys.path; the moderation and blocked user list -# callbacks import enterprise.enterprise_hooks from it) +# working directory on sys.path; the moderation, blocked user list, and +# banned keywords callbacks import enterprise.enterprise_hooks from it) COPY --from=builder /app/enterprise /app/enterprise COPY --from=builder /app/litellm-proxy-extras /app/litellm-proxy-extras # Prisma CLI + engines are baked under /opt/prisma, a fixed path every runtime diff --git a/enterprise/litellm_enterprise/proxy/hooks/managed_vector_stores.py b/enterprise/litellm_enterprise/proxy/hooks/managed_vector_stores.py index fcaf58099b9..b95a112674d 100644 --- a/enterprise/litellm_enterprise/proxy/hooks/managed_vector_stores.py +++ b/enterprise/litellm_enterprise/proxy/hooks/managed_vector_stores.py @@ -12,6 +12,9 @@ from litellm._uuid import uuid from litellm.integrations.custom_logger import CustomLogger from litellm.llms.base_llm.managed_resources import BaseManagedResource from litellm.llms.base_llm.managed_resources.utils import ( + extract_model_id_from_unified_id, + extract_provider_resource_id_from_unified_id, + extract_target_model_names_from_unified_id, generate_unified_id_string, is_base64_encoded_unified_id, ) @@ -38,6 +41,27 @@ else: PrismaClient = Any +_ROUTED_VECTOR_STORE_CALL_TYPES: Final = frozenset( + {"avector_store_search", "avector_store_retrieve", "avector_store_delete"} +) + + +def _route_to_provider_vector_store(data: dict[str, object]) -> dict[str, object]: + vector_store_id: Final = data.get("vector_store_id") + decoded_id: Final = is_base64_encoded_unified_id(vector_store_id) if isinstance(vector_store_id, str) else False + if not decoded_id: + return data + routing_model: Final = extract_model_id_from_unified_id(decoded_id) or next( + iter(extract_target_model_names_from_unified_id(decoded_id)), None + ) + provider_resource_id: Final = extract_provider_resource_id_from_unified_id(decoded_id) + return { + **data, + **({"model": routing_model} if routing_model else {}), + **({"vector_store_id": provider_resource_id} if provider_resource_id else {}), + } + + class _PROXY_LiteLLMManagedVectorStores( CustomLogger, BaseManagedResource[VectorStoreCreateResponse] ): @@ -322,78 +346,10 @@ class _PROXY_LiteLLMManagedVectorStores( Returns: Modified request data or None """ - from litellm.llms.base_llm.managed_resources.utils import ( - is_base64_encoded_unified_id, - parse_unified_id, - ) - - # Handle vector store search operations - if call_type == "avector_store_search": - vector_store_id = data.get("vector_store_id") - - if vector_store_id: - # Check if it's a managed vector store ID - decoded_id = is_base64_encoded_unified_id(vector_store_id) - - if decoded_id: - verbose_logger.debug( - f"Processing managed vector store search: {vector_store_id}" - ) - - # Check access - has_access = await self.can_user_access_unified_resource_id( - vector_store_id, user_api_key_dict - ) - - if not has_access: - raise HTTPException( - status_code=403, - detail=f"User {user_api_key_dict.user_id} does not have access to vector store {vector_store_id}", - ) - - # Parse the unified ID to extract components - parsed_id = parse_unified_id(vector_store_id) - - if parsed_id: - # Extract the model ID and provider resource ID - model_id = parsed_id.get("model_id") - provider_resource_id = parsed_id.get("provider_resource_id") - target_model_names = parsed_id.get("target_model_names", []) - - verbose_logger.debug( - f"Decoded vector store - model_id: {model_id}, provider_resource_id: {provider_resource_id}, target_model_names: {target_model_names}" - ) - - # Determine which model to use for routing - # Priority: model_id (deployment ID) > first target_model_name - routing_model = None - if model_id: - routing_model = model_id - elif target_model_names and len(target_model_names) > 0: - routing_model = target_model_names[0] - - # Set the model for routing - if routing_model: - data["model"] = routing_model - verbose_logger.info( - f"Routing vector store search to model: {routing_model}" - ) - - # Replace the unified ID with the provider-specific ID - if provider_resource_id: - data["vector_store_id"] = provider_resource_id - verbose_logger.debug( - f"Replaced unified ID with provider resource ID: {provider_resource_id}" - ) - - # Handle vector store retrieve/delete operations - elif call_type in ("avector_store_retrieve", "avector_store_delete"): - await self.check_managed_vector_store_access(data, user_api_key_dict) - - # If it's a managed vector store, we'll handle it in the endpoint - # No need to transform here as the endpoint will route to the hook - - return data + if call_type not in _ROUTED_VECTOR_STORE_CALL_TYPES: + return data + await self.check_managed_vector_store_access(data, user_api_key_dict) + return _route_to_provider_vector_store(data) # ============================================================================ # POST-CALL HOOK (For Response Transformation) diff --git a/tests/unit/enterprise/proxy/hooks/test_managed_vector_stores.py b/tests/unit/enterprise/proxy/hooks/test_managed_vector_stores.py new file mode 100644 index 00000000000..169064716d9 --- /dev/null +++ b/tests/unit/enterprise/proxy/hooks/test_managed_vector_stores.py @@ -0,0 +1,105 @@ +from collections.abc import Mapping +from dataclasses import dataclass +from typing import Final + +import pytest +from fastapi import HTTPException +from litellm_enterprise.proxy.hooks.managed_vector_stores import PROXY_LiteLLMManagedVectorStores + +from litellm.caching import DualCache +from litellm.proxy._types import UserAPIKeyAuth +from litellm.types.vector_stores import VectorStoreCreateResponse + +_OWNER: Final = "vector-store-owner" +_PROVIDER_ID: Final = "vs_provider_store" +_ROUTED_CALL_TYPES: Final = ("avector_store_search", "avector_store_retrieve", "avector_store_delete") + + +@dataclass(frozen=True, slots=True) +class _StoreRow: + created_by: str + + def model_dump(self) -> dict[str, str | None]: + return {"created_by": self.created_by, "team_id": None} + + +class _SingleStoreTable: + async def find_first(self, where: Mapping[str, object]) -> _StoreRow: + return _StoreRow(created_by=_OWNER) + + +class _Database: + def __init__(self) -> None: + self.litellm_managedvectorstoretable: Final = _SingleStoreTable() + + +class _PrismaClient: + def __init__(self) -> None: + self.db: Final = _Database() + + +def _hook_and_unified_id() -> tuple[PROXY_LiteLLMManagedVectorStores, str]: + hook: Final = PROXY_LiteLLMManagedVectorStores(DualCache(), prisma_client=_PrismaClient()) + created: Final = VectorStoreCreateResponse( + id=_PROVIDER_ID, + object="vector_store", + created_at=1700000000, + name="managed-store", + bytes=0, + file_counts={"in_progress": 0, "completed": 0, "failed": 0, "cancelled": 0, "total": 0}, + status="completed", + expires_after=None, + expires_at=None, + last_active_at=None, + metadata={}, + ) + unified_id: Final = hook.generate_unified_resource_id( + resource_objects=[created], target_model_names_list=["model-a", "model-b"] + ) + return hook, unified_id + + +@pytest.mark.asyncio +@pytest.mark.parametrize("call_type", _ROUTED_CALL_TYPES) +async def test_owner_request_goes_to_provider_store_on_first_target_model(call_type: str) -> None: + hook, unified_id = _hook_and_unified_id() + + routed: Final = await hook.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(user_id=_OWNER), + cache=DualCache(), + data={"vector_store_id": unified_id}, + call_type=call_type, + ) + + assert routed == {"vector_store_id": _PROVIDER_ID, "model": "model-a"} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("call_type", _ROUTED_CALL_TYPES) +async def test_other_user_is_refused_before_reaching_provider(call_type: str) -> None: + hook, unified_id = _hook_and_unified_id() + + with pytest.raises(HTTPException) as refused: + await hook.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(user_id="someone-else"), + cache=DualCache(), + data={"vector_store_id": unified_id}, + call_type=call_type, + ) + + assert refused.value.status_code == 403 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("call_type", _ROUTED_CALL_TYPES) +async def test_provider_store_id_passes_through_unchanged(call_type: str) -> None: + hook, _ = _hook_and_unified_id() + + routed: Final = await hook.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(user_id="someone-else"), + cache=DualCache(), + data={"vector_store_id": "vs_raw_store"}, + call_type=call_type, + ) + + assert routed == {"vector_store_id": "vs_raw_store"} From d031d90b900fd91b2fbe2c965372f9547c9f0dca Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 28 Oct 2026 20:02:26 -0700 Subject: [PATCH 07/11] fix: route vector store update by managed id, refuse multi-model writes, and keep the enterprise warning off SDK installs Update by managed id now reaches the provider store with the ownership check, like retrieve, search and delete. Update or delete by the managed id of a store created on several models answers 400, since it would only change the first model's copy. The missing litellm_enterprise message logs at debug when the package is not installed. The Dockerfile comments now say only why enterprise/ is copied. --- Dockerfile | 4 +- docker/Dockerfile.database | 4 +- docker/Dockerfile.non_root | 4 +- .../proxy/hooks/managed_vector_stores.py | 32 +++++++--- litellm/proxy/hooks/__init__.py | 7 ++- .../test_managed_vector_store_unified_id.py | 63 +++++++++++++++++++ .../proxy/hooks/test_managed_vector_stores.py | 45 +++++++++++-- .../unit/proxy/hooks/test_proxy_hooks_init.py | 30 +++++++++ 8 files changed, 165 insertions(+), 24 deletions(-) diff --git a/Dockerfile b/Dockerfile index 46a69f5d910..a7874460dfb 100644 --- a/Dockerfile +++ b/Dockerfile @@ -161,9 +161,7 @@ COPY --from=liteadmin-builder /opt/liteadmin /opt/liteadmin COPY --from=builder /app/docker /app/docker COPY --from=builder /app/schema.prisma /app/schema.prisma COPY --from=builder /app/litellm/proxy/prisma_migration.py /app/litellm/proxy/prisma_migration.py -# enterprise/ is imported by source path at runtime (proxy_cli puts the -# working directory on sys.path; the moderation, blocked user list, and -# banned keywords callbacks import enterprise.enterprise_hooks from it) +# enterprise/ is imported by source path at runtime (proxy_cli puts the working directory on sys.path) COPY --from=builder /app/enterprise /app/enterprise COPY --from=builder /app/litellm-proxy-extras /app/litellm-proxy-extras # Prisma CLI + engines are baked under /opt/prisma, a fixed path every diff --git a/docker/Dockerfile.database b/docker/Dockerfile.database index e2726b347b1..f7314f700e4 100644 --- a/docker/Dockerfile.database +++ b/docker/Dockerfile.database @@ -141,9 +141,7 @@ COPY --from=builder /app/.venv /app/.venv COPY --from=builder /app/docker /app/docker COPY --from=builder /app/schema.prisma /app/schema.prisma COPY --from=builder /app/litellm/proxy/prisma_migration.py /app/litellm/proxy/prisma_migration.py -# enterprise/ is imported by source path at runtime (proxy_cli puts the -# working directory on sys.path; the moderation, blocked user list, and -# banned keywords callbacks import enterprise.enterprise_hooks from it) +# enterprise/ is imported by source path at runtime (proxy_cli puts the working directory on sys.path) COPY --from=builder /app/enterprise /app/enterprise COPY --from=builder /app/litellm-proxy-extras /app/litellm-proxy-extras # Prisma CLI + engines are baked under /opt/prisma, a fixed path every diff --git a/docker/Dockerfile.non_root b/docker/Dockerfile.non_root index 183c5141744..0631bd14496 100644 --- a/docker/Dockerfile.non_root +++ b/docker/Dockerfile.non_root @@ -147,9 +147,7 @@ COPY --from=builder /app/.venv /app/.venv COPY --from=builder /app/docker /app/docker COPY --from=builder /app/schema.prisma /app/schema.prisma COPY --from=builder /app/litellm/proxy/prisma_migration.py /app/litellm/proxy/prisma_migration.py -# enterprise/ is imported by source path at runtime (proxy_cli puts the -# working directory on sys.path; the moderation, blocked user list, and -# banned keywords callbacks import enterprise.enterprise_hooks from it) +# enterprise/ is imported by source path at runtime (proxy_cli puts the working directory on sys.path) COPY --from=builder /app/enterprise /app/enterprise COPY --from=builder /app/litellm-proxy-extras /app/litellm-proxy-extras # Prisma CLI + engines are baked under /opt/prisma, a fixed path every runtime diff --git a/enterprise/litellm_enterprise/proxy/hooks/managed_vector_stores.py b/enterprise/litellm_enterprise/proxy/hooks/managed_vector_stores.py index b95a112674d..7ff50363521 100644 --- a/enterprise/litellm_enterprise/proxy/hooks/managed_vector_stores.py +++ b/enterprise/litellm_enterprise/proxy/hooks/managed_vector_stores.py @@ -41,16 +41,14 @@ else: PrismaClient = Any -_ROUTED_VECTOR_STORE_CALL_TYPES: Final = frozenset( - {"avector_store_search", "avector_store_retrieve", "avector_store_delete"} -) +_WRITE_VECTOR_STORE_CALL_TYPES: Final = frozenset({"avector_store_update", "avector_store_delete"}) +_ROUTED_VECTOR_STORE_CALL_TYPES: Final = _WRITE_VECTOR_STORE_CALL_TYPES | { + "avector_store_search", + "avector_store_retrieve", +} -def _route_to_provider_vector_store(data: dict[str, object]) -> dict[str, object]: - vector_store_id: Final = data.get("vector_store_id") - decoded_id: Final = is_base64_encoded_unified_id(vector_store_id) if isinstance(vector_store_id, str) else False - if not decoded_id: - return data +def _route_to_provider_vector_store(data: dict[str, object], decoded_id: str) -> dict[str, object]: routing_model: Final = extract_model_id_from_unified_id(decoded_id) or next( iter(extract_target_model_names_from_unified_id(decoded_id)), None ) @@ -71,7 +69,7 @@ class _PROXY_LiteLLMManagedVectorStores( This class provides functionality to: - Create vector stores across multiple models - Retrieve vector stores by unified ID - - Delete vector stores from all models + - Update or delete a vector store created on one model by unified ID - List vector stores created by a user """ @@ -349,7 +347,21 @@ class _PROXY_LiteLLMManagedVectorStores( if call_type not in _ROUTED_VECTOR_STORE_CALL_TYPES: return data await self.check_managed_vector_store_access(data, user_api_key_dict) - return _route_to_provider_vector_store(data) + vector_store_id: Final = data.get("vector_store_id") + decoded_id: Final = is_base64_encoded_unified_id(vector_store_id) if isinstance(vector_store_id, str) else False + if not decoded_id: + return data + created_on_several_models: Final = len(extract_target_model_names_from_unified_id(decoded_id)) > 1 + if call_type in _WRITE_VECTOR_STORE_CALL_TYPES and created_on_several_models: + raise HTTPException( + status_code=400, + detail=( + f"Vector store {vector_store_id} was created on several models, so updating or deleting it by " + "this id would only change the first model's store. Update or delete each model's store by its " + "provider id instead." + ), + ) + return _route_to_provider_vector_store(data, decoded_id) # ============================================================================ # POST-CALL HOOK (For Response Transformation) diff --git a/litellm/proxy/hooks/__init__.py b/litellm/proxy/hooks/__init__.py index 444bc086d35..089dca7ec11 100644 --- a/litellm/proxy/hooks/__init__.py +++ b/litellm/proxy/hooks/__init__.py @@ -70,7 +70,12 @@ try: from litellm_enterprise.proxy.hooks.managed_files import PROXY_LiteLLMManagedFiles from litellm_enterprise.proxy.hooks.managed_vector_stores import PROXY_LiteLLMManagedVectorStores except ImportError as litellm_enterprise_import_error: - verbose_proxy_logger.warning( + log_disabled_managed_hooks: Final = ( + verbose_proxy_logger.debug + if litellm_enterprise_import_error.name == "litellm_enterprise" + else verbose_proxy_logger.warning + ) + log_disabled_managed_hooks( "Managed files and managed vector stores are disabled because litellm_enterprise failed to import: %s", litellm_enterprise_import_error, ) diff --git a/tests/integration/management/test_managed_vector_store_unified_id.py b/tests/integration/management/test_managed_vector_store_unified_id.py index 0b0e51c0871..1c3e4a5a9d2 100644 --- a/tests/integration/management/test_managed_vector_store_unified_id.py +++ b/tests/integration/management/test_managed_vector_store_unified_id.py @@ -265,3 +265,66 @@ def test_managed_vector_store_id_routes_to_creator_deployment_and_enforces_owner db_row_after_denial: Final = read_rows(_MANAGED_STORE_ROW, (unified_id,)) assert db_row_after_denial == [{"created_by": creator_user, "team_id": creator_team}] + + update_refused: Final = gateway.request( + "POST", f"/v1/vector_stores/{unified_id}", {"name": "renamed"}, key=creator_key + ) + assert update_refused.status_code == 400, update_refused.text + delete_refused: Final = gateway.request("DELETE", f"/v1/vector_stores/{unified_id}", key=creator_key) + assert delete_refused.status_code == 400, delete_refused.text + assert wire_a.drain() == () + assert wire_b.drain() == () + + +def test_managed_vector_store_id_retrieves_updates_and_deletes_the_provider_store(gateway: Gateway) -> None: + store_name: Final = f"managed-store-{uuid.uuid4().hex}" + provider_id: Final = f"vs_provider_{uuid.uuid4().hex}" + provider_key: Final = f"provider-{uuid.uuid4().hex}" + store_path: Final = f"/v1/vector_stores/{provider_id}" + + def respond(request: Request) -> Reply: + path: Final = urlsplit(request.target).path + assert request.headers["authorization"] == f"Bearer {provider_key}" + if request.method == "GET" and path == "/v1/models": + return _models_reply() + if request.method == "POST" and path == "/v1/vector_stores": + _assert_create_request(request, store_name) + return _create_reply(provider_id, store_name) + assert path == store_path, request.target + if request.method == "GET": + return _create_reply(provider_id, store_name) + if request.method == "POST": + assert JSON_OBJECT.validate_json(request.body)["name"] == "renamed" + return _create_reply(provider_id, "renamed") + assert request.method == "DELETE" + return Reply(body=json.dumps({"id": provider_id, "object": "vector_store.deleted", "deleted": True}).encode()) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(api_base=f"{wire.url}/v1", api_key=provider_key) + _wait_until_every_worker_serves(gateway, model) + team: Final = scenario.team(models=[model]) + user: Final = scenario.member(team) + key: Final = scenario.key(team_id=team, user_id=user, models=[model]) + + with OpenAI( + base_url=f"{str(gateway.client.base_url).rstrip('/')}/v1", + api_key=key, + http_client=httpx.Client(trust_env=False), + max_retries=0, + ) as client: + created: Final = client.vector_stores.create(name=store_name, extra_body={"target_model_names": model}) + scenario.cleanups.callback(_delete_managed_store_row, created.id) + assert len(_store_creation_requests(wire)) == 1 + + retrieved: Final = client.vector_stores.retrieve(created.id) + updated: Final = client.vector_stores.update(created.id, name="renamed") + deleted: Final = client.vector_stores.delete(created.id) + + assert (retrieved.id, retrieved.name) == (provider_id, store_name) + assert (updated.id, updated.name) == (provider_id, "renamed") + assert (deleted.id, deleted.deleted) == (provider_id, True) + assert [(request.method, urlsplit(request.target).path) for request in wire.drain()] == [ + ("GET", store_path), + ("POST", store_path), + ("DELETE", store_path), + ] diff --git a/tests/unit/enterprise/proxy/hooks/test_managed_vector_stores.py b/tests/unit/enterprise/proxy/hooks/test_managed_vector_stores.py index 169064716d9..cdc094973b6 100644 --- a/tests/unit/enterprise/proxy/hooks/test_managed_vector_stores.py +++ b/tests/unit/enterprise/proxy/hooks/test_managed_vector_stores.py @@ -12,7 +12,9 @@ from litellm.types.vector_stores import VectorStoreCreateResponse _OWNER: Final = "vector-store-owner" _PROVIDER_ID: Final = "vs_provider_store" -_ROUTED_CALL_TYPES: Final = ("avector_store_search", "avector_store_retrieve", "avector_store_delete") +_READ_CALL_TYPES: Final = ("avector_store_search", "avector_store_retrieve") +_WRITE_CALL_TYPES: Final = ("avector_store_update", "avector_store_delete") +_ROUTED_CALL_TYPES: Final = _READ_CALL_TYPES + _WRITE_CALL_TYPES @dataclass(frozen=True, slots=True) @@ -38,7 +40,9 @@ class _PrismaClient: self.db: Final = _Database() -def _hook_and_unified_id() -> tuple[PROXY_LiteLLMManagedVectorStores, str]: +def _hook_and_unified_id( + target_model_names: tuple[str, ...] = ("model-a",), +) -> tuple[PROXY_LiteLLMManagedVectorStores, str]: hook: Final = PROXY_LiteLLMManagedVectorStores(DualCache(), prisma_client=_PrismaClient()) created: Final = VectorStoreCreateResponse( id=_PROVIDER_ID, @@ -54,14 +58,14 @@ def _hook_and_unified_id() -> tuple[PROXY_LiteLLMManagedVectorStores, str]: metadata={}, ) unified_id: Final = hook.generate_unified_resource_id( - resource_objects=[created], target_model_names_list=["model-a", "model-b"] + resource_objects=[created], target_model_names_list=list(target_model_names) ) return hook, unified_id @pytest.mark.asyncio @pytest.mark.parametrize("call_type", _ROUTED_CALL_TYPES) -async def test_owner_request_goes_to_provider_store_on_first_target_model(call_type: str) -> None: +async def test_owner_request_goes_to_provider_store_on_its_model(call_type: str) -> None: hook, unified_id = _hook_and_unified_id() routed: Final = await hook.async_pre_call_hook( @@ -74,6 +78,39 @@ async def test_owner_request_goes_to_provider_store_on_first_target_model(call_t assert routed == {"vector_store_id": _PROVIDER_ID, "model": "model-a"} +@pytest.mark.asyncio +@pytest.mark.parametrize("call_type", _READ_CALL_TYPES) +async def test_reading_store_created_on_several_models_goes_to_first_target_model(call_type: str) -> None: + hook, unified_id = _hook_and_unified_id(("model-a", "model-b")) + + routed: Final = await hook.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(user_id=_OWNER), + cache=DualCache(), + data={"vector_store_id": unified_id}, + call_type=call_type, + ) + + assert routed == {"vector_store_id": _PROVIDER_ID, "model": "model-a"} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("call_type", _WRITE_CALL_TYPES) +async def test_changing_store_created_on_several_models_is_refused_instead_of_changing_one_copy( + call_type: str, +) -> None: + hook, unified_id = _hook_and_unified_id(("model-a", "model-b")) + + with pytest.raises(HTTPException) as refused: + await hook.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(user_id=_OWNER), + cache=DualCache(), + data={"vector_store_id": unified_id}, + call_type=call_type, + ) + + assert refused.value.status_code == 400 + + @pytest.mark.asyncio @pytest.mark.parametrize("call_type", _ROUTED_CALL_TYPES) async def test_other_user_is_refused_before_reaching_provider(call_type: str) -> None: diff --git a/tests/unit/proxy/hooks/test_proxy_hooks_init.py b/tests/unit/proxy/hooks/test_proxy_hooks_init.py index 82c2a24f4c1..ce4559e4c2a 100644 --- a/tests/unit/proxy/hooks/test_proxy_hooks_init.py +++ b/tests/unit/proxy/hooks/test_proxy_hooks_init.py @@ -96,3 +96,33 @@ def test_proxy_hooks_still_import_and_warn_when_litellm_enterprise_cannot_load( assert any( "Managed files and managed vector stores are disabled" in record.getMessage() for record in caplog.records ) + + +class _LitellmEnterpriseNotInstalled: + def find_spec(self, fullname: str, path: object, target: object = None) -> None: + if fullname.partition(".")[0] == "litellm_enterprise": + raise ModuleNotFoundError(f"No module named {fullname!r}", name=fullname) + return None + + +def test_proxy_hooks_import_without_a_warning_when_litellm_enterprise_is_not_installed( + monkeypatch: pytest.MonkeyPatch, + caplog: pytest.LogCaptureFixture, +): + for enterprise_module in [name for name in sys.modules if name.partition(".")[0] == "litellm_enterprise"]: + monkeypatch.delitem(sys.modules, enterprise_module) + monkeypatch.setattr(sys, "meta_path", [_LitellmEnterpriseNotInstalled(), *sys.meta_path]) + monkeypatch.delitem(sys.modules, "litellm.proxy.hooks") + monkeypatch.delattr(litellm.proxy, "hooks") + verbose_proxy_logger.addHandler(caplog.handler) + try: + with caplog.at_level(logging.WARNING, logger=verbose_proxy_logger.name): + sdk_only_hooks: Final = importlib.import_module("litellm.proxy.hooks") + finally: + verbose_proxy_logger.removeHandler(caplog.handler) + + assert "managed_files" not in sdk_only_hooks.PROXY_HOOKS + assert "parallel_request_limiter" in sdk_only_hooks.PROXY_HOOKS + assert not any( + "Managed files and managed vector stores are disabled" in record.getMessage() for record in caplog.records + ) From a8fe887b6414f2c38d735ea0af88d58864ac3616 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 28 Oct 2026 22:07:33 -0700 Subject: [PATCH 08/11] fix: keep the managed id in vector store answers, forget a deleted store, and name the database in hook-missing errors --- .../proxy/hooks/managed_vector_stores.py | 28 +++--- litellm/proxy/hooks/batch_rate_limiter.py | 3 +- .../openai_files_endpoints/files_endpoints.py | 8 +- .../proxy/vector_store_endpoints/endpoints.py | 38 +++++++- litellm/types/proxy/common_proxy_errors.py | 8 ++ .../test_managed_vector_store_unified_id.py | 14 ++- .../proxy/hooks/test_managed_vector_stores.py | 41 +++++++-- .../test_files_endpoint.py | 11 +++ .../test_vector_store_endpoints.py | 90 +++++++++++++++++++ 9 files changed, 209 insertions(+), 32 deletions(-) diff --git a/enterprise/litellm_enterprise/proxy/hooks/managed_vector_stores.py b/enterprise/litellm_enterprise/proxy/hooks/managed_vector_stores.py index 7ff50363521..1f68e0b188a 100644 --- a/enterprise/litellm_enterprise/proxy/hooks/managed_vector_stores.py +++ b/enterprise/litellm_enterprise/proxy/hooks/managed_vector_stores.py @@ -11,6 +11,7 @@ from litellm import Router, verbose_logger from litellm._uuid import uuid from litellm.integrations.custom_logger import CustomLogger from litellm.llms.base_llm.managed_resources import BaseManagedResource +from litellm.llms.base_llm.managed_resources.isolation import can_access_resource from litellm.llms.base_llm.managed_resources.utils import ( extract_model_id_from_unified_id, extract_provider_resource_id_from_unified_id, @@ -303,18 +304,21 @@ class _PROXY_LiteLLMManagedVectorStores( else False ) - if is_unified_id and vector_store_id: - if await self.can_user_access_unified_resource_id( - vector_store_id, user_api_key_dict - ): - return True - else: - raise HTTPException( - status_code=403, - detail=f"User {user_api_key_dict.user_id} does not have access to vector store {vector_store_id}", - ) - - return False + if not (is_unified_id and vector_store_id): + return False + managed_store: Final = await self.get_unified_resource_id(vector_store_id, user_api_key_dict.parent_otel_span) + if managed_store is None: + raise HTTPException(status_code=404, detail=f"Vector store not found: {vector_store_id}") + if can_access_resource( + user_api_key_dict=user_api_key_dict, + created_by=managed_store.get("created_by"), + resource_team_id=managed_store.get("team_id"), + ): + return True + raise HTTPException( + status_code=403, + detail=f"User {user_api_key_dict.user_id} does not have access to vector store {vector_store_id}", + ) # ============================================================================ # PRE-CALL HOOK (For Router Integration) diff --git a/litellm/proxy/hooks/batch_rate_limiter.py b/litellm/proxy/hooks/batch_rate_limiter.py index 94762232e0b..7896010d89d 100644 --- a/litellm/proxy/hooks/batch_rate_limiter.py +++ b/litellm/proxy/hooks/batch_rate_limiter.py @@ -40,6 +40,7 @@ from litellm.constants import BATCH_TPD_DESCRIPTOR_SUFFIX, BATCH_TPD_WINDOW_SECO from litellm.exceptions import RateLimitErrorCategory from litellm.integrations.custom_logger import CustomLogger from litellm.proxy._types import ( + CommonProxyErrors, ProxyErrorTypes, ProxyException, SpecialModelNames, @@ -1221,7 +1222,7 @@ class _PROXY_BatchRateLimiter(CustomLogger): managed_files_obj: Final = proxy_logging_obj.get_proxy_hook("managed_files") if managed_files_obj is None: - raise ValueError("Managed files hook not found. Cannot access managed file.") + raise ValueError(CommonProxyErrors.managed_files_unavailable.value) if not isinstance(managed_files_obj, BaseFileEndpoints): raise ValueError("Managed files hook is not a BaseFileEndpoints instance.") diff --git a/litellm/proxy/openai_files_endpoints/files_endpoints.py b/litellm/proxy/openai_files_endpoints/files_endpoints.py index d2ceefffe47..0fc8f155379 100644 --- a/litellm/proxy/openai_files_endpoints/files_endpoints.py +++ b/litellm/proxy/openai_files_endpoints/files_endpoints.py @@ -482,7 +482,7 @@ async def route_create_file( managed_files_obj: Final = proxy_logging_obj.get_proxy_hook("managed_files") if managed_files_obj is None: raise ProxyException( - message="Managed files hook not found", + message=CommonProxyErrors.managed_files_unavailable.value, type=ProxyErrorTypes.internal_server_error.value, param=None, code=500, @@ -1027,7 +1027,7 @@ async def get_file_content( managed_files_obj: Final = proxy_logging_obj.get_proxy_hook("managed_files") if managed_files_obj is None: raise ProxyException( - message="Managed files hook not found", + message=CommonProxyErrors.managed_files_unavailable.value, type=ProxyErrorTypes.internal_server_error.value, param=None, code=500, @@ -1369,7 +1369,7 @@ async def get_file( managed_files_obj: Final = proxy_logging_obj.get_proxy_hook("managed_files") if managed_files_obj is None: raise ProxyException( - message="Managed files hook not found", + message=CommonProxyErrors.managed_files_unavailable.value, type=ProxyErrorTypes.internal_server_error.value, param=None, code=500, @@ -1579,7 +1579,7 @@ async def delete_file( managed_files_obj: Final = proxy_logging_obj.get_proxy_hook("managed_files") if managed_files_obj is None: raise ProxyException( - message="Managed files hook not found", + message=CommonProxyErrors.managed_files_unavailable.value, type=ProxyErrorTypes.internal_server_error.value, param=None, code=500, diff --git a/litellm/proxy/vector_store_endpoints/endpoints.py b/litellm/proxy/vector_store_endpoints/endpoints.py index 1b54dddcf4f..4c7f0a2124e 100644 --- a/litellm/proxy/vector_store_endpoints/endpoints.py +++ b/litellm/proxy/vector_store_endpoints/endpoints.py @@ -10,9 +10,15 @@ from typing import ( from fastapi import APIRouter, Depends, HTTPException, Request, Response from litellm.constants import VECTOR_STORE_CALLER_METADATA_KWARG +from litellm.integrations.custom_logger import CustomLogger from litellm.integrations.vector_store_integrations.vector_store_pre_call_hook import ( LiteLLM_ManagedVectorStore, ) +from litellm.llms.base_llm.managed_resources import BaseManagedResource +from litellm.llms.base_llm.managed_resources.utils import ( + extract_provider_resource_id_from_unified_id, + is_base64_encoded_unified_id, +) from litellm.proxy._types import CommonProxyErrors, UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing @@ -240,7 +246,7 @@ async def vector_store_create( if managed_vector_stores is None: raise HTTPException( status_code=500, - detail="Managed vector stores not configured. Please ensure the proxy is initialized with database support.", + detail=CommonProxyErrors.managed_vector_stores_unavailable.value, ) if llm_router is None: @@ -289,6 +295,21 @@ async def vector_store_create( ) +async def _answer_with_managed_vector_store_id( + response: object, + vector_store_id: str, + managed_vector_stores: CustomLogger | None, +) -> object: + if not isinstance(managed_vector_stores, BaseManagedResource) or not isinstance(response, dict): + return response + decoded_id: Final = is_base64_encoded_unified_id(vector_store_id) + if not decoded_id or response.get("id") != extract_provider_resource_id_from_unified_id(decoded_id): + return response + if response.get("deleted") is True: + await managed_vector_stores.delete_unified_resource_id(vector_store_id) + return {**response, "id": vector_store_id} + + @router.get("/v1/vector_stores/{vector_store_id}", dependencies=[Depends(user_api_key_auth)]) @router.get("/vector_stores/{vector_store_id}", dependencies=[Depends(user_api_key_auth)]) async def vector_store_retrieve( @@ -325,7 +346,7 @@ async def vector_store_retrieve( processor: Final = ProxyBaseLLMRequestProcessing(data=data) try: - return await processor.base_process_llm_request( + response: Final = await processor.base_process_llm_request( request=request, fastapi_response=fastapi_response, user_api_key_dict=user_api_key_dict, @@ -343,6 +364,9 @@ async def vector_store_retrieve( user_api_base=user_api_base, version=version, ) + return await _answer_with_managed_vector_store_id( + response, vector_store_id, proxy_logging_obj.get_proxy_hook("managed_vector_stores") + ) except Exception as e: raise await processor.handle_llm_api_exception( e=e, @@ -464,7 +488,7 @@ async def vector_store_update( processor: Final = ProxyBaseLLMRequestProcessing(data=data) try: - return await processor.base_process_llm_request( + response: Final = await processor.base_process_llm_request( request=request, fastapi_response=fastapi_response, user_api_key_dict=user_api_key_dict, @@ -482,6 +506,9 @@ async def vector_store_update( user_api_base=user_api_base, version=version, ) + return await _answer_with_managed_vector_store_id( + response, vector_store_id, proxy_logging_obj.get_proxy_hook("managed_vector_stores") + ) except Exception as e: raise await processor.handle_llm_api_exception( e=e, @@ -527,7 +554,7 @@ async def vector_store_delete( processor: Final = ProxyBaseLLMRequestProcessing(data=data) try: - return await processor.base_process_llm_request( + response: Final = await processor.base_process_llm_request( request=request, fastapi_response=fastapi_response, user_api_key_dict=user_api_key_dict, @@ -545,6 +572,9 @@ async def vector_store_delete( user_api_base=user_api_base, version=version, ) + return await _answer_with_managed_vector_store_id( + response, vector_store_id, proxy_logging_obj.get_proxy_hook("managed_vector_stores") + ) except Exception as e: raise await processor.handle_llm_api_exception( e=e, diff --git a/litellm/types/proxy/common_proxy_errors.py b/litellm/types/proxy/common_proxy_errors.py index 51426220c82..0fbbcc5fb2e 100644 --- a/litellm/types/proxy/common_proxy_errors.py +++ b/litellm/types/proxy/common_proxy_errors.py @@ -13,3 +13,11 @@ class CommonProxyErrors(str, enum.Enum): max_parallel_request_limit_reached = "Crossed TPM / RPM / Max Parallel Request Limit" missing_enterprise_package = "Missing litellm-enterprise package. Please install it to use this feature. Run `pip install litellm-enterprise`" missing_enterprise_package_docker = "This uses the enterprise folder - only available on the Docker image." + managed_files_unavailable = ( + "Managed files need a database (set DATABASE_URL) and the litellm-enterprise package, " + "and this proxy is missing one of them" + ) + managed_vector_stores_unavailable = ( + "Managed vector stores need a database (set DATABASE_URL) and the litellm-enterprise package, " + "and this proxy is missing one of them" + ) diff --git a/tests/integration/management/test_managed_vector_store_unified_id.py b/tests/integration/management/test_managed_vector_store_unified_id.py index 1c3e4a5a9d2..440a9f94fa7 100644 --- a/tests/integration/management/test_managed_vector_store_unified_id.py +++ b/tests/integration/management/test_managed_vector_store_unified_id.py @@ -1,5 +1,6 @@ import base64 import json +import threading import uuid from concurrent.futures import ThreadPoolExecutor from functools import partial @@ -281,6 +282,7 @@ def test_managed_vector_store_id_retrieves_updates_and_deletes_the_provider_stor provider_id: Final = f"vs_provider_{uuid.uuid4().hex}" provider_key: Final = f"provider-{uuid.uuid4().hex}" store_path: Final = f"/v1/vector_stores/{provider_id}" + store_deleted: Final = threading.Event() def respond(request: Request) -> Reply: path: Final = urlsplit(request.target).path @@ -291,12 +293,15 @@ def test_managed_vector_store_id_retrieves_updates_and_deletes_the_provider_stor _assert_create_request(request, store_name) return _create_reply(provider_id, store_name) assert path == store_path, request.target + if request.method == "GET" and store_deleted.is_set(): + return Reply(status=404, body=json.dumps({"error": {"message": "No vector store found"}}).encode()) if request.method == "GET": return _create_reply(provider_id, store_name) if request.method == "POST": assert JSON_OBJECT.validate_json(request.body)["name"] == "renamed" return _create_reply(provider_id, "renamed") assert request.method == "DELETE" + store_deleted.set() return Reply(body=json.dumps({"id": provider_id, "object": "vector_store.deleted", "deleted": True}).encode()) with wire_server(respond) as wire, gateway.scenario() as scenario: @@ -320,11 +325,14 @@ def test_managed_vector_store_id_retrieves_updates_and_deletes_the_provider_stor updated: Final = client.vector_stores.update(created.id, name="renamed") deleted: Final = client.vector_stores.delete(created.id) - assert (retrieved.id, retrieved.name) == (provider_id, store_name) - assert (updated.id, updated.name) == (provider_id, "renamed") - assert (deleted.id, deleted.deleted) == (provider_id, True) + assert (retrieved.id, retrieved.name) == (created.id, store_name) + assert (updated.id, updated.name) == (created.id, "renamed") + assert (deleted.id, deleted.deleted) == (created.id, True) assert [(request.method, urlsplit(request.target).path) for request in wire.drain()] == [ ("GET", store_path), ("POST", store_path), ("DELETE", store_path), ] + assert read_rows(_MANAGED_STORE_ROW, (created.id,)) == [] + retrieved_after_delete: Final = gateway.request("GET", f"/v1/vector_stores/{created.id}", key=key) + assert retrieved_after_delete.status_code == 404, retrieved_after_delete.text diff --git a/tests/unit/enterprise/proxy/hooks/test_managed_vector_stores.py b/tests/unit/enterprise/proxy/hooks/test_managed_vector_stores.py index cdc094973b6..101fd635db4 100644 --- a/tests/unit/enterprise/proxy/hooks/test_managed_vector_stores.py +++ b/tests/unit/enterprise/proxy/hooks/test_managed_vector_stores.py @@ -25,25 +25,34 @@ class _StoreRow: return {"created_by": self.created_by, "team_id": None} -class _SingleStoreTable: - async def find_first(self, where: Mapping[str, object]) -> _StoreRow: - return _StoreRow(created_by=_OWNER) +_OWNER_ROW: Final = _StoreRow(created_by=_OWNER) +@dataclass(frozen=True, slots=True) +class _StoreTable: + row: _StoreRow | None + + async def find_first(self, where: Mapping[str, object]) -> _StoreRow | None: + return self.row + + +@dataclass(frozen=True, slots=True) class _Database: - def __init__(self) -> None: - self.litellm_managedvectorstoretable: Final = _SingleStoreTable() + litellm_managedvectorstoretable: _StoreTable +@dataclass(frozen=True, slots=True) class _PrismaClient: - def __init__(self) -> None: - self.db: Final = _Database() + db: _Database def _hook_and_unified_id( target_model_names: tuple[str, ...] = ("model-a",), + stored_row: _StoreRow | None = _OWNER_ROW, ) -> tuple[PROXY_LiteLLMManagedVectorStores, str]: - hook: Final = PROXY_LiteLLMManagedVectorStores(DualCache(), prisma_client=_PrismaClient()) + hook: Final = PROXY_LiteLLMManagedVectorStores( + DualCache(), prisma_client=_PrismaClient(_Database(_StoreTable(stored_row))) + ) created: Final = VectorStoreCreateResponse( id=_PROVIDER_ID, object="vector_store", @@ -140,3 +149,19 @@ async def test_provider_store_id_passes_through_unchanged(call_type: str) -> Non ) assert routed == {"vector_store_id": "vs_raw_store"} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("call_type", _ROUTED_CALL_TYPES) +async def test_deleted_managed_store_answers_not_found_instead_of_access_denied(call_type: str) -> None: + hook, unified_id = _hook_and_unified_id(stored_row=None) + + with pytest.raises(HTTPException) as missing: + await hook.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(user_id=_OWNER), + cache=DualCache(), + data={"vector_store_id": unified_id}, + call_type=call_type, + ) + + assert missing.value.status_code == 404 diff --git a/tests/unit/proxy/openai_files_endpoint/test_files_endpoint.py b/tests/unit/proxy/openai_files_endpoint/test_files_endpoint.py index d47db69257b..c2a50470401 100644 --- a/tests/unit/proxy/openai_files_endpoint/test_files_endpoint.py +++ b/tests/unit/proxy/openai_files_endpoint/test_files_endpoint.py @@ -6263,3 +6263,14 @@ def test_create_file_passthrough_fails_closed_when_guardrails_would_scan_the_bat assert error["param"] == "passthrough" assert "guardrails" in error["message"] assert forwarded_calls == [] + + +def test_upload_with_target_model_names_on_a_proxy_without_a_database_says_what_it_needs(batch_upload_seams): + stored, provider_upload, _ = batch_upload_seams + + response = _upload_batch_file({}, {"target_model_names": "gemini-2.0-flash"}) + + assert response.status_code == 500, response.text + assert "DATABASE_URL" in response.json()["error"]["message"] + stored.assert_not_awaited() + provider_upload.assert_not_awaited() diff --git a/tests/unit/proxy/vector_store_endpoints/test_vector_store_endpoints.py b/tests/unit/proxy/vector_store_endpoints/test_vector_store_endpoints.py index 48642a9b50d..3a0b453a2d3 100644 --- a/tests/unit/proxy/vector_store_endpoints/test_vector_store_endpoints.py +++ b/tests/unit/proxy/vector_store_endpoints/test_vector_store_endpoints.py @@ -3413,3 +3413,93 @@ def test_proxy_vector_store_writes_send_only_the_callers_metadata_when_the_raw_b assert response.status_code == 200, response.text assert len(sent_bodies) == 1 assert sent_bodies[0].get("metadata") == CALLERS_METADATA + + +class _ManagedStoreRow: + def __init__(self, resource_object: dict) -> None: + self.resource_object = resource_object + + def model_dump(self) -> dict: + return {"created_by": "store-owner", "team_id": None} + + +class _ManagedStoreTable: + def __init__(self) -> None: + self.rows = {} + + async def find_first(self, where): + return self.rows.get(where["unified_resource_id"]) + + async def delete(self, where): + return self.rows.pop(where["unified_resource_id"]) + + +def test_proxy_answers_by_managed_vector_store_id_and_forgets_the_store_once_the_provider_deletes_it(monkeypatch): + from types import SimpleNamespace + + import httpx + import respx + from fastapi.testclient import TestClient + from litellm_enterprise.proxy.hooks.managed_vector_stores import PROXY_LiteLLMManagedVectorStores + + from litellm.caching import DualCache + from litellm.proxy import proxy_server + from litellm.proxy.proxy_server import app + + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + litellm.in_memory_llm_clients_cache.flush_cache() + monkeypatch.setattr(proxy_server, "master_key", "sk-master-test") + monkeypatch.setattr( + proxy_server, + "llm_router", + litellm.Router( + model_list=[{"model_name": "model-a", "litellm_params": {"model": "openai/gpt-5", "api_key": "sk-test"}}] + ), + ) + table = _ManagedStoreTable() + hook = PROXY_LiteLLMManagedVectorStores( + DualCache(), prisma_client=SimpleNamespace(db=SimpleNamespace(litellm_managedvectorstoretable=table)) + ) + provider_store = { + "id": "vs_provider", + "object": "vector_store", + "created_at": 0, + "name": "store", + "usage_bytes": 0, + "file_counts": {"in_progress": 0, "completed": 0, "failed": 0, "cancelled": 0, "total": 0}, + "status": "completed", + "metadata": {}, + } + unified_id = hook.generate_unified_resource_id(resource_objects=[provider_store], target_model_names_list=["model-a"]) + table.rows[unified_id] = _ManagedStoreRow(provider_store) + monkeypatch.setitem(proxy_server.proxy_logging_obj.proxy_hook_mapping, "managed_vector_stores", hook) + monkeypatch.setattr(litellm, "callbacks", [hook]) + headers = {"Authorization": "Bearer sk-master-test"} + client = TestClient(app) + + with respx.mock(assert_all_called=False) as upstream: + provider_route = "https://api.openai.com/v1/vector_stores/vs_provider" + upstream.get(provider_route).mock(return_value=httpx.Response(200, json=provider_store)) + upstream.post(provider_route).mock(return_value=httpx.Response(200, json={**provider_store, "name": "renamed"})) + upstream.delete(provider_route).mock( + return_value=httpx.Response(200, json={"id": "vs_provider", "object": "vector_store.deleted", "deleted": True}) + ) + + retrieved = client.get(f"/v1/vector_stores/{unified_id}", headers=headers) + updated = client.post(f"/v1/vector_stores/{unified_id}", json={"name": "renamed"}, headers=headers) + deleted = client.delete(f"/v1/vector_stores/{unified_id}", headers=headers) + retrieved_after_delete = client.get(f"/v1/vector_stores/{unified_id}", headers=headers) + upstream_requests = [(call.request.method, call.request.url.path) for call in upstream.calls] + + assert retrieved.status_code == 200, retrieved.text + assert (retrieved.json()["id"], retrieved.json()["name"]) == (unified_id, "store") + assert updated.status_code == 200, updated.text + assert (updated.json()["id"], updated.json()["name"]) == (unified_id, "renamed") + assert deleted.status_code == 200, deleted.text + assert (deleted.json()["id"], deleted.json()["deleted"]) == (unified_id, True) + assert retrieved_after_delete.status_code == 404, retrieved_after_delete.text + assert upstream_requests == [ + ("GET", "/v1/vector_stores/vs_provider"), + ("POST", "/v1/vector_stores/vs_provider"), + ("DELETE", "/v1/vector_stores/vs_provider"), + ] From 8cc51b31d7c0711eb326f66108dcd9a71ede75f2 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 29 Oct 2026 07:22:49 -0700 Subject: [PATCH 09/11] fix(managed files): answer 400 for a chat file part with no file object instead of 500 --- .../proxy/hooks/managed_files.py | 7 ++++--- .../proxy/hooks/test_managed_files.py | 20 +++++++++++++++++++ 2 files changed, 24 insertions(+), 3 deletions(-) diff --git a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py index 416444664a3..01aa8465567 100644 --- a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py +++ b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py @@ -1742,10 +1742,11 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): continue for c in content: if c.get("type") == "file": - file_object = cast(ChatCompletionFileObject, c) - file_object_file_field = file_object["file"] + file_object_file_field = cast(Dict[str, object], c).get("file") + if not isinstance(file_object_file_field, dict): + continue file_id = file_object_file_field.get("file_id") - if file_id: + if isinstance(file_id, str) and file_id: file_ids.append(file_id) return file_ids diff --git a/tests/unit/enterprise/proxy/hooks/test_managed_files.py b/tests/unit/enterprise/proxy/hooks/test_managed_files.py index 7f6e955e831..33d7289c893 100644 --- a/tests/unit/enterprise/proxy/hooks/test_managed_files.py +++ b/tests/unit/enterprise/proxy/hooks/test_managed_files.py @@ -332,6 +332,26 @@ def test_get_file_ids_from_messages_skips_bedrock_content_blocks_without_type(): assert file_ids == ["file-keep"] +def test_get_file_ids_from_messages_leaves_malformed_file_parts_to_request_validation(): + proxy_managed_files = PROXY_LiteLLMManagedFiles( + DualCache(), prisma_client=MagicMock() + ) + messages = [ + { + "role": "user", + "content": [ + {"type": "file"}, + {"type": "file", "file": "file-raw"}, + {"type": "file", "file": None}, + {"type": "file", "file": {"file_id": 7}}, + {"type": "file", "file": {"file_id": "file-keep"}}, + ], + } + ] + file_ids = proxy_managed_files.get_file_ids_from_messages(messages) + assert file_ids == ["file-keep"] + + @pytest.mark.asyncio async def test_async_pre_call_hook_batch_retrieve(): from litellm.proxy._types import UserAPIKeyAuth From 15e6516fe110d6b6fba58070f497a3b06218bfbd Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 29 Oct 2026 09:30:58 -0700 Subject: [PATCH 10/11] fix(managed files): skip chat messages and content items that are not objects --- enterprise/litellm_enterprise/proxy/hooks/managed_files.py | 4 ++-- tests/unit/enterprise/proxy/hooks/test_managed_files.py | 2 ++ 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py index 01aa8465567..171be00ef6b 100644 --- a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py +++ b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py @@ -1735,13 +1735,13 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): """ file_ids = [] for message in messages: - if message.get("role") == "user": + if isinstance(message, dict) and message.get("role") == "user": content = message.get("content") if content: if isinstance(content, str): continue for c in content: - if c.get("type") == "file": + if isinstance(c, dict) and c.get("type") == "file": file_object_file_field = cast(Dict[str, object], c).get("file") if not isinstance(file_object_file_field, dict): continue diff --git a/tests/unit/enterprise/proxy/hooks/test_managed_files.py b/tests/unit/enterprise/proxy/hooks/test_managed_files.py index 33d7289c893..7d9b328fc9a 100644 --- a/tests/unit/enterprise/proxy/hooks/test_managed_files.py +++ b/tests/unit/enterprise/proxy/hooks/test_managed_files.py @@ -337,9 +337,11 @@ def test_get_file_ids_from_messages_leaves_malformed_file_parts_to_request_valid DualCache(), prisma_client=MagicMock() ) messages = [ + "not-a-message", { "role": "user", "content": [ + "not-a-part", {"type": "file"}, {"type": "file", "file": "file-raw"}, {"type": "file", "file": None}, From 0f160e9042a31547a69a3f78ffa1fbe847ac0c37 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 29 Oct 2026 10:20:52 -0700 Subject: [PATCH 11/11] fix(vector stores): confirm a provider delete when a concurrent delete already removed the managed row The base managed resource delete returns None for a missing row, as its docstring says, so the route answers the provider's confirmation instead of 500. The managed files chat guard checks the message as an object, which keeps the runtime check without a redundant isinstance for the type gate --- .../proxy/hooks/managed_files.py | 2 +- .../base_managed_resource.py | 2 +- .../test_vector_store_endpoints.py | 53 +++++++++++++------ 3 files changed, 40 insertions(+), 17 deletions(-) diff --git a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py index 171be00ef6b..6b4c0519b27 100644 --- a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py +++ b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py @@ -1735,7 +1735,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): """ file_ids = [] for message in messages: - if isinstance(message, dict) and message.get("role") == "user": + if isinstance(cast(object, message), dict) and message.get("role") == "user": content = message.get("content") if content: if isinstance(content, str): diff --git a/litellm/llms/base_llm/managed_resources/base_managed_resource.py b/litellm/llms/base_llm/managed_resources/base_managed_resource.py index 2663a3c42e7..b36770876d5 100644 --- a/litellm/llms/base_llm/managed_resources/base_managed_resource.py +++ b/litellm/llms/base_llm/managed_resources/base_managed_resource.py @@ -302,7 +302,7 @@ class BaseManagedResource(ABC, Generic[ResourceObjectType]): initial_value: Final = await table.find_first(where={"unified_resource_id": unified_resource_id}) if initial_value is None: - raise Exception(f"LiteLLM Managed {self.resource_type} with id={unified_resource_id} not found") + return None # Delete from cache await self.internal_usage_cache.async_set_cache( diff --git a/tests/unit/proxy/vector_store_endpoints/test_vector_store_endpoints.py b/tests/unit/proxy/vector_store_endpoints/test_vector_store_endpoints.py index 3a0b453a2d3..ba204eb2913 100644 --- a/tests/unit/proxy/vector_store_endpoints/test_vector_store_endpoints.py +++ b/tests/unit/proxy/vector_store_endpoints/test_vector_store_endpoints.py @@ -3434,11 +3434,13 @@ class _ManagedStoreTable: return self.rows.pop(where["unified_resource_id"]) -def test_proxy_answers_by_managed_vector_store_id_and_forgets_the_store_once_the_provider_deletes_it(monkeypatch): +_PROVIDER_STORE_ROUTE = "https://api.openai.com/v1/vector_stores/vs_provider" +_PROVIDER_DELETED_STORE = {"id": "vs_provider", "object": "vector_store.deleted", "deleted": True} + + +def _proxy_serving_one_managed_vector_store(monkeypatch): from types import SimpleNamespace - import httpx - import respx from fastapi.testclient import TestClient from litellm_enterprise.proxy.hooks.managed_vector_stores import PROXY_LiteLLMManagedVectorStores @@ -3474,21 +3476,24 @@ def test_proxy_answers_by_managed_vector_store_id_and_forgets_the_store_once_the table.rows[unified_id] = _ManagedStoreRow(provider_store) monkeypatch.setitem(proxy_server.proxy_logging_obj.proxy_hook_mapping, "managed_vector_stores", hook) monkeypatch.setattr(litellm, "callbacks", [hook]) - headers = {"Authorization": "Bearer sk-master-test"} - client = TestClient(app) + return TestClient(app, headers={"Authorization": "Bearer sk-master-test"}), table, unified_id, provider_store + + +def test_proxy_answers_by_managed_vector_store_id_and_forgets_the_store_once_the_provider_deletes_it(monkeypatch): + import httpx + import respx + + client, _, unified_id, provider_store = _proxy_serving_one_managed_vector_store(monkeypatch) with respx.mock(assert_all_called=False) as upstream: - provider_route = "https://api.openai.com/v1/vector_stores/vs_provider" - upstream.get(provider_route).mock(return_value=httpx.Response(200, json=provider_store)) - upstream.post(provider_route).mock(return_value=httpx.Response(200, json={**provider_store, "name": "renamed"})) - upstream.delete(provider_route).mock( - return_value=httpx.Response(200, json={"id": "vs_provider", "object": "vector_store.deleted", "deleted": True}) - ) + upstream.get(_PROVIDER_STORE_ROUTE).mock(return_value=httpx.Response(200, json=provider_store)) + upstream.post(_PROVIDER_STORE_ROUTE).mock(return_value=httpx.Response(200, json={**provider_store, "name": "renamed"})) + upstream.delete(_PROVIDER_STORE_ROUTE).mock(return_value=httpx.Response(200, json=_PROVIDER_DELETED_STORE)) - retrieved = client.get(f"/v1/vector_stores/{unified_id}", headers=headers) - updated = client.post(f"/v1/vector_stores/{unified_id}", json={"name": "renamed"}, headers=headers) - deleted = client.delete(f"/v1/vector_stores/{unified_id}", headers=headers) - retrieved_after_delete = client.get(f"/v1/vector_stores/{unified_id}", headers=headers) + retrieved = client.get(f"/v1/vector_stores/{unified_id}") + updated = client.post(f"/v1/vector_stores/{unified_id}", json={"name": "renamed"}) + deleted = client.delete(f"/v1/vector_stores/{unified_id}") + retrieved_after_delete = client.get(f"/v1/vector_stores/{unified_id}") upstream_requests = [(call.request.method, call.request.url.path) for call in upstream.calls] assert retrieved.status_code == 200, retrieved.text @@ -3503,3 +3508,21 @@ def test_proxy_answers_by_managed_vector_store_id_and_forgets_the_store_once_the ("POST", "/v1/vector_stores/vs_provider"), ("DELETE", "/v1/vector_stores/vs_provider"), ] + + +def test_proxy_confirms_a_managed_vector_store_delete_when_a_concurrent_delete_already_forgot_the_store(monkeypatch): + import httpx + import respx + + client, table, unified_id, _ = _proxy_serving_one_managed_vector_store(monkeypatch) + + def provider_deletes_after_a_concurrent_delete_forgot_the_store(request): + table.rows.pop(unified_id) + return httpx.Response(200, json=_PROVIDER_DELETED_STORE) + + with respx.mock(assert_all_called=False) as upstream: + upstream.delete(_PROVIDER_STORE_ROUTE).mock(side_effect=provider_deletes_after_a_concurrent_delete_forgot_the_store) + deleted = client.delete(f"/v1/vector_stores/{unified_id}") + + assert deleted.status_code == 200, deleted.text + assert (deleted.json()["id"], deleted.json()["deleted"]) == (unified_id, True)