chore: add open in studio flag

lavague-integrations/exporters/lavague-exporters-python-selenium/lavague/exporters/python_selenium/qa.py

@@ -92,6 +92,7 @@
 Element outer HTML: {outer_html}
 Test:"""
 
+
 class QASeleniumExporter(PythonSeleniumExporter):
     def __init__(self, model: str = "gpt-4o", time_between_actions: float = 2.5):
         self.model = model
@@ -128,19 +129,23 @@ def export(self, trajectory: TrajectoryData, scenario: str) -> str:
                     outer_html: str = output.outer_html
 
                     prompt = PROMPT_TEMPLATE.format(
-                        test_specs=scenario, 
+                        test_specs=scenario,
                         context=context,
                         description=description,
                         text=text,
                         outer_html=outer_html,
                     )
-                    test_response = completion(
-                        model = self.model,
-                        messages = [
-                            {"role": "system", "content": SYSTEM_PROMPT},
-                            {"role": "user", "content": prompt}
-                        ]
-                    ).choices[0].message.content           
+                    test_response = (
+                        completion(
+                            model=self.model,
+                            messages=[
+                                {"role": "system", "content": SYSTEM_PROMPT},
+                                {"role": "user", "content": prompt},
+                            ],
+                        )
+                        .choices[0]
+                        .message.content
+                    )
 
                     generated_asserts: str = extract_code_block(test_response)
                     translated_action_lines.append(generated_asserts)
@@ -154,11 +159,12 @@ def export(self, trajectory: TrajectoryData, scenario: str) -> str:
         translated_actions_str: str = self.merge_code(*translated_actions)
         return self.merge_code(setup, translated_actions_str, teardown)
 
+
 def extract_code_block(code_block):
     pattern = r"```(?:python)?\n(.*?)\n```"
     match = re.search(pattern, code_block, re.DOTALL)
 
     if match:
         return match.group(1).strip()
     else:
-        return "No code block found"
+        return "No code block found"

lavague-sdk/lavague/sdk/agent.py

@@ -25,7 +25,12 @@ def __init__(
         self.create_public_runs = create_public_runs
 
     def run(
-        self, url: str, objective: str, async_run=False, viewport_size=(1096, 1096)
+        self,
+        url: str,
+        objective: str,
+        async_run=False,
+        viewport_size=(1096, 1096),
+        open_in_studio=False,
     ) -> Trajectory:
         request = RunRequest(
             url=url,
@@ -35,6 +40,8 @@ def run(
             viewport_size=viewport_size,
         )
         trajectory = self.client.run(request)
+        if open_in_studio:
+            trajectory.open_in_studio()
         if not async_run:
             trajectory.run_to_completion()
         return trajectory

lavague-sdk/lavague/sdk/trajectory/model.py

@@ -5,6 +5,7 @@
 from lavague.sdk.action.base import DEFAULT_PARSER
 from pydantic import model_validator
 
+
 class RunStatus(str, Enum):
     STARTING = "starting"
     RUNNING = "running"
@@ -31,32 +32,35 @@ class TrajectoryData(BaseModel):
     actions: List[SerializeAsAny[Action]]
     error_msg: Optional[str] = None
 
-    @model_validator(mode='before')
+    @model_validator(mode="before")
     @classmethod
     def deserialize_actions(cls, values: Dict[str, Any]) -> Dict[str, Any]:
-        if 'actions' in values:
-            actions = values['actions']
+        if "actions" in values:
+            actions = values["actions"]
             deserialized_actions = []
             for action_data in actions:
                 if isinstance(action_data, Action):
                     deserialized_actions.append(action_data)
                     continue
-                action_type = action_data.get('action_type')
+                action_type = action_data.get("action_type")
                 if action_type:
-                    action_class = DEFAULT_PARSER.engine_action_builders.get(action_type, Action)
+                    action_class = DEFAULT_PARSER.engine_action_builders.get(
+                        action_type, Action
+                    )
                     deserialized_action = action_class.parse(action_data)
                     deserialized_actions.append(deserialized_action)
                 else:
                     deserialized_actions.append(Action.parse(action_data))
-            values['actions'] = deserialized_actions
+            values["actions"] = deserialized_actions
         return values
 
     def write_to_file(self, file_path: str):
         json_model = self.model_dump_json(indent=2)
         with open(file_path, "w", encoding="utf-8") as file:
             file.write(json_model)
 
+
 class StepCompletion(BaseModel):
     run_status: RunStatus
     action: Optional[Action]
-    run_mode: RunMode
+    run_mode: RunMode