fix(llm): replace retry with tenacity in ollama.py (apache#367)

bitflicker64 · web-flow · commit e58a098d2e7c · 2026-06-23T12:33:29.000+08:00
## Summary `ollama.py` was using the abandoned `retry` PyPI package (last maintained 2016) while `openai.py` and `litellm.py` in the same directory already use `tenacity`. This PR brings `ollama.py` in line with the established patterns. Closes apache#365
diff --git a/hugegraph-llm/pyproject.toml b/hugegraph-llm/pyproject.toml
@@ -47,7 +47,7 @@ dependencies = [
     # LLM specific dependencies
     "openai",
     "ollama",
-    "retry",
+    "tenacity",
     "tiktoken",
     "nltk",
     "gradio",
diff --git a/hugegraph-llm/src/hugegraph_llm/models/llms/ollama.py b/hugegraph-llm/src/hugegraph_llm/models/llms/ollama.py
@@ -19,8 +19,14 @@
 import json
 from typing import Any, AsyncGenerator, Callable, Dict, Generator, List, Optional
 
+import httpx
 import ollama
-from retry import retry
+from tenacity import (
+    retry,
+    retry_if_exception_type,
+    stop_after_attempt,
+    wait_exponential,
+)
 
 from hugegraph_llm.models.llms.base import BaseLLM
 from hugegraph_llm.utils.log import log
@@ -34,13 +40,17 @@ def __init__(self, model: str, host: str = "127.0.0.1", port: int = 11434, **kwa
         self.client = ollama.Client(host=f"http://{host}:{port}", **kwargs)
         self.async_client = ollama.AsyncClient(host=f"http://{host}:{port}", **kwargs)
 
-    @retry(tries=3, delay=1)
+    @retry(
+        stop=stop_after_attempt(3),
+        wait=wait_exponential(multiplier=1, min=4, max=10),
+        retry=retry_if_exception_type((ollama.ResponseError, httpx.ConnectError, httpx.TimeoutException)),
+    )
     def generate(
         self,
         messages: Optional[List[Dict[str, Any]]] = None,
         prompt: Optional[str] = None,
     ) -> str:
-        """Comment"""
+        """Generate a response to the query messages/prompt."""
         if messages is None:
             assert prompt is not None, "Messages or prompt must be provided."
             messages = [{"role": "user", "content": prompt}]
@@ -56,17 +66,21 @@ def generate(
             }
             log.info("Token usage: %s", json.dumps(usage))
             return response["message"]["content"]
-        except Exception as e:
-            print(f"Retrying LLM call {e}")
-            raise e
-
-    @retry(tries=3, delay=1)
+        except (ollama.ResponseError, httpx.ConnectError, httpx.TimeoutException) as e:
+            log.error("Retrying LLM call %s", e)
+            raise
+
+    @retry(
+        stop=stop_after_attempt(3),
+        wait=wait_exponential(multiplier=1, min=4, max=10),
+        retry=retry_if_exception_type((ollama.ResponseError, httpx.ConnectError, httpx.TimeoutException)),
+    )
     async def agenerate(
         self,
         messages: Optional[List[Dict[str, Any]]] = None,
         prompt: Optional[str] = None,
     ) -> str:
-        """Comment"""
+        """Generate a response to the query messages/prompt asynchronously."""
         if messages is None:
             assert prompt is not None, "Messages or prompt must be provided."
             messages = [{"role": "user", "content": prompt}]
@@ -82,17 +96,17 @@ async def agenerate(
             }
             log.info("Token usage: %s", json.dumps(usage))
             return response["message"]["content"]
-        except Exception as e:
-            print(f"Retrying LLM call {e}")
-            raise e
+        except (ollama.ResponseError, httpx.ConnectError, httpx.TimeoutException) as e:
+            log.error("Retrying LLM call %s", e)
+            raise
 
     def generate_streaming(
         self,
         messages: Optional[List[Dict[str, Any]]] = None,
         prompt: Optional[str] = None,
         on_token_callback: Optional[Callable] = None,
     ) -> Generator[str, None, None]:
-        """Comment"""
+        """Stream response tokens one by one."""
         if messages is None:
             assert prompt is not None, "Messages or prompt must be provided."
             messages = [{"role": "user", "content": prompt}]
@@ -112,7 +126,7 @@ async def agenerate_streaming(
         prompt: Optional[str] = None,
         on_token_callback: Optional[Callable] = None,
     ) -> AsyncGenerator[str, None]:
-        """Comment"""
+        """Stream response tokens one by one."""
         if messages is None:
             assert prompt is not None, "Messages or prompt must be provided."
             messages = [{"role": "user", "content": prompt}]
@@ -124,9 +138,9 @@ async def agenerate_streaming(
                 if on_token_callback:
                     on_token_callback(token)
                 yield token
-        except Exception as e:
-            print(f"Retrying LLM call {e}")
-            raise e
+        except (ollama.ResponseError, httpx.ConnectError, httpx.TimeoutException) as e:
+            log.error("Error in agenerate_streaming: %s", e)
+            raise
 
     def num_tokens_from_string(
         self,
diff --git a/hugegraph-llm/src/tests/models/llms/test_ollama_client.py b/hugegraph-llm/src/tests/models/llms/test_ollama_client.py
@@ -15,27 +15,212 @@
 # specific language governing permissions and limitations
 # under the License.
 
+import asyncio
 import os
 import unittest
+from unittest.mock import AsyncMock, MagicMock, patch
+
+import httpx
+import ollama
+import pytest
+from tenacity import RetryError, wait_none
 
 from hugegraph_llm.models.llms.ollama import OllamaClient
 
+pytestmark = pytest.mark.contract
+
+# Minimal dict response matching the structure ollama.Client.chat() returns
+_MOCK_RESPONSE = {
+    "prompt_eval_count": 10,
+    "eval_count": 5,
+    "message": {"content": "Paris"},
+}
+
+
+class TestOllamaClientRetryPolicy(unittest.TestCase):
+    """Mock-based contract tests for the Tenacity retry policy in OllamaClient.
+
+    These tests do not require a running Ollama service.  They verify:
+      - retryable exceptions (ollama.ResponseError, httpx.ConnectError,
+        httpx.TimeoutException) trigger the configured number of attempts;
+      - non-retryable exceptions (e.g. ValueError) are NOT retried;
+      - a transient failure followed by success resolves correctly.
+    """
+
+    def setUp(self):
+        # Zero out exponential wait so retry tests complete in milliseconds.
+        # Tenacity exposes the Retrying object on the decorated function via
+        # the .retry attribute; its .wait field is mutable.
+        self._orig_generate_wait = OllamaClient.generate.retry.wait
+        self._orig_agenerate_wait = OllamaClient.agenerate.retry.wait
+        OllamaClient.generate.retry.wait = wait_none()
+        OllamaClient.agenerate.retry.wait = wait_none()
+
+    def tearDown(self):
+        OllamaClient.generate.retry.wait = self._orig_generate_wait
+        OllamaClient.agenerate.retry.wait = self._orig_agenerate_wait
+
+    # ------------------------------------------------------------------ #
+    # generate()                                                           #
+    # ------------------------------------------------------------------ #
+
+    @patch("hugegraph_llm.models.llms.ollama.ollama.Client")
+    def test_generate_returns_content_on_success(self, mock_client_class):
+        """Happy path: generate() returns the message content string."""
+        mock_client = MagicMock()
+        mock_client.chat.return_value = _MOCK_RESPONSE
+        mock_client_class.return_value = mock_client
+
+        result = OllamaClient(model="llama3").generate(prompt="hello")
+
+        self.assertEqual(result, "Paris")
+        mock_client.chat.assert_called_once()
+
+    @patch("hugegraph_llm.models.llms.ollama.ollama.Client")
+    def test_generate_retries_response_error_exhausts_all_attempts(self, mock_client_class):
+        """ollama.ResponseError is retryable; all 3 attempts are made."""
+        mock_client = MagicMock()
+        mock_client.chat.side_effect = ollama.ResponseError("model not found")
+        mock_client_class.return_value = mock_client
+
+        with self.assertRaises(RetryError):
+            OllamaClient(model="llama3").generate(prompt="hello")
+
+        self.assertEqual(mock_client.chat.call_count, 3)
+
+    @patch("hugegraph_llm.models.llms.ollama.ollama.Client")
+    def test_generate_retries_connect_error_exhausts_all_attempts(self, mock_client_class):
+        """httpx.ConnectError is retryable; all 3 attempts are made."""
+        mock_client = MagicMock()
+        mock_client.chat.side_effect = httpx.ConnectError("connection refused")
+        mock_client_class.return_value = mock_client
+
+        with self.assertRaises(RetryError):
+            OllamaClient(model="llama3").generate(prompt="hello")
+
+        self.assertEqual(mock_client.chat.call_count, 3)
+
+    @patch("hugegraph_llm.models.llms.ollama.ollama.Client")
+    def test_generate_does_not_retry_non_retriable_error(self, mock_client_class):
+        """ValueError is not in the retry predicate; only 1 attempt is made."""
+        mock_client = MagicMock()
+        mock_client.chat.side_effect = ValueError("unexpected")
+        mock_client_class.return_value = mock_client
+
+        with self.assertRaises(ValueError):
+            OllamaClient(model="llama3").generate(prompt="hello")
+
+        mock_client.chat.assert_called_once()
+
+    @patch("hugegraph_llm.models.llms.ollama.ollama.Client")
+    def test_generate_succeeds_on_second_attempt(self, mock_client_class):
+        """Transient ResponseError on attempt 1, success on attempt 2."""
+        mock_client = MagicMock()
+        mock_client.chat.side_effect = [
+            ollama.ResponseError("transient"),
+            _MOCK_RESPONSE,
+        ]
+        mock_client_class.return_value = mock_client
+
+        result = OllamaClient(model="llama3").generate(prompt="hello")
+
+        self.assertEqual(result, "Paris")
+        self.assertEqual(mock_client.chat.call_count, 2)
+
+    # ------------------------------------------------------------------ #
+    # agenerate()                                                          #
+    # ------------------------------------------------------------------ #
+
+    @patch("hugegraph_llm.models.llms.ollama.ollama.AsyncClient")
+    def test_agenerate_retries_connect_error_exhausts_all_attempts(self, mock_async_client_class):
+        """httpx.ConnectError is retryable in agenerate(); all 3 attempts made."""
+        mock_async_client = MagicMock()
+        mock_async_client.chat = AsyncMock(side_effect=httpx.ConnectError("connection refused"))
+        mock_async_client_class.return_value = mock_async_client
+
+        async def run():
+            with self.assertRaises(RetryError):
+                await OllamaClient(model="llama3").agenerate(prompt="hello")
+            self.assertEqual(mock_async_client.chat.call_count, 3)
+
+        asyncio.run(run())
+
+    @patch("hugegraph_llm.models.llms.ollama.ollama.AsyncClient")
+    def test_agenerate_retries_timeout_exception_exhausts_all_attempts(self, mock_async_client_class):
+        """httpx.TimeoutException is retryable in agenerate(); all 3 attempts made."""
+        mock_async_client = MagicMock()
+        mock_async_client.chat = AsyncMock(side_effect=httpx.TimeoutException("timed out"))
+        mock_async_client_class.return_value = mock_async_client
+
+        async def run():
+            with self.assertRaises(RetryError):
+                await OllamaClient(model="llama3").agenerate(prompt="hello")
+            self.assertEqual(mock_async_client.chat.call_count, 3)
+
+        asyncio.run(run())
+
+    @patch("hugegraph_llm.models.llms.ollama.ollama.AsyncClient")
+    def test_agenerate_does_not_retry_non_retriable_error(self, mock_async_client_class):
+        """ValueError is not retryable in agenerate(); only 1 attempt is made."""
+        mock_async_client = MagicMock()
+        mock_async_client.chat = AsyncMock(side_effect=ValueError("unexpected"))
+        mock_async_client_class.return_value = mock_async_client
+
+        async def run():
+            with self.assertRaises(ValueError):
+                await OllamaClient(model="llama3").agenerate(prompt="hello")
+            mock_async_client.chat.assert_called_once()
+
+        asyncio.run(run())
+
+    @patch("hugegraph_llm.models.llms.ollama.ollama.AsyncClient")
+    def test_agenerate_succeeds_on_second_attempt(self, mock_async_client_class):
+        """Transient ResponseError on attempt 1, success on attempt 2."""
+        mock_async_client = MagicMock()
+        mock_async_client.chat = AsyncMock(side_effect=[ollama.ResponseError("transient"), _MOCK_RESPONSE])
+        mock_async_client_class.return_value = mock_async_client
+
+        async def run():
+            result = await OllamaClient(model="llama3").agenerate(prompt="hello")
+            self.assertEqual(result, "Paris")
+            self.assertEqual(mock_async_client.chat.call_count, 2)
+
+        asyncio.run(run())
+
+
+class TestOllamaClientExternalService(unittest.TestCase):
+    """Integration tests that require a live Ollama service.
+
+    Skipped in CI via SKIP_EXTERNAL_SERVICES=true (set in conftest.py).
+    """
 
-class TestOllamaClient(unittest.TestCase):
     def setUp(self):
         self.skip_external = os.getenv("SKIP_EXTERNAL_SERVICES", "false").lower() == "true"
 
-    @unittest.skipIf(os.getenv("SKIP_EXTERNAL_SERVICES", "false").lower() == "true", "Skipping external service tests")
+    @unittest.skipIf(
+        os.getenv("SKIP_EXTERNAL_SERVICES", "false").lower() == "true",
+        "Skipping external service tests",
+    )
     def test_generate(self):
         ollama_client = OllamaClient(model="llama3:8b-instruct-fp16")
         response = ollama_client.generate(prompt="What is the capital of France?")
         print(response)
 
-    @unittest.skipIf(os.getenv("SKIP_EXTERNAL_SERVICES", "false").lower() == "true", "Skipping external service tests")
+    @unittest.skipIf(
+        os.getenv("SKIP_EXTERNAL_SERVICES", "false").lower() == "true",
+        "Skipping external service tests",
+    )
     def test_stream_generate(self):
         ollama_client = OllamaClient(model="llama3:8b-instruct-fp16")
 
         def on_token_callback(chunk):
             print(chunk, end="", flush=True)
 
-        ollama_client.generate_streaming(prompt="What is the capital of France?", on_token_callback=on_token_callback)
+        ollama_client.generate_streaming(
+            prompt="What is the capital of France?",
+            on_token_callback=on_token_callback,
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/pyproject.toml b/pyproject.toml
@@ -114,7 +114,7 @@ constraint-dependencies = [
     # LLM dependencies
     "openai~=1.61.0",
     "ollama~=0.4.8",
-    "retry~=0.9.2",
+    "tenacity~=8.5.0",
     "tiktoken~=0.7.0",
     "nltk~=3.9.1",
     "gradio~=5.20.0",
