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