feat(langchain-sdk): Support authentication in LangChain Toolbox SDK. (#133)

This PR adds support for authentication in the `ToolboxClient`.

Note: This is a resubmission of #117

Author: anubhav756

src/toolbox_langchain_sdk/client.py

@@ -1,5 +1,6 @@
 import asyncio
-from typing import Any, Optional, Type
+import warnings
+from typing import Any, Callable, Optional, Type
 
 from aiohttp import ClientSession
 from langchain_core.tools import StructuredTool
@@ -20,6 +21,8 @@ def __init__(self, url: str, session: Optional[ClientSession] = None):
         """
         self._url: str = url
         self._should_close_session: bool = session is None
+        self._id_token_getters: dict[str, Callable[[], str]] = {}
+        self._tool_param_auth: dict[str, dict[str, list[str]]] = {}
         self._session: ClientSession = session or ClientSession()
 
     async def close(self) -> None:
@@ -77,6 +80,35 @@ async def _load_toolset_manifest(
         url = f"{self._url}/api/toolset/{toolset_name or ''}"
         return await _load_yaml(url, self._session)
 
+    def _validate_auth(self, tool_name: str) -> bool:
+        """
+        Helper method that validates the authentication requirements of the tool
+        with the given tool_name. We consider the validation to pass if at least
+        one auth sources of each of the auth parameters, of the given tool, is
+        registered.
+
+        Args:
+            tool_name: Name of the tool to validate auth sources for.
+
+        Returns:
+            True if at least one permitted auth source of each of the auth
+            params, of the given tool, is registered. Also returns True if the
+            given tool does not require any auth sources.
+        """
+
+        if tool_name not in self._tool_param_auth:
+            return True
+
+        for permitted_auth_sources in self._tool_param_auth[tool_name].values():
+            found_match = False
+            for registered_auth_source in self._id_token_getters:
+                if registered_auth_source in permitted_auth_sources:
+                    found_match = True
+                    break
+            if not found_match:
+                return False
+        return True
+
     def _generate_tool(
         self, tool_name: str, manifest: ManifestSchema
     ) -> StructuredTool:
@@ -96,8 +128,16 @@ def _generate_tool(
             model_name=tool_name, schema=tool_schema.parameters
         )
 
+        # If the tool had parameters that require authentication, then right
+        # before invoking that tool, we validate whether all these required
+        # authentication sources have been registered or not.
         async def _tool_func(**kwargs: Any) -> dict:
-            return await _invoke_tool(self._url, self._session, tool_name, kwargs)
+            if not self._validate_auth(tool_name):
+                raise PermissionError(f"Login required before invoking {tool_name}.")
+
+            return await _invoke_tool(
+                self._url, self._session, tool_name, kwargs, self._id_token_getters
+            )
 
         return StructuredTool.from_function(
             coroutine=_tool_func,
@@ -106,21 +146,89 @@ async def _tool_func(**kwargs: Any) -> dict:
             args_schema=tool_model,
         )
 
-    async def load_tool(self, tool_name: str) -> StructuredTool:
+    def _process_auth_params(self, manifest: ManifestSchema) -> None:
+        """
+        Extracts parameters requiring authentication from the manifest.
+        Verifies each parameter has at least one valid auth source.
+
+        Args:
+            manifest: The manifest to validate and modify.
+
+        Warns:
+            UserWarning: If a parameter in the manifest has no valid sources.
+        """
+        for tool_name, tool_schema in manifest.tools.items():
+            non_auth_params = []
+            for param in tool_schema.parameters:
+
+                # Extract auth params from the tool schema.
+                #
+                # These parameters are removed from the manifest to prevent data
+                # validation errors since their values are inferred by the
+                # Toolbox service, not provided by the user.
+                #
+                # Store the permitted authentication sources for each parameter
+                # in '_tool_param_auth' for efficient validation in
+                # '_validate_auth'.
+                if not param.authSources:
+                    non_auth_params.append(param)
+                    continue
+
+                self._tool_param_auth.setdefault(tool_name, {})[
+                    param.name
+                ] = param.authSources
+
+            tool_schema.parameters = non_auth_params
+
+            # If none of the permitted auth sources of a parameter are
+            # registered, raise a warning message to the user.
+            if not self._validate_auth(tool_name):
+                warnings.warn(
+                    f"Some parameters of tool {tool_name} require authentication, but no valid auth sources are registered. Please register the required sources before use."
+                )
+
+    def add_auth_header(
+        self, auth_source: str, get_id_token: Callable[[], str]
+    ) -> None:
+        """
+        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.
+        """
+        self._id_token_getters[auth_source] = get_id_token
+
+    async def load_tool(
+        self, tool_name: str, auth_headers: dict[str, Callable[[], str]] = {}
+    ) -> StructuredTool:
         """
         Loads the tool, with the given tool name, from the Toolbox service.
 
         Args:
             tool_name: The name of the tool to load.
+            auth_headers: A mapping of authentication source names to
+                functions that retrieve ID tokens. If provided, these will
+                override or be added to the existing ID token getters.
+                Default: Empty.
 
         Returns:
             A tool loaded from the Toolbox
         """
+        for auth_source, get_id_token in auth_headers.items():
+            self.add_auth_header(auth_source, get_id_token)
+
         manifest: ManifestSchema = await self._load_tool_manifest(tool_name)
+
+        self._process_auth_params(manifest)
+
         return self._generate_tool(tool_name, manifest)
 
     async def load_toolset(
-        self, toolset_name: Optional[str] = None
+        self,
+        toolset_name: Optional[str] = None,
+        auth_headers: dict[str, Callable[[], str]] = {},
     ) -> list[StructuredTool]:
         """
         Loads tools from the Toolbox service, optionally filtered by toolset
@@ -129,12 +237,22 @@ async def load_toolset(
         Args:
             toolset_name: The name of the toolset to load.
                 Default: None. If not provided, then all the tools are loaded.
+            auth_headers: A mapping of authentication source names to
+                functions that retrieve ID tokens. If provided, these will
+                override or be added to the existing ID token getters.
+                Default: Empty.
 
         Returns:
             A list of all tools loaded from the Toolbox.
         """
+        for auth_source, get_id_token in auth_headers.items():
+            self.add_auth_header(auth_source, get_id_token)
+
         tools: list[StructuredTool] = []
         manifest: ManifestSchema = await self._load_toolset_manifest(toolset_name)
+
+        self._process_auth_params(manifest)
+
         for tool_name in manifest.tools:
             tools.append(self._generate_tool(tool_name, manifest))
         return tools

src/toolbox_langchain_sdk/utils.py

@@ -1,4 +1,5 @@
-from typing import Any, Optional, Type, cast
+import warnings
+from typing import Any, Callable, Optional, Type, cast
 
 import yaml
 from aiohttp import ClientSession
@@ -9,6 +10,7 @@ class ParameterSchema(BaseModel):
     name: str
     type: str
     description: str
+    authSources: Optional[list[str]] = None
 
 
 class ToolSchema(BaseModel):
@@ -34,8 +36,14 @@ async def _load_yaml(url: str, session: ClientSession) -> ManifestSchema:
     """
     async with session.get(url) as response:
         response.raise_for_status()
-        parsed_yaml = yaml.safe_load(await response.text())
-        return ManifestSchema(**parsed_yaml)
+        try:
+            parsed_yaml = yaml.safe_load(await response.text())
+        except yaml.YAMLError as e:
+            raise yaml.YAMLError(f"Failed to parse YAML from {url}: {e}") from e
+        try:
+            return ManifestSchema(**parsed_yaml)
+        except ValueError as e:
+            raise ValueError(f"Invalid YAML data from {url}: {e}") from e
 
 
 def _schema_to_model(model_name: str, schema: list[ParameterSchema]) -> Type[BaseModel]:
@@ -54,7 +62,8 @@ def _schema_to_model(model_name: str, schema: list[ParameterSchema]) -> Type[Bas
         field_definitions[field.name] = cast(
             Any,
             (
-                # TODO: Remove the hardcoded optional types once optional fields are supported by Toolbox.
+                # TODO: Remove the hardcoded optional types once optional fields
+                # are supported by Toolbox.
                 Optional[_parse_type(field.type)],
                 Field(description=field.description),
             ),
@@ -88,8 +97,30 @@ def _parse_type(type_: str) -> Any:
         raise ValueError(f"Unsupported schema type: {type_}")
 
 
+def _get_auth_headers(id_token_getters: dict[str, Callable[[], str]]) -> dict[str, str]:
+    """
+    Gets id tokens for the given auth sources in the getters map and returns
+    headers to be included in tool invocation.
+
+    Args:
+          id_token_getters: A dict that maps auth source names to the functions
+              that return its ID token.
+
+    Returns:
+      A dictionary of headers to be included in the tool invocation.
+    """
+    auth_headers = {}
+    for auth_source, get_id_token in id_token_getters.items():
+        auth_headers[f"{auth_source}_token"] = get_id_token()
+    return auth_headers
+
+
 async def _invoke_tool(
-    url: str, session: ClientSession, tool_name: str, data: dict
+    url: str,
+    session: ClientSession,
+    tool_name: str,
+    data: dict,
+    id_token_getters: dict[str, Callable[[], str]],
 ) -> dict:
     """
     Asynchronously makes an API call to the Toolbox service to invoke a tool.
@@ -99,12 +130,29 @@ async def _invoke_tool(
         session: The HTTP client session.
         tool_name: The name of the tool to invoke.
         data: The input data for the tool.
+        id_token_getters: A dict that maps auth source names to the functions
+            that return its ID token.
 
     Returns:
-        A dictionary containing the parsed JSON response from the tool invocation.
+        A dictionary containing the parsed JSON response from the tool
+        invocation.
     """
     url = f"{url}/api/tool/{tool_name}/invoke"
-    async with session.post(url, json=_convert_none_to_empty_string(data)) as response:
+    auth_headers = _get_auth_headers(id_token_getters)
+
+    # ID tokens contain sensitive user information (claims). Transmitting these
+    # over HTTP exposes the data to interception and unauthorized access. Always
+    # use HTTPS to ensure secure communication and protect user privacy.
+    if auth_headers and not url.startswith("https://"):
+        warnings.warn(
+            "Sending ID token over HTTP. User data may be exposed. Use HTTPS for secure communication."
+        )
+
+    async with session.post(
+        url,
+        json=_convert_none_to_empty_string(data),
+        headers=auth_headers,
+    ) as response:
         response.raise_for_status()
         return await response.json()