Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 6 additions & 0 deletions litellm/proxy/vector_store_endpoints/endpoints.py
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@
_resolve_embedding_config,
)
from litellm.proxy.vector_store_endpoints.utils import (
assert_proxy_admin_for_vector_store_index_management,
assert_user_can_access_vector_store,
get_litellm_managed_vector_store,
)
Expand Down Expand Up @@ -575,6 +576,11 @@ async def index_create(
"""
from litellm.proxy.proxy_server import prisma_client

assert_proxy_admin_for_vector_store_index_management(
user_api_key_dict,
operation="create",
)

if prisma_client is None:
raise HTTPException(
status_code=500,
Expand Down
84 changes: 83 additions & 1 deletion litellm/proxy/vector_store_endpoints/utils.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
import json
import re
from typing import Any, Dict, Literal, Optional

from fastapi import HTTPException, Request
Expand Down Expand Up @@ -37,6 +38,64 @@ def _is_proxy_admin(user_api_key_dict: UserAPIKeyAuth) -> bool:
)


def assert_proxy_admin_for_vector_store_index_management(
user_api_key_dict: UserAPIKeyAuth,
*,
operation: Literal["create", "delete", "update"] = "create",
) -> None:
"""Raise 403 unless the caller is a proxy admin."""
if _is_proxy_admin(user_api_key_dict):
return
raise HTTPException(
status_code=403,
detail=(
f"Only proxy admins can {operation} vector store indexes. "
"Contact your LiteLLM administrator."
),
)


def _suffix_after_index_name(request_path: str, index_name: str) -> Optional[str]:
"""Return the path suffix after ``/indexes/{index_name}``, or None if absent."""
match = re.search(rf"/indexes/{re.escape(index_name)}(?=$|[/?])", request_path)
if match is None:
return None
return request_path[match.end() :]


def _is_vector_store_index_lifecycle_request(
request_method: str,
request_path: str,
index_name: str,
) -> bool:
"""
True when the request creates or deletes a search index itself (not documents).

Examples (admin-only):
- DELETE /azure_ai/indexes/my-index
- PUT /azure_ai/indexes/my-index
- POST /azure_ai/indexes
"""
if request_method not in ("POST", "PUT", "DELETE", "PATCH"):
return False

suffix = _suffix_after_index_name(request_path, index_name)
if suffix is not None:
# Document operations live under /indexes/{name}/docs/...
if suffix.startswith("/docs"):
return False
# DELETE/PUT/PATCH on /indexes/{name} itself is index lifecycle.
if suffix == "" or suffix.startswith("?"):
return True

# POST /indexes (create index at service level; no index name in path).
normalized = request_path.rstrip("/")
if request_method == "POST" and normalized.endswith("/indexes"):
return True

return False


def _object_permission_allows_vector_store(
object_permission: Optional[LiteLLM_ObjectPermissionTable],
vector_store_id: str,
Expand Down Expand Up @@ -335,6 +394,22 @@ def is_allowed_to_call_vector_store_endpoint(

request_route = get_request_route(request)

if _is_vector_store_index_lifecycle_request(
request_method=request.method,
request_path=request_route,
index_name=index_name,
):
operation_label: Literal["create", "delete", "update"] = "create"
if request.method == "DELETE":
operation_label = "delete"
elif request.method in ("PUT", "PATCH"):
operation_label = "update"
assert_proxy_admin_for_vector_store_index_management(
user_api_key_dict,
operation=operation_label,
)
Comment thread
greptile-apps[bot] marked this conversation as resolved.
return True

# Determine the permission type based on the request
permission_type = None
for endpoint in provider_vector_store_endpoints["read"]:
Expand All @@ -353,7 +428,14 @@ def is_allowed_to_call_vector_store_endpoint(
break

if permission_type is None:
return None
raise HTTPException(
status_code=403,
detail=(
f"User does not have permission to call vector store endpoint "
f"{index_name}. Ask your administrator to add the necessary "
"permissions to your API key/Team."
),
)
Comment thread
greptile-apps[bot] marked this conversation as resolved.

# Check if key has specific permission for allowed_vector_store_indexes
has_permission = check_vector_store_permission(
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -16,9 +16,10 @@
from litellm.integrations.vector_store_integrations.vector_store_pre_call_hook import (
LiteLLM_ManagedVectorStore,
)
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
from litellm.proxy.vector_store_endpoints.endpoints import (
_update_request_data_with_litellm_managed_vector_store_registry,
index_create,
)
from litellm.proxy.vector_store_endpoints.management_endpoints import (
_check_vector_store_access,
Expand All @@ -33,6 +34,7 @@
is_allowed_to_call_vector_store_endpoint,
is_allowed_to_call_vector_store_files_endpoint,
)
from litellm.types.vector_stores import IndexCreateRequest
from litellm.types.utils import LlmProviders


Expand Down Expand Up @@ -674,18 +676,151 @@ def test_endpoint_not_recognized(self):
"write": [("POST", "/create")],
}

with patch(
"litellm.proxy.vector_store_endpoints.utils.ProviderConfigManager.get_provider_vector_stores_config",
return_value=mock_provider_config,
):
with pytest.raises(HTTPException) as exc_info:
is_allowed_to_call_vector_store_endpoint(
provider=LlmProviders.OPENAI,
index_name="my-index",
request=mock_request,
user_api_key_dict=mock_user_api_key,
)

assert exc_info.value.status_code == 403

def test_delete_index_requires_admin(self):
"""Non-admin users must not delete managed search indexes via pass-through."""
mock_request = MagicMock(spec=Request)
mock_request.method = "DELETE"
mock_request.url.path = "/azure_ai/indexes/my-index"

mock_user_api_key = MagicMock(spec=UserAPIKeyAuth)
mock_user_api_key.user_role = None
mock_user_api_key.metadata = {
"allowed_vector_store_indexes": [
{"index_name": "my-index", "index_permissions": ["read", "write"]}
]
}
mock_user_api_key.team_metadata = None

mock_provider_config = MagicMock()
mock_provider_config.get_vector_store_endpoints_by_type.return_value = {
"read": [("GET", "/docs/search"), ("POST", "/docs/search")],
"write": [("PUT", "/docs")],
}

with patch(
"litellm.proxy.vector_store_endpoints.utils.ProviderConfigManager.get_provider_vector_stores_config",
return_value=mock_provider_config,
):
with pytest.raises(HTTPException) as exc_info:
is_allowed_to_call_vector_store_endpoint(
provider=LlmProviders.AZURE_AI,
index_name="my-index",
request=mock_request,
user_api_key_dict=mock_user_api_key,
)

assert exc_info.value.status_code == 403
assert "Only proxy admins can delete" in exc_info.value.detail

def test_delete_index_allowed_for_admin(self):
"""Proxy admins can delete managed search indexes via pass-through."""
mock_request = MagicMock(spec=Request)
mock_request.method = "DELETE"
mock_request.url.path = "/azure_ai/indexes/my-index"

mock_user_api_key = MagicMock(spec=UserAPIKeyAuth)
mock_user_api_key.user_role = LitellmUserRoles.PROXY_ADMIN

mock_provider_config = MagicMock()
mock_provider_config.get_vector_store_endpoints_by_type.return_value = {
"read": [("GET", "/docs/search"), ("POST", "/docs/search")],
"write": [("PUT", "/docs")],
}

with patch(
"litellm.proxy.vector_store_endpoints.utils.ProviderConfigManager.get_provider_vector_stores_config",
return_value=mock_provider_config,
):
result = is_allowed_to_call_vector_store_endpoint(
provider=LlmProviders.OPENAI,
provider=LlmProviders.AZURE_AI,
index_name="my-index",
request=mock_request,
user_api_key_dict=mock_user_api_key,
)

assert result is None
assert result is True

def test_update_index_requires_admin_with_update_message(self):
"""Non-admin users get an update-specific message for index replacement."""
mock_request = MagicMock(spec=Request)
mock_request.method = "PUT"
mock_request.url.path = "/azure_ai/indexes/my-index"

mock_user_api_key = MagicMock(spec=UserAPIKeyAuth)
mock_user_api_key.user_role = None
mock_user_api_key.metadata = {
"allowed_vector_store_indexes": [
{"index_name": "my-index", "index_permissions": ["read", "write"]}
]
}
mock_user_api_key.team_metadata = None

mock_provider_config = MagicMock()
mock_provider_config.get_vector_store_endpoints_by_type.return_value = {
"read": [("GET", "/docs/search"), ("POST", "/docs/search")],
"write": [("PUT", "/docs")],
}

with patch(
"litellm.proxy.vector_store_endpoints.utils.ProviderConfigManager.get_provider_vector_stores_config",
return_value=mock_provider_config,
):
with pytest.raises(HTTPException) as exc_info:
is_allowed_to_call_vector_store_endpoint(
provider=LlmProviders.AZURE_AI,
index_name="my-index",
request=mock_request,
user_api_key_dict=mock_user_api_key,
)

assert exc_info.value.status_code == 403
assert "Only proxy admins can update" in exc_info.value.detail

def test_index_name_prefix_does_not_match_lifecycle_request(self):
"""An index name that is only a path prefix must not trigger lifecycle checks."""
mock_request = MagicMock(spec=Request)
mock_request.method = "DELETE"
mock_request.url.path = "/azure_ai/indexes/my-index-archive"

mock_user_api_key = MagicMock(spec=UserAPIKeyAuth)
mock_user_api_key.user_role = None
mock_user_api_key.metadata = None
mock_user_api_key.team_metadata = None

mock_provider_config = MagicMock()
mock_provider_config.get_vector_store_endpoints_by_type.return_value = {
"read": [],
"write": [],
}

with patch(
"litellm.proxy.vector_store_endpoints.utils.ProviderConfigManager.get_provider_vector_stores_config",
return_value=mock_provider_config,
):
with pytest.raises(HTTPException) as exc_info:
is_allowed_to_call_vector_store_endpoint(
provider=LlmProviders.AZURE_AI,
index_name="my-index",
request=mock_request,
user_api_key_dict=mock_user_api_key,
)

assert exc_info.value.status_code == 403
assert "Only proxy admins" not in exc_info.value.detail

def test_team_metadata_permissions(self):
"""Test that team metadata permissions work."""
Expand Down Expand Up @@ -800,6 +935,81 @@ def test_azure_ai_permission_allowed(self):
assert exc_info.value.status_code == 403


class TestIndexCreate:
@pytest.mark.asyncio
async def test_index_create_requires_admin(self):
"""Non-admin users must not register managed vector store indexes."""
request = IndexCreateRequest(
index_name="test-index",
litellm_params={
"vector_store_index": "real-index",
"vector_store_name": "azure-ai-search",
},
)
mock_request = MagicMock(spec=Request)
mock_response = MagicMock()

with pytest.raises(HTTPException) as exc_info:
await index_create(
request=mock_request,
index_create_request=request,
fastapi_response=mock_response,
user_api_key_dict=UserAPIKeyAuth(
token="sk-test",
key_name="sk-...test",
user_role=LitellmUserRoles.INTERNAL_USER,
),
)

assert exc_info.value.status_code == 403
assert "Only proxy admins can create" in exc_info.value.detail

@pytest.mark.asyncio
async def test_index_create_allowed_for_admin(self):
"""Proxy admins can register managed vector store indexes."""
create_request = IndexCreateRequest(
index_name="test-index",
litellm_params={
"vector_store_index": "real-index",
"vector_store_name": "azure-ai-search",
},
)
mock_request = MagicMock(spec=Request)
mock_response = MagicMock()
mock_row = MagicMock()
mock_row.model_dump.return_value = {
"index_name": "test-index",
"litellm_params": create_request.litellm_params.model_dump(),
}

mock_prisma = MagicMock()
mock_prisma.db.litellm_managedvectorstoreindextable.find_unique = AsyncMock(
return_value=None
)
mock_prisma.db.litellm_managedvectorstoreindextable.create = AsyncMock(
return_value=mock_row
)

with patch(
"litellm.proxy.proxy_server.prisma_client",
mock_prisma,
):
result = await index_create(
request=mock_request,
index_create_request=create_request,
fastapi_response=mock_response,
user_api_key_dict=UserAPIKeyAuth(
token="sk-test",
key_name="sk-...test",
user_role=LitellmUserRoles.PROXY_ADMIN,
user_id="admin-user",
),
)

assert result["index_name"] == "test-index"
mock_prisma.db.litellm_managedvectorstoreindextable.create.assert_awaited_once()


class TestIsAllowedToCallVectorStoreFilesEndpoint:
def _mock_provider_config(self):
provider_config = MagicMock()
Expand Down
Loading