diff --git a/libs/partners/ollama/langchain_ollama/chat_models.py b/libs/partners/ollama/langchain_ollama/chat_models.py index a7046ef79ca..0bd1f9fd381 100644 --- a/libs/partners/ollama/langchain_ollama/chat_models.py +++ b/libs/partners/ollama/langchain_ollama/chat_models.py @@ -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( diff --git a/libs/partners/ollama/langchain_ollama/embeddings.py b/libs/partners/ollama/langchain_ollama/embeddings.py index c6327343d48..15643529cda 100644 --- a/libs/partners/ollama/langchain_ollama/embeddings.py +++ b/libs/partners/ollama/langchain_ollama/embeddings.py @@ -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]]: diff --git a/libs/partners/ollama/langchain_ollama/llms.py b/libs/partners/ollama/langchain_ollama/llms.py index 5298fec0e3c..4a40af3e99d 100644 --- a/libs/partners/ollama/langchain_ollama/llms.py +++ b/libs/partners/ollama/langchain_ollama/llms.py @@ -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(