From 5c2d8c4cfa86debc8f769cc9a8838a21ae1a4f72 Mon Sep 17 00:00:00 2001
From: Mish Ushakov <10400064+mishushakov@users.noreply.github.com>
Date: Mon, 20 May 2024 08:58:25 +0000
Subject: [PATCH] updated browserbase tool

---
 .../browserbase_load_tool.py                  | 29 ++++++++++++++-----
 1 file changed, 22 insertions(+), 7 deletions(-)

diff --git a/src/crewai_tools/tools/browserbase_load_tool/browserbase_load_tool.py b/src/crewai_tools/tools/browserbase_load_tool/browserbase_load_tool.py
index 48f3dacef..52722520d 100644
--- a/src/crewai_tools/tools/browserbase_load_tool/browserbase_load_tool.py
+++ b/src/crewai_tools/tools/browserbase_load_tool/browserbase_load_tool.py
@@ -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
+        )