wehub-resource-sync
309 lines
9.9 KiB
Python
309 lines
9.9 KiB
Python
import json
|
|
from unittest.mock import MagicMock
|
|
|
|
import pytest
|
|
|
|
from livekit.agents.beta.toolsets.tool_proxy import ToolProxyToolset
|
|
from livekit.agents.llm import ToolContext, ToolError, Toolset, function_tool
|
|
|
|
pytestmark = [pytest.mark.unit, pytest.mark.concurrent]
|
|
|
|
|
|
@function_tool
|
|
async def weather_tool(city: str) -> str:
|
|
"""Get weather for a city.
|
|
|
|
Args:
|
|
city: City name or region
|
|
"""
|
|
return f"sunny in {city}"
|
|
|
|
|
|
@function_tool
|
|
async def forecast_tool(city: str, days: int) -> str:
|
|
"""Get weather forecast.
|
|
|
|
Args:
|
|
city: City name or region
|
|
days: Number of days to forecast
|
|
"""
|
|
return f"{days}-day forecast for {city}: mostly sunny"
|
|
|
|
|
|
@function_tool
|
|
async def stock_tool(symbol: str) -> str:
|
|
"""Get stock price for a symbol.
|
|
|
|
Args:
|
|
symbol: Stock ticker symbol
|
|
"""
|
|
return f"{symbol}: $100"
|
|
|
|
|
|
@function_tool(
|
|
raw_schema={
|
|
"name": "raw_search",
|
|
"description": "Search with raw schema",
|
|
"parameters": {
|
|
"type": "object",
|
|
"properties": {
|
|
"query": {"type": "string", "description": "Search query"},
|
|
"limit": {"type": "integer", "description": "Max results"},
|
|
},
|
|
},
|
|
}
|
|
)
|
|
async def raw_search(raw_arguments: dict[str, object]) -> str:
|
|
return f"results for {raw_arguments.get('query')}"
|
|
|
|
|
|
class WeatherToolset(Toolset):
|
|
def __init__(self):
|
|
super().__init__(id="weather", tools=[weather_tool, forecast_tool])
|
|
|
|
|
|
class FinanceToolset(Toolset):
|
|
def __init__(self):
|
|
super().__init__(id="finance", tools=[stock_tool])
|
|
|
|
|
|
def _mock_ctx() -> MagicMock:
|
|
return MagicMock()
|
|
|
|
|
|
class TestToolProxyToolset:
|
|
def test_tools_always_two(self):
|
|
"""ToolProxyToolset always exposes exactly search_tools and call_tool."""
|
|
ts = ToolProxyToolset(
|
|
id="proxy",
|
|
tools=[WeatherToolset(), FinanceToolset()],
|
|
)
|
|
assert len(ts.tools) == 2
|
|
tool_names = [t.id for t in ts.tools]
|
|
assert "tool_search" in tool_names
|
|
assert "call_tool" in tool_names
|
|
|
|
async def test_tools_constant_after_search(self):
|
|
"""Tool list stays constant even after search — no dynamic loading."""
|
|
ts = ToolProxyToolset(
|
|
id="proxy",
|
|
tools=[WeatherToolset(), FinanceToolset()],
|
|
)
|
|
await ts.setup()
|
|
|
|
assert len(ts.tools) == 2
|
|
await ts._handle_search({"query": "weather"})
|
|
assert len(ts.tools) == 2 # still just search + call
|
|
|
|
async def test_search_returns_full_schemas(self):
|
|
"""search_tools returns tool schemas with full JSON schema for parameters."""
|
|
ts = ToolProxyToolset(
|
|
id="proxy",
|
|
tools=[WeatherToolset(), FinanceToolset()],
|
|
)
|
|
await ts.setup()
|
|
|
|
result = await ts._handle_search({"query": "weather"})
|
|
parsed = [json.loads(schema) for schema in result.split("\n")]
|
|
|
|
assert isinstance(parsed, list)
|
|
assert len(parsed) >= 1
|
|
|
|
names = [t["name"] for t in parsed]
|
|
assert "weather_tool" in names
|
|
|
|
# Check schema has full JSON schema with types
|
|
weather = next(t for t in parsed if t["name"] == "weather_tool")
|
|
assert "Get weather for a city" in weather["description"]
|
|
assert "parameters" in weather
|
|
params = weather["parameters"]
|
|
assert "properties" in params
|
|
assert "city" in params["properties"]
|
|
assert params["properties"]["city"]["type"] == "string"
|
|
|
|
async def test_search_returns_raw_tool_schema(self):
|
|
"""search_tools returns full raw_schema parameters for RawFunctionTool."""
|
|
ts = ToolProxyToolset(
|
|
id="proxy",
|
|
tools=[raw_search],
|
|
)
|
|
await ts.setup()
|
|
|
|
result = await ts._handle_search({"query": "search"})
|
|
parsed = [json.loads(schema) for schema in result.split("\n")]
|
|
|
|
assert len(parsed) >= 1
|
|
tool = parsed[0]
|
|
assert tool["name"] == "raw_search"
|
|
params = tool["parameters"]
|
|
assert params["type"] == "object"
|
|
assert "query" in params["properties"]
|
|
assert params["properties"]["query"]["type"] == "string"
|
|
assert params["properties"]["limit"]["type"] == "integer"
|
|
|
|
async def test_search_no_results_raises_tool_error(self):
|
|
ts = ToolProxyToolset(
|
|
id="proxy",
|
|
tools=[WeatherToolset()],
|
|
)
|
|
await ts.setup()
|
|
|
|
with pytest.raises(ToolError, match="No tools found"):
|
|
await ts._handle_search({"query": "xyznonexistent"})
|
|
|
|
async def test_search_empty_query_raises_tool_error(self):
|
|
ts = ToolProxyToolset(
|
|
id="proxy",
|
|
tools=[WeatherToolset()],
|
|
)
|
|
await ts.setup()
|
|
|
|
with pytest.raises(ToolError, match="cannot be empty"):
|
|
await ts._handle_search({"query": ""})
|
|
|
|
async def test_call_tool_function_tool(self):
|
|
"""call_tool can execute a FunctionTool by name."""
|
|
ts = ToolProxyToolset(
|
|
id="proxy",
|
|
tools=[WeatherToolset(), FinanceToolset()],
|
|
)
|
|
await ts.setup()
|
|
|
|
result = await ts._handle_call(
|
|
_mock_ctx(), {"name": "weather_tool", "parameters": {"city": "London"}}
|
|
)
|
|
assert "sunny in London" in str(result)
|
|
|
|
async def test_call_tool_with_multiple_args(self):
|
|
ts = ToolProxyToolset(
|
|
id="proxy",
|
|
tools=[WeatherToolset()],
|
|
)
|
|
await ts.setup()
|
|
|
|
result = await ts._handle_call(
|
|
_mock_ctx(), {"name": "forecast_tool", "parameters": {"city": "Tokyo", "days": 3}}
|
|
)
|
|
assert "3-day forecast for Tokyo" in str(result)
|
|
|
|
async def test_call_tool_raw_function_tool(self):
|
|
"""call_tool can execute a RawFunctionTool by name."""
|
|
ts = ToolProxyToolset(
|
|
id="proxy",
|
|
tools=[raw_search],
|
|
)
|
|
await ts.setup()
|
|
|
|
result = await ts._handle_call(
|
|
_mock_ctx(), {"name": "raw_search", "parameters": {"query": "test", "limit": 5}}
|
|
)
|
|
assert "results for test" in str(result)
|
|
|
|
async def test_call_unknown_tool_raises_tool_error(self):
|
|
ts = ToolProxyToolset(
|
|
id="proxy",
|
|
tools=[WeatherToolset()],
|
|
)
|
|
await ts.setup()
|
|
|
|
with pytest.raises(ToolError, match="unknown tool"):
|
|
await ts._handle_call(_mock_ctx(), {"name": "nonexistent_tool", "parameters": {}})
|
|
|
|
async def test_call_missing_required_arg_raises_tool_error(self, capsys):
|
|
"""Missing required argument produces a detailed ToolError for the LLM."""
|
|
ts = ToolProxyToolset(
|
|
id="proxy",
|
|
tools=[WeatherToolset()],
|
|
)
|
|
await ts.setup()
|
|
|
|
with pytest.raises(ToolError, match="Error parsing arguments") as exc_info:
|
|
await ts._handle_call(_mock_ctx(), {"name": "weather_tool", "parameters": {}})
|
|
|
|
error_msg = exc_info.value.message
|
|
print(f"Missing arg error: {error_msg}")
|
|
|
|
# Error message should contain the missing field name and indicate it's required
|
|
assert "city" in error_msg
|
|
assert "type=missing" in error_msg
|
|
|
|
async def test_call_wrong_type_arg_raises_tool_error(self, capsys):
|
|
"""Wrong argument type produces a detailed ToolError for the LLM."""
|
|
ts = ToolProxyToolset(
|
|
id="proxy",
|
|
tools=[WeatherToolset()],
|
|
)
|
|
await ts.setup()
|
|
|
|
with pytest.raises(ToolError, match="Error parsing arguments") as exc_info:
|
|
await ts._handle_call(
|
|
_mock_ctx(),
|
|
{"name": "forecast_tool", "parameters": {"city": "Tokyo", "days": "not_a_number"}},
|
|
)
|
|
|
|
error_msg = exc_info.value.message
|
|
print(f"Wrong type error: {error_msg}")
|
|
|
|
# Error message should contain the field name and indicate a type parsing issue
|
|
assert "days" in error_msg
|
|
assert "int_parsing" in error_msg
|
|
|
|
async def test_call_tool_propagates_tool_error(self):
|
|
"""ToolError raised inside a tool is re-raised as-is."""
|
|
|
|
@function_tool
|
|
async def tool_with_tool_error(x: str) -> str:
|
|
"""A tool that raises ToolError.
|
|
|
|
Args:
|
|
x: Some input
|
|
"""
|
|
raise ToolError("custom error message")
|
|
|
|
ts = ToolProxyToolset(
|
|
id="proxy",
|
|
tools=[tool_with_tool_error],
|
|
)
|
|
await ts.setup()
|
|
|
|
with pytest.raises(ToolError, match="custom error message"):
|
|
await ts._handle_call(
|
|
_mock_ctx(), {"name": "tool_with_tool_error", "parameters": {"x": "test"}}
|
|
)
|
|
|
|
async def test_tool_context_sees_only_proxy_tools(self):
|
|
"""ToolContext built from ToolProxyToolset only sees search + call."""
|
|
ts = ToolProxyToolset(
|
|
id="proxy",
|
|
tools=[WeatherToolset(), FinanceToolset()],
|
|
)
|
|
await ts.setup()
|
|
|
|
ctx = ToolContext([ts])
|
|
assert "tool_search" in ctx.function_tools
|
|
assert "call_tool" in ctx.function_tools
|
|
assert len(ctx.function_tools) == 2
|
|
assert "weather_tool" not in ctx.function_tools
|
|
|
|
async def test_nested_toolsets(self):
|
|
"""Nested toolsets are accessible via call_tool."""
|
|
|
|
class OuterToolset(Toolset):
|
|
def __init__(self):
|
|
super().__init__(id="outer", tools=[WeatherToolset(), stock_tool])
|
|
|
|
ts = ToolProxyToolset(id="proxy", tools=[OuterToolset()])
|
|
await ts.setup()
|
|
|
|
# Search should find tools from nested toolsets
|
|
result = await ts._handle_search({"query": "weather stock"})
|
|
parsed = [json.loads(schema) for schema in result.split("\n")]
|
|
names = [t["name"] for t in parsed]
|
|
assert "weather_tool" in names or "stock_tool" in names
|
|
|
|
# call_tool should work for nested tools
|
|
result = await ts._handle_call(
|
|
_mock_ctx(), {"name": "weather_tool", "parameters": {"city": "Paris"}}
|
|
)
|
|
assert "sunny in Paris" in str(result)
|