Spaces:
Build error
Build error
File size: 2,607 Bytes
a3a55bb |
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 48 49 50 51 52 53 54 55 56 57 58 |
import logging
import re
from smolagents import Tool
from smolagents.default_tools import DuckDuckGoSearchTool, WikipediaSearchTool
logger = logging.getLogger(__name__)
class SmartSearchTool(Tool):
name = "smart_search"
description = """A smart search tool that first performs a web search and then, if a Wikipedia article is found,
uses Wikipedia search for more reliable information."""
inputs = {"query": {"type": "string", "description": "The search query to find information"}}
output_type = "string"
def __init__(self):
super().__init__()
self.web_search_tool = DuckDuckGoSearchTool(max_results=1)
self.wiki_tool = WikipediaSearchTool(
user_agent="SmartSearchTool ([email protected])",
language="en",
# content_type="summary",
content_type="text",
extract_format="WIKI"
)
def forward(self, query: str) -> str:
logger.info(f"Starting smart search for query: {query}")
# First perform a web search with a single result
web_result = self.web_search_tool.forward(query)
logger.info(f"Web search result: {web_result[:100]}...")
# Check if the result contains a Wikipedia link
if "wikipedia.org" in web_result.lower():
logger.info("Wikipedia link found in web search results")
# Extract the Wikipedia page title from the URL using regex
wiki_match = re.search(r'wikipedia\.org/wiki/([^)\s]+)', web_result)
if wiki_match:
wiki_title = wiki_match.group(1)
logger.info(f"Extracted Wikipedia title: {wiki_title}")
# Use Wikipedia search for more reliable information
wiki_result = self.wiki_tool.forward(wiki_title)
logger.info(f"Wikipedia search result: {wiki_result[:100]}...")
if wiki_result and "No Wikipedia page found" not in wiki_result:
logger.info("Successfully retrieved Wikipedia content")
return f"Web search result:\n{web_result}\n\nWikipedia result:\n{wiki_result}"
else:
logger.warning("Wikipedia search failed or returned no results")
else:
logger.warning("Could not extract Wikipedia title from URL")
# If no Wikipedia link was found or Wikipedia search failed, return the web search result
logger.info("Returning web search result only")
return f"Web search result:\n{web_result}"
|