scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-10 08:28:25 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 108

fix: keep the managed id in vector store answers, forget a deleted store, and name the database in hook-missing errors

Browse source
This commit is contained in:
mateo-berri 2026-10-28 22:07:33 -07:00
parent d031d90b90
commit a8fe887b64
9 changed files with 209 additions and 32 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

28
enterprise/litellm_enterprise/proxy/hooks/managed_vector_stores.py
View file

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

3
litellm/proxy/hooks/batch_rate_limiter.py
View file

@ -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.")

8
litellm/proxy/openai_files_endpoints/files_endpoints.py
View file

@ -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,

38
litellm/proxy/vector_store_endpoints/endpoints.py
View file

@ -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,

8
litellm/types/proxy/common_proxy_errors.py
View file

@ -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"
)

14
tests/integration/management/test_managed_vector_store_unified_id.py
View file

@ -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

41
tests/unit/enterprise/proxy/hooks/test_managed_vector_stores.py
View file

@ -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

11
tests/unit/proxy/openai_files_endpoint/test_files_endpoint.py
View file

@ -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()

90
tests/unit/proxy/vector_store_endpoints/test_vector_store_endpoints.py
View file

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