feat: Wrap toolbox-langchain's AsyncToolboxTool over toolbox-core's ToolboxTool.

This simplifies code and increases maintainability while removing duplicate code.

Author: anubhav756

packages/toolbox-core/src/toolbox_core/client.py

@@ -13,14 +13,13 @@
 # limitations under the License.
 
 
-import types
 from typing import Any, Callable, Mapping, Optional, Union
 
 from aiohttp import ClientSession
 
 from .protocol import ManifestSchema, ToolSchema
 from .tool import ToolboxTool, identify_required_authn_params
-
+from .utils import parse_tool
 
 class ToolboxClient:
     """
@@ -55,42 +54,6 @@ def __init__(
             session = ClientSession()
         self.__session = session
 
-    def __parse_tool(
-        self,
-        name: str,
-        schema: ToolSchema,
-        auth_token_getters: dict[str, Callable[[], str]],
-        all_bound_params: Mapping[str, Union[Callable[[], Any], Any]],
-    ) -> ToolboxTool:
-        """Internal helper to create a callable tool from its schema."""
-        # sort into reg, authn, and bound params
-        params = []
-        authn_params: dict[str, list[str]] = {}
-        bound_params: dict[str, Callable[[], str]] = {}
-        for p in schema.parameters:
-            if p.authSources:  # authn parameter
-                authn_params[p.name] = p.authSources
-            elif p.name in all_bound_params:  # bound parameter
-                bound_params[p.name] = all_bound_params[p.name]
-            else:  # regular parameter
-                params.append(p)
-
-        authn_params = identify_required_authn_params(
-            authn_params, auth_token_getters.keys()
-        )
-
-        tool = ToolboxTool(
-            session=self.__session,
-            base_url=self.__base_url,
-            name=name,
-            description=schema.description,
-            params=params,
-            # create a read-only values for the maps to prevent mutation
-            required_authn_params=types.MappingProxyType(authn_params),
-            auth_service_token_getters=types.MappingProxyType(auth_token_getters),
-            bound_params=types.MappingProxyType(bound_params),
-        )
-        return tool
 
     async def __aenter__(self):
         """
@@ -163,8 +126,8 @@ async def load_tool(
         if name not in manifest.tools:
             # TODO: Better exception
             raise Exception(f"Tool '{name}' not found!")
-        tool = self.__parse_tool(
-            name, manifest.tools[name], auth_token_getters, bound_params
+        tool = parse_tool(
+            self.__session, self.__base_url, name, manifest.tools[name], auth_token_getters, bound_params
         )
 
         return tool
@@ -199,7 +162,7 @@ async def load_toolset(
 
         # parse each tools name and schema into a list of ToolboxTools
         tools = [
-            self.__parse_tool(n, s, auth_token_getters, bound_params)
+            parse_tool(self.__session, self.__base_url, n, s, auth_token_getters, bound_params)
             for n, s in manifest.tools.items()
         ]
         return tools

packages/toolbox-core/src/toolbox_core/utils.py

@@ -25,10 +25,12 @@
     Union,
     cast,
 )
-
+from aiohttp import ClientSession
+from types import MappingProxyType
 from pydantic import BaseModel, Field, create_model
 
-from toolbox_core.protocol import ParameterSchema
+from toolbox_core.protocol import ParameterSchema, ToolSchema
+from toolbox_core.tool import ToolboxTool
 
 
 def create_func_docstring(description: str, params: Sequence[ParameterSchema]) -> str:
@@ -110,3 +112,42 @@ async def resolve_value(
     elif callable(source):
         return source()
     return source
+
+
+def parse_tool(
+    session: ClientSession,
+    base_url: str,
+    name: str,
+    schema: ToolSchema,
+    auth_token_getters: dict[str, Callable[[], str]],
+    all_bound_params: Mapping[str, Union[Callable[[], Any], Any]],
+) -> ToolboxTool:
+    """Helper to create a callable tool from its schema."""
+    # sort into reg, authn, and bound params
+    params = []
+    authn_params: dict[str, list[str]] = {}
+    bound_params: dict[str, Callable[[], str]] = {}
+    for p in schema.parameters:
+        if p.authSources:  # authn parameter
+            authn_params[p.name] = p.authSources
+        elif p.name in all_bound_params:  # bound parameter
+            bound_params[p.name] = all_bound_params[p.name]
+        else:  # regular parameter
+            params.append(p)
+
+    authn_params = identify_required_authn_params(
+        authn_params, auth_token_getters.keys()
+    )
+
+    tool = ToolboxTool(
+        session=session,
+        base_url=base_url,
+        name=name,
+        description=schema.description,
+        params=params,
+        # create a read-only values for the maps to prevent mutation
+        required_authn_params=MappingProxyType(authn_params),
+        auth_service_token_getters=MappingProxyType(auth_token_getters),
+        bound_params=MappingProxyType(bound_params),
+    )
+    return tool