SupplyAgent/agents/generate_demand_agent/tools/__init__.py

"""
generate_demand_agent 工具包
"""
from __future__ import annotations

from collections.abc import Callable
from typing import Any

from agents.generate_demand_agent.tools.batch_save_generated_demands import (
    batch_save_generated_demands,
)
from agents.generate_demand_agent.tools.query_category_leaves_by_dim import (
    query_category_leaves_by_dim,
)
from agents.generate_demand_agent.tools.query_category_path import query_category_path
from agents.generate_demand_agent.tools.query_category_tree_by_dim import (
    query_category_tree_by_dim,
)
from agents.generate_demand_agent.tools.query_latest_biz_dt import query_latest_biz_dt
from agents.generate_demand_agent.tools.query_demand_words_by_category import (
    query_demand_words_by_category,
)
from supply_agent.tools.registry import ToolRegistry

ALL_TOOLS: list[Callable[..., Any]] = [
    query_latest_biz_dt,
    query_category_tree_by_dim,
    query_category_leaves_by_dim,
    query_demand_words_by_category,
    query_category_path,
    batch_save_generated_demands,
]

__all__ = [
    "ALL_TOOLS",
    "batch_save_generated_demands",
    "query_category_leaves_by_dim",
    "query_category_path",
    "query_category_tree_by_dim",
    "query_demand_words_by_category",
    "query_latest_biz_dt",
    "register_all_tools",
]


def register_all_tools(registry: ToolRegistry) -> ToolRegistry:
    """将 generate_demand_agent 包内的所有工具注册到 ToolRegistry。"""
    return registry.from_decorated(*ALL_TOOLS)