dlaima commited on
Commit
88ce352
·
verified ·
1 Parent(s): 2d05710

Create tools.py

Browse files
Files changed (1) hide show
  1. tools.py +52 -0
tools.py ADDED
@@ -0,0 +1,52 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import requests
2
+
3
+ def fetch_wikipedia_article(query: str) -> str:
4
+ """
5
+ Fetch the plain text of a Wikipedia article using the Wikipedia API.
6
+ """
7
+ url = "https://en.wikipedia.org/w/api.php"
8
+ params = {
9
+ "action": "query",
10
+ "prop": "extracts",
11
+ "explaintext": True,
12
+ "format": "json",
13
+ "titles": query,
14
+ }
15
+ response = requests.get(url, params=params)
16
+ data = response.json()
17
+ pages = data.get("query", {}).get("pages", {})
18
+ page = next(iter(pages.values()))
19
+ return page.get("extract", "No content found for this query.")
20
+
21
+
22
+ def search_wikipedia(query: str) -> str:
23
+ """
24
+ Search for a Wikipedia article title using the Wikipedia API.
25
+ """
26
+ url = "https://en.wikipedia.org/w/api.php"
27
+ params = {
28
+ "action": "query",
29
+ "list": "search",
30
+ "srsearch": query,
31
+ "format": "json",
32
+ }
33
+ response = requests.get(url, params=params)
34
+ data = response.json()
35
+ results = data.get("query", {}).get("search", [])
36
+ if results:
37
+ return "\n".join([f"{r['title']}" for r in results])
38
+ return "No search results found."
39
+
40
+
41
+ class WikipediaTool:
42
+ description = "Fetches plain text from a Wikipedia article. Input should be the article title."
43
+
44
+ def __call__(self, query: str) -> str:
45
+ return fetch_wikipedia_article(query)
46
+
47
+
48
+ class WikipediaSearchTool:
49
+ description = "Searches Wikipedia for article titles related to a query. Input should be a question or topic."
50
+
51
+ def __call__(self, query: str) -> str:
52
+ return search_wikipedia(query)