lisaterumi commited on
Commit
552b51a
·
verified ·
1 Parent(s): f0a2cc6

Create tools.py

Browse files
Files changed (1) hide show
  1. tools.py +137 -0
tools.py ADDED
@@ -0,0 +1,137 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from langchain_core.tools import tool
2
+ from langchain_community.document_loaders import WikipediaLoader
3
+ from langchain_community.document_loaders import ArxivLoader
4
+ from langchain_community.tools.tavily_search import TavilySearchResults
5
+ from difflib import SequenceMatcher
6
+
7
+ @tool
8
+ def add(a: float, b: float) -> float:
9
+ """
10
+ Adds two numbers.
11
+ Args:
12
+ a (float): the first number
13
+ b (float): the second number
14
+ """
15
+ return a + b
16
+
17
+ @tool
18
+ def subtract(a: float, b: float) -> float:
19
+ """
20
+ Subtracts two numbers.
21
+ Args:
22
+ a (float): the first number
23
+ b (float): the second number
24
+ """
25
+ return a - b
26
+
27
+ @tool
28
+ def multiply(a: float, b: float) -> float:
29
+ """
30
+ Multiplies two numbers.
31
+ Args:
32
+ a (float): the first number
33
+ b (float): the second number
34
+ """
35
+ return a * b
36
+
37
+ @tool
38
+ def divide(a: float, b: float) -> float:
39
+ """
40
+ Divides two numbers.
41
+ Args:
42
+ a (float): the first float number
43
+ b (float): the second float number
44
+ """
45
+ if b == 0:
46
+ raise ValueError("Cannot divided by zero.")
47
+ return a / b
48
+
49
+ @tool
50
+ def modulus(a: int, b: int) -> int:
51
+ """
52
+ Get the modulus of two numbers.
53
+ Args:
54
+ a (int): the first number
55
+ b (int): the second number
56
+ """
57
+ return a % b
58
+
59
+ @tool
60
+ def wiki_search(query: str) -> str:
61
+ """Search Wikipedia for a query and return maximum 2 results.
62
+
63
+ Args:
64
+ query: The search query."""
65
+ search_docs = WikipediaLoader(query=query, load_max_docs=2).load()
66
+ formatted_search_docs = "\n\n---\n\n".join(
67
+ [
68
+ f'<Document source="{doc.metadata["source"]}" page="{doc.metadata.get("page", "")}"/>\n{doc.page_content}\n</Document>'
69
+ for doc in search_docs
70
+ ])
71
+ return {"wiki_results": formatted_search_docs}
72
+
73
+ @tool
74
+ def web_search(query: str) -> str:
75
+ """Search Tavily for a query and return maximum 3 results.
76
+
77
+ Args:
78
+ query: The search query."""
79
+ search_docs = TavilySearchResults(max_results=3).invoke(query=query)
80
+ formatted_search_docs = "\n\n---\n\n".join(
81
+ [
82
+ f'<Document source="{doc.metadata["source"]}" page="{doc.metadata.get("page", "")}"/>\n{doc.page_content}\n</Document>'
83
+ for doc in search_docs
84
+ ])
85
+ return {"web_results": formatted_search_docs}
86
+
87
+ @tool
88
+ def arvix_search(query: str) -> str:
89
+ """Search Arxiv for a query and return maximum 3 result.
90
+
91
+ Args:
92
+ query: The search query."""
93
+ search_docs = ArxivLoader(query=query, load_max_docs=3).load()
94
+ formatted_search_docs = "\n\n---\n\n".join(
95
+ [
96
+ f'<Document source="{doc.metadata["source"]}" page="{doc.metadata.get("page", "")}"/>\n{doc.page_content[:1000]}\n</Document>'
97
+ for doc in search_docs
98
+ ])
99
+ return {"arvix_results": formatted_search_docs}
100
+
101
+ @tool
102
+ def search_metadata(query: str) -> str:
103
+ """Search through metadata.jsonl file for matching questions and answers.
104
+
105
+ Args:
106
+ query: The search query to match against questions in the metadata file.
107
+ """
108
+ import json
109
+
110
+ results = []
111
+ with open("metadata.jsonl", "r", encoding="utf-8") as f:
112
+ for line in f:
113
+ try:
114
+ data = json.loads(line)
115
+ # Calculate similarity ratio
116
+ similarity = SequenceMatcher(None, query.lower(), data["Question"].lower()).ratio()
117
+ if similarity > 0.6: # Threshold for similarity
118
+ results.append({
119
+ "question": data["Question"],
120
+ "answer": data["Final answer"],
121
+ "steps": data["Annotator Metadata"]["Steps"],
122
+ "similarity": similarity
123
+ })
124
+ except json.JSONDecodeError:
125
+ continue
126
+
127
+ if not results:
128
+ return "No matching results found in metadata."
129
+
130
+ # Sort by similarity
131
+ results.sort(key=lambda x: x["similarity"], reverse=True)
132
+
133
+ formatted_results = "\n\n---\n\n".join(
134
+ [f"Question: {r['question']}\nAnswer: {r['answer']}\nSteps: {r['steps']}"
135
+ for r in results[:3]] # Return top 3 most similar results
136
+ )
137
+ return formatted_results