File size: 2,249 Bytes
c15b707
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
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