mirror of
https://github.com/crewAIInc/crewAI.git
synced 2026-07-02 21:58:11 +00:00
feat: Add WikipediaSearchTool for free research
Closes #5823 - Add WikipediaSearchTool that searches Wikipedia and returns concise summaries - Handle disambiguation errors by listing possible matches - Handle page not found errors gracefully - Support configurable sentences, language, and top_k results - Add wikipedia as optional dependency - Add 14 unit tests with full mock coverage Co-Authored-By: João <joao@crewai.com>
This commit is contained in:
@@ -147,6 +147,9 @@ e2b = [
|
||||
"e2b~=2.20.0",
|
||||
"e2b-code-interpreter~=2.6.0",
|
||||
]
|
||||
wikipedia = [
|
||||
"wikipedia>=1.4.0",
|
||||
]
|
||||
|
||||
|
||||
[tool.uv]
|
||||
|
||||
@@ -208,6 +208,9 @@ from crewai_tools.tools.txt_search_tool.txt_search_tool import TXTSearchTool
|
||||
from crewai_tools.tools.vision_tool.vision_tool import VisionTool
|
||||
from crewai_tools.tools.weaviate_tool.vector_search import WeaviateVectorSearchTool
|
||||
from crewai_tools.tools.website_search.website_search_tool import WebsiteSearchTool
|
||||
from crewai_tools.tools.wikipedia_search_tool.wikipedia_search_tool import (
|
||||
WikipediaSearchTool,
|
||||
)
|
||||
from crewai_tools.tools.xml_search_tool.xml_search_tool import XMLSearchTool
|
||||
from crewai_tools.tools.youtube_channel_search_tool.youtube_channel_search_tool import (
|
||||
YoutubeChannelSearchTool,
|
||||
@@ -323,6 +326,7 @@ __all__ = [
|
||||
"VisionTool",
|
||||
"WeaviateVectorSearchTool",
|
||||
"WebsiteSearchTool",
|
||||
"WikipediaSearchTool",
|
||||
"XMLSearchTool",
|
||||
"YoutubeChannelSearchTool",
|
||||
"YoutubeVideoSearchTool",
|
||||
|
||||
@@ -195,6 +195,9 @@ from crewai_tools.tools.txt_search_tool.txt_search_tool import TXTSearchTool
|
||||
from crewai_tools.tools.vision_tool.vision_tool import VisionTool
|
||||
from crewai_tools.tools.weaviate_tool.vector_search import WeaviateVectorSearchTool
|
||||
from crewai_tools.tools.website_search.website_search_tool import WebsiteSearchTool
|
||||
from crewai_tools.tools.wikipedia_search_tool.wikipedia_search_tool import (
|
||||
WikipediaSearchTool,
|
||||
)
|
||||
from crewai_tools.tools.xml_search_tool.xml_search_tool import XMLSearchTool
|
||||
from crewai_tools.tools.youtube_channel_search_tool.youtube_channel_search_tool import (
|
||||
YoutubeChannelSearchTool,
|
||||
@@ -306,6 +309,7 @@ __all__ = [
|
||||
"VisionTool",
|
||||
"WeaviateVectorSearchTool",
|
||||
"WebsiteSearchTool",
|
||||
"WikipediaSearchTool",
|
||||
"XMLSearchTool",
|
||||
"YoutubeChannelSearchTool",
|
||||
"YoutubeVideoSearchTool",
|
||||
|
||||
@@ -0,0 +1,104 @@
|
||||
import logging
|
||||
from typing import Any, ClassVar
|
||||
|
||||
from crewai.tools import BaseTool
|
||||
from pydantic import BaseModel, ConfigDict, Field
|
||||
|
||||
|
||||
logger = logging.getLogger(__file__)
|
||||
|
||||
try:
|
||||
import wikipedia
|
||||
except ImportError:
|
||||
wikipedia = None
|
||||
|
||||
|
||||
class WikipediaSearchToolInput(BaseModel):
|
||||
"""Input for WikipediaSearchTool."""
|
||||
|
||||
search_query: str = Field(
|
||||
..., description="The topic or query to search for on Wikipedia"
|
||||
)
|
||||
|
||||
|
||||
class WikipediaSearchTool(BaseTool):
|
||||
"""Tool for searching Wikipedia and retrieving article summaries.
|
||||
|
||||
Uses the `wikipedia` Python library to search for topics and return
|
||||
concise summaries. Handles disambiguation and missing pages gracefully.
|
||||
"""
|
||||
|
||||
name: str = "Wikipedia Search"
|
||||
description: str = (
|
||||
"A tool that searches Wikipedia for a given topic and returns "
|
||||
"a concise summary. Useful for quick factual lookups and research."
|
||||
)
|
||||
args_schema: type[BaseModel] = WikipediaSearchToolInput
|
||||
sentences: int = 5
|
||||
language: str = "en"
|
||||
top_k: int = 1
|
||||
|
||||
model_config: ClassVar[ConfigDict] = ConfigDict(extra="allow")
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
sentences: int = 5,
|
||||
language: str = "en",
|
||||
top_k: int = 1,
|
||||
**kwargs: Any,
|
||||
):
|
||||
super().__init__(**kwargs)
|
||||
self.sentences = sentences
|
||||
self.language = language
|
||||
self.top_k = top_k
|
||||
|
||||
def _run(self, search_query: str, **kwargs: Any) -> str:
|
||||
if wikipedia is None:
|
||||
return (
|
||||
"The 'wikipedia' package is required to use WikipediaSearchTool. "
|
||||
"Install it with: pip install wikipedia"
|
||||
)
|
||||
|
||||
wikipedia.set_lang(self.language)
|
||||
|
||||
try:
|
||||
results = wikipedia.search(search_query, results=self.top_k)
|
||||
except Exception as e:
|
||||
logger.error(f"Wikipedia search error: {e}")
|
||||
return f"Error searching Wikipedia: {e}"
|
||||
|
||||
if not results:
|
||||
return f"No Wikipedia articles found for '{search_query}'."
|
||||
|
||||
output_parts = [self._fetch_article(title) for title in results]
|
||||
|
||||
return "\n\n---\n\n".join(output_parts)
|
||||
|
||||
def _fetch_article(self, title: str) -> str:
|
||||
try:
|
||||
summary = wikipedia.summary(title, sentences=self.sentences)
|
||||
page = wikipedia.page(title, auto_suggest=False)
|
||||
return (
|
||||
f"Title: {page.title}\n"
|
||||
f"URL: {page.url}\n"
|
||||
f"Summary: {summary}"
|
||||
)
|
||||
except wikipedia.exceptions.DisambiguationError as e:
|
||||
options = ", ".join(e.options[:5])
|
||||
return (
|
||||
f"Title: {title}\n"
|
||||
f"Disambiguation: The term '{title}' is ambiguous. "
|
||||
f"Possible options: {options}. "
|
||||
f"Please refine your search query."
|
||||
)
|
||||
except wikipedia.exceptions.PageError:
|
||||
return (
|
||||
f"Title: {title}\n"
|
||||
f"Error: No Wikipedia page found for '{title}'."
|
||||
)
|
||||
except Exception as e:
|
||||
logger.error(f"Error retrieving Wikipedia page '{title}': {e}")
|
||||
return (
|
||||
f"Title: {title}\n"
|
||||
f"Error: Failed to retrieve page: {e}"
|
||||
)
|
||||
208
lib/crewai-tools/tests/tools/test_wikipedia_search_tool.py
Normal file
208
lib/crewai-tools/tests/tools/test_wikipedia_search_tool.py
Normal file
@@ -0,0 +1,208 @@
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
from crewai_tools import WikipediaSearchTool
|
||||
|
||||
MODULE_PATH = "crewai_tools.tools.wikipedia_search_tool.wikipedia_search_tool"
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def tool():
|
||||
return WikipediaSearchTool()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def tool_short():
|
||||
return WikipediaSearchTool(sentences=2)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def tool_multi():
|
||||
return WikipediaSearchTool(top_k=3)
|
||||
|
||||
|
||||
def _make_mock_wikipedia():
|
||||
mock_wiki = MagicMock()
|
||||
mock_wiki.exceptions = MagicMock()
|
||||
mock_wiki.exceptions.DisambiguationError = type(
|
||||
"DisambiguationError", (Exception,), {}
|
||||
)
|
||||
mock_wiki.exceptions.PageError = type("PageError", (Exception,), {})
|
||||
return mock_wiki
|
||||
|
||||
|
||||
class FakePage:
|
||||
def __init__(self, title, url="https://en.wikipedia.org/wiki/Test"):
|
||||
self.title = title
|
||||
self.url = url
|
||||
|
||||
|
||||
class TestWikipediaSearchToolInit:
|
||||
def test_default_params(self, tool):
|
||||
assert tool.sentences == 5
|
||||
assert tool.language == "en"
|
||||
assert tool.top_k == 1
|
||||
|
||||
def test_custom_params(self):
|
||||
tool = WikipediaSearchTool(sentences=3, language="fr", top_k=2)
|
||||
assert tool.sentences == 3
|
||||
assert tool.language == "fr"
|
||||
assert tool.top_k == 2
|
||||
|
||||
def test_tool_metadata(self, tool):
|
||||
assert tool.name == "Wikipedia Search"
|
||||
assert "Wikipedia" in tool.description
|
||||
|
||||
|
||||
class TestWikipediaSearchToolRun:
|
||||
def test_basic_search(self, tool):
|
||||
mock_wiki = _make_mock_wikipedia()
|
||||
mock_wiki.search.return_value = ["Python (programming language)"]
|
||||
mock_wiki.summary.return_value = "Python is a programming language."
|
||||
mock_wiki.page.return_value = FakePage(
|
||||
"Python (programming language)",
|
||||
"https://en.wikipedia.org/wiki/Python_(programming_language)",
|
||||
)
|
||||
|
||||
with patch(f"{MODULE_PATH}.wikipedia", mock_wiki):
|
||||
result = tool._run(search_query="Python programming")
|
||||
|
||||
mock_wiki.set_lang.assert_called_once_with("en")
|
||||
mock_wiki.search.assert_called_once_with("Python programming", results=1)
|
||||
assert "Python (programming language)" in result
|
||||
assert "Python is a programming language." in result
|
||||
assert "https://en.wikipedia.org/wiki/Python_(programming_language)" in result
|
||||
|
||||
def test_no_results(self, tool):
|
||||
mock_wiki = _make_mock_wikipedia()
|
||||
mock_wiki.search.return_value = []
|
||||
|
||||
with patch(f"{MODULE_PATH}.wikipedia", mock_wiki):
|
||||
result = tool._run(search_query="xyznonexistenttopic123")
|
||||
|
||||
assert "No Wikipedia articles found" in result
|
||||
|
||||
def test_disambiguation_error(self, tool):
|
||||
mock_wiki = _make_mock_wikipedia()
|
||||
mock_wiki.search.return_value = ["Mercury"]
|
||||
|
||||
exc = mock_wiki.exceptions.DisambiguationError(
|
||||
"Mercury", ["Mercury (planet)", "Mercury (element)", "Freddie Mercury"]
|
||||
)
|
||||
exc.options = ["Mercury (planet)", "Mercury (element)", "Freddie Mercury"]
|
||||
mock_wiki.summary.side_effect = exc
|
||||
|
||||
with patch(f"{MODULE_PATH}.wikipedia", mock_wiki):
|
||||
result = tool._run(search_query="Mercury")
|
||||
|
||||
assert "Disambiguation" in result
|
||||
assert "Mercury (planet)" in result
|
||||
|
||||
def test_page_not_found(self, tool):
|
||||
mock_wiki = _make_mock_wikipedia()
|
||||
mock_wiki.search.return_value = ["NonexistentPage12345"]
|
||||
mock_wiki.summary.side_effect = mock_wiki.exceptions.PageError(
|
||||
"NonexistentPage12345"
|
||||
)
|
||||
|
||||
with patch(f"{MODULE_PATH}.wikipedia", mock_wiki):
|
||||
result = tool._run(search_query="NonexistentPage12345")
|
||||
|
||||
assert "No Wikipedia page found" in result
|
||||
|
||||
def test_generic_exception(self, tool):
|
||||
mock_wiki = _make_mock_wikipedia()
|
||||
mock_wiki.search.return_value = ["Test"]
|
||||
mock_wiki.summary.side_effect = RuntimeError("Network error")
|
||||
|
||||
with patch(f"{MODULE_PATH}.wikipedia", mock_wiki):
|
||||
result = tool._run(search_query="Test")
|
||||
|
||||
assert "Failed to retrieve page" in result
|
||||
|
||||
def test_search_error(self, tool):
|
||||
mock_wiki = _make_mock_wikipedia()
|
||||
mock_wiki.search.side_effect = Exception("API error")
|
||||
|
||||
with patch(f"{MODULE_PATH}.wikipedia", mock_wiki):
|
||||
result = tool._run(search_query="Test")
|
||||
|
||||
assert "Error searching Wikipedia" in result
|
||||
|
||||
def test_custom_sentences(self, tool_short):
|
||||
mock_wiki = _make_mock_wikipedia()
|
||||
mock_wiki.search.return_value = ["Test"]
|
||||
mock_wiki.summary.return_value = "Short summary."
|
||||
mock_wiki.page.return_value = FakePage("Test")
|
||||
|
||||
with patch(f"{MODULE_PATH}.wikipedia", mock_wiki):
|
||||
tool_short._run(search_query="Test")
|
||||
|
||||
mock_wiki.summary.assert_called_once_with("Test", sentences=2)
|
||||
|
||||
def test_multiple_results(self, tool_multi):
|
||||
mock_wiki = _make_mock_wikipedia()
|
||||
mock_wiki.search.return_value = ["AI", "ML", "DL"]
|
||||
mock_wiki.summary.side_effect = [
|
||||
"AI summary.",
|
||||
"ML summary.",
|
||||
"DL summary.",
|
||||
]
|
||||
mock_wiki.page.side_effect = [
|
||||
FakePage("Artificial intelligence"),
|
||||
FakePage("Machine learning"),
|
||||
FakePage("Deep learning"),
|
||||
]
|
||||
|
||||
with patch(f"{MODULE_PATH}.wikipedia", mock_wiki):
|
||||
result = tool_multi._run(search_query="AI")
|
||||
|
||||
mock_wiki.search.assert_called_once_with("AI", results=3)
|
||||
assert "AI summary." in result
|
||||
assert "ML summary." in result
|
||||
assert "DL summary." in result
|
||||
assert result.count("---") == 2
|
||||
|
||||
def test_language_setting(self):
|
||||
tool = WikipediaSearchTool(language="fr")
|
||||
mock_wiki = _make_mock_wikipedia()
|
||||
mock_wiki.search.return_value = ["Test"]
|
||||
mock_wiki.summary.return_value = "Résumé du test."
|
||||
mock_wiki.page.return_value = FakePage("Test")
|
||||
|
||||
with patch(f"{MODULE_PATH}.wikipedia", mock_wiki):
|
||||
tool._run(search_query="Test")
|
||||
|
||||
mock_wiki.set_lang.assert_called_once_with("fr")
|
||||
|
||||
def test_missing_wikipedia_package(self, tool):
|
||||
with patch(f"{MODULE_PATH}.wikipedia", None):
|
||||
result = tool._run(search_query="Test")
|
||||
|
||||
assert "wikipedia" in result.lower()
|
||||
assert "pip install" in result
|
||||
|
||||
|
||||
class TestWikipediaSearchToolMixedResults:
|
||||
def test_mixed_success_and_disambiguation(self):
|
||||
tool = WikipediaSearchTool(top_k=2)
|
||||
mock_wiki = _make_mock_wikipedia()
|
||||
mock_wiki.search.return_value = ["Python", "Mercury"]
|
||||
|
||||
exc = mock_wiki.exceptions.DisambiguationError(
|
||||
"Mercury", ["Mercury (planet)", "Mercury (element)"]
|
||||
)
|
||||
exc.options = ["Mercury (planet)", "Mercury (element)"]
|
||||
|
||||
mock_wiki.summary.side_effect = [
|
||||
"Python is a language.",
|
||||
exc,
|
||||
]
|
||||
mock_wiki.page.return_value = FakePage("Python")
|
||||
|
||||
with patch(f"{MODULE_PATH}.wikipedia", mock_wiki):
|
||||
result = tool._run(search_query="Python Mercury")
|
||||
|
||||
assert "Python is a language." in result
|
||||
assert "Disambiguation" in result
|
||||
Reference in New Issue
Block a user