mirror of
https://github.com/ollama/ollama-python.git
synced 2026-05-01 11:48:17 +08:00
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
https://github.com/ollama/ollama-python/pull/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.
This commit is contained in:
parent
dbccf192ac
commit
9cc01492ef
@ -1,4 +1,5 @@
|
||||
import contextlib
|
||||
import inspect
|
||||
import ipaddress
|
||||
import json
|
||||
import os
|
||||
@ -93,8 +94,18 @@ class BaseClient(contextlib.AbstractContextManager, contextlib.AbstractAsyncCont
|
||||
- `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 @@ CONNECTION_ERROR_MESSAGE = 'Failed to connect to Ollama. Please check that Ollam
|
||||
|
||||
|
||||
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 @@ class Client(BaseClient):
|
||||
|
||||
|
||||
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()
|
||||
|
||||
Loading…
Reference in New Issue
Block a user