72 lines
2.9 KiB
Python
72 lines
2.9 KiB
Python
|
import asyncio
|
||
|
|
import unittest
|
||
|
|
|
||
|
|
from backend.agent.registry import (
|
||
|
|
NativeToolProvider,
|
||
|
|
SchemaValidationError,
|
||
|
|
ToolDefinition,
|
||
|
|
ToolExecutionContext,
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
INPUT = {
|
||
|
|
"type": "object",
|
||
|
|
"properties": {"value": {"type": "string"}},
|
||
|
|
"required": ["value"],
|
||
|
|
"additionalProperties": False,
|
||
|
|
}
|
||
|
|
OUTPUT = {
|
||
|
|
"type": "object",
|
||
|
|
"properties": {"result": {"type": "string"}},
|
||
|
|
"required": ["result"],
|
||
|
|
"additionalProperties": False,
|
||
|
|
}
|
||
|
|
|
||
|
|
|
||
|
|
def context(registry):
|
||
|
|
return ToolExecutionContext(
|
||
|
|
run_id="run", workflow_id="workflow", node_id="node", session_id=None,
|
||
|
|
selection_mode="explicit", explicit_user_action=True,
|
||
|
|
emit=lambda *_args: asyncio.sleep(0), cancellation_event=asyncio.Event(),
|
||
|
|
db_factory=lambda: None, registry=registry,
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
class RegistryTests(unittest.IsolatedAsyncioTestCase):
|
||
|
|
async def test_duplicate_names_are_rejected(self):
|
||
|
|
async def handler(arguments, _context): return {"result": arguments["value"]}
|
||
|
|
registry = NativeToolProvider()
|
||
|
|
definition = ToolDefinition("test.echo", "echo", INPUT, OUTPUT, handler)
|
||
|
|
registry.register(definition)
|
||
|
|
with self.assertRaises(ValueError):
|
||
|
|
registry.register(definition)
|
||
|
|
|
||
|
|
async def test_arguments_and_outputs_are_validated(self):
|
||
|
|
async def handler(arguments, _context): return {"result": arguments["value"]}
|
||
|
|
registry = NativeToolProvider()
|
||
|
|
registry.register(ToolDefinition("test.echo", "echo", INPUT, OUTPUT, handler))
|
||
|
|
self.assertEqual((await registry.call_tool("test.echo", {"value": "ok"}, context(registry)))["result"], "ok")
|
||
|
|
with self.assertRaises(SchemaValidationError):
|
||
|
|
await registry.call_tool("test.echo", {"value": 1}, context(registry))
|
||
|
|
|
||
|
|
async def invalid_output(_arguments, _context): return {"result": 1}
|
||
|
|
invalid = NativeToolProvider()
|
||
|
|
invalid.register(ToolDefinition("test.invalid", "invalid", INPUT, OUTPUT, invalid_output))
|
||
|
|
with self.assertRaises(SchemaValidationError):
|
||
|
|
await invalid.call_tool("test.invalid", {"value": "x"}, context(invalid))
|
||
|
|
|
||
|
|
async def test_timeout_and_result_size_are_enforced(self):
|
||
|
|
async def slow(_arguments, _context):
|
||
|
|
await asyncio.sleep(0.05)
|
||
|
|
return {"result": "ok"}
|
||
|
|
registry = NativeToolProvider()
|
||
|
|
registry.register(ToolDefinition("test.slow", "slow", INPUT, OUTPUT, slow, timeout_seconds=0.01))
|
||
|
|
with self.assertRaises(TimeoutError):
|
||
|
|
await registry.call_tool("test.slow", {"value": "x"}, context(registry))
|
||
|
|
|
||
|
|
async def large(_arguments, _context): return {"result": "x" * 100}
|
||
|
|
limited = NativeToolProvider()
|
||
|
|
limited.register(ToolDefinition("test.large", "large", INPUT, OUTPUT, large, result_size_limit=20))
|
||
|
|
with self.assertRaises(ValueError):
|
||
|
|
await limited.call_tool("test.large", {"value": "x"}, context(limited))
|