feat(toolbox-llamaindex)!: Update tool and client codebase

Author: anubhav756

packages/toolbox-llamaindex/src/toolbox_llamaindex/async_client.py

@@ -16,9 +16,9 @@
 from warnings import warn
 
 from aiohttp import ClientSession
+from toolbox_core.client import ToolboxClient as ToolboxCoreClient
 
-from .tools import AsyncToolboxTool
-from .utils import ManifestSchema, _load_manifest
+from .async_tools import AsyncToolboxTool
 
 
 # This class is an internal implementation detail and is not exposed to the
@@ -38,67 +38,72 @@ def __init__(
             url: The base URL of the Toolbox service.
             session: An HTTP client session.
         """
-        self.__url = url
-        self.__session = session
+        self.__core_client = ToolboxCoreClient(url=url, session=session)
 
     async def aload_tool(
         self,
         tool_name: str,
-        auth_tokens: dict[str, Callable[[], str]] = {},
+        auth_token_getters: dict[str, Callable[[], str]] = {},
+        auth_tokens: Optional[dict[str, Callable[[], str]]] = None,
         auth_headers: Optional[dict[str, Callable[[], str]]] = None,
         bound_params: dict[str, Union[Any, Callable[[], Any]]] = {},
-        strict: bool = True,
     ) -> AsyncToolboxTool:
         """
         Loads the tool with the given tool name from the Toolbox service.
 
         Args:
             tool_name: The name of the tool to load.
-            auth_tokens: An optional mapping of authentication source names to
-                functions that retrieve ID tokens.
-            auth_headers: Deprecated. Use `auth_tokens` instead.
+            auth_token_getters: An optional mapping of authentication source
+                names to functions that retrieve ID tokens.
+            auth_tokens: Deprecated. Use `auth_token_getters` instead.
+            auth_headers: Deprecated. Use `auth_token_getters` 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.
         """
+        if auth_tokens:
+            if auth_token_getters:
+                warn(
+                    "Both `auth_token_getters` and `auth_tokens` are provided. `auth_tokens` is deprecated, and `auth_token_getters` will be used.",
+                    DeprecationWarning,
+                )
+            else:
+                warn(
+                    "Argument `auth_tokens` is deprecated. Use `auth_token_getters` instead.",
+                    DeprecationWarning,
+                )
+                auth_token_getters = auth_tokens
+
         if auth_headers:
-            if auth_tokens:
+            if auth_token_getters:
                 warn(
-                    "Both `auth_tokens` and `auth_headers` are provided. `auth_headers` is deprecated, and `auth_tokens` will be used.",
+                    "Both `auth_token_getters` and `auth_headers` are provided. `auth_headers` is deprecated, and `auth_token_getters` will be used.",
                     DeprecationWarning,
                 )
             else:
                 warn(
-                    "Argument `auth_headers` is deprecated. Use `auth_tokens` instead.",
+                    "Argument `auth_headers` is deprecated. Use `auth_token_getters` instead.",
                     DeprecationWarning,
                 )
-                auth_tokens = auth_headers
-
-        url = f"{self.__url}/api/tool/{tool_name}"
-        manifest: ManifestSchema = await _load_manifest(url, self.__session)
-
-        return AsyncToolboxTool(
-            tool_name,
-            manifest.tools[tool_name],
-            self.__url,
-            self.__session,
-            auth_tokens,
-            bound_params,
-            strict,
+                auth_token_getters = auth_headers
+
+        core_tool = await self.__core_client.load_tool(
+            name=tool_name,
+            auth_token_getters=auth_token_getters,
+            bound_params=bound_params,
         )
+        return AsyncToolboxTool(core_tool=core_tool)
 
     async def aload_toolset(
         self,
         toolset_name: Optional[str] = None,
-        auth_tokens: dict[str, Callable[[], str]] = {},
+        auth_token_getters: dict[str, Callable[[], str]] = {},
+        auth_tokens: Optional[dict[str, Callable[[], str]]] = None,
         auth_headers: Optional[dict[str, Callable[[], str]]] = None,
         bound_params: dict[str, Union[Any, Callable[[], Any]]] = {},
-        strict: bool = True,
+        strict: bool = False,
     ) -> list[AsyncToolboxTool]:
         """
         Loads tools from the Toolbox service, optionally filtered by toolset
@@ -107,65 +112,76 @@ async def aload_toolset(
         Args:
             toolset_name: The name of the toolset to load. If not provided,
                 all tools are loaded.
-            auth_tokens: An optional mapping of authentication source names to
-                functions that retrieve ID tokens.
-            auth_headers: Deprecated. Use `auth_tokens` instead.
+            auth_token_getters: An optional mapping of authentication source
+                names to functions that retrieve ID tokens.
+            auth_tokens: Deprecated. Use `auth_token_getters` instead.
+            auth_headers: Deprecated. Use `auth_token_getters` 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.
+            strict: If True, raises an error if *any* loaded tool instance fails
+                to utilize at least one provided parameter or auth token (if any
+                provided). If False (default), raises an error only if a
+                user-provided parameter or auth token cannot be applied to *any*
+                loaded tool across the set.
 
         Returns:
             A list of all tools loaded from the Toolbox.
         """
-        if auth_headers:
-            if auth_tokens:
+        if auth_tokens:
+            if auth_token_getters:
                 warn(
-                    "Both `auth_tokens` and `auth_headers` are provided. `auth_headers` is deprecated, and `auth_tokens` will be used.",
+                    "Both `auth_token_getters` and `auth_tokens` are provided. `auth_tokens` is deprecated, and `auth_token_getters` will be used.",
                     DeprecationWarning,
                 )
             else:
                 warn(
-                    "Argument `auth_headers` is deprecated. Use `auth_tokens` instead.",
+                    "Argument `auth_tokens` is deprecated. Use `auth_token_getters` instead.",
                     DeprecationWarning,
                 )
-                auth_tokens = auth_headers
-
-        url = f"{self.__url}/api/toolset/{toolset_name or ''}"
-        manifest: ManifestSchema = await _load_manifest(url, self.__session)
-        tools: list[AsyncToolboxTool] = []
-
-        for tool_name, tool_schema in manifest.tools.items():
-            tools.append(
-                AsyncToolboxTool(
-                    tool_name,
-                    tool_schema,
-                    self.__url,
-                    self.__session,
-                    auth_tokens,
-                    bound_params,
-                    strict,
+                auth_token_getters = auth_tokens
+
+        if auth_headers:
+            if auth_token_getters:
+                warn(
+                    "Both `auth_token_getters` and `auth_headers` are provided. `auth_headers` is deprecated, and `auth_token_getters` will be used.",
+                    DeprecationWarning,
+                )
+            else:
+                warn(
+                    "Argument `auth_headers` is deprecated. Use `auth_token_getters` instead.",
+                    DeprecationWarning,
                 )
-            )
+                auth_token_getters = auth_headers
+
+        core_tools = await self.__core_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(AsyncToolboxTool(core_tool=core_tool))
         return tools
 
     def load_tool(
         self,
         tool_name: str,
-        auth_tokens: dict[str, Callable[[], str]] = {},
+        auth_token_getters: dict[str, Callable[[], str]] = {},
+        auth_tokens: Optional[dict[str, Callable[[], str]]] = None,
         auth_headers: Optional[dict[str, Callable[[], str]]] = None,
         bound_params: dict[str, Union[Any, Callable[[], Any]]] = {},
-        strict: bool = True,
     ) -> AsyncToolboxTool:
         raise NotImplementedError("Synchronous methods not supported by async client.")
 
     def load_toolset(
         self,
         toolset_name: Optional[str] = None,
-        auth_tokens: dict[str, Callable[[], str]] = {},
+        auth_token_getters: dict[str, Callable[[], str]] = {},
+        auth_tokens: Optional[dict[str, Callable[[], str]]] = None,
         auth_headers: Optional[dict[str, Callable[[], str]]] = None,
         bound_params: dict[str, Union[Any, Callable[[], Any]]] = {},
-        strict: bool = True,
+        strict: bool = False,
     ) -> list[AsyncToolboxTool]:
         raise NotImplementedError("Synchronous methods not supported by async client.")