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.

Author: anubhav756

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`.
 

pyproject.toml

@@ -8,10 +8,10 @@ authors = [
     {name = "Google LLC", email = "[email protected]"}
 ]
 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 = [

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

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"})