services/tools.py

import logging

from smolagents import DuckDuckGoSearchTool, Tool, WikipediaSearchTool

logger = logging.getLogger(__name__)


class GeneralSearchTool(Tool):
    name = "search"
    description = """Performs a general web search using both DuckDuckGo and Wikipedia, then returns the combined search results."""
    inputs = {
        "query": {"type": "string", "description": "The search query to perform."}
    }
    output_type = "string"

    def __init__(self, max_results=10, **kwargs):
        super().__init__()
        self.max_results = max_results
        self.ddg_tool = DuckDuckGoSearchTool()
        self.wiki_tool = WikipediaSearchTool()

    def forward(self, query: str) -> str:
        # Get DuckDuckGo results
        try:
            ddg_results = self.ddg_tool.forward(query)
        except Exception as e:
            ddg_results = "No DuckDuckGo results found."
            logger.warning(f"DuckDuckGo search failed: {str(e)}")

        # Get Wikipedia results
        try:
            wiki_results = self.wiki_tool.forward(query)
        except Exception as e:
            wiki_results = "No Wikipedia results found."
            logger.warning(f"Wikipedia search failed: {str(e)}")

        # Combine and format results
        output = []
        if ddg_results and ddg_results != "No DuckDuckGo results found.":
            output.append("## DuckDuckGo Search Results\n\n" + ddg_results)
        if wiki_results and wiki_results != "No Wikipedia results found.":
            output.append("## Wikipedia Results\n\n" + wiki_results)

        if not output:
            raise Exception("No results found! Try a less restrictive/shorter query.")

        return "\n\n---\n\n".join(output)


# Export all tools
tools = [
    # DuckDuckGoSearchTool(),
    GeneralSearchTool(),
    # WikipediaSearchTool(),
]