diff --git a/litellm/llms/openai/openai.py b/litellm/llms/openai/openai.py index c7524925bd0..44c168bf286 100644 --- a/litellm/llms/openai/openai.py +++ b/litellm/llms/openai/openai.py @@ -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 diff --git a/tests/test_litellm/llms/openai/test_openai_client_cache_cleanup.py b/tests/test_litellm/llms/openai/test_openai_client_cache_cleanup.py new file mode 100644 index 00000000000..921de3f2062 --- /dev/null +++ b/tests/test_litellm/llms/openai/test_openai_client_cache_cleanup.py @@ -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()