feat(toolbox-llamaindex): Add context manager support for sync and async clients

Author: anubhav756

packages/toolbox-llamaindex/src/toolbox_llamaindex/async_client.py

@@ -190,3 +190,28 @@ def load_toolset(
         strict: bool = False,
     ) -> list[AsyncToolboxTool]:
         raise NotImplementedError("Synchronous methods not supported by async client.")
+
+    async def close(self):
+        """Close the underlying synchronous client."""
+        await self.__core_client.close()
+
+    async def __aenter__(self):
+        """
+        Enter the runtime context related to this client instance.
+
+        Allows the client to be used as an asynchronous context manager
+        (e.g., `async with AsyncToolboxClient(...) as client:`).
+
+        Returns:
+            self: The client instance itself.
+        """
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        """
+        Exit the runtime context and close the internally managed session.
+
+        Allows the client to be used as an asynchronous context manager
+        (e.g., `async with AsyncToolboxClient(...) as client:`).
+        """
+        await self.close()

packages/toolbox-llamaindex/src/toolbox_llamaindex/client.py

@@ -292,3 +292,49 @@ def load_toolset(
         for core_sync_tool in core_sync_tools:
             tools.append(ToolboxTool(core_tool=core_sync_tool))
         return tools
+
+    def close(self):
+        """Close the underlying synchronous client."""
+        self.__core_client.close()
+
+    async def __aenter__(self):
+        """
+        Enter the runtime context related to this client instance.
+
+        Allows the client to be used as an asynchronous context manager
+        (e.g., `async with ToolboxClient(...) as client:`).
+
+        Returns:
+            self: The client instance itself.
+        """
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        """
+        Exit the runtime context and close the internally managed session.
+
+        Allows the client to be used as an asynchronous context manager
+        (e.g., `async with ToolboxClient(...) as client:`).
+        """
+        self.close()
+
+    def __enter__(self):
+        """
+        Enter the runtime context related to this client instance.
+
+        Allows the client to be used as a context manager
+        (e.g., `with ToolboxClient(...) as client:`).
+
+        Returns:
+            self: The client instance itself.
+        """
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        """
+        Exit the runtime context and close the internally managed session.
+
+        Allows the client to be used as a context manager
+        (e.g., `with ToolboxClient(...) as client:`).
+        """
+        self.close()