Skip to content

Commit 9cc0149

Browse files
committed
Add support for custom httpx client instances
Allow passing a pre-configured `httpx.Client` or `httpx.AsyncClient` instance to reuse connections and custom configurations. ```python import httpx from ollama import Client custom_httpx_client = httpx.Client(timeout=30.0) client = Client(client=custom_httpx_client) messages = [ { 'role': 'user', 'content': 'Why is the sky blue?', }, ] for part in client.chat('gpt-oss:120b-cloud', messages=messages, stream=True): print(part.message.content, end='', flush=True) ``` Note: this is an updated minimal version of ollama#380 to support a custom `httpx` client. I took the approach of checking for the class instead of isinstance so that nothing has to be a direct instance of an httpx.Client or httpx.AsyncClient. It's up to the user to provide the right custom client.
1 parent dbccf19 commit 9cc0149

1 file changed

Lines changed: 15 additions & 4 deletions

File tree

ollama/_client.py

Lines changed: 15 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
import contextlib
2+
import inspect
23
import ipaddress
34
import json
45
import os
@@ -93,8 +94,18 @@ def __init__(
9394
- `follow_redirects`: True
9495
- `timeout`: None
9596
`kwargs` are passed to the httpx client.
97+
98+
Args:
99+
client: Either a httpx.Client/AsyncClient class, or an instance
96100
"""
97101

102+
if not inspect.isclass(client):
103+
assert follow_redirects is True, 'Cannot provide `follow_redirects` with custom client instance'
104+
assert timeout is None, 'Cannot provide `timeout` with custom client instance'
105+
assert not kwargs, 'Cannot provide additional kwargs with custom client instance'
106+
self._client = client
107+
return
108+
98109
headers = {
99110
k.lower(): v
100111
for k, v in {
@@ -128,8 +139,8 @@ async def __aexit__(self, exc_type, exc_val, exc_tb):
128139

129140

130141
class Client(BaseClient):
131-
def __init__(self, host: Optional[str] = None, **kwargs) -> None:
132-
super().__init__(httpx.Client, host, **kwargs)
142+
def __init__(self, host: Optional[str] = None, *, client: Optional[httpx.Client] = None, **kwargs) -> None:
143+
super().__init__(client or httpx.Client, host, **kwargs)
133144

134145
def close(self):
135146
self._client.close()
@@ -721,8 +732,8 @@ def web_fetch(self, url: str) -> WebFetchResponse:
721732

722733

723734
class AsyncClient(BaseClient):
724-
def __init__(self, host: Optional[str] = None, **kwargs) -> None:
725-
super().__init__(httpx.AsyncClient, host, **kwargs)
735+
def __init__(self, host: Optional[str] = None, *, client: Optional[httpx.AsyncClient] = None, **kwargs) -> None:
736+
super().__init__(client or httpx.AsyncClient, host, **kwargs)
726737

727738
async def close(self):
728739
await self._client.aclose()

0 commit comments

Comments
 (0)