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

Merge 0f160e9042 into f32ac94ad1

Browse source
This commit is contained in:
norkelbyte 2026-10-29 23:33:37 -07:00 • committed by GitHub
commit c05428cc98
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
19 changed files with 601 additions and 157 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

4
Dockerfile
View file

@ -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; litellm/proxy/hooks resolves
# 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

4
docker/Dockerfile.database
View file

@ -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; litellm/proxy/hooks resolves
# 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

4
docker/Dockerfile.non_root
View file

@ -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; litellm/proxy/hooks resolves
# 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

33
enterprise/enterprise_hooks/__init__.py
View file

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

11
enterprise/litellm_enterprise/proxy/hooks/managed_files.py
View file

@ -1735,17 +1735,18 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
"""
file_ids = []
for message in messages:
if 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):
continue
for c in content:
if c.get("type") == "file":
file_object = cast(ChatCompletionFileObject, c)
file_object_file_field = file_object["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
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

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

@ -11,7 +11,11 @@ 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,
extract_target_model_names_from_unified_id,
generate_unified_id_string,
is_base64_encoded_unified_id,
)
@ -38,6 +42,25 @@ else:
PrismaClient = Any
_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], 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
)
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]
):
@ -47,7 +70,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
"""
@ -281,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)
@ -322,78 +348,24 @@ 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)
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)

2
litellm/llms/base_llm/managed_resources/base_managed_resource.py
View file

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

34
litellm/proxy/hooks/__init__.py
View file

@ -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,19 @@ 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:
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,
)
else:
PROXY_HOOKS["managed_files"] = PROXY_LiteLLMManagedFiles
PROXY_HOOKS["managed_vector_stores"] = PROXY_LiteLLMManagedVectorStores

6
litellm/proxy/hooks/autorouter_baseline_cache.py
View file

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

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

@ -489,7 +489,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,
@ -1034,7 +1034,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,
@ -1376,7 +1376,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,
@ -1586,7 +1586,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
@ -266,7 +272,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:
@ -315,6 +321,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(
@ -351,7 +372,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,
@ -369,6 +390,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,
@ -490,7 +514,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,
@ -508,6 +532,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,
@ -553,7 +580,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,
@ -571,6 +598,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"
)

71
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
@ -265,3 +266,73 @@ 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}"
store_deleted: Final = threading.Event()
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" 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:
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) == (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

22
tests/unit/enterprise/proxy/hooks/test_managed_files.py
View file

@ -332,6 +332,28 @@ 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 = [
"not-a-message",
{
"role": "user",
"content": [
"not-a-part",
{"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

167
tests/unit/enterprise/proxy/hooks/test_managed_vector_stores.py Normal file
View file

@ -0,0 +1,167 @@
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"
_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)
class _StoreRow:
created_by: str
def model_dump(self) -> dict[str, str | None]:
return {"created_by": self.created_by, "team_id": None}
_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:
litellm_managedvectorstoretable: _StoreTable
@dataclass(frozen=True, slots=True)
class _PrismaClient:
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(_Database(_StoreTable(stored_row)))
)
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=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_its_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", _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:
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"}
@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

77
tests/unit/proxy/hooks/test_proxy_hooks_init.py
View file

@ -14,8 +14,15 @@ silently swallowed the ImportError in `hooks/__init__.py`, leaving
``managed_files`` unregistered and the /files endpoint returning 500.
"""
import importlib
import logging
import sys
from typing import Final
import pytest
import litellm.proxy
from litellm._logging import verbose_proxy_logger
from litellm.proxy.hooks import PROXY_HOOKS, get_proxy_hook
@ -49,3 +56,73 @@ 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: 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
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: Final = 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
)
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
)

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

@ -6306,6 +6306,17 @@ def test_create_file_passthrough_fails_closed_when_guardrails_would_scan_the_bat
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()
def _callers_files_on_every_deployment() -> list[MagicMock]:
return [
_managed_file_row("unified-gpt-newer", model_mappings={"gpt-3.5-turbo-id": "file-gpt-newer"}),

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

@ -3466,3 +3466,116 @@ 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"])
_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
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])
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:
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}")
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
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"),
]
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)