Arseny Boykov
8c71144c6d
Running `pytest --unit` with these changes goes from 3 minutes 44 seconds to just **30 seconds**.
### How does it work
#### Virtual time
Takes advantage of the fact that many of the tests rely on `asyncio.sleep()` to synchronize events in the loop. It unties it from the real time so that `await asyncio.sleep(60)` returns instantly with `loop.time()` reporting 60 seconds passed. `time.time()` is then patched to report loop time. This brings down `test_agent_session.py` test time form 90 seconds to just 1 second without any modifications to the test itself.
- To disable globally `--real-time` option can be used.
- To enable/disable per-module or per-test `pytest.mark.virtual_time`/`.real_time` can be used.
#### Concurrent tests
Tests with real I/O are groupped by their module and executed in a shared event loop. This shrinks wall time of the whole module to wall time of the slowest test + raw compute. `contextvars` are used to preserve input capture and track leaked tasks.
- To disable concurrent execution, `--no-concurrent` option can be used.
- To enable/disable per-module or per-test, `pytest.mark.no_concurrent`/`.no_concurrent` can be used.
### Changes
- Add `pytest-asyncio-concurrent` — gives the speedup for I/O bound tests
- Add `tests/concurrent.py` that extends on `pytest-asyncio-concurrent` and adds crucial features:
- Captures stdout/stderr and reports on failure (same as `pytest`)
- Allows leaked tasks detection in concurrent mode (so that `fail_on_leaked_tasks` continues to work)
- adds `--concurrent`/`--no-concurrent` options
- Adds live progress
- Add `tests/virtual_time.py` that provides instant and deterministic `asyncio` clock for time-bound execution
#### Changes to the tests sources
- https://github.com/livekit/agents/pull/5980/commits/04fe6cbb55c32d1058a01a4c0ddda790b5834770: Add `concurrent`/`virtual_time` markers to tests modules
- https://github.com/livekit/agents/pull/5980/changes/a8f5be0be8b76efefaa55743d8ad0a5d77b15b1b: Tighten timings in `test_audio_decoder.py`, `test_ipc.pt` and `test_room.py`
- https://github.com/livekit/agents/pull/5980/changes/396bce8cbf4618a292ab28a257ea13f58bd3c5a6: Make `test_async_channel` and `test_proc_pool_launch_job_raises_when_all_spawns_fail` in `test_ipc.py` safe to run concurrently
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)
|