test: WIP for E2E tests

anubhav756 · anubhav756 · commit 0fc894ba7353 · 2025-05-16T19:33:53.000+05:30
diff --git a/packages/toolbox-core/src/toolbox_core/sync_client.py b/packages/toolbox-core/src/toolbox_core/sync_client.py
@@ -20,6 +20,7 @@
 
 from .client import ToolboxClient
 from .sync_tool import ToolboxSyncTool
+from concurrent.futures import Future
 
 
 class ToolboxSyncClient:
@@ -85,6 +86,46 @@ def close(self):
         coro = self.__async_client.close()
         asyncio.run_coroutine_threadsafe(coro, self.__loop).result()
 
+    def load_tool_future(
+        self,
+        name: str,
+        auth_token_getters: dict[str, Callable[[], str]] = {},
+        bound_params: Mapping[str, Union[Callable[[], Any], Any]] = {},
+    ) -> Future[ToolboxSyncTool]:
+        """
+        Returns a future that loads a tool from the server.
+        """
+        if not self.__loop or not self.__thread:
+            raise ValueError("Background loop or thread cannot be None.")
+
+        async def async_worker() -> ToolboxSyncTool:
+            async_tool = self.__async_client.load_tool(name, auth_token_getters, bound_params)
+            return ToolboxSyncTool(async_tool, self.__loop, self.__thread)
+        return asyncio.run_coroutine_threadsafe(async_worker(), self.__loop)
+
+    def load_toolset_future(
+        self,
+        name: Optional[str] = None,
+        auth_token_getters: dict[str, Callable[[], str]] = {},
+        bound_params: Mapping[str, Union[Callable[[], Any], Any]] = {},
+        strict: bool = False,
+    ) -> Future[list[ToolboxSyncTool]]:
+        """
+        Returns a future that fetches a toolset and loads all tools defined within it.
+        """
+        if not self.__loop or not self.__thread:
+            raise ValueError("Background loop or thread cannot be None.")
+
+        async def async_worker() -> list[ToolboxSyncTool]:
+            async_tools = await self.__async_client.load_toolset(
+                name, auth_token_getters, bound_params, strict
+            )
+            return [
+                ToolboxSyncTool(async_tool, self.__loop, self.__thread)
+                for async_tool in async_tools
+            ]
+        return asyncio.run_coroutine_threadsafe(async_worker(), self.__loop)
+
     def load_tool(
         self,
         name: str,
@@ -110,13 +151,7 @@ def load_tool(
                 for execution. The specific arguments and behavior of the callable
                 depend on the tool itself.
         """
-        coro = self.__async_client.load_tool(name, auth_token_getters, bound_params)
-
-        if not self.__loop or not self.__thread:
-            raise ValueError("Background loop or thread cannot be None.")
-
-        async_tool = asyncio.run_coroutine_threadsafe(coro, self.__loop).result()
-        return ToolboxSyncTool(async_tool, self.__loop, self.__thread)
+        return self.load_tool_future(name, auth_token_getters, bound_params).result()
 
     def load_toolset(
         self,
@@ -147,18 +182,7 @@ def load_toolset(
         Raises:
             ValueError: If validation fails based on the `strict` flag.
         """
-        coro = self.__async_client.load_toolset(
-            name, auth_token_getters, bound_params, strict
-        )
-
-        if not self.__loop or not self.__thread:
-            raise ValueError("Background loop or thread cannot be None.")
-
-        async_tools = asyncio.run_coroutine_threadsafe(coro, self.__loop).result()
-        return [
-            ToolboxSyncTool(async_tool, self.__loop, self.__thread)
-            for async_tool in async_tools
-        ]
+        return self.load_toolset_future(name, auth_token_getters, bound_params, strict).result()
 
     def add_headers(
         self, headers: Mapping[str, Union[Callable, Coroutine, str]]
diff --git a/packages/toolbox-core/src/toolbox_core/sync_tool.py b/packages/toolbox-core/src/toolbox_core/sync_tool.py
@@ -21,6 +21,7 @@
 
 from .protocol import ParameterSchema
 from .tool import ToolboxTool
+from concurrent.futures import Future
 
 
 class ToolboxSyncTool:
@@ -129,6 +130,13 @@ def _auth_service_token_getters(self) -> Mapping[str, Callable[[], str]]:
     def _client_headers(self) -> Mapping[str, Union[Callable, Coroutine, str]]:
         return self.__async_tool._client_headers
 
+    def call_future(self, *args: Any, **kwargs: Any) -> Future[str]:
+        """
+        Returns future that calls the remote tool with the provided arguments.
+        """
+        coro = self.__async_tool(*args, **kwargs)
+        return asyncio.run_coroutine_threadsafe(coro, self.__loop)
+
     def __call__(self, *args: Any, **kwargs: Any) -> str:
         """
         Synchronously calls the remote tool with the provided arguments.
@@ -143,8 +151,7 @@ def __call__(self, *args: Any, **kwargs: Any) -> str:
         Returns:
             The string result returned by the remote tool execution.
         """
-        coro = self.__async_tool(*args, **kwargs)
-        return asyncio.run_coroutine_threadsafe(coro, self.__loop).result()
+        return self.call_future(*args, **kwargs).result()
 
     def add_auth_token_getters(
         self,
diff --git a/packages/toolbox-langchain/src/toolbox_langchain/client.py b/packages/toolbox-langchain/src/toolbox_langchain/client.py
@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import asyncio
+from asyncio import wrap_future
 from typing import Any, Callable, Optional, Union
 from warnings import warn
 
@@ -85,30 +85,12 @@ async def aload_tool(
                 )
                 auth_token_getters = auth_headers
 
-        coro = self.__core_client._async_client.load_tool(
+        core_tool = await wrap_future(self.__core_client.aload_tool(
             name=tool_name,
             auth_token_getters=auth_token_getters,
             bound_params=bound_params,
-        )
-
-        if not self.__core_client._loop:
-            # If a loop has not been provided, attempt to run in current thread.
-            core_tool = await coro
-        else:
-            # Otherwise, run in the background thread.
-            core_tool = await asyncio.wrap_future(
-                asyncio.run_coroutine_threadsafe(coro, self.__core_client._loop)
-            )
-
-        if not self.__core_client._loop or not self.__core_client._thread:
-            raise ValueError("Background loop or thread cannot be None.")
-
-        core_sync_tool = ToolboxSyncTool(
-            core_tool,
-            self.__core_client._loop,
-            self.__core_client._thread,
-        )
-        return ToolboxTool(core_tool=core_sync_tool)
+        ))
+        return ToolboxTool(core_tool=core_tool)
 
     async def aload_toolset(
         self,
@@ -167,36 +149,16 @@ async def aload_toolset(
                 )
                 auth_token_getters = auth_headers
 
-        coro = self.__core_client._async_client.load_toolset(
+        core_tools = await wrap_future(self.__core_client.aload_toolset(
             name=toolset_name,
             auth_token_getters=auth_token_getters,
             bound_params=bound_params,
             strict=strict,
-        )
-
-        if not self.__core_client._loop:
-            # If a loop has not been provided, attempt to run in current thread.
-            core_tools = await coro
-        else:
-            # Otherwise, run in the background thread.
-            core_tools = await asyncio.wrap_future(
-                asyncio.run_coroutine_threadsafe(coro, self.__core_client._loop)
-            )
-
-        if not self.__core_client._loop or not self.__core_client._thread:
-            raise ValueError("Background loop or thread cannot be None.")
+        ))
 
-        core_sync_tools = [
-            ToolboxSyncTool(
-                core_tool,
-                self.__core_client._loop,
-                self.__core_client._thread,
-            )
-            for core_tool in core_tools
-        ]
         tools = []
-        for core_sync_tool in core_sync_tools:
-            tools.append(ToolboxTool(core_tool=core_sync_tool))
+        for core_tool in core_tools:
+            tools.append(ToolboxTool(core_tool=core_tool))
         return tools
 
     def load_tool(
diff --git a/packages/toolbox-langchain/src/toolbox_langchain/tools.py b/packages/toolbox-langchain/src/toolbox_langchain/tools.py
@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import asyncio
+from asyncio import wrap_future
 from typing import Any, Callable, Union
 
 from deprecated import deprecated
@@ -50,16 +50,7 @@ def _run(self, **kwargs: Any) -> str:
         return self.__core_tool(**kwargs)
 
     async def _arun(self, **kwargs: Any) -> str:
-        coro = self.__core_tool._async_tool(**kwargs)
-
-        # If a loop has not been provided, attempt to run in current thread.
-        if not self.__core_tool._loop:
-            return await coro
-
-        # Otherwise, run in the background thread.
-        return await asyncio.wrap_future(
-            asyncio.run_coroutine_threadsafe(coro, self.__core_tool._loop)
-        )
+        return await wrap_future(self.__core_tool._async_tool.arun(**kwargs))
 
     def add_auth_token_getters(
         self, auth_token_getters: dict[str, Callable[[], str]]
