Add WikipediaSearchTool to default tools (#514)

Co-authored-by: Aymeric Roucher <[email protected]>

Author: touseefahmed96

pyproject.toml

@@ -88,6 +88,7 @@ test = [
   "python-dotenv>=1.0.1", # For test_all_docs
   "smolagents[all]",
   "rank-bm25", # For test_all_docs
+  "Wikipedia-API>=0.8.1",
 ]
 dev = [
   "smolagents[quality,test]",

src/smolagents/default_tools.py

@@ -261,6 +261,102 @@ def forward(self, url: str) -> str:
             return f"An unexpected error occurred: {str(e)}"
 
 
+class WikipediaSearchTool(Tool):
+    """
+    WikipediaSearchTool searches Wikipedia and returns a summary or full text of the given topic, along with the page URL.
+
+    Attributes:
+        user_agent (str): A custom user-agent string to identify the project. This is required as per Wikipedia API policies, read more here: http://github.com/martin-majlis/Wikipedia-API/blob/master/README.rst
+        language (str): The language in which to retrieve Wikipedia articles.
+                http://meta.wikimedia.org/wiki/List_of_Wikipedias
+        content_type (str): Defines the content to fetch. Can be "summary" for a short summary or "text" for the full article.
+        extract_format (str): Defines the output format. Can be `"WIKI"` or `"HTML"`.
+
+    Example:
+        >>> from smolagents import CodeAgent, HfApiModel, WikipediaSearchTool
+        >>> agent = CodeAgent(
+        >>>     tools=[
+        >>>            WikipediaSearchTool(
+        >>>                user_agent="MyResearchBot ([email protected])",
+        >>>                language="en",
+        >>>                content_type="summary",  # or "text"
+        >>>                extract_format="WIKI",
+        >>>            )
+        >>>        ],
+        >>>     model=HfApiModel(),
+        >>> )
+        >>> agent.run("Python_(programming_language)")
+    """
+
+    name = "wikipedia_search"
+    description = "Searches Wikipedia and returns a summary or full text of the given topic, along with the page URL."
+    inputs = {
+        "query": {
+            "type": "string",
+            "description": "The topic to search on Wikipedia.",
+        }
+    }
+    output_type = "string"
+
+    def __init__(
+        self,
+        user_agent: str = "Smolagents ([email protected])",
+        language: str = "en",
+        content_type: str = "text",
+        extract_format: str = "WIKI",
+    ):
+        super().__init__()
+        try:
+            import wikipediaapi
+        except ImportError as e:
+            raise ImportError(
+                "You must install `wikipedia-api` to run this tool: for instance run `pip install wikipedia-api`"
+            ) from e
+        if not user_agent:
+            raise ValueError("User-agent is required. Provide a meaningful identifier for your project.")
+
+        self.user_agent = user_agent
+        self.language = language
+        self.content_type = content_type
+
+        # Map string format to wikipediaapi.ExtractFormat
+        extract_format_map = {
+            "WIKI": wikipediaapi.ExtractFormat.WIKI,
+            "HTML": wikipediaapi.ExtractFormat.HTML,
+        }
+
+        if extract_format not in extract_format_map:
+            raise ValueError("Invalid extract_format. Choose between 'WIKI' or 'HTML'.")
+
+        self.extract_format = extract_format_map[extract_format]
+
+        self.wiki = wikipediaapi.Wikipedia(
+            user_agent=self.user_agent, language=self.language, extract_format=self.extract_format
+        )
+
+    def forward(self, query: str) -> str:
+        try:
+            page = self.wiki.page(query)
+
+            if not page.exists():
+                return f"No Wikipedia page found for '{query}'. Try a different query."
+
+            title = page.title
+            url = page.fullurl
+
+            if self.content_type == "summary":
+                text = page.summary
+            elif self.content_type == "text":
+                text = page.text
+            else:
+                return "⚠️ Invalid `content_type`. Use either 'summary' or 'text'."
+
+            return f"✅ **Wikipedia Page:** {title}\n\n**Content:** {text}\n\n🔗 **Read more:** {url}"
+
+        except Exception as e:
+            return f"Error fetching Wikipedia summary: {str(e)}"
+
+
 class SpeechToTextTool(PipelineTool):
     default_checkpoint = "openai/whisper-large-v3-turbo"
     description = "This is a tool that transcribes an audio into text. It returns the transcribed text."
@@ -312,5 +408,6 @@ def decode(self, outputs):
     "DuckDuckGoSearchTool",
     "GoogleSearchTool",
     "VisitWebpageTool",
+    "WikipediaSearchTool",
     "SpeechToTextTool",
 ]

tests/test_default_tools.py

@@ -17,7 +17,13 @@
 import pytest
 
 from smolagents.agent_types import _AGENT_TYPE_MAPPING
-from smolagents.default_tools import DuckDuckGoSearchTool, PythonInterpreterTool, SpeechToTextTool, VisitWebpageTool
+from smolagents.default_tools import (
+    DuckDuckGoSearchTool,
+    PythonInterpreterTool,
+    SpeechToTextTool,
+    VisitWebpageTool,
+    WikipediaSearchTool,
+)
 
 from .test_tools import ToolTesterMixin
 
@@ -87,3 +93,32 @@ def test_new_instance(self):
         assert tool is not None
         assert tool.pre_processor_class == WhisperProcessor
         assert tool.model_class == WhisperForConditionalGeneration
+
+
+@pytest.mark.parametrize(
+    "language, content_type, extract_format, query",
+    [
+        ("en", "summary", "HTML", "Python_(programming_language)"),  # English, Summary Mode, HTML format
+        ("en", "text", "WIKI", "Python_(programming_language)"),  # English, Full Text Mode, WIKI format
+        ("es", "summary", "HTML", "Python_(lenguaje_de_programación)"),  # Spanish, Summary Mode, HTML format
+        ("es", "text", "WIKI", "Python_(lenguaje_de_programación)"),  # Spanish, Full Text Mode, WIKI format
+    ],
+)
+def test_wikipedia_search(language, content_type, extract_format, query):
+    tool = WikipediaSearchTool(
+        user_agent="TestAgent ([email protected])",
+        language=language,
+        content_type=content_type,
+        extract_format=extract_format,
+    )
+
+    result = tool.forward(query)
+
+    assert isinstance(result, str), "Output should be a string"
+    assert "✅ **Wikipedia Page:**" in result, "Response should contain Wikipedia page title"
+    assert "🔗 **Read more:**" in result, "Response should contain Wikipedia page URL"
+
+    if content_type == "summary":
+        assert len(result.split()) < 1000, "Summary mode should return a shorter text"
+    if content_type == "text":
+        assert len(result.split()) > 1000, "Full text mode should return a longer text"