Unfucked gemini's mess
This commit is contained in:
+92
-164
@@ -1,181 +1,109 @@
|
||||
"""Tool registry - defines and registers all available tools for the agent."""
|
||||
|
||||
import logging
|
||||
from collections.abc import Callable
|
||||
from dataclasses import dataclass
|
||||
from typing import Any
|
||||
|
||||
from .tools import api as api_tools
|
||||
from .tools import filesystem as fs_tools
|
||||
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.
|
||||
|
||||
Attributes:
|
||||
name: Unique identifier for the tool.
|
||||
description: Human-readable description for the LLM.
|
||||
func: The callable that implements the tool.
|
||||
parameters: JSON Schema describing the tool's parameters.
|
||||
"""
|
||||
|
||||
"""Represents a tool that can be used by the agent."""
|
||||
name: str
|
||||
description: str
|
||||
func: Callable[..., dict[str, Any]]
|
||||
parameters: dict[str, Any]
|
||||
func: Callable[..., Dict[str, Any]]
|
||||
parameters: Dict[str, Any]
|
||||
|
||||
|
||||
def make_tools() -> dict[str, Tool]:
|
||||
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.
|
||||
|
||||
Tools access memory via get_memory() context.
|
||||
|
||||
|
||||
Returns:
|
||||
Dictionary mapping tool names to Tool instances.
|
||||
Dictionary mapping tool names to Tool objects
|
||||
"""
|
||||
tools = [
|
||||
# Filesystem tools
|
||||
Tool(
|
||||
name="set_path_for_folder",
|
||||
description=(
|
||||
"Sets a path in the configuration "
|
||||
"(download_folder, tvshow_folder, movie_folder, or torrent_folder)."
|
||||
),
|
||||
func=fs_tools.set_path_for_folder,
|
||||
parameters={
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"folder_name": {
|
||||
"type": "string",
|
||||
"description": "Name of folder to set",
|
||||
"enum": ["download", "tvshow", "movie", "torrent"],
|
||||
},
|
||||
"path_value": {
|
||||
"type": "string",
|
||||
"description": "Absolute path to the folder",
|
||||
},
|
||||
},
|
||||
"required": ["folder_name", "path_value"],
|
||||
},
|
||||
),
|
||||
Tool(
|
||||
name="list_folder",
|
||||
description="Lists the contents of a configured folder.",
|
||||
func=fs_tools.list_folder,
|
||||
parameters={
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"folder_type": {
|
||||
"type": "string",
|
||||
"description": "Type of folder to list",
|
||||
"enum": ["download", "tvshow", "movie", "torrent"],
|
||||
},
|
||||
"path": {
|
||||
"type": "string",
|
||||
"description": "Relative path within the folder",
|
||||
"default": ".",
|
||||
},
|
||||
},
|
||||
"required": ["folder_type"],
|
||||
},
|
||||
),
|
||||
# Media search tools
|
||||
Tool(
|
||||
name="find_media_imdb_id",
|
||||
description=(
|
||||
"Finds the IMDb ID for a given media title using TMDB API. "
|
||||
"Use this to get information about a movie or TV show."
|
||||
),
|
||||
func=api_tools.find_media_imdb_id,
|
||||
parameters={
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"media_title": {
|
||||
"type": "string",
|
||||
"description": "Title of the media to search for",
|
||||
},
|
||||
},
|
||||
"required": ["media_title"],
|
||||
},
|
||||
),
|
||||
# Torrent tools
|
||||
Tool(
|
||||
name="find_torrents",
|
||||
description=(
|
||||
"Finds torrents for a given media title. "
|
||||
"Results are numbered (1, 2, 3...) so the user can select by number."
|
||||
),
|
||||
func=api_tools.find_torrent,
|
||||
parameters={
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"media_title": {
|
||||
"type": "string",
|
||||
"description": "Title to search for (include quality if specified)",
|
||||
},
|
||||
},
|
||||
"required": ["media_title"],
|
||||
},
|
||||
),
|
||||
Tool(
|
||||
name="add_torrent_by_index",
|
||||
description=(
|
||||
"Adds a torrent from the previous search results by its number. "
|
||||
"Use when the user says 'download the 3rd one' or 'take number 2'."
|
||||
),
|
||||
func=api_tools.add_torrent_by_index,
|
||||
parameters={
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"index": {
|
||||
"type": "integer",
|
||||
"description": "Number of the torrent in search results (1, 2, 3...)",
|
||||
},
|
||||
},
|
||||
"required": ["index"],
|
||||
},
|
||||
),
|
||||
Tool(
|
||||
name="add_torrent_to_qbittorrent",
|
||||
description=(
|
||||
"Adds a torrent to qBittorrent using a magnet link directly. "
|
||||
"Use add_torrent_by_index if user selected from search results."
|
||||
),
|
||||
func=api_tools.add_torrent_to_qbittorrent,
|
||||
parameters={
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"magnet_link": {
|
||||
"type": "string",
|
||||
"description": "The magnet link of the torrent",
|
||||
},
|
||||
},
|
||||
"required": ["magnet_link"],
|
||||
},
|
||||
),
|
||||
Tool(
|
||||
name="get_torrent_by_index",
|
||||
description=(
|
||||
"Gets details of a torrent from search results by its number, "
|
||||
"without downloading it."
|
||||
),
|
||||
func=api_tools.get_torrent_by_index,
|
||||
parameters={
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"index": {
|
||||
"type": "integer",
|
||||
"description": "Number of the torrent in search results (1, 2, 3...)",
|
||||
},
|
||||
},
|
||||
"required": ["index"],
|
||||
},
|
||||
),
|
||||
# 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,
|
||||
]
|
||||
|
||||
logger.info(f"Registered {len(tools)} tools")
|
||||
return {t.name: t for t in tools}
|
||||
|
||||
# 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
|
||||
|
||||
Reference in New Issue
Block a user