From b0ba4258aa5564ab42bcd0b5db535eab244500c3 Mon Sep 17 00:00:00 2001
From: ParthSareen <parth.sareen@ollama.com>
Date: Thu, 13 Feb 2025 12:08:27 -0800
Subject: [PATCH 1/3] client: expose resource cleanup methods

---
 ollama/_client.py    | 18 ++++++++++++++++++
 tests/test_client.py | 30 ++++++++++++++++++++++++++++++
 2 files changed, 48 insertions(+)

diff --git a/ollama/_client.py b/ollama/_client.py
index cbe43c94..2b0b8b82 100644
--- a/ollama/_client.py
+++ b/ollama/_client.py
@@ -113,6 +113,15 @@ class Client(BaseClient):
   def __init__(self, host: Optional[str] = None, **kwargs) -> None:
     super().__init__(httpx.Client, host, **kwargs)
 
+  def close(self):
+    self._client.close()
+
+  def __enter__(self):
+    return self
+
+  def __exit__(self, exc_type, exc_val, exc_tb):
+    self.close()
+
   def _request_raw(self, *args, **kwargs):
     try:
       r = self._client.request(*args, **kwargs)
@@ -617,6 +626,15 @@ class AsyncClient(BaseClient):
   def __init__(self, host: Optional[str] = None, **kwargs) -> None:
     super().__init__(httpx.AsyncClient, host, **kwargs)
 
+  async def close(self):
+    await self._client.aclose()
+
+  async def __aenter__(self):
+    return self
+
+  async def __aexit__(self, exc_type, exc_val, exc_tb):
+    await self.close()
+
   async def _request_raw(self, *args, **kwargs):
     try:
       r = await self._client.request(*args, **kwargs)
diff --git a/tests/test_client.py b/tests/test_client.py
index dacb953d..c93060aa 100644
--- a/tests/test_client.py
+++ b/tests/test_client.py
@@ -1140,3 +1140,33 @@ async def test_async_client_connection_error():
   with pytest.raises(ConnectionError) as exc_info:
     await client.show('model')
   assert str(exc_info.value) == 'Failed to connect to Ollama. Please check that Ollama is downloaded, running and accessible. https://ollama.com/download'
+
+
+def test_client_close():
+  client = Client()
+  client.close()
+  assert client._client.is_closed
+
+
+@pytest.mark.asyncio
+async def test_async_client_close():
+  client = AsyncClient()
+  await client.close()
+  assert client._client.is_closed
+
+
+def test_client_context_manager():
+  with Client() as client:
+    assert isinstance(client, Client)
+    assert not client._client.is_closed
+
+  assert client._client.is_closed
+
+
+@pytest.mark.asyncio
+async def test_async_client_context_manager():
+  async with AsyncClient() as client:
+    assert isinstance(client, AsyncClient)
+    assert not client._client.is_closed
+
+  assert client._client.is_closed

From 1a047b0670880b70e631d5a7060f18522918324c Mon Sep 17 00:00:00 2001
From: ParthSareen <parth.sareen@ollama.com>
Date: Thu, 13 Feb 2025 22:27:37 -0800
Subject: [PATCH 2/3] client: use contextlib for resource management

---
 ollama/_client.py | 24 +++++++++++-------------
 1 file changed, 11 insertions(+), 13 deletions(-)

diff --git a/ollama/_client.py b/ollama/_client.py
index 2b0b8b82..ce43c341 100644
--- a/ollama/_client.py
+++ b/ollama/_client.py
@@ -1,3 +1,4 @@
+import contextlib
 import ipaddress
 import json
 import os
@@ -70,7 +71,7 @@
 T = TypeVar('T')
 
 
-class BaseClient:
+class BaseClient(contextlib.AbstractContextManager):
   def __init__(
     self,
     client,
@@ -105,6 +106,15 @@ def __init__(
       **kwargs,
     )
 
+  def __exit__(self, exc_type, exc_val, exc_tb):
+    self.close()
+
+  async def __aenter__(self) -> Any:
+    return self
+
+  async def __aexit__(self, exc_type, exc_val, exc_tb):
+    await self.close()
+
 
 CONNECTION_ERROR_MESSAGE = 'Failed to connect to Ollama. Please check that Ollama is downloaded, running and accessible. https://ollama.com/download'
 
@@ -116,12 +126,6 @@ def __init__(self, host: Optional[str] = None, **kwargs) -> None:
   def close(self):
     self._client.close()
 
-  def __enter__(self):
-    return self
-
-  def __exit__(self, exc_type, exc_val, exc_tb):
-    self.close()
-
   def _request_raw(self, *args, **kwargs):
     try:
       r = self._client.request(*args, **kwargs)
@@ -629,12 +633,6 @@ def __init__(self, host: Optional[str] = None, **kwargs) -> None:
   async def close(self):
     await self._client.aclose()
 
-  async def __aenter__(self):
-    return self
-
-  async def __aexit__(self, exc_type, exc_val, exc_tb):
-    await self.close()
-
   async def _request_raw(self, *args, **kwargs):
     try:
       r = await self._client.request(*args, **kwargs)

From 71ffad6579158302eeeec03cd8ca12aaff8a7613 Mon Sep 17 00:00:00 2001
From: ParthSareen <parth.sareen@ollama.com>
Date: Tue, 18 Feb 2025 15:48:39 -0800
Subject: [PATCH 3/3] add AbstractAsyncContextManager to BaseClient

---
 ollama/_client.py | 5 +----
 1 file changed, 1 insertion(+), 4 deletions(-)

diff --git a/ollama/_client.py b/ollama/_client.py
index ce43c341..839cc253 100644
--- a/ollama/_client.py
+++ b/ollama/_client.py
@@ -71,7 +71,7 @@
 T = TypeVar('T')
 
 
-class BaseClient(contextlib.AbstractContextManager):
+class BaseClient(contextlib.AbstractContextManager, contextlib.AbstractAsyncContextManager):
   def __init__(
     self,
     client,
@@ -109,9 +109,6 @@ def __init__(
   def __exit__(self, exc_type, exc_val, exc_tb):
     self.close()
 
-  async def __aenter__(self) -> Any:
-    return self
-
   async def __aexit__(self, exc_type, exc_val, exc_tb):
     await self.close()