mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-07 08:26:10 +00:00
feat(vector-stores): add retrieve/list/update/delete handlers
- Add vector_store_retrieve/list/update/delete handlers in llm_http_handler - Fix AsyncHTTPHandler.get() timeout arg (not supported) - Fix update/delete URL (api_base already includes /vector_stores) - Clean metadata for update to avoid UserAPIKeyAuth JSON serialization Made-with: Cursor
This commit is contained in:
Sameer Kankute
parent
c4aa15b4e2
commit
18a05f7a40
2 changed files with 572 additions and 0 deletions
|
|
@ -7625,6 +7625,542 @@ 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(
|
||||
cast(Optional[Dict[str, Any]], 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(
|
||||
cast(Optional[Dict[str, Any]], 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 ########################
|
||||
#####################################################################
|
||||
|
|
|
|||
|
|
@ -215,3 +215,39 @@ 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"},
|
||||
)
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue