Skip to content
This repository was archived by the owner on Nov 10, 2025. It is now read-only.

Commit e7a182a

Browse files
mplachtamplachta
authored
feat: add required env var validation for Bright Data API key in tools (#476)
1 parent d1fee6e commit e7a182a

File tree

3 files changed

+100
-19
lines changed

3 files changed

+100
-19
lines changed

‎crewai_tools/tools/brightdata_tool/brightdata_dataset.py‎

Lines changed: 36 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,12 @@
11
import asyncio
22
import os
3-
from typing import Any, Dict, Optional, Type
3+
from typing import Any, Dict, List, Optional, Type
44

55
import aiohttp
6-
from crewai.tools import BaseTool
6+
from crewai.tools import BaseTool, EnvVar
77
from pydantic import BaseModel, Field
88

9+
910
class BrightDataConfig(BaseModel):
1011
API_URL: str = "https://api.brightdata.com"
1112
DEFAULT_TIMEOUT: int = 600
@@ -16,8 +17,12 @@ def from_env(cls):
1617
return cls(
1718
API_URL=os.environ.get("BRIGHTDATA_API_URL", "https://api.brightdata.com"),
1819
DEFAULT_TIMEOUT=int(os.environ.get("BRIGHTDATA_DEFAULT_TIMEOUT", "600")),
19-
DEFAULT_POLLING_INTERVAL=int(os.environ.get("BRIGHTDATA_DEFAULT_POLLING_INTERVAL", "1"))
20+
DEFAULT_POLLING_INTERVAL=int(
21+
os.environ.get("BRIGHTDATA_DEFAULT_POLLING_INTERVAL", "1")
22+
),
2023
)
24+
25+
2126
class BrightDataDatasetToolException(Exception):
2227
"""Exception raised for custom error in the application."""
2328

@@ -52,6 +57,7 @@ class BrightDataDatasetToolSchema(BaseModel):
5257
default=None, description="Additional params if any"
5358
)
5459

60+
5561
config = BrightDataConfig.from_env()
5662

5763
BRIGHTDATA_API_URL = config.API_URL
@@ -410,8 +416,22 @@ class BrightDataDatasetTool(BaseTool):
410416
format: str = "json"
411417
zipcode: Optional[str] = None
412418
additional_params: Optional[Dict[str, Any]] = None
419+
env_vars: List[EnvVar] = [
420+
EnvVar(
421+
name="BRIGHT_DATA_API_KEY",
422+
description="API key for Bright Data",
423+
required=True,
424+
),
425+
]
413426

