jnath-devops/AlfredAgent

1from typing import Any, Optional2from smolagents.tools import Tool3import requests4import html5import xml6 7class WebSearchTool(Tool):8    name = "web_search"9    description = "Performs a web search for a query and returns a string of the top search results formatted as markdown with titles, links, and descriptions."10    inputs = {'query': {'type': 'string', 'description': 'The search query to perform.'}}11    output_type = "string"12 13    def __init__(self, max_results: int = 10, engine: str = "duckduckgo"):14        super().__init__()15        self.max_results = max_results16        self.engine = engine17 18    def forward(self, query: str) -> str:19        results = self.search(query)20        if len(results) == 0:21            raise Exception("No results found! Try a less restrictive/shorter query.")22        return self.parse_results(results)23 24    def search(self, query: str) -> list:25        if self.engine == "duckduckgo":26            return self.search_duckduckgo(query)27        elif self.engine == "bing":28            return self.search_bing(query)29        else:30            raise ValueError(f"Unsupported engine: {self.engine}")31 32    def parse_results(self, results: list) -> str:33        return "## Search Results\n\n" + "\n\n".join(34            [f"[{result['title']}]({result['link']})\n{result['description']}" for result in results]35        )36 37    def search_duckduckgo(self, query: str) -> list:38        import requests39 40        response = requests.get(41            "https://lite.duckduckgo.com/lite/",42            params={"q": query},43            headers={"User-Agent": "Mozilla/5.0"},44        )45        response.raise_for_status()46        parser = self._create_duckduckgo_parser()47        parser.feed(response.text)48        return parser.results49 50    def _create_duckduckgo_parser(self):51        from html.parser import HTMLParser52 53        class SimpleResultParser(HTMLParser):54            def __init__(self):55                super().__init__()56                self.results = []57                self.current = {}58                self.capture_title = False59                self.capture_description = False60                self.capture_link = False61 62            def handle_starttag(self, tag, attrs):63                attrs = dict(attrs)64                if tag == "a" and attrs.get("class") == "result-link":65                    self.capture_title = True66                elif tag == "td" and attrs.get("class") == "result-snippet":67                    self.capture_description = True68                elif tag == "span" and attrs.get("class") == "link-text":69                    self.capture_link = True70 71            def handle_endtag(self, tag):72                if tag == "a" and self.capture_title:73                    self.capture_title = False74                elif tag == "td" and self.capture_description:75                    self.capture_description = False76                elif tag == "span" and self.capture_link:77                    self.capture_link = False78                elif tag == "tr":79                    # Store current result if all parts are present80                    if {"title", "description", "link"} <= self.current.keys():81                        self.current["description"] = " ".join(self.current["description"])82                        self.results.append(self.current)83                        self.current = {}84 85            def handle_data(self, data):86                if self.capture_title:87                    self.current["title"] = data.strip()88                elif self.capture_description:89                    self.current.setdefault("description", [])90                    self.current["description"].append(data.strip())91                elif self.capture_link:92                    self.current["link"] = "https://" + data.strip()93 94        return SimpleResultParser()95 96    def search_bing(self, query: str) -> list:97        import xml.etree.ElementTree as ET98 99        import requests100 101        response = requests.get(102            "https://www.bing.com/search",103            params={"q": query, "format": "rss"},104        )105        response.raise_for_status()106        root = ET.fromstring(response.text)107        items = root.findall(".//item")108        results = [109            {110                "title": item.findtext("title"),111                "link": item.findtext("link"),112                "description": item.findtext("description"),113            }114            for item in items[: self.max_results]115        ]116        return results117