Pythonregistry.py
Copy
1
from dataclasses import dataclass, field
2
from typing import Callable, Any
3
from functools import wraps
4
 
5
@dataclass
6
class ToolSpec:
7
    name: str
8
    description: str
9
    input_schema: dict
10
    fn: Callable
11
    tags: list[str] = field(default_factory=list)
12
 
13
class ToolRegistry:
14
    def __init__(self):
15
        self._tools: dict[str, ToolSpec] = {}
16
 
17
    def register(self, description: str, tags: list[str] = None):
18
        def decorator(fn: Callable) -> Callable:
19
            hints = {k: v for k, v in fn.__annotations__.items() if k != "return"}
20
            schema = {
21
                "type": "object",
22
                "properties": {
23
                    name: {"type": "string" if t == str else "number" if t in (int, float) else "boolean" if t == bool else "object"}
24
                    for name, t in hints.items()
25
                },
26
                "required": list(hints.keys())
27
            }
28
            self._tools[fn.__name__] = ToolSpec(
29
                name=fn.__name__, description=description,
30
                input_schema=schema, fn=fn, tags=tags or []
31
            )
32
            @wraps(fn)
33
            def wrapper(*args, **kwargs):
34
                return fn(*args, **kwargs)
35
            return wrapper
36
        return decorator
37
 
38
    def get(self, name: str) -> ToolSpec | None:
39
        return self._tools.get(name)
40
 
41
    def openai_schemas(self) -> list[dict]:
42
        return [{"type": "function", "function": {"name": s.name, "description": s.description, "parameters": s.input_schema}} for s in self._tools.values()]
43
 
44
    def by_tag(self, tag: str) -> list[ToolSpec]:
45
        return [s for s in self._tools.values() if tag in s.tags]
46
 
47
registry = ToolRegistry()
48
 
49
@registry.register("Search the web for real-time information", tags=["search"])
50
def web_search(query: str) -> str:
51
    return f"Results for: {query}"
52
 
53
@registry.register("Calculate a math expression", tags=["compute"])
54
def calculate(expression: str) -> float:
55
    return eval(expression)
56
 
57
@registry.register("Write content to a file", tags=["io", "file"])
58
def write_file(path: str, content: str) -> str:
59
    with open(path, "w") as f:
60
        f.write(content)
61
    return f"Written to {path}"
62
 
63
# Test
64
print("All tools:", [s.name for s in registry._tools.values()])
65
print("Search tools:", [s.name for s in registry.by_tag("search")])
66
print("Calculate:", registry.get("calculate").fn("2 ** 10"))
67