414-
def __init__(self, dataset_type: str = None, url: str = None, format: str = "json", zipcode: str = None, additional_params: Dict[str, Any] = None):
427+
def __init__(
428+
self,
429+
dataset_type: str = None,
430+
url: str = None,
431+
format: str = "json",
432+
zipcode: str = None,
433+
additional_params: Dict[str, Any] = None,
434+
):
415435
super().__init__()
416436
self.dataset_type = dataset_type
417437
self.url = url
@@ -530,15 +550,25 @@ async def get_dataset_data_async(
530550

531551
return await snapshot_response.text()
532552

533-
def _run(self, url: str = None, dataset_type: str = None, format: str = None, zipcode: str = None, additional_params: Dict[str, Any] = None, **kwargs: Any) -> Any:
553+
def _run(
554+
self,
555+
url: str = None,
556+
dataset_type: str = None,
557+
format: str = None,
558+
zipcode: str = None,
559+
additional_params: Dict[str, Any] = None,
560+
**kwargs: Any,
561+
) -> Any:
534562
dataset_type = dataset_type or self.dataset_type
535563
output_format = format or self.format
536564
url = url or self.url
537565
zipcode = zipcode or self.zipcode
538566
additional_params = additional_params or self.additional_params
539567

540568
if not dataset_type:
541-
raise ValueError("dataset_type is required either in constructor or method call")
569+
raise ValueError(
570+
"dataset_type is required either in constructor or method call"
571+
)
542572
if not url:
543573
raise ValueError("url is required either in constructor or method call")
544574

‎crewai_tools/tools/brightdata_tool/brightdata_serp.py‎

Lines changed: 39 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1,20 +1,24 @@
11
import os
22
import urllib.parse
3-
from typing import Any, Optional, Type
3+
from typing import Any, List, Optional, Type
44

55
import requests
6-
from crewai.tools import BaseTool
6+
from crewai.tools import BaseTool, EnvVar
77
from pydantic import BaseModel, Field
88

9+
910
class BrightDataConfig(BaseModel):
1011
API_URL: str = "https://api.brightdata.com/request"
1112

1213
@classmethod
1314
def from_env(cls):
1415
return cls(
15-
API_URL=os.environ.get("BRIGHTDATA_API_URL", "https://api.brightdata.com/request")
16+
API_URL=os.environ.get(
17+
"BRIGHTDATA_API_URL", "https://api.brightdata.com/request"
18+
)
1619
)
1720

21+
1822
class BrightDataSearchToolSchema(BaseModel):
1923
"""
2024
Schema that defines the input arguments for the BrightDataSearchToolSchema.
@@ -87,8 +91,24 @@ class BrightDataSearchTool(BaseTool):
8791
search_type: Optional[str] = None
8892
device_type: str = "desktop"
8993
parse_results: bool = True
90-
91-
def __init__(self, query: str = None, search_engine: str = "google", country: str = "us", language: str = "en", search_type: str = None, device_type: str = "desktop", parse_results: bool = True):
94+
env_vars: List[EnvVar] = [
95+
EnvVar(
96+
name="BRIGHT_DATA_API_KEY",
97+
description="API key for Bright Data",
98+
required=True,
99+
),
100+
]
101+
102+
def __init__(
103+
self,
104+
query: str = None,
105+
search_engine: str = "google",
106+
country: str = "us",
107+
language: str = "en",
108+
search_type: str = None,
109+
device_type: str = "desktop",
110+
parse_results: bool = True,
111+
):
92112
super().__init__()
93113
self.base_url = self._config.API_URL
94114
self.query = query
@@ -113,7 +133,17 @@ def get_search_url(self, engine: str, query: str):
113133
return f"https://www.bing.com/search?q=${query}"
114134
return f"https://www.google.com/search?q=${query}"
115135

116-
def _run(self, query: str = None, search_engine: str = None, country: str = None, language: str = None, search_type: str = None, device_type: str = None, parse_results: bool = None, **kwargs) -> Any:
136+
def _run(
137+
self,
138+
query: str = None,
139+
search_engine: str = None,
140+
country: str = None,
141+
language: str = None,
142+
search_type: str = None,
143+
device_type: str = None,
144+
parse_results: bool = None,
145+
**kwargs,
146+
) -> Any:
117147
"""
118148
Executes a search query using Bright Data SERP API and returns results.
119149
@@ -137,7 +167,9 @@ def _run(self, query: str = None, search_engine: str = None, country: str = None
137167
language = language or self.language
138168
search_type = search_type or self.search_type
139169
device_type = device_type or self.device_type
140-
parse_results = parse_results if parse_results is not None else self.parse_results
170+
parse_results = (
171+
parse_results if parse_results is not None else self.parse_results
172+
)
141173
results_count = kwargs.get("results_count", "10")
142174

143175
# Validate required parameters

‎crewai_tools/tools/brightdata_tool/brightdata_unlocker.py‎

Lines changed: 25 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1,19 +1,23 @@
11
import os
2-
from typing import Any, Optional, Type
2+
from typing import Any, List, Optional, Type
33

44
import requests
5-
from crewai.tools import BaseTool
5+
from crewai.tools import BaseTool, EnvVar
66
from pydantic import BaseModel, Field
77

8+
89
class BrightDataConfig(BaseModel):
910
API_URL: str = "https://api.brightdata.com/request"
1011

1112
@classmethod
1213
def from_env(cls):
1314
return cls(
14-
API_URL=os.environ.get("BRIGHTDATA_API_URL", "https://api.brightdata.com/request")
15+
API_URL=os.environ.get(
16+
"BRIGHTDATA_API_URL", "https://api.brightdata.com/request"
17+
)
1518
)
1619

20+
1721
class BrightDataUnlockerToolSchema(BaseModel):
1822
"""
1923
Pydantic schema for input parameters used by the BrightDataWebUnlockerTool.
@@ -67,8 +71,17 @@ class BrightDataWebUnlockerTool(BaseTool):
6771
url: Optional[str] = None
6872
format: str = "raw"
6973
data_format: str = "markdown"
70-
71-
def __init__(self, url: str = None, format: str = "raw", data_format: str = "markdown"):
74+
env_vars: List[EnvVar] = [
75+
EnvVar(
76+
name="BRIGHT_DATA_API_KEY",
77+
description="API key for Bright Data",
78+
required=True,
79+
),
80+
]
81+
82+
def __init__(
83+
self, url: str = None, format: str = "raw", data_format: str = "markdown"
84+
):
7285
super().__init__()
7386
self.base_url = self._config.API_URL
7487
self.url = url
@@ -82,7 +95,13 @@ def __init__(self, url: str = None, format: str = "raw", data_format: str = "mar
8295
if not self.zone:
8396
raise ValueError("BRIGHT_DATA_ZONE environment variable is required.")
8497

85-
def _run(self, url: str = None, format: str = None, data_format: str = None, **kwargs: Any) -> Any:
98+
def _run(
99+
self,
100+
url: str = None,
101+
format: str = None,
102+
data_format: str = None,
103+
**kwargs: Any,
104+
) -> Any:
86105
url = url or self.url
87106
format = format or self.format
88107
data_format = data_format or self.data_format

0 commit comments

Comments
 (0)