110 lines
3.0 KiB
Python
110 lines
3.0 KiB
Python
"""Tool registry - defines and registers all available tools for the agent."""
|
|
from dataclasses import dataclass
|
|
from typing import Callable, Any, Dict
|
|
import logging
|
|
import inspect
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
@dataclass
|
|
class Tool:
|
|
"""Represents a tool that can be used by the agent."""
|
|
name: str
|
|
description: str
|
|
func: Callable[..., Dict[str, Any]]
|
|
parameters: Dict[str, Any]
|
|
|
|
|
|
def _create_tool_from_function(func: Callable) -> Tool:
|
|
"""
|
|
Create a Tool object from a function.
|
|
|
|
Args:
|
|
func: Function to convert to a tool
|
|
|
|
Returns:
|
|
Tool object with metadata extracted from function
|
|
"""
|
|
sig = inspect.signature(func)
|
|
doc = inspect.getdoc(func)
|
|
|
|
# Extract description from docstring (first line)
|
|
description = doc.strip().split('\n')[0] if doc else func.__name__
|
|
|
|
# Build JSON schema from function signature
|
|
properties = {}
|
|
required = []
|
|
|
|
for param_name, param in sig.parameters.items():
|
|
if param_name == "self":
|
|
continue
|
|
|
|
# Map Python types to JSON schema types
|
|
param_type = "string" # default
|
|
if param.annotation != inspect.Parameter.empty:
|
|
if param.annotation == str:
|
|
param_type = "string"
|
|
elif param.annotation == int:
|
|
param_type = "integer"
|
|
elif param.annotation == float:
|
|
param_type = "number"
|
|
elif param.annotation == bool:
|
|
param_type = "boolean"
|
|
|
|
properties[param_name] = {
|
|
"type": param_type,
|
|
"description": f"Parameter {param_name}"
|
|
}
|
|
|
|
# Add to required if no default value
|
|
if param.default == inspect.Parameter.empty:
|
|
required.append(param_name)
|
|
|
|
parameters = {
|
|
"type": "object",
|
|
"properties": properties,
|
|
"required": required,
|
|
}
|
|
|
|
return Tool(
|
|
name=func.__name__,
|
|
description=description,
|
|
func=func,
|
|
parameters=parameters,
|
|
)
|
|
|
|
|
|
def make_tools() -> Dict[str, Tool]:
|
|
"""
|
|
Create and register all available tools.
|
|
|
|
Returns:
|
|
Dictionary mapping tool names to Tool objects
|
|
"""
|
|
# Import tools here to avoid circular dependencies
|
|
from .tools import filesystem as fs_tools
|
|
from .tools import api as api_tools
|
|
from .tools import language as lang_tools
|
|
|
|
# List of all tool functions
|
|
tool_functions = [
|
|
fs_tools.set_path_for_folder,
|
|
fs_tools.list_folder,
|
|
api_tools.find_media_imdb_id,
|
|
api_tools.find_torrent,
|
|
api_tools.add_torrent_by_index,
|
|
api_tools.add_torrent_to_qbittorrent,
|
|
api_tools.get_torrent_by_index,
|
|
lang_tools.set_language,
|
|
]
|
|
|
|
# Create Tool objects from functions
|
|
tools = {}
|
|
for func in tool_functions:
|
|
tool = _create_tool_from_function(func)
|
|
tools[tool.name] = tool
|
|
|
|
logger.info(f"Registered {len(tools)} tools: {list(tools.keys())}")
|
|
return tools
|