ollama: Add separate kwargs parameter for async client (#31209)

**Description**:

Add a `async_client_kwargs` field to ollama chat/llm/embeddings adapters
that is passed to async httpx client constructor.

**Motivation:**

In my use-case:
- chat/embedding model adapters may be created frequently, sometimes to
be called just once or to never be called at all
- they may be used in bots sunc and async mode (not known at the moment
they are created)

So, I want to keep a static transport instance maintaining connection
pool, so model adapters can be created and destroyed freely. But that
doesn't work when both sync and async functions are in use as I can only
pass one transport instance for both sync and async client, while
transport types must be different for them. So I can't make both sync
and async calls use shared transport with current model adapter
interfaces.

In this PR I add a separate `async_client_kwargs` that gets passed to
async client constructor, so it will be possible to pass a separate
transport instance. For sake of backwards compatibility, it is merged
with `client_kwargs`, so nothing changes when it is not set.

I am unable to run linter right now, but the changes look ok.
This commit is contained in:
Alexey Bondarenko 2025-05-16 02:10:10 +06:00 committed by GitHub
parent 03adca6c44
commit 9efafe3337
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 81 additions and 12 deletions

View File

@ -436,8 +436,22 @@ class ChatOllama(BaseChatModel):
"""Base url the model is hosted under."""
client_kwargs: Optional[dict] = {}
"""Additional kwargs to pass to the httpx Client.
For a full list of the params, see [this link](https://pydoc.dev/httpx/latest/httpx.Client.html)
"""Additional kwargs to pass to the httpx clients.
These arguments are passed to both synchronous and async clients.
Use sync_client_kwargs and async_client_kwargs to pass different arguments
to synchronous and asynchronous clients.
"""
async_client_kwargs: Optional[dict] = {}
"""Additional kwargs to merge with client_kwargs before
passing to the httpx AsyncClient.
For a full list of the params, see [this link](https://www.python-httpx.org/api/#asyncclient)
"""
sync_client_kwargs: Optional[dict] = {}
"""Additional kwargs to merge with client_kwargs before
passing to the httpx Client.
For a full list of the params, see [this link](https://www.python-httpx.org/api/#client)
"""
_client: Client = PrivateAttr(default=None) # type: ignore
@ -503,8 +517,17 @@ class ChatOllama(BaseChatModel):
def _set_clients(self) -> Self:
"""Set clients to use for ollama."""
client_kwargs = self.client_kwargs or {}
self._client = Client(host=self.base_url, **client_kwargs)
self._async_client = AsyncClient(host=self.base_url, **client_kwargs)
sync_client_kwargs = client_kwargs
if self.sync_client_kwargs:
sync_client_kwargs = {**sync_client_kwargs, **self.sync_client_kwargs}
async_client_kwargs = client_kwargs
if self.async_client_kwargs:
async_client_kwargs = {**async_client_kwargs, **self.async_client_kwargs}
self._client = Client(host=self.base_url, **sync_client_kwargs)
self._async_client = AsyncClient(host=self.base_url, **async_client_kwargs)
return self
def _convert_messages_to_ollama_messages(

View File

@ -127,8 +127,22 @@ class OllamaEmbeddings(BaseModel, Embeddings):
"""Base url the model is hosted under."""
client_kwargs: Optional[dict] = {}
"""Additional kwargs to pass to the httpx Client.
For a full list of the params, see [this link](https://pydoc.dev/httpx/latest/httpx.Client.html)
"""Additional kwargs to pass to the httpx clients.
These arguments are passed to both synchronous and async clients.
Use sync_client_kwargs and async_client_kwargs to pass different arguments
to synchronous and asynchronous clients.
"""
async_client_kwargs: Optional[dict] = {}
"""Additional kwargs to merge with client_kwargs before
passing to the httpx AsyncClient.
For a full list of the params, see [this link](https://www.python-httpx.org/api/#asyncclient)
"""
sync_client_kwargs: Optional[dict] = {}
"""Additional kwargs to merge with client_kwargs before
passing to the httpx Client.
For a full list of the params, see [this link](https://www.python-httpx.org/api/#client)
"""
_client: Client = PrivateAttr(default=None) # type: ignore
@ -233,8 +247,17 @@ class OllamaEmbeddings(BaseModel, Embeddings):
def _set_clients(self) -> Self:
"""Set clients to use for ollama."""
client_kwargs = self.client_kwargs or {}
self._client = Client(host=self.base_url, **client_kwargs)
self._async_client = AsyncClient(host=self.base_url, **client_kwargs)
sync_client_kwargs = client_kwargs
if self.sync_client_kwargs:
sync_client_kwargs = {**sync_client_kwargs, **self.sync_client_kwargs}
async_client_kwargs = client_kwargs
if self.async_client_kwargs:
async_client_kwargs = {**async_client_kwargs, **self.async_client_kwargs}
self._client = Client(host=self.base_url, **sync_client_kwargs)
self._async_client = AsyncClient(host=self.base_url, **async_client_kwargs)
return self
def embed_documents(self, texts: list[str]) -> list[list[float]]:

View File

@ -113,8 +113,22 @@ class OllamaLLM(BaseLLM):
"""Base url the model is hosted under."""
client_kwargs: Optional[dict] = {}
"""Additional kwargs to pass to the httpx Client.
For a full list of the params, see [this link](https://pydoc.dev/httpx/latest/httpx.Client.html)
"""Additional kwargs to pass to the httpx clients.
These arguments are passed to both synchronous and async clients.
Use sync_client_kwargs and async_client_kwargs to pass different arguments
to synchronous and asynchronous clients.
"""
async_client_kwargs: Optional[dict] = {}
"""Additional kwargs to merge with client_kwargs before
passing to the httpx AsyncClient.
For a full list of the params, see [this link](https://www.python-httpx.org/api/#asyncclient)
"""
sync_client_kwargs: Optional[dict] = {}
"""Additional kwargs to merge with client_kwargs before
passing to the httpx Client.
For a full list of the params, see [this link](https://www.python-httpx.org/api/#client)
"""
_client: Client = PrivateAttr(default=None) # type: ignore
@ -189,8 +203,17 @@ class OllamaLLM(BaseLLM):
def _set_clients(self) -> Self:
"""Set clients to use for ollama."""
client_kwargs = self.client_kwargs or {}
self._client = Client(host=self.base_url, **client_kwargs)
self._async_client = AsyncClient(host=self.base_url, **client_kwargs)
sync_client_kwargs = client_kwargs
if self.sync_client_kwargs:
sync_client_kwargs = {**sync_client_kwargs, **self.sync_client_kwargs}
async_client_kwargs = client_kwargs
if self.async_client_kwargs:
async_client_kwargs = {**async_client_kwargs, **self.async_client_kwargs}
self._client = Client(host=self.base_url, **sync_client_kwargs)
self._async_client = AsyncClient(host=self.base_url, **async_client_kwargs)
return self
async def _acreate_generate_stream(