feat(langchain-sdk): Add features for binding parameters to ToolboxTool class. (#192)

The newly implemented `ToolboxTool` class manages tool state and
supports this new feature of binding parameters along wiith the existing
OAuth.

> [!NOTE]
> These changes are done in the LlamaIndex SDK as well in #203, along
with documentation updates in #193.

Author: anubhav756

src/toolbox_langchain_sdk/client.py

@@ -99,6 +99,8 @@ async def load_tool(
         tool_name: str,
         auth_tokens: dict[str, Callable[[], str]] = {},
         auth_headers: Optional[dict[str, Callable[[], str]]] = None,
+        bound_params: dict[str, Union[Any, Callable[[], Any]]] = {},
+        strict: bool = True,
     ) -> ToolboxTool:
         """
         Loads the tool with the given tool name from the Toolbox service.
@@ -108,6 +110,11 @@ async def load_tool(
             auth_tokens: An optional mapping of authentication source names to
                 functions that retrieve ID tokens.
             auth_headers: Deprecated. Use `auth_tokens` instead.
+            bound_params: An optional mapping of parameter names to their
+                bound values.
+            strict: If True, raises a ValueError if any of the given bound
+                parameters are missing from the schema or require
+                authentication. If False, only issues a warning.
 
         Returns:
             A tool loaded from the Toolbox.
@@ -132,13 +139,17 @@ async def load_tool(
             self._url,
             self._session,
             auth_tokens,
+            bound_params,
+            strict,
         )
 
     async def load_toolset(
         self,
         toolset_name: Optional[str] = None,
         auth_tokens: dict[str, Callable[[], str]] = {},
         auth_headers: Optional[dict[str, Callable[[], str]]] = None,
+        bound_params: dict[str, Union[Any, Callable[[], Any]]] = {},
+        strict: bool = True,
     ) -> list[ToolboxTool]:
         """
         Loads tools from the Toolbox service, optionally filtered by toolset
@@ -150,6 +161,11 @@ async def load_toolset(
             auth_tokens: An optional mapping of authentication source names to
                 functions that retrieve ID tokens.
             auth_headers: Deprecated. Use `auth_tokens` instead.
+            bound_params: An optional mapping of parameter names to their
+                bound values.
+            strict: If True, raises a ValueError if any of the given bound
+                parameters are missing from the schema or require
+                authentication. If False, only issues a warning.
 
         Returns:
             A list of all tools loaded from the Toolbox.
@@ -178,6 +194,8 @@ async def load_toolset(
                     self._url,
                     self._session,
                     auth_tokens,
+                    bound_params,
+                    strict,
                 )
             )
         return tools

src/toolbox_langchain_sdk/tools.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 
 from copy import deepcopy
-from typing import Any, Callable
+from typing import Any, Callable, Union
 from warnings import warn
 
 from aiohttp import ClientSession
@@ -24,6 +24,7 @@
     ParameterSchema,
     ToolSchema,
     _find_auth_params,
+    _find_bound_params,
     _invoke_tool,
     _schema_to_model,
 )
@@ -32,7 +33,7 @@
 class ToolboxTool(StructuredTool):
     """
     A subclass of LangChain's StructuredTool that supports features specific to
-    Toolbox, like authenticated tools.
+    Toolbox, like bound parameters and authenticated tools.
     """
 
     def __init__(
@@ -42,6 +43,8 @@ def __init__(
         url: str,
         session: ClientSession,
         auth_tokens: dict[str, Callable[[], str]] = {},
+        bound_params: dict[str, Union[Any, Callable[[], Any]]] = {},
+        strict: bool = True,
     ) -> None:
         """
         Initializes a ToolboxTool instance.
@@ -53,6 +56,11 @@ def __init__(
             session: The HTTP client session.
             auth_tokens: A mapping of authentication source names to functions
                 that retrieve ID tokens.
+            bound_params: A mapping of parameter names to their bound
+                values.
+            strict: If True, raises a ValueError if any of the given bound
+                parameters are missing from the schema or require
+                authentication. If False, only issues a warning.
         """
 
         # If the schema is not already a ToolSchema instance, we create one from
@@ -63,10 +71,51 @@ def __init__(
             schema = ToolSchema(**schema)
 
         auth_params, non_auth_params = _find_auth_params(schema.parameters)
+        non_auth_bound_params, non_auth_non_bound_params = _find_bound_params(
+            non_auth_params, list(bound_params)
+        )
+
+        # Check if the user is trying to bind a param that is authenticated or
+        # is missing from the given schema.
+        auth_bound_params: list[str] = []
+        missing_bound_params: list[str] = []
+        for bound_param in bound_params:
+            if bound_param in [param.name for param in auth_params]:
+                auth_bound_params.append(bound_param)
+            elif bound_param not in [param.name for param in non_auth_params]:
+                missing_bound_params.append(bound_param)
+
+        # Create error messages for any params that are found to be
+        # authenticated or missing.
+        messages: list[str] = []
+        if auth_bound_params:
+            messages.append(
+                f"Parameter(s) {', '.join(auth_bound_params)} already authenticated and cannot be bound."
+            )
+        if missing_bound_params:
+            messages.append(
+                f"Parameter(s) {', '.join(missing_bound_params)} missing and cannot be bound."
+            )
+
+        # Join any error messages and raise them as an error or warning,
+        # depending on the value of the strict flag.
+        if messages:
+            message = "\n\n".join(messages)
+            if strict:
+                raise ValueError(message)
+            warn(message)
+
+        # Bind values for parameters present in the schema that don't require
+        # authentication.
+        bound_params = {
+            param_name: param_value
+            for param_name, param_value in bound_params.items()
+            if param_name in [param.name for param in non_auth_bound_params]
+        }
 
         # Update the tools schema to validate only the presence of parameters
-        # that do not require authentication.
-        schema.parameters = non_auth_params
+        # that neither require authentication nor are bound.
+        schema.parameters = non_auth_non_bound_params
 
         # Due to how pydantic works, we must initialize the underlying
         # StructuredTool class before assigning values to member variables.
@@ -84,6 +133,7 @@ def __init__(
         self._session: ClientSession = session
         self._auth_tokens: dict[str, Callable[[], str]] = auth_tokens
         self._auth_params: list[ParameterSchema] = auth_params
+        self._bound_params: dict[str, Union[Any, Callable[[], Any]]] = bound_params
 
         # Warn users about any missing authentication so they can add it before
         # tool invocation.
@@ -106,6 +156,17 @@ async def __tool_func(self, **kwargs: Any) -> dict:
         # authentication sources have been registered or not.
         self.__validate_auth()
 
+        # Evaluate dynamic parameter values if any
+        evaluated_params = {}
+        for param_name, param_value in self._bound_params.items():
+            if callable(param_value):
+                evaluated_params[param_name] = param_value()
+            else:
+                evaluated_params[param_name] = param_value
+
+        # Merge bound parameters with the provided arguments
+        kwargs.update(evaluated_params)
+
         return await _invoke_tool(
             self._url, self._session, self._name, kwargs, self._auth_tokens
         )
@@ -154,42 +215,66 @@ def __create_copy(
         self,
         *,
         auth_tokens: dict[str, Callable[[], str]] = {},
+        bound_params: dict[str, Union[Any, Callable[[], Any]]] = {},
+        strict: bool,
     ) -> Self:
         """
         Creates a deep copy of the current ToolboxTool instance, allowing for
-        modification of auth tokens.
+        modification of auth tokens and bound params.
 
         This method enables the creation of new tool instances with inherited
         properties from the current instance, while optionally updating the auth
-        tokens. This is useful for creating variations of the tool with
-        additional auth tokens without modifying the original instance, ensuring
-        immutability.
+        tokens and bound params. This is useful for creating variations of the
+        tool with additional auth tokens or bound params without modifying the
+        original instance, ensuring immutability.
 
         Args:
             auth_tokens: A dictionary of auth source names to functions that
                 retrieve ID tokens. These tokens will be merged with the
                 existing auth tokens.
+            bound_params: A dictionary of parameter names to their
+                bound values or functions to retrieve the values. These params
+                will be merged with the existing bound params.
+            strict: If True, raises a ValueError if any of the given bound
+                parameters are missing from the schema or require
+                authentication. If False, only issues a warning.
 
         Returns:
             A new ToolboxTool instance that is a deep copy of the current
-            instance, with optionally updated auth tokens.
+            instance, with added auth tokens or bound params.
         """
+        new_schema = deepcopy(self._schema)
+
+        # Reconstruct the complete parameter schema by merging the auth
+        # parameters back with the non-auth parameters. This is necessary to
+        # accurately validate the new combination of auth tokens and bound
+        # params in the constructor of the new ToolboxTool instance, ensuring
+        # that any overlaps or conflicts are correctly identified and reported
+        # as errors or warnings, depending on the given `strict` flag.
+        new_schema.parameters += self._auth_params
         return type(self)(
             name=self._name,
-            schema=deepcopy(self._schema),
+            schema=new_schema,
             url=self._url,
             session=self._session,
             auth_tokens={**self._auth_tokens, **auth_tokens},
+            bound_params={**self._bound_params, **bound_params},
+            strict=strict,
         )
 
-    def add_auth_tokens(self, auth_tokens: dict[str, Callable[[], str]]) -> Self:
+    def add_auth_tokens(
+        self, auth_tokens: dict[str, Callable[[], str]], strict: bool = True
+    ) -> Self:
         """
         Registers functions to retrieve ID tokens for the corresponding
         authentication sources.
 
         Args:
             auth_tokens: 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
+                tokens are already registered, or are already bound. If False,
+                only a warning is issued.
 
         Returns:
             A new ToolboxTool instance that is a deep copy of the current
@@ -207,19 +292,82 @@ def add_auth_tokens(self, auth_tokens: dict[str, Callable[[], str]]) -> Self:
                 f"Authentication source(s) `{', '.join(dupe_tokens)}` already registered in tool `{self._name}`."
             )
 
-        return self.__create_copy(auth_tokens=auth_tokens)
+        return self.__create_copy(auth_tokens=auth_tokens, strict=strict)
 
-    def add_auth_token(self, auth_source: str, get_id_token: Callable[[], str]) -> Self:
+    def add_auth_token(
+        self, auth_source: str, get_id_token: Callable[[], str], strict: bool = True
+    ) -> Self:
         """
         Registers a function to retrieve an ID token for a given authentication
         source.
 
         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 any of the provided auth
+                tokens are already registered, or are 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.
         """
-        return self.add_auth_tokens({auth_source: get_id_token})
+        return self.add_auth_tokens({auth_source: get_id_token}, strict=strict)
+
+    def bind_params(
+        self,
+        bound_params: dict[str, Union[Any, Callable[[], Any]]],
+        strict: bool = True,
+    ) -> Self:
+        """
+        Registers values or functions to retrieve the value for the
+        corresponding bound parameters.
+
+        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 are already bound, not defined in the tool's schema, or
+                require 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.
+        """
+
+        # Check if the parameter is already bound.
+        dupe_params: list[str] = []
+        for param_name, _ in bound_params.items():
+            if param_name in self._bound_params:
+                dupe_params.append(param_name)
+
+        if dupe_params:
+            raise ValueError(
+                f"Parameter(s) `{', '.join(dupe_params)}` already bound in tool `{self._name}`."
+            )
+
+        return self.__create_copy(bound_params=bound_params, strict=strict)
+
+    def bind_param(
+        self,
+        param_name: str,
+        param_value: Union[Any, Callable[[], Any]],
+        strict: bool = True,
+    ) -> Self:
+        """
+        Registers a value or a function to retrieve the value for a given
+        bound parameter.
+
+        Args:
+            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 any of the provided bound
+                params are already bound, not defined in the tool's schema, or
+                require 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.
+        """
+        return self.bind_params({param_name: param_value}, strict)

src/toolbox_langchain_sdk/utils.py

@@ -238,3 +238,18 @@ def _find_auth_params(
             _non_auth_params.append(param)
 
     return (_auth_params, _non_auth_params)
+
+
+def _find_bound_params(
+    params: list[ParameterSchema], bound_params: list[str]
+) -> tuple[list[ParameterSchema], list[ParameterSchema]]:
+    _bound_params: list[ParameterSchema] = []
+    _non_bound_params: list[ParameterSchema] = []
+
+    for param in params:
+        if param.name in bound_params:
+            _bound_params.append(param)
+        else:
+            _non_bound_params.append(param)
+
+    return (_bound_params, _non_bound_params)