diff --git a/ollama/_client.py b/ollama/_client.py index 18cb0fb4..afaeb2ac 100644 --- a/ollama/_client.py +++ b/ollama/_client.py @@ -1,4 +1,5 @@ import contextlib +import inspect import ipaddress import json import os @@ -93,8 +94,18 @@ def __init__( - `follow_redirects`: True - `timeout`: None `kwargs` are passed to the httpx client. + + Args: + client: Either a httpx.Client/AsyncClient class, or an instance """ + if not inspect.isclass(client): + assert follow_redirects is True, 'Cannot provide `follow_redirects` with custom client instance' + assert timeout is None, 'Cannot provide `timeout` with custom client instance' + assert not kwargs, 'Cannot provide additional kwargs with custom client instance' + self._client = client + return + headers = { k.lower(): v for k, v in { @@ -128,8 +139,8 @@ async def __aexit__(self, exc_type, exc_val, exc_tb): class Client(BaseClient): - def __init__(self, host: Optional[str] = None, **kwargs) -> None: - super().__init__(httpx.Client, host, **kwargs) + def __init__(self, host: Optional[str] = None, *, client: Optional[httpx.Client] = None, **kwargs) -> None: + super().__init__(client or httpx.Client, host, **kwargs) def close(self): self._client.close() @@ -721,8 +732,8 @@ def web_fetch(self, url: str) -> WebFetchResponse: class AsyncClient(BaseClient): - def __init__(self, host: Optional[str] = None, **kwargs) -> None: - super().__init__(httpx.AsyncClient, host, **kwargs) + def __init__(self, host: Optional[str] = None, *, client: Optional[httpx.AsyncClient] = None, **kwargs) -> None: + super().__init__(client or httpx.AsyncClient, host, **kwargs) async def close(self): await self._client.aclose()