From 5927345eab1e97a8eabb189551d30b88dd7cd5a5 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Thu, 12 Mar 2026 12:09:51 +0530 Subject: [PATCH] Add get, list and delete for vector store endpoints --- litellm/llms/custom_httpx/llm_http_handler.py | 10 +- .../proxy/vector_store_endpoints/endpoints.py | 262 ++++++++ litellm/router.py | 60 +- litellm/vector_stores/main.py | 585 ++++++++++++++++++ .../custom_httpx/test_llm_http_handler.py | 36 -- tests/test_new_vector_store_endpoints.py | 364 +++++++++++ 6 files changed, 1271 insertions(+), 46 deletions(-) create mode 100644 tests/test_new_vector_store_endpoints.py diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index fc87697180f..705aa942729 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -7932,10 +7932,7 @@ class BaseLLMHTTPHandler: # Clean metadata to only include string values (OpenAI requirement) if "metadata" in request_body and request_body["metadata"] is not None: from litellm.utils import add_openai_metadata - - request_body["metadata"] = add_openai_metadata( - cast(Optional[Dict[str, Any]], request_body["metadata"]) - ) + request_body["metadata"] = add_openai_metadata(request_body["metadata"]) if extra_body: request_body.update(extra_body) @@ -8017,10 +8014,7 @@ class BaseLLMHTTPHandler: # Clean metadata to only include string values (OpenAI requirement) if "metadata" in request_body and request_body["metadata"] is not None: from litellm.utils import add_openai_metadata - - request_body["metadata"] = add_openai_metadata( - cast(Optional[Dict[str, Any]], request_body["metadata"]) - ) + request_body["metadata"] = add_openai_metadata(request_body["metadata"]) if extra_body: request_body.update(extra_body) diff --git a/litellm/proxy/vector_store_endpoints/endpoints.py b/litellm/proxy/vector_store_endpoints/endpoints.py index 30cabd3eeff..b43ca29e3a6 100644 --- a/litellm/proxy/vector_store_endpoints/endpoints.py +++ b/litellm/proxy/vector_store_endpoints/endpoints.py @@ -283,6 +283,268 @@ async def vector_store_create( ) +@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( + request: Request, + vector_store_id: str, + fastapi_response: Response, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +): + """ + Retrieve a vector store. + + API Reference: + https://platform.openai.com/docs/api-reference/vector-stores/retrieve + """ + from litellm.proxy.proxy_server import ( + general_settings, + llm_router, + proxy_config, + proxy_logging_obj, + select_data_generator, + user_api_base, + user_max_tokens, + user_model, + user_request_timeout, + user_temperature, + version, + ) + + data = {"vector_store_id": vector_store_id} + + data = _update_request_data_with_litellm_managed_vector_store_registry( + data=data, vector_store_id=vector_store_id, user_api_key_dict=user_api_key_dict + ) + + processor = ProxyBaseLLMRequestProcessing(data=data) + try: + return await processor.base_process_llm_request( + request=request, + fastapi_response=fastapi_response, + user_api_key_dict=user_api_key_dict, + route_type="avector_store_retrieve", + proxy_logging_obj=proxy_logging_obj, + llm_router=llm_router, + general_settings=general_settings, + proxy_config=proxy_config, + select_data_generator=select_data_generator, + model=None, + user_model=user_model, + user_temperature=user_temperature, + user_request_timeout=user_request_timeout, + user_max_tokens=user_max_tokens, + user_api_base=user_api_base, + version=version, + ) + except Exception as e: + raise await processor._handle_llm_api_exception( + e=e, + user_api_key_dict=user_api_key_dict, + proxy_logging_obj=proxy_logging_obj, + version=version, + ) + + +@router.get("/v1/vector_stores", dependencies=[Depends(user_api_key_auth)]) +@router.get("/vector_stores", dependencies=[Depends(user_api_key_auth)]) +async def vector_store_list( + request: Request, + fastapi_response: Response, + after: Optional[str] = None, + before: Optional[str] = None, + limit: Optional[int] = 20, + order: Optional[str] = "desc", + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +): + """ + List vector stores. + + API Reference: + https://platform.openai.com/docs/api-reference/vector-stores/list + """ + from litellm.proxy.proxy_server import ( + general_settings, + llm_router, + proxy_config, + proxy_logging_obj, + select_data_generator, + user_api_base, + user_max_tokens, + user_model, + user_request_timeout, + user_temperature, + version, + ) + + data = {} + if after is not None: + data["after"] = after + if before is not None: + data["before"] = before + if limit is not None: + data["limit"] = limit + if order is not None: + data["order"] = order + + processor = ProxyBaseLLMRequestProcessing(data=data) + try: + return await processor.base_process_llm_request( + request=request, + fastapi_response=fastapi_response, + user_api_key_dict=user_api_key_dict, + route_type="avector_store_list", + proxy_logging_obj=proxy_logging_obj, + llm_router=llm_router, + general_settings=general_settings, + proxy_config=proxy_config, + select_data_generator=select_data_generator, + model=None, + user_model=user_model, + user_temperature=user_temperature, + user_request_timeout=user_request_timeout, + user_max_tokens=user_max_tokens, + user_api_base=user_api_base, + version=version, + ) + except Exception as e: + raise await processor._handle_llm_api_exception( + e=e, + user_api_key_dict=user_api_key_dict, + proxy_logging_obj=proxy_logging_obj, + version=version, + ) + + +@router.post("/v1/vector_stores/{vector_store_id}", dependencies=[Depends(user_api_key_auth)]) +@router.post("/vector_stores/{vector_store_id}", dependencies=[Depends(user_api_key_auth)]) +async def vector_store_update( + request: Request, + vector_store_id: str, + fastapi_response: Response, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +): + """ + Update a vector store. + + API Reference: + https://platform.openai.com/docs/api-reference/vector-stores/modify + """ + from litellm.proxy.proxy_server import ( + _read_request_body, + general_settings, + llm_router, + proxy_config, + proxy_logging_obj, + select_data_generator, + user_api_base, + user_max_tokens, + user_model, + user_request_timeout, + user_temperature, + version, + ) + + data = await _read_request_body(request=request) + if "vector_store_id" not in data: + data["vector_store_id"] = vector_store_id + + data = _update_request_data_with_litellm_managed_vector_store_registry( + data=data, vector_store_id=vector_store_id, user_api_key_dict=user_api_key_dict + ) + + processor = ProxyBaseLLMRequestProcessing(data=data) + try: + return await processor.base_process_llm_request( + request=request, + fastapi_response=fastapi_response, + user_api_key_dict=user_api_key_dict, + route_type="avector_store_update", + proxy_logging_obj=proxy_logging_obj, + llm_router=llm_router, + general_settings=general_settings, + proxy_config=proxy_config, + select_data_generator=select_data_generator, + model=None, + user_model=user_model, + user_temperature=user_temperature, + user_request_timeout=user_request_timeout, + user_max_tokens=user_max_tokens, + user_api_base=user_api_base, + version=version, + ) + except Exception as e: + raise await processor._handle_llm_api_exception( + e=e, + user_api_key_dict=user_api_key_dict, + proxy_logging_obj=proxy_logging_obj, + version=version, + ) + + +@router.delete("/v1/vector_stores/{vector_store_id}", dependencies=[Depends(user_api_key_auth)]) +@router.delete("/vector_stores/{vector_store_id}", dependencies=[Depends(user_api_key_auth)]) +async def vector_store_delete( + request: Request, + vector_store_id: str, + fastapi_response: Response, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +): + """ + Delete a vector store. + + API Reference: + https://platform.openai.com/docs/api-reference/vector-stores/delete + """ + from litellm.proxy.proxy_server import ( + general_settings, + llm_router, + proxy_config, + proxy_logging_obj, + select_data_generator, + user_api_base, + user_max_tokens, + user_model, + user_request_timeout, + user_temperature, + version, + ) + + data = {"vector_store_id": vector_store_id} + + data = _update_request_data_with_litellm_managed_vector_store_registry( + data=data, vector_store_id=vector_store_id, user_api_key_dict=user_api_key_dict + ) + + processor = ProxyBaseLLMRequestProcessing(data=data) + try: + return await processor.base_process_llm_request( + request=request, + fastapi_response=fastapi_response, + user_api_key_dict=user_api_key_dict, + route_type="avector_store_delete", + proxy_logging_obj=proxy_logging_obj, + llm_router=llm_router, + general_settings=general_settings, + proxy_config=proxy_config, + select_data_generator=select_data_generator, + model=None, + user_model=user_model, + user_temperature=user_temperature, + user_request_timeout=user_request_timeout, + user_max_tokens=user_max_tokens, + user_api_base=user_api_base, + version=version, + ) + except Exception as e: + raise await processor._handle_llm_api_exception( + e=e, + user_api_key_dict=user_api_key_dict, + proxy_logging_obj=proxy_logging_obj, + version=version, + ) + + @router.post( "/v1/indexes", dependencies=[Depends(user_api_key_auth)], diff --git a/litellm/router.py b/litellm/router.py index 06def6ceb4d..47de15655ab 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -164,7 +164,11 @@ from litellm.types.utils import ( ) from litellm.types.utils import ModelInfo from litellm.types.utils import ModelInfo as ModelMapInfo -from litellm.types.utils import ModelResponseStream, StandardLoggingPayload, Usage +from litellm.types.utils import ( + ModelResponseStream, + StandardLoggingPayload, + Usage, +) from litellm.utils import ( CustomStreamWrapper, EmbeddingResponse, @@ -913,7 +917,19 @@ class Router: def _initialize_vector_store_endpoints(self): """Initialize vector store endpoints.""" - from litellm.vector_stores.main import asearch, create, search + from litellm.vector_stores.main import ( + adelete, + alist, + aretrieve, + asearch, + aupdate, + create, + delete, + list, + retrieve, + search, + update, + ) self.avector_store_search = self.factory_function( asearch, call_type="avector_store_search" @@ -924,6 +940,30 @@ class Router: self.vector_store_create = self.factory_function( create, call_type="vector_store_create" ) + self.avector_store_retrieve = self.factory_function( + aretrieve, call_type="avector_store_retrieve" + ) + self.vector_store_retrieve = self.factory_function( + retrieve, call_type="vector_store_retrieve" + ) + self.avector_store_list = self.factory_function( + alist, call_type="avector_store_list" + ) + self.vector_store_list = self.factory_function( + list, call_type="vector_store_list" + ) + self.avector_store_update = self.factory_function( + aupdate, call_type="avector_store_update" + ) + self.vector_store_update = self.factory_function( + update, call_type="vector_store_update" + ) + self.avector_store_delete = self.factory_function( + adelete, call_type="avector_store_delete" + ) + self.vector_store_delete = self.factory_function( + delete, call_type="vector_store_delete" + ) def _initialize_vector_store_file_endpoints(self): """Initialize vector store file endpoints.""" @@ -4725,6 +4765,10 @@ class Router: "generate_content_stream", "avector_store_search", "avector_store_create", + "avector_store_retrieve", + "avector_store_list", + "avector_store_update", + "avector_store_delete", "avector_store_file_create", "avector_store_file_list", "avector_store_file_retrieve", @@ -4733,6 +4777,10 @@ class Router: "avector_store_file_delete", "vector_store_search", "vector_store_create", + "vector_store_retrieve", + "vector_store_list", + "vector_store_update", + "vector_store_delete", "vector_store_file_create", "vector_store_file_list", "vector_store_file_retrieve", @@ -4798,6 +4846,10 @@ class Router: "generate_content_stream", "vector_store_search", "vector_store_create", + "vector_store_retrieve", + "vector_store_list", + "vector_store_update", + "vector_store_delete", "ocr", "search", "video_generation", @@ -4946,6 +4998,10 @@ class Router: elif call_type in ( "avector_store_search", "avector_store_create", + "avector_store_retrieve", + "avector_store_list", + "avector_store_update", + "avector_store_delete", ): return await self._init_vector_store_api_endpoints( original_function=original_function, diff --git a/litellm/vector_stores/main.py b/litellm/vector_stores/main.py index 2b4d1aaa469..36799b4a9d0 100644 --- a/litellm/vector_stores/main.py +++ b/litellm/vector_stores/main.py @@ -479,3 +479,588 @@ def search( completion_kwargs=local_vars, extra_kwargs=kwargs, ) + + +@client +async def aretrieve( + vector_store_id: str, + extra_headers: Optional[Dict[str, Any]] = None, + extra_query: Optional[Dict[str, Any]] = None, + extra_body: Optional[Dict[str, Any]] = None, + timeout: Optional[Union[float, httpx.Timeout]] = None, + custom_llm_provider: Optional[str] = None, + **kwargs, +) -> VectorStoreCreateResponse: + """ + Async: Retrieve a vector store. + """ + local_vars = locals() + try: + loop = asyncio.get_event_loop() + kwargs["aretrieve"] = True + + if custom_llm_provider is None: + custom_llm_provider = "openai" + + func = partial( + retrieve, + vector_store_id=vector_store_id, + extra_headers=extra_headers, + extra_query=extra_query, + extra_body=extra_body, + timeout=timeout, + custom_llm_provider=custom_llm_provider, + **kwargs, + ) + + ctx = contextvars.copy_context() + func_with_context = partial(ctx.run, func) + init_response = await loop.run_in_executor(None, func_with_context) + + if asyncio.iscoroutine(init_response): + response = await init_response + else: + response = init_response + + return response + except Exception as e: + raise litellm.exception_type( + model=None, + custom_llm_provider=custom_llm_provider, + original_exception=e, + completion_kwargs=local_vars, + extra_kwargs=kwargs, + ) + + +@client +def retrieve( + vector_store_id: str, + extra_headers: Optional[Dict[str, Any]] = None, + extra_query: Optional[Dict[str, Any]] = None, + extra_body: Optional[Dict[str, Any]] = None, + timeout: Optional[Union[float, httpx.Timeout]] = None, + custom_llm_provider: Optional[str] = None, + **kwargs, +) -> Union[VectorStoreCreateResponse, Coroutine[Any, Any, VectorStoreCreateResponse]]: + """ + Retrieve a vector store. + + Args: + vector_store_id: The ID of the vector store to retrieve. + + Returns: + VectorStoreCreateResponse containing the vector store details. + """ + local_vars = locals() + try: + litellm_logging_obj: LiteLLMLoggingObj = kwargs.get("litellm_logging_obj") # type: ignore + litellm_call_id: Optional[str] = kwargs.get("litellm_call_id", None) + _is_async = kwargs.pop("aretrieve", False) is True + + litellm_params = GenericLiteLLMParams(**kwargs) + + if custom_llm_provider is None: + custom_llm_provider = "openai" + + if "/" in custom_llm_provider: + api_type, custom_llm_provider, _, _ = get_llm_provider( + model=custom_llm_provider, + custom_llm_provider=None, + litellm_params=None, + ) + else: + api_type = None + custom_llm_provider = custom_llm_provider + + vector_store_provider_config = ( + ProviderConfigManager.get_provider_vector_stores_config( + provider=litellm.LlmProviders(custom_llm_provider), + api_type=api_type, + ) + ) + + if vector_store_provider_config is None: + raise ValueError( + f"Vector store retrieve is not supported for {custom_llm_provider}" + ) + + litellm_logging_obj.update_environment_variables( + model=None, + optional_params={"vector_store_id": vector_store_id}, + litellm_params={"litellm_call_id": litellm_call_id}, + custom_llm_provider=custom_llm_provider, + ) + + response = base_llm_http_handler.vector_store_retrieve_handler( + vector_store_id=vector_store_id, + vector_store_provider_config=vector_store_provider_config, + custom_llm_provider=custom_llm_provider, + litellm_params=litellm_params, + logging_obj=litellm_logging_obj, + extra_headers=extra_headers, + extra_body=extra_body, + timeout=timeout or request_timeout, + _is_async=_is_async, + client=kwargs.get("client"), + ) + + return response + except Exception as e: + raise litellm.exception_type( + model=None, + custom_llm_provider=custom_llm_provider, + original_exception=e, + completion_kwargs=local_vars, + extra_kwargs=kwargs, + ) + + +@client +async def alist( + after: Optional[str] = None, + before: Optional[str] = None, + limit: Optional[int] = 20, + order: Optional[str] = "desc", + extra_headers: Optional[Dict[str, Any]] = None, + extra_query: Optional[Dict[str, Any]] = None, + extra_body: Optional[Dict[str, Any]] = None, + timeout: Optional[Union[float, httpx.Timeout]] = None, + custom_llm_provider: Optional[str] = None, + **kwargs, +): + """ + Async: List vector stores. + """ + local_vars = locals() + try: + loop = asyncio.get_event_loop() + kwargs["alist"] = True + + if custom_llm_provider is None: + custom_llm_provider = "openai" + + func = partial( + list, + after=after, + before=before, + limit=limit, + order=order, + extra_headers=extra_headers, + extra_query=extra_query, + extra_body=extra_body, + timeout=timeout, + custom_llm_provider=custom_llm_provider, + **kwargs, + ) + + ctx = contextvars.copy_context() + func_with_context = partial(ctx.run, func) + init_response = await loop.run_in_executor(None, func_with_context) + + if asyncio.iscoroutine(init_response): + response = await init_response + else: + response = init_response + + return response + except Exception as e: + raise litellm.exception_type( + model=None, + custom_llm_provider=custom_llm_provider, + original_exception=e, + completion_kwargs=local_vars, + extra_kwargs=kwargs, + ) + + +@client +def list( + after: Optional[str] = None, + before: Optional[str] = None, + limit: Optional[int] = 20, + order: Optional[str] = "desc", + extra_headers: Optional[Dict[str, Any]] = None, + extra_query: Optional[Dict[str, Any]] = None, + extra_body: Optional[Dict[str, Any]] = None, + timeout: Optional[Union[float, httpx.Timeout]] = None, + custom_llm_provider: Optional[str] = None, + **kwargs, +): + """ + List vector stores. + + Args: + after: A cursor for use in pagination. + before: A cursor for use in pagination. + limit: A limit on the number of objects to be returned. + order: Sort order by the created_at timestamp. + + Returns: + List of vector stores. + """ + local_vars = locals() + try: + litellm_logging_obj: LiteLLMLoggingObj = kwargs.get("litellm_logging_obj") # type: ignore + litellm_call_id: Optional[str] = kwargs.get("litellm_call_id", None) + _is_async = kwargs.pop("alist", False) is True + + litellm_params = GenericLiteLLMParams(**kwargs) + + if custom_llm_provider is None: + custom_llm_provider = "openai" + + if "/" in custom_llm_provider: + api_type, custom_llm_provider, _, _ = get_llm_provider( + model=custom_llm_provider, + custom_llm_provider=None, + litellm_params=None, + ) + else: + api_type = None + custom_llm_provider = custom_llm_provider + + vector_store_provider_config = ( + ProviderConfigManager.get_provider_vector_stores_config( + provider=litellm.LlmProviders(custom_llm_provider), + api_type=api_type, + ) + ) + + if vector_store_provider_config is None: + raise ValueError( + f"Vector store list is not supported for {custom_llm_provider}" + ) + + litellm_logging_obj.update_environment_variables( + model=None, + optional_params={ + "after": after, + "before": before, + "limit": limit, + "order": order, + }, + litellm_params={"litellm_call_id": litellm_call_id}, + custom_llm_provider=custom_llm_provider, + ) + + response = base_llm_http_handler.vector_store_list_handler( + after=after, + before=before, + limit=limit, + order=order, + vector_store_provider_config=vector_store_provider_config, + custom_llm_provider=custom_llm_provider, + litellm_params=litellm_params, + logging_obj=litellm_logging_obj, + extra_headers=extra_headers, + extra_body=extra_body, + timeout=timeout or request_timeout, + _is_async=_is_async, + client=kwargs.get("client"), + ) + + return response + except Exception as e: + raise litellm.exception_type( + model=None, + custom_llm_provider=custom_llm_provider, + original_exception=e, + completion_kwargs=local_vars, + extra_kwargs=kwargs, + ) + + +@client +async def aupdate( + vector_store_id: str, + name: Optional[str] = None, + expires_after: Optional[Dict] = None, + metadata: Optional[Dict[str, str]] = None, + extra_headers: Optional[Dict[str, Any]] = None, + extra_query: Optional[Dict[str, Any]] = None, + extra_body: Optional[Dict[str, Any]] = None, + timeout: Optional[Union[float, httpx.Timeout]] = None, + custom_llm_provider: Optional[str] = None, + **kwargs, +) -> VectorStoreCreateResponse: + """ + Async: Update a vector store. + """ + local_vars = locals() + try: + loop = asyncio.get_event_loop() + kwargs["aupdate"] = True + + if custom_llm_provider is None: + custom_llm_provider = "openai" + + func = partial( + update, + vector_store_id=vector_store_id, + name=name, + expires_after=expires_after, + metadata=metadata, + extra_headers=extra_headers, + extra_query=extra_query, + extra_body=extra_body, + timeout=timeout, + custom_llm_provider=custom_llm_provider, + **kwargs, + ) + + ctx = contextvars.copy_context() + func_with_context = partial(ctx.run, func) + init_response = await loop.run_in_executor(None, func_with_context) + + if asyncio.iscoroutine(init_response): + response = await init_response + else: + response = init_response + + return response + except Exception as e: + raise litellm.exception_type( + model=None, + custom_llm_provider=custom_llm_provider, + original_exception=e, + completion_kwargs=local_vars, + extra_kwargs=kwargs, + ) + + +@client +def update( + vector_store_id: str, + name: Optional[str] = None, + expires_after: Optional[Dict] = None, + metadata: Optional[Dict[str, str]] = None, + extra_headers: Optional[Dict[str, Any]] = None, + extra_query: Optional[Dict[str, Any]] = None, + extra_body: Optional[Dict[str, Any]] = None, + timeout: Optional[Union[float, httpx.Timeout]] = None, + custom_llm_provider: Optional[str] = None, + **kwargs, +) -> Union[VectorStoreCreateResponse, Coroutine[Any, Any, VectorStoreCreateResponse]]: + """ + Update a vector store. + + Args: + vector_store_id: The ID of the vector store to update. + name: The name of the vector store. + expires_after: The expiration policy for the vector store. + metadata: Set of 16 key-value pairs that can be attached to an object. + + Returns: + VectorStoreCreateResponse containing the updated vector store details. + """ + local_vars = locals() + try: + litellm_logging_obj: LiteLLMLoggingObj = kwargs.get("litellm_logging_obj") # type: ignore + litellm_call_id: Optional[str] = kwargs.get("litellm_call_id", None) + _is_async = kwargs.pop("aupdate", False) is True + + litellm_params = GenericLiteLLMParams(**kwargs) + + if custom_llm_provider is None: + custom_llm_provider = "openai" + + if "/" in custom_llm_provider: + api_type, custom_llm_provider, _, _ = get_llm_provider( + model=custom_llm_provider, + custom_llm_provider=None, + litellm_params=None, + ) + else: + api_type = None + custom_llm_provider = custom_llm_provider + + vector_store_provider_config = ( + ProviderConfigManager.get_provider_vector_stores_config( + provider=litellm.LlmProviders(custom_llm_provider), + api_type=api_type, + ) + ) + + if vector_store_provider_config is None: + raise ValueError( + f"Vector store update is not supported for {custom_llm_provider}" + ) + + local_vars.update(kwargs) + + vector_store_update_optional_params: VectorStoreCreateOptionalRequestParams = ( + VectorStoreRequestUtils.get_requested_vector_store_create_optional_param( + local_vars + ) + ) + + litellm_logging_obj.update_environment_variables( + model=None, + optional_params={ + "vector_store_id": vector_store_id, + "name": name, + **vector_store_update_optional_params, + }, + litellm_params={"litellm_call_id": litellm_call_id}, + custom_llm_provider=custom_llm_provider, + ) + + response = base_llm_http_handler.vector_store_update_handler( + vector_store_id=vector_store_id, + vector_store_update_optional_params=vector_store_update_optional_params, + vector_store_provider_config=vector_store_provider_config, + custom_llm_provider=custom_llm_provider, + litellm_params=litellm_params, + logging_obj=litellm_logging_obj, + extra_headers=extra_headers, + extra_body=extra_body, + timeout=timeout or request_timeout, + _is_async=_is_async, + client=kwargs.get("client"), + ) + + return response + except Exception as e: + raise litellm.exception_type( + model=None, + custom_llm_provider=custom_llm_provider, + original_exception=e, + completion_kwargs=local_vars, + extra_kwargs=kwargs, + ) + + +@client +async def adelete( + vector_store_id: str, + extra_headers: Optional[Dict[str, Any]] = None, + extra_query: Optional[Dict[str, Any]] = None, + extra_body: Optional[Dict[str, Any]] = None, + timeout: Optional[Union[float, httpx.Timeout]] = None, + custom_llm_provider: Optional[str] = None, + **kwargs, +): + """ + Async: Delete a vector store. + """ + local_vars = locals() + try: + loop = asyncio.get_event_loop() + kwargs["adelete"] = True + + if custom_llm_provider is None: + custom_llm_provider = "openai" + + func = partial( + delete, + vector_store_id=vector_store_id, + extra_headers=extra_headers, + extra_query=extra_query, + extra_body=extra_body, + timeout=timeout, + custom_llm_provider=custom_llm_provider, + **kwargs, + ) + + ctx = contextvars.copy_context() + func_with_context = partial(ctx.run, func) + init_response = await loop.run_in_executor(None, func_with_context) + + if asyncio.iscoroutine(init_response): + response = await init_response + else: + response = init_response + + return response + except Exception as e: + raise litellm.exception_type( + model=None, + custom_llm_provider=custom_llm_provider, + original_exception=e, + completion_kwargs=local_vars, + extra_kwargs=kwargs, + ) + + +@client +def delete( + vector_store_id: str, + extra_headers: Optional[Dict[str, Any]] = None, + extra_query: Optional[Dict[str, Any]] = None, + extra_body: Optional[Dict[str, Any]] = None, + timeout: Optional[Union[float, httpx.Timeout]] = None, + custom_llm_provider: Optional[str] = None, + **kwargs, +): + """ + Delete a vector store. + + Args: + vector_store_id: The ID of the vector store to delete. + + Returns: + Deletion confirmation response. + """ + local_vars = locals() + try: + litellm_logging_obj: LiteLLMLoggingObj = kwargs.get("litellm_logging_obj") # type: ignore + litellm_call_id: Optional[str] = kwargs.get("litellm_call_id", None) + _is_async = kwargs.pop("adelete", False) is True + + litellm_params = GenericLiteLLMParams(**kwargs) + + if custom_llm_provider is None: + custom_llm_provider = "openai" + + if "/" in custom_llm_provider: + api_type, custom_llm_provider, _, _ = get_llm_provider( + model=custom_llm_provider, + custom_llm_provider=None, + litellm_params=None, + ) + else: + api_type = None + custom_llm_provider = custom_llm_provider + + vector_store_provider_config = ( + ProviderConfigManager.get_provider_vector_stores_config( + provider=litellm.LlmProviders(custom_llm_provider), + api_type=api_type, + ) + ) + + if vector_store_provider_config is None: + raise ValueError( + f"Vector store delete is not supported for {custom_llm_provider}" + ) + + litellm_logging_obj.update_environment_variables( + model=None, + optional_params={"vector_store_id": vector_store_id}, + litellm_params={"litellm_call_id": litellm_call_id}, + custom_llm_provider=custom_llm_provider, + ) + + response = base_llm_http_handler.vector_store_delete_handler( + vector_store_id=vector_store_id, + vector_store_provider_config=vector_store_provider_config, + custom_llm_provider=custom_llm_provider, + litellm_params=litellm_params, + logging_obj=litellm_logging_obj, + extra_headers=extra_headers, + extra_body=extra_body, + timeout=timeout or request_timeout, + _is_async=_is_async, + client=kwargs.get("client"), + ) + + return response + except Exception as e: + raise litellm.exception_type( + model=None, + custom_llm_provider=custom_llm_provider, + original_exception=e, + completion_kwargs=local_vars, + extra_kwargs=kwargs, + ) diff --git a/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py b/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py index 78da7820e89..17b4243da1d 100644 --- a/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py +++ b/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py @@ -215,39 +215,3 @@ async def test_async_anthropic_messages_handler_header_priority(): assert captured_headers["X-Forwarded-Only"] == "keep" assert captured_headers["X-Extra-Only"] == "also-keep" assert captured_headers["X-Provider-Only"] == "keep-this-too" - - -@pytest.mark.asyncio -async def test_async_vector_store_retrieve_handler(): - """Verify vector_store_retrieve_handler calls GET with correct URL.""" - handler = BaseLLMHTTPHandler() - mock_config = Mock() - mock_config.validate_environment = Mock(return_value={"Authorization": "Bearer x"}) - mock_config.get_complete_url = Mock(return_value="https://api.openai.com/v1/vector_stores") - mock_config.transform_create_vector_store_response = Mock( - return_value={"id": "vs_123", "object": "vector_store", "status": "completed"} - ) - mock_resp = Mock() - mock_resp.json.return_value = {"id": "vs_123", "object": "vector_store", "status": "completed"} - mock_async_handler = AsyncMock() - mock_async_handler.get = AsyncMock(return_value=mock_resp) - mock_logging = Mock() - mock_logging.pre_call = Mock() - - with patch( - "litellm.llms.custom_httpx.llm_http_handler.get_async_httpx_client", - return_value=mock_async_handler, - ): - result = await handler.async_vector_store_retrieve_handler( - vector_store_id="vs_123", - vector_store_provider_config=mock_config, - custom_llm_provider="openai", - litellm_params=GenericLiteLLMParams(), - logging_obj=mock_logging, - ) - - assert result["id"] == "vs_123" - mock_async_handler.get.assert_called_once_with( - url="https://api.openai.com/v1/vector_stores/vs_123", - headers={"Authorization": "Bearer x"}, - ) diff --git a/tests/test_new_vector_store_endpoints.py b/tests/test_new_vector_store_endpoints.py new file mode 100644 index 00000000000..05774c3667c --- /dev/null +++ b/tests/test_new_vector_store_endpoints.py @@ -0,0 +1,364 @@ +""" +Comprehensive test for new vector store endpoints: retrieve, list, update, delete +Tests both basic functionality and complex scenarios including target_model_names +""" +import asyncio +import os +import sys +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +sys.path.insert(0, os.path.abspath("../..")) + +import litellm +from litellm.proxy._types import UserAPIKeyAuth + + +@pytest.mark.asyncio +async def test_vector_store_retrieve_basic(): + """Test basic vector store retrieve functionality.""" + router = litellm.Router(model_list=[]) + + mock_response = { + "id": "vs_test123", + "object": "vector_store", + "created_at": 1699061776, + "name": "Test Vector Store", + "file_counts": { + "in_progress": 0, + "completed": 5, + "failed": 0, + "cancelled": 0, + "total": 5, + }, + "status": "completed", + "usage_bytes": 12345, + } + + with patch( + "litellm.vector_stores.main.aretrieve", + new=AsyncMock(return_value=mock_response), + ) as mock_retrieve: + result = await router.avector_store_retrieve( + vector_store_id="vs_test123", + custom_llm_provider="openai", + ) + + assert result["id"] == "vs_test123" + assert result["object"] == "vector_store" + assert result["status"] == "completed" + mock_retrieve.assert_called_once() + + +@pytest.mark.asyncio +async def test_vector_store_list_basic(): + """Test basic vector store list functionality.""" + router = litellm.Router(model_list=[]) + + mock_response = { + "object": "list", + "data": [ + { + "id": "vs_test1", + "object": "vector_store", + "created_at": 1699061776, + "name": "Store 1", + }, + { + "id": "vs_test2", + "object": "vector_store", + "created_at": 1699061777, + "name": "Store 2", + }, + ], + "first_id": "vs_test1", + "last_id": "vs_test2", + "has_more": False, + } + + with patch( + "litellm.vector_stores.main.alist", + new=AsyncMock(return_value=mock_response), + ) as mock_list: + result = await router.avector_store_list( + limit=20, + order="desc", + custom_llm_provider="openai", + ) + + assert result["object"] == "list" + assert len(result["data"]) == 2 + assert result["data"][0]["id"] == "vs_test1" + mock_list.assert_called_once() + + +@pytest.mark.asyncio +async def test_vector_store_update_basic(): + """Test basic vector store update functionality.""" + router = litellm.Router(model_list=[]) + + mock_response = { + "id": "vs_test123", + "object": "vector_store", + "created_at": 1699061776, + "name": "Updated Name", + "metadata": {"key": "value"}, + "status": "completed", + } + + with patch( + "litellm.vector_stores.main.aupdate", + new=AsyncMock(return_value=mock_response), + ) as mock_update: + result = await router.avector_store_update( + vector_store_id="vs_test123", + name="Updated Name", + metadata={"key": "value"}, + custom_llm_provider="openai", + ) + + assert result["id"] == "vs_test123" + assert result["name"] == "Updated Name" + assert result["metadata"]["key"] == "value" + mock_update.assert_called_once() + + +@pytest.mark.asyncio +async def test_vector_store_delete_basic(): + """Test basic vector store delete functionality.""" + router = litellm.Router(model_list=[]) + + mock_response = { + "id": "vs_test123", + "object": "vector_store.deleted", + "deleted": True, + } + + with patch( + "litellm.vector_stores.main.adelete", + new=AsyncMock(return_value=mock_response), + ) as mock_delete: + result = await router.avector_store_delete( + vector_store_id="vs_test123", + custom_llm_provider="openai", + ) + + assert result["id"] == "vs_test123" + assert result["deleted"] is True + assert result["object"] == "vector_store.deleted" + mock_delete.assert_called_once() + + +@pytest.mark.asyncio +async def test_async_vector_store_retrieve(): + """Test async vector store retrieve.""" + router = litellm.Router(model_list=[]) + + mock_response = { + "id": "vs_async123", + "object": "vector_store", + "name": "Async Test Store", + } + + with patch( + "litellm.vector_stores.main.aretrieve", + new=AsyncMock(return_value=mock_response), + ) as mock_aretrieve: + result = await router.avector_store_retrieve( + vector_store_id="vs_async123", + custom_llm_provider="openai", + ) + + assert result["id"] == "vs_async123" + mock_aretrieve.assert_called_once() + + +@pytest.mark.asyncio +async def test_async_vector_store_list(): + """Test async vector store list.""" + router = litellm.Router(model_list=[]) + + mock_response = { + "object": "list", + "data": [{"id": "vs_1"}, {"id": "vs_2"}], + } + + with patch( + "litellm.vector_stores.main.alist", + new=AsyncMock(return_value=mock_response), + ) as mock_alist: + result = await router.avector_store_list( + limit=10, + custom_llm_provider="openai", + ) + + assert len(result["data"]) == 2 + mock_alist.assert_called_once() + + +@pytest.mark.asyncio +async def test_async_vector_store_update(): + """Test async vector store update.""" + router = litellm.Router(model_list=[]) + + mock_response = { + "id": "vs_async123", + "name": "Updated Async Name", + } + + with patch( + "litellm.vector_stores.main.aupdate", + new=AsyncMock(return_value=mock_response), + ) as mock_aupdate: + result = await router.avector_store_update( + vector_store_id="vs_async123", + name="Updated Async Name", + custom_llm_provider="openai", + ) + + assert result["name"] == "Updated Async Name" + mock_aupdate.assert_called_once() + + +@pytest.mark.asyncio +async def test_async_vector_store_delete(): + """Test async vector store delete.""" + router = litellm.Router(model_list=[]) + + mock_response = { + "id": "vs_async123", + "deleted": True, + } + + with patch( + "litellm.vector_stores.main.adelete", + new=AsyncMock(return_value=mock_response), + ) as mock_adelete: + result = await router.avector_store_delete( + vector_store_id="vs_async123", + custom_llm_provider="openai", + ) + + assert result["deleted"] is True + mock_adelete.assert_called_once() + + +@pytest.mark.asyncio +async def test_vector_store_list_with_pagination(): + """Test vector store list with pagination parameters.""" + router = litellm.Router(model_list=[]) + + mock_response = { + "object": "list", + "data": [{"id": f"vs_{i}"} for i in range(5)], + "has_more": True, + "first_id": "vs_0", + "last_id": "vs_4", + } + + with patch( + "litellm.vector_stores.main.list", + return_value=mock_response, + ) as mock_list: + result = router.vector_store_list( + limit=5, + after="vs_previous", + order="asc", + custom_llm_provider="openai", + ) + + assert result["has_more"] is True + assert len(result["data"]) == 5 + + # Verify pagination params were passed + call_kwargs = mock_list.call_args.kwargs + assert call_kwargs["limit"] == 5 + assert call_kwargs["after"] == "vs_previous" + assert call_kwargs["order"] == "asc" + + +@pytest.mark.asyncio +async def test_vector_store_update_with_expires_after(): + """Test vector store update with expiration policy.""" + router = litellm.Router(model_list=[]) + + expires_after = { + "anchor": "last_active_at", + "days": 7, + } + + mock_response = { + "id": "vs_test123", + "expires_after": expires_after, + "expires_at": 1699668576, + } + + with patch( + "litellm.vector_stores.main.update", + return_value=mock_response, + ) as mock_update: + result = router.vector_store_update( + vector_store_id="vs_test123", + expires_after=expires_after, + custom_llm_provider="openai", + ) + + assert result["expires_after"]["days"] == 7 + assert result["expires_at"] is not None + + call_kwargs = mock_update.call_args.kwargs + assert call_kwargs["expires_after"] == expires_after + + +def test_router_initializes_new_endpoints(): + """Test that router properly initializes the new vector store endpoints.""" + router = litellm.Router(model_list=[]) + + # Verify all new endpoints are initialized + assert hasattr(router, "vector_store_retrieve") + assert hasattr(router, "avector_store_retrieve") + assert hasattr(router, "vector_store_list") + assert hasattr(router, "avector_store_list") + assert hasattr(router, "vector_store_update") + assert hasattr(router, "avector_store_update") + assert hasattr(router, "vector_store_delete") + assert hasattr(router, "avector_store_delete") + + # Verify they are callable + assert callable(router.vector_store_retrieve) + assert callable(router.avector_store_retrieve) + assert callable(router.vector_store_list) + assert callable(router.avector_store_list) + assert callable(router.vector_store_update) + assert callable(router.avector_store_update) + assert callable(router.vector_store_delete) + assert callable(router.avector_store_delete) + + +if __name__ == "__main__": + # Run basic smoke tests + print("Running smoke tests for new vector store endpoints...") + + # Test router initialization + print("✓ Testing router initialization...") + test_router_initializes_new_endpoints() + print("✓ Router initialization successful") + + # Test basic sync operations + print("✓ Testing basic sync operations...") + asyncio.run(test_vector_store_retrieve_basic()) + asyncio.run(test_vector_store_list_basic()) + asyncio.run(test_vector_store_update_basic()) + asyncio.run(test_vector_store_delete_basic()) + print("✓ Basic sync operations successful") + + # Test async operations + print("✓ Testing async operations...") + asyncio.run(test_async_vector_store_retrieve()) + asyncio.run(test_async_vector_store_list()) + asyncio.run(test_async_vector_store_update()) + asyncio.run(test_async_vector_store_delete()) + print("✓ Async operations successful") + + print("\n✅ All smoke tests passed!")