Add Bright Data Dataset Tool

app.py

@@ -0,0 +1,5 @@
+from smolagents import launch_gradio_demo
+from tool import BrightDataDatasetTool
+
+tool = BrightDataDatasetTool()
+launch_gradio_demo(tool)

requirements.txt

@@ -0,0 +1,2 @@
+requests
+smolagents

tool.py

@@ -0,0 +1,195 @@
+from typing import Any, Optional
+from smolagents.tools import Tool
+import time
+import json
+import requests
+import os
+
+class BrightDataDatasetTool(Tool):
+    name = "brightdata_dataset_fetch"
+    description = "Trigger a Bright Data dataset collection and poll until the snapshot is ready. Choose a dataset key (e.g., amazon_product, linkedin_company_profile, google_maps_reviews). For most datasets, you only need to provide the URL parameter. For example: brightdata_dataset_fetch(dataset='linkedin_person_profile', url='https://linkedin.com/in/...')"
+    output_type = "string"
+
+    def __init__(self):
+        # Keep dataset catalogue on the instance and build the inputs schema dynamically to satisfy tool validation.
+        self.datasets = globals().get("DATASETS", {})
+        if not self.datasets:
+            raise ValueError("Dataset catalogue is not available.")
+        self.inputs = {
+            "dataset": {
+                "type": "string",
+                "description": f"Dataset key. Options: {', '.join(sorted(self.datasets.keys()))}",
+            },
+            "url": {
+                "type": "string",
+                "description": "URL for the dataset (required for most datasets)",
+                "nullable": True,
+            },
+            "keyword": {
+                "type": "string",
+                "description": "Search keyword (for search datasets like amazon_product_search)",
+                "nullable": True,
+            },
+            "first_name": {
+                "type": "string",
+                "description": "First name (for datasets like linkedin_people_search)",
+                "nullable": True,
+            },
+            "last_name": {
+                "type": "string",
+                "description": "Last name (for datasets like linkedin_people_search)",
+                "nullable": True,
+            },
+            "days_limit": {
+                "type": "string",
+                "description": "Days limit (for datasets like google_maps_reviews, default: 3)",
+                "nullable": True,
+            },
+            "num_of_reviews": {
+                "type": "string",
+                "description": "Number of reviews (for datasets like facebook_company_reviews)",
+                "nullable": True,
+            },
+            "num_of_comments": {
+                "type": "string",
+                "description": "Number of comments (for datasets like youtube_comments, default: 10)",
+                "nullable": True,
+            },
+        }
+        super().__init__()
+
+    def _prepare_payload(self, dataset_key: str, params):
+        """Validate required fields, apply defaults, and merge fixed values."""
+        config = self.datasets[dataset_key]
+        payload = {}
+
+        defaults = config.get("defaults", {})
+        fixed_values = config.get("fixed_values", {})
+
+        for field in config["inputs"]:
+            if field in params:
+                payload[field] = params[field]
+            elif field in defaults:
+                payload[field] = defaults[field]
+            else:
+                raise ValueError(f"Missing required field '{field}' for dataset '{dataset_key}'")
+
+        # Apply fixed values that should always be sent
+        payload.update(fixed_values)
+        return payload
+
+    def forward(
+        self,
+        dataset: str,
+        url: str = None,
+        keyword: str = None,
+        first_name: str = None,
+        last_name: str = None,
+        days_limit: str = None,
+        num_of_reviews: str = None,
+        num_of_comments: str = None,
+    ) -> str:
+        """
+        Trigger a dataset run and poll until results are ready.
+
+        Args:
+            dataset: The dataset key from DATASETS.
+            url: URL for the dataset (required for most datasets).
+            keyword: Search keyword (for search datasets).
+            first_name: First name (for people search datasets).
+            last_name: Last name (for people search datasets).
+            days_limit: Days limit (for time-based datasets).
+            num_of_reviews: Number of reviews to fetch.
+            num_of_comments: Number of comments to fetch.
+
+        Returns:
+            JSON string of the snapshot data once ready.
+        """
+        import os
+        import json
+        import time
+        import requests
+
+        api_token = os.getenv("BRIGHT_DATA_API_TOKEN")
+        if not api_token:
+            raise ValueError("BRIGHT_DATA_API_TOKEN not found in environment variables")
+
+        if dataset not in self.datasets:
+            raise ValueError(f"Unknown dataset '{dataset}'. Valid options: {', '.join(sorted(self.datasets.keys()))}")
+
+        # Build params dict from provided arguments
+        params = {}
+        if url is not None:
+            params["url"] = url
+        if keyword is not None:
+            params["keyword"] = keyword
+        if first_name is not None:
+            params["first_name"] = first_name
+        if last_name is not None:
+            params["last_name"] = last_name
+        if days_limit is not None:
+            params["days_limit"] = days_limit
+        if num_of_reviews is not None:
+            params["num_of_reviews"] = num_of_reviews
+        if num_of_comments is not None:
+            params["num_of_comments"] = num_of_comments
+
+        payload = self._prepare_payload(dataset, params)
+        dataset_id = self.datasets[dataset]["dataset_id"]
+
+        trigger_url = "https://api.brightdata.com/datasets/v3/trigger"
+        trigger_headers = {
+            "Authorization": f"Bearer {api_token}",
+            "Content-Type": "application/json",
+        }
+
+        trigger_response = requests.post(
+            trigger_url,
+            params={"dataset_id": dataset_id, "include_errors": "true"},
+            json=[payload],
+            headers=trigger_headers,
+            timeout=60,
+        )
+        trigger_response.raise_for_status()
+        snapshot_id = trigger_response.json().get("snapshot_id")
+
+        if not snapshot_id:
+            raise RuntimeError("No snapshot ID returned from Bright Data.")
+
+        # Poll for completion (up to 10 minutes, matching MCP logic)
+        snapshot_url = f"https://api.brightdata.com/datasets/v3/snapshot/{snapshot_id}"
+        max_attempts = 600
+        attempts = 0
+
+        while attempts < max_attempts:
+            try:
+                response = requests.get(
+                    snapshot_url,
+                    params={"format": "json"},
+                    headers={"Authorization": f"Bearer {api_token}"},
+                    timeout=30,
+                )
+
+                # If Bright Data returns an error response we don't want to loop forever
+                if response.status_code == 400:
+                    response.raise_for_status()
+
+                data = response.json()
+                if isinstance(data, list):
+                    return json.dumps(data, indent=2)
+
+                status = data.get("status") if isinstance(data, dict) else None
+                if status not in {"running", "building"}:
+                    return json.dumps(data, indent=2)
+
+                attempts += 1
+                time.sleep(1)
+
+            except requests.exceptions.RequestException as exc:
+                # Mirror JS logic: tolerate transient failures, but break on 400
+                if getattr(getattr(exc, "response", None), "status_code", None) == 400:
+                    raise
+                attempts += 1
+                time.sleep(1)
+
+        raise TimeoutError(f"Timeout waiting for snapshot {snapshot_id} after {max_attempts} seconds")