dlaima commited on
Commit
34f6dce
·
verified ·
1 Parent(s): 710c4a0

Update tools.py

Browse files
Files changed (1) hide show
  1. tools.py +37 -4
tools.py CHANGED
@@ -31,6 +31,7 @@ def search_wikipedia(query: str) -> str:
31
  return "\n".join([r['title'] for r in results])
32
  return "No search results found."
33
 
 
34
  class WikipediaTool(Tool):
35
  name = "WikipediaTool"
36
  description = "Fetches plain text from a Wikipedia article. Input should be the article title."
@@ -38,8 +39,7 @@ class WikipediaTool(Tool):
38
  output_type = "string"
39
 
40
  def __call__(self, inputs: dict) -> str:
41
- query = inputs["query"]
42
- return fetch_wikipedia_article(query)
43
 
44
  class WikipediaSearchTool(Tool):
45
  name = "WikipediaSearchTool"
@@ -48,7 +48,40 @@ class WikipediaSearchTool(Tool):
48
  output_type = "string"
49
 
50
  def __call__(self, inputs: dict) -> str:
51
- query = inputs["query"]
52
- return search_wikipedia(query)
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
53
 
54
 
 
31
  return "\n".join([r['title'] for r in results])
32
  return "No search results found."
33
 
34
+ # ✅ Actual tools expecting dict inputs
35
  class WikipediaTool(Tool):
36
  name = "WikipediaTool"
37
  description = "Fetches plain text from a Wikipedia article. Input should be the article title."
 
39
  output_type = "string"
40
 
41
  def __call__(self, inputs: dict) -> str:
42
+ return fetch_wikipedia_article(inputs["query"])
 
43
 
44
  class WikipediaSearchTool(Tool):
45
  name = "WikipediaSearchTool"
 
48
  output_type = "string"
49
 
50
  def __call__(self, inputs: dict) -> str:
51
+ return search_wikipedia(inputs["query"])
52
+
53
+ # ✅ Wrapper Tools (these are called by CodeAgent)
54
+ class WikipediaToolWrapper(Tool):
55
+ name = "WikipediaToolWrapper"
56
+ description = WikipediaTool.description
57
+ inputs = {"query": "string"}
58
+ output_type = "string"
59
+
60
+ def __init__(self):
61
+ self.tool = WikipediaTool()
62
+
63
+ def __call__(self, inputs: dict) -> str:
64
+ if isinstance(inputs, str):
65
+ inputs = {"query": inputs}
66
+ elif isinstance(inputs, dict) and "query" not in inputs:
67
+ inputs = {"query": next(iter(inputs.values()))}
68
+ return self.tool(inputs)
69
+
70
+ class WikipediaSearchToolWrapper(Tool):
71
+ name = "WikipediaSearchToolWrapper"
72
+ description = WikipediaSearchTool.description
73
+ inputs = {"query": "string"}
74
+ output_type = "string"
75
+
76
+ def __init__(self):
77
+ self.tool = WikipediaSearchTool()
78
+
79
+ def __call__(self, inputs: dict) -> str:
80
+ if isinstance(inputs, str):
81
+ inputs = {"query": inputs}
82
+ elif isinstance(inputs, dict) and "query" not in inputs:
83
+ inputs = {"query": next(iter(inputs.values()))}
84
+ return self.tool(inputs)
85
+
86
 
87