From 18a05f7a408c3e9901ce154ed103ec6ed672484a Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Thu, 12 Mar 2026 11:58:44 +0530 Subject: [PATCH] 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 --- litellm/llms/custom_httpx/llm_http_handler.py | 536 ++++++++++++++++++ .../custom_httpx/test_llm_http_handler.py | 36 ++ 2 files changed, 572 insertions(+) diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 1cef3e9ce15..fc87697180f 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -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 ######################## ##################################################################### diff --git a/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py b/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py index 17b4243da1d..78da7820e89 100644 --- a/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py +++ b/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py @@ -215,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"}, + )