Skip to content

Commit c1f4cf6

Browse files
constkclaude
andcommitted
feat: tool-registry pattern + example echo_tool (#20)
Add src/tools/registry.py: a generic dispatcher mapping tool name -> (input schema, callable). Each tool's input + output are StrictModel subclasses. Three-method API matches the issue spec: - register(name, input_schema) -> decorator - dispatch(name, raw_input) -> validates raw dict against input schema, invokes the tool, returns typed output - names() -> sorted list of registered tool names UnknownToolError (KeyError subclass) raises on dispatch with a missing name; Pydantic's ValidationError propagates on bad input (wrong type or unknown keys via StrictModel's extra="forbid"). Module-global `registry` singleton; `echo_tool` (input/output pair EchoToolInput/EchoToolOutput) self-registers at module load to demonstrate the pattern. Layer hygiene: registry.py imports only from src/models/. Verified by lint-imports — both contracts still kept. 7 unit tests cover: module-global resolves echo, happy-path dispatch, unknown-tool error, bad input rejection, unknown-key rejection, duplicate registration, and registry isolation. Closes #20 Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
1 parent 8116629 commit c1f4cf6

2 files changed

Lines changed: 169 additions & 0 deletions

File tree

src/tools/registry.py

Lines changed: 109 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,109 @@
1+
"""Generic tool registry — dispatcher for typed agent tools.
2+
3+
Each tool is a function that takes a single ``StrictModel`` input and returns
4+
a single ``StrictModel`` output. The registry holds a mapping from tool name
5+
to ``(input_schema, callable)`` and provides:
6+
7+
- ``register(name, input_schema)`` — decorator to register a callable.
8+
- ``dispatch(name, raw_input)`` — validate the dict-shaped ``raw_input``
9+
against the input schema, call the
10+
tool, return the typed output.
11+
12+
Layer-wise the registry sits below ``agent`` / ``api`` / ``eval`` (it doesn't
13+
import from them) and above ``models``. Verified by the import-linter
14+
contract in ``pyproject.toml``.
15+
"""
16+
17+
from __future__ import annotations
18+
19+
from collections.abc import Callable
20+
from typing import Any
21+
22+
from src.models._base import StrictModel
23+
24+
ToolFn = Callable[[StrictModel], StrictModel]
25+
26+
27+
class UnknownToolError(KeyError):
28+
"""Raised when ``dispatch`` is called with an unregistered tool name."""
29+
30+
31+
class Registry:
32+
"""Maps a tool name to its input schema and callable implementation."""
33+
34+
def __init__(self) -> None:
35+
self._tools: dict[str, tuple[type[StrictModel], ToolFn]] = {}
36+
37+
def register(
38+
self,
39+
name: str,
40+
input_schema: type[StrictModel],
41+
) -> Callable[[ToolFn], ToolFn]:
42+
"""Register a tool implementation.
43+
44+
Returns a decorator so callers can use either of:
45+
46+
@registry.register("echo", EchoToolInput)
47+
def echo_tool(payload: EchoToolInput) -> EchoToolOutput: ...
48+
49+
registry.register("echo", EchoToolInput)(echo_tool)
50+
"""
51+
52+
def decorator(fn: ToolFn) -> ToolFn:
53+
if name in self._tools:
54+
msg = f"Tool {name!r} is already registered."
55+
raise ValueError(msg)
56+
self._tools[name] = (input_schema, fn)
57+
return fn
58+
59+
return decorator
60+
61+
def dispatch(self, name: str, raw_input: dict[str, Any]) -> StrictModel:
62+
"""Validate ``raw_input`` and call the tool.
63+
64+
Raises ``UnknownToolError`` when *name* isn't registered. Pydantic's
65+
``ValidationError`` propagates when ``raw_input`` doesn't match the
66+
registered input schema.
67+
"""
68+
if name not in self._tools:
69+
registered = sorted(self._tools)
70+
msg = f"Unknown tool {name!r}. Registered: {registered}"
71+
raise UnknownToolError(msg)
72+
input_schema, fn = self._tools[name]
73+
payload = input_schema.model_validate(raw_input)
74+
return fn(payload)
75+
76+
def names(self) -> list[str]:
77+
"""Return the sorted list of registered tool names."""
78+
return sorted(self._tools)
79+
80+
81+
# Module-global singleton — agent / eval consumers import this directly so
82+
# tools self-register at module load via the decorator below.
83+
registry = Registry()
84+
85+
86+
# ---------------------------------------------------------------------------
87+
# Example tool: echo — exercises the layer + demonstrates the contract shape.
88+
# ---------------------------------------------------------------------------
89+
90+
91+
class EchoToolInput(StrictModel, strict=True):
92+
"""Input contract for the example echo tool."""
93+
94+
msg: str
95+
96+
97+
class EchoToolOutput(StrictModel, strict=True):
98+
"""Output contract for the example echo tool."""
99+
100+
echoed: str
101+
102+
103+
@registry.register("echo", EchoToolInput)
104+
def echo_tool(payload: StrictModel) -> StrictModel:
105+
"""Return the input string wrapped in ``EchoToolOutput``."""
106+
if not isinstance(payload, EchoToolInput): # pragma: no cover — defensive
107+
msg = f"echo_tool got unexpected payload type: {type(payload)!r}"
108+
raise TypeError(msg)
109+
return EchoToolOutput(echoed=payload.msg)

tests/test_tools.py

Lines changed: 60 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,60 @@
1+
"""Tests for ``src.tools.registry`` — happy path, unknown tool, bad input."""
2+
3+
from __future__ import annotations
4+
5+
import pytest
6+
from pydantic import ValidationError
7+
8+
from src.tools.registry import (
9+
EchoToolInput,
10+
EchoToolOutput,
11+
Registry,
12+
UnknownToolError,
13+
echo_tool,
14+
registry,
15+
)
16+
17+
18+
def test_module_registry_resolves_echo() -> None:
19+
"""The module-global registry has the echo tool wired at import."""
20+
assert "echo" in registry.names()
21+
22+
23+
def test_dispatch_happy_path() -> None:
24+
output = registry.dispatch("echo", {"msg": "hello"})
25+
assert isinstance(output, EchoToolOutput)
26+
assert output.echoed == "hello"
27+
28+
29+
def test_dispatch_unknown_tool_raises() -> None:
30+
with pytest.raises(UnknownToolError, match="Unknown tool 'nope'"):
31+
registry.dispatch("nope", {})
32+
33+
34+
def test_dispatch_rejects_bad_input() -> None:
35+
"""Wrong-typed payload triggers Pydantic ValidationError, not a runtime crash."""
36+
with pytest.raises(ValidationError):
37+
registry.dispatch("echo", {"msg": 123}) # msg must be str
38+
39+
40+
def test_dispatch_rejects_unknown_keys() -> None:
41+
"""StrictModel input schema rejects unknown keys — extra='forbid' propagates."""
42+
with pytest.raises(ValidationError):
43+
registry.dispatch("echo", {"msg": "hi", "boom": True})
44+
45+
46+
def test_register_rejects_duplicate_names() -> None:
47+
"""Registering twice under the same name is a programmer error."""
48+
local = Registry()
49+
local.register("echo", EchoToolInput)(echo_tool)
50+
with pytest.raises(ValueError, match="already registered"):
51+
local.register("echo", EchoToolInput)(echo_tool)
52+
53+
54+
def test_local_registry_isolation() -> None:
55+
"""Multiple registries don't share state."""
56+
a = Registry()
57+
b = Registry()
58+
a.register("echo", EchoToolInput)(echo_tool)
59+
assert a.names() == ["echo"]
60+
assert b.names() == []

0 commit comments

Comments
 (0)