updated browserbase tool

This commit is contained in:
Mish Ushakov
2024-05-20 08:58:25 +00:00
parent 53c7d815ae
commit 5c2d8c4cfa

View File

@@ -10,20 +10,35 @@ class BrowserbaseLoadTool(BaseTool):
description: str = "Load webpages url in a headless browser using Browserbase and return the contents"
args_schema: Type[BaseModel] = BrowserbaseLoadToolSchema
api_key: Optional[str] = None
project_id: Optional[str] = None
text_content: Optional[bool] = False
session_id: Optional[str] = None
proxy: Optional[bool] = None
browserbase: Optional[Any] = None
def __init__(self, api_key: Optional[str] = None, text_content: Optional[bool] = False, **kwargs):
def __init__(
self,
api_key: Optional[str] = None,
project_id: Optional[str] = None,
text_content: Optional[bool] = False,
session_id: Optional[str] = None,
proxy: Optional[bool] = None,
**kwargs,
):
super().__init__(**kwargs)
try:
from browserbase import Browserbase # type: ignore
from browserbase import Browserbase # type: ignore
except ImportError:
raise ImportError(
"`browserbase` package not found, please run `pip install browserbase`"
)
raise ImportError(
"`browserbase` package not found, please run `pip install browserbase`"
)
self.browserbase = Browserbase(api_key=api_key)
self.browserbase = Browserbase(api_key, project_id)
self.text_content = text_content
self.session_id = session_id
self.proxy = proxy
def _run(self, url: str):
return self.browserbase.load_url(url, text_content=self.text_content)
return self.browserbase.load_url(
url, self.text_content, self.session_id, self.proxy
)