tools/tool_collection_common.py

from langchain.tools import StructuredTool
from langchain_community.tools import DuckDuckGoSearchRun

from agents.wiki_agent import WikipediaEnToolAgentInput, wikipedia_en_tool_agent
from config import settings
from tools.encyclopedia import EncyclopediaRetrieveInput, EncyclopediaRetriever
from tools.get_hub_stats import HubStatInput, get_hub_stats
from tools.get_weather import WeatherInput, get_weather


class ToolsCollection:
    @staticmethod
    def get_tools(needed_tools: list[str]) -> list:
        tools = []

        for nt in needed_tools:
            if nt == "search_tool":
                tools.append(DuckDuckGoSearchRun())
            elif nt == "get_weather":
                tools.append(
                    StructuredTool(
                        name="get_weather",
                        func=get_weather,
                        description=get_weather.__doc__,
                        args_schema=WeatherInput,
                    )
                )
            elif nt == "hub_stats_tool":
                tools.append(
                    StructuredTool(
                        name="hub_stats_tool",
                        func=get_hub_stats,
                        description=get_hub_stats.__doc__,
                        args_schema=HubStatInput,
                    )
                )
            elif nt == "wikipedia_en_tool_agent":
                tools.append(
                    StructuredTool(
                        name="wikipedia_en_tool_agent",
                        func=wikipedia_en_tool_agent,
                        description=wikipedia_en_tool_agent.__doc__,
                        args_schema=WikipediaEnToolAgentInput,
                    )
                )
            elif nt == "EncyclopediaRetriever":
                retriver = EncyclopediaRetriever(["gaia"], settings.PROJ_PATH)

                tools.append(
                    StructuredTool(
                        name="get_related_question_from_encyclopedia",
                        func=retriver.get_related_question,
                        description=retriver.get_related_question.__doc__,
                        args_schema=EncyclopediaRetrieveInput,
                    )
                )

        return tools