mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-10 22:41:41 +00:00
Add get, list and delete for vector store endpoints
This commit is contained in:
Sameer Kankute
parent
18a05f7a40
commit
5927345eab
6 changed files with 1271 additions and 46 deletions
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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)],
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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"},
|
||||
)
|
||||
|
|
|
|||
364
tests/test_new_vector_store_endpoints.py
Normal file
364
tests/test_new_vector_store_endpoints.py
Normal file
|
|
@ -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!")
|
||||
Loading…
Add table
Reference in a new issue