mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-13 08:52:02 +00:00
Merge 40810f2a76 into bc18e10f6b
This commit is contained in:
devin-ai-integration[bot]
GitHub
commit
a865f15a68
12 changed files with 991 additions and 66 deletions
|
|
@ -518,7 +518,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 #######
|
||||
|
|
|
|||
|
|
@ -124,6 +124,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 []
|
||||
|
||||
|
|
|
|||
105
litellm/llms/custom_httpx/http2_connection_spread.py
Normal file
105
litellm/llms/custom_httpx/http2_connection_spread.py
Normal file
|
|
@ -0,0 +1,105 @@
|
|||
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 = release
|
||||
|
||||
async def __aiter__(self) -> AsyncIterator[bytes]:
|
||||
async for chunk in self._stream:
|
||||
yield chunk
|
||||
|
||||
async def aclose(self) -> None:
|
||||
self._release()
|
||||
await self._stream.aclose()
|
||||
|
||||
|
||||
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(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))
|
||||
|
||||
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
|
||||
|
||||
|
|
@ -42,6 +43,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:
|
||||
|
|
@ -72,7 +74,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]]]
|
||||
|
||||
|
|
@ -83,6 +85,31 @@ def http2_enabled() -> bool:
|
|||
return litellm.http2 is True or str_to_bool(os.getenv("LITELLM_HTTP2", "False")) is True
|
||||
|
||||
|
||||
def http2_disabled() -> bool:
|
||||
from litellm.secret_managers.main import str_to_bool
|
||||
|
||||
return litellm.http2 is False or str_to_bool(os.getenv("LITELLM_HTTP2")) is False
|
||||
|
||||
|
||||
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]:
|
||||
|
|
@ -279,9 +306,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(
|
||||
|
|
@ -374,6 +399,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.
|
||||
|
|
@ -422,7 +448,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:
|
||||
|
|
@ -633,6 +659,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
|
||||
|
|
@ -640,6 +667,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,
|
||||
|
|
@ -682,8 +710,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
|
||||
|
|
@ -697,14 +728,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,
|
||||
|
|
@ -712,7 +752,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):
|
||||
|
|
@ -1202,7 +1243,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"
|
||||
|
|
@ -1220,7 +1264,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,
|
||||
|
|
@ -1230,7 +1274,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
|
||||
|
||||
|
|
@ -1377,28 +1426,52 @@ 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
|
||||
if http2 or http2_enabled():
|
||||
return AsyncHTTP2ConnectionSpread(
|
||||
lambda lane_limits: AsyncHTTPTransport(
|
||||
local_address=local_address, verify=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=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
|
||||
|
|
@ -1424,8 +1497,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
|
||||
|
|
@ -1434,17 +1508,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
|
||||
|
|
@ -1719,27 +1795,29 @@ 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=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)
|
||||
)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -25,6 +25,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
|
||||
|
|
@ -342,6 +343,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 _signed_body_kwargs(*, signed_body: bytes | None, data: dict[str, object]) -> _SignedBodyKwargs:
|
||||
if signed_body is not None:
|
||||
return {"data": signed_body}
|
||||
|
|
@ -10458,7 +10487,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
|
||||
|
|
@ -10530,11 +10562,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(
|
||||
|
|
@ -10667,11 +10696,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(
|
||||
|
|
|
|||
|
|
@ -418,12 +418,7 @@ class ContextCachingEndpoints(VertexBase):
|
|||
headers.update(extra_headers)
|
||||
|
||||
if client is None or not isinstance(client, HTTPHandler):
|
||||
_params: Final = {}
|
||||
if timeout is not None:
|
||||
if isinstance(timeout, float) or isinstance(timeout, int):
|
||||
timeout = httpx.Timeout(timeout)
|
||||
_params["timeout"] = timeout
|
||||
client = HTTPHandler(**_params)
|
||||
client = HTTPHandler(timeout=httpx.Timeout(timeout) if isinstance(timeout, float | int) else timeout)
|
||||
else:
|
||||
client = client
|
||||
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -10234,10 +10234,10 @@ class ProviderConfigManager:
|
|||
return VertexVectorStoreConfig()
|
||||
elif 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
|
||||
elif litellm.LlmProviders.BEDROCK == provider:
|
||||
from litellm.llms.bedrock.vector_stores.transformation import (
|
||||
BedrockVectorStoreConfig,
|
||||
|
|
|
|||
|
|
@ -1,30 +1,57 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import contextlib
|
||||
import datetime
|
||||
import ipaddress
|
||||
import json
|
||||
import queue
|
||||
import socket
|
||||
import ssl
|
||||
import threading
|
||||
import time
|
||||
from collections.abc import Iterator
|
||||
from collections.abc import Callable, Generator, Iterator
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from contextlib import contextmanager
|
||||
from dataclasses import dataclass
|
||||
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
|
||||
_HELD_ARRIVALS: Final = queue.SimpleQueue[int]()
|
||||
_TOKEN_REQUESTS: Final = queue.SimpleQueue[str]()
|
||||
_COUNTED_CLIENT_PORTS: Final = queue.SimpleQueue[int]()
|
||||
VERTEX_SERVING_CONFIG_PATH: Final = (
|
||||
"/v1/projects/p/locations/global/collections/default_collection/engines/e/servingConfigs/default_search"
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class TlsPeer:
|
||||
url: str
|
||||
cert_file: Path
|
||||
key_file: Path
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
|
|
@ -76,6 +103,22 @@ 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(client[1] if client else 0)}]})
|
||||
return
|
||||
if scope["path"].endswith("/_oauth/token"):
|
||||
_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"]
|
||||
_COUNTED_CLIENT_PORTS.put(counted_client[1] if counted_client else 0)
|
||||
await _send_json(send, {"results": [{"id": "counted"}]})
|
||||
return
|
||||
if scope["path"].endswith(":search"):
|
||||
await _send_json(send, {"results": [{"id": f"http/{version}"}]})
|
||||
return
|
||||
if scope["path"] == "/stream":
|
||||
await send(
|
||||
HTTPResponseStartEvent(
|
||||
|
|
@ -90,18 +133,26 @@ 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(client_port: int) -> str:
|
||||
_HELD_ARRIVALS.put(client_port)
|
||||
deadline: Final = time.monotonic() + HOLD_SECONDS
|
||||
while _HELD_ARRIVALS.qsize() < HELD_SEARCHES and time.monotonic() < deadline:
|
||||
await asyncio.sleep(0.01)
|
||||
return "held-together" if _HELD_ARRIVALS.qsize() >= HELD_SEARCHES 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]:
|
||||
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()
|
||||
|
|
@ -126,7 +177,7 @@ 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)
|
||||
shutdown.set()
|
||||
thread.join(timeout=10)
|
||||
|
||||
|
|
@ -185,15 +236,437 @@ 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
|
||||
) -> 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,
|
||||
)
|
||||
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_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"]
|
||||
|
||||
|
||||
def test_explicit_http2_off_keeps_vertex_search_on_http1(
|
||||
monkeypatch: pytest.MonkeyPatch, http2_tls_peer: TlsPeer
|
||||
) -> None:
|
||||
import litellm
|
||||
|
||||
_trust_peer_without_global_http2(monkeypatch, http2_tls_peer)
|
||||
monkeypatch.setenv("LITELLM_HTTP2", "False")
|
||||
assert _vertex_search_file_ids(http2_tls_peer) == ["http/1.1"]
|
||||
assert asyncio.run(_avertex_search_file_ids(http2_tls_peer)) == ["http/1.1"]
|
||||
|
||||
monkeypatch.delenv("LITELLM_HTTP2")
|
||||
monkeypatch.setattr(litellm, "http2", False)
|
||||
litellm.in_memory_llm_clients_cache.flush_cache()
|
||||
assert _vertex_search_file_ids(http2_tls_peer) == ["http/1.1"]
|
||||
assert asyncio.run(_avertex_search_file_ids(http2_tls_peer)) == ["http/1.1"]
|
||||
|
||||
|
||||
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(_HELD_ARRIVALS.get_nowait() for _ in range(_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"}
|
||||
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 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"]
|
||||
|
|
|
|||
186
tests/unit/llms/custom_httpx/test_http2_connection_spread.py
Normal file
186
tests/unit/llms/custom_httpx/test_http2_connection_spread.py
Normal file
|
|
@ -0,0 +1,186 @@
|
|||
import asyncio
|
||||
from collections.abc import Callable
|
||||
from typing import Final
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
|
||||
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
|
||||
|
|
@ -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
|
||||
|
|
@ -5542,3 +5544,54 @@ async def test_anthropic_messages_dropped_safeguards_send_neither_the_field_nor_
|
|||
sent: Final = received[0]
|
||||
assert "safeguards" not in json.loads(sent.content)
|
||||
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)
|
||||
|
||||
|
||||
@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
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue