Add support for context manager (#435)

* Add support for context manager

* Fix

Author: billytrend-cohere

src/cohere/client.py

@@ -87,6 +87,14 @@ def __init__(
 
         validate_args(self, "chat", throw_if_stream_is_true)
 
+    # support context manager until Fern upstreams
+    # https://linear.app/buildwithfern/issue/FER-1242/expose-a-context-manager-interface-or-the-http-client-easily
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        self._client_wrapper.httpx_client.httpx_client.close()
+
     wait = wait
 
     """
@@ -161,6 +169,14 @@ def __init__(
 
         validate_args(self, "chat", throw_if_stream_is_true)
 
+    # support context manager until Fern upstreams
+    # https://linear.app/buildwithfern/issue/FER-1242/expose-a-context-manager-interface-or-the-http-client-easily
+    async def __aenter__(self):
+        return self
+
+    async def __aexit__(self, exc_type, exc_value, traceback):
+        await self._client_wrapper.httpx_client.httpx_client.aclose()
+
     wait = async_wait
 
     """

tests/test_async_client.py

@@ -19,6 +19,10 @@ async def test_token_falls_back_on_env_variable(self) -> None:
         cohere.AsyncClient(api_key=None)
         cohere.AsyncClient(None)
 
+    async def test_context_manager(self) -> None:
+        async with cohere.AsyncClient(api_key="xxx") as client:
+            self.assertIsNotNone(client)
+
     async def test_chat(self) -> None:
         chat = await self.co.chat(
             chat_history=[

tests/test_client.py

@@ -17,6 +17,10 @@ def test_token_falls_back_on_env_variable(self) -> None:
         cohere.Client(api_key=None)
         cohere.Client(None)
 
+    def test_context_manager(self) -> None:
+        with cohere.Client(api_key="xxx") as client:
+            self.assertIsNotNone(client)
+
     def test_chat(self) -> None:
         chat = co.chat(
             chat_history=[