Merge 98998c207d into 1c61c2606e

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 []
 

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(

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
 

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]:
         """

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",

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."""

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()

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()
 

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" },