(fix) httpx handler - bind to ipv4 for httpx handler (#6785)

* bind to ipv4 on httpx handler

* add force_ipv4

* use helper for _create_async_transport

* fix circular import

* document force_ipv4

* test_async_http_handler_force_ipv4
This commit is contained in:
Ishaan Jaff 2024-11-18 12:22:51 -08:00 committed by GitHub
parent b854f6c07b
commit f43768d617
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 87 additions and 1 deletions

View file

@ -855,6 +855,7 @@ def test_async_http_handler(mock_async_client):
mock_async_client.assert_called_with(
cert="/client.pem",
transport=None,
event_hooks=event_hooks,
headers=headers,
limits=httpx.Limits(
@ -866,6 +867,52 @@ def test_async_http_handler(mock_async_client):
)
@mock.patch("httpx.AsyncClient")
@mock.patch.dict(os.environ, {}, clear=True)
def test_async_http_handler_force_ipv4(mock_async_client):
"""
Test AsyncHTTPHandler when litellm.force_ipv4 is True
This is prod test - we need to ensure that httpx always uses ipv4 when litellm.force_ipv4 is True
"""
import httpx
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
# Set force_ipv4 to True
litellm.force_ipv4 = True
try:
timeout = 120
event_hooks = {"request": [lambda r: r]}
concurrent_limit = 2
AsyncHTTPHandler(timeout, event_hooks, concurrent_limit)
# Get the call arguments
call_args = mock_async_client.call_args[1]
############# IMPORTANT ASSERTION #################
# Assert transport exists and is configured correctly for using ipv4
assert isinstance(call_args["transport"], httpx.AsyncHTTPTransport)
print(call_args["transport"])
assert call_args["transport"]._pool._local_address == "0.0.0.0"
####################################
# Assert other parameters match
assert call_args["event_hooks"] == event_hooks
assert call_args["headers"] == headers
assert isinstance(call_args["limits"], httpx.Limits)
assert call_args["limits"].max_connections == concurrent_limit
assert call_args["limits"].max_keepalive_connections == concurrent_limit
assert call_args["timeout"] == timeout
assert call_args["verify"] is True
assert call_args["cert"] is None
finally:
# Reset force_ipv4 to default
litellm.force_ipv4 = False
@pytest.mark.parametrize(
"model, expected_bool", [("gpt-3.5-turbo", False), ("gpt-4o-audio-preview", True)]
)