Spaces:
Sleeping
Sleeping
File size: 5,727 Bytes
4ea4bca | 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 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 | from typing import Any, Optional
from smolagents.tools import Tool
import os
import requests
import html
import xml
class WebSearchTool(Tool):
name = "web_search"
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."
inputs = {'query': {'type': 'string', 'description': 'The search query to perform.'}}
output_type = "string"
def __init__(self, max_results: int = 10, engine: str = "duckduckgo"):
super().__init__()
self.max_results = max_results
self.engine = engine
def forward(self, query: str) -> str:
results = self.search(query)
if len(results) == 0:
raise Exception("No results found! Try a less restrictive/shorter query.")
return self.parse_results(results)
def search(self, query: str) -> list:
if self.engine == "duckduckgo":
return self.search_duckduckgo(query)
elif self.engine == "bing":
return self.search_bing(query)
elif self.engine == "exa":
return self.search_exa(query)
else:
raise ValueError(f"Unsupported engine: {self.engine}")
def parse_results(self, results: list) -> str:
return "## Search Results\n\n" + "\n\n".join(
[f"[{result['title']}]({result['link']})\n{result['description']}" for result in results]
)
def search_duckduckgo(self, query: str) -> list:
import requests
response = requests.get(
"https://lite.duckduckgo.com/lite/",
params={"q": query},
headers={"User-Agent": "Mozilla/5.0"},
)
response.raise_for_status()
parser = self._create_duckduckgo_parser()
parser.feed(response.text)
return parser.results
def _create_duckduckgo_parser(self):
from html.parser import HTMLParser
class SimpleResultParser(HTMLParser):
def __init__(self):
super().__init__()
self.results = []
self.current = {}
self.capture_title = False
self.capture_description = False
self.capture_link = False
def handle_starttag(self, tag, attrs):
attrs = dict(attrs)
if tag == "a" and attrs.get("class") == "result-link":
self.capture_title = True
elif tag == "td" and attrs.get("class") == "result-snippet":
self.capture_description = True
elif tag == "span" and attrs.get("class") == "link-text":
self.capture_link = True
def handle_endtag(self, tag):
if tag == "a" and self.capture_title:
self.capture_title = False
elif tag == "td" and self.capture_description:
self.capture_description = False
elif tag == "span" and self.capture_link:
self.capture_link = False
elif tag == "tr":
# Store current result if all parts are present
if {"title", "description", "link"} <= self.current.keys():
self.current["description"] = " ".join(self.current["description"])
self.results.append(self.current)
self.current = {}
def handle_data(self, data):
if self.capture_title:
self.current["title"] = data.strip()
elif self.capture_description:
self.current.setdefault("description", [])
self.current["description"].append(data.strip())
elif self.capture_link:
self.current["link"] = "https://" + data.strip()
return SimpleResultParser()
def search_bing(self, query: str) -> list:
import xml.etree.ElementTree as ET
import requests
response = requests.get(
"https://www.bing.com/search",
params={"q": query, "format": "rss"},
)
response.raise_for_status()
root = ET.fromstring(response.text)
items = root.findall(".//item")
results = [
{
"title": item.findtext("title"),
"link": item.findtext("link"),
"description": item.findtext("description"),
}
for item in items[: self.max_results]
]
return results
def search_exa(self, query: str) -> list:
"""Search using the Exa API. Requires an EXA_API_KEY environment variable."""
import os
import requests
api_key = os.getenv("EXA_API_KEY")
if not api_key:
raise ValueError("Missing API key. Make sure you have 'EXA_API_KEY' in your env variables.")
response = requests.post(
"https://api.exa.ai/search",
headers={
"x-api-key": api_key,
"Content-Type": "application/json",
"x-exa-integration": "smolagents",
},
json={
"query": query,
"numResults": self.max_results,
"contents": {"highlights": True},
},
timeout=getattr(self, "timeout", 30),
)
response.raise_for_status()
data = response.json()
return [
{
"title": result.get("title", ""),
"link": result["url"],
"description": " ".join(result.get("highlights") or []),
}
for result in data.get("results", [])
]
|