feat(langchain-sdk): Add Toolbox SDK for LangChain (#22)

Adds a initial python SDK for interacting with Toolbox from LangChain.
This commit is contained in:
Anubhav Dhawan
2024-10-28 21:20:41 +05:30
committed by GitHub
parent 5f565e43df
commit 0bcd4b6e41
5 changed files with 1224 additions and 0 deletions

View File

@@ -0,0 +1,57 @@
[project]
name = "toolbox_langchain_sdk"
version="0.0.1"
description = "Python SDK for interacting with the Toolbox service with LangChain"
license = {file = "LICENSE"}
requires-python = ">=3.9"
authors = [
{name = "Google LLC", email = "googleapis-packages@google.com"}
]
dependencies = [
"aiohttp",
"PyYAML",
"langchain-core",
"pydantic",
]
classifiers = [
"Intended Audience :: Developers",
"License :: OSI Approved :: Apache Software License",
"Programming Language :: Python",
"Programming Language :: Python :: 3",
"Programming Language :: Python :: 3.9",
"Programming Language :: Python :: 3.10",
"Programming Language :: Python :: 3.11",
"Programming Language :: Python :: 3.12",
]
[project.urls]
Homepage = "https://github.com/googleapis/genai-toolbox"
Repository = "https://github.com/googleapis/genai-toolbox.git"
"Bug Tracker" = "https://github.com/googleapis/genai-toolbox/issues"
[project.optional-dependencies]
test = [
"black[jupyter]",
"isort",
"mypy",
"pytest-asyncio",
"pytest",
"pytest-cov",
"Pillow"
]
[build-system]
requires = ["setuptools"]
build-backend = "setuptools.build_meta"
[tool.black]
target-version = ['py39']
[tool.isort]
profile = "black"
[tool.mypy]
python_version = "3.9"
warn_unused_configs = true
disallow_incomplete_defs = true

View File

@@ -0,0 +1,3 @@
from .client import ToolboxClient
__all__ = ["ToolboxClient"]

View File

@@ -0,0 +1,110 @@
from typing import Optional
from aiohttp import ClientSession
from langchain_core.tools import StructuredTool
from pydantic import BaseModel
from .utils import ManifestSchema, _invoke_tool, _load_yaml, _schema_to_model
class ToolboxClient:
def __init__(self, url: str, session: ClientSession):
"""
Initializes the ToolboxClient for the Toolbox service at the given URL.
Args:
url: The base URL of the Toolbox service.
session: The HTTP client session.
"""
self._url: str = url
self._session = session
async def _load_tool_manifest(self, tool_name: str) -> ManifestSchema:
"""
Fetches and parses the YAML manifest for the given tool from the Toolbox service.
Args:
tool_name: The name of the tool to load.
Returns:
The parsed Toolbox manifest.
"""
url = f"{self._url}/api/tool/{tool_name}"
return await _load_yaml(url, self._session)
async def _load_toolset_manifest(
self, toolset_name: Optional[str] = None
) -> ManifestSchema:
"""
Fetches and parses the YAML manifest from the Toolbox service.
Args:
toolset_name: The name of the toolset to load.
Default: None. If not provided, then all the available tools are loaded.
Returns:
The parsed Toolbox manifest.
"""
url = f"{self._url}/api/toolset/{toolset_name or ''}"
return await _load_yaml(url, self._session)
def _generate_tool(
self, tool_name: str, manifest: ManifestSchema
) -> StructuredTool:
"""
Creates a StructuredTool object and a dynamically generated BaseModel for the given tool.
Args:
tool_name: The name of the tool to generate.
manifest: The parsed Toolbox manifest.
Returns:
The generated tool.
"""
tool_schema = manifest.tools[tool_name]
tool_model: BaseModel = _schema_to_model(
model_name=tool_name, schema=tool_schema.parameters
)
async def _tool_func(**kwargs) -> dict:
return await _invoke_tool(self._url, self._session, tool_name, kwargs)
return StructuredTool.from_function(
coroutine=_tool_func,
name=tool_name,
description=tool_schema.description,
args_schema=tool_model,
)
async def load_tool(self, tool_name: str) -> StructuredTool:
"""
Loads the tool, with the given tool name, from the Toolbox service.
Args:
toolset_name: The name of the toolset to load.
Default: None. If not provided, then all the tools are loaded.
Returns:
A tool loaded from the Toolbox
"""
manifest: ManifestSchema = await self._load_tool_manifest(tool_name)
return self._generate_tool(tool_name, manifest)
async def load_toolset(
self, toolset_name: Optional[str] = None
) -> list[StructuredTool]:
"""
Loads tools from the Toolbox service, optionally filtered by toolset name.
Args:
toolset_name: The name of the toolset to load.
Default: None. If not provided, then all the tools are loaded.
Returns:
A list of all tools loaded from the Toolbox.
"""
tools: list[StructuredTool] = []
manifest: ManifestSchema = await self._load_toolset_manifest(toolset_name)
for tool_name in manifest.tools:
tools.append(self._generate_tool(tool_name, manifest))
return tools

View File

@@ -0,0 +1,117 @@
from typing import Any, Type, Optional
import yaml
from aiohttp import ClientSession
from pydantic import BaseModel, Field, create_model
class ParameterSchema(BaseModel):
name: str
type: str
description: str
class ToolSchema(BaseModel):
description: str
parameters: list[ParameterSchema]
class ManifestSchema(BaseModel):
serverVersion: str
tools: dict[str, ToolSchema]
async def _load_yaml(url: str, session: ClientSession) -> ManifestSchema:
"""
Asynchronously fetches and parses the YAML data from the given URL.
Args:
url: The base URL to fetch the YAML from.
session: The HTTP client session
Returns:
The parsed Toolbox manifest.
"""
async with session.get(url) as response:
response.raise_for_status()
parsed_yaml = yaml.safe_load(await response.text())
return ManifestSchema(**parsed_yaml)
def _schema_to_model(model_name: str, schema: list[ParameterSchema]) -> Type[BaseModel]:
"""
Converts a schema (from the YAML manifest) to a Pydantic BaseModel class.
Args:
model_name: The name of the model to create.
schema: The schema to convert.
Returns:
A Pydantic BaseModel class.
"""
field_definitions = {}
for field in schema:
field_definitions[field.name] = (
# TODO: Remove the hardcoded optional types once optional fields are supported by Toolbox.
Optional[_parse_type(field.type)],
Field(description=field.description),
)
return create_model(model_name, **field_definitions)
def _parse_type(type_: str) -> Any:
"""
Converts a schema type to a JSON type.
Args:
type_: The type name to convert.
Returns:
A valid JSON type.
"""
if type_ == "string":
return str
elif type_ == "integer":
return int
elif type_ == "number":
return float
elif type_ == "boolean":
return bool
elif type_ == "array":
return list
else:
raise ValueError(f"Unsupported schema type: {type_}")
async def _invoke_tool(
url: str, session: ClientSession, tool_name: str, data: dict
) -> dict:
"""
Asynchronously makes an API call to the Toolbox service to invoke a tool.
Args:
url: The base URL of the Toolbox service.
session: The HTTP client session.
tool_name: The name of the tool to invoke.
data: The input data for the tool.
Returns:
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:
response.raise_for_status()
return await response.json()
# TODO: Remove this temporary fix once optional fields are supported by Toolbox.
def _convert_none_to_empty_string(input_dict):
new_dict = {}
for key, value in input_dict.items():
if value is None:
new_dict[key] = ""
else:
new_dict[key] = value
return new_dict

View File

@@ -0,0 +1,937 @@
import asyncio
from unittest.mock import AsyncMock, Mock, call, patch
import aiohttp
import pytest
from langchain_core.tools import StructuredTool
from pydantic import ValidationError
from toolbox_langchain_sdk import ToolboxClient
@pytest.mark.asyncio
@patch("toolbox_langchain_sdk.utils.aiohttp.ClientSession.get")
async def test_load_tool_manifest(mock_get):
client = ToolboxClient("https://my-toolbox.com")
mock_response = aiohttp.ClientResponse(
method="GET",
url=aiohttp.client.URL("https://my-toolbox.com/api/tool/test_tool"),
writer=None,
continue100=None,
timer=None,
request_info=None,
traces=None,
session=None,
loop=asyncio.get_event_loop(),
)
mock_response.raise_for_status = Mock()
mock_response.text = AsyncMock(
return_value="""
{
"serverVersion": "0.0.1",
"tools": {
"search_airport": {
"description": "Use this tool to list all airports matching search criteria.\nTakes a country and a city and returns all matching airports.\nThe agent can decide to return the results directly to the user.\nInput of this tool must be in JSON format and include both inputs - country and city.\nExample:\n{{\n \"country\": \"United States\",\n \"city\": \"San Francisco\",\n}}\n",
"parameters": [
{
"name": "country",
"type": "string",
"description": "Country"
},
{
"name": "city",
"type": "string",
"description": "City"
}
]
}
}
}
"""
)
mock_get.return_value = mock_response
await client._load_tool_manifest("test_tool")
mock_get.assert_called_once_with("https://my-toolbox.com/api/tool/test_tool")
assert client._manifest["serverVersion"] == "0.0.1"
assert "tools" in client._manifest
assert len(client._manifest["tools"].keys()) == 1
assert "test_tool" in client._manifest["tools"]
tool = client._manifest["tools"]["test_tool"]
assert "summary" in tool
assert "description" in tool
assert "parameters" in tool
assert tool["summary"] == "Test Tool"
assert tool["description"] == "This is a test tool."
assert len(tool["parameters"].keys()) == 2
assert "param1" in tool["parameters"]
assert "type" in tool["parameters"]["param1"]
assert "description" in tool["parameters"]["param1"]
assert tool["parameters"]["param1"]["type"] == "string"
assert tool["parameters"]["param1"]["description"] == "Parameter 1"
assert "param2" in tool["parameters"]
assert "type" in tool["parameters"]["param2"]
assert "description" in tool["parameters"]["param2"]
assert tool["parameters"]["param2"]["type"] == "integer"
assert tool["parameters"]["param2"]["description"] == "Parameter 2"
@pytest.mark.asyncio
@patch("toolbox_langchain_sdk.utils.aiohttp.ClientSession.get")
async def test_load_multiple_tool_manifest(mock_get):
client = ToolboxClient("https://my-toolbox.com")
mock_response = aiohttp.ClientResponse(
method="GET",
url=aiohttp.client.URL("https://my-toolbox.com/api/tool/test_tool1"),
writer=None,
continue100=None,
timer=None,
request_info=None,
traces=None,
session=None,
loop=asyncio.get_event_loop(),
)
mock_response.raise_for_status = Mock()
mock_response.text = AsyncMock(
return_value="""
serverVersion: 0.0.1
tools:
test_tool1:
summary: Test Tool 1
description: This is a test tool 1.
parameters:
param1:
type: string
description: Parameter 1
param2:
type: integer
description: Parameter 2
"""
)
mock_get.return_value = mock_response
await client._load_tool_manifest("test_tool1")
mock_get.assert_called_once_with("https://my-toolbox.com/api/tool/test_tool1")
assert client._manifest["serverVersion"] == "0.0.1"
assert "tools" in client._manifest
assert len(client._manifest["tools"].keys()) == 1
assert "test_tool1" in client._manifest["tools"]
tool1 = client._manifest["tools"]["test_tool1"]
assert "summary" in tool1
assert "description" in tool1
assert "parameters" in tool1
assert tool1["summary"] == "Test Tool 1"
assert tool1["description"] == "This is a test tool 1."
assert len(tool1["parameters"].keys()) == 2
assert "param1" in tool1["parameters"]
assert "type" in tool1["parameters"]["param1"]
assert "description" in tool1["parameters"]["param1"]
assert tool1["parameters"]["param1"]["type"] == "string"
assert tool1["parameters"]["param1"]["description"] == "Parameter 1"
assert "param2" in tool1["parameters"]
assert "type" in tool1["parameters"]["param2"]
assert "description" in tool1["parameters"]["param2"]
assert tool1["parameters"]["param2"]["type"] == "integer"
assert tool1["parameters"]["param2"]["description"] == "Parameter 2"
mock_response = aiohttp.ClientResponse(
method="GET",
url=aiohttp.client.URL("https://my-toolbox.com/api/tool/test_tool2"),
writer=None,
continue100=None,
timer=None,
request_info=None,
traces=None,
session=None,
loop=asyncio.get_event_loop(),
)
mock_response.raise_for_status = Mock()
mock_response.text = AsyncMock(
return_value="""
serverVersion: 0.0.1
tools:
test_tool2:
summary: Test Tool 2
description: This is a test tool 2.
parameters:
param1:
type: integer
description: Parameter 1
param2:
type: string
description: Parameter 2
"""
)
mock_get.return_value = mock_response
await client._load_tool_manifest("test_tool2")
assert mock_get.call_count == 2
mock_get.assert_has_calls(
[
call("https://my-toolbox.com/api/tool/test_tool1"),
call("https://my-toolbox.com/api/tool/test_tool2"),
]
)
assert client._manifest["serverVersion"] == "0.0.1"
assert "tools" in client._manifest
assert len(client._manifest["tools"].keys()) == 2
assert "test_tool1" in client._manifest["tools"]
tool1 = client._manifest["tools"]["test_tool1"]
assert "summary" in tool1
assert "description" in tool1
assert "parameters" in tool1
assert tool1["summary"] == "Test Tool 1"
assert tool1["description"] == "This is a test tool 1."
assert len(tool1["parameters"].keys()) == 2
assert "param1" in tool1["parameters"]
assert "type" in tool1["parameters"]["param1"]
assert "description" in tool1["parameters"]["param1"]
assert tool1["parameters"]["param1"]["type"] == "string"
assert tool1["parameters"]["param1"]["description"] == "Parameter 1"
assert "param2" in tool1["parameters"]
assert "type" in tool1["parameters"]["param2"]
assert "description" in tool1["parameters"]["param2"]
assert tool1["parameters"]["param2"]["type"] == "integer"
assert tool1["parameters"]["param2"]["description"] == "Parameter 2"
assert "test_tool2" in client._manifest["tools"]
tool2 = client._manifest["tools"]["test_tool2"]
assert "summary" in tool2
assert "description" in tool2
assert "parameters" in tool2
assert tool2["summary"] == "Test Tool 2"
assert tool2["description"] == "This is a test tool 2."
assert len(tool2["parameters"].keys()) == 2
assert "param1" in tool2["parameters"]
assert "type" in tool2["parameters"]["param1"]
assert "description" in tool2["parameters"]["param1"]
assert tool2["parameters"]["param1"]["type"] == "integer"
assert tool2["parameters"]["param1"]["description"] == "Parameter 1"
assert "param2" in tool2["parameters"]
assert "type" in tool2["parameters"]["param2"]
assert "description" in tool2["parameters"]["param2"]
assert tool2["parameters"]["param2"]["type"] == "string"
assert tool2["parameters"]["param2"]["description"] == "Parameter 2"
@pytest.mark.asyncio
@patch("toolbox_langchain_sdk.utils.aiohttp.ClientSession.get")
async def test_load_tool_manifest_invalid_yaml(mock_get):
client = ToolboxClient("https://my-toolbox.com")
mock_response = aiohttp.ClientResponse(
method="GET",
url=aiohttp.client.URL("https://my-toolbox.com/api/tool/test_tool"),
writer=None,
continue100=None,
timer=None,
request_info=None,
traces=None,
session=None,
loop=asyncio.get_event_loop(),
)
mock_response.raise_for_status = Mock()
mock_response.text = AsyncMock(return_value="invalid yaml")
mock_get.return_value = mock_response
await client._load_tool_manifest("test_tool")
mock_get.assert_called_once_with("https://my-toolbox.com/api/tool/test_tool")
assert client._manifest == "invalid yaml"
@pytest.mark.asyncio
@patch("toolbox_langchain_sdk.utils.aiohttp.ClientSession.get")
async def test_load_tool_manifest_api_error(mock_get):
client = ToolboxClient("https://my-toolbox.com")
mock_response = aiohttp.ClientResponse(
method="GET",
url=aiohttp.client.URL("https://my-toolbox.com/api/tool/test_tool"),
writer=None,
continue100=None,
timer=None,
request_info=None,
traces=None,
session=None,
loop=asyncio.get_event_loop(),
)
error = aiohttp.ClientError("Simulated HTTP Error")
mock_response.raise_for_status = Mock()
mock_response.text = AsyncMock(side_effect=error)
mock_get.return_value = mock_response
with pytest.raises(aiohttp.ClientError) as exc_info:
await client._load_tool_manifest("test_tool")
mock_get.assert_called_once_with("https://my-toolbox.com/api/tool/test_tool")
assert exc_info.value == error
@pytest.mark.asyncio
@patch("toolbox_langchain_sdk.utils.aiohttp.ClientSession.get")
async def test_load_tool_manifest_valid_then_invalid_yaml(mock_get):
client = ToolboxClient("https://my-toolbox.com")
mock_response = aiohttp.ClientResponse(
method="GET",
url=aiohttp.client.URL("https://my-toolbox.com/api/tool/test_tool1"),
writer=None,
continue100=None,
timer=None,
request_info=None,
traces=None,
session=None,
loop=asyncio.get_event_loop(),
)
mock_response.raise_for_status = Mock()
mock_response.text = AsyncMock(
return_value="""
serverVersion: 0.0.1
tools:
test_tool1:
summary: Test Tool 1
description: This is a test tool 1.
parameters:
param1:
type: integer
description: Parameter 1
param2:
type: string
description: Parameter 2
"""
)
mock_get.return_value = mock_response
await client._load_tool_manifest("test_tool1")
mock_get.assert_called_once_with("https://my-toolbox.com/api/tool/test_tool1")
assert client._manifest["serverVersion"] == "0.0.1"
assert "tools" in client._manifest
assert len(client._manifest["tools"].keys()) == 1
assert "test_tool1" in client._manifest["tools"]
tool1 = client._manifest["tools"]["test_tool1"]
assert "summary" in tool1
assert "description" in tool1
assert "parameters" in tool1
assert tool1["summary"] == "Test Tool 1"
assert tool1["description"] == "This is a test tool 1."
assert len(tool1["parameters"].keys()) == 2
assert "param1" in tool1["parameters"]
assert "type" in tool1["parameters"]["param1"]
assert "description" in tool1["parameters"]["param1"]
assert tool1["parameters"]["param1"]["type"] == "integer"
assert tool1["parameters"]["param1"]["description"] == "Parameter 1"
assert "param2" in tool1["parameters"]
assert "type" in tool1["parameters"]["param2"]
assert "description" in tool1["parameters"]["param2"]
assert tool1["parameters"]["param2"]["type"] == "string"
assert tool1["parameters"]["param2"]["description"] == "Parameter 2"
mock_response = aiohttp.ClientResponse(
method="GET",
url=aiohttp.client.URL("https://my-toolbox.com/api/tool/test_tool2"),
writer=None,
continue100=None,
timer=None,
request_info=None,
traces=None,
session=None,
loop=asyncio.get_event_loop(),
)
mock_response.raise_for_status = Mock()
mock_response.text = AsyncMock(return_value="invalid yaml")
mock_get.return_value = mock_response
await client._load_tool_manifest("test_tool2")
assert mock_get.call_count == 2
mock_get.assert_has_calls(
[
call("https://my-toolbox.com/api/tool/test_tool1"),
call("https://my-toolbox.com/api/tool/test_tool2"),
]
)
assert client._manifest == "invalid yaml"
@pytest.mark.asyncio
@patch("toolbox_langchain_sdk.utils.aiohttp.ClientSession.get")
async def test_load_toolset_manifest(mock_get):
client = ToolboxClient("https://my-toolbox.com")
mock_response = aiohttp.ClientResponse(
method="GET",
url=aiohttp.client.URL("https://my-toolbox.com/api/toolset/test_toolset"),
writer=None,
continue100=None,
timer=None,
request_info=None,
traces=None,
session=None,
loop=asyncio.get_event_loop(),
)
mock_response.raise_for_status = Mock()
mock_response.text = AsyncMock(
return_value="""
serverVersion: 0.0.1
tools:
test_tool:
summary: Test Tool
description: This is a test tool.
parameters:
param1:
type: string
description: Parameter 1
param2:
type: integer
description: Parameter 2
"""
)
mock_get.return_value = mock_response
await client._load_toolset_manifest("test_toolset")
mock_get.assert_called_once_with("https://my-toolbox.com/api/toolset/test_toolset")
assert client._manifest["serverVersion"] == "0.0.1"
assert "tools" in client._manifest
assert len(client._manifest["tools"].keys()) == 1
assert "test_tool" in client._manifest["tools"]
tool = client._manifest["tools"]["test_tool"]
assert "summary" in tool
assert "description" in tool
assert "parameters" in tool
assert tool["summary"] == "Test Tool"
assert tool["description"] == "This is a test tool."
assert len(tool["parameters"].keys()) == 2
assert "param1" in tool["parameters"]
assert "type" in tool["parameters"]["param1"]
assert "description" in tool["parameters"]["param1"]
assert tool["parameters"]["param1"]["type"] == "string"
assert tool["parameters"]["param1"]["description"] == "Parameter 1"
assert "param2" in tool["parameters"]
assert "type" in tool["parameters"]["param2"]
assert "description" in tool["parameters"]["param2"]
assert tool["parameters"]["param2"]["type"] == "integer"
assert tool["parameters"]["param2"]["description"] == "Parameter 2"
@pytest.mark.asyncio
@patch("toolbox_langchain_sdk.utils.aiohttp.ClientSession.get")
async def test_load_toolset_manifest_all_toolsets(mock_get):
client = ToolboxClient("https://my-toolbox.com")
mock_response = aiohttp.ClientResponse(
method="GET",
url=aiohttp.client.URL("https://my-toolbox.com/api/toolset"),
writer=None,
continue100=None,
timer=None,
request_info=None,
traces=None,
session=None,
loop=asyncio.get_event_loop(),
)
mock_response.raise_for_status = Mock()
mock_response.text = AsyncMock(
return_value="""
serverVersion: 0.0.1
tools:
test_tool1:
summary: Test Tool 1
description: This is a test tool 1.
parameters:
param1:
type: string
description: Parameter 1
test_tool2:
summary: Test Tool 2
description: This is a test tool 2.
parameters:
param2:
type: integer
description: Parameter 2
"""
)
mock_get.return_value = mock_response
await client._load_toolset_manifest()
mock_get.assert_called_once_with("https://my-toolbox.com/api/toolset")
assert client._manifest["serverVersion"] == "0.0.1"
assert "tools" in client._manifest
assert len(client._manifest["tools"].keys()) == 2
assert "test_tool1" in client._manifest["tools"]
assert "test_tool2" in client._manifest["tools"]
tool1 = client._manifest["tools"]["test_tool1"]
assert "summary" in tool1
assert "description" in tool1
assert "parameters" in tool1
assert tool1["summary"] == "Test Tool 1"
assert tool1["description"] == "This is a test tool 1."
assert len(tool1["parameters"].keys()) == 1
assert "param1" in tool1["parameters"]
assert "type" in tool1["parameters"]["param1"]
assert "description" in tool1["parameters"]["param1"]
assert tool1["parameters"]["param1"]["type"] == "string"
assert tool1["parameters"]["param1"]["description"] == "Parameter 1"
tool2 = client._manifest["tools"]["test_tool2"]
assert "summary" in tool2
assert "description" in tool2
assert "parameters" in tool2
assert tool2["summary"] == "Test Tool 2"
assert tool2["description"] == "This is a test tool 2."
assert len(tool2["parameters"].keys()) == 1
assert "param2" in tool2["parameters"]
assert "type" in tool2["parameters"]["param2"]
assert "description" in tool2["parameters"]["param2"]
assert tool2["parameters"]["param2"]["type"] == "integer"
assert tool2["parameters"]["param2"]["description"] == "Parameter 2"
@pytest.mark.asyncio
@patch("toolbox_langchain_sdk.utils.aiohttp.ClientSession.get")
async def test_load_toolset_manifest_invalid_yaml(mock_get):
client = ToolboxClient("https://my-toolbox.com")
mock_response = aiohttp.ClientResponse(
method="GET",
url=aiohttp.client.URL("https://my-toolbox.com/api/toolset/test_toolset"),
writer=None,
continue100=None,
timer=None,
request_info=None,
traces=None,
session=None,
loop=asyncio.get_event_loop(),
)
mock_response.raise_for_status = Mock()
mock_response.text = AsyncMock(return_value="invalid yaml")
mock_get.return_value = mock_response
await client._load_toolset_manifest("test_toolset")
mock_get.assert_called_once_with("https://my-toolbox.com/api/toolset/test_toolset")
assert client._manifest == "invalid yaml"
@pytest.mark.asyncio
@patch("toolbox_langchain_sdk.utils.aiohttp.ClientSession.get")
async def test_load_toolset_manifest_api_error(mock_get):
client = ToolboxClient("https://my-toolbox.com")
mock_response = aiohttp.ClientResponse(
method="GET",
url=aiohttp.client.URL("https://my-toolbox.com/api/toolset/test_toolset"),
writer=None,
continue100=None,
timer=None,
request_info=None,
traces=None,
session=None,
loop=asyncio.get_event_loop(),
)
error = aiohttp.ClientError("Simulated HTTP Error")
mock_response.raise_for_status = Mock(side_effect=error)
mock_get.return_value = mock_response
with pytest.raises(aiohttp.ClientError) as exc_info:
await client._load_toolset_manifest("test_toolset")
mock_get.assert_called_once_with("https://my-toolbox.com/api/toolset/test_toolset")
assert exc_info.value == error
@pytest.mark.asyncio
@patch("toolbox_langchain_sdk.utils.aiohttp.ClientSession.post")
async def test_generate_tool_success(mock_post):
client = ToolboxClient("https://my-toolbox.com")
client._manifest = {
"tools": {
"test_tool": {
"summary": "Test Tool",
"description": "This is a test tool.",
"parameters": {
"param1": {"type": "string", "description": "Parameter 1"},
"param2": {"type": "integer", "description": "Parameter 2"},
},
},
},
}
mock_response = aiohttp.ClientResponse(
method="POST",
url=aiohttp.client.URL("https://my-toolbox.com/api/tool/test_tool"),
writer=None,
continue100=None,
timer=None,
request_info=None,
traces=None,
session=None,
loop=asyncio.get_event_loop(),
)
mock_response.status = 200
mock_response.json = AsyncMock(return_value={"result": "mocked_result"})
mock_post.return_value = mock_response
client._generate_tool("test_tool")
assert len(client._tools) == 1
tool = client._tools[0]
assert isinstance(tool, StructuredTool)
assert tool.name == "Test Tool"
assert tool.description == "This is a test tool."
assert tool.args_schema.model_fields.keys() == {"param1", "param2"}
assert tool.args_schema.model_fields["param1"].annotation == str
assert tool.args_schema.model_fields["param2"].annotation == int
assert tool.args_schema.model_fields["param1"].description == "Parameter 1"
assert tool.args_schema.model_fields["param2"].description == "Parameter 2"
params = {"param1": "value1", "param2": 123}
result = await tool.arun(params)
mock_post.assert_called_once_with(
"https://my-toolbox.com/api/tool/test_tool", json=params
)
assert result == mock_response.json.return_value
@pytest.mark.asyncio
@patch("toolbox_langchain_sdk.utils.aiohttp.ClientSession.post")
async def test_generate_tool_api_error(mock_post):
client = ToolboxClient("https://my-toolbox.com")
client._manifest = {
"tools": {
"test_tool": {
"summary": "Test Tool",
"description": "This is a test tool.",
"parameters": {
"param1": {"type": "string", "description": "Parameter 1"},
"param2": {"type": "integer", "description": "Parameter 2"},
},
},
},
}
mock_response = aiohttp.ClientResponse(
method="POST",
url=aiohttp.client.URL("https://my-toolbox.com/api/tool/test_tool"),
writer=None,
continue100=None,
timer=None,
request_info=None,
traces=None,
session=None,
loop=asyncio.get_event_loop(),
)
mock_response.status = 200
error = aiohttp.ClientError("Simulated HTTP Error")
mock_response.json = AsyncMock(side_effect=error)
mock_post.return_value = mock_response
client._generate_tool("test_tool")
assert len(client._tools) == 1
tool = client._tools[0]
assert isinstance(tool, StructuredTool)
assert tool.name == "Test Tool"
assert tool.description == "This is a test tool."
assert tool.args_schema.model_fields.keys() == {"param1", "param2"}
assert tool.args_schema.model_fields["param1"].annotation == str
assert tool.args_schema.model_fields["param2"].annotation == int
assert tool.args_schema.model_fields["param1"].description == "Parameter 1"
assert tool.args_schema.model_fields["param2"].description == "Parameter 2"
with pytest.raises(aiohttp.ClientError) as exc_info:
params = {"param1": "test", "param2": 123}
await tool.arun(params)
mock_post.assert_called_once_with(
"https://my-toolbox.com/api/tool/test_tool",
json=params,
)
assert exc_info.value == error
def test_generate_tool_missing_schema_fields():
client = ToolboxClient("https://my-toolbox.com")
client._manifest = {"tools": {"test_tool": {"summary": "Test Tool"}}}
with pytest.raises(ValidationError) as exc_info:
client._generate_tool("test_tool")
errors = exc_info.value.errors()
assert len(errors) == 2
assert errors[0]["input"] == client._manifest["tools"]["test_tool"]
assert errors[0]["loc"] == ("description",)
assert errors[0]["msg"] == "Field required"
assert errors[0]["type"] == "missing"
assert errors[1]["input"] == client._manifest["tools"]["test_tool"]
assert errors[1]["loc"] == ("parameters",)
assert errors[1]["msg"] == "Field required"
assert errors[1]["type"] == "missing"
def test_generate_tool_invalid_schema_types():
client = ToolboxClient("https://my-toolbox.com")
client._manifest = {
"tools": {
"test_tool": {
"summary": 123,
"description": "This is a test tool.",
"parameters": {
"param1": {"type": "string", "description": "Parameter 1"},
"param2": {"type": "integer", "description": "Parameter 2"},
},
},
},
}
with pytest.raises(ValidationError) as exc_info:
client._generate_tool("test_tool")
errors = exc_info.value.errors()
assert len(errors) == 1
assert errors[0]["loc"] == ("summary",)
assert errors[0]["input"] == 123
assert errors[0]["msg"] == "Input should be a valid string"
@pytest.mark.asyncio
@patch("toolbox_langchain_sdk.utils.aiohttp.ClientSession.post")
async def test_generate_tool_invalid_parameter_types(mock_post):
client = ToolboxClient("https://my-toolbox.com")
client._manifest = {
"tools": {
"test_tool": {
"summary": "Test Tool",
"description": "This is a test tool.",
"parameters": {
"param1": {"type": "string", "description": "Parameter 1"},
"param2": {"type": "integer", "description": "Parameter 2"},
},
},
},
}
client._generate_tool("test_tool")
assert len(client._tools) == 1
tool = client._tools[0]
assert isinstance(tool, StructuredTool)
assert tool.name == "Test Tool"
assert tool.description == "This is a test tool."
assert tool.args_schema.model_fields.keys() == {"param1", "param2"}
assert tool.args_schema.model_fields["param1"].annotation == str
assert tool.args_schema.model_fields["param2"].annotation == int
assert tool.args_schema.model_fields["param1"].description == "Parameter 1"
assert tool.args_schema.model_fields["param2"].description == "Parameter 2"
with pytest.raises(ValidationError) as exc_info:
await tool.arun({"param1": "test", "param2": "abc"})
mock_post.assert_not_called()
errors = exc_info.value.errors()
assert len(errors) == 1
assert errors[0]["loc"] == ("param2",)
assert errors[0]["input"] == "abc"
assert (
errors[0]["msg"]
== "Input should be a valid integer, unable to parse string as an integer"
)
@pytest.mark.asyncio
@patch.object(ToolboxClient, "_load_tool_manifest")
async def test_load_tool(mock_load_tool_manifest):
client = ToolboxClient("https://my-toolbox.com")
mock_load_tool_manifest.side_effect = lambda _: setattr(
client,
"_manifest",
{
"serverVersion": "0.0.1",
"tools": {
"test_tool": {
"summary": "Test Tool",
"description": "This is a test tool.",
"parameters": {
"param1": {"type": "string", "description": "Parameter 1"},
"param2": {"type": "integer", "description": "Parameter 2"},
},
},
},
},
)
tool = await client.load_tool("test_tool")
mock_load_tool_manifest.assert_called_once_with("test_tool")
assert isinstance(tool, StructuredTool)
assert tool.name == "Test Tool"
assert tool.description == "This is a test tool."
assert tool.args == {
"param1": {"title": "Param1", "description": "Parameter 1", "type": "string"},
"param2": {"title": "Param2", "description": "Parameter 2", "type": "integer"},
}
@pytest.mark.asyncio
@patch.object(ToolboxClient, "_load_tool_manifest")
async def test_load_multiple_tools(mock_load_tool_manifest):
client = ToolboxClient("https://my-toolbox.com")
mock_load_tool_manifest.side_effect = lambda _: setattr(
client,
"_manifest",
{
"serverVersion": "0.0.1",
"tools": {
"test_tool1": {
"summary": "Test Tool 1",
"description": "This is a test tool 1.",
"parameters": {
"param1": {"type": "string", "description": "Parameter 1"},
"param2": {"type": "integer", "description": "Parameter 2"},
},
},
},
},
)
tool1 = await client.load_tool("test_tool1")
mock_load_tool_manifest.assert_called_once_with("test_tool1")
assert isinstance(tool1, StructuredTool)
assert tool1.name == "Test Tool 1"
assert tool1.description == "This is a test tool 1."
assert tool1.args == {
"param1": {"title": "Param1", "description": "Parameter 1", "type": "string"},
"param2": {"title": "Param2", "description": "Parameter 2", "type": "integer"},
}
mock_load_tool_manifest.side_effect = lambda _: setattr(
client,
"_manifest",
{
"serverVersion": "0.0.1",
"tools": {
"test_tool1": {
"summary": "Test Tool 1",
"description": "This is a test tool 1.",
"parameters": {
"param1": {"type": "string", "description": "Parameter 1"},
"param2": {"type": "integer", "description": "Parameter 2"},
},
},
"test_tool2": {
"summary": "Test Tool 2",
"description": "This is a test tool 2.",
"parameters": {
"param1": {"type": "integer", "description": "Parameter 1"},
"param2": {"type": "string", "description": "Parameter 2"},
},
},
},
},
)
tool2 = await client.load_tool("test_tool2")
mock_load_tool_manifest.assert_called_with("test_tool2")
assert isinstance(tool2, StructuredTool)
assert tool2.name == "Test Tool 2"
assert tool2.description == "This is a test tool 2."
assert tool2.args == {
"param1": {"title": "Param1", "description": "Parameter 1", "type": "integer"},
"param2": {"title": "Param2", "description": "Parameter 2", "type": "string"},
}
assert client._tools == [tool1, tool2]
@pytest.mark.asyncio
@patch.object(ToolboxClient, "_load_toolset_manifest")
async def test_load_toolset(mock_load_toolset_manifest):
client = ToolboxClient("https://my-toolbox.com")
mock_load_toolset_manifest.side_effect = lambda _: setattr(
client,
"_manifest",
{
"serverVersion": "0.0.1",
"tools": {
"test_tool1": {
"summary": "Test Tool 1",
"description": "This is a test tool 1.",
"parameters": {
"param1": {"type": "string", "description": "Parameter 1"},
"param2": {"type": "integer", "description": "Parameter 2"},
},
},
"test_tool2": {
"summary": "Test Tool 2",
"description": "This is a test tool 2.",
"parameters": {
"param1": {"type": "integer", "description": "Parameter 1"},
"param2": {"type": "string", "description": "Parameter 2"},
},
},
},
},
)
[tool1, tool2] = await client.load_toolset("test_toolset")
mock_load_toolset_manifest.assert_called_once_with("test_toolset")
assert isinstance(tool1, StructuredTool)
assert isinstance(tool2, StructuredTool)
assert tool1.name == "Test Tool 1"
assert tool1.description == "This is a test tool 1."
assert tool1.args == {
"param1": {"title": "Param1", "description": "Parameter 1", "type": "string"},
"param2": {"title": "Param2", "description": "Parameter 2", "type": "integer"},
}
assert tool2.name == "Test Tool 2"
assert tool2.description == "This is a test tool 2."
assert tool2.args == {
"param1": {"title": "Param1", "description": "Parameter 1", "type": "integer"},
"param2": {"title": "Param2", "description": "Parameter 2", "type": "string"},
}
@pytest.mark.asyncio
@patch.object(ToolboxClient, "_load_toolset_manifest")
async def test_load_default_toolset(mock_load_toolset_manifest):
client = ToolboxClient("https://my-toolbox.com")
mock_load_toolset_manifest.side_effect = lambda _: setattr(
client,
"_manifest",
{
"serverVersion": "0.0.1",
"tools": {
"test_tool1": {
"summary": "Test Tool 1",
"description": "This is a test tool 1.",
"parameters": {
"param1": {"type": "string", "description": "Parameter 1"},
"param2": {"type": "integer", "description": "Parameter 2"},
},
},
"test_tool2": {
"summary": "Test Tool 2",
"description": "This is a test tool 2.",
"parameters": {
"param1": {"type": "integer", "description": "Parameter 1"},
"param2": {"type": "string", "description": "Parameter 2"},
},
},
},
},
)
[tool1, tool2] = await client.load_toolset()
mock_load_toolset_manifest.assert_called_once_with(None)
assert isinstance(tool1, StructuredTool)
assert isinstance(tool2, StructuredTool)
assert tool1.name == "Test Tool 1"
assert tool1.description == "This is a test tool 1."
assert tool1.args == {
"param1": {"title": "Param1", "description": "Parameter 1", "type": "string"},
"param2": {"title": "Param2", "description": "Parameter 2", "type": "integer"},
}
assert tool2.name == "Test Tool 2"
assert tool2.description == "This is a test tool 2."
assert tool2.args == {
"param1": {"title": "Param1", "description": "Parameter 1", "type": "integer"},
"param2": {"title": "Param2", "description": "Parameter 2", "type": "string"},
}