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)


class MathTool(Tool):
    name = "math"
    description = """Performs mathematical calculations and returns the result."""
    inputs = {
        "expression": {
            "type": "string",
            "description": "The mathematical expression to evaluate.",
        }
    }
    output_type = "string"

    def forward(self, expression: str) -> str:
        try:
            # Use eval with a restricted set of builtins for safety
            safe_dict = {
                "__builtins__": {
                    "abs": abs,
                    "round": round,
                    "min": min,
                    "max": max,
                    "sum": sum,
                }
            }
            result = eval(expression, safe_dict)
            return str(result)
        except Exception as e:
            raise Exception(f"Error evaluating expression: {str(e)}")


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