This is an automated email from the ASF dual-hosted git repository.
imbajin pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/hugegraph-ai.git
The following commit(s) were added to refs/heads/main by this push:
new 5f4b54ac fix(llm): replace retry with tenacity in ollama.py (#367)
5f4b54ac is described below
commit 5f4b54ac3c6aaa60e124c8e5b4c76fbd744d66d4
Author: KAI <[email protected]>
AuthorDate: Tue Jun 23 10:03:29 2026 +0530
fix(llm): replace retry with tenacity in ollama.py (#367)
## 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 #365
---
hugegraph-llm/pyproject.toml | 2 +-
.../src/hugegraph_llm/models/llms/ollama.py | 48 +++--
.../src/tests/models/llms/test_ollama_client.py | 193 ++++++++++++++++++++-
pyproject.toml | 2 +-
4 files changed, 222 insertions(+), 23 deletions(-)
diff --git a/hugegraph-llm/pyproject.toml b/hugegraph-llm/pyproject.toml
index cd2452d9..5327e813 100644
--- 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
index 515d8d69..bfffb689 100644
--- 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 @@ class OllamaClient(BaseLLM):
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 @@ class OllamaClient(BaseLLM):
}
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,9 +96,9 @@ class OllamaClient(BaseLLM):
}
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,
@@ -92,7 +106,7 @@ class OllamaClient(BaseLLM):
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 @@ class OllamaClient(BaseLLM):
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 @@ class OllamaClient(BaseLLM):
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
index bbc014fd..39d97f78 100644
--- 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
index bb63371f..a3d0bc37 100644
--- 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",
