mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-01 07:03:02 +00:00
feat(vertex_ai)!: use HTTP/2 httpx client for search_api vector store (#40631)
* feat(vertex_ai): use HTTP/2 httpx client for search_api vector store Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(httpx): exercise configured request path Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(httpx): verify HTTP/2 request negotiation Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(httpx): preserve HTTP/2 for sync proxy mounts Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * build(deps): add h2 to proxy extra so vertex search HTTP/2 is on by default Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(http_handler): isolate ssl context for http2 clients Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(http): add an HTTP/2 opt-out and harden the Vertex AI Search HTTP/2 client litellm.http2 is now unset by default, so http2: false or LITELLM_HTTP2=False turns the provider's HTTP/2 default off. A caller's own SSLContext keeps the search on HTTP/1.1, the HTTP/2 async client honors DISABLE_AIOHTTP_TRUST_ENV, forced IPv4 transports get the CA bundle and client cert, and vector store search retries once when the server drops the connection with a GOAWAY. * fix(http): retry vector store search only on HTTP/2 GOAWAY and keep aiohttp pool and proxy behavior Scope the one retry to a connection Google closes with an HTTP/2 GOAWAY and log it. When HTTP/2 replaces the aiohttp transport, keep aiohttp's idle keepalive and connection cap, and only then skip env proxies under DISABLE_AIOHTTP_TRUST_ENV, so the global LITELLM_HTTP2 toggle keeps HTTPS_PROXY * test(http): cover the vector store search GOAWAY retry in unit tests and drop duplicate SSL cache comments * fix(http): open another HTTP/2 connection once 100 streams are in use httpcore keeps every request to an origin on one HTTP/2 connection and caps it at 100 concurrent streams, so request 101 and later waited with no timeout. The handlers now spread requests over extra connections, opening one only when every open connection has 100 streams in flight. The force_ipv4 wire tests now cover ssl_verify=False and the client certificate. * fix(vertex_ai): reuse one Vertex search OAuth token across searches Every vector store search built a new search config, so the token cache it carries was thrown away and each search fetched a new Google OAuth token with a blocking call on the event loop. Under a burst of searches that blocking fetch held back every HTTP/2 response. The provider config is now shared, so searches with the same credentials reuse one token until it expires. * fix(http): keep the HTTP/2 connection spread within the connection cap * fix(http): keep sync vertex search on HTTP/1.1, since httpcore's sync HTTP/2 is not thread-safe httpcore 1.0.9 hands two threads the same HTTP/2 stream id, so threaded sync searches failed with protocol errors. Only the async client, which the proxy uses, now defaults to HTTP/2, and the sync connection spread is removed. * fix(http): split the connection cap across HTTP/2 lanes Each lane had its own pool with the full connection cap, so an origin answering HTTP/1.1, or several origins on one client, could exceed max_connections once more than 100 requests were in flight. Lanes now open only while lanes times 100 streams stay under the cap, and each lane pool gets at most 100 of it. * test(http2): scope wire-test peer logs to each fixture instance * fix(http2): keep TE, force_ipv4 trust, evicted-client and opt-out behavior on the HTTP/2 path - drop TE values other than trailers on https requests through the HTTP/2 spread, since h2 refuses them - force_ipv4 transports keep httpx's default trust store (SSL_CERT_DIR) unless SSL is configured - the evicted client closer asks the HTTP/2 spread whether a request is in flight - LITELLM_HTTP2 set to 0, no or off, and a quoted "false" http2 setting, opt out of HTTP/2 - the wire test's held-search barrier waits on an event instead of polling * fix(http2): count proxy-mounted transports when checking an evicted client for in-flight requests * test(http2): move the HTTP/2 eviction cases that open a loopback server into the integration wire tests * fix(http2): read an unquoted http2: 0 as off and free the stream slot of a response dropped without closing * fix(http2): free the stream slot of a partly read response dropped without closing * test(http2): check that the evicted client and its stream close after a partly read response is dropped --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Co-authored-by: mrinal <mrinal@berri.ai> Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
This commit is contained in:
devin-ai-integration[bot]
GitHub
parent
ead221995e
commit
c213234783
12 changed files with 1278 additions and 79 deletions
|
|
@ -520,7 +520,7 @@ aiohttp_trust_env: bool = False # set to true to use HTTP_ Proxy settings
|
|||
disable_aiohttp_transport: bool = False # Set this to true to use httpx instead
|
||||
disable_aiohttp_trust_env: bool = False # When False, aiohttp will respect HTTP(S)_PROXY env vars
|
||||
force_ipv4: bool = False # when True, litellm will force ipv4 for all LLM requests. Some users have seen httpx ConnectionError when using ipv6.
|
||||
http2: bool = False
|
||||
http2: bool | None = None
|
||||
network_mock: bool = False # When True, use mock transport — no real network calls
|
||||
|
||||
####### STOP SEQUENCE LIMIT #######
|
||||
|
|
|
|||
|
|
@ -44,6 +44,8 @@ from collections.abc import Awaitable, Callable, Iterator
|
|||
from dataclasses import dataclass, replace
|
||||
from typing import Final
|
||||
|
||||
import httpx
|
||||
|
||||
from litellm.constants import (
|
||||
EVICTED_LLM_CLIENT_CLOSE_GRACE_SECONDS,
|
||||
EVICTED_LLM_CLIENT_CLOSE_MAX_PENDING,
|
||||
|
|
@ -94,13 +96,24 @@ def _close_function(client: object) -> Callable[[], object] | None:
|
|||
return close_fn
|
||||
|
||||
|
||||
def _transport_of(client: object) -> object:
|
||||
"""The httpx transport behind an SDK wrapper, a litellm handler, or a bare client."""
|
||||
def _mounted_transports(holder: object) -> tuple[object, ...]:
|
||||
mounts: Final[object] = getattr(holder, "_mounts", None)
|
||||
if not isinstance(mounts, dict):
|
||||
return ()
|
||||
return tuple(
|
||||
mounted
|
||||
for mounted in mounts.values() # pyright: ignore[reportUnknownVariableType] # untyped httpx mount map
|
||||
if isinstance(mounted, (httpx.AsyncBaseTransport, httpx.BaseTransport))
|
||||
)
|
||||
|
||||
|
||||
def _transports_of(client: object) -> tuple[object, ...]:
|
||||
"""The httpx transports behind an SDK wrapper, a litellm handler, or a bare client, proxy mounts included."""
|
||||
for holder in (getattr(client, "_client", None), getattr(client, "client", None), client):
|
||||
transport: object = getattr(holder, "_transport", None)
|
||||
if transport is not None:
|
||||
return transport
|
||||
return None
|
||||
return (transport, *_mounted_transports(holder))
|
||||
return ()
|
||||
|
||||
|
||||
def _connection_is_idle(connection: object) -> bool:
|
||||
|
|
@ -123,6 +136,17 @@ def _pool_has_busy_connection(transport: object) -> bool | None:
|
|||
)
|
||||
|
||||
|
||||
def _transport_in_flight(transport: object) -> bool:
|
||||
spread_in_flight: Final[object] = getattr(transport, "has_request_in_flight", None)
|
||||
if callable(spread_in_flight):
|
||||
return bool(spread_in_flight())
|
||||
pooled_busy: Final = _pool_has_busy_connection(transport)
|
||||
if pooled_busy is not None:
|
||||
return pooled_busy
|
||||
session: Final[object] = getattr(transport, "client", None)
|
||||
return bool(getattr(getattr(session, "connector", None), "_acquired", None))
|
||||
|
||||
|
||||
def _has_connection_in_flight(client: object) -> bool:
|
||||
"""Whether the client is servicing a request right now.
|
||||
|
||||
|
|
@ -138,12 +162,7 @@ def _has_connection_in_flight(client: object) -> bool:
|
|||
try:
|
||||
if getattr(getattr(client, "connection_pool", None), "_in_use_connections", None):
|
||||
return True
|
||||
transport: Final = _transport_of(client)
|
||||
pooled_busy: Final = _pool_has_busy_connection(transport)
|
||||
if pooled_busy is not None:
|
||||
return pooled_busy
|
||||
session: Final[object] = getattr(transport, "client", None)
|
||||
return bool(getattr(getattr(session, "connector", None), "_acquired", None))
|
||||
return any(_transport_in_flight(transport) for transport in _transports_of(client))
|
||||
except Exception: # noqa: BLE001 - a client that cannot report its state is treated as idle
|
||||
return False
|
||||
|
||||
|
|
|
|||
|
|
@ -127,6 +127,9 @@ class BaseVectorStoreConfig:
|
|||
def validate_create_vector_store(self) -> None:
|
||||
return None
|
||||
|
||||
def get_async_httpx_client_params(self) -> Mapping[str, object]:
|
||||
return MappingProxyType({})
|
||||
|
||||
def get_supported_openai_params(self, model: str) -> list[VECTOR_STORE_OPENAI_PARAMS]:
|
||||
return []
|
||||
|
||||
|
|
|
|||
128
litellm/llms/custom_httpx/http2_connection_spread.py
Normal file
128
litellm/llms/custom_httpx/http2_connection_spread.py
Normal file
|
|
@ -0,0 +1,128 @@
|
|||
import weakref
|
||||
from collections.abc import AsyncIterator, Callable
|
||||
from typing import Final
|
||||
|
||||
import httpx
|
||||
|
||||
HTTP2_STREAMS_PER_CONNECTION: Final = 100
|
||||
|
||||
|
||||
class _Lane:
|
||||
__slots__ = ("in_flight", "transport")
|
||||
|
||||
def __init__(self, transport: httpx.AsyncBaseTransport) -> None:
|
||||
self.transport: Final = transport
|
||||
self.in_flight = 0
|
||||
|
||||
def acquire(self) -> None:
|
||||
self.in_flight += 1
|
||||
|
||||
def release(self) -> None:
|
||||
self.in_flight -= 1
|
||||
|
||||
|
||||
class _Lanes:
|
||||
def __init__(
|
||||
self,
|
||||
new_transport: Callable[[httpx.Limits], httpx.AsyncBaseTransport],
|
||||
limits: httpx.Limits,
|
||||
streams_per_connection: int,
|
||||
) -> None:
|
||||
self._new_transport: Final = new_transport
|
||||
self._limits: Final = limits
|
||||
self._streams_per_connection: Final = streams_per_connection
|
||||
self.lanes: tuple[_Lane, ...] = ()
|
||||
|
||||
def acquire(self) -> _Lane:
|
||||
free: Final = next((lane for lane in self.lanes if lane.in_flight < self._streams_per_connection), None)
|
||||
lane: Final = free if free is not None else self._open_or_share()
|
||||
lane.acquire()
|
||||
return lane
|
||||
|
||||
def _open_or_share(self) -> _Lane:
|
||||
cap: Final = self._limits.max_connections
|
||||
if cap is not None and len(self.lanes) * self._streams_per_connection >= cap:
|
||||
return min(self.lanes, key=lambda lane: lane.in_flight)
|
||||
return self._open()
|
||||
|
||||
def _open(self) -> _Lane:
|
||||
lane: Final = _Lane(self._new_transport(self._next_lane_limits()))
|
||||
self.lanes = (*self.lanes, lane)
|
||||
return lane
|
||||
|
||||
def _next_lane_limits(self) -> httpx.Limits:
|
||||
cap: Final = self._limits.max_connections
|
||||
if cap is None:
|
||||
return self._limits
|
||||
return httpx.Limits(
|
||||
max_connections=min(self._streams_per_connection, cap - len(self.lanes) * self._streams_per_connection),
|
||||
max_keepalive_connections=self._limits.max_keepalive_connections,
|
||||
keepalive_expiry=self._limits.keepalive_expiry,
|
||||
)
|
||||
|
||||
|
||||
class _AsyncReleasingStream(httpx.AsyncByteStream):
|
||||
def __init__(self, stream: httpx.AsyncByteStream, release: Callable[[], None]) -> None:
|
||||
self._stream: Final = stream
|
||||
self._release: Final = weakref.finalize(self, release)
|
||||
|
||||
async def __aiter__(self) -> AsyncIterator[bytes]:
|
||||
try:
|
||||
async for chunk in self._stream:
|
||||
yield chunk
|
||||
except BaseException:
|
||||
self._release()
|
||||
raise
|
||||
|
||||
async def aclose(self) -> None:
|
||||
self._release()
|
||||
await self._stream.aclose()
|
||||
|
||||
|
||||
def _allowed_over_http2(name: bytes, value: bytes) -> bool:
|
||||
return name.lower() != b"te" or value.strip().lower() == b"trailers"
|
||||
|
||||
|
||||
def _without_te_codings(request: httpx.Request) -> httpx.Request:
|
||||
if request.url.scheme != "https":
|
||||
return request
|
||||
kept: Final = tuple((name, value) for name, value in request.headers.raw if _allowed_over_http2(name, value))
|
||||
if len(kept) == len(request.headers.raw):
|
||||
return request
|
||||
return httpx.Request(
|
||||
request.method, request.url, headers=kept, stream=request.stream, extensions=request.extensions
|
||||
)
|
||||
|
||||
|
||||
def _with_stream(response: httpx.Response, stream: httpx.AsyncByteStream) -> httpx.Response:
|
||||
return httpx.Response(response.status_code, headers=response.headers, stream=stream, extensions=response.extensions)
|
||||
|
||||
|
||||
class AsyncHTTP2ConnectionSpread(httpx.AsyncBaseTransport):
|
||||
def __init__(
|
||||
self,
|
||||
new_transport: Callable[[httpx.Limits], httpx.AsyncBaseTransport],
|
||||
limits: httpx.Limits,
|
||||
streams_per_connection: int = HTTP2_STREAMS_PER_CONNECTION,
|
||||
) -> None:
|
||||
self._lanes: Final = _Lanes(new_transport, limits, streams_per_connection)
|
||||
|
||||
async def handle_async_request(self, request: httpx.Request) -> httpx.Response:
|
||||
lane: Final = self._lanes.acquire()
|
||||
try:
|
||||
response: Final = await lane.transport.handle_async_request(_without_te_codings(request))
|
||||
except BaseException:
|
||||
lane.release()
|
||||
raise
|
||||
stream: Final = response.stream
|
||||
if not isinstance(stream, httpx.AsyncByteStream):
|
||||
lane.release()
|
||||
return response
|
||||
return _with_stream(response, _AsyncReleasingStream(stream, lane.release))
|
||||
|
||||
def has_request_in_flight(self) -> bool:
|
||||
return any(lane.in_flight > 0 for lane in self._lanes.lanes)
|
||||
|
||||
async def aclose(self) -> None:
|
||||
for lane in self._lanes.lanes:
|
||||
await lane.transport.aclose()
|
||||
|
|
@ -18,6 +18,7 @@ import certifi
|
|||
import httpx
|
||||
from aiohttp import ClientSession, DummyCookieJar, TCPConnector
|
||||
from httpx import USE_CLIENT_DEFAULT, AsyncHTTPTransport, HTTPTransport
|
||||
from httpx._config import DEFAULT_LIMITS
|
||||
from httpx._types import CertTypes, RequestFiles
|
||||
from httpx._utils import get_environment_proxies
|
||||
from pydantic import TypeAdapter
|
||||
|
|
@ -43,6 +44,7 @@ from litellm.litellm_core_utils.logging_utils import track_llm_api_timing
|
|||
from litellm.litellm_core_utils.request_timeout_resolver import (
|
||||
get_configured_request_timeout,
|
||||
)
|
||||
from litellm.llms.custom_httpx.http2_connection_spread import AsyncHTTP2ConnectionSpread
|
||||
from litellm.types.llms.custom_http import *
|
||||
|
||||
if TYPE_CHECKING:
|
||||
|
|
@ -73,7 +75,7 @@ _RequestContent: TypeAlias = str | bytes | Iterable[bytes] | AsyncIterable[bytes
|
|||
|
||||
_IPV4_LOCAL_ADDRESS: Final = "0.0.0.0"
|
||||
|
||||
_HttpxTransportT = TypeVar("_HttpxTransportT", HTTPTransport, AsyncHTTPTransport)
|
||||
_HttpxTransportT: Final = TypeVar("_HttpxTransportT", httpx.BaseTransport, httpx.AsyncBaseTransport)
|
||||
|
||||
HTTPEventHooks: TypeAlias = Mapping[str, list[Callable[..., object]]]
|
||||
|
||||
|
|
@ -84,6 +86,36 @@ def http2_enabled() -> bool:
|
|||
return litellm.http2 is True or str_to_bool(os.getenv("LITELLM_HTTP2", "False")) is True
|
||||
|
||||
|
||||
_HTTP2_OFF_VALUES: Final = frozenset({"false", "0", "no", "off"})
|
||||
|
||||
|
||||
def _is_http2_off(value: bool | int | str | None) -> bool:
|
||||
return str(value).strip().lower() in _HTTP2_OFF_VALUES
|
||||
|
||||
|
||||
def http2_disabled() -> bool:
|
||||
return _is_http2_off(litellm.http2) or _is_http2_off(os.getenv("LITELLM_HTTP2"))
|
||||
|
||||
|
||||
def _resolve_http2(provider_default: bool, ssl_verify: VerifyTypes | None) -> bool:
|
||||
if http2_enabled():
|
||||
return True
|
||||
return provider_default and not http2_disabled() and not isinstance(ssl_verify, ssl.SSLContext)
|
||||
|
||||
|
||||
def _aiohttp_env_proxies_disabled() -> bool:
|
||||
from litellm.secret_managers.main import str_to_bool
|
||||
|
||||
return litellm.disable_aiohttp_trust_env or str_to_bool(os.getenv("DISABLE_AIOHTTP_TRUST_ENV", "False")) is True
|
||||
|
||||
|
||||
_AIOHTTP_EQUIVALENT_LIMITS: Final = httpx.Limits(
|
||||
max_connections=AIOHTTP_CONNECTOR_LIMIT if AIOHTTP_CONNECTOR_LIMIT > 0 else None,
|
||||
max_keepalive_connections=None,
|
||||
keepalive_expiry=AIOHTTP_KEEPALIVE_TIMEOUT,
|
||||
)
|
||||
|
||||
|
||||
def _environment_proxy_mounts(
|
||||
build_proxy_transport: Callable[[str], _HttpxTransportT],
|
||||
) -> Mapping[str, _HttpxTransportT | None]:
|
||||
|
|
@ -145,7 +177,7 @@ def build_aiohttp_keepalive_socket_factory() -> Callable[[_AddrInfo], socket.soc
|
|||
_build_aiohttp_keepalive_socket_factory = build_aiohttp_keepalive_socket_factory
|
||||
|
||||
|
||||
def get_default_headers() -> dict:
|
||||
def get_default_headers() -> dict[str, str]:
|
||||
"""
|
||||
Get default headers for HTTP requests.
|
||||
|
||||
|
|
@ -280,9 +312,7 @@ 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)
|
||||
# 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(
|
||||
|
|
@ -375,6 +405,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.
|
||||
|
|
@ -397,6 +428,8 @@ def get_ssl_configuration(
|
|||
- False: Disable SSL verification
|
||||
- True: Enable SSL verification
|
||||
- str: Path to CA bundle file
|
||||
http2: Whether the context is for an HTTP/2 client. httpx sets ALPN on the context it is given,
|
||||
so HTTP/2 clients get their own cached context instead of sharing the HTTP/1.1 one
|
||||
|
||||
Returns:
|
||||
Union[bool, str, ssl.SSLContext]: Appropriate SSL configuration
|
||||
|
|
@ -423,7 +456,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:
|
||||
|
|
@ -439,6 +472,11 @@ def get_ssl_configuration(
|
|||
return ssl_verify
|
||||
|
||||
|
||||
def _httpx_default_verify_unless_configured(verify: VerifyTypes) -> VerifyTypes:
|
||||
litellm_defaults: Final = (get_ssl_configuration(True), get_ssl_configuration(True, http2=True))
|
||||
return True if any(verify is default for default in litellm_defaults) else verify
|
||||
|
||||
|
||||
_shared_realtime_ssl_context: bool | str | ssl.SSLContext | None = None
|
||||
|
||||
|
||||
|
|
@ -662,6 +700,7 @@ class AsyncHTTPHandler:
|
|||
shared_session: Optional["ClientSession"] = None,
|
||||
transport: httpx.AsyncBaseTransport | None = None,
|
||||
follow_redirects: bool = True,
|
||||
http2: bool = False,
|
||||
):
|
||||
self.timeout = timeout
|
||||
self.event_hooks = event_hooks
|
||||
|
|
@ -669,6 +708,7 @@ class AsyncHTTPHandler:
|
|||
self.shared_session = shared_session
|
||||
self.transport = transport
|
||||
self.follow_redirects = follow_redirects
|
||||
self.http2 = http2
|
||||
self._owns_client = True
|
||||
self._client = self.create_client(
|
||||
timeout=timeout,
|
||||
|
|
@ -711,8 +751,11 @@ class AsyncHTTPHandler:
|
|||
follow_redirects=self.follow_redirects,
|
||||
trust_env=False,
|
||||
)
|
||||
use_http2: Final = _resolve_http2(self.http2, ssl_verify)
|
||||
replaces_aiohttp: Final = use_http2 and AsyncHTTPHandler.should_use_aiohttp_transport()
|
||||
limits: Final = _AIOHTTP_EQUIVALENT_LIMITS if replaces_aiohttp else DEFAULT_LIMITS
|
||||
# Get unified SSL configuration
|
||||
ssl_config: Final = get_ssl_configuration(ssl_verify)
|
||||
ssl_config: Final = get_ssl_configuration(ssl_verify, http2=use_http2)
|
||||
|
||||
# An SSL certificate used by the requested host to authenticate the client.
|
||||
# /path/to/client.pem
|
||||
|
|
@ -726,14 +769,23 @@ class AsyncHTTPHandler:
|
|||
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,
|
||||
http2=use_http2,
|
||||
cert=cert,
|
||||
limits=limits,
|
||||
)
|
||||
|
||||
# Get default headers (User-Agent, overridable via LITELLM_USER_AGENT)
|
||||
default_headers: Final = get_default_headers()
|
||||
trust_env: Final = not (replaces_aiohttp and _aiohttp_env_proxies_disabled())
|
||||
|
||||
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=use_http2, limits=limits
|
||||
)
|
||||
if trust_env
|
||||
else None,
|
||||
limits=limits,
|
||||
event_hooks=event_hooks,
|
||||
timeout=timeout,
|
||||
verify=ssl_config,
|
||||
|
|
@ -741,7 +793,8 @@ class AsyncHTTPHandler:
|
|||
headers=default_headers,
|
||||
cookies=blocked_cookie_jar(),
|
||||
follow_redirects=self.follow_redirects,
|
||||
http2=http2_enabled(),
|
||||
http2=use_http2,
|
||||
trust_env=trust_env,
|
||||
)
|
||||
|
||||
async def close(self):
|
||||
|
|
@ -1231,7 +1284,10 @@ class AsyncHTTPHandler:
|
|||
ssl_context: ssl.SSLContext | None = None,
|
||||
ssl_verify: bool | None = None,
|
||||
shared_session: Optional["ClientSession"] = None,
|
||||
) -> LiteLLMAiohttpTransport | AsyncHTTPTransport | None:
|
||||
http2: bool = False,
|
||||
cert: CertTypes | None = None,
|
||||
limits: httpx.Limits = DEFAULT_LIMITS,
|
||||
) -> LiteLLMAiohttpTransport | AsyncHTTPTransport | AsyncHTTP2ConnectionSpread | None:
|
||||
"""
|
||||
- Creates a transport for httpx.AsyncClient
|
||||
- if litellm.force_ipv4 is True, it will return AsyncHTTPTransport with local_address="0.0.0.0"
|
||||
|
|
@ -1249,7 +1305,7 @@ class AsyncHTTPHandler:
|
|||
#########################################################
|
||||
# AIOHTTP TRANSPORT is off by default
|
||||
#########################################################
|
||||
if AsyncHTTPHandler.should_use_aiohttp_transport():
|
||||
if not http2 and AsyncHTTPHandler.should_use_aiohttp_transport():
|
||||
return AsyncHTTPHandler.create_aiohttp_transport(
|
||||
ssl_context=ssl_context,
|
||||
ssl_verify=ssl_verify,
|
||||
|
|
@ -1259,7 +1315,12 @@ class AsyncHTTPHandler:
|
|||
#########################################################
|
||||
# HTTPX TRANSPORT is used when aiohttp is not installed
|
||||
#########################################################
|
||||
return AsyncHTTPHandler._create_httpx_transport()
|
||||
return AsyncHTTPHandler._create_httpx_transport(
|
||||
http2=http2,
|
||||
verify=ssl_context if ssl_context is not None else ssl_verify is not False,
|
||||
cert=cert,
|
||||
limits=limits,
|
||||
)
|
||||
|
||||
_create_async_transport = create_async_transport
|
||||
|
||||
|
|
@ -1406,28 +1467,57 @@ class AsyncHTTPHandler:
|
|||
_create_aiohttp_transport = create_aiohttp_transport
|
||||
|
||||
@staticmethod
|
||||
def _create_httpx_transport() -> AsyncHTTPTransport | None:
|
||||
def _create_httpx_transport(
|
||||
http2: bool = False,
|
||||
verify: ssl.SSLContext | bool = True,
|
||||
cert: CertTypes | None = None,
|
||||
limits: httpx.Limits = DEFAULT_LIMITS,
|
||||
) -> AsyncHTTPTransport | AsyncHTTP2ConnectionSpread | None:
|
||||
"""
|
||||
Creates an AsyncHTTPTransport
|
||||
|
||||
- If HTTP/2 is on, it spreads requests over as many HTTP/2 connections as the stream limit needs
|
||||
- If force_ipv4 is True, it will create an AsyncHTTPTransport with local_address set to "0.0.0.0"
|
||||
- [Default] If force_ipv4 is False, it will return None
|
||||
- [Default] Otherwise it will return None
|
||||
"""
|
||||
local_address: Final = _IPV4_LOCAL_ADDRESS if litellm.force_ipv4 else None
|
||||
transport_verify: Final = (
|
||||
_httpx_default_verify_unless_configured(verify)
|
||||
if litellm.force_ipv4 and not AsyncHTTPHandler.should_use_aiohttp_transport()
|
||||
else verify
|
||||
)
|
||||
if http2 or http2_enabled():
|
||||
return AsyncHTTP2ConnectionSpread(
|
||||
lambda lane_limits: AsyncHTTPTransport(
|
||||
local_address=local_address, verify=transport_verify, cert=cert, http2=True, limits=lane_limits
|
||||
),
|
||||
limits,
|
||||
)
|
||||
if litellm.force_ipv4:
|
||||
return AsyncHTTPTransport(local_address=_IPV4_LOCAL_ADDRESS, http2=http2_enabled())
|
||||
else:
|
||||
return None
|
||||
return AsyncHTTPTransport(local_address=local_address, verify=transport_verify, cert=cert, limits=limits)
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
def create_httpx_proxy_mounts(
|
||||
transport: LiteLLMAiohttpTransport | AsyncHTTPTransport | None,
|
||||
transport: LiteLLMAiohttpTransport | AsyncHTTPTransport | AsyncHTTP2ConnectionSpread | None,
|
||||
verify: VerifyTypes,
|
||||
cert: CertTypes | None,
|
||||
) -> Mapping[str, AsyncHTTPTransport | None] | None:
|
||||
if not isinstance(transport, AsyncHTTPTransport):
|
||||
http2: bool = False,
|
||||
limits: httpx.Limits = DEFAULT_LIMITS,
|
||||
) -> Mapping[str, httpx.AsyncBaseTransport | None] | None:
|
||||
if not isinstance(transport, (AsyncHTTPTransport, AsyncHTTP2ConnectionSpread)):
|
||||
return None
|
||||
if http2 or http2_enabled():
|
||||
return _environment_proxy_mounts(
|
||||
lambda proxy_url: AsyncHTTP2ConnectionSpread(
|
||||
lambda lane_limits: AsyncHTTPTransport(
|
||||
proxy=proxy_url, verify=verify, cert=cert, http2=True, limits=lane_limits
|
||||
),
|
||||
limits,
|
||||
)
|
||||
)
|
||||
return _environment_proxy_mounts(
|
||||
lambda proxy_url: AsyncHTTPTransport(proxy=proxy_url, verify=verify, cert=cert, http2=http2_enabled())
|
||||
lambda proxy_url: AsyncHTTPTransport(proxy=proxy_url, verify=verify, cert=cert, limits=limits)
|
||||
)
|
||||
|
||||
_create_httpx_proxy_mounts = create_httpx_proxy_mounts
|
||||
|
|
@ -1453,8 +1543,9 @@ class HTTPHandler:
|
|||
self._client = self.create_client() if client is None else client
|
||||
|
||||
def create_client(self) -> httpx.Client:
|
||||
use_http2: Final = http2_enabled()
|
||||
# Get unified SSL configuration
|
||||
ssl_config: Final = get_ssl_configuration(self.ssl_verify)
|
||||
ssl_config: Final = get_ssl_configuration(self.ssl_verify, http2=use_http2)
|
||||
|
||||
# An SSL certificate used by the requested host to authenticate the client.
|
||||
# /path/to/client.pem
|
||||
|
|
@ -1463,17 +1554,19 @@ class HTTPHandler:
|
|||
# Get default headers (User-Agent, overridable via LITELLM_USER_AGENT)
|
||||
default_headers: Final = get_default_headers() if not self.disable_default_headers else None
|
||||
|
||||
transport: Final = self._create_sync_transport(verify=ssl_config, cert=cert, http2=use_http2)
|
||||
|
||||
# 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),
|
||||
transport=transport,
|
||||
mounts=self._create_sync_proxy_mounts(transport, verify=ssl_config, cert=cert, http2=use_http2),
|
||||
timeout=self.timeout if self.timeout is not None else _DEFAULT_TIMEOUT,
|
||||
verify=ssl_config,
|
||||
cert=cert,
|
||||
headers=default_headers,
|
||||
cookies=blocked_cookie_jar(),
|
||||
follow_redirects=self.follow_redirects,
|
||||
http2=http2_enabled(),
|
||||
http2=use_http2,
|
||||
)
|
||||
|
||||
@property
|
||||
|
|
@ -1769,27 +1862,34 @@ class HTTPHandler:
|
|||
except Exception:
|
||||
pass
|
||||
|
||||
def _create_sync_transport(self) -> HTTPTransport | None:
|
||||
@staticmethod
|
||||
def _create_sync_transport(verify: VerifyTypes, cert: CertTypes | None, http2: bool) -> HTTPTransport | None:
|
||||
"""
|
||||
Create an HTTP transport with IPv4 only if litellm.force_ipv4 is True.
|
||||
Otherwise, return None.
|
||||
Otherwise, return litellm.sync_transport or None.
|
||||
|
||||
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, http2=http2_enabled())
|
||||
else:
|
||||
return getattr(litellm, "sync_transport", None)
|
||||
return HTTPTransport(
|
||||
local_address=_IPV4_LOCAL_ADDRESS,
|
||||
verify=_httpx_default_verify_unless_configured(verify),
|
||||
cert=cert,
|
||||
http2=http2,
|
||||
)
|
||||
return getattr(litellm, "sync_transport", None)
|
||||
|
||||
@staticmethod
|
||||
def _create_sync_proxy_mounts(
|
||||
transport: httpx.BaseTransport | None,
|
||||
verify: VerifyTypes,
|
||||
cert: CertTypes | None,
|
||||
) -> Mapping[str, HTTPTransport | None] | None:
|
||||
if not litellm.force_ipv4:
|
||||
http2: bool,
|
||||
) -> Mapping[str, httpx.BaseTransport | None] | None:
|
||||
if transport is None or transport is getattr(litellm, "sync_transport", None):
|
||||
return None
|
||||
return _environment_proxy_mounts(
|
||||
lambda proxy_url: HTTPTransport(proxy=proxy_url, verify=verify, cert=cert, http2=http2_enabled())
|
||||
lambda proxy_url: HTTPTransport(proxy=proxy_url, verify=verify, cert=cert, http2=http2)
|
||||
)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -26,6 +26,7 @@ from typing import (
|
|||
from urllib.parse import parse_qs, urlencode, urlparse, urlunparse
|
||||
|
||||
import httpx
|
||||
from h2.events import ConnectionTerminated
|
||||
from httpx import USE_CLIENT_DEFAULT
|
||||
from httpx._types import FileContent
|
||||
from openai.types.file_deleted import FileDeleted
|
||||
|
|
@ -371,6 +372,34 @@ class _SignedBodyKwargs(TypedDict, total=False):
|
|||
json: ReadOnly[dict[str, object]]
|
||||
|
||||
|
||||
_GOAWAY_RETRY_MESSAGE: Final = "Vector store search connection was closed by an HTTP/2 GOAWAY, retrying once"
|
||||
|
||||
|
||||
def _closed_by_http2_goaway(error: httpx.RemoteProtocolError) -> bool:
|
||||
cause: Final = error.__cause__
|
||||
return cause is not None and any(isinstance(arg, ConnectionTerminated) for arg in cause.args)
|
||||
|
||||
|
||||
async def _asend_with_goaway_retry(send: Callable[[], Awaitable[httpx.Response]]) -> httpx.Response:
|
||||
try:
|
||||
return await send()
|
||||
except httpx.RemoteProtocolError as error:
|
||||
if not _closed_by_http2_goaway(error):
|
||||
raise
|
||||
verbose_logger.warning(_GOAWAY_RETRY_MESSAGE)
|
||||
return await send()
|
||||
|
||||
|
||||
def _send_with_goaway_retry(send: Callable[[], httpx.Response]) -> httpx.Response:
|
||||
try:
|
||||
return send()
|
||||
except httpx.RemoteProtocolError as error:
|
||||
if not _closed_by_http2_goaway(error):
|
||||
raise
|
||||
verbose_logger.warning(_GOAWAY_RETRY_MESSAGE)
|
||||
return send()
|
||||
|
||||
|
||||
def _is_sized_json_body(response: httpx.Response) -> bool:
|
||||
headers: Final = httpx.Headers(response.headers)
|
||||
media_types: Final = (value.split(";")[0].strip().lower() for value in headers.get_list("content-type"))
|
||||
|
|
@ -10627,7 +10656,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_async_httpx_client_params(),
|
||||
},
|
||||
)
|
||||
else:
|
||||
async_httpx_client = client
|
||||
|
|
@ -10699,11 +10731,8 @@ class BaseLLMHTTPHandler:
|
|||
request_data: Final = json.dumps(request_body) if signed_json_body is None else signed_json_body
|
||||
|
||||
try:
|
||||
response: Final = await async_httpx_client.post(
|
||||
url=url,
|
||||
headers=headers,
|
||||
data=request_data,
|
||||
timeout=timeout,
|
||||
response: Final = await _asend_with_goaway_retry(
|
||||
partial(async_httpx_client.post, url=url, headers=headers, data=request_data, timeout=timeout)
|
||||
)
|
||||
except httpx.TimeoutException:
|
||||
raise vector_store_provider_config.get_error_class(
|
||||
|
|
@ -10836,11 +10865,8 @@ class BaseLLMHTTPHandler:
|
|||
request_data: Final = json.dumps(request_body) if signed_json_body is None else signed_json_body
|
||||
|
||||
try:
|
||||
response: Final = sync_httpx_client.post(
|
||||
url=url,
|
||||
headers=headers,
|
||||
data=request_data,
|
||||
timeout=timeout,
|
||||
response: Final = _send_with_goaway_retry(
|
||||
partial(sync_httpx_client.post, url=url, headers=headers, data=request_data, timeout=timeout)
|
||||
)
|
||||
except httpx.TimeoutException:
|
||||
raise vector_store_provider_config.get_error_class(
|
||||
|
|
|
|||
|
|
@ -234,6 +234,9 @@ class VertexSearchAPIVectorStoreConfig(BaseVectorStoreConfig, VertexBase):
|
|||
def __init__(self):
|
||||
super().__init__()
|
||||
|
||||
def get_async_httpx_client_params(self) -> Mapping[str, object]:
|
||||
return MappingProxyType({"http2": True})
|
||||
|
||||
@staticmethod
|
||||
def get_supported_extra_body_fields(is_engine: bool = False) -> frozenset[str]:
|
||||
"""
|
||||
|
|
@ -458,3 +461,6 @@ class VertexSearchAPIVectorStoreConfig(BaseVectorStoreConfig, VertexBase):
|
|||
|
||||
input_cost_per_query: Final = model_info.get("input_cost_per_query") or 0.0
|
||||
return input_cost_per_query, 0.0
|
||||
|
||||
|
||||
vertex_search_api_vector_store_config: Final = VertexSearchAPIVectorStoreConfig()
|
||||
|
|
|
|||
|
|
@ -10572,10 +10572,10 @@ class ProviderConfigManager:
|
|||
return VertexVectorStoreConfig()
|
||||
if api_type == "search_api":
|
||||
from litellm.llms.vertex_ai.vector_stores.search_api.transformation import (
|
||||
VertexSearchAPIVectorStoreConfig,
|
||||
vertex_search_api_vector_store_config,
|
||||
)
|
||||
|
||||
return VertexSearchAPIVectorStoreConfig()
|
||||
return vertex_search_api_vector_store_config
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
|
|
|
|||
|
|
@ -1,30 +1,66 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import contextlib
|
||||
import datetime
|
||||
import functools
|
||||
import ipaddress
|
||||
import json
|
||||
import queue
|
||||
import socket
|
||||
import ssl
|
||||
import subprocess
|
||||
import threading
|
||||
import time
|
||||
from collections.abc import Iterator
|
||||
from dataclasses import dataclass
|
||||
from collections.abc import Callable, Generator, Iterator
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from contextlib import contextmanager
|
||||
from dataclasses import dataclass, field
|
||||
from pathlib import Path
|
||||
from types import MappingProxyType
|
||||
from typing import Final
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
from cryptography import x509
|
||||
from cryptography.hazmat.primitives import hashes, serialization
|
||||
from cryptography.hazmat.primitives.asymmetric import rsa
|
||||
from cryptography.x509.oid import NameOID
|
||||
from h2.config import H2Configuration
|
||||
from h2.connection import H2Connection
|
||||
from h2.events import StreamEnded
|
||||
from httpx._config import DEFAULT_LIMITS
|
||||
from hypercorn.asyncio import (
|
||||
serve, # pyright: ignore[reportUnknownVariableType] # hypercorn's serve signature passes through untyped worker hooks
|
||||
)
|
||||
from hypercorn.config import Config
|
||||
from hypercorn.typing import ASGIReceiveCallable, ASGISendCallable, HTTPResponseBodyEvent, HTTPResponseStartEvent, Scope
|
||||
from integration._support.ports import allocate_port
|
||||
from integration._support.vertex import service_account_json
|
||||
|
||||
STREAM_CHUNKS: Final = 3
|
||||
HELD_SEARCHES: Final = 150
|
||||
HOLD_SECONDS: Final = 10.0
|
||||
PROBED_HEADERS: Final = frozenset({b"te", b"keep-alive", b"connection", b"proxy-connection", b"upgrade", b"x-probe"})
|
||||
VERTEX_SERVING_CONFIG_PATH: Final = (
|
||||
"/v1/projects/p/locations/global/collections/default_collection/engines/e/servingConfigs/default_search"
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class PeerLog:
|
||||
held_arrivals: queue.SimpleQueue[int] = field(default_factory=queue.SimpleQueue)
|
||||
all_held_arrived: asyncio.Event = field(default_factory=asyncio.Event)
|
||||
token_requests: queue.SimpleQueue[str] = field(default_factory=queue.SimpleQueue)
|
||||
counted_client_ports: queue.SimpleQueue[int] = field(default_factory=queue.SimpleQueue)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class TlsPeer:
|
||||
url: str
|
||||
cert_file: Path
|
||||
key_file: Path
|
||||
log: PeerLog
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
|
|
@ -66,7 +102,7 @@ def _write_self_signed_cert(cert_dir: Path) -> tuple[Path, Path]:
|
|||
return cert_file, key_file
|
||||
|
||||
|
||||
async def _peer(scope: Scope, receive: ASGIReceiveCallable, send: ASGISendCallable) -> None:
|
||||
async def _peer(log: PeerLog, scope: Scope, receive: ASGIReceiveCallable, send: ASGISendCallable) -> None:
|
||||
if scope["type"] != "http":
|
||||
return
|
||||
while True:
|
||||
|
|
@ -76,6 +112,26 @@ async def _peer(scope: Scope, receive: ASGIReceiveCallable, send: ASGISendCallab
|
|||
if message["type"] == "http.request" and not message["more_body"]:
|
||||
break
|
||||
version: Final = scope["http_version"]
|
||||
if scope["path"] == "/hold:search":
|
||||
client: Final = scope["client"]
|
||||
await _send_json(send, {"results": [{"id": await _hold_until_all_arrive(log, client[1] if client else 0)}]})
|
||||
return
|
||||
if scope["path"].endswith("/_oauth/token"):
|
||||
log.token_requests.put(scope["path"])
|
||||
await _send_json(send, {"access_token": "scripted", "expires_in": 3600, "token_type": "Bearer"})
|
||||
return
|
||||
if scope["path"] == "/counted:search":
|
||||
counted_client: Final = scope["client"]
|
||||
log.counted_client_ports.put(counted_client[1] if counted_client else 0)
|
||||
await _send_json(send, {"results": [{"id": "counted"}]})
|
||||
return
|
||||
if scope["path"].startswith("/headers/") and scope["path"].endswith(":search"):
|
||||
seen: Final = sorted(name.decode() for name, _ in scope["headers"] if name in PROBED_HEADERS)
|
||||
await _send_json(send, {"results": [{"id": f"http/{version}"}, *({"id": name} for name in seen)]})
|
||||
return
|
||||
if scope["path"].endswith(":search"):
|
||||
await _send_json(send, {"results": [{"id": f"http/{version}"}]})
|
||||
return
|
||||
if scope["path"] == "/stream":
|
||||
await send(
|
||||
HTTPResponseStartEvent(
|
||||
|
|
@ -90,21 +146,32 @@ async def _peer(scope: Scope, receive: ASGIReceiveCallable, send: ASGISendCallab
|
|||
)
|
||||
await send(HTTPResponseBodyEvent(type="http.response.body", body=b"", more_body=False))
|
||||
return
|
||||
await _send_json(send, {"http_version": version})
|
||||
|
||||
|
||||
async def _hold_until_all_arrive(log: PeerLog, client_port: int) -> str:
|
||||
log.held_arrivals.put(client_port)
|
||||
if log.held_arrivals.qsize() >= HELD_SEARCHES:
|
||||
log.all_held_arrived.set()
|
||||
waiter: Final = asyncio.ensure_future(log.all_held_arrived.wait())
|
||||
await asyncio.wait((waiter,), timeout=HOLD_SECONDS)
|
||||
waiter.cancel()
|
||||
return "held-together" if log.all_held_arrived.is_set() else "released-alone"
|
||||
|
||||
|
||||
async def _send_json(send: ASGISendCallable, body: dict[str, object]) -> None:
|
||||
await send(
|
||||
HTTPResponseStartEvent(type="http.response.start", status=200, headers=[(b"content-type", b"application/json")])
|
||||
)
|
||||
await send(
|
||||
HTTPResponseBodyEvent(
|
||||
type="http.response.body", body=json.dumps({"http_version": version}).encode(), more_body=False
|
||||
)
|
||||
)
|
||||
await send(HTTPResponseBodyEvent(type="http.response.body", body=json.dumps(body).encode(), more_body=False))
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def http2_tls_peer(tmp_path_factory: pytest.TempPathFactory) -> Iterator[str]:
|
||||
@pytest.fixture
|
||||
def http2_tls_peer(tmp_path_factory: pytest.TempPathFactory) -> Iterator[TlsPeer]:
|
||||
cert_file, key_file = _write_self_signed_cert(tmp_path_factory.mktemp("h2certs"))
|
||||
port: Final = allocate_port()
|
||||
shutdown: Final = threading.Event()
|
||||
log: Final = PeerLog()
|
||||
|
||||
def _serve() -> None:
|
||||
loop: Final = asyncio.new_event_loop()
|
||||
|
|
@ -113,7 +180,9 @@ def http2_tls_peer(tmp_path_factory: pytest.TempPathFactory) -> Iterator[str]:
|
|||
config.certfile = str(cert_file)
|
||||
config.keyfile = str(key_file)
|
||||
config.alpn_protocols = ["h2", "http/1.1"]
|
||||
loop.run_until_complete(serve(_peer, config, shutdown_trigger=lambda: asyncio.to_thread(shutdown.wait)))
|
||||
loop.run_until_complete(
|
||||
serve(functools.partial(_peer, log), config, shutdown_trigger=lambda: asyncio.to_thread(shutdown.wait))
|
||||
)
|
||||
loop.close()
|
||||
|
||||
thread: Final = threading.Thread(target=_serve, daemon=True)
|
||||
|
|
@ -126,16 +195,16 @@ def http2_tls_peer(tmp_path_factory: pytest.TempPathFactory) -> Iterator[str]:
|
|||
time.sleep(0.05)
|
||||
else:
|
||||
pytest.fail("hypercorn peer did not start")
|
||||
yield f"https://127.0.0.1:{port}"
|
||||
yield TlsPeer(url=f"https://127.0.0.1:{port}", cert_file=cert_file, key_file=key_file, log=log)
|
||||
shutdown.set()
|
||||
thread.join(timeout=10)
|
||||
|
||||
|
||||
def _async_exchange(base_url: str) -> Observed:
|
||||
def _async_exchange(base_url: str, ssl_verify: bool | None = False) -> Observed:
|
||||
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
|
||||
|
||||
async def _run() -> Observed:
|
||||
handler: Final = AsyncHTTPHandler(ssl_verify=False)
|
||||
handler: Final = AsyncHTTPHandler(ssl_verify=ssl_verify)
|
||||
try:
|
||||
response: Final = await handler.client.post(f"{base_url}/echo", json={"ping": "pong"})
|
||||
async with handler.client.stream("POST", f"{base_url}/stream", json={}) as stream_response:
|
||||
|
|
@ -151,10 +220,10 @@ def _async_exchange(base_url: str) -> Observed:
|
|||
return asyncio.run(_run())
|
||||
|
||||
|
||||
def _sync_exchange(base_url: str) -> Observed:
|
||||
def _sync_exchange(base_url: str, ssl_verify: bool | None = False) -> Observed:
|
||||
from litellm.llms.custom_httpx.http_handler import HTTPHandler
|
||||
|
||||
handler: Final = HTTPHandler(ssl_verify=False)
|
||||
handler: Final = HTTPHandler(ssl_verify=ssl_verify)
|
||||
try:
|
||||
response: Final = handler.client.post(f"{base_url}/echo", json={"ping": "pong"})
|
||||
with handler.client.stream("POST", f"{base_url}/stream", json={}) as stream_response:
|
||||
|
|
@ -185,15 +254,546 @@ def _assert_negotiated(observed: Observed, enabled: bool) -> None:
|
|||
|
||||
|
||||
@pytest.mark.covers("other.sdk_wire.http2.async_handler_negotiates_h2_only_when_enabled")
|
||||
def test_async_handler_negotiates_http2_only_when_enabled(monkeypatch: pytest.MonkeyPatch, http2_tls_peer: str) -> None:
|
||||
def test_async_handler_negotiates_http2_only_when_enabled(
|
||||
monkeypatch: pytest.MonkeyPatch, http2_tls_peer: TlsPeer
|
||||
) -> None:
|
||||
monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True")
|
||||
for enabled in (False, True):
|
||||
_set_http2(monkeypatch, enabled)
|
||||
_assert_negotiated(_async_exchange(http2_tls_peer), enabled)
|
||||
_assert_negotiated(_async_exchange(http2_tls_peer.url), enabled)
|
||||
|
||||
|
||||
@pytest.mark.covers("other.sdk_wire.http2.sync_handler_negotiates_h2_only_when_enabled")
|
||||
def test_sync_handler_negotiates_http2_only_when_enabled(monkeypatch: pytest.MonkeyPatch, http2_tls_peer: str) -> None:
|
||||
def test_sync_handler_negotiates_http2_only_when_enabled(
|
||||
monkeypatch: pytest.MonkeyPatch, http2_tls_peer: TlsPeer
|
||||
) -> None:
|
||||
for enabled in (False, True):
|
||||
_set_http2(monkeypatch, enabled)
|
||||
_assert_negotiated(_sync_exchange(http2_tls_peer), enabled)
|
||||
_assert_negotiated(_sync_exchange(http2_tls_peer.url), enabled)
|
||||
|
||||
|
||||
def _trust_peer_without_global_http2(monkeypatch: pytest.MonkeyPatch, peer: TlsPeer) -> None:
|
||||
import litellm
|
||||
|
||||
_set_http2(monkeypatch, False)
|
||||
monkeypatch.delenv("DISABLE_AIOHTTP_TRANSPORT", raising=False)
|
||||
monkeypatch.setenv("SSL_VERIFY", str(peer.cert_file))
|
||||
monkeypatch.setenv("REQUESTS_CA_BUNDLE", str(peer.cert_file))
|
||||
litellm.in_memory_llm_clients_cache.flush_cache()
|
||||
|
||||
|
||||
def _vertex_search_file_ids(
|
||||
peer: TlsPeer, search_url: str | None = None, ssl_verify: ssl.SSLContext | None = None
|
||||
) -> list[str | None]:
|
||||
import litellm
|
||||
|
||||
response: Final = litellm.vector_stores.search(
|
||||
vector_store_id="default_search",
|
||||
query="ping",
|
||||
custom_llm_provider="vertex_ai/search_api",
|
||||
api_base=f"{search_url or peer.url}{VERTEX_SERVING_CONFIG_PATH}",
|
||||
vertex_credentials=service_account_json("p", peer.url),
|
||||
vertex_project="p",
|
||||
ssl_verify=ssl_verify,
|
||||
)
|
||||
assert isinstance(response, dict)
|
||||
return [result.get("file_id") for result in response.get("data") or ()]
|
||||
|
||||
|
||||
async def _avertex_search_file_ids(
|
||||
peer: TlsPeer,
|
||||
search_url: str | None = None,
|
||||
ssl_verify: ssl.SSLContext | None = None,
|
||||
extra_headers: dict[str, object] | None = None,
|
||||
) -> list[str | None]:
|
||||
import litellm
|
||||
|
||||
response: Final = await litellm.vector_stores.asearch(
|
||||
vector_store_id="default_search",
|
||||
query="ping",
|
||||
custom_llm_provider="vertex_ai/search_api",
|
||||
api_base=f"{search_url or peer.url}{VERTEX_SERVING_CONFIG_PATH}",
|
||||
vertex_credentials=service_account_json("p", peer.url),
|
||||
vertex_project="p",
|
||||
ssl_verify=ssl_verify,
|
||||
extra_headers=extra_headers,
|
||||
)
|
||||
return [result.get("file_id") for result in response.get("data") or ()]
|
||||
|
||||
|
||||
def test_sync_vertex_search_stays_on_http1_without_global_flag(
|
||||
monkeypatch: pytest.MonkeyPatch, http2_tls_peer: TlsPeer
|
||||
) -> None:
|
||||
_trust_peer_without_global_http2(monkeypatch, http2_tls_peer)
|
||||
assert _vertex_search_file_ids(http2_tls_peer) == ["http/1.1"]
|
||||
|
||||
|
||||
def test_async_vertex_search_negotiates_http2_without_global_flag(
|
||||
monkeypatch: pytest.MonkeyPatch, http2_tls_peer: TlsPeer
|
||||
) -> None:
|
||||
_trust_peer_without_global_http2(monkeypatch, http2_tls_peer)
|
||||
assert asyncio.run(_avertex_search_file_ids(http2_tls_peer)) == ["http/2"]
|
||||
|
||||
|
||||
def test_default_async_client_stays_http1_after_vertex_search_used_http2(
|
||||
monkeypatch: pytest.MonkeyPatch, http2_tls_peer: TlsPeer
|
||||
) -> None:
|
||||
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
|
||||
|
||||
_trust_peer_without_global_http2(monkeypatch, http2_tls_peer)
|
||||
|
||||
async def _run() -> tuple[list[str | None], str]:
|
||||
file_ids: Final = await _avertex_search_file_ids(http2_tls_peer)
|
||||
handler: Final = AsyncHTTPHandler(ssl_verify=str(http2_tls_peer.cert_file))
|
||||
try:
|
||||
response: Final = await handler.post(f"{http2_tls_peer.url}/echo", json={"ping": "pong"})
|
||||
return file_ids, response.json()["http_version"]
|
||||
finally:
|
||||
await handler.close()
|
||||
|
||||
assert asyncio.run(_run()) == (["http/2"], "1.1")
|
||||
|
||||
|
||||
def test_vertex_search_fetches_one_token_for_many_searches(
|
||||
monkeypatch: pytest.MonkeyPatch, http2_tls_peer: TlsPeer
|
||||
) -> None:
|
||||
import litellm
|
||||
|
||||
_trust_peer_without_global_http2(monkeypatch, http2_tls_peer)
|
||||
credentials: Final = service_account_json("p", f"{http2_tls_peer.url}/token-reuse")
|
||||
search: Final = MappingProxyType(
|
||||
{
|
||||
"vector_store_id": "default_search",
|
||||
"query": "ping",
|
||||
"custom_llm_provider": "vertex_ai/search_api",
|
||||
"api_base": f"{http2_tls_peer.url}{VERTEX_SERVING_CONFIG_PATH}",
|
||||
"vertex_credentials": credentials,
|
||||
"vertex_project": "p",
|
||||
}
|
||||
)
|
||||
|
||||
async def _searches() -> None:
|
||||
await litellm.vector_stores.asearch(**search)
|
||||
await asyncio.gather(*(litellm.vector_stores.asearch(**search) for _ in range(5)))
|
||||
|
||||
asyncio.run(_searches())
|
||||
litellm.vector_stores.search(**search)
|
||||
|
||||
token_requests: Final = http2_tls_peer.log.token_requests
|
||||
token_paths: Final = tuple(iter(lambda: None if token_requests.empty() else token_requests.get(), None))
|
||||
assert token_paths.count("/token-reuse/_oauth/token") == 1
|
||||
|
||||
|
||||
def test_vertex_search_over_forced_ipv4_trusts_the_configured_ca(
|
||||
monkeypatch: pytest.MonkeyPatch, http2_tls_peer: TlsPeer
|
||||
) -> None:
|
||||
import litellm
|
||||
|
||||
_trust_peer_without_global_http2(monkeypatch, http2_tls_peer)
|
||||
monkeypatch.setattr(litellm, "force_ipv4", True)
|
||||
|
||||
assert _vertex_search_file_ids(http2_tls_peer) == ["http/1.1"]
|
||||
assert asyncio.run(_avertex_search_file_ids(http2_tls_peer)) == ["http/2"]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("off", ["False", "0", "no", " OFF "])
|
||||
def test_explicit_http2_off_keeps_vertex_search_on_http1(
|
||||
monkeypatch: pytest.MonkeyPatch, http2_tls_peer: TlsPeer, off: str
|
||||
) -> None:
|
||||
_trust_peer_without_global_http2(monkeypatch, http2_tls_peer)
|
||||
monkeypatch.setenv("LITELLM_HTTP2", off)
|
||||
assert asyncio.run(_avertex_search_file_ids(http2_tls_peer)) == ["http/1.1"]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("off", [False, "false", 0])
|
||||
def test_http2_setting_off_keeps_vertex_search_on_http1(
|
||||
monkeypatch: pytest.MonkeyPatch, http2_tls_peer: TlsPeer, off: bool | int | str
|
||||
) -> None:
|
||||
import litellm
|
||||
|
||||
_trust_peer_without_global_http2(monkeypatch, http2_tls_peer)
|
||||
monkeypatch.setattr(litellm, "http2", off)
|
||||
assert asyncio.run(_avertex_search_file_ids(http2_tls_peer)) == ["http/1.1"]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("te", "forwarded"), [("gzip", ["http/2", "x-probe"]), ("trailers", ["http/2", "te", "x-probe"])]
|
||||
)
|
||||
def test_async_vertex_search_over_http2_sends_no_connection_specific_headers(
|
||||
monkeypatch: pytest.MonkeyPatch, http2_tls_peer: TlsPeer, te: str, forwarded: list[str]
|
||||
) -> None:
|
||||
_trust_peer_without_global_http2(monkeypatch, http2_tls_peer)
|
||||
received: Final = asyncio.run(
|
||||
_avertex_search_file_ids(
|
||||
http2_tls_peer,
|
||||
search_url=f"{http2_tls_peer.url}/headers",
|
||||
extra_headers={
|
||||
"TE": te,
|
||||
"Keep-Alive": "timeout=5",
|
||||
"Connection": "keep-alive",
|
||||
"Proxy-Connection": "keep-alive",
|
||||
"Upgrade": "h2c",
|
||||
"X-Probe": "kept",
|
||||
},
|
||||
)
|
||||
)
|
||||
assert received == forwarded
|
||||
|
||||
|
||||
def _alpn_offered_by(context: ssl.SSLContext, peer: TlsPeer) -> str | None:
|
||||
host, port = peer.url.removeprefix("https://").split(":")
|
||||
with (
|
||||
socket.create_connection((host, int(port)), timeout=5) as raw,
|
||||
context.wrap_socket(raw, server_hostname=host) as tls,
|
||||
):
|
||||
return tls.selected_alpn_protocol()
|
||||
|
||||
|
||||
def test_vertex_search_leaves_a_caller_ssl_context_on_http1(
|
||||
monkeypatch: pytest.MonkeyPatch, http2_tls_peer: TlsPeer
|
||||
) -> None:
|
||||
_trust_peer_without_global_http2(monkeypatch, http2_tls_peer)
|
||||
sync_context: Final = ssl.create_default_context(cafile=str(http2_tls_peer.cert_file))
|
||||
async_context: Final = ssl.create_default_context(cafile=str(http2_tls_peer.cert_file))
|
||||
|
||||
assert _vertex_search_file_ids(http2_tls_peer, ssl_verify=sync_context) == ["http/1.1"]
|
||||
assert asyncio.run(_avertex_search_file_ids(http2_tls_peer, ssl_verify=async_context)) == ["http/1.1"]
|
||||
assert _alpn_offered_by(sync_context, http2_tls_peer) != "h2"
|
||||
assert _alpn_offered_by(async_context, http2_tls_peer) != "h2"
|
||||
|
||||
|
||||
def _route_https_through_a_dead_proxy(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
for name in ("https_proxy", "HTTP_PROXY", "http_proxy", "ALL_PROXY", "all_proxy", "NO_PROXY", "no_proxy"):
|
||||
monkeypatch.delenv(name, raising=False)
|
||||
monkeypatch.delenv("DISABLE_AIOHTTP_TRANSPORT", raising=False)
|
||||
monkeypatch.setenv("HTTPS_PROXY", f"http://127.0.0.1:{allocate_port()}")
|
||||
monkeypatch.setenv("DISABLE_AIOHTTP_TRUST_ENV", "True")
|
||||
|
||||
|
||||
async def _async_peer_http_version(peer: TlsPeer, http2: bool) -> str:
|
||||
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
|
||||
|
||||
handler: Final = AsyncHTTPHandler(ssl_verify=str(peer.cert_file), http2=http2)
|
||||
try:
|
||||
response: Final = await handler.client.post(f"{peer.url}/echo", json={"ping": "pong"})
|
||||
return response.http_version
|
||||
finally:
|
||||
await handler.close()
|
||||
|
||||
|
||||
def test_http2_client_skips_env_proxies_when_aiohttp_trust_env_is_disabled(
|
||||
monkeypatch: pytest.MonkeyPatch, http2_tls_peer: TlsPeer
|
||||
) -> None:
|
||||
import litellm
|
||||
|
||||
_set_http2(monkeypatch, False)
|
||||
_route_https_through_a_dead_proxy(monkeypatch)
|
||||
|
||||
for force_ipv4 in (False, True):
|
||||
monkeypatch.setattr(litellm, "force_ipv4", force_ipv4)
|
||||
assert asyncio.run(_async_peer_http_version(http2_tls_peer, http2=True)) == "HTTP/2"
|
||||
|
||||
|
||||
def test_global_http2_toggle_keeps_env_proxies_when_aiohttp_trust_env_is_disabled(
|
||||
monkeypatch: pytest.MonkeyPatch, http2_tls_peer: TlsPeer
|
||||
) -> None:
|
||||
import litellm
|
||||
|
||||
_set_http2(monkeypatch, True)
|
||||
_route_https_through_a_dead_proxy(monkeypatch)
|
||||
|
||||
for force_ipv4 in (False, True):
|
||||
monkeypatch.setattr(litellm, "force_ipv4", force_ipv4)
|
||||
with pytest.raises(httpx.ConnectError):
|
||||
asyncio.run(_async_peer_http_version(http2_tls_peer, http2=False))
|
||||
|
||||
|
||||
def _answer_or_drain(connection: H2Connection, stream_id: int, drain: bool) -> None:
|
||||
if drain:
|
||||
connection.close_connection(last_stream_id=stream_id)
|
||||
return
|
||||
connection.send_headers(stream_id, [(":status", "200"), ("content-type", "application/json")])
|
||||
connection.send_data(stream_id, json.dumps({"results": [{"id": "answered"}]}).encode(), end_stream=True)
|
||||
|
||||
|
||||
def _reply_to_ended_streams(connection: H2Connection, tls: ssl.SSLSocket, data: bytes, drain: bool) -> bool:
|
||||
ended: Final = tuple(event.stream_id for event in connection.receive_data(data) if isinstance(event, StreamEnded))
|
||||
for stream_id in ended:
|
||||
_answer_or_drain(connection, stream_id, drain)
|
||||
tls.sendall(connection.data_to_send())
|
||||
return bool(ended)
|
||||
|
||||
|
||||
def _serve_h2(context: ssl.SSLContext, raw: socket.socket, drain: bool) -> None:
|
||||
with contextlib.suppress(OSError), raw, context.wrap_socket(raw, server_side=True) as tls:
|
||||
connection: Final = H2Connection(config=H2Configuration(client_side=False))
|
||||
connection.initiate_connection()
|
||||
tls.sendall(connection.data_to_send())
|
||||
while data := tls.recv(65535):
|
||||
if _reply_to_ended_streams(connection, tls, data, drain) and drain:
|
||||
return
|
||||
|
||||
|
||||
def _read_http1_request(tls: ssl.SSLSocket) -> bytes:
|
||||
with tls.makefile("rb") as stream:
|
||||
head: Final = tuple(iter(stream.readline, b"\r\n"))
|
||||
length: Final = next(
|
||||
(int(line.split(b":", 1)[1]) for line in head if line.lower().startswith(b"content-length:")), 0
|
||||
)
|
||||
return b"".join(head) + stream.read(length)
|
||||
|
||||
|
||||
def _hang_up_after_request(context: ssl.SSLContext, raw: socket.socket, connection_number: int) -> None:
|
||||
with contextlib.suppress(OSError), raw, context.wrap_socket(raw, server_side=True) as tls:
|
||||
_ = _read_http1_request(tls)
|
||||
|
||||
|
||||
def _drain_every_other_connection(context: ssl.SSLContext, raw: socket.socket, connection_number: int) -> None:
|
||||
_serve_h2(context, raw, drain=connection_number % 2 == 1)
|
||||
|
||||
|
||||
def _answer_every_connection(context: ssl.SSLContext, raw: socket.socket, connection_number: int) -> None:
|
||||
_serve_h2(context, raw, drain=False)
|
||||
|
||||
|
||||
@contextmanager
|
||||
def _raw_tls_peer(
|
||||
peer: TlsPeer,
|
||||
alpn: str,
|
||||
serve_connection: Callable[[ssl.SSLContext, socket.socket, int], None],
|
||||
client_ca: Path | None = None,
|
||||
) -> Generator[tuple[str, queue.SimpleQueue[int]]]:
|
||||
context: Final = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER)
|
||||
context.load_cert_chain(peer.cert_file, peer.key_file)
|
||||
context.set_alpn_protocols([alpn])
|
||||
if client_ca is not None:
|
||||
context.verify_mode = ssl.CERT_REQUIRED
|
||||
context.load_verify_locations(client_ca)
|
||||
stop: Final = threading.Event()
|
||||
accepted: Final = queue.SimpleQueue[int]()
|
||||
|
||||
with socket.create_server(("127.0.0.1", 0)) as listener:
|
||||
listener.settimeout(0.1)
|
||||
|
||||
def _accept_one() -> None:
|
||||
raw: Final = listener.accept()[0]
|
||||
accepted.put(accepted.qsize() + 1)
|
||||
threading.Thread(target=serve_connection, args=(context, raw, accepted.qsize()), daemon=True).start()
|
||||
|
||||
def _accept() -> None:
|
||||
while not stop.is_set():
|
||||
with contextlib.suppress(TimeoutError):
|
||||
_accept_one()
|
||||
|
||||
acceptor: Final = threading.Thread(target=_accept, daemon=True)
|
||||
acceptor.start()
|
||||
try:
|
||||
yield f"https://127.0.0.1:{listener.getsockname()[1]}", accepted
|
||||
finally:
|
||||
stop.set()
|
||||
acceptor.join(timeout=5)
|
||||
|
||||
|
||||
def test_vertex_search_retries_a_connection_the_server_drains_with_goaway(
|
||||
monkeypatch: pytest.MonkeyPatch, http2_tls_peer: TlsPeer
|
||||
) -> None:
|
||||
_trust_peer_without_global_http2(monkeypatch, http2_tls_peer)
|
||||
|
||||
with _raw_tls_peer(http2_tls_peer, "h2", _drain_every_other_connection) as (search_url, _):
|
||||
assert asyncio.run(_avertex_search_file_ids(http2_tls_peer, search_url=search_url)) == ["answered"]
|
||||
_set_http2(monkeypatch, True)
|
||||
assert _vertex_search_file_ids(http2_tls_peer, search_url=search_url) == ["answered"]
|
||||
|
||||
|
||||
def test_vertex_search_sends_once_when_an_http1_server_hangs_up(
|
||||
monkeypatch: pytest.MonkeyPatch, http2_tls_peer: TlsPeer
|
||||
) -> None:
|
||||
import litellm
|
||||
|
||||
_trust_peer_without_global_http2(monkeypatch, http2_tls_peer)
|
||||
monkeypatch.setenv("LITELLM_HTTP2", "False")
|
||||
monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True")
|
||||
|
||||
with _raw_tls_peer(http2_tls_peer, "http/1.1", _hang_up_after_request) as (search_url, accepted):
|
||||
with pytest.raises(litellm.InternalServerError, match="Server disconnected"):
|
||||
_vertex_search_file_ids(http2_tls_peer, search_url=search_url)
|
||||
assert accepted.qsize() == 1
|
||||
with pytest.raises(litellm.InternalServerError, match="Server disconnected"):
|
||||
asyncio.run(_avertex_search_file_ids(http2_tls_peer, search_url=search_url))
|
||||
assert accepted.qsize() == 2
|
||||
|
||||
|
||||
@pytest.mark.parametrize("force_ipv4", [False, True])
|
||||
def test_vertex_search_keeps_its_http2_connection_across_an_idle_gap(
|
||||
monkeypatch: pytest.MonkeyPatch, http2_tls_peer: TlsPeer, force_ipv4: bool
|
||||
) -> None:
|
||||
import litellm
|
||||
from litellm.constants import AIOHTTP_KEEPALIVE_TIMEOUT
|
||||
|
||||
idle_gap: Final = DEFAULT_LIMITS.keepalive_expiry + 1
|
||||
assert idle_gap < AIOHTTP_KEEPALIVE_TIMEOUT
|
||||
_trust_peer_without_global_http2(monkeypatch, http2_tls_peer)
|
||||
monkeypatch.setattr(litellm, "force_ipv4", force_ipv4)
|
||||
|
||||
async def _search_twice(search_url: str) -> list[str | None]:
|
||||
first: Final = await _avertex_search_file_ids(http2_tls_peer, search_url=search_url)
|
||||
await asyncio.sleep(idle_gap)
|
||||
return [*first, *await _avertex_search_file_ids(http2_tls_peer, search_url=search_url)]
|
||||
|
||||
with _raw_tls_peer(http2_tls_peer, "h2", _answer_every_connection) as (search_url, accepted):
|
||||
assert asyncio.run(_search_twice(search_url)) == ["answered", "answered"]
|
||||
assert accepted.qsize() == 1
|
||||
|
||||
|
||||
def test_async_http2_handler_opens_another_connection_once_the_streams_run_out(
|
||||
monkeypatch: pytest.MonkeyPatch, http2_tls_peer: TlsPeer
|
||||
) -> None:
|
||||
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
|
||||
|
||||
_trust_peer_without_global_http2(monkeypatch, http2_tls_peer)
|
||||
|
||||
async def _burst() -> list[str]:
|
||||
handler: Final = AsyncHTTPHandler(ssl_verify=str(http2_tls_peer.cert_file), http2=True)
|
||||
try:
|
||||
responses: Final = await asyncio.gather(
|
||||
*(handler.post(f"{http2_tls_peer.url}/hold:search", json={}) for _ in range(HELD_SEARCHES))
|
||||
)
|
||||
return [response.json()["results"][0]["id"] for response in responses]
|
||||
finally:
|
||||
await handler.close()
|
||||
|
||||
search_ids: Final = asyncio.run(_burst())
|
||||
client_ports: Final = frozenset(
|
||||
http2_tls_peer.log.held_arrivals.get_nowait() for _ in range(http2_tls_peer.log.held_arrivals.qsize())
|
||||
)
|
||||
|
||||
assert search_ids == ["held-together"] * HELD_SEARCHES
|
||||
assert len(client_ports) == 2
|
||||
|
||||
|
||||
def test_http2_transport_opens_no_more_connections_than_its_cap(http2_tls_peer: TlsPeer) -> None:
|
||||
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
|
||||
|
||||
async def _burst() -> frozenset[str]:
|
||||
transport: Final = AsyncHTTPHandler.create_async_transport(
|
||||
ssl_context=ssl.create_default_context(cafile=str(http2_tls_peer.cert_file)),
|
||||
http2=True,
|
||||
limits=httpx.Limits(max_connections=1),
|
||||
)
|
||||
async with httpx.AsyncClient(transport=transport) as client:
|
||||
responses: Final = await asyncio.gather(
|
||||
*(client.post(f"{http2_tls_peer.url}/counted:search", json={}) for _ in range(HELD_SEARCHES))
|
||||
)
|
||||
return frozenset(response.json()["results"][0]["id"] for response in responses)
|
||||
|
||||
assert asyncio.run(_burst()) == {"counted"}
|
||||
counted_client_ports: Final = http2_tls_peer.log.counted_client_ports
|
||||
client_ports: Final = frozenset(
|
||||
iter(lambda: None if counted_client_ports.empty() else counted_client_ports.get(), None)
|
||||
)
|
||||
assert len(client_ports) == 1
|
||||
|
||||
|
||||
def test_forced_ipv4_handlers_honour_ssl_verify_false(monkeypatch: pytest.MonkeyPatch, http2_tls_peer: TlsPeer) -> None:
|
||||
import litellm
|
||||
|
||||
monkeypatch.setattr(litellm, "force_ipv4", True)
|
||||
monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True")
|
||||
for enabled in (False, True):
|
||||
_set_http2(monkeypatch, enabled)
|
||||
_assert_negotiated(_async_exchange(http2_tls_peer.url), enabled)
|
||||
_assert_negotiated(_sync_exchange(http2_tls_peer.url), enabled)
|
||||
|
||||
|
||||
def _ca_dir_holding(cert_file: Path, ca_dir: Path) -> Path:
|
||||
subject_hash: Final = subprocess.run(
|
||||
["openssl", "x509", "-hash", "-noout", "-in", str(cert_file)], capture_output=True, text=True, check=True
|
||||
).stdout.strip()
|
||||
(ca_dir / f"{subject_hash}.0").write_bytes(cert_file.read_bytes())
|
||||
return ca_dir
|
||||
|
||||
|
||||
def test_forced_ipv4_handlers_trust_a_ca_found_only_in_ssl_cert_dir(
|
||||
monkeypatch: pytest.MonkeyPatch, http2_tls_peer: TlsPeer, tmp_path: Path
|
||||
) -> None:
|
||||
import litellm
|
||||
|
||||
monkeypatch.setattr(litellm, "force_ipv4", True)
|
||||
monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True")
|
||||
monkeypatch.delenv("SSL_VERIFY", raising=False)
|
||||
monkeypatch.delenv("SSL_CERT_FILE", raising=False)
|
||||
monkeypatch.setenv("SSL_CERT_DIR", str(_ca_dir_holding(http2_tls_peer.cert_file, tmp_path)))
|
||||
for enabled in (False, True):
|
||||
_set_http2(monkeypatch, enabled)
|
||||
_assert_negotiated(_async_exchange(http2_tls_peer.url, ssl_verify=None), enabled)
|
||||
_assert_negotiated(_sync_exchange(http2_tls_peer.url, ssl_verify=None), enabled)
|
||||
|
||||
|
||||
def test_vertex_search_over_forced_ipv4_presents_the_client_certificate(
|
||||
monkeypatch: pytest.MonkeyPatch, http2_tls_peer: TlsPeer, tmp_path: Path
|
||||
) -> None:
|
||||
import litellm
|
||||
|
||||
_trust_peer_without_global_http2(monkeypatch, http2_tls_peer)
|
||||
monkeypatch.setattr(litellm, "force_ipv4", True)
|
||||
client_pem: Final = tmp_path / "client.pem"
|
||||
client_pem.write_bytes(http2_tls_peer.cert_file.read_bytes() + http2_tls_peer.key_file.read_bytes())
|
||||
monkeypatch.setenv("SSL_CERTIFICATE", str(client_pem))
|
||||
|
||||
with _raw_tls_peer(http2_tls_peer, "h2", _answer_every_connection, client_ca=http2_tls_peer.cert_file) as (
|
||||
search_url,
|
||||
_,
|
||||
):
|
||||
assert asyncio.run(_avertex_search_file_ids(http2_tls_peer, search_url=search_url)) == ["answered"]
|
||||
_set_http2(monkeypatch, True)
|
||||
assert _vertex_search_file_ids(http2_tls_peer, search_url=search_url) == ["answered"]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("proxied", [False, True], ids=["direct", "behind-http-proxy"])
|
||||
def test_an_evicted_http2_handler_is_not_closed_mid_request(monkeypatch: pytest.MonkeyPatch, proxied: bool) -> None:
|
||||
from litellm.caching.evicted_client_closer import EvictedClientCloser
|
||||
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
|
||||
|
||||
for name in ("http_proxy", "HTTP_PROXY", "all_proxy", "ALL_PROXY", "no_proxy", "NO_PROXY"):
|
||||
monkeypatch.delenv(name, raising=False)
|
||||
|
||||
async def _evict_mid_request() -> None:
|
||||
arrived: Final = asyncio.Event()
|
||||
release: Final = asyncio.Event()
|
||||
|
||||
async def _upstream(reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None:
|
||||
await reader.read(4096)
|
||||
writer.write(b"HTTP/1.1 200 OK\r\nTransfer-Encoding: chunked\r\n\r\n5\r\nhello\r\n")
|
||||
await writer.drain()
|
||||
arrived.set()
|
||||
await release.wait()
|
||||
writer.write(b"0\r\n\r\n")
|
||||
await writer.drain()
|
||||
|
||||
server: Final = await asyncio.start_server(_upstream, "127.0.0.1", 0)
|
||||
port: Final[int] = server.sockets[0].getsockname()[1]
|
||||
if proxied:
|
||||
monkeypatch.setenv("HTTP_PROXY", f"http://127.0.0.1:{port}")
|
||||
closer: Final = EvictedClientCloser(grace_seconds=0.0)
|
||||
handler: Final = AsyncHTTPHandler(http2=True)
|
||||
evicted: Final = handler.client
|
||||
closer.mark_owned(handler)
|
||||
closer.schedule(handler)
|
||||
request: Final = asyncio.create_task(
|
||||
handler.get("http://upstream.invalid/" if proxied else f"http://127.0.0.1:{port}/")
|
||||
)
|
||||
|
||||
await arrived.wait()
|
||||
closer.reap()
|
||||
await asyncio.sleep(0)
|
||||
assert evicted.is_closed is False, "closed a handler that was serving a request"
|
||||
|
||||
release.set()
|
||||
assert (await request).text == "hello"
|
||||
closer.reap()
|
||||
await asyncio.sleep(0)
|
||||
assert evicted.is_closed is True
|
||||
assert handler.client.is_closed is False, "a held handler must self-heal after its evicted client is closed"
|
||||
server.close()
|
||||
|
||||
asyncio.run(_evict_mid_request())
|
||||
|
|
|
|||
264
tests/unit/llms/custom_httpx/test_http2_connection_spread.py
Normal file
264
tests/unit/llms/custom_httpx/test_http2_connection_spread.py
Normal file
|
|
@ -0,0 +1,264 @@
|
|||
import asyncio
|
||||
import gc
|
||||
from collections.abc import AsyncIterator, Callable, Iterator
|
||||
from typing import Final
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
|
||||
from litellm.caching.evicted_client_closer import EvictedClientCloser
|
||||
from litellm.llms.custom_httpx.http2_connection_spread import (
|
||||
HTTP2_STREAMS_PER_CONNECTION,
|
||||
AsyncHTTP2ConnectionSpread,
|
||||
)
|
||||
|
||||
URL: Final = "https://search.example/v1/search"
|
||||
BURST: Final = HTTP2_STREAMS_PER_CONNECTION + HTTP2_STREAMS_PER_CONNECTION // 2
|
||||
UNCAPPED: Final = httpx.Limits(max_connections=None)
|
||||
|
||||
|
||||
class _GatedAsyncConnection(httpx.AsyncBaseTransport):
|
||||
def __init__(self, gate: asyncio.Event) -> None:
|
||||
self._gate: Final = gate
|
||||
self.in_flight = 0
|
||||
self.peak = 0
|
||||
self.served = 0
|
||||
self.closed = False
|
||||
|
||||
async def handle_async_request(self, request: httpx.Request) -> httpx.Response:
|
||||
self.in_flight += 1
|
||||
self.peak = max(self.peak, self.in_flight)
|
||||
await self._gate.wait()
|
||||
self.in_flight -= 1
|
||||
self.served += 1
|
||||
return httpx.Response(200, stream=httpx.ByteStream(b"ok"))
|
||||
|
||||
async def aclose(self) -> None:
|
||||
self.closed = True
|
||||
|
||||
|
||||
class _AsyncConnections:
|
||||
def __init__(self, gate: asyncio.Event) -> None:
|
||||
self._gate: Final = gate
|
||||
self.opened: tuple[_GatedAsyncConnection, ...] = ()
|
||||
|
||||
def __call__(self, _limits: httpx.Limits) -> _GatedAsyncConnection:
|
||||
connection: Final = _GatedAsyncConnection(self._gate)
|
||||
self.opened = (*self.opened, connection)
|
||||
return connection
|
||||
|
||||
def in_flight(self) -> int:
|
||||
return sum(connection.in_flight for connection in self.opened)
|
||||
|
||||
|
||||
class _RefusedConnection(httpx.AsyncBaseTransport):
|
||||
async def handle_async_request(self, request: httpx.Request) -> httpx.Response:
|
||||
raise httpx.ConnectError("refused", request=request)
|
||||
|
||||
|
||||
class _RefusedConnections:
|
||||
def __init__(self) -> None:
|
||||
self.opened = 0
|
||||
|
||||
def __call__(self, _limits: httpx.Limits) -> _RefusedConnection:
|
||||
self.opened += 1
|
||||
return _RefusedConnection()
|
||||
|
||||
|
||||
class _HTTP1Connections:
|
||||
def __init__(self, gate: asyncio.Event) -> None:
|
||||
self._gate: Final = gate
|
||||
self.requested = 0
|
||||
self.open = 0
|
||||
|
||||
def __call__(self, limits: httpx.Limits) -> "_HTTP1Pool":
|
||||
assert limits.max_connections is not None
|
||||
return _HTTP1Pool(self, asyncio.Semaphore(limits.max_connections))
|
||||
|
||||
async def serve(self, pool_slots: asyncio.Semaphore) -> None:
|
||||
self.requested += 1
|
||||
async with pool_slots:
|
||||
self.open += 1
|
||||
await self._gate.wait()
|
||||
self.open -= 1
|
||||
|
||||
|
||||
class _HTTP1Pool(httpx.AsyncBaseTransport):
|
||||
def __init__(self, connections: _HTTP1Connections, slots: asyncio.Semaphore) -> None:
|
||||
self._connections: Final = connections
|
||||
self._slots: Final = slots
|
||||
|
||||
async def handle_async_request(self, request: httpx.Request) -> httpx.Response:
|
||||
await self._connections.serve(self._slots)
|
||||
return httpx.Response(200, stream=httpx.ByteStream(b"ok"))
|
||||
|
||||
|
||||
async def _until(condition: Callable[[], bool]) -> None:
|
||||
for _ in range(1000):
|
||||
if condition():
|
||||
return
|
||||
await asyncio.sleep(0)
|
||||
pytest.fail("condition never held")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_spread_opens_another_connection_only_once_streams_run_out():
|
||||
gate: Final = asyncio.Event()
|
||||
connections: Final = _AsyncConnections(gate)
|
||||
async with httpx.AsyncClient(transport=AsyncHTTP2ConnectionSpread(connections, UNCAPPED)) as client:
|
||||
pending: Final = tuple(asyncio.create_task(client.get(URL)) for _ in range(BURST))
|
||||
await _until(lambda: connections.in_flight() == BURST)
|
||||
assert [connection.peak for connection in connections.opened] == [
|
||||
HTTP2_STREAMS_PER_CONNECTION,
|
||||
BURST - HTTP2_STREAMS_PER_CONNECTION,
|
||||
]
|
||||
|
||||
gate.set()
|
||||
assert {response.status_code for response in await asyncio.gather(*pending)} == {200}
|
||||
assert (await client.get(URL)).text == "ok"
|
||||
assert [connection.served for connection in connections.opened] == [
|
||||
HTTP2_STREAMS_PER_CONNECTION + 1,
|
||||
BURST - HTTP2_STREAMS_PER_CONNECTION,
|
||||
]
|
||||
assert all(connection.closed for connection in connections.opened)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_spread_shares_the_least_busy_connection_once_the_connection_cap_is_reached():
|
||||
gate: Final = asyncio.Event()
|
||||
connections: Final = _AsyncConnections(gate)
|
||||
burst: Final = 2 * HTTP2_STREAMS_PER_CONNECTION + HTTP2_STREAMS_PER_CONNECTION // 2
|
||||
two_connections_of_streams: Final = httpx.Limits(max_connections=2 * HTTP2_STREAMS_PER_CONNECTION)
|
||||
async with httpx.AsyncClient(
|
||||
transport=AsyncHTTP2ConnectionSpread(connections, two_connections_of_streams)
|
||||
) as client:
|
||||
pending: Final = tuple(asyncio.create_task(client.get(URL)) for _ in range(burst))
|
||||
await _until(lambda: connections.in_flight() == burst)
|
||||
assert [connection.peak for connection in connections.opened] == [burst // 2, burst // 2]
|
||||
|
||||
gate.set()
|
||||
assert {response.status_code for response in await asyncio.gather(*pending)} == {200}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_spread_keeps_http1_connections_within_the_connection_cap():
|
||||
gate: Final = asyncio.Event()
|
||||
connections: Final = _HTTP1Connections(gate)
|
||||
cap: Final = HTTP2_STREAMS_PER_CONNECTION + HTTP2_STREAMS_PER_CONNECTION // 2
|
||||
burst: Final = 2 * cap
|
||||
async with httpx.AsyncClient(
|
||||
transport=AsyncHTTP2ConnectionSpread(connections, httpx.Limits(max_connections=cap))
|
||||
) as client:
|
||||
pending: Final = tuple(asyncio.create_task(client.get(URL)) for _ in range(burst))
|
||||
await _until(lambda: connections.requested == burst)
|
||||
assert connections.open == cap
|
||||
|
||||
gate.set()
|
||||
assert {response.status_code for response in await asyncio.gather(*pending)} == {200}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_spread_holds_a_stream_until_its_response_is_closed():
|
||||
gate: Final = asyncio.Event()
|
||||
gate.set()
|
||||
connections: Final = _AsyncConnections(gate)
|
||||
async with httpx.AsyncClient(transport=AsyncHTTP2ConnectionSpread(connections, UNCAPPED)) as client:
|
||||
open_responses: Final = tuple(
|
||||
[
|
||||
await client.send(client.build_request("GET", URL), stream=True)
|
||||
for _ in range(HTTP2_STREAMS_PER_CONNECTION)
|
||||
]
|
||||
)
|
||||
await client.get(URL)
|
||||
assert len(connections.opened) == 2
|
||||
|
||||
for response in open_responses:
|
||||
await response.aclose()
|
||||
await client.get(URL)
|
||||
assert [connection.served for connection in connections.opened] == [HTTP2_STREAMS_PER_CONNECTION + 1, 1]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_spread_frees_the_stream_when_the_request_fails():
|
||||
connections: Final = _RefusedConnections()
|
||||
async with httpx.AsyncClient(transport=AsyncHTTP2ConnectionSpread(connections, UNCAPPED)) as client:
|
||||
for _ in range(BURST):
|
||||
with pytest.raises(httpx.ConnectError):
|
||||
await client.get(URL)
|
||||
assert connections.opened == 1
|
||||
|
||||
|
||||
async def _open_streams_without_closing(client: httpx.AsyncClient, count: int) -> int:
|
||||
responses: Final = tuple([await client.send(client.build_request("GET", URL), stream=True) for _ in range(count)])
|
||||
return sum(not response.is_closed for response in responses)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_spread_frees_the_streams_of_responses_dropped_without_closing():
|
||||
gate: Final = asyncio.Event()
|
||||
gate.set()
|
||||
connections: Final = _AsyncConnections(gate)
|
||||
spread: Final = AsyncHTTP2ConnectionSpread(connections, UNCAPPED)
|
||||
async with httpx.AsyncClient(transport=spread) as client:
|
||||
assert await _open_streams_without_closing(client, HTTP2_STREAMS_PER_CONNECTION) == HTTP2_STREAMS_PER_CONNECTION
|
||||
gc.collect()
|
||||
|
||||
assert spread.has_request_in_flight() is False
|
||||
assert (await client.get(URL)).text == "ok"
|
||||
assert len(connections.opened) == 1
|
||||
|
||||
|
||||
class _TwoChunks(httpx.AsyncByteStream):
|
||||
def __init__(self) -> None:
|
||||
self.closed = False
|
||||
|
||||
async def __aiter__(self) -> AsyncIterator[bytes]:
|
||||
try:
|
||||
yield b"o"
|
||||
yield b"k"
|
||||
finally:
|
||||
self.closed = True
|
||||
|
||||
|
||||
class _TwoChunkConnection(httpx.AsyncBaseTransport):
|
||||
def __init__(self) -> None:
|
||||
self.stream: Final = _TwoChunks()
|
||||
self.closed = False
|
||||
|
||||
async def handle_async_request(self, request: httpx.Request) -> httpx.Response:
|
||||
return httpx.Response(200, stream=self.stream)
|
||||
|
||||
async def aclose(self) -> None:
|
||||
self.closed = True
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def cyclic_gc_off() -> Iterator[None]:
|
||||
gc.disable()
|
||||
yield
|
||||
gc.enable()
|
||||
|
||||
|
||||
async def _read_one_chunk_and_drop(client: httpx.AsyncClient) -> bytes:
|
||||
response: Final = await client.send(client.build_request("GET", URL), stream=True)
|
||||
return await anext(response.aiter_raw())
|
||||
|
||||
|
||||
def _reaped(closer: EvictedClientCloser, client: httpx.AsyncClient) -> bool:
|
||||
closer.reap()
|
||||
return client.is_closed
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.usefixtures("cyclic_gc_off")
|
||||
async def test_an_evicted_client_closes_once_a_partly_read_response_is_dropped_without_closing():
|
||||
connection: Final = _TwoChunkConnection()
|
||||
client: Final = httpx.AsyncClient(transport=AsyncHTTP2ConnectionSpread(lambda _limits: connection, UNCAPPED))
|
||||
closer: Final = EvictedClientCloser(grace_seconds=0.0)
|
||||
closer.mark_owned(client)
|
||||
closer.schedule(client)
|
||||
|
||||
assert await _read_one_chunk_and_drop(client) == b"o"
|
||||
|
||||
await _until(lambda: _reaped(closer, client))
|
||||
await _until(lambda: connection.closed and connection.stream.closed)
|
||||
|
|
@ -1917,12 +1917,12 @@ async def test_http2_flag_bypasses_aiohttp_transport(monkeypatch: pytest.MonkeyP
|
|||
|
||||
monkeypatch.setattr(litellm, "http2", True)
|
||||
assert AsyncHTTPHandler.should_use_aiohttp_transport() is False
|
||||
assert AsyncHTTPHandler.create_async_transport() is None
|
||||
assert not isinstance(AsyncHTTPHandler.create_async_transport(), LiteLLMAiohttpTransport)
|
||||
|
||||
monkeypatch.setattr(litellm, "http2", False)
|
||||
monkeypatch.setenv("LITELLM_HTTP2", "True")
|
||||
assert AsyncHTTPHandler.should_use_aiohttp_transport() is False
|
||||
assert AsyncHTTPHandler.create_async_transport() is None
|
||||
assert not isinstance(AsyncHTTPHandler.create_async_transport(), LiteLLMAiohttpTransport)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
|
|||
|
|
@ -10,10 +10,12 @@ from types import MappingProxyType
|
|||
from typing import Final
|
||||
from unittest.mock import AsyncMock, Mock, patch
|
||||
|
||||
import httpcore
|
||||
import httpx
|
||||
import pytest
|
||||
import respx
|
||||
from botocore.credentials import RefreshableCredentials
|
||||
from h2.events import ConnectionTerminated
|
||||
|
||||
import litellm
|
||||
from litellm._logging import verbose_logger
|
||||
|
|
@ -5788,6 +5790,39 @@ async def test_anthropic_messages_dropped_safeguards_send_neither_the_field_nor_
|
|||
assert "dangerous-tool-use-2026-09-03" not in sent.headers.get("anthropic-beta", ""), dict(sent.headers)
|
||||
|
||||
|
||||
_EMPTY_SEARCH_PAGE: Final = {
|
||||
"object": "vector_store.search_results.page",
|
||||
"search_query": "q",
|
||||
"data": [],
|
||||
"has_more": False,
|
||||
"next_page": None,
|
||||
}
|
||||
|
||||
|
||||
def _close_with_goaway(request: httpx.Request) -> httpx.Response:
|
||||
raise httpx.RemoteProtocolError("GOAWAY", request=request) from httpcore.RemoteProtocolError(ConnectionTerminated())
|
||||
|
||||
|
||||
def _hang_up(request: httpx.Request) -> httpx.Response:
|
||||
raise httpx.RemoteProtocolError("Server disconnected without sending a response.", request=request)
|
||||
|
||||
|
||||
def _answer_search(request: httpx.Request) -> httpx.Response:
|
||||
return httpx.Response(200, json=_EMPTY_SEARCH_PAGE)
|
||||
|
||||
|
||||
async def _search_through(transport: httpx.MockTransport, is_async: bool) -> object:
|
||||
search_kwargs: Final = {"vector_store_id": "vs_1", "query": "q", "custom_llm_provider": "openai", "api_key": "sk-test"}
|
||||
if not is_async:
|
||||
with httpx.Client(transport=transport) as sync_client:
|
||||
return litellm.vector_stores.search(**search_kwargs, client=HTTPHandler(client=sync_client))
|
||||
async with httpx.AsyncClient(transport=transport) as async_client:
|
||||
client: Final = AsyncHTTPHandler()
|
||||
await client.close()
|
||||
client.client = async_client
|
||||
return await litellm.vector_stores.asearch(**search_kwargs, client=client)
|
||||
|
||||
|
||||
class TestAgenticFollowUpKeepsTheProviderPrefix:
|
||||
@staticmethod
|
||||
async def _run_followup(
|
||||
|
|
@ -5956,6 +5991,24 @@ def test_vector_store_retrieve_handler_sends_id_before_azure_api_version_query()
|
|||
assert len(received) == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("is_async", (False, True))
|
||||
async def test_vector_store_search_sends_again_once_after_an_http2_goaway(is_async: bool) -> None:
|
||||
sends: Final = iter((_close_with_goaway, _answer_search))
|
||||
response: Final = await _search_through(httpx.MockTransport(lambda request: next(sends)(request)), is_async)
|
||||
assert response["search_query"] == "q"
|
||||
assert next(sends, None) is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("is_async", (False, True))
|
||||
async def test_vector_store_search_sends_once_when_the_connection_drops_without_a_goaway(is_async: bool) -> None:
|
||||
sends: Final = iter((_hang_up, _answer_search))
|
||||
with pytest.raises(litellm.InternalServerError, match="Server disconnected"):
|
||||
await _search_through(httpx.MockTransport(lambda request: next(sends)(request)), is_async)
|
||||
assert next(sends) is _answer_search
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("is_async", (False, True))
|
||||
@pytest.mark.parametrize("provider", ("openai", "azure"))
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue