from dataclasses import dataclass
from typing import Any, Callable
from athena_tools import (
execute_query,
get_table_schema,
list_tables,
)
@dataclass
class ToolConfig:
"""Configuration for a single MCP tool."""
name: str
description: str
input_schema: dict[str, Any]
execute: Callable[..., dict[str, Any]]
TOOL_REGISTRY: list[ToolConfig] = [
ToolConfig(
name=execute_query.TOOL_NAME,
description=execute_query.TOOL_DESCRIPTION,
input_schema=execute_query.TOOL_INPUT_SCHEMA,
execute=execute_query.execute,
),
ToolConfig(
name=list_tables.TOOL_NAME,
description=list_tables.TOOL_DESCRIPTION,
input_schema=list_tables.TOOL_INPUT_SCHEMA,
execute=list_tables.execute,
),
ToolConfig(
name=get_table_schema.TOOL_NAME,
description=get_table_schema.TOOL_DESCRIPTION,
input_schema=get_table_schema.TOOL_INPUT_SCHEMA,
execute=get_table_schema.execute,
),
]
TOOL_REGISTRY_DICT: dict[str, Callable[..., dict[str, Any]]] = {
tool.name: tool.execute for tool in TOOL_REGISTRY
}
