diff --git a/poetry.lock b/poetry.lock index ab83a15..25e0bbd 100644 --- a/poetry.lock +++ b/poetry.lock @@ -1060,6 +1060,17 @@ files = [ marshmallow = ">=3.18.0,<4.0.0" typing-inspect = ">=0.4.0,<1" +[[package]] +name = "decorator" +version = "5.1.1" +description = "Decorators for Humans" +optional = false +python-versions = ">=3.5" +files = [ + {file = "decorator-5.1.1-py3-none-any.whl", hash = "sha256:b8c3f85900b9dc423225913c5aace94729fe1fa9763b38939a95226f02d37186"}, + {file = "decorator-5.1.1.tar.gz", hash = "sha256:637996211036b6385ef91435e4fae22989472f9d571faba8927ba8253acbc330"}, +] + [[package]] name = "deprecated" version = "1.2.14" @@ -5015,6 +5026,20 @@ requests = ">=2.0.0" [package.extras] rsa = ["oauthlib[signedtoken] (>=3.0.0)"] +[[package]] +name = "retry-async" +version = "0.1.4" +description = "" +optional = false +python-versions = ">=3.10,<4.0" +files = [ + {file = "retry_async-0.1.4-py3-none-any.whl", hash = "sha256:21b383c7bc52013478337b894f476c9f106485cfeeb5d449abe5f745be2da219"}, + {file = "retry_async-0.1.4.tar.gz", hash = "sha256:8414d69b20920a1d700de34b68c0f972fa36a0158450a6f6abc5b45a241ac6b6"}, +] + +[package.dependencies] +decorator = ">=5.1.1,<6.0.0" + [[package]] name = "rich" version = "13.7.1" @@ -6737,4 +6762,4 @@ vector-stores-qdrant = ["llama-index-vector-stores-qdrant"] [metadata] lock-version = "2.0" python-versions = ">=3.11,<3.12" -content-hash = "119ebc4623e38c745677ca78dee1133087cce3b04ac210bc7a774e6c5cce26c6" +content-hash = "3fa6ef447847895b1a16b8b0422dd9e4fda1aaaadef3af71971eb412da89bf67" diff --git a/private_gpt/utils/ollama.py b/private_gpt/utils/ollama.py index 9c75a87..da9107b 100644 --- a/private_gpt/utils/ollama.py +++ b/private_gpt/utils/ollama.py @@ -3,10 +3,13 @@ from collections import deque 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`" @@ -14,13 +17,25 @@ except ImportError as 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 diff --git a/private_gpt/utils/retry.py b/private_gpt/utils/retry.py new file mode 100644 index 0000000..614b680 --- /dev/null +++ b/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 diff --git a/pyproject.toml b/pyproject.toml index 733690e..afbb83c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -67,6 +67,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"]