feat: add retry connection to ollama

private_gpt/utils/ollama.py

@@ -3,24 +3,39 @@
 from collections.abc import Iterator, Mapping
 from typing import Any
 
+from httpx import ConnectError
 from tqdm import tqdm  # type: ignore
 
+from private_gpt.utils.retry import retry
+
 try:
-    from ollama import Client  # type: ignore
+    from ollama import Client, ResponseError  # type: ignore
 except ImportError as e:
     raise ImportError(
         "Ollama dependencies not found, install with `poetry install --extras llms-ollama or embeddings-ollama`"
     ) from e
 
 logger = logging.getLogger(__name__)
 
+_MAX_RETRIES = 5
+_JITTER = (3.0, 10.0)
+
 
+@retry(
+    is_async=False,
+    exceptions=(ConnectError, ResponseError),
+    tries=_MAX_RETRIES,
+    jitter=_JITTER,
+    logger=logger,
+)
 def check_connection(client: Client) -> bool:
     try:
         client.list()
         return True
+    except (ConnectError, ResponseError) as e:
+        raise e
     except Exception as e:
-        logger.error(f"Failed to connect to Ollama: {e!s}")
+        logger.error(f"Failed to connect to Ollama: {type(e).__name__}: {e!s}")
         return False
 
 

private_gpt/utils/retry.py

@@ -0,0 +1,31 @@
+import logging
+from collections.abc import Callable
+from typing import Any
+
+from retry_async import retry as retry_untyped  # type: ignore
+
+retry_logger = logging.getLogger(__name__)
+
+
+def retry(
+    exceptions: Any = Exception,
+    *,
+    is_async: bool = False,
+    tries: int = -1,
+    delay: float = 0,
+    max_delay: float | None = None,
+    backoff: float = 1,
+    jitter: float | tuple[float, float] = 0,
+    logger: logging.Logger = retry_logger,
+) -> Callable[..., Any]:
+    wrapped = retry_untyped(
+        exceptions=exceptions,
+        is_async=is_async,
+        tries=tries,
+        delay=delay,
+        max_delay=max_delay,
+        backoff=backoff,
+        jitter=jitter,
+        logger=logger,
+    )
+    return wrapped  # type: ignore

pyproject.toml

@@ -66,6 +66,7 @@ ollama = {version ="^0.3.0", optional = true}
 
 # Optional HF Transformers
 einops = {version = "^0.8.0", optional = true}
+retry-async = "^0.1.4"
 
 [tool.poetry.extras]
 ui = ["gradio", "ffmpy"]