dlaima commited on
Commit
3becd38
·
verified ·
1 Parent(s): 021130d

Update tools.py

Browse files
Files changed (1) hide show
  1. tools.py +10 -6
tools.py CHANGED
@@ -1,3 +1,4 @@
 
1
  import requests
2
 
3
  def fetch_wikipedia_article(query: str) -> str:
@@ -18,7 +19,6 @@ def fetch_wikipedia_article(query: str) -> str:
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.
@@ -34,19 +34,23 @@ def search_wikipedia(query: str) -> str:
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)
 
1
+ from smolagents import Tool
2
  import requests
3
 
4
  def fetch_wikipedia_article(query: str) -> str:
 
19
  page = next(iter(pages.values()))
20
  return page.get("extract", "No content found for this query.")
21
 
 
22
  def search_wikipedia(query: str) -> str:
23
  """
24
  Search for a Wikipedia article title using the Wikipedia API.
 
34
  data = response.json()
35
  results = data.get("query", {}).get("search", [])
36
  if results:
37
+ return "\n".join([r['title'] for r in results])
38
  return "No search results found."
39
 
40
+ class WikipediaTool(Tool):
 
41
  description = "Fetches plain text from a Wikipedia article. Input should be the article title."
42
 
43
+ def __init__(self):
44
+ super().__init__(name="WikipediaTool", description=self.description)
45
+
46
  def __call__(self, query: str) -> str:
47
  return fetch_wikipedia_article(query)
48
 
49
+ class WikipediaSearchTool(Tool):
 
50
  description = "Searches Wikipedia for article titles related to a query. Input should be a question or topic."
51
 
52
+ def __init__(self):
53
+ super().__init__(name="WikipediaSearchTool", description=self.description)
54
+
55
  def __call__(self, query: str) -> str:
56
  return search_wikipedia(query)