diff --git a/lavague-integrations/exporters/lavague-exporters-python-selenium/lavague/exporters/python_selenium/qa.py b/lavague-integrations/exporters/lavague-exporters-python-selenium/lavague/exporters/python_selenium/qa.py index c1f055a7..a85d4164 100644 --- a/lavague-integrations/exporters/lavague-exporters-python-selenium/lavague/exporters/python_selenium/qa.py +++ b/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,6 +159,7 @@ 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) @@ -161,4 +167,4 @@ def extract_code_block(code_block): if match: return match.group(1).strip() else: - return "No code block found" \ No newline at end of file + return "No code block found" diff --git a/lavague-sdk/lavague/sdk/agent.py b/lavague-sdk/lavague/sdk/agent.py index c472e41e..14f566d3 100644 --- a/lavague-sdk/lavague/sdk/agent.py +++ b/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 diff --git a/lavague-sdk/lavague/sdk/trajectory/model.py b/lavague-sdk/lavague/sdk/trajectory/model.py index 4e0ebe3c..35774825 100644 --- a/lavague-sdk/lavague/sdk/trajectory/model.py +++ b/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,24 +32,26 @@ 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): @@ -56,7 +59,8 @@ def write_to_file(self, file_path: str): 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 \ No newline at end of file + run_mode: RunMode