fix: Base toolbox sync & async tools to toolbox core counterparts

Author: anubhav756

packages/toolbox-langchain/src/toolbox_langchain/async_client.py

@@ -17,8 +17,7 @@
 
 from aiohttp import ClientSession
 
-from .tools import AsyncToolboxTool
-
+from .async_tools import AsyncToolboxTool
 from toolbox_core.client import ToolboxClient as ToolboxCoreClient
 
 

packages/toolbox-langchain/src/toolbox_langchain/async_tools.py

@@ -12,16 +12,12 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from copy import deepcopy
-from typing import Any, Callable, TypeVar, Union
-from warnings import warn
+from typing import Any, Callable, Union
 
 from langchain_core.tools import BaseTool
 from toolbox_core.tool import ToolboxTool as ToolboxCoreTool
 
 
-T = TypeVar("T")
-
 
 # This class is an internal implementation detail and is not exposed to the
 # end-user. It should not be used directly by external code. Changes to this
@@ -40,14 +36,7 @@ def __init__(
         Initializes an AsyncToolboxTool instance.
 
         Args:
-            name: The name of the tool.
-            schema: The tool schema.
-            url: The base URL of the Toolbox service.
-            session: The HTTP client session.
-            auth_token_getters: A mapping of authentication source names to
-                functions that retrieve ID tokens.
-            bound_params: A mapping of parameter names to their bound
-                values.
+            core_tool: The underlying core async ToolboxTool instance.
         """
 
         self.__core_tool = core_tool
@@ -88,7 +77,7 @@ def add_auth_token_getters(
 
         Returns:
             A new AsyncToolboxTool instance that is a deep copy of the current
-            instance, with added auth tokens.
+            instance, with added auth token getters.
 
         Raises:
             ValueError: If any of the provided auth parameters is already

packages/toolbox-langchain/src/toolbox_langchain/client.py

@@ -95,11 +95,11 @@ async def aload_tool(
             return await coro
 
         # Otherwise, run in the background thread.
-        core_tool = await asyncio.wrap_future(
+        core_sync_tool = await asyncio.wrap_future(
             asyncio.run_coroutine_threadsafe(coro, self.__core_sync_client._ToolboxSyncClient__loop)
         )
 
-        return ToolboxTool(core_tool=core_tool)
+        return ToolboxTool(core_sync_tool=core_sync_tool)
 
     async def aload_toolset(
         self,
@@ -170,13 +170,13 @@ async def aload_toolset(
             return await coro
 
         # Otherwise, run in the background thread.
-        core_tools = await asyncio.wrap_future(
+        core_sync_tools = await asyncio.wrap_future(
             asyncio.run_coroutine_threadsafe(coro, self.__core_sync_client._ToolboxSyncClient__loop)
         )
 
         tools = []
-        for core_tool in core_tools:
-            tools.append(ToolboxTool(core_tool=core_tool))
+        for core_sync_tool in core_sync_tools:
+            tools.append(ToolboxTool(core_sync_tool=core_sync_tool))
         return tools
 
     def load_tool(
@@ -228,12 +228,12 @@ def load_tool(
                 )
                 auth_token_getters = auth_tokens
 
-        core_tool = self.__core_sync_client.load_tool(
+        core_sync_tool = self.__core_sync_client.load_tool(
             name=tool_name,
             auth_token_getters=auth_token_getters,
             bound_params=bound_params
         )
-        return ToolboxTool(core_tool=core_tool)
+        return ToolboxTool(core_sync_tool=core_sync_tool)
 
     def load_toolset(
         self,
@@ -292,14 +292,14 @@ def load_toolset(
                 )
                 auth_token_getters = auth_tokens
 
-        core_tools = self.__core_sync_client.load_toolset(
+        core_sync_tools = self.__core_sync_client.load_toolset(
             name=toolset_name,
             auth_token_getters=auth_token_getters,
             bound_params=bound_params,
             strict=strict
         )
 
         tools = []
-        for core_tool in core_tools:
-            tools.append(ToolboxTool(core_tool=core_tool))
+        for core_sync_tool in core_sync_tools:
+            tools.append(ToolboxTool(core_sync_tool=core_sync_tool))
         return tools

packages/toolbox-langchain/src/toolbox_langchain/tools.py

@@ -13,15 +13,11 @@
 # limitations under the License.
 
 import asyncio
-from asyncio import AbstractEventLoop
-from threading import Thread
-from typing import Any, Awaitable, Callable, TypeVar, Union
+from typing import Any, Callable, Union
 
 from langchain_core.tools import BaseTool
+from toolbox_core.sync_tool import ToolboxSyncTool as ToolboxCoreSyncTool
 
-from .async_tools import AsyncToolboxTool
-
-T = TypeVar("T")
 
 
 class ToolboxTool(BaseTool):
@@ -32,56 +28,37 @@ class ToolboxTool(BaseTool):
 
     def __init__(
         self,
-        async_tool: AsyncToolboxTool,
-        loop: AbstractEventLoop,
-        thread: Thread,
+        core_sync_tool: ToolboxCoreSyncTool,
     ) -> None:
         """
         Initializes a ToolboxTool instance.
 
         Args:
-            async_tool: The underlying AsyncToolboxTool instance.
-            loop: The event loop used to run asynchronous tasks.
-            thread: The thread to run blocking operations in.
+            core_sync_tool: The underlying core sync ToolboxTool instance.
         """
 
-        # Due to how pydantic works, we must initialize the underlying
-        # BaseTool class before assigning values to member variables.
+        self.__core_sync_tool = core_sync_tool
         super().__init__(
-            name=async_tool.name,
-            description=async_tool.description,
-            args_schema=async_tool.args_schema,
+            name=self.__core_sync_tool.__name__,
+            description=self.__core_sync_tool.__doc__,
+            args_schema=self.__core_sync_tool._ToolboxSyncTool__pydantic_model,
         )
 
-        self.__async_tool = async_tool
-        self.__loop = loop
-        self.__thread = thread
-
-    def __run_as_sync(self, coro: Awaitable[T]) -> T:
-        """Run an async coroutine synchronously"""
-        if not self.__loop:
-            raise Exception(
-                "Cannot call synchronous methods before the background loop is initialized."
-            )
-        return asyncio.run_coroutine_threadsafe(coro, self.__loop).result()
+    def _run(self, **kwargs: Any) -> dict[str, Any]:
+        return self.__core_sync_tool(**kwargs)
 
-    async def __run_as_async(self, coro: Awaitable[T]) -> T:
-        """Run an async coroutine asynchronously"""
+    async def _arun(self, **kwargs: Any) -> dict[str, Any]:
+        coro = self.__core_sync_tool._ToolboxSyncTool__async_tool(**kwargs)
 
         # If a loop has not been provided, attempt to run in current thread.
-        if not self.__loop:
+        if not self.__core_sync_client._ToolboxSyncClient__loop:
             return await coro
 
         # Otherwise, run in the background thread.
-        return await asyncio.wrap_future(
-            asyncio.run_coroutine_threadsafe(coro, self.__loop)
+        await asyncio.wrap_future(
+            asyncio.run_coroutine_threadsafe(coro, self.__core_sync_client._ToolboxSyncTool__loop)
         )
 
-    def _run(self, **kwargs: Any) -> dict[str, Any]:
-        return self.__run_as_sync(self.__async_tool._arun(**kwargs))
-
-    async def _arun(self, **kwargs: Any) -> dict[str, Any]:
-        return await self.__run_as_async(self.__async_tool._arun(**kwargs))
 
     def add_auth_token_getters(
         self, auth_token_getters: dict[str, Callable[[], str]], strict: bool = True
@@ -93,27 +70,21 @@ def add_auth_token_getters(
         Args:
             auth_token_getters: A dictionary of authentication source names to
                 the functions that return corresponding ID token.
-            strict: If True, a ValueError is raised if any of the provided auth
-                parameters is already bound. If False, only a warning is issued.
 
         Returns:
             A new ToolboxTool instance that is a deep copy of the current
-            instance, with added auth tokens.
+            instance, with added auth token getters.
 
         Raises:
             ValueError: If any of the provided auth parameters is already
                 registered.
-            ValueError: If any of the provided auth parameters is already bound
-                and strict is True.
         """
-        return ToolboxTool(
-            self.__async_tool.add_auth_token_getters(auth_token_getters, strict),
-            self.__loop,
-            self.__thread,
-        )
+        new_core_sync_tool = self.__core_sync_tool.add_auth_token_getters(auth_token_getters)
+        return ToolboxTool(core_sync_tool=new_core_sync_tool)
+
 
     def add_auth_token_getter(
-        self, auth_source: str, get_id_token: Callable[[], str], strict: bool = True
+        self, auth_source: str, get_id_token: Callable[[], str]
     ) -> "ToolboxTool":
         """
         Registers a function to retrieve an ID token for a given authentication
@@ -122,28 +93,19 @@ def add_auth_token_getter(
         Args:
             auth_source: The name of the authentication source.
             get_id_token: A function that returns the ID token.
-            strict: If True, a ValueError is raised if the provided auth
-                parameter is already bound. If False, only a warning is issued.
 
         Returns:
             A new ToolboxTool instance that is a deep copy of the current
             instance, with added auth token.
 
         Raises:
             ValueError: If the provided auth parameter is already registered.
-            ValueError: If the provided auth parameter is already bound and
-                strict is True.
         """
-        return ToolboxTool(
-            self.__async_tool.add_auth_token_getter(auth_source, get_id_token, strict),
-            self.__loop,
-            self.__thread,
-        )
+        return self.add_auth_token_getters({auth_source: get_id_token})
 
     def bind_params(
         self,
         bound_params: dict[str, Union[Any, Callable[[], Any]]],
-        strict: bool = True,
     ) -> "ToolboxTool":
         """
         Registers values or functions to retrieve the value for the
@@ -152,25 +114,16 @@ def bind_params(
         Args:
             bound_params: A dictionary of the bound parameter name to the
                 value or function of the bound value.
-            strict: If True, a ValueError is raised if any of the provided bound
-                params is not defined in the tool's schema, or requires
-                authentication. If False, only a warning is issued.
 
         Returns:
             A new ToolboxTool instance that is a deep copy of the current
             instance, with added bound params.
 
         Raises:
             ValueError: If any of the provided bound params is already bound.
-            ValueError: if any of the provided bound params is not defined in
-                the tool's schema, or require authentication, and strict is
-                True.
         """
-        return ToolboxTool(
-            self.__async_tool.bind_params(bound_params, strict),
-            self.__loop,
-            self.__thread,
-        )
+        new_core_sync_tool = self.__core_sync_tool.bind_params(bound_params)
+        return ToolboxTool(core_sync_tool=new_core_sync_tool)
 
     def bind_param(
         self,
@@ -186,21 +139,12 @@ def bind_param(
             param_name: The name of the bound parameter.
             param_value: The value of the bound parameter, or a callable that
                 returns the value.
-            strict: If True, a ValueError is raised if the provided bound
-                param is not defined in the tool's schema, or requires
-                authentication. If False, only a warning is issued.
 
         Returns:
             A new ToolboxTool instance that is a deep copy of the current
             instance, with added bound param.
 
         Raises:
             ValueError: If the provided bound param is already bound.
-            ValueError: if the provided bound param is not defined in the tool's
-                schema, or requires authentication, and strict is True.
         """
-        return ToolboxTool(
-            self.__async_tool.bind_param(param_name, param_value, strict),
-            self.__loop,
-            self.__thread,
-        )
+        return self.bind_params({param_name: param_value})