File size: 2,798 Bytes
5acc791
8fe992b
 
5ba26df
5acc791
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
8fe992b
 
 
 
 
 
 
 
 
 
5acc791
 
 
 
8fe992b
5acc791
8fe992b
 
5acc791
 
 
 
 
5ba26df
 
 
 
 
 
 
 
 
 
5acc791
 
 
 
5ba26df
5acc791
5ba26df
5acc791
 
 
 
 
 
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
from typing import Any, Optional, Callable, List
from smolagents.tools import Tool
import duckduckgo_search
import googlesearch
from functools import wraps

def setup_search_dependency(package_name: str, import_func: Callable):
    """Utility function to handle search dependency setup"""
    try:
        return import_func()
    except ImportError as e:
        raise ImportError(
            f"You must install package `{package_name}` to run this tool: "
            f"for instance run `pip install {package_name}`."
        ) from e

def handle_search_errors(func):
    """Decorator to handle common search error cases"""
    @wraps(func)
    def wrapper(*args, **kwargs):
        try:
            results = func(*args, **kwargs)
            if not results:
                raise Exception("No results found! Try a less restrictive/shorter query.")
            return results
        except Exception as e:
            return f"Error performing search: {str(e)}"
    return wrapper

class DuckDuckGoSearchTool(Tool):
    name = "web_search"
    description = "Performs a duckduckgo web search based on your query (think a Google search) then returns the top 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.ddgs = setup_search_dependency(
            'duckduckgo-search',
            lambda: duckduckgo_search.DDGS(**kwargs)
        )

    @handle_search_errors
    def forward(self, query: str) -> str:
        results = self.ddgs.text(query, max_results=self.max_results)
        formatted_results = [
            f"[{result['title']}]({result['href']})\n{result['body']}" 
            for result in results
        ]
        return "## Search Results\n\n" + "\n\n".join(formatted_results)

class GoogleSearchTool(Tool):
    name = "google_search"
    description = "Performs a Google web search based on your query and returns the top 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.search = setup_search_dependency(
            'googlesearch-python',
            lambda: googlesearch.search
        )

    @handle_search_errors
    def forward(self, query: str) -> str:
        search_results = list(self.search(query, num_results=self.max_results))
        formatted_results = [
            f"[Result {i+1}]({url})\n{url}" 
            for i, url in enumerate(search_results)
        ]
        return "## Search Results\n\n" + "\n\n".join(formatted_results)