refactor and patch litellm newer version dependency

Author: miguelg719

pyproject.toml

@@ -19,7 +19,7 @@ dependencies = [
  "rich>=13.7.0",
  "openai>=1.83.0,<1.99.6",
  "anthropic>=0.51.0",
- "litellm>=1.72.0",
+ "litellm>=1.72.0,<1.75.0",
  "nest-asyncio>=1.6.0",
 ]
 [[project.authors]]

stagehand/main.py

@@ -29,7 +29,7 @@
 from .logging import StagehandLogger, default_log_handler
 from .metrics import StagehandFunctionName, StagehandMetrics
 from .page import StagehandPage
-from .utils import make_serializable
+from .utils import get_download_path, make_serializable
 
 load_dotenv()
 
@@ -512,27 +512,6 @@ async def init(self):
                 )
                 self._playwright_page = self._page._page
 
-                # Set up download behavior via CDP
-                try:
-                    # Create CDP session for the page
-                    cdp_session = await self._context.new_cdp_session(
-                        self._playwright_page
-                    )
-                    # Enable download behavior
-                    await cdp_session.send(
-                        "Browser.setDownloadBehavior",
-                        {
-                            "behavior": "allow",
-                            "downloadPath": "downloads",
-                            "eventsEnabled": True,
-                        },
-                    )
-                    self.logger.debug("Set up CDP download behavior")
-                except Exception as e:
-                    self.logger.warning(
-                        f"Failed to set up CDP download behavior: {str(e)}"
-                    )
-
             except Exception:
                 await self.close()
                 raise
@@ -561,6 +540,23 @@ async def init(self):
             # Should not happen due to __init__ validation
             raise RuntimeError(f"Invalid env value: {self.env}")
 
+        # Set up download behavior via CDP
+        try:
+            # Create CDP session for the page
+            cdp_session = await self._context.new_cdp_session(self._playwright_page)
+            # Enable download behavior
+            await cdp_session.send(
+                "Browser.setDownloadBehavior",
+                {
+                    "behavior": "allow",
+                    "downloadPath": get_download_path(self),
+                    "eventsEnabled": True,
+                },
+            )
+            self.logger.debug("Set up CDP download behavior")
+        except Exception as e:
+            self.logger.warning(f"Failed to set up CDP download behavior: {str(e)}")
+
         self._initialized = True
 
     def agent(self, **kwargs) -> Agent:

stagehand/utils.py

@@ -1,4 +1,5 @@
 import inspect
+import os
 from typing import Any, Union, get_args, get_origin
 
 from pydantic import AnyUrl, BaseModel, Field, HttpUrl, create_model
@@ -530,3 +531,15 @@ def make_serializable(obj):
         elif isinstance(obj, dict):
             return {key: make_serializable(value) for key, value in obj.items()}
     return obj
+
+
+def get_download_path(stagehand):
+    if stagehand.env == "BROWSERBASE":
+        return "downloads"
+    else:
+        if stagehand.local_browser_launch_options.get("downloadPath"):
+            return stagehand.local_browser_launch_options["downloadPath"]
+        else:
+            path = os.path.join(os.getcwd(), "downloads")
+            os.makedirs(path, exist_ok=True)
+            return path