Spaces:
Sleeping
Sleeping
| from __future__ import annotations | |
| import json | |
| import os | |
| from typing import Any, Dict, List, Tuple | |
| import requests | |
| from smolagents.tools import Tool | |
| class BrightDataSearchTool(Tool): | |
| name = "brightdata_search_engine" | |
| description = """ | |
| Search Google, Bing, or Yandex and get structured results. | |
| Returns search results with URLs, titles, and descriptions. | |
| Ideal for gathering current information and news. | |
| """ | |
| output_type = "string" | |
| engines: Tuple[str, ...] = ("google", "bing", "yandex") | |
| def __init__(self) -> None: | |
| self.inputs = { | |
| "query": { | |
| "type": "string", | |
| "description": "The search query", | |
| }, | |
| "engine": { | |
| "type": "string", | |
| "description": "Search engine to use", | |
| "enum": list(self.engines), | |
| "nullable": True, | |
| "default": "google", | |
| }, | |
| } | |
| super().__init__() | |
| def forward(self, query: str, engine: str = "google") -> str: | |
| """ | |
| Search using Bright Data's search API. | |
| Args: | |
| query: The search query. | |
| engine: Search engine to use (google, bing, or yandex). | |
| Returns: | |
| JSON string with search results or markdown for non-Google engines. | |
| """ | |
| api_token = os.getenv("BRIGHT_DATA_API_TOKEN") | |
| unlocker_zone = os.getenv("BRIGHT_DATA_UNLOCKER_ZONE", "web_unlocker1") | |
| if not api_token: | |
| raise ValueError("BRIGHT_DATA_API_TOKEN not found in environment variables") | |
| normalized_engine = engine.strip().lower() | |
| if normalized_engine not in self.engines: | |
| normalized_engine = "google" | |
| search_urls = self._build_search_urls(query) | |
| search_url = search_urls[normalized_engine] | |
| is_google = normalized_engine == "google" | |
| api_url = "https://api.brightdata.com/request" | |
| headers = { | |
| "Authorization": f"Bearer {api_token}", | |
| "Content-Type": "application/json", | |
| } | |
| payload = { | |
| "url": search_url, | |
| "zone": unlocker_zone, | |
| "format": "raw", | |
| } | |
| if not is_google: | |
| payload["data_format"] = "markdown" | |
| try: | |
| response = requests.post(api_url, json=payload, headers=headers, timeout=120) | |
| response.raise_for_status() | |
| if is_google: | |
| return self._format_google_results(response.json()) | |
| return response.text | |
| except requests.exceptions.RequestException as e: | |
| return json.dumps({"error": str(e)}) | |
| def _build_search_urls(self, query: str) -> Dict[str, str]: | |
| encoded_query = requests.utils.quote(query) | |
| return { | |
| "google": f"https://www.google.com/search?q={encoded_query}&brd_json=1", | |
| "bing": f"https://www.bing.com/search?q={encoded_query}", | |
| "yandex": f"https://yandex.com/search/?text={encoded_query}", | |
| } | |
| def _format_google_results(self, data: Dict[str, Any]) -> str: | |
| results: Dict[str, Any] = { | |
| "organic": data.get("organic", []), | |
| "images": self._extract_image_links(data.get("images", [])), | |
| "related": data.get("related", []), | |
| "ai_overview": data.get("ai_overview"), | |
| } | |
| return json.dumps(results, indent=2) | |
| def _extract_image_links(self, images: List[Dict[str, Any]]) -> List[str]: | |
| return [img.get("link", "") for img in images if img.get("link")] | |