scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-10 22:41:41 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 82

Fix: handle closed async httpx clients in OpenAI client cache

Browse source
This commit is contained in:
Stoyan Yanev 2026-02-20 15:15:00 +02:00
parent 6134984008
commit 35ce975b1d
2 changed files with 313 additions and 81 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

188
litellm/llms/openai/openai.py
View file

@ -350,6 +350,13 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM):
if max_retries is not None:
client.max_retries = max_retries
def _httpx_is_closed(self, openai_client) -> bool:
"""Return True if underlying httpx client is missing or closed."""
httpx_client = getattr(openai_client, "_client", None)
if httpx_client is None:
return True
return bool(getattr(httpx_client, "is_closed", True))
def _get_openai_client(
self,
is_async: bool,
@ -362,55 +369,8 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM):
client: Optional[Union[OpenAI, AsyncOpenAI]] = None,
shared_session: Optional["ClientSession"] = None,
) -> Optional[Union[OpenAI, AsyncOpenAI]]:
client_initialization_params: Dict = locals()
if client is None:
if not isinstance(max_retries, int):
raise OpenAIError(
status_code=422,
message="max retries must be an int. Passed in value: {}".format(
max_retries
),
)
cached_client = self.get_cached_openai_client(
client_initialization_params=client_initialization_params,
client_type="openai",
)
if cached_client:
if isinstance(cached_client, OpenAI) or isinstance(
cached_client, AsyncOpenAI
):
return cached_client
if is_async:
_new_client: Union[OpenAI, AsyncOpenAI] = AsyncOpenAI(
api_key=api_key,
base_url=api_base,
http_client=OpenAIChatCompletion._get_async_http_client(
shared_session=shared_session
),
timeout=timeout,
max_retries=max_retries,
organization=organization,
)
else:
_new_client = OpenAI(
api_key=api_key,
base_url=api_base,
http_client=OpenAIChatCompletion._get_sync_http_client(),
timeout=timeout,
max_retries=max_retries,
organization=organization,
)
## SAVE CACHE KEY
self.set_cached_openai_client(
openai_client=_new_client,
client_initialization_params=client_initialization_params,
client_type="openai",
)
return _new_client
else:
# If caller passed a client explicitly, do not use cache
if client is not None:
self._set_dynamic_params_on_client(
client=client,
organization=organization,
@ -418,6 +378,66 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM):
)
return client
# IMPORTANT: compute params AFTER explicit-client early return so `client` is always None in cache-key
client_initialization_params: Dict = locals()
# Cache path
if not isinstance(max_retries, int):
raise OpenAIError(
status_code=422,
message="max retries must be an int. Passed in value: {}".format(
max_retries
),
)
cached_client = self.get_cached_openai_client(
client_initialization_params=client_initialization_params,
client_type="openai",
)
if cached_client and isinstance(cached_client, (AsyncOpenAI, OpenAI)):
if self._httpx_is_closed(cached_client):
raw_key = BaseOpenAILLM.get_openai_client_cache_key(
client_initialization_params=client_initialization_params,
client_type="openai",
)
key = litellm.in_memory_llm_clients_cache.update_cache_key_with_event_loop(
raw_key
)
litellm.in_memory_llm_clients_cache._remove_key(key)
else:
return cached_client
# Create new client
if is_async:
_new_client: Union[OpenAI, AsyncOpenAI] = AsyncOpenAI(
api_key=api_key,
base_url=api_base,
http_client=OpenAIChatCompletion._get_async_http_client(
shared_session=shared_session
),
timeout=timeout,
max_retries=max_retries,
organization=organization,
)
else:
_new_client = OpenAI(
api_key=api_key,
base_url=api_base,
http_client=OpenAIChatCompletion._get_sync_http_client(),
timeout=timeout,
max_retries=max_retries,
organization=organization,
)
## SAVE CACHE KEY
self.set_cached_openai_client(
openai_client=_new_client,
client_initialization_params=client_initialization_params,
client_type="openai",
)
return _new_client
@track_llm_api_timing()
async def make_openai_chat_completion_request(
self,
@ -522,17 +542,14 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM):
from litellm._logging import verbose_logger
from litellm.integrations.custom_logger import CustomLogger
callbacks = litellm.callbacks + (
logging_obj.dynamic_success_callbacks or []
)
callbacks = litellm.callbacks + (logging_obj.dynamic_success_callbacks or [])
# Avoid logging full callback objects to prevent leaking sensitive data
verbose_logger.debug(
"LiteLLM.AgenticHooks: callbacks_count=%s", len(callbacks)
)
verbose_logger.debug("LiteLLM.AgenticHooks: callbacks_count=%s", len(callbacks))
tools = optional_params.get("tools", [])
# Avoid logging full tools payloads; they may contain sensitive parameters
verbose_logger.debug(
"LiteLLM.AgenticHooks: tools_count=%s", len(tools) if isinstance(tools, list) else 1 if tools else 0
"LiteLLM.AgenticHooks: tools_count=%s",
len(tools) if isinstance(tools, list) else 1 if tools else 0,
)
# Get custom_llm_provider from litellm_params
custom_llm_provider = litellm_params.get("custom_llm_provider", "openai")
@ -541,37 +558,46 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM):
try:
if isinstance(callback, CustomLogger):
# Check if the callback has the chat completion agentic loop methods
if not hasattr(callback, 'async_should_run_chat_completion_agentic_loop'):
if not hasattr(
callback, "async_should_run_chat_completion_agentic_loop"
):
continue
# First: Check if agentic loop should run (using chat completion method)
should_run, tool_calls = (
await callback.async_should_run_chat_completion_agentic_loop(
response=response,
model=model,
messages=messages,
tools=tools,
stream=stream,
custom_llm_provider=custom_llm_provider,
kwargs=litellm_params,
)
(
should_run,
tool_calls,
) = await callback.async_should_run_chat_completion_agentic_loop(
response=response,
model=model,
messages=messages,
tools=tools,
stream=stream,
custom_llm_provider=custom_llm_provider,
kwargs=litellm_params,
)
if should_run:
# Second: Execute agentic loop
kwargs_with_provider = litellm_params.copy() if litellm_params else {}
kwargs_with_provider["custom_llm_provider"] = custom_llm_provider
kwargs_with_provider = (
litellm_params.copy() if litellm_params else {}
)
kwargs_with_provider[
"custom_llm_provider"
] = custom_llm_provider
# For OpenAI Chat Completions, use the chat completion agentic loop method
agentic_response = await callback.async_run_chat_completion_agentic_loop(
tools=tool_calls,
model=model,
messages=messages,
response=response,
optional_params=optional_params,
logging_obj=logging_obj,
stream=stream,
kwargs=kwargs_with_provider,
agentic_response = (
await callback.async_run_chat_completion_agentic_loop(
tools=tool_calls,
model=model,
messages=messages,
response=response,
optional_params=optional_params,
logging_obj=logging_obj,
stream=stream,
kwargs=kwargs_with_provider,
)
)
# First hook that runs agentic loop wins
return agentic_response
@ -951,7 +977,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM):
stream=False,
litellm_params=litellm_params,
)
if agentic_response is not None:
final_response_obj = agentic_response

206
tests/test_litellm/llms/openai/test_openai_client_cache_cleanup.py Normal file
View file

@ -0,0 +1,206 @@
import os
import sys
from types import SimpleNamespace
from unittest.mock import MagicMock, patch
import pytest
sys.path.insert(0, os.path.abspath("../../.."))
import litellm
from litellm.llms.openai.openai import OpenAIChatCompletion
from openai import AsyncOpenAI, OpenAI
@pytest.fixture(autouse=True)
def _clear_in_memory_cache():
"""Clear global cache to avoid cross-test contamination."""
litellm.in_memory_llm_clients_cache.cache_dict.clear()
yield
litellm.in_memory_llm_clients_cache.cache_dict.clear()
class TestHttpClientClosedHandling:
def test_httpx_is_closed_returns_true_when_client_is_none(self):
chat = OpenAIChatCompletion()
assert chat._httpx_is_closed(None) is True
def test_httpx_is_closed_returns_true_when_client_attribute_missing(self):
chat = OpenAIChatCompletion()
client_without__client = SimpleNamespace() # no _client attribute
assert chat._httpx_is_closed(client_without__client) is True
def test_httpx_is_closed_returns_true_when_httpx_is_closed(self):
chat = OpenAIChatCompletion()
openai_client = SimpleNamespace(_client=SimpleNamespace(is_closed=True))
assert chat._httpx_is_closed(openai_client) is True
def test_httpx_is_closed_returns_false_when_httpx_is_open(self):
chat = OpenAIChatCompletion()
openai_client = SimpleNamespace(_client=SimpleNamespace(is_closed=False))
assert chat._httpx_is_closed(openai_client) is False
class TestOpenAIClientCacheLogic:
@patch.object(OpenAIChatCompletion, "_set_dynamic_params_on_client")
def test_explicit_client_bypasses_cache(self, mock_set_dynamic):
chat = OpenAIChatCompletion()
explicit_client = MagicMock(spec=OpenAI)
with patch.object(
OpenAIChatCompletion, "get_cached_openai_client"
) as mock_get_cache:
result = chat._get_openai_client(
is_async=False,
api_key="test-key",
client=explicit_client,
max_retries=3,
)
assert result is explicit_client
mock_get_cache.assert_not_called()
mock_set_dynamic.assert_called_once_with(
client=explicit_client, organization=None, max_retries=3
)
def test_max_retries_must_be_int_raises(self):
chat = OpenAIChatCompletion()
with pytest.raises(Exception) as excinfo:
chat._get_openai_client(
is_async=False,
api_key="test-key",
max_retries="3", # not int
)
assert "max retries must be an int" in str(excinfo.value)
@patch.object(OpenAIChatCompletion, "set_cached_openai_client")
@patch.object(OpenAIChatCompletion, "get_cached_openai_client")
def test_cached_async_client_open_is_returned(self, mock_get_cache, mock_set_cache):
chat = OpenAIChatCompletion()
cached_async = object.__new__(AsyncOpenAI)
cached_async._client = SimpleNamespace(is_closed=False)
mock_get_cache.return_value = cached_async
with patch.object(
AsyncOpenAI, "__init__", return_value=None
) as mock_async_init:
result = chat._get_openai_client(
is_async=True,
api_key="test-key",
max_retries=3,
)
assert result is cached_async
mock_async_init.assert_not_called()
mock_set_cache.assert_not_called()
@patch.object(OpenAIChatCompletion, "set_cached_openai_client")
@patch.object(OpenAIChatCompletion, "get_cached_openai_client")
def test_cached_async_client_closed_is_evicted_and_new_created(
self, mock_get_cache, mock_set_cache
):
chat = OpenAIChatCompletion()
cached_async = object.__new__(AsyncOpenAI)
cached_async._client = SimpleNamespace(is_closed=True)
mock_get_cache.return_value = cached_async
# Store it under the *event-loop-suffixed* key, like set_cache() does
raw_key = "fake-cache-key"
loop_key = litellm.in_memory_llm_clients_cache.update_cache_key_with_event_loop(
raw_key
)
litellm.in_memory_llm_clients_cache.cache_dict[loop_key] = cached_async
with patch(
"litellm.llms.openai.openai.BaseOpenAILLM.get_openai_client_cache_key",
return_value=raw_key,
), patch.object(
OpenAIChatCompletion, "_get_async_http_client", return_value=MagicMock()
), patch.object(
AsyncOpenAI, "__init__", return_value=None
) as mock_async_init, patch.object(
litellm.in_memory_llm_clients_cache,
"_remove_key",
wraps=litellm.in_memory_llm_clients_cache._remove_key,
) as mock_remove_key:
result = chat._get_openai_client(
is_async=True,
api_key="test-key",
max_retries=3,
)
assert isinstance(result, AsyncOpenAI)
assert result is not cached_async
# Evicted using the loop-suffixed key
assert loop_key not in litellm.in_memory_llm_clients_cache.cache_dict
mock_remove_key.assert_called_once_with(loop_key)
mock_set_cache.assert_called_once()
mock_async_init.assert_called_once()
@patch.object(OpenAIChatCompletion, "set_cached_openai_client")
@patch.object(OpenAIChatCompletion, "get_cached_openai_client")
def test_cached_sync_client_open_is_returned(self, mock_get_cache, mock_set_cache):
chat = OpenAIChatCompletion()
cached_sync = object.__new__(OpenAI)
cached_sync._client = SimpleNamespace(is_closed=False)
mock_get_cache.return_value = cached_sync
with patch.object(OpenAI, "__init__", return_value=None) as mock_sync_init:
result = chat._get_openai_client(
is_async=False,
api_key="test-key",
max_retries=3,
)
assert result is cached_sync
mock_sync_init.assert_not_called()
mock_set_cache.assert_not_called()
@patch.object(OpenAIChatCompletion, "set_cached_openai_client")
@patch.object(OpenAIChatCompletion, "get_cached_openai_client")
def test_cached_sync_client_closed_is_evicted_and_new_created(
self, mock_get_cache, mock_set_cache
):
chat = OpenAIChatCompletion()
cached_sync = object.__new__(OpenAI)
cached_sync._client = SimpleNamespace(is_closed=True)
mock_get_cache.return_value = cached_sync
raw_key = "fake-cache-key"
loop_key = litellm.in_memory_llm_clients_cache.update_cache_key_with_event_loop(
raw_key
)
litellm.in_memory_llm_clients_cache.cache_dict[loop_key] = cached_sync
with patch(
"litellm.llms.openai.openai.BaseOpenAILLM.get_openai_client_cache_key",
return_value=raw_key,
), patch.object(
OpenAIChatCompletion, "_get_sync_http_client", return_value=MagicMock()
), patch.object(
OpenAI, "__init__", return_value=None
) as mock_sync_init, patch.object(
litellm.in_memory_llm_clients_cache,
"_remove_key",
wraps=litellm.in_memory_llm_clients_cache._remove_key,
) as mock_remove_key:
result = chat._get_openai_client(
is_async=False,
api_key="test-key",
max_retries=3,
)
assert isinstance(result, OpenAI)
assert result is not cached_sync
assert loop_key not in litellm.in_memory_llm_clients_cache.cache_dict
mock_remove_key.assert_called_once_with(loop_key)
mock_set_cache.assert_called_once()
mock_sync_init.assert_called_once()