feat(llamaindex-sdk): Implement OAuth support for LlamaIndex. (#159)

* Update docs for Authenticated Tools. * Implement test cases for OAuth in LlamaIndex. * Sync SDK dependency versions b/w LangChain and LlamaIndex.
googleapis · Dec 23, 2024 · 003ce51 · 003ce51
1 parent 5ae30c2
commit 003ce51
Show file tree

Hide file tree

Showing 12 changed files with 1,057 additions and 133 deletions.
diff --git a/sdks/langchain/README.md b/sdks/langchain/README.md
@@ -4,29 +4,32 @@ This SDK allows you to seamlessly integrate the functionalities of
 [Toolbox](https://github.com/googleapis/genai-toolbox) into your LLM
 applications, enabling advanced orchestration and interaction with GenAI models.
 
+<!-- TOC ignore:true -->
+## Table of Contents
 <!-- TOC -->
 
-- [GenAI Toolbox SDK](#genai-toolbox-sdk)
-    - [Installation](#installation)
-    - [Usage](#usage)
-    - [Load a toolset](#load-a-toolset)
-    - [Load a single tool](#load-a-single-tool)
-    - [Use with LangChain](#use-with-langchain)
-    - [Use with LangGraph](#use-with-langgraph)
-        - [Represent Tools as Nodes](#represent-tools-as-nodes)
-        - [Connect Tools with LLM](#connect-tools-with-llm)
-    - [Manual usage](#manual-usage)
-    - [Authenticating Tools](#authenticating-tools)
-        - [Supported Authentication Mechanisms](#supported-authentication-mechanisms)
-        - [Configuring Tools for Authentication](#configuring-tools-for-authentication)
-        - [Configure SDK for Authentication](#configure-sdk-for-authentication)
-        - [Complete Example](#complete-example)
+- [Installation](#installation)
+- [Usage](#usage)
+- [Load a toolset](#load-a-toolset)
+- [Load a single tool](#load-a-single-tool)
+- [Use with LangChain](#use-with-langchain)
+- [Use with LangGraph](#use-with-langgraph)
+    - [Represent Tools as Nodes](#represent-tools-as-nodes)
+    - [Connect Tools with LLM](#connect-tools-with-llm)
+- [Manual usage](#manual-usage)
+- [Authenticating Tools](#authenticating-tools)
+    - [Supported Authentication Mechanisms](#supported-authentication-mechanisms)
+    - [Configuring Tools for Authentication](#configuring-tools-for-authentication)
+    - [Configure SDK for Authentication](#configure-sdk-for-authentication)
+    - [Complete Example](#complete-example)
 
 <!-- /TOC -->
+
 ## Installation
 
 > [!IMPORTANT]
-> This SDK is not yet available on PyPI. For now, install it from source by following these [installation instructions](DEVELOPER.md).
+> This SDK is not yet available on PyPI. For now, install it from source by
+> following these [installation instructions](DEVELOPER.md).
 
 You can install the Toolbox SDK for LangChain using `pip`.
 

diff --git a/sdks/langchain/pyproject.toml b/sdks/langchain/pyproject.toml
@@ -8,10 +8,10 @@ authors = [
     {name = "Google LLC", email = "googleapis-packages@google.com"}
 ]
 dependencies = [
-    "aiohttp>=3.7.0,<4.0.0",
     "langchain-core>=0.2.23,<1.0.0",
-    "PyYAML>=6.0.0,<7.0.0",
-    "pydantic>=2.0,<3.0.0",
+    "PyYAML>=6.0.1,<7.0.0",
+    "pydantic>=2.7.0,<3.0.0",
+    "aiohttp>=3.8.6,<4.0.0",
 ]
 
 classifiers = [

diff --git a/sdks/langchain/tests/test_client.py b/sdks/langchain/tests/test_client.py
@@ -5,6 +5,7 @@
 import aiohttp
 import pytest
 from langchain_core.tools import StructuredTool
+from pydantic import BaseModel
 
 from toolbox_langchain_sdk import ToolboxClient
 from toolbox_langchain_sdk.utils import ManifestSchema, ParameterSchema, ToolSchema
@@ -166,7 +167,7 @@ async def test_load_tool_success(mock_generate_tool, mock_load_manifest):
     mock_generate_tool.return_value = StructuredTool(
         name="test_tool",
         description="This is test tool.",
-        args_schema=None,
+        args_schema=BaseModel,
         coroutine=AsyncMock(),
     )
 
@@ -201,13 +202,13 @@ async def test_load_toolset_success(mock_generate_tool, mock_load_manifest):
         StructuredTool(
             name="test_tool",
             description="This is test tool.",
-            args_schema=None,
+            args_schema=BaseModel,
             coroutine=AsyncMock(),
         ),
         StructuredTool(
             name="test_tool2",
             description="This is test tool 2.",
-            args_schema=None,
+            args_schema=BaseModel,
             coroutine=AsyncMock(),
         ),
     ] * 2

diff --git a/sdks/langchain/tests/test_utils.py b/sdks/langchain/tests/test_utils.py
@@ -1,12 +1,12 @@
 import asyncio
 import json
+import re
 import warnings
 from typing import Union
 from unittest.mock import AsyncMock, Mock, patch
 
 import aiohttp
 import pytest
-from aiohttp import ClientSession
 from pydantic import BaseModel
 
 from toolbox_langchain_sdk.utils import (
@@ -91,8 +91,6 @@ async def test_load_manifest_invalid_json(self, mock_get, mock_manifest):
         with pytest.raises(Exception) as e:
             session = aiohttp.ClientSession()
             await _load_manifest(URL, session)
-            await session.close()
-            mock_get.assert_called_once_with(URL)
 
         mock_get.assert_called_once_with(URL)
         assert isinstance(e.value, json.JSONDecodeError)
@@ -111,14 +109,12 @@ async def test_load_manifest_invalid_manifest(self, mock_get, mock_manifest):
         with pytest.raises(Exception) as e:
             session = aiohttp.ClientSession()
             await _load_manifest(URL, session)
-            await session.close()
-            mock_get.assert_called_once_with(URL)
 
         mock_get.assert_called_once_with(URL)
         assert isinstance(e.value, ValueError)
-        assert (
-            str(e.value)
-            == "Invalid JSON data from https://my-toolbox.com/test: 2 validation errors for ManifestSchema\nserverVersion\n  Field required [type=missing, input_value={'something': 'invalid'}, input_type=dict]\n    For further information visit https://errors.pydantic.dev/2.10/v/missing\ntools\n  Field required [type=missing, input_value={'something': 'invalid'}, input_type=dict]\n    For further information visit https://errors.pydantic.dev/2.10/v/missing"
+        assert re.match(
+            r"Invalid JSON data from https://my-toolbox.com/test: 2 validation errors for ManifestSchema\nserverVersion\n  Field required \[type=missing, input_value={'something': 'invalid'}, input_type=dict]\n    For further information visit https://errors.pydantic.dev/\d+\.\d+/v/missing\ntools\n  Field required \[type=missing, input_value={'something': 'invalid'}, input_type=dict]\n    For further information visit https://errors.pydantic.dev/\d+\.\d+/v/missing",
+            str(e.value),
         )
 
     @pytest.mark.asyncio
@@ -132,7 +128,6 @@ async def test_load_manifest_api_error(self, mock_get, mock_manifest):
         with pytest.raises(aiohttp.ClientError) as exc_info:
             session = aiohttp.ClientSession()
             await _load_manifest(URL, session)
-            await session.close()
         mock_get.assert_called_once_with(URL)
         assert exc_info.value == error
 
@@ -180,7 +175,11 @@ async def test_invoke_tool(self, mock_post):
         mock_post.return_value.__aenter__.return_value = mock_response
 
         result = await _invoke_tool(
-            "http://localhost:8000", ClientSession(), "tool_name", {"input": "data"}, {}
+            "http://localhost:8000",
+            aiohttp.ClientSession(),
+            "tool_name",
+            {"input": "data"},
+            {},
         )
 
         mock_post.assert_called_once_with(
@@ -204,7 +203,7 @@ async def test_invoke_tool_unsecure_with_auth(self, mock_post):
         ):
             result = await _invoke_tool(
                 "http://localhost:8000",
-                ClientSession(),
+                aiohttp.ClientSession(),
                 "tool_name",
                 {"input": "data"},
                 {"my_test_auth": lambda: "fake_id_token"},
@@ -220,7 +219,7 @@ async def test_invoke_tool_unsecure_with_auth(self, mock_post):
     @pytest.mark.asyncio
     @patch("aiohttp.ClientSession.post")
     async def test_invoke_tool_secure_with_auth(self, mock_post):
-        session = ClientSession()
+        session = aiohttp.ClientSession()
         mock_response = Mock()
         mock_response.raise_for_status = Mock()
         mock_response.json = AsyncMock(return_value={"key": "value"})

diff --git a/sdks/llamaindex/README.md b/sdks/llamaindex/README.md
@@ -2,13 +2,10 @@
 
 This SDK allows you to seamlessly integrate the functionalities of
 [Toolbox](https://github.com/googleapis/genai-toolbox) into your LLM
-applications, enabling advanced orchestration and interaction with GenAI
-models.
+applications, enabling advanced orchestration and interaction with GenAI models.
 
 <!-- TOC ignore:true -->
-
 ## Table of Contents
-
 <!-- TOC -->
 
 - [Installation](#installation)
@@ -17,26 +14,31 @@ models.
 - [Load a single tool](#load-a-single-tool)
 - [Use with LlamaIndex](#use-with-llamaindex)
 - [Manual usage](#manual-usage)
+- [Authenticating Tools](#authenticating-tools)
+    - [Supported Authentication Mechanisms](#supported-authentication-mechanisms)
+    - [Configuring Tools for Authentication](#configuring-tools-for-authentication)
+    - [Configure SDK for Authentication](#configure-sdk-for-authentication)
+    - [Complete Example](#complete-example)
 
 <!-- /TOC -->
 
 ## Installation
 
+> [!IMPORTANT]
+> This SDK is not yet available on PyPI. For now, install it from source by
+> following these [installation instructions](DEVELOPER.md).
+
 You can install the Toolbox SDK for LlamaIndex using `pip`.
 
 ```bash
 pip install toolbox-llamaindex-sdk
 ```
 
-> [!IMPORTANT]
-> This SDK is not yet available on PyPI. For now, install it from source by
-following these [instructions](DEVELOPER.md#setting-up-a-development-environment).
-
 ## Usage
 
 Import and initialize the toolbox client.
 
-```python
+```py
 from toolbox_llamaindex_sdk import ToolboxClient
 
 # Replace with your Toolbox service's URL
@@ -51,16 +53,16 @@ toolbox = ToolboxClient("http://127.0.0.1:5000")
 > [!TIP]
 > You can also pass your own `ClientSession` so that the `ToolboxClient` can
 > reuse the same session.
-> ```
+> ```py
 > async with ClientSession() as session:
->   client = ToolboxClient(http://localhost:5000, session)
+>   toolbox = ToolboxClient("http://localhost:5000", session)
 > ```
 
 ## Load a toolset
 
 You can load a toolset, a collection of related tools.
 
-```python
+```py
 # Load all tools
 tools = await toolbox.load_toolset()
 
@@ -72,21 +74,21 @@ tools = await toolbox.load_toolset("my-toolset")
 
 You can also load a single tool.
 
-```python
+```py
 tool = await toolbox.load_tool("my-tool")
 ```
 
 ## Use with LlamaIndex
 
 LlamaIndex agents can dynamically choose and execute tools based on the user
-input. The user can include the tools loaded from the Toolbox SDK in the
-agent's toolkit.
+input. The user can include the tools loaded from the Toolbox SDK in the agent's
+toolkit.
 
-```python
+```py
 from llama_index.llms.vertex import Vertex
 from llama_index.core.agent import ReActAgent
 
-model = Vertex(model="gemini-1.5-pro")
+model = Vertex(model="gemini-1.5-flash")
 
 # Initialize agent with tools
 agent = ReActAgent.from_tools(tools, llm=model, verbose=True)
@@ -99,6 +101,97 @@ response = agent.query("Get some response from the agent.")
 
 You can also execute a tool manually using the `acall` method.
 
-```python
-result = await tools[0].acall({"param1": "value1", "param2": "value2"})
+```py
+result = await tools[0].acall({ "name": "Alice", "age": 30 })
+```
+
+## Authenticating Tools
+
+> [!WARNING]
+> Always use HTTPS to connect your application with the Toolbox service,
+> especially when using tools with authentication configured. Using HTTP exposes
+> your application to serious security risks, including unauthorized access to
+> user information and man-in-the-middle attacks, where sensitive data can be
+> intercepted.
+
+Some tools in your Toolbox configuration might require user authentication to
+access sensitive data. This section guides you on how to configure tools for
+authentication and use them with the SDK.
+
+### Supported Authentication Mechanisms
+The Toolbox SDK currently supports authentication using [OIDC
+protocol](https://openid.net/specs/openid-connect-core-1_0.html). Specifically,
+it uses [ID
+tokens](https://openid.net/specs/openid-connect-core-1_0.html#IDToken) and *not*
+access tokens for [Google OAuth
+2.0](https://cloud.google.com/apigee/docs/api-platform/security/oauth/oauth-home).
+
+### Configuring Tools for Authentication
+
+Refer to [these
+instructions](../../docs/tools/README.md#authenticated-parameters) on
+configuring tools for authenticated parameters.
+
+### Configure SDK for Authentication
+
+Provide the `auth_headers` parameter to the `load_tool` or `load_toolset` calls
+with a dictionary. The keys of this dictionary should match the names of the
+authentication sources configured in your tools file (e.g., `my_auth_service`),
+and the values should be callable functions (e.g., lambdas or regular functions)
+that return the ID token of the logged-in user.
+
+Here's an example:
+
+```py
+def get_auth_header():
+    # ... Logic to retrieve ID token (e.g., from local storage, OAuth flow)
+    # This example just returns a placeholder. Replace with your actual token retrieval.
+    return "YOUR_ID_TOKEN"
+
+toolbox = ToolboxClient("http://localhost:5000")
+
+tools = toolbox.load_toolset(auth_headers={ "my_auth_service": get_auth_header })
+
+# OR
+
+tool = toolbox.load_tool("my_tool", auth_headers={ "my_auth_service": get_auth_header })
 ```
+
+Alternatively, you can call the `add_auth_header` method to configure
+authentication separately.
+
+```py
+toolbox.add_auth_header("my_auth_service", get_auth_header)
+```
+
+> [!NOTE]
+> Authentication headers added via `load_tool`, `load_toolset`, or
+> `add_auth_header` apply to all subsequent tool invocations, regardless of when
+> the tool was loaded. This ensures a consistent authentication context.
+
+### Complete Example
+
+```py
+import asyncio
+from toolbox_llamaindex_sdk import ToolboxClient
+
+async def get_auth_header():
+    # Replace with your actual ID token retrieval logic.
+    # For example, using a library like google-auth
+    # from google.oauth2 import id_token
+    # from google.auth.transport import requests
+    # request = requests.Request()
+    # id_token_string = id_token.fetch_id_token(request, "YOUR_AUDIENCE")# Replace with your audience
+    # return id_token_string
+    return "YOUR_ACTUAL_ID_TOKEN" # placeholder
+
+async def main():
+    toolbox = ToolboxClient("http://localhost:5000")
+    toolbox.add_auth_header("my_auth_service", get_auth_header)
+    tools = await toolbox.load_toolset()
+    result = await tools[0].acall({"input": "some input"})
+    print(result)
+
+if __name__ == "__main__":
+    asyncio.run(main())
+```
diff --git a/sdks/llamaindex/pyproject.toml b/sdks/llamaindex/pyproject.toml
@@ -1,7 +1,7 @@
 [project]
 name = "toolbox-llamaindex-sdk"
 version="0.0.1"
-description = "Python SDK for interacting with the Toolbox service with Llamaindex"
+description = "Python SDK for interacting with the Toolbox service with LlamaIndex"
 license = {file = "LICENSE"}
 requires-python = ">=3.9"
 authors = [

diff --git a/sdks/llamaindex/requirements.txt b/sdks/llamaindex/requirements.txt
@@ -1,4 +1,4 @@
 llama-index==0.12.2
 PyYAML==6.0.2
-pydantic==2.9.2
-aiohttp==3.11.7
+pydantic==2.10.2
+aiohttp==3.11.7
diff --git a/sdks/llamaindex/src/toolbox_llamaindex_sdk/__init__.py b/sdks/llamaindex/src/toolbox_llamaindex_sdk/__init__.py
@@ -1,5 +1,3 @@
 from .client import ToolboxClient
 
-# import utils
-
 __all__ = ["ToolboxClient"]