diff --git a/Dockerfile b/Dockerfile index 4cef63c8d71..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; 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 diff --git a/docker/Dockerfile.database b/docker/Dockerfile.database index 20f810686a3..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; 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 diff --git a/docker/Dockerfile.non_root b/docker/Dockerfile.non_root index 90bcc523d9e..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; 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 diff --git a/enterprise/enterprise_hooks/__init__.py b/enterprise/enterprise_hooks/__init__.py index 6de3325f113..e69de29bb2d 100644 --- a/enterprise/enterprise_hooks/__init__.py +++ b/enterprise/enterprise_hooks/__init__.py @@ -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] diff --git a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py index 5b54e7d72a0..7878895447d 100644 --- a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py +++ b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py @@ -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 diff --git a/enterprise/litellm_enterprise/proxy/hooks/managed_vector_stores.py b/enterprise/litellm_enterprise/proxy/hooks/managed_vector_stores.py index fcaf58099b9..1f68e0b188a 100644 --- a/enterprise/litellm_enterprise/proxy/hooks/managed_vector_stores.py +++ b/enterprise/litellm_enterprise/proxy/hooks/managed_vector_stores.py @@ -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) diff --git a/litellm/llms/base_llm/managed_resources/base_managed_resource.py b/litellm/llms/base_llm/managed_resources/base_managed_resource.py index 2663a3c42e7..b36770876d5 100644 --- a/litellm/llms/base_llm/managed_resources/base_managed_resource.py +++ b/litellm/llms/base_llm/managed_resources/base_managed_resource.py @@ -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( diff --git a/litellm/proxy/hooks/__init__.py b/litellm/proxy/hooks/__init__.py index 138050ccc9a..089dca7ec11 100644 --- a/litellm/proxy/hooks/__init__.py +++ b/litellm/proxy/hooks/__init__.py @@ -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 diff --git a/litellm/proxy/hooks/autorouter_baseline_cache.py b/litellm/proxy/hooks/autorouter_baseline_cache.py index b7dbebbe8d8..55b2949dd45 100644 --- a/litellm/proxy/hooks/autorouter_baseline_cache.py +++ b/litellm/proxy/hooks/autorouter_baseline_cache.py @@ -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 = ( diff --git a/litellm/proxy/hooks/batch_rate_limiter.py b/litellm/proxy/hooks/batch_rate_limiter.py index 747dda100bf..74663e5b6c4 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 c1d4534ffa2..866c4bcd790 100644 --- a/litellm/proxy/openai_files_endpoints/files_endpoints.py +++ b/litellm/proxy/openai_files_endpoints/files_endpoints.py @@ -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, diff --git a/litellm/proxy/vector_store_endpoints/endpoints.py b/litellm/proxy/vector_store_endpoints/endpoints.py index b6a0ebe3c58..110c2d1034a 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 @@ -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, 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 0b0e51c0871..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 @@ -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 diff --git a/tests/unit/enterprise/proxy/hooks/test_managed_files.py b/tests/unit/enterprise/proxy/hooks/test_managed_files.py index 1755ab97689..f79ff4d2b1f 100644 --- a/tests/unit/enterprise/proxy/hooks/test_managed_files.py +++ b/tests/unit/enterprise/proxy/hooks/test_managed_files.py @@ -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 diff --git a/tests/unit/enterprise/proxy/hooks/test_managed_vector_stores.py b/tests/unit/enterprise/proxy/hooks/test_managed_vector_stores.py new file mode 100644 index 00000000000..101fd635db4 --- /dev/null +++ b/tests/unit/enterprise/proxy/hooks/test_managed_vector_stores.py @@ -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 diff --git a/tests/unit/proxy/hooks/test_proxy_hooks_init.py b/tests/unit/proxy/hooks/test_proxy_hooks_init.py index a6edd3db944..ce4559e4c2a 100644 --- a/tests/unit/proxy/hooks/test_proxy_hooks_init.py +++ b/tests/unit/proxy/hooks/test_proxy_hooks_init.py @@ -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 + ) 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 ff63e37bcd8..4dce37de6e9 100644 --- a/tests/unit/proxy/openai_files_endpoint/test_files_endpoint.py +++ b/tests/unit/proxy/openai_files_endpoint/test_files_endpoint.py @@ -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"}), 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 926566605a5..bdfc3f7b8f0 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 @@ -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)