diff --git a/agent/core/llm.py b/agent/core/llm.py index 4168ec1..2a744db 100644 --- a/agent/core/llm.py +++ b/agent/core/llm.py @@ -21,6 +21,8 @@ class LLMClient: def __init__(self, config: LLMConfig): self.config = config self.base_url = config.base_url.rstrip("/") + self.max_retries = 3 + self.retry_delay = 1.0 def _validate_response(self, response: dict, required_fields: list[str]) -> None: """Validate that response contains expected fields. @@ -76,36 +78,49 @@ def _request( ) -> dict: url = f"{self.base_url}{endpoint}" payload = json.dumps(data).encode("utf-8") - req = urllib.request.Request( - url, - data=payload, - headers={"Content-Type": "application/json"}, - method="POST", - ) effective_timeout = timeout if timeout is not None else self.config.timeout - try: - resp = urllib.request.urlopen(req, timeout=effective_timeout) - if stream: - return resp # Return response object for streaming - return json.loads(resp.read().decode("utf-8")) - except socket.timeout: - raise OllamaError( - f"Timeout connecting to Ollama at {self.base_url} " - f"(waited {self.config.timeout}s). Make sure Ollama is running." - ) - except urllib.error.URLError as e: - if isinstance(e.reason, (socket.timeout, socket.error, ConnectionRefusedError, OSError)): + + for attempt in range(self.max_retries): + try: + req = urllib.request.Request( + url, + data=payload, + headers={"Content-Type": "application/json"}, + method="POST", + ) + resp = urllib.request.urlopen(req, timeout=effective_timeout) + if stream: + return resp + return json.loads(resp.read().decode("utf-8")) + except socket.timeout as e: + # A timeout means Ollama is still working, not that the + # connection failed - retrying would just queue up more + # concurrent work against an already-struggling server. + raise OllamaError( + f"Timeout connecting to Ollama at {self.base_url} " + f"(waited {effective_timeout}s). Make sure Ollama is running." + ) from e + except urllib.error.HTTPError as e: + # Must be checked before URLError: HTTPError subclasses it. + body = e.read().decode("utf-8", errors="replace") + raise OllamaError(f"Ollama API error ({e.code}): {body}") from e + except urllib.error.URLError as e: + if isinstance(e.reason, (ConnectionRefusedError, socket.gaierror)): + if attempt < self.max_retries - 1: + time.sleep(self.retry_delay * (2 ** attempt)) + continue + raise OllamaError( + f"Cannot connect to Ollama at {self.base_url} " + f"after {self.max_retries} attempts. " + f"Make sure Ollama is running: {e}" + ) from e raise OllamaError( f"Cannot connect to Ollama at {self.base_url}. " f"Make sure Ollama is running: {e}" - ) - raise OllamaError( - f"Cannot connect to Ollama at {self.base_url}. " - f"Make sure Ollama is running: {e}" - ) - except urllib.error.HTTPError as e: - body = e.read().decode("utf-8", errors="replace") - raise OllamaError(f"Ollama API error ({e.code}): {body}") + ) from e + + # Unreachable: every loop iteration either returns or raises above. + raise OllamaError(f"Failed to reach Ollama at {self.base_url}.") def check_health(self) -> bool: """Check if Ollama is running and the model is available.""" diff --git a/tests/test_llm.py b/tests/test_llm.py index e9955ef..85ebd38 100644 --- a/tests/test_llm.py +++ b/tests/test_llm.py @@ -1,8 +1,11 @@ """Tests for the Ollama LLM client response validation.""" +import io import json +import socket import sys import unittest +import urllib.error from pathlib import Path from unittest.mock import patch @@ -166,5 +169,62 @@ def test_non_dict_message_is_skipped(self): self.assertEqual(result["message"]["content"], "ok") +class TestRequestRetry(unittest.TestCase): + """Test the retry behavior in LLMClient._request.""" + + def test_connection_refused_is_retried_and_eventually_raises(self): + client = make_client() + error = urllib.error.URLError(ConnectionRefusedError()) + with patch("urllib.request.urlopen", side_effect=error) as mock_urlopen, \ + patch("time.sleep") as mock_sleep: + with self.assertRaises(OllamaError): + client._request("/api/chat", {}) + self.assertEqual(mock_urlopen.call_count, client.max_retries) + self.assertEqual(mock_sleep.call_count, client.max_retries - 1) + + def test_dns_failure_is_retried_and_eventually_raises(self): + client = make_client() + error = urllib.error.URLError(socket.gaierror()) + with patch("urllib.request.urlopen", side_effect=error) as mock_urlopen, \ + patch("time.sleep"): + with self.assertRaises(OllamaError): + client._request("/api/chat", {}) + self.assertEqual(mock_urlopen.call_count, client.max_retries) + + def test_socket_timeout_fails_fast_without_retry(self): + """A timeout means Ollama is still generating - retrying would only + make a slow response slower, so it must not be retried.""" + client = make_client() + with patch("urllib.request.urlopen", side_effect=socket.timeout()) as mock_urlopen, \ + patch("time.sleep") as mock_sleep: + with self.assertRaises(OllamaError): + client._request("/api/chat", {}) + self.assertEqual(mock_urlopen.call_count, 1) + mock_sleep.assert_not_called() + + def test_non_connection_url_error_fails_fast_without_retry(self): + client = make_client() + error = urllib.error.URLError("ssl handshake failed") + with patch("urllib.request.urlopen", side_effect=error) as mock_urlopen, \ + patch("time.sleep") as mock_sleep: + with self.assertRaises(OllamaError): + client._request("/api/chat", {}) + self.assertEqual(mock_urlopen.call_count, 1) + mock_sleep.assert_not_called() + + def test_http_error_surfaces_response_body(self): + """HTTPError subclasses URLError, so its handler must be checked + first or the actual API error body never reaches the caller.""" + client = make_client() + http_error = urllib.error.HTTPError( + "http://x", 400, "Bad Request", {}, io.BytesIO(b'{"error": "bad model"}') + ) + with patch("urllib.request.urlopen", side_effect=http_error): + with self.assertRaises(OllamaError) as ctx: + client._request("/api/chat", {}) + self.assertIn("bad model", str(ctx.exception)) + self.assertIn("400", str(ctx.exception)) + + if __name__ == "__main__": unittest.main()