Alfred828 commited on
Commit
10e6230
·
verified ·
1 Parent(s): c15b707

Create tools/tool_collection_wiki.py

Browse files
Files changed (1) hide show
  1. tools/tool_collection_wiki.py +45 -0
tools/tool_collection_wiki.py ADDED
@@ -0,0 +1,45 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from langchain.tools import StructuredTool
2
+
3
+ from tools.wiki_search import (
4
+ GetPageSectionContentInput,
5
+ GetPageTitleExcerptSectionsInput,
6
+ SearchWikipediaEn,
7
+ WikipediaOpensearchInput,
8
+ )
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 == "wikipedia_opensearch":
18
+ tools.append(
19
+ StructuredTool(
20
+ name="wikipedia_opensearch",
21
+ func=SearchWikipediaEn.wikipedia_opensearch,
22
+ description=SearchWikipediaEn.wikipedia_opensearch.__doc__,
23
+ args_schema=WikipediaOpensearchInput,
24
+ )
25
+ )
26
+ elif nt == "get_page_title_excerpt_sections":
27
+ tools.append(
28
+ StructuredTool(
29
+ name="get_page_title_excerpt_sections",
30
+ func=SearchWikipediaEn.get_page_title_excerpt_sections,
31
+ description=SearchWikipediaEn.get_page_title_excerpt_sections.__doc__,
32
+ args_schema=GetPageTitleExcerptSectionsInput,
33
+ )
34
+ )
35
+ elif nt == "get_page_section_content":
36
+ tools.append(
37
+ StructuredTool(
38
+ name="get_page_section_content",
39
+ func=SearchWikipediaEn.get_page_section_content,
40
+ description=SearchWikipediaEn.get_page_section_content.__doc__,
41
+ args_schema=GetPageSectionContentInput,
42
+ )
43
+ )
44
+
45
+ return tools