File size: 3,470 Bytes
a3a55bb ea0c151 |
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 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 |
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}"
def main(query: str) -> str:
"""
Test function to run the SmartSearchTool directly.
Args:
query: The search query to test
Returns:
The search results
"""
# Configure logging
logging.basicConfig(
level=logging.INFO,
format="%(asctime)s - %(name)s - %(levelname)s - %(message)s"
)
# Create and run the tool
tool = SmartSearchTool()
result = tool.forward(query)
# Print the result
print("\nSearch Results:")
print("-" * 80)
print(result)
print("-" * 80)
return result
if __name__ == "__main__":
import sys
if len(sys.argv) > 1:
query = " ".join(sys.argv[1:])
main(query)
else:
print("Usage: python tool.py <search query>")
print("Example: python tool.py 'Mercedes Sosa discography'")
|