ORromu commited on
Commit
9400406
·
verified ·
1 Parent(s): 3254e76

Updating tool.py

Browse files

Adding the first tools:
- Some basic algebra tools
- Some basic internet / databases (wikipedia, pubmed, arxiv) searching tools

Files changed (1) hide show
  1. tool.py +108 -0
tool.py CHANGED
@@ -0,0 +1,108 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from langchain_core.tools import tool
2
+ from langchain_community.utilities.duckduckgo_search import DuckDuckGoSearchAPIWrapper
3
+ from langchain_community.document_loaders.wikipedia import WikipediaLoader
4
+ from langchain_community.document_loaders.arxiv import ArxivLoader
5
+ from langchain_community.document_loaders.pubmed import PubMedLoader
6
+
7
+
8
+ ## Simple algebra tools
9
+ @tool
10
+ def add(a: float, b: float) -> float:
11
+ """Add two numbers.
12
+
13
+ Args:
14
+ a: first float
15
+ b: second float
16
+ """
17
+ return a + b
18
+
19
+ @tool
20
+ def substract(a: float, b: float) -> float:
21
+ """Substract two numbers.
22
+
23
+ Args:
24
+ a: first float
25
+ b: second float
26
+ """
27
+ return a - b
28
+
29
+ @tool
30
+ def multiply(a: float, b: float) -> float:
31
+ """Multiply two numbers.
32
+
33
+ Args:
34
+ a: first float
35
+ b: second float
36
+ """
37
+ return a * b
38
+
39
+ @tool
40
+ def divide(a: float, b: float) -> float:
41
+ """Divide two numbers.
42
+
43
+ Args:
44
+ a: first float
45
+ b: second float
46
+ """
47
+ if b == 0:
48
+ raise ValueError("Cannot divide any number by zero.")
49
+ return a / b
50
+
51
+ # Maybe add algebra tools???
52
+
53
+ ## Search Tools
54
+ def DuckDuckGoSearchTool(query: str) -> dict:
55
+ """Search DuckDuckGo for a query and return maximum 3 results.
56
+
57
+ Args:
58
+ query: The search query.
59
+ """
60
+ results = DuckDuckGoSearchAPIWrapper().results(query = query, max_results=3)
61
+ formatted_search_docs = "\n\n---\n\n".join(
62
+ [
63
+ f'<Document source="{res["link"]}" title="{res["title"]}">\n{res["snippet"]}\n</Document>'
64
+ for res in results
65
+ ])
66
+ return {"web_results": formatted_search_docs}
67
+
68
+ @tool
69
+ def WikipediaSearchTool(query: str) -> str:
70
+ """Search Wikipedia for a query and return maximum 2 results.
71
+
72
+ Args:
73
+ query: The search query."""
74
+ search_docs = WikipediaLoader(query=query, load_max_docs=2).load()
75
+ formatted_search_docs = "\n\n---\n\n".join(
76
+ [
77
+ f'<Document source="{doc.metadata["source"]}" page="{doc.metadata.get("page", "")}"/>\n{doc.page_content}\n</Document>'
78
+ for doc in search_docs
79
+ ])
80
+ return {"wiki_results": formatted_search_docs}
81
+
82
+ @tool
83
+ def ArxivSearchTool(query: str) -> str:
84
+ """Search Arxiv for a query and return maximum 3 result.
85
+
86
+ Args:
87
+ query: The search query."""
88
+ search_docs = ArxivLoader(query=query, load_max_docs=3).load()
89
+ formatted_search_docs = "\n\n---\n\n".join(
90
+ [
91
+ f'<Document source="{doc.metadata["source"]}" page="{doc.metadata.get("page", "")}"/>\n{doc.page_content[:1000]}\n</Document>'
92
+ for doc in search_docs
93
+ ])
94
+ return {"arvix_results": formatted_search_docs}
95
+
96
+ @tool
97
+ def PubmedSearchTool(query: str) -> str:
98
+ """Search Arxiv for a query and return maximum 3 result.
99
+
100
+ Args:
101
+ query: The search query."""
102
+ search_docs = PubMedLoader(query=query, load_max_docs=3).load()
103
+ formatted_search_docs = "\n\n---\n\n".join(
104
+ [
105
+ f'<Document source="{doc.metadata["uid"]}" title="{doc.metadata["Title"]}"/>\n{doc.page_content[:1000]}\n</Document>'
106
+ for doc in search_docs
107
+ ])
108
+ return {"pubmed_results": formatted_search_docs}