client: expose resource cleanup methods

Author: ParthSareen

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)

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