From 547dee4b2747f28aa0a27d83f680d1faa4eb38e2 Mon Sep 17 00:00:00 2001
From: Ghraven <115199279+Ghraven@users.noreply.github.com>
Date: Sun, 3 May 2026 16:57:25 +0800
Subject: [PATCH] feat: add explicit __enter__ and __aenter__ to BaseClient
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Fixes #642

While `BaseClient` inherits from both `AbstractContextManager` and
`AbstractAsyncContextManager`, it only defined `__exit__` and `__aexit__`.
`AbstractContextManager` provides a default `__enter__` that returns `self`,
but `AbstractAsyncContextManager.__aenter__` is abstract and must be
implemented explicitly — meaning `async with AsyncClient() as c:` raised
`TypeError` at runtime.

This commit adds explicit `__enter__` and `__aenter__` methods that both
return `self`, completing the context manager protocol for both sync and
async clients and making `with Client() as c:` and `async with AsyncClient() as c:`
work as expected.
---
 ollama/_client.py | 10 ++++++++--
 1 file changed, 8 insertions(+), 2 deletions(-)

diff --git a/ollama/_client.py b/ollama/_client.py
index 18cb0fb4..7cf15dd6 100644
--- a/ollama/_client.py
+++ b/ollama/_client.py
@@ -117,10 +117,16 @@ def __init__(
       **kwargs,
     )
 
-  def __exit__(self, exc_type, exc_val, exc_tb):
+  def __enter__(self) -> "BaseClient":
+    return self
+
+  def __exit__(self, exc_type, exc_val, exc_tb) -> None:
     self.close()
 
-  async def __aexit__(self, exc_type, exc_val, exc_tb):
+  async def __aenter__(self) -> "BaseClient":
+    return self
+
+  async def __aexit__(self, exc_type, exc_val, exc_tb) -> None:
     await self.close()