diff --git a/hugegraph-llm/pyproject.toml b/hugegraph-llm/pyproject.toml index cd2452d96..5327e8136 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 515d8d69d..bfffb689f 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 @@ 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,9 +96,9 @@ 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, @@ -92,7 +106,7 @@ def generate_streaming( 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 index bbc014fd2..39d97f78c 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 bb63371f6..a3d0bc375 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",