Spaces:
Sleeping
Sleeping
Create tools/tool_collection_common.py
Browse files
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
|