Coverage for .venv/lib/python3.13/site-packages/litellm/proxy/vector_store_endpoints/endpoints.py: 82%
142 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-10-10 12:01 +0000
1from collections.abc import Mapping
2from types import MappingProxyType
3from typing import (
4 Annotated,
5 Any, # noqa: TID251 # jsonify_object in proxy/utils.py is annotated with a bare dict
6 Final,
7 cast, # noqa: TID251 # jsonify_object in proxy/utils.py is annotated with a bare dict
8)
10from fastapi import APIRouter, Depends, HTTPException, Request, Response
12from litellm.integrations.vector_store_integrations.vector_store_pre_call_hook import (
13 LiteLLM_ManagedVectorStore,
14)
15from litellm.proxy._types import CommonProxyErrors, UserAPIKeyAuth
16from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
17from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing
18from litellm.proxy.utils import jsonify_object
19from litellm.proxy.vector_store_endpoints.utils import (
20 assert_proxy_admin_for_vector_store_index_management,
21 assert_user_can_access_vector_store,
22 get_litellm_managed_vector_store,
23)
24from litellm.repositories.table_repositories import ManagedVectorStoreIndexRepository
25from litellm.types.vector_stores import IndexCreateRequest, IndexListResponse
26from litellm.vector_stores.vector_store_registry import VectorStoreIndexRegistry
28router: Final = APIRouter()
30BLOCKED_QUERY_EMBEDDING_SELECTION_PARAMS: Final = frozenset(
31 {
32 "embedding_model",
33 "litellm_embedding_model",
34 "litellm_embedding_config",
35 "litellm_credential_name",
36 }
37)
40def reject_caller_embedding_selection_params(payload: Mapping[str, object], source: str) -> None:
41 blocked: Final = sorted(BLOCKED_QUERY_EMBEDDING_SELECTION_PARAMS & payload.keys())
42 if blocked: 42 ↛ 43line 42 didn't jump to line 43 because the condition on line 42 was never true
43 raise HTTPException(
44 status_code=400,
45 detail={
46 "error": f"'{blocked[0]}' cannot be set in {source}. "
47 "Embedding configuration comes from the vector store's server-side registration."
48 },
49 )
52########################################################
53# OpenAI Compatible Endpoints
54########################################################
57def build_request_data_from_managed_vector_store(
58 vector_store: LiteLLM_ManagedVectorStore,
59) -> Mapping[str, object]:
60 top_level: Final = MappingProxyType(
61 {
62 key: vector_store.get(key)
63 for key in ("custom_llm_provider", "litellm_credential_name")
64 if key in vector_store
65 }
66 )
67 litellm_params: Final = vector_store.get("litellm_params") or MappingProxyType({})
68 return MappingProxyType({**top_level, **litellm_params})
71async def _update_request_data_with_litellm_managed_vector_store_registry(
72 data: dict,
73 vector_store_id: str,
74 user_api_key_dict: UserAPIKeyAuth | None = None,
75) -> dict:
76 """
77 Update the request data with the litellm managed vector store registry.
79 Args:
80 data: Request data to update
81 vector_store_id: ID of the vector store
82 user_api_key_dict: User API key authentication info for access control
84 Raises:
85 HTTPException: If user doesn't have access to the vector store
86 """
87 vector_store_to_run: Final[LiteLLM_ManagedVectorStore | None] = await get_litellm_managed_vector_store(
88 vector_store_id=vector_store_id
89 )
90 if vector_store_to_run is None:
91 return data
92 if user_api_key_dict is not None: 92 ↛ 97line 92 didn't jump to line 97 because the condition on line 92 was always true
93 await assert_user_can_access_vector_store(
94 vector_store=vector_store_to_run,
95 user_api_key_dict=user_api_key_dict,
96 )
97 return {**data, **build_request_data_from_managed_vector_store(vector_store_to_run)}
100@router.post(
101 "/v1/vector_stores/{vector_store_id:path}/search",
102 dependencies=[Depends(user_api_key_auth)],
103)
104@router.post(
105 "/vector_stores/{vector_store_id:path}/search",
106 dependencies=[Depends(user_api_key_auth)],
107)
108async def vector_store_search(
109 request: Request,
110 vector_store_id: str,
111 fastapi_response: Response,
112 user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
113):
114 """
115 Search a vector store.
117 API Reference:
118 https://platform.openai.com/docs/api-reference/vector-stores/search
119 """
120 from litellm.proxy.proxy_server import (
121 _read_request_body,
122 general_settings,
123 llm_router,
124 proxy_config,
125 proxy_logging_obj,
126 select_data_generator,
127 user_api_base,
128 user_max_tokens,
129 user_model,
130 user_request_timeout,
131 user_temperature,
132 version,
133 )
135 data = await _read_request_body(request=request)
136 reject_caller_embedding_selection_params(payload=data, source="the search request body")
137 data["vector_store_id"] = vector_store_id
139 # Check for legacy vector store registry (non-managed vector stores)
140 data = await _update_request_data_with_litellm_managed_vector_store_registry(
141 data=data, vector_store_id=vector_store_id, user_api_key_dict=user_api_key_dict
142 )
144 # The managed_vector_stores pre-call hook will handle:
145 # 1. Decoding managed vector store IDs
146 # 2. Extracting model and provider resource ID
147 # 3. Setting up proper routing
148 # 4. Authentication checks
150 processor: Final = ProxyBaseLLMRequestProcessing(data=data)
151 try:
152 return await processor.base_process_llm_request(
153 request=request,
154 fastapi_response=fastapi_response,
155 user_api_key_dict=user_api_key_dict,
156 route_type="avector_store_search",
157 proxy_logging_obj=proxy_logging_obj,
158 llm_router=llm_router,
159 general_settings=general_settings,
160 proxy_config=proxy_config,
161 select_data_generator=select_data_generator,
162 model=None,
163 user_model=user_model,
164 user_temperature=user_temperature,
165 user_request_timeout=user_request_timeout,
166 user_max_tokens=user_max_tokens,
167 user_api_base=user_api_base,
168 version=version,
169 )
170 except Exception as e:
171 raise await processor._handle_llm_api_exception(
172 e=e,
173 user_api_key_dict=user_api_key_dict,
174 proxy_logging_obj=proxy_logging_obj,
175 version=version,
176 )
179@router.post("/v1/vector_stores", dependencies=[Depends(user_api_key_auth)])
180@router.post("/vector_stores", dependencies=[Depends(user_api_key_auth)])
181async def vector_store_create(
182 request: Request,
183 fastapi_response: Response,
184 user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
185):
186 """
187 Create a vector store.
189 API Reference:
190 https://platform.openai.com/docs/api-reference/vector-stores/create
192 Supports target_model_names parameter for creating vector stores across multiple models:
193 ```json
194 {
195 "name": "my-vector-store",
196 "target_model_names": "gpt-4,gemini-2.0"
197 }
198 ```
199 """
200 from litellm.proxy.proxy_server import (
201 _read_request_body,
202 general_settings,
203 llm_router,
204 proxy_config,
205 proxy_logging_obj,
206 select_data_generator,
207 user_api_base,
208 user_max_tokens,
209 user_model,
210 user_request_timeout,
211 user_temperature,
212 version,
213 )
215 data: Final = await _read_request_body(request=request)
217 # Check for target_model_names parameter
218 target_model_names: Final = data.pop("target_model_names", None)
220 if target_model_names: 220 ↛ 222line 220 didn't jump to line 222 because the condition on line 220 was never true
221 # Use managed vector stores for multi-model support
222 if isinstance(target_model_names, str):
223 target_model_names_list = [m.strip() for m in target_model_names.split(",")]
224 elif isinstance(target_model_names, list):
225 target_model_names_list = target_model_names
226 else:
227 raise HTTPException(
228 status_code=400,
229 detail="target_model_names must be a comma-separated string or list of model names",
230 )
232 # Get managed vector stores hook
233 managed_vector_stores: Final[Any] = proxy_logging_obj.get_proxy_hook("managed_vector_stores")
234 if managed_vector_stores is None:
235 raise HTTPException(
236 status_code=500,
237 detail="Managed vector stores not configured. Please ensure the proxy is initialized with database support.",
238 )
240 if llm_router is None:
241 raise HTTPException(
242 status_code=500,
243 detail="LLM Router not initialized. Ensure models are added to proxy.",
244 )
246 # Create vector store across multiple models
247 response: Final[object] = await managed_vector_stores.acreate_vector_store(
248 create_request=data,
249 llm_router=llm_router,
250 target_model_names_list=target_model_names_list,
251 litellm_parent_otel_span=user_api_key_dict.parent_otel_span,
252 user_api_key_dict=user_api_key_dict,
253 )
255 return response
257 processor: Final = ProxyBaseLLMRequestProcessing(data=data)
258 try:
259 return await processor.base_process_llm_request(
260 request=request,
261 fastapi_response=fastapi_response,
262 user_api_key_dict=user_api_key_dict,
263 route_type="avector_store_create",
264 proxy_logging_obj=proxy_logging_obj,
265 llm_router=llm_router,
266 general_settings=general_settings,
267 proxy_config=proxy_config,
268 select_data_generator=select_data_generator,
269 model=None,
270 user_model=user_model,
271 user_temperature=user_temperature,
272 user_request_timeout=user_request_timeout,
273 user_max_tokens=user_max_tokens,
274 user_api_base=user_api_base,
275 version=version,
276 )
277 except Exception as e:
278 raise await processor._handle_llm_api_exception(
279 e=e,
280 user_api_key_dict=user_api_key_dict,
281 proxy_logging_obj=proxy_logging_obj,
282 version=version,
283 )
286@router.get("/v1/vector_stores/{vector_store_id}", dependencies=[Depends(user_api_key_auth)])
287@router.get("/vector_stores/{vector_store_id}", dependencies=[Depends(user_api_key_auth)])
288async def vector_store_retrieve(
289 request: Request,
290 vector_store_id: str,
291 fastapi_response: Response,
292 user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
293):
294 """
295 Retrieve a vector store.
297 API Reference:
298 https://platform.openai.com/docs/api-reference/vector-stores/retrieve
299 """
300 from litellm.proxy.proxy_server import (
301 general_settings,
302 llm_router,
303 proxy_config,
304 proxy_logging_obj,
305 select_data_generator,
306 user_api_base,
307 user_max_tokens,
308 user_model,
309 user_request_timeout,
310 user_temperature,
311 version,
312 )
314 data = {"vector_store_id": vector_store_id}
316 data = await _update_request_data_with_litellm_managed_vector_store_registry(
317 data=data, vector_store_id=vector_store_id, user_api_key_dict=user_api_key_dict
318 )
320 processor: Final = ProxyBaseLLMRequestProcessing(data=data)
321 try:
322 return await processor.base_process_llm_request(
323 request=request,
324 fastapi_response=fastapi_response,
325 user_api_key_dict=user_api_key_dict,
326 route_type="avector_store_retrieve",
327 proxy_logging_obj=proxy_logging_obj,
328 llm_router=llm_router,
329 general_settings=general_settings,
330 proxy_config=proxy_config,
331 select_data_generator=select_data_generator,
332 model=None,
333 user_model=user_model,
334 user_temperature=user_temperature,
335 user_request_timeout=user_request_timeout,
336 user_max_tokens=user_max_tokens,
337 user_api_base=user_api_base,
338 version=version,
339 )
340 except Exception as e:
341 raise await processor._handle_llm_api_exception(
342 e=e,
343 user_api_key_dict=user_api_key_dict,
344 proxy_logging_obj=proxy_logging_obj,
345 version=version,
346 )
349@router.get("/v1/vector_stores", dependencies=[Depends(user_api_key_auth)])
350@router.get("/vector_stores", dependencies=[Depends(user_api_key_auth)])
351async def vector_store_list(
352 request: Request,
353 fastapi_response: Response,
354 after: str | None = None,
355 before: str | None = None,
356 limit: int | None = 20,
357 order: str | None = "desc",
358 user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
359):
360 """
361 List vector stores.
363 API Reference:
364 https://platform.openai.com/docs/api-reference/vector-stores/list
365 """
366 from litellm.proxy.proxy_server import (
367 general_settings,
368 llm_router,
369 proxy_config,
370 proxy_logging_obj,
371 select_data_generator,
372 user_api_base,
373 user_max_tokens,
374 user_model,
375 user_request_timeout,
376 user_temperature,
377 version,
378 )
380 data: Final[dict] = {}
381 if after is not None:
382 data["after"] = after
383 if before is not None:
384 data["before"] = before
385 if limit is not None: 385 ↛ 387line 385 didn't jump to line 387 because the condition on line 385 was always true
386 data["limit"] = limit
387 if order is not None: 387 ↛ 390line 387 didn't jump to line 390 because the condition on line 387 was always true
388 data["order"] = order
390 processor: Final = ProxyBaseLLMRequestProcessing(data=data)
391 try:
392 return await processor.base_process_llm_request(
393 request=request,
394 fastapi_response=fastapi_response,
395 user_api_key_dict=user_api_key_dict,
396 route_type="avector_store_list",
397 proxy_logging_obj=proxy_logging_obj,
398 llm_router=llm_router,
399 general_settings=general_settings,
400 proxy_config=proxy_config,
401 select_data_generator=select_data_generator,
402 model=None,
403 user_model=user_model,
404 user_temperature=user_temperature,
405 user_request_timeout=user_request_timeout,
406 user_max_tokens=user_max_tokens,
407 user_api_base=user_api_base,
408 version=version,
409 )
410 except Exception as e:
411 raise await processor._handle_llm_api_exception(
412 e=e,
413 user_api_key_dict=user_api_key_dict,
414 proxy_logging_obj=proxy_logging_obj,
415 version=version,
416 )
419@router.post("/v1/vector_stores/{vector_store_id}", dependencies=[Depends(user_api_key_auth)])
420@router.post("/vector_stores/{vector_store_id}", dependencies=[Depends(user_api_key_auth)])
421async def vector_store_update(
422 request: Request,
423 vector_store_id: str,
424 fastapi_response: Response,
425 user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
426):
427 """
428 Update a vector store.
430 API Reference:
431 https://platform.openai.com/docs/api-reference/vector-stores/modify
432 """
433 from litellm.proxy.proxy_server import (
434 _read_request_body,
435 general_settings,
436 llm_router,
437 proxy_config,
438 proxy_logging_obj,
439 select_data_generator,
440 user_api_base,
441 user_max_tokens,
442 user_model,
443 user_request_timeout,
444 user_temperature,
445 version,
446 )
448 data = await _read_request_body(request=request)
449 if "vector_store_id" not in data: 449 ↛ 452line 449 didn't jump to line 452 because the condition on line 449 was always true
450 data["vector_store_id"] = vector_store_id
452 data = await _update_request_data_with_litellm_managed_vector_store_registry(
453 data=data, vector_store_id=vector_store_id, user_api_key_dict=user_api_key_dict
454 )
456 processor: Final = ProxyBaseLLMRequestProcessing(data=data)
457 try:
458 return await processor.base_process_llm_request(
459 request=request,
460 fastapi_response=fastapi_response,
461 user_api_key_dict=user_api_key_dict,
462 route_type="avector_store_update",
463 proxy_logging_obj=proxy_logging_obj,
464 llm_router=llm_router,
465 general_settings=general_settings,
466 proxy_config=proxy_config,
467 select_data_generator=select_data_generator,
468 model=None,
469 user_model=user_model,
470 user_temperature=user_temperature,
471 user_request_timeout=user_request_timeout,
472 user_max_tokens=user_max_tokens,
473 user_api_base=user_api_base,
474 version=version,
475 )
476 except Exception as e:
477 raise await processor._handle_llm_api_exception(
478 e=e,
479 user_api_key_dict=user_api_key_dict,
480 proxy_logging_obj=proxy_logging_obj,
481 version=version,
482 )
485@router.delete("/v1/vector_stores/{vector_store_id}", dependencies=[Depends(user_api_key_auth)])
486@router.delete("/vector_stores/{vector_store_id}", dependencies=[Depends(user_api_key_auth)])
487async def vector_store_delete(
488 request: Request,
489 vector_store_id: str,
490 fastapi_response: Response,
491 user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
492):
493 """
494 Delete a vector store.
496 API Reference:
497 https://platform.openai.com/docs/api-reference/vector-stores/delete
498 """
499 from litellm.proxy.proxy_server import (
500 general_settings,
501 llm_router,
502 proxy_config,
503 proxy_logging_obj,
504 select_data_generator,
505 user_api_base,
506 user_max_tokens,
507 user_model,
508 user_request_timeout,
509 user_temperature,
510 version,
511 )
513 data = {"vector_store_id": vector_store_id}
515 data = await _update_request_data_with_litellm_managed_vector_store_registry(
516 data=data, vector_store_id=vector_store_id, user_api_key_dict=user_api_key_dict
517 )
519 processor: Final = ProxyBaseLLMRequestProcessing(data=data)
520 try:
521 return await processor.base_process_llm_request(
522 request=request,
523 fastapi_response=fastapi_response,
524 user_api_key_dict=user_api_key_dict,
525 route_type="avector_store_delete",
526 proxy_logging_obj=proxy_logging_obj,
527 llm_router=llm_router,
528 general_settings=general_settings,
529 proxy_config=proxy_config,
530 select_data_generator=select_data_generator,
531 model=None,
532 user_model=user_model,
533 user_temperature=user_temperature,
534 user_request_timeout=user_request_timeout,
535 user_max_tokens=user_max_tokens,
536 user_api_base=user_api_base,
537 version=version,
538 )
539 except Exception as e:
540 raise await processor._handle_llm_api_exception(
541 e=e,
542 user_api_key_dict=user_api_key_dict,
543 proxy_logging_obj=proxy_logging_obj,
544 version=version,
545 )
548@router.post(
549 "/v1/indexes",
550 dependencies=[Depends(user_api_key_auth)],
551)
552async def index_create(
553 request: Request,
554 index_create_request: IndexCreateRequest,
555 fastapi_response: Response,
556 user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
557):
558 """
559 Create an index. Just writes the index to the database.
561 ```bash
562 curl -L -X POST 'http://0.0.0.0:4000/v1/indexes' \
563 -H 'Content-Type: application/json' \
564 -H 'Authorization: Bearer sk-1234' \
565 -d '{
566 "index_name": "dall-e-3",
567 "litellm_params": {
568 "vector_store_index": "real-index-name",
569 "vector_store_name": "azure-ai-search"
570 }
571 }'
572 ```
573 """
574 from litellm.proxy.proxy_server import prisma_client
576 assert_proxy_admin_for_vector_store_index_management(
577 user_api_key_dict,
578 operation="create",
579 )
581 if prisma_client is None: 581 ↛ 582line 581 didn't jump to line 582 because the condition on line 581 was never true
582 raise HTTPException(
583 status_code=500,
584 detail=CommonProxyErrors.db_not_connected_error.value,
585 )
586 ## 1. check if index already exists
587 existing_index: Final = await ManagedVectorStoreIndexRepository(prisma_client).table.find_unique(
588 where={"index_name": index_create_request.index_name}
589 )
591 ## 2. set created_by and updated_by
593 if existing_index is not None:
594 raise HTTPException(
595 status_code=400,
596 detail=f"Index {index_create_request.index_name} already exists",
597 )
599 ## 2. create index
600 index_data: Final = index_create_request.model_dump(exclude_none=True)
601 index_data["created_by"] = user_api_key_dict.user_id
602 index_data["updated_by"] = user_api_key_dict.user_id
603 new_index = await ManagedVectorStoreIndexRepository(prisma_client).table.create(
604 data=cast( # cast-ok: jsonify_object deep-copies a model_dump, so keys are str and values plain objects
605 "dict[str, object]", jsonify_object(index_data)
606 )
607 )
609 return new_index.model_dump()
612@router.get(
613 "/v1/indexes",
614 dependencies=[Depends(user_api_key_auth)],
615 response_model=IndexListResponse,
616)
617async def index_list(
618 user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)],
619) -> IndexListResponse:
620 """
621 List all vector store indexes. Proxy admin only.
623 ```bash
624 curl -L -X GET 'http://0.0.0.0:4000/v1/indexes' \
625 -H 'Authorization: Bearer sk-1234'
626 ```
627 """
628 from litellm.proxy.proxy_server import prisma_client
630 assert_proxy_admin_for_vector_store_index_management(
631 user_api_key_dict,
632 operation="list",
633 )
635 if prisma_client is None: 635 ↛ 636line 635 didn't jump to line 636 because the condition on line 635 was never true
636 raise HTTPException(
637 status_code=500,
638 detail=CommonProxyErrors.db_not_connected_error.value,
639 )
641 indexes: Final = await VectorStoreIndexRegistry._get_vector_store_indexes_from_db(prisma_client)
642 return IndexListResponse(data=indexes)