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] 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 + )