Alfred828 commited on
Commit
c15b707
·
verified ·
1 Parent(s): 51933cb

Create tools/tool_collection_common.py

Browse files
Files changed (1) hide show
  1. tools/tool_collection_common.py +58 -0
tools/tool_collection_common.py ADDED
@@ -0,0 +1,58 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from langchain.tools import StructuredTool
2
+ from langchain_community.tools import DuckDuckGoSearchRun
3
+
4
+ from agents.wiki_agent import WikipediaEnToolAgentInput, wikipedia_en_tool_agent
5
+ from config import settings
6
+ from tools.encyclopedia import EncyclopediaRetrieveInput, EncyclopediaRetriever
7
+ from tools.get_hub_stats import HubStatInput, get_hub_stats
8
+ from tools.get_weather import WeatherInput, get_weather
9
+
10
+
11
+ class ToolsCollection:
12
+ @staticmethod
13
+ def get_tools(needed_tools: list[str]) -> list:
14
+ tools = []
15
+
16
+ for nt in needed_tools:
17
+ if nt == "search_tool":
18
+ tools.append(DuckDuckGoSearchRun())
19
+ elif nt == "get_weather":
20
+ tools.append(
21
+ StructuredTool(
22
+ name="get_weather",
23
+ func=get_weather,
24
+ description=get_weather.__doc__,
25
+ args_schema=WeatherInput,
26
+ )
27
+ )
28
+ elif nt == "hub_stats_tool":
29
+ tools.append(
30
+ StructuredTool(
31
+ name="hub_stats_tool",
32
+ func=get_hub_stats,
33
+ description=get_hub_stats.__doc__,
34
+ args_schema=HubStatInput,
35
+ )
36
+ )
37
+ elif nt == "wikipedia_en_tool_agent":
38
+ tools.append(
39
+ StructuredTool(
40
+ name="wikipedia_en_tool_agent",
41
+ func=wikipedia_en_tool_agent,
42
+ description=wikipedia_en_tool_agent.__doc__,
43
+ args_schema=WikipediaEnToolAgentInput,
44
+ )
45
+ )
46
+ elif nt == "EncyclopediaRetriever":
47
+ retriver = EncyclopediaRetriever(["gaia"], settings.PROJ_PATH)
48
+
49
+ tools.append(
50
+ StructuredTool(
51
+ name="get_related_question_from_encyclopedia",
52
+ func=retriver.get_related_question,
53
+ description=retriver.get_related_question.__doc__,
54
+ args_schema=EncyclopediaRetrieveInput,
55
+ )
56
+ )
57
+
58
+ return tools