diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 1cef3e9ce15..705aa942729 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -7625,6 +7625,536 @@ class BaseLLMHTTPHandler: response=response, ) + async def async_vector_store_retrieve_handler( + self, + vector_store_id: str, + vector_store_provider_config: BaseVectorStoreConfig, + custom_llm_provider: str, + litellm_params: GenericLiteLLMParams, + logging_obj: LiteLLMLoggingObj, + extra_headers: Optional[Dict[str, Any]] = None, + extra_body: Optional[Dict[str, Any]] = None, + timeout: Optional[Union[float, httpx.Timeout]] = None, + client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, + ) -> VectorStoreCreateResponse: + if client is None or not isinstance(client, AsyncHTTPHandler): + async_httpx_client = get_async_httpx_client( + llm_provider=litellm.LlmProviders(custom_llm_provider), + params={"ssl_verify": litellm_params.get("ssl_verify", None)}, + ) + else: + async_httpx_client = client + + headers = vector_store_provider_config.validate_environment( + headers=extra_headers or {}, litellm_params=litellm_params + ) + + if extra_headers: + headers.update(extra_headers) + + api_base = vector_store_provider_config.get_complete_url( + api_base=litellm_params.api_base, + litellm_params=dict(litellm_params), + ) + + url = f"{api_base}/{vector_store_id}" + + logging_obj.pre_call( + input="", + api_key="", + additional_args={ + "api_base": api_base, + "headers": headers, + }, + ) + + try: + response = await async_httpx_client.get( + url=url, headers=headers + ) + except Exception as e: + raise self._handle_error(e=e, provider_config=vector_store_provider_config) + + return vector_store_provider_config.transform_create_vector_store_response( + response=response, + ) + + def vector_store_retrieve_handler( + self, + vector_store_id: str, + vector_store_provider_config: BaseVectorStoreConfig, + custom_llm_provider: str, + litellm_params: GenericLiteLLMParams, + logging_obj: LiteLLMLoggingObj, + extra_headers: Optional[Dict[str, Any]] = None, + extra_body: Optional[Dict[str, Any]] = None, + timeout: Optional[Union[float, httpx.Timeout]] = None, + client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, + _is_async: bool = False, + ) -> Union[ + VectorStoreCreateResponse, Coroutine[Any, Any, VectorStoreCreateResponse] + ]: + if _is_async: + return self.async_vector_store_retrieve_handler( + vector_store_id=vector_store_id, + vector_store_provider_config=vector_store_provider_config, + litellm_params=litellm_params, + logging_obj=logging_obj, + custom_llm_provider=custom_llm_provider, + extra_headers=extra_headers, + extra_body=extra_body, + timeout=timeout, + client=client, + ) + + if client is None or not isinstance(client, HTTPHandler): + sync_httpx_client = _get_httpx_client( + params={"ssl_verify": litellm_params.get("ssl_verify", None)} + ) + else: + sync_httpx_client = client + + headers = vector_store_provider_config.validate_environment( + headers=extra_headers or {}, litellm_params=litellm_params + ) + + if extra_headers: + headers.update(extra_headers) + + api_base = vector_store_provider_config.get_complete_url( + api_base=litellm_params.api_base, + litellm_params=dict(litellm_params), + ) + + url = f"{api_base}/{vector_store_id}" + + logging_obj.pre_call( + input="", + api_key="", + additional_args={ + "api_base": api_base, + "headers": headers, + }, + ) + + try: + response = sync_httpx_client.get(url=url, headers=headers) + except Exception as e: + raise self._handle_error(e=e, provider_config=vector_store_provider_config) + + return vector_store_provider_config.transform_create_vector_store_response( + response=response, + ) + + async def async_vector_store_list_handler( + self, + after: Optional[str], + before: Optional[str], + limit: Optional[int], + order: Optional[str], + vector_store_provider_config: BaseVectorStoreConfig, + custom_llm_provider: str, + litellm_params: GenericLiteLLMParams, + logging_obj: LiteLLMLoggingObj, + extra_headers: Optional[Dict[str, Any]] = None, + extra_body: Optional[Dict[str, Any]] = None, + timeout: Optional[Union[float, httpx.Timeout]] = None, + client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, + ): + if client is None or not isinstance(client, AsyncHTTPHandler): + async_httpx_client = get_async_httpx_client( + llm_provider=litellm.LlmProviders(custom_llm_provider), + params={"ssl_verify": litellm_params.get("ssl_verify", None)}, + ) + else: + async_httpx_client = client + + headers = vector_store_provider_config.validate_environment( + headers=extra_headers or {}, litellm_params=litellm_params + ) + + if extra_headers: + headers.update(extra_headers) + + api_base = vector_store_provider_config.get_complete_url( + api_base=litellm_params.api_base, + litellm_params=dict(litellm_params), + ) + + url = api_base + + params = {} + if after is not None: + params["after"] = after + if before is not None: + params["before"] = before + if limit is not None: + params["limit"] = limit + if order is not None: + params["order"] = order + + logging_obj.pre_call( + input="", + api_key="", + additional_args={ + "api_base": api_base, + "headers": headers, + "params": params, + }, + ) + + try: + response = await async_httpx_client.get( + url=url, headers=headers, params=params + ) + except Exception as e: + raise self._handle_error(e=e, provider_config=vector_store_provider_config) + + return response.json() + + def vector_store_list_handler( + self, + after: Optional[str], + before: Optional[str], + limit: Optional[int], + order: Optional[str], + vector_store_provider_config: BaseVectorStoreConfig, + custom_llm_provider: str, + litellm_params: GenericLiteLLMParams, + logging_obj: LiteLLMLoggingObj, + extra_headers: Optional[Dict[str, Any]] = None, + extra_body: Optional[Dict[str, Any]] = None, + timeout: Optional[Union[float, httpx.Timeout]] = None, + client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, + _is_async: bool = False, + ): + if _is_async: + return self.async_vector_store_list_handler( + after=after, + before=before, + limit=limit, + order=order, + vector_store_provider_config=vector_store_provider_config, + litellm_params=litellm_params, + logging_obj=logging_obj, + custom_llm_provider=custom_llm_provider, + extra_headers=extra_headers, + extra_body=extra_body, + timeout=timeout, + client=client, + ) + + if client is None or not isinstance(client, HTTPHandler): + sync_httpx_client = _get_httpx_client( + params={"ssl_verify": litellm_params.get("ssl_verify", None)} + ) + else: + sync_httpx_client = client + + headers = vector_store_provider_config.validate_environment( + headers=extra_headers or {}, litellm_params=litellm_params + ) + + if extra_headers: + headers.update(extra_headers) + + api_base = vector_store_provider_config.get_complete_url( + api_base=litellm_params.api_base, + litellm_params=dict(litellm_params), + ) + + url = api_base + + params = {} + if after is not None: + params["after"] = after + if before is not None: + params["before"] = before + if limit is not None: + params["limit"] = limit + if order is not None: + params["order"] = order + + logging_obj.pre_call( + input="", + api_key="", + additional_args={ + "api_base": api_base, + "headers": headers, + "params": params, + }, + ) + + try: + response = sync_httpx_client.get(url=url, headers=headers, params=params) + except Exception as e: + raise self._handle_error(e=e, provider_config=vector_store_provider_config) + + return response.json() + + async def async_vector_store_update_handler( + self, + vector_store_id: str, + vector_store_update_optional_params: VectorStoreCreateOptionalRequestParams, + vector_store_provider_config: BaseVectorStoreConfig, + custom_llm_provider: str, + litellm_params: GenericLiteLLMParams, + logging_obj: LiteLLMLoggingObj, + extra_headers: Optional[Dict[str, Any]] = None, + extra_body: Optional[Dict[str, Any]] = None, + timeout: Optional[Union[float, httpx.Timeout]] = None, + client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, + ) -> VectorStoreCreateResponse: + if client is None or not isinstance(client, AsyncHTTPHandler): + async_httpx_client = get_async_httpx_client( + llm_provider=litellm.LlmProviders(custom_llm_provider), + params={"ssl_verify": litellm_params.get("ssl_verify", None)}, + ) + else: + async_httpx_client = client + + headers = vector_store_provider_config.validate_environment( + headers=extra_headers or {}, litellm_params=litellm_params + ) + + if extra_headers: + headers.update(extra_headers) + + api_base = vector_store_provider_config.get_complete_url( + api_base=litellm_params.api_base, + litellm_params=dict(litellm_params), + ) + + url = f"{api_base}/{vector_store_id}" + + request_body = dict(vector_store_update_optional_params) + + # 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(request_body["metadata"]) + + if extra_body: + request_body.update(extra_body) + + logging_obj.pre_call( + input="", + api_key="", + additional_args={ + "complete_input_dict": request_body, + "api_base": api_base, + "headers": headers, + }, + ) + + try: + response = await async_httpx_client.post( + url=url, headers=headers, json=request_body, timeout=timeout + ) + except Exception as e: + raise self._handle_error(e=e, provider_config=vector_store_provider_config) + + return vector_store_provider_config.transform_create_vector_store_response( + response=response, + ) + + def vector_store_update_handler( + self, + vector_store_id: str, + vector_store_update_optional_params: VectorStoreCreateOptionalRequestParams, + vector_store_provider_config: BaseVectorStoreConfig, + custom_llm_provider: str, + litellm_params: GenericLiteLLMParams, + logging_obj: LiteLLMLoggingObj, + extra_headers: Optional[Dict[str, Any]] = None, + extra_body: Optional[Dict[str, Any]] = None, + timeout: Optional[Union[float, httpx.Timeout]] = None, + client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, + _is_async: bool = False, + ) -> Union[ + VectorStoreCreateResponse, Coroutine[Any, Any, VectorStoreCreateResponse] + ]: + if _is_async: + return self.async_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, + litellm_params=litellm_params, + logging_obj=logging_obj, + custom_llm_provider=custom_llm_provider, + extra_headers=extra_headers, + extra_body=extra_body, + timeout=timeout, + client=client, + ) + + if client is None or not isinstance(client, HTTPHandler): + sync_httpx_client = _get_httpx_client( + params={"ssl_verify": litellm_params.get("ssl_verify", None)} + ) + else: + sync_httpx_client = client + + headers = vector_store_provider_config.validate_environment( + headers=extra_headers or {}, litellm_params=litellm_params + ) + + if extra_headers: + headers.update(extra_headers) + + api_base = vector_store_provider_config.get_complete_url( + api_base=litellm_params.api_base, + litellm_params=dict(litellm_params), + ) + + url = f"{api_base}/{vector_store_id}" + + request_body = dict(vector_store_update_optional_params) + + # 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(request_body["metadata"]) + + if extra_body: + request_body.update(extra_body) + + logging_obj.pre_call( + input="", + api_key="", + additional_args={ + "complete_input_dict": request_body, + "api_base": api_base, + "headers": headers, + }, + ) + + try: + response = sync_httpx_client.post( + url=url, headers=headers, json=request_body + ) + except Exception as e: + raise self._handle_error(e=e, provider_config=vector_store_provider_config) + + return vector_store_provider_config.transform_create_vector_store_response( + response=response, + ) + + async def async_vector_store_delete_handler( + self, + vector_store_id: str, + vector_store_provider_config: BaseVectorStoreConfig, + custom_llm_provider: str, + litellm_params: GenericLiteLLMParams, + logging_obj: LiteLLMLoggingObj, + extra_headers: Optional[Dict[str, Any]] = None, + extra_body: Optional[Dict[str, Any]] = None, + timeout: Optional[Union[float, httpx.Timeout]] = None, + client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, + ): + if client is None or not isinstance(client, AsyncHTTPHandler): + async_httpx_client = get_async_httpx_client( + llm_provider=litellm.LlmProviders(custom_llm_provider), + params={"ssl_verify": litellm_params.get("ssl_verify", None)}, + ) + else: + async_httpx_client = client + + headers = vector_store_provider_config.validate_environment( + headers=extra_headers or {}, litellm_params=litellm_params + ) + + if extra_headers: + headers.update(extra_headers) + + api_base = vector_store_provider_config.get_complete_url( + api_base=litellm_params.api_base, + litellm_params=dict(litellm_params), + ) + + url = f"{api_base}/{vector_store_id}" + + logging_obj.pre_call( + input="", + api_key="", + additional_args={ + "api_base": api_base, + "headers": headers, + }, + ) + + try: + response = await async_httpx_client.delete( + url=url, headers=headers, timeout=timeout + ) + except Exception as e: + raise self._handle_error(e=e, provider_config=vector_store_provider_config) + + return response.json() + + def vector_store_delete_handler( + self, + vector_store_id: str, + vector_store_provider_config: BaseVectorStoreConfig, + custom_llm_provider: str, + litellm_params: GenericLiteLLMParams, + logging_obj: LiteLLMLoggingObj, + extra_headers: Optional[Dict[str, Any]] = None, + extra_body: Optional[Dict[str, Any]] = None, + timeout: Optional[Union[float, httpx.Timeout]] = None, + client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, + _is_async: bool = False, + ): + if _is_async: + return self.async_vector_store_delete_handler( + vector_store_id=vector_store_id, + vector_store_provider_config=vector_store_provider_config, + litellm_params=litellm_params, + logging_obj=logging_obj, + custom_llm_provider=custom_llm_provider, + extra_headers=extra_headers, + extra_body=extra_body, + timeout=timeout, + client=client, + ) + + if client is None or not isinstance(client, HTTPHandler): + sync_httpx_client = _get_httpx_client( + params={"ssl_verify": litellm_params.get("ssl_verify", None)} + ) + else: + sync_httpx_client = client + + headers = vector_store_provider_config.validate_environment( + headers=extra_headers or {}, litellm_params=litellm_params + ) + + if extra_headers: + headers.update(extra_headers) + + api_base = vector_store_provider_config.get_complete_url( + api_base=litellm_params.api_base, + litellm_params=dict(litellm_params), + ) + + url = f"{api_base}/{vector_store_id}" + + logging_obj.pre_call( + input="", + api_key="", + additional_args={ + "api_base": api_base, + "headers": headers, + }, + ) + + try: + response = sync_httpx_client.delete(url=url, headers=headers) + except Exception as e: + raise self._handle_error(e=e, provider_config=vector_store_provider_config) + + return response.json() + ##################################################################### ################ Vector Store Files HANDLERS ######################## ##################################################################### 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_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!")