98 lines
3.4 KiB
Python
Executable File
98 lines
3.4 KiB
Python
Executable File
"""
|
|
Wikipedia Tool
|
|
Search and retrieve Wikipedia articles.
|
|
"""
|
|
from typing import Dict, Any, Optional
|
|
import wikipedia
|
|
from loguru import logger
|
|
|
|
from tools.base import BaseTool, ToolResult
|
|
|
|
|
|
class WikipediaTool(BaseTool):
|
|
"""Wikipedia search and retrieval."""
|
|
|
|
@property
|
|
def name(self) -> str:
|
|
return "wikipedia_search"
|
|
|
|
@property
|
|
def description(self) -> str:
|
|
return "Search Wikipedia for encyclopedia articles. Best for factual information, definitions, and historical topics."
|
|
|
|
@property
|
|
def parameters(self) -> Dict[str, Any]:
|
|
return {
|
|
"type": "object",
|
|
"properties": {
|
|
"query": {
|
|
"type": "string",
|
|
"description": "The search query"
|
|
},
|
|
"sentences": {
|
|
"type": "integer",
|
|
"description": "Number of sentences to return (default: 5)",
|
|
"default": 5
|
|
}
|
|
},
|
|
"required": ["query"]
|
|
}
|
|
|
|
async def execute(self, query: str, sentences: int = 5, **kwargs) -> ToolResult:
|
|
"""Execute Wikipedia search."""
|
|
self._log_execution({"query": query, "sentences": sentences})
|
|
|
|
try:
|
|
# Search for the page
|
|
search_results = wikipedia.search(query, results=3)
|
|
|
|
if not search_results:
|
|
return ToolResult(
|
|
success=True,
|
|
data="No Wikipedia articles found for this query."
|
|
)
|
|
|
|
# Try to get the first result
|
|
for title in search_results:
|
|
try:
|
|
page = wikipedia.page(title, auto_suggest=False)
|
|
summary = wikipedia.summary(title, sentences=sentences, auto_suggest=False)
|
|
|
|
output = (
|
|
f"Wikipedia Article: {page.title}\n"
|
|
f"URL: {page.url}\n\n"
|
|
f"Summary:\n{summary}"
|
|
)
|
|
|
|
self._log_success(output[:100])
|
|
return ToolResult(success=True, data=output)
|
|
|
|
except wikipedia.exceptions.DisambiguationError as e:
|
|
# Try the first option
|
|
try:
|
|
page = wikipedia.page(e.options[0], auto_suggest=False)
|
|
summary = wikipedia.summary(e.options[0], sentences=sentences, auto_suggest=False)
|
|
|
|
output = (
|
|
f"Wikipedia Article: {page.title}\n"
|
|
f"URL: {page.url}\n\n"
|
|
f"Summary:\n{summary}"
|
|
)
|
|
|
|
self._log_success(output[:100])
|
|
return ToolResult(success=True, data=output)
|
|
except Exception:
|
|
continue
|
|
|
|
except wikipedia.exceptions.PageError:
|
|
continue
|
|
|
|
return ToolResult(
|
|
success=True,
|
|
data="Could not find a specific Wikipedia article. Try a more specific query."
|
|
)
|
|
|
|
except Exception as e:
|
|
self._log_error(str(e))
|
|
return ToolResult(success=False, error=str(e))
|