Coverage for .venv/lib/python3.13/site-packages/litellm/proxy/vector_store_endpoints/utils.py: 22%
225 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-10-10 12:01 +0000
1import json
2import re
3from collections.abc import Iterable, Mapping
4from types import MappingProxyType
5from typing import Final, Literal
7from fastapi import HTTPException, Request
9import litellm
10from litellm._logging import verbose_proxy_logger
11from litellm.proxy._experimental.mcp_server.ui_session_utils import (
12 is_ui_session_credential,
13 resolve_ui_session_team_ids,
14)
15from litellm.proxy._types import (
16 LiteLLM_ObjectPermissionTable,
17 LitellmUserRoles,
18 UserAPIKeyAuth,
19)
20from litellm.types.utils import LlmProviders
21from litellm.types.vector_stores import LiteLLM_ManagedVectorStore
22from litellm.utils import ProviderConfigManager
25def _normalize_litellm_params(
26 vector_store: LiteLLM_ManagedVectorStore,
27) -> LiteLLM_ManagedVectorStore:
28 litellm_params: Final = vector_store.get("litellm_params")
29 if isinstance(litellm_params, str): 29 ↛ 30line 29 didn't jump to line 30 because the condition on line 29 was never true
30 normalized: Final = LiteLLM_ManagedVectorStore(**dict(vector_store))
31 try:
32 parsed: Final = json.loads(litellm_params)
33 normalized["litellm_params"] = parsed if isinstance(parsed, dict) else {}
34 except (TypeError, ValueError):
35 normalized["litellm_params"] = {}
36 return normalized
37 return vector_store
40def _is_proxy_admin(user_api_key_dict: UserAPIKeyAuth) -> bool:
41 return (
42 user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN
43 or user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value
44 )
47def assert_proxy_admin_for_vector_store_index_management(
48 user_api_key_dict: UserAPIKeyAuth,
49 *,
50 operation: Literal["create", "delete", "update", "list"] = "create",
51) -> None:
52 """Raise 403 unless the caller is a proxy admin."""
53 if _is_proxy_admin(user_api_key_dict): 53 ↛ 55line 53 didn't jump to line 55 because the condition on line 53 was always true
54 return
55 raise HTTPException(
56 status_code=403,
57 detail=(f"Only proxy admins can {operation} vector store indexes. Contact your LiteLLM administrator."),
58 )
61def _suffix_after_index_name(request_path: str, index_name: str) -> str | None:
62 """Return the path suffix after ``/indexes/{index_name}``, or None if absent."""
63 match: Final = re.search(rf"/indexes/{re.escape(index_name)}(?=$|[/?])", request_path)
64 if match is None:
65 return None
66 return request_path[match.end() :]
69def _is_vector_store_index_lifecycle_request(
70 request_method: str,
71 request_path: str,
72 index_name: str,
73) -> bool:
74 """
75 True when the request creates or deletes a search index itself (not documents).
77 Examples (admin-only):
78 - DELETE /azure_ai/indexes/my-index
79 - PUT /azure_ai/indexes/my-index
80 - POST /azure_ai/indexes
81 """
82 if request_method not in ("POST", "PUT", "DELETE", "PATCH"):
83 return False
85 suffix: Final = _suffix_after_index_name(request_path, index_name)
86 if suffix is not None:
87 # Document operations live under /indexes/{name}/docs/...
88 if suffix.startswith("/docs"):
89 return False
90 # DELETE/PUT/PATCH on /indexes/{name} itself is index lifecycle.
91 if suffix == "" or suffix.startswith("?"):
92 return True
94 # POST /indexes (create index at service level; no index name in path).
95 normalized: Final = request_path.split("?", 1)[0].rstrip("/")
96 if request_method == "POST" and normalized.endswith("/indexes"):
97 return True
99 return False
102def _object_permission_allows_vector_store(
103 object_permission: LiteLLM_ObjectPermissionTable | None,
104 vector_store_id: str,
105) -> bool:
106 """Returns True if an object permission explicitly allowlists the vector store."""
107 if object_permission is None:
108 return False
109 allowed: Final = object_permission.vector_stores
110 if not allowed:
111 return False
112 return vector_store_id in allowed
115async def _get_object_permission_for_id(
116 object_permission_id: str | None,
117) -> LiteLLM_ObjectPermissionTable | None:
118 """Load an object permission record by id, using the shared cache/DB helper."""
119 if not object_permission_id:
120 return None
122 from litellm.proxy.auth.auth_checks import get_object_permission
123 from litellm.proxy.proxy_server import (
124 prisma_client,
125 proxy_logging_obj,
126 user_api_key_cache,
127 )
129 if prisma_client is None:
130 return None
132 try:
133 return await get_object_permission(
134 object_permission_id=object_permission_id,
135 prisma_client=prisma_client,
136 user_api_key_cache=user_api_key_cache,
137 proxy_logging_obj=proxy_logging_obj,
138 )
139 except Exception as e:
140 verbose_proxy_logger.debug(
141 "Failed to load object_permission id=%s: %s",
142 object_permission_id,
143 e,
144 )
145 return None
148async def can_user_access_vector_store(
149 vector_store: LiteLLM_ManagedVectorStore,
150 user_api_key_dict: UserAPIKeyAuth,
151) -> bool:
152 """
153 Returns True if the caller is allowed to access this managed vector store.
155 Access is granted (first match wins) when any of the following is true:
156 1. The caller's role is PROXY_ADMIN.
157 2. The vector store has no team_id (legacy behavior - accessible to all).
158 3. The caller's key-level object_permission.vector_stores explicitly lists
159 this vector store id.
160 4. The caller's team-level object_permission.vector_stores explicitly lists
161 this vector store id.
162 5. The caller's team_id matches the vector store's team_id.
164 A dashboard session credential is evaluated against the same effective
165 contexts as listing (its own grants plus each real team of the user).
166 Otherwise access is denied.
167 """
168 if _is_proxy_admin(user_api_key_dict): 168 ↛ 171line 168 didn't jump to line 171 because the condition on line 168 was always true
169 return True
171 if vector_store.get("team_id") is None:
172 return True
174 auth_contexts: Final = await _vector_store_auth_contexts(user_api_key_dict)
175 return await _is_vector_store_granted_to_any(vector_store, auth_contexts)
178async def _is_vector_store_granted(
179 vector_store: LiteLLM_ManagedVectorStore,
180 user_api_key_dict: UserAPIKeyAuth,
181) -> bool:
182 vector_store_id: Final = vector_store.get("vector_store_id") or ""
184 key_object_permission = user_api_key_dict.object_permission
185 if key_object_permission is None:
186 key_object_permission = await _get_object_permission_for_id(user_api_key_dict.object_permission_id)
187 if _object_permission_allows_vector_store(key_object_permission, vector_store_id):
188 return True
190 team_object_permission: LiteLLM_ObjectPermissionTable | None = user_api_key_dict.team_object_permission
191 if team_object_permission is None:
192 team_object_permission = await _get_object_permission_for_id(user_api_key_dict.team_object_permission_id)
193 if _object_permission_allows_vector_store(team_object_permission, vector_store_id):
194 return True
196 return user_api_key_dict.team_id is not None and user_api_key_dict.team_id == vector_store.get("team_id")
199async def _team_auth_context(team_id: str, user_api_key_dict: UserAPIKeyAuth) -> UserAPIKeyAuth:
200 from litellm.proxy.auth.auth_checks import get_team_object
201 from litellm.proxy.proxy_server import (
202 prisma_client,
203 proxy_logging_obj,
204 user_api_key_cache,
205 )
207 team: Final = await get_team_object(
208 team_id=team_id,
209 prisma_client=prisma_client,
210 user_api_key_cache=user_api_key_cache,
211 parent_otel_span=user_api_key_dict.parent_otel_span,
212 proxy_logging_obj=proxy_logging_obj,
213 )
214 return user_api_key_dict.model_copy(
215 update=MappingProxyType(
216 {
217 "team_id": team_id,
218 "team_object_permission": team.object_permission,
219 "team_object_permission_id": team.object_permission_id,
220 }
221 )
222 )
225async def _vector_store_auth_contexts(
226 user_api_key_dict: UserAPIKeyAuth,
227) -> tuple[UserAPIKeyAuth, ...]:
228 if not is_ui_session_credential(user_api_key_dict):
229 return (user_api_key_dict,)
230 session_key_context: Final = user_api_key_dict.model_copy(
231 update=MappingProxyType({"team_id": None, "team_object_permission": None, "team_object_permission_id": None})
232 )
233 team_ids: Final = await resolve_ui_session_team_ids(user_api_key_dict)
234 team_contexts: Final = tuple([await _team_auth_context(team_id, user_api_key_dict) for team_id in team_ids])
235 return (session_key_context, *team_contexts)
238async def _is_vector_store_granted_to_any(
239 vector_store: LiteLLM_ManagedVectorStore,
240 auth_contexts: tuple[UserAPIKeyAuth, ...],
241) -> bool:
242 for auth_context in auth_contexts:
243 if await _is_vector_store_granted(vector_store, auth_context):
244 return True
245 return False
248async def filter_listable_vector_stores(
249 vector_stores: Iterable[LiteLLM_ManagedVectorStore],
250 user_api_key_dict: UserAPIKeyAuth,
251) -> tuple[LiteLLM_ManagedVectorStore, ...]:
252 """Non-admins only see stores their key, one of their teams' object_permission, or team ownership grants."""
253 if _is_proxy_admin(user_api_key_dict): 253 ↛ 256line 253 didn't jump to line 256 because the condition on line 253 was always true
254 return tuple(vector_stores)
256 auth_contexts: Final = await _vector_store_auth_contexts(user_api_key_dict)
257 return tuple([vs for vs in vector_stores if await _is_vector_store_granted_to_any(vs, auth_contexts)])
260async def get_litellm_managed_vector_store(
261 vector_store_id: str,
262) -> LiteLLM_ManagedVectorStore | None:
263 """
264 Resolve a LiteLLM-managed vector store from the registry or shared cache.
266 Provider-native vector store IDs will not be present in either location and
267 return None, preserving direct provider behavior while still protecting
268 LiteLLM-managed multi-tenant stores.
269 """
270 if not vector_store_id: 270 ↛ 271line 270 didn't jump to line 271 because the condition on line 270 was never true
271 return None
273 if litellm.vector_store_registry is not None:
274 try:
275 vector_store: Final = litellm.vector_store_registry.get_litellm_managed_vector_store_from_registry(
276 vector_store_id=vector_store_id
277 )
278 if vector_store is not None:
279 return _normalize_litellm_params(vector_store)
280 except Exception as e:
281 verbose_proxy_logger.warning(
282 "Failed to resolve vector store id=%s from registry: %s",
283 vector_store_id,
284 e,
285 )
286 raise HTTPException(
287 status_code=500,
288 detail="Unable to validate vector store access",
289 ) from e
291 try:
292 from litellm.proxy.auth.auth_checks import (
293 get_managed_vector_store_rows_by_uuids,
294 )
295 from litellm.proxy.proxy_server import (
296 prisma_client,
297 proxy_logging_obj,
298 user_api_key_cache,
299 )
301 if prisma_client is None: 301 ↛ 302line 301 didn't jump to line 302 because the condition on line 301 was never true
302 return None
303 rows: Final = await get_managed_vector_store_rows_by_uuids(
304 uuids=[vector_store_id],
305 prisma_client=prisma_client,
306 user_api_key_cache=user_api_key_cache,
307 proxy_logging_obj=proxy_logging_obj,
308 )
309 if not rows:
310 return None
311 return _normalize_litellm_params(LiteLLM_ManagedVectorStore(**rows[0].model_dump()))
312 except Exception as e:
313 verbose_proxy_logger.warning(
314 "Failed to resolve vector store id=%s from shared cache: %s",
315 vector_store_id,
316 e,
317 )
318 raise HTTPException(
319 status_code=500,
320 detail="Unable to validate vector store access",
321 ) from e
324async def assert_user_can_access_vector_store(
325 vector_store: LiteLLM_ManagedVectorStore,
326 user_api_key_dict: UserAPIKeyAuth,
327 detail: str = "Access denied: You do not have permission to access this vector store",
328) -> None:
329 """Raise 403 unless the caller can access the resolved vector store."""
330 if not await can_user_access_vector_store(vector_store, user_api_key_dict): 330 ↛ 331line 330 didn't jump to line 331 because the condition on line 330 was never true
331 raise HTTPException(status_code=403, detail=detail)
334async def assert_user_can_access_vector_store_id(
335 vector_store_id: str,
336 user_api_key_dict: UserAPIKeyAuth,
337 detail: str = "Access denied: You do not have permission to access this vector store",
338) -> LiteLLM_ManagedVectorStore | None:
339 """
340 Resolve a managed vector store id and enforce ownership if it exists.
342 Unknown ids are treated as provider-native ids and are not rejected here.
343 """
344 vector_store: Final = await get_litellm_managed_vector_store(vector_store_id=vector_store_id)
345 if vector_store is not None:
346 await assert_user_can_access_vector_store(
347 vector_store=vector_store,
348 user_api_key_dict=user_api_key_dict,
349 detail=detail,
350 )
351 return vector_store
354def _does_endpoint_match(endpoint_path: str, request_path: str) -> bool:
355 if endpoint_path in request_path:
356 return True
357 if "{" in endpoint_path:
358 prefix: Final = endpoint_path.split("{", 1)[0]
359 if prefix and prefix in request_path:
360 return True
361 return False
364def check_vector_store_permission(
365 index_name: str,
366 permission: str,
367 key_metadata: Mapping[str, object] | None,
368 team_metadata: Mapping[str, object] | None,
369) -> bool:
370 """
371 Check if a specific permission is allowed for a given vector store index.
373 Args:
374 index_name: The name of the vector store index
375 permission: The permission to check (e.g., "read", "write")
376 key_metadata: Metadata from the API key
377 team_metadata: Metadata from the team
379 Returns:
380 True if the permission is allowed, False otherwise
382 Example metadata format:
383 "metadata": {
384 "allowed_vector_store_indexes": [
385 {
386 "index_name": "dall-e-3",
387 "index_permissions": ["write"]
388 }
389 ]
390 }
391 """
392 # Check both key_metadata and team_metadata
393 for metadata in [key_metadata, team_metadata]:
394 if metadata is None:
395 continue
397 allowed_indexes = metadata.get("allowed_vector_store_indexes")
398 if not allowed_indexes or not isinstance(allowed_indexes, list):
399 continue
401 # Look for matching index
402 for index_config in allowed_indexes:
403 if not isinstance(index_config, dict):
404 continue
406 if index_config.get("index_name") == index_name:
407 index_permissions = index_config.get("index_permissions", [])
408 if isinstance(index_permissions, list) and permission in index_permissions:
409 return True
411 return False
414def is_allowed_to_call_vector_store_endpoint(
415 provider: LlmProviders,
416 index_name: str,
417 request: Request,
418 user_api_key_dict: UserAPIKeyAuth,
419) -> Literal[True] | None:
420 """
421 Check if the user is allowed to call the vector store endpoint.
423 Cover:
424 1. Creating a vector store index
425 2. Reading a vector store index (Search / List / Get)
426 """
427 if (
428 user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN
429 or user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value
430 ):
431 return True
432 # check what allowed permissions are for the key
433 key_metadata: Final = user_api_key_dict.metadata
434 team_metadata: Final = user_api_key_dict.team_metadata
436 provider_config: Final = ProviderConfigManager.get_provider_vector_stores_config(provider=provider)
437 if provider_config is None:
438 return None
440 provider_vector_store_endpoints: Final = provider_config.get_vector_store_endpoints_by_type()
442 # Inline import — auth_utils participates in a proxy import cycle.
443 from litellm.proxy.auth.auth_utils import get_request_route # noqa: PLC0415
445 request_route: Final = get_request_route(request)
447 if _is_vector_store_index_lifecycle_request(
448 request_method=request.method,
449 request_path=request_route,
450 index_name=index_name,
451 ):
452 operation_label: Literal["create", "delete", "update"] = "create"
453 if request.method == "DELETE":
454 operation_label = "delete"
455 elif request.method in ("PUT", "PATCH"):
456 operation_label = "update"
457 assert_proxy_admin_for_vector_store_index_management(
458 user_api_key_dict,
459 operation=operation_label,
460 )
461 return True
463 # Writes are classified before reads so a path matching both patterns
464 # requires the stronger grant (e.g. the azure batch write on an index
465 # named "analyze*" also contains the "/analyze" read fragment)
466 permission_type = None
467 for endpoint in provider_vector_store_endpoints["write"]:
468 if request.method == endpoint[0] and _does_endpoint_match(endpoint[1], request_route):
469 permission_type = "write"
470 break
472 if permission_type is None:
473 for endpoint in provider_vector_store_endpoints["read"]:
474 if request.method == endpoint[0] and _does_endpoint_match(endpoint[1], request_route):
475 permission_type = "read"
476 break
478 if permission_type is None:
479 raise HTTPException(
480 status_code=403,
481 detail=(
482 f"User does not have permission to call vector store endpoint "
483 f"{index_name}. Ask your administrator to add the necessary "
484 "permissions to your API key/Team."
485 ),
486 )
488 # Check if key has specific permission for allowed_vector_store_indexes
489 has_permission: Final = check_vector_store_permission(
490 index_name=index_name,
491 permission=permission_type,
492 key_metadata=key_metadata,
493 team_metadata=team_metadata,
494 )
496 if not has_permission:
497 raise HTTPException(
498 status_code=403,
499 detail=f"User does not have permission to call vector store endpoint {index_name}. Ask your administrator to add the necessary permissions to your API key/Team.",
500 )
502 return has_permission
505def is_allowed_to_call_vector_store_files_endpoint(
506 provider: LlmProviders,
507 vector_store_id: str,
508 request: Request,
509 user_api_key_dict: UserAPIKeyAuth,
510) -> Literal[True] | None:
511 if (
512 user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN
513 or user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value
514 ):
515 return True
517 key_metadata: Final = user_api_key_dict.metadata
518 team_metadata: Final = user_api_key_dict.team_metadata
520 provider_config: Final = ProviderConfigManager.get_provider_vector_store_files_config(provider=provider)
521 if provider_config is None:
522 return None
524 provider_vector_store_endpoints: Final = provider_config.get_vector_store_file_endpoints_by_type()
526 # Inline import — auth_utils participates in a proxy import cycle.
527 from litellm.proxy.auth.auth_utils import get_request_route # noqa: PLC0415
529 request_route: Final = get_request_route(request)
531 permission_type: str | None = None
532 for endpoint in provider_vector_store_endpoints.get("write", ()):
533 if request.method == endpoint[0] and _does_endpoint_match(endpoint[1], request_route):
534 permission_type = "write"
535 break
537 if permission_type is None:
538 for endpoint in provider_vector_store_endpoints.get("read", ()):
539 if request.method == endpoint[0] and _does_endpoint_match(endpoint[1], request_route):
540 permission_type = "read"
541 break
543 if permission_type is None:
544 return None
546 has_permission: Final = check_vector_store_permission(
547 index_name=vector_store_id,
548 permission=permission_type,
549 key_metadata=key_metadata,
550 team_metadata=team_metadata,
551 )
553 if not has_permission:
554 raise HTTPException(
555 status_code=403,
556 detail=f"User does not have permission to call vector store file endpoint {vector_store_id}. Ask your administrator to add the necessary permissions to your API key/Team.",
557 )
559 return has_permission