Merge pull request #1 from browserbase/main

merge to latest

Author: Neolithic

.changeset/fortunate-friendly-avocet.md

@@ -0,0 +1,5 @@
+---
+"stagehand": patch
+---
+
+Fix ability to pass raw JSON to Extract schema

.changeset/smiling-snobbish-cockle.md

@@ -0,0 +1,5 @@
+---
+"stagehand": patch
+---
+
+Pass api_timeout param to Stagehand API correctly

stagehand/api.py

@@ -27,6 +27,15 @@ async def _create_session(self):
         else None
     )
 
+    # API requires timeout while python sdk uses api_timeout which is converted to apiTimeout
+    if (
+        browserbase_session_create_params
+        and "apiTimeout" in browserbase_session_create_params
+    ):
+        browserbase_session_create_params["timeout"] = (
+            browserbase_session_create_params.pop("apiTimeout")
+        )
+
     payload = {
         "modelName": self.model_name,
         "verbose": 2 if self.verbose == 3 else self.verbose,

stagehand/handlers/extract_handler.py

@@ -151,9 +151,7 @@ async def extract(
 
         processed_data_payload = raw_data_dict  # Default to the raw dictionary
 
-        if schema and isinstance(
-            raw_data_dict, dict
-        ):  # schema is the Pydantic model type
+        if schema and isinstance(schema, type) and issubclass(schema, BaseModel):
             # Try direct validation first
             try:
                 validated_model_instance = schema.model_validate(raw_data_dict)

stagehand/llm/inference.py

@@ -169,11 +169,19 @@ async def extract(
     start_time = time.time()
 
     # Determine if we need to use schema-based response format
-    # TODO: if schema is json, return json
     response_format = {"type": "json_object"}
     if schema:
-        # If schema is a Pydantic model, use it directly
-        response_format = schema
+        if isinstance(schema, dict):
+            response_format = {
+                "type": "json_schema",
+                "json_schema": {
+                    "name": "extraction_schema",
+                    "strict": False,
+                    "schema": schema,
+                },
+            }
+        else:
+            response_format = schema
 
     # Call the LLM with appropriate parameters
     try: