diff --git a/litellm/llms/base_llm/vector_store/transformation.py b/litellm/llms/base_llm/vector_store/transformation.py index 07b60cb4b72..f7044249ffc 100644 --- a/litellm/llms/base_llm/vector_store/transformation.py +++ b/litellm/llms/base_llm/vector_store/transformation.py @@ -124,6 +124,9 @@ class BaseVectorStoreConfig: def validate_create_vector_store(self) -> None: return None + def get_httpx_client_params(self) -> Mapping[str, object]: + return MappingProxyType({}) + def get_supported_openai_params(self, model: str) -> list[VECTOR_STORE_OPENAI_PARAMS]: return [] diff --git a/litellm/llms/custom_httpx/http_handler.py b/litellm/llms/custom_httpx/http_handler.py index f4883b57fbc..ecfb1184943 100644 --- a/litellm/llms/custom_httpx/http_handler.py +++ b/litellm/llms/custom_httpx/http_handler.py @@ -236,9 +236,9 @@ def _prepare_request_data_and_content( # Cache for SSL contexts to avoid creating duplicate contexts with the same configuration -# Key: tuple of (cafile, ssl_security_level, ssl_ecdh_curve) +# Key: tuple of (cafile, ssl_security_level, ssl_ecdh_curve, http2) # Value: ssl.SSLContext -_ssl_context_cache: Final[dict[tuple[str | None, str | None, str | None], ssl.SSLContext]] = {} +_ssl_context_cache: Final[dict[tuple[str | None, str | None, str | None, bool], ssl.SSLContext]] = {} def _create_ssl_context( @@ -331,6 +331,7 @@ def get_ssl_verify( def get_ssl_configuration( ssl_verify: VerifyTypes | None = None, + http2: bool = False, ) -> bool | str | ssl.SSLContext: """ Unified SSL configuration function that handles ssl_context and ssl_verify logic. @@ -346,6 +347,7 @@ def get_ssl_configuration( SSL contexts are cached to avoid creating duplicate contexts with the same configuration, which reduces memory allocation and improves performance. + http2 clients get their own cached context because httpcore sets ALPN protocols on it. Args: ssl_verify: SSL verification setting. Can be: @@ -379,7 +381,7 @@ def get_ssl_configuration( if ssl_verify is not False: # Create cache key from configuration parameters - cache_key: Final = (cafile, ssl_security_level, ssl_ecdh_curve) + cache_key: Final = (cafile, ssl_security_level, ssl_ecdh_curve, http2) # Check if we have a cached SSL context for this configuration if cache_key not in _ssl_context_cache: @@ -571,17 +573,20 @@ class AsyncHTTPHandler: client_alias: str | None = None, # name for client in logs ssl_verify: VerifyTypes | None = None, shared_session: Optional["ClientSession"] = None, + http2: bool = False, ): self.timeout = timeout self.event_hooks = event_hooks self.ssl_verify = ssl_verify self.shared_session = shared_session + self.http2 = http2 self._owns_client = True self._client = self.create_client( timeout=timeout, event_hooks=event_hooks, ssl_verify=ssl_verify, shared_session=shared_session, + http2=http2, ) self.client_alias = client_alias @@ -593,6 +598,7 @@ class AsyncHTTPHandler: event_hooks=self.event_hooks, ssl_verify=self.ssl_verify, shared_session=self.shared_session, + http2=self.http2, ) return self._client @@ -607,9 +613,10 @@ class AsyncHTTPHandler: event_hooks: Mapping[str, list[Callable[..., object]]] | None, ssl_verify: VerifyTypes | None = None, shared_session: Optional["ClientSession"] = None, + http2: bool = False, ) -> httpx.AsyncClient: # Get unified SSL configuration - ssl_config: Final = get_ssl_configuration(ssl_verify) + ssl_config: Final = get_ssl_configuration(ssl_verify, http2=http2) # An SSL certificate used by the requested host to authenticate the client. # /path/to/client.pem @@ -619,10 +626,19 @@ class AsyncHTTPHandler: timeout = _DEFAULT_TIMEOUT # Create a client with a connection pool - transport: Final = AsyncHTTPHandler._create_async_transport( - ssl_context=ssl_config if isinstance(ssl_config, ssl.SSLContext) else None, - ssl_verify=ssl_config if isinstance(ssl_config, bool) else None, - shared_session=shared_session, + transport: Final = ( + httpx.AsyncHTTPTransport( + http2=True, + verify=ssl_config, + cert=cert, + local_address=_IPV4_LOCAL_ADDRESS if litellm.force_ipv4 else None, + ) + if http2 + else AsyncHTTPHandler._create_async_transport( + ssl_context=ssl_config if isinstance(ssl_config, ssl.SSLContext) else None, + ssl_verify=ssl_config if isinstance(ssl_config, bool) else None, + shared_session=shared_session, + ) ) # Get default headers (User-Agent, overridable via LITELLM_USER_AGENT) @@ -630,7 +646,7 @@ class AsyncHTTPHandler: return httpx.AsyncClient( transport=transport, - mounts=AsyncHTTPHandler._create_httpx_proxy_mounts(transport, verify=ssl_config, cert=cert), + mounts=AsyncHTTPHandler._create_httpx_proxy_mounts(transport, verify=ssl_config, cert=cert, http2=http2), event_hooks=event_hooks, timeout=timeout, verify=ssl_config, @@ -1296,11 +1312,12 @@ class AsyncHTTPHandler: transport: LiteLLMAiohttpTransport | AsyncHTTPTransport | None, verify: VerifyTypes, cert: CertTypes | None, + http2: bool = False, ) -> Mapping[str, AsyncHTTPTransport | None] | None: if not isinstance(transport, AsyncHTTPTransport): return None return _environment_proxy_mounts( - lambda proxy_url: AsyncHTTPTransport(proxy=proxy_url, verify=verify, cert=cert) + lambda proxy_url: AsyncHTTPTransport(proxy=proxy_url, verify=verify, cert=cert, http2=http2) ) @@ -1313,17 +1330,19 @@ class HTTPHandler: ssl_verify: bool | str | None = None, disable_default_headers: bool | None = False, # arize phoenix returns different API responses when user agent header in request + http2: bool = False, ): self.timeout = timeout self.ssl_verify = ssl_verify self.disable_default_headers = disable_default_headers + self.http2 = http2 self._owns_client = client is None self._heal_lock = threading.Lock() self._client = self.create_client() if client is None else client def create_client(self) -> httpx.Client: # Get unified SSL configuration - ssl_config: Final = get_ssl_configuration(self.ssl_verify) + ssl_config: Final = get_ssl_configuration(self.ssl_verify, http2=self.http2) # An SSL certificate used by the requested host to authenticate the client. # /path/to/client.pem @@ -1335,7 +1354,8 @@ class HTTPHandler: # Create a client with a connection pool return httpx.Client( transport=self._create_sync_transport(), - mounts=self._create_sync_proxy_mounts(verify=ssl_config, cert=cert), + mounts=self._create_sync_proxy_mounts(verify=ssl_config, cert=cert, http2=self.http2), + http2=self.http2, timeout=self.timeout if self.timeout is not None else _DEFAULT_TIMEOUT, verify=ssl_config, cert=cert, @@ -1616,7 +1636,7 @@ class HTTPHandler: Some users have seen httpx ConnectionError when using ipv6 - forcing ipv4 resolves the issue for them """ if litellm.force_ipv4: - return HTTPTransport(local_address=_IPV4_LOCAL_ADDRESS) + return HTTPTransport(http2=self.http2, local_address=_IPV4_LOCAL_ADDRESS) else: return getattr(litellm, "sync_transport", None) @@ -1624,10 +1644,13 @@ class HTTPHandler: def _create_sync_proxy_mounts( verify: VerifyTypes, cert: CertTypes | None, + http2: bool = False, ) -> Mapping[str, HTTPTransport | None] | None: if not litellm.force_ipv4: return None - return _environment_proxy_mounts(lambda proxy_url: HTTPTransport(proxy=proxy_url, verify=verify, cert=cert)) + return _environment_proxy_mounts( + lambda proxy_url: HTTPTransport(proxy=proxy_url, verify=verify, cert=cert, http2=http2) + ) def get_async_httpx_client( diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 7109e6942d1..05df5741029 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -9821,7 +9821,10 @@ class BaseLLMHTTPHandler: 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)}, + params={ + "ssl_verify": litellm_params.get("ssl_verify", None), + **vector_store_provider_config.get_httpx_client_params(), + }, ) else: async_httpx_client = client @@ -9959,7 +9962,12 @@ class BaseLLMHTTPHandler: ) if client is None or not isinstance(client, HTTPHandler): - sync_httpx_client = _get_httpx_client(params={"ssl_verify": litellm_params.get("ssl_verify", None)}) + sync_httpx_client = _get_httpx_client( + params={ + "ssl_verify": litellm_params.get("ssl_verify", None), + **vector_store_provider_config.get_httpx_client_params(), + } + ) else: sync_httpx_client = client diff --git a/litellm/llms/vertex_ai/vector_stores/search_api/transformation.py b/litellm/llms/vertex_ai/vector_stores/search_api/transformation.py index 0bcf16ee06f..2578b44b8ef 100644 --- a/litellm/llms/vertex_ai/vector_stores/search_api/transformation.py +++ b/litellm/llms/vertex_ai/vector_stores/search_api/transformation.py @@ -1,4 +1,6 @@ +import importlib.util from collections.abc import Mapping +from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Protocol import httpx @@ -51,6 +53,8 @@ VERTEX_SEARCH_DATASTORE_EXTRA_BODY_FIELDS: Final = frozenset(VertexSearchDataSto VERTEX_SEARCH_ENGINE_EXTRA_BODY_FIELDS: Final = frozenset(VertexSearchEngineExtraBody.__annotations__) +_H2_AVAILABLE: Final = importlib.util.find_spec("h2") is not None + class VertexSearchSnippet(TypedDict, total=False): snippet: ReadOnly[str] @@ -108,6 +112,11 @@ class VertexSearchAPIVectorStoreConfig(BaseVectorStoreConfig, VertexBase): def __init__(self): super().__init__() + def get_httpx_client_params(self) -> Mapping[str, object]: + if _H2_AVAILABLE: + return MappingProxyType({"http2": True}) + return MappingProxyType({}) + @staticmethod def get_supported_extra_body_fields(is_engine: bool = False) -> frozenset[str]: """ diff --git a/pyproject.toml b/pyproject.toml index 448451f7f93..0a15b517ad8 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -45,6 +45,7 @@ Documentation = "https://docs.litellm.ai" proxy = [ "gunicorn>=23.0.0,<24.0", "uvicorn>=0.33.0,<1.0", + "h2>=4.1.0,<5.0", "granian>=2.7.4,<3.0", "uvloop>=0.22.1,<1.0; sys_platform != 'win32'", "fastapi>=0.136.3,<1.0", diff --git a/tests/test_litellm/llms/custom_httpx/test_http_handler.py b/tests/test_litellm/llms/custom_httpx/test_http_handler.py index f8868cfaf83..6c7c98d2b90 100644 --- a/tests/test_litellm/llms/custom_httpx/test_http_handler.py +++ b/tests/test_litellm/llms/custom_httpx/test_http_handler.py @@ -21,6 +21,7 @@ from litellm.llms.custom_httpx.http_handler import ( HTTPHandler, MaskedHTTPStatusError, _get_httpx_client, + get_async_httpx_client, get_ssl_configuration, ) @@ -305,6 +306,23 @@ def test_get_ssl_configuration(): assert result == mock_ssl_context +def test_get_ssl_configuration_http2_uses_separate_context(): + from litellm.llms.custom_httpx.http_handler import _ssl_context_cache + + _ssl_context_cache.clear() + + default_context = get_ssl_configuration() + http2_context = get_ssl_configuration(http2=True) + + assert default_context is not http2_context + assert get_ssl_configuration() is default_context + assert get_ssl_configuration(http2=True) is http2_context + + custom_context = ssl.create_default_context() + assert get_ssl_configuration(custom_context) is custom_context + assert get_ssl_configuration(custom_context, http2=True) is custom_context + + def test_get_ssl_configuration_integration(): """Integration test that _get_ssl_context() returns a working SSL context""" # Call the static method without mocking @@ -645,6 +663,89 @@ def test_get_httpx_client_applies_httpx_timeout_object_without_mocking_handler() handler.close() +@pytest.mark.asyncio +async def test_async_http_handler_http2_transport(): + monkeypatch = pytest.MonkeyPatch() + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + http2_handler = AsyncHTTPHandler(http2=True) + default_handler = AsyncHTTPHandler() + try: + assert isinstance(http2_handler.client._transport, httpx.AsyncHTTPTransport) + assert http2_handler.client._transport._pool._http2 is True + assert isinstance(default_handler.client._transport, httpx.AsyncHTTPTransport) + assert default_handler.client._transport._pool._http2 is False + finally: + await http2_handler.close() + await default_handler.close() + monkeypatch.undo() + + +@pytest.mark.asyncio +async def test_async_http_handler_http2_request(): + requests: list[httpx.Request] = [] + + async def mock_handler(request: httpx.Request) -> httpx.Response: + requests.append(request) + return httpx.Response(200, request=request, json={"ok": True}) + + handler = AsyncHTTPHandler(http2=True) + await handler.client.aclose() + handler.client = httpx.AsyncClient(transport=httpx.MockTransport(mock_handler)) + try: + response = await handler.get("https://example.com/search") + assert response.json() == {"ok": True} + assert requests[0].url == "https://example.com/search" + finally: + await handler.close() + + +def test_http_handler_http2_transport(): + http2_handler = HTTPHandler(http2=True) + default_handler = HTTPHandler() + try: + assert http2_handler.client._transport._pool._http2 is True + assert default_handler.client._transport._pool._http2 is False + finally: + http2_handler.close() + default_handler.close() + + +def test_http_handler_http2_request(): + requests: list[httpx.Request] = [] + + def mock_handler(request: httpx.Request) -> httpx.Response: + requests.append(request) + return httpx.Response(200, request=request, json={"ok": True}) + + handler = HTTPHandler(http2=True) + handler.client.close() + handler.client = httpx.Client(transport=httpx.MockTransport(mock_handler)) + try: + response = handler.get("https://example.com/search") + assert response.json() == {"ok": True} + assert requests[0].url == "https://example.com/search" + finally: + handler.close() + + +@pytest.mark.asyncio +async def test_get_async_httpx_client_http2_cache_key(): + from litellm.caching.llm_caching_handler import LLMClientCache + from litellm.types.utils import LlmProviders + + monkeypatch = pytest.MonkeyPatch() + monkeypatch.setattr(litellm, "in_memory_llm_clients_cache", LLMClientCache()) + http2_handler = get_async_httpx_client(llm_provider=LlmProviders.VERTEX_AI, params={"http2": True}) + default_handler = get_async_httpx_client(llm_provider=LlmProviders.VERTEX_AI) + try: + assert http2_handler is not default_handler + assert http2_handler.client._transport._pool._http2 is True + finally: + await http2_handler.close() + await default_handler.close() + monkeypatch.undo() + + def test_sync_get_forwards_per_request_timeout(): """HTTPHandler.get(timeout=...) must apply the timeout to that request, overriding the client default rather than silently ignoring it.""" 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 c39779972c0..26861fecd41 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 @@ -2665,6 +2665,40 @@ def test_vector_store_search_handler_direct_config_sync_skips_http(): assert pre_call_args["vector_store_id"] == "vs_direct" +@pytest.mark.asyncio +async def test_async_vector_store_search_handler_passes_provider_httpx_params(): + from litellm.llms.vertex_ai.vector_stores.search_api.transformation import ( + VertexSearchAPIVectorStoreConfig, + ) + + response = httpx.Response(200, json={"results": []}) + client = AsyncMock(spec=AsyncHTTPHandler) + client.post.return_value = response + config = VertexSearchAPIVectorStoreConfig() + logging_obj = Mock(model_call_details={}) + + with ( + patch.object(config, "validate_environment", return_value={}), + patch.object(config, "get_complete_url", return_value="https://discoveryengine.googleapis.com/v1/search"), + patch( # test-quality-ok: verifies vector-store transport configuration reaches the HTTP client factory + "litellm.llms.custom_httpx.llm_http_handler.get_async_httpx_client", + return_value=client, + ) as get_client, + ): + result = await BaseLLMHTTPHandler().async_vector_store_search_handler( + vector_store_id="vs", + query="q", + vector_store_search_optional_params={}, + vector_store_provider_config=config, + custom_llm_provider="vertex_ai", + litellm_params=GenericLiteLLMParams(), + logging_obj=logging_obj, + ) + + assert result["data"] == [] + assert get_client.call_args.kwargs["params"] == {"ssl_verify": None, "http2": True} + + @pytest.mark.asyncio async def test_vector_store_search_handler_direct_config_async_skips_http(): handler = BaseLLMHTTPHandler() diff --git a/tests/test_litellm/llms/vertex_ai/test_vertex_ai_search_vector_store_transformation.py b/tests/test_litellm/llms/vertex_ai/test_vertex_ai_search_vector_store_transformation.py index 034f85f5a0b..02b5b964b68 100644 --- a/tests/test_litellm/llms/vertex_ai/test_vertex_ai_search_vector_store_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/test_vertex_ai_search_vector_store_transformation.py @@ -3,11 +3,17 @@ from types import SimpleNamespace import pytest from litellm.exceptions import BadRequestError +from litellm.llms.base_llm.vector_store.transformation import BaseVectorStoreConfig from litellm.llms.vertex_ai.vector_stores.search_api.transformation import ( VertexSearchAPIVectorStoreConfig, ) +def test_vector_store_httpx_client_params(): + assert VertexSearchAPIVectorStoreConfig().get_httpx_client_params() == {"http2": True} + assert BaseVectorStoreConfig().get_httpx_client_params() == {} + + def test_should_encode_vertex_search_vector_store_id_in_complete_url(): config = VertexSearchAPIVectorStoreConfig() diff --git a/uv.lock b/uv.lock index 88f558221e7..65c37b5c9df 100644 --- a/uv.lock +++ b/uv.lock @@ -4427,6 +4427,7 @@ proxy = [ { name = "fastapi-sso" }, { name = "granian" }, { name = "gunicorn" }, + { name = "h2" }, { name = "hiredis" }, { name = "inquirerpy" }, { name = "litellm-enterprise" }, @@ -4614,6 +4615,7 @@ requires-dist = [ { name = "grpcio", marker = "extra == 'grpc'", specifier = "==1.78.0" }, { name = "grpcio", marker = "extra == 'proxy-runtime'", specifier = "==1.78.0" }, { name = "gunicorn", marker = "extra == 'proxy'", specifier = ">=23.0.0,<24.0" }, + { name = "h2", marker = "extra == 'proxy'", specifier = ">=4.1.0,<5.0" }, { name = "hiredis", marker = "extra == 'proxy'", specifier = ">=3.0.0,<4.0" }, { name = "httpx", specifier = ">=0.28.0,<1.0" }, { name = "importlib-metadata", specifier = ">=8.0.0,<9.0" },