Spaces:
Runtime error
Runtime error
File size: 1,825 Bytes
cd5abb0 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 |
from typing import Any, Optional
from smolagents.tools import Tool
from langchain_community.document_loaders import WikipediaLoader
class WikipediaSearchTool(Tool):
name = "wikipedia_search"
description = "Search Wikipedia articles based on a query and return the content. Useful for gathering information about topics, people, places, events, etc."
inputs = {'query': {'type': 'string', 'description': 'The search query to look up on Wikipedia.'}}
output_type = "string"
def __init__(self, load_max_docs=2, **kwargs):
super().__init__()
self.load_max_docs = load_max_docs
self.is_initialized = True
def forward(self, query: str) -> str:
try:
# Use WikipediaLoader from langchain_community to load articles
loader = WikipediaLoader(
query=query,
load_max_docs=self.load_max_docs,
load_all_available_meta=True
)
# Get the documents (articles)
docs = loader.load()
if not docs:
return f"No Wikipedia articles found for the query: {query}"
# Format the results nicely
results = []
for doc in docs:
# Extract metadata and content
metadata = doc.metadata
title = metadata.get('title', 'Untitled')
summary = doc.page_content[:500] + "..." if len(doc.page_content) > 500 else doc.page_content
# Format each article
article = f"## {title}\n\n{summary}\n\n"
results.append(article)
return "\n".join(results)
except Exception as e:
return f"Error searching Wikipedia: {str(e)}"
|