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