diff --git a/rlm/clients/anthropic.py b/rlm/clients/anthropic.py index a4ab53407..5c747de4b 100644 --- a/rlm/clients/anthropic.py +++ b/rlm/clients/anthropic.py @@ -20,8 +20,8 @@ def __init__( **kwargs, ): super().__init__(model_name=model_name, **kwargs) - self.client = anthropic.Anthropic(api_key=api_key) - self.async_client = anthropic.AsyncAnthropic(api_key=api_key) + self.client = anthropic.Anthropic(api_key=api_key, timeout=self.timeout) + self.async_client = anthropic.AsyncAnthropic(api_key=api_key, timeout=self.timeout) self.model_name = model_name self.max_tokens = max_tokens diff --git a/rlm/clients/azure_openai.py b/rlm/clients/azure_openai.py index 90e04da3f..f8e0d57ee 100644 --- a/rlm/clients/azure_openai.py +++ b/rlm/clients/azure_openai.py @@ -53,12 +53,14 @@ def __init__( azure_endpoint=azure_endpoint, api_version=api_version, azure_deployment=azure_deployment, + timeout=self.timeout, ) self.async_client = openai.AsyncAzureOpenAI( api_key=api_key, azure_endpoint=azure_endpoint, api_version=api_version, azure_deployment=azure_deployment, + timeout=self.timeout, ) self.model_name = model_name self.azure_deployment = azure_deployment diff --git a/rlm/clients/base_lm.py b/rlm/clients/base_lm.py index 365536621..e1ec8f70c 100644 --- a/rlm/clients/base_lm.py +++ b/rlm/clients/base_lm.py @@ -3,6 +3,9 @@ from rlm.core.types import ModelUsageSummary, UsageSummary +# Default timeout for LM API calls (in seconds) +DEFAULT_TIMEOUT: float = 300.0 + class BaseLM(ABC): """ @@ -10,8 +13,9 @@ class BaseLM(ABC): does so in a model-agnostic way, so this class provides a base interface for all language models. """ - def __init__(self, model_name: str, **kwargs): + def __init__(self, model_name: str, timeout: float = DEFAULT_TIMEOUT, **kwargs): self.model_name = model_name + self.timeout = timeout self.kwargs = kwargs @abstractmethod diff --git a/rlm/clients/gemini.py b/rlm/clients/gemini.py index 7f6dc1520..8a211611d 100644 --- a/rlm/clients/gemini.py +++ b/rlm/clients/gemini.py @@ -36,7 +36,9 @@ def __init__( "Gemini API key is required. Set GEMINI_API_KEY env var or pass api_key." ) - self.client = genai.Client(api_key=api_key) + # Configure HTTP options with timeout + http_options = types.HttpOptions(timeout=int(self.timeout * 1000)) # milliseconds + self.client = genai.Client(api_key=api_key, http_options=http_options) self.model_name = model_name # Per-model usage tracking diff --git a/rlm/clients/litellm.py b/rlm/clients/litellm.py index a4247159a..1efec752d 100644 --- a/rlm/clients/litellm.py +++ b/rlm/clients/litellm.py @@ -43,7 +43,7 @@ def completion(self, prompt: str | list[dict[str, Any]], model: str | None = Non if not model: raise ValueError("Model name is required for LiteLLM client.") - kwargs = {"model": model, "messages": messages} + kwargs = {"model": model, "messages": messages, "timeout": self.timeout} if self.api_key: kwargs["api_key"] = self.api_key if self.api_base: @@ -67,7 +67,7 @@ async def acompletion( if not model: raise ValueError("Model name is required for LiteLLM client.") - kwargs = {"model": model, "messages": messages} + kwargs = {"model": model, "messages": messages, "timeout": self.timeout} if self.api_key: kwargs["api_key"] = self.api_key if self.api_base: diff --git a/rlm/clients/openai.py b/rlm/clients/openai.py index 7f60a1eae..b235d4b71 100644 --- a/rlm/clients/openai.py +++ b/rlm/clients/openai.py @@ -40,8 +40,10 @@ def __init__( api_key = DEFAULT_VERCEL_API_KEY # For vLLM, set base_url to local vLLM server address. - self.client = openai.OpenAI(api_key=api_key, base_url=base_url) - self.async_client = openai.AsyncOpenAI(api_key=api_key, base_url=base_url) + self.client = openai.OpenAI(api_key=api_key, base_url=base_url, timeout=self.timeout) + self.async_client = openai.AsyncOpenAI( + api_key=api_key, base_url=base_url, timeout=self.timeout + ) self.model_name = model_name # Per-model usage tracking diff --git a/rlm/clients/portkey.py b/rlm/clients/portkey.py index 6835c7404..b8a310e29 100644 --- a/rlm/clients/portkey.py +++ b/rlm/clients/portkey.py @@ -21,8 +21,8 @@ def __init__( **kwargs, ): super().__init__(model_name=model_name, **kwargs) - self.client = Portkey(api_key=api_key, base_url=base_url) - self.async_client = AsyncPortkey(api_key=api_key, base_url=base_url) + self.client = Portkey(api_key=api_key, base_url=base_url, timeout=self.timeout) + self.async_client = AsyncPortkey(api_key=api_key, base_url=base_url, timeout=self.timeout) self.model_name = model_name # Per-model usage tracking diff --git a/tests/clients/test_timeout.py b/tests/clients/test_timeout.py new file mode 100644 index 000000000..c0ee0903d --- /dev/null +++ b/tests/clients/test_timeout.py @@ -0,0 +1,182 @@ +"""Tests for client timeout functionality.""" + +from unittest.mock import MagicMock, patch + +import httpx +import pytest + +from rlm.clients.base_lm import DEFAULT_TIMEOUT + + +class TestDefaultTimeout: + """Tests for the default timeout constant.""" + + def test_default_timeout_value(self): + """Default timeout should be 300 seconds.""" + assert DEFAULT_TIMEOUT == 300.0 + + def test_base_lm_stores_timeout(self): + """BaseLM should store timeout in instance.""" + from rlm.clients.openai import OpenAIClient + + with patch("rlm.clients.openai.openai.OpenAI"): + with patch("rlm.clients.openai.openai.AsyncOpenAI"): + client = OpenAIClient(api_key="test-key", model_name="gpt-4o") + assert client.timeout == DEFAULT_TIMEOUT + + def test_custom_timeout_override(self): + """Custom timeout should override default.""" + from rlm.clients.openai import OpenAIClient + + with patch("rlm.clients.openai.openai.OpenAI"): + with patch("rlm.clients.openai.openai.AsyncOpenAI"): + client = OpenAIClient(api_key="test-key", model_name="gpt-4o", timeout=60.0) + assert client.timeout == 60.0 + + +class TestOpenAIClientTimeout: + """Tests for OpenAI client timeout.""" + + def test_timeout_passed_to_client(self): + """Timeout should be passed to OpenAI client.""" + from rlm.clients.openai import OpenAIClient + + with patch("rlm.clients.openai.openai.OpenAI") as mock_openai: + with patch("rlm.clients.openai.openai.AsyncOpenAI") as mock_async: + OpenAIClient(api_key="test-key", model_name="gpt-4o", timeout=120.0) + + # Check sync client received timeout + mock_openai.assert_called_once() + call_kwargs = mock_openai.call_args[1] + assert call_kwargs["timeout"] == 120.0 + + # Check async client received timeout + mock_async.assert_called_once() + async_call_kwargs = mock_async.call_args[1] + assert async_call_kwargs["timeout"] == 120.0 + + def test_timeout_raises_exception(self): + """Timeout should raise appropriate exception.""" + from rlm.clients.openai import OpenAIClient + + # Create a mock client that raises timeout + mock_client = MagicMock() + mock_client.chat.completions.create.side_effect = httpx.TimeoutException( + "Connection timed out" + ) + + with patch("rlm.clients.openai.openai.OpenAI", return_value=mock_client): + with patch("rlm.clients.openai.openai.AsyncOpenAI"): + client = OpenAIClient(api_key="test-key", model_name="gpt-4o", timeout=0.001) + + with pytest.raises(httpx.TimeoutException): + client.completion("Hello") + + +class TestAnthropicClientTimeout: + """Tests for Anthropic client timeout.""" + + def test_timeout_passed_to_client(self): + """Timeout should be passed to Anthropic client.""" + from rlm.clients.anthropic import AnthropicClient + + with patch("rlm.clients.anthropic.anthropic.Anthropic") as mock_anthropic: + with patch("rlm.clients.anthropic.anthropic.AsyncAnthropic") as mock_async: + AnthropicClient( + api_key="test-key", model_name="claude-sonnet-4-20250514", timeout=120.0 + ) + + mock_anthropic.assert_called_once() + call_kwargs = mock_anthropic.call_args[1] + assert call_kwargs["timeout"] == 120.0 + + mock_async.assert_called_once() + async_call_kwargs = mock_async.call_args[1] + assert async_call_kwargs["timeout"] == 120.0 + + +class TestAzureOpenAIClientTimeout: + """Tests for Azure OpenAI client timeout.""" + + def test_timeout_passed_to_client(self): + """Timeout should be passed to Azure OpenAI client.""" + from rlm.clients.azure_openai import AzureOpenAIClient + + with patch("rlm.clients.azure_openai.openai.AzureOpenAI") as mock_azure: + with patch("rlm.clients.azure_openai.openai.AsyncAzureOpenAI") as mock_async: + AzureOpenAIClient( + api_key="test-key", + model_name="gpt-4o", + azure_endpoint="https://test.openai.azure.com", + timeout=120.0, + ) + + mock_azure.assert_called_once() + call_kwargs = mock_azure.call_args[1] + assert call_kwargs["timeout"] == 120.0 + + mock_async.assert_called_once() + async_call_kwargs = mock_async.call_args[1] + assert async_call_kwargs["timeout"] == 120.0 + + +class TestPortkeyClientTimeout: + """Tests for Portkey client timeout.""" + + def test_timeout_passed_to_client(self): + """Timeout should be passed to Portkey client.""" + from rlm.clients.portkey import PortkeyClient + + with patch("rlm.clients.portkey.Portkey") as mock_portkey: + with patch("rlm.clients.portkey.AsyncPortkey") as mock_async: + PortkeyClient(api_key="test-key", model_name="gpt-4o", timeout=120.0) + + mock_portkey.assert_called_once() + call_kwargs = mock_portkey.call_args[1] + assert call_kwargs["timeout"] == 120.0 + + mock_async.assert_called_once() + async_call_kwargs = mock_async.call_args[1] + assert async_call_kwargs["timeout"] == 120.0 + + +class TestLiteLLMClientTimeout: + """Tests for LiteLLM client timeout.""" + + def test_timeout_passed_to_completion(self): + """Timeout should be passed to litellm.completion call.""" + pytest.importorskip("litellm") + from rlm.clients.litellm import LiteLLMClient + + mock_response = MagicMock() + mock_response.choices = [MagicMock()] + mock_response.choices[0].message.content = "Hello" + mock_response.usage.prompt_tokens = 10 + mock_response.usage.completion_tokens = 5 + mock_response.usage.total_tokens = 15 + + with patch( + "rlm.clients.litellm.litellm.completion", return_value=mock_response + ) as mock_completion: + client = LiteLLMClient(model_name="gpt-4o", timeout=120.0) + client.completion("Hello") + + mock_completion.assert_called_once() + call_kwargs = mock_completion.call_args[1] + assert call_kwargs["timeout"] == 120.0 + + +class TestGeminiClientTimeout: + """Tests for Gemini client timeout.""" + + def test_timeout_passed_to_client(self): + """Timeout should be passed to Gemini client via http_options.""" + from rlm.clients.gemini import GeminiClient + + with patch("rlm.clients.gemini.genai.Client") as mock_genai: + GeminiClient(api_key="test-key", model_name="gemini-2.5-flash", timeout=120.0) + + mock_genai.assert_called_once() + call_kwargs = mock_genai.call_args[1] + # Gemini uses milliseconds + assert call_kwargs["http_options"].timeout == 120000