From 2804a03d82af8cc86c29c076af2240d48d3943c8 Mon Sep 17 00:00:00 2001 From: Michael Yang Date: Wed, 10 Jan 2024 09:49:09 -0800 Subject: [PATCH] httpx: kwargs --- ollama/_client.py | 25 ++++++++++++++++++------- 1 file changed, 18 insertions(+), 7 deletions(-) diff --git a/ollama/_client.py b/ollama/_client.py index cc444f4..09b4a47 100644 --- a/ollama/_client.py +++ b/ollama/_client.py @@ -20,14 +20,25 @@ from ollama._types import Message, Options class BaseClient: - def __init__(self, client, base_url: Optional[str] = None) -> None: - base_url = base_url or os.getenv('OLLAMA_HOST', 'http://127.0.0.1:11434') - self._client = client(base_url=base_url, follow_redirects=True, timeout=None) + def __init__( + self, + client, + base_url: Optional[str] = None, + follow_redirects: bool = True, + timeout: Any = None, + **kwargs, + ) -> None: + self._client = client( + base_url=base_url or os.getenv('OLLAMA_HOST', 'http://127.0.0.1:11434'), + follow_redirects=follow_redirects, + timeout=timeout, + **kwargs, + ) class Client(BaseClient): - def __init__(self, base_url: Optional[str] = None) -> None: - super().__init__(httpx.Client, base_url) + def __init__(self, base_url: Optional[str] = None, **kwargs) -> None: + super().__init__(httpx.Client, base_url, **kwargs) def _request(self, method: str, url: str, **kwargs) -> httpx.Response: response = self._client.request(method, url, **kwargs) @@ -247,8 +258,8 @@ class Client(BaseClient): class AsyncClient(BaseClient): - def __init__(self, base_url: Optional[str] = None) -> None: - super().__init__(httpx.AsyncClient, base_url) + def __init__(self, base_url: Optional[str] = None, **kwargs) -> None: + super().__init__(httpx.AsyncClient, base_url, **kwargs) async def _request(self, method: str, url: str, **kwargs) -> httpx.Response: response = await self._client.request(method, url, **kwargs)