diff --git a/ .env.example b/ .env.example new file mode 100644 index 00000000..fec6db9c --- /dev/null +++ b/ .env.example @@ -0,0 +1,4 @@ +MODEL_API_KEY = "anthropic-or-openai-api-key" +BROWSERBASE_API_KEY = "browserbase-api-key" +BROWSERBASE_PROJECT_ID = "browserbase-project-id" +STAGEHAND_SERVER_URL = "api_url" diff --git a/MANIFEST.in b/MANIFEST.in deleted file mode 100644 index 7240e718..00000000 --- a/MANIFEST.in +++ /dev/null @@ -1,6 +0,0 @@ -include README.md -include requirements.txt -global-exclude *.pyc -global-exclude __pycache__ -global-exclude .DS_Store -global-exclude */node_modules/* \ No newline at end of file diff --git a/evals/act/google_jobs.py b/evals/act/google_jobs.py index eda1be00..5b78122d 100644 --- a/evals/act/google_jobs.py +++ b/evals/act/google_jobs.py @@ -1,7 +1,9 @@ import asyncio import traceback -from typing import Optional, Any, Dict +from typing import Any, Dict, Optional + from pydantic import BaseModel + from evals.init_stagehand import init_stagehand from stagehand.schemas import ActOptions, ExtractOptions @@ -49,12 +51,12 @@ async def google_jobs(model_name: str, logger, use_text_extract: bool) -> dict: - Clicking on the search button - Clicking on the first job link 4. Extracting job posting details using an AI-driven extraction schema. - + The extraction schema requires: - applicationDeadline: The opening date until which applications are accepted. - minimumQualifications: An object with degree and yearsOfExperience. - preferredQualifications: An object with degree and yearsOfExperience. - + Returns a dictionary containing: - _success (bool): Whether valid job details were extracted. - jobDetails (dict): The extracted job details. @@ -77,7 +79,7 @@ async def google_jobs(model_name: str, logger, use_text_extract: bool) -> dict: try: await stagehand.page.navigate("https://www.google.com/") - await asyncio.sleep(3) + await asyncio.sleep(3) await stagehand.page.act(ActOptions(action="click on the about page")) await stagehand.page.act(ActOptions(action="click on the careers page")) await stagehand.page.act(ActOptions(action="input data scientist into role")) @@ -85,15 +87,17 @@ async def google_jobs(model_name: str, logger, use_text_extract: bool) -> dict: await stagehand.page.act(ActOptions(action="click on the search button")) await stagehand.page.act(ActOptions(action="click on the first job link")) - job_details = await stagehand.page.extract(ExtractOptions( - instruction=( - "Extract the following details from the job posting: application deadline, " - "minimum qualifications (degree and years of experience), and preferred qualifications " - "(degree and years of experience)" - ), - schemaDefinition=JobDetails.model_json_schema(), - useTextExtract=use_text_extract - )) + job_details = await stagehand.page.extract( + ExtractOptions( + instruction=( + "Extract the following details from the job posting: application deadline, " + "minimum qualifications (degree and years of experience), and preferred qualifications " + "(degree and years of experience)" + ), + schemaDefinition=JobDetails.model_json_schema(), + useTextExtract=use_text_extract, + ) + ) valid = is_job_details_valid(job_details) @@ -104,19 +108,21 @@ async def google_jobs(model_name: str, logger, use_text_extract: bool) -> dict: "jobDetails": job_details, "debugUrl": debug_url, "sessionUrl": session_url, - "logs": logger.get_logs() if hasattr(logger, "get_logs") else [] + "logs": logger.get_logs() if hasattr(logger, "get_logs") else [], } except Exception as e: err_message = str(e) err_trace = traceback.format_exc() - logger.error({ - "message": "error in google_jobs function", - "level": 0, - "auxiliary": { - "error": {"value": err_message, "type": "string"}, - "trace": {"value": err_trace, "type": "string"} + logger.error( + { + "message": "error in google_jobs function", + "level": 0, + "auxiliary": { + "error": {"value": err_message, "type": "string"}, + "trace": {"value": err_trace, "type": "string"}, + }, } - }) + ) await stagehand.close() @@ -125,31 +131,37 @@ async def google_jobs(model_name: str, logger, use_text_extract: bool) -> dict: "debugUrl": debug_url, "sessionUrl": session_url, "error": {"message": err_message, "trace": err_trace}, - "logs": logger.get_logs() if hasattr(logger, "get_logs") else [] - } - + "logs": logger.get_logs() if hasattr(logger, "get_logs") else [], + } + + # For quick local testing if __name__ == "__main__": - import os import asyncio import logging + logging.basicConfig(level=logging.INFO) - + class SimpleLogger: def __init__(self): self._logs = [] + def info(self, message): self._logs.append(message) print("INFO:", message) + def error(self, message): self._logs.append(message) print("ERROR:", message) + def get_logs(self): return self._logs async def main(): logger = SimpleLogger() - result = await google_jobs("gpt-4o-mini", logger, use_text_extract=False) # TODO - use text extract + result = await google_jobs( + "gpt-4o-mini", logger, use_text_extract=False + ) # TODO - use text extract print("Result:", result) - - asyncio.run(main()) \ No newline at end of file + + asyncio.run(main()) diff --git a/evals/extract/extract_press_releases.py b/evals/extract/extract_press_releases.py index 7c504052..800fed5f 100644 --- a/evals/extract/extract_press_releases.py +++ b/evals/extract/extract_press_releases.py @@ -1,26 +1,31 @@ import asyncio + from pydantic import BaseModel -from stagehand.schemas import ExtractOptions + from evals.init_stagehand import init_stagehand from evals.utils import compare_strings +from stagehand.schemas import ExtractOptions + # Define Pydantic models for validating press release data class PressRelease(BaseModel): title: str publish_date: str + class PressReleases(BaseModel): items: list[PressRelease] + async def extract_press_releases(model_name: str, logger, use_text_extract: bool): """ Extract press releases from the dummy press releases page using the Stagehand client. - + Args: model_name (str): Name of the AI model to use. logger: A custom logger that provides .error() and .get_logs() methods. use_text_extract (bool): Flag to control text extraction behavior. - + Returns: dict: A result object containing: - _success (bool): Whether the eval was successful. @@ -34,12 +39,16 @@ async def extract_press_releases(model_name: str, logger, use_text_extract: bool session_url = None try: # Initialize Stagehand (mimicking the TS initStagehand) - stagehand, init_response = await init_stagehand(model_name, logger, dom_settle_timeout_ms=3000) + stagehand, init_response = await init_stagehand( + model_name, logger, dom_settle_timeout_ms=3000 + ) debug_url = init_response["debugUrl"] session_url = init_response["sessionUrl"] # Navigate to the dummy press releases page # TODO - choose a different page - await stagehand.page.navigate("https://dummy-press-releases.surge.sh/news", wait_until="networkidle") + await stagehand.page.navigate( + "https://dummy-press-releases.surge.sh/news", wait_until="networkidle" + ) # Wait for 5 seconds to ensure content has loaded await asyncio.sleep(5) @@ -49,7 +58,7 @@ async def extract_press_releases(model_name: str, logger, use_text_extract: bool ExtractOptions( instruction="extract the title and corresponding publish date of EACH AND EVERY press releases on this page. DO NOT MISS ANY PRESS RELEASES.", schemaDefinition=PressReleases.model_json_schema(), - useTextExtract=use_text_extract + useTextExtract=use_text_extract, ) ) print("Raw result:", raw_result) @@ -73,19 +82,21 @@ async def extract_press_releases(model_name: str, logger, use_text_extract: bool expected_length = 28 expected_first = PressRelease( title="UAW Region 9A Endorses Brad Lander for Mayor", - publish_date="Dec 4, 2024" + publish_date="Dec 4, 2024", ) expected_last = PressRelease( title="Fox Sued by New York City Pension Funds Over Election Falsehoods", - publish_date="Nov 12, 2023" + publish_date="Nov 12, 2023", ) if len(items) <= expected_length: - logger.error({ - "message": "Not enough items extracted", - "expected": f"> {expected_length}", - "actual": len(items) - }) + logger.error( + { + "message": "Not enough items extracted", + "expected": f"> {expected_length}", + "actual": len(items), + } + ) return { "_success": False, "error": "Not enough items extracted", @@ -111,10 +122,9 @@ def is_item_match(item: PressRelease, expected: PressRelease) -> bool: await stagehand.close() return result except Exception as e: - logger.error({ - "message": "Error in extract_press_releases function", - "error": str(e) - }) + logger.error( + {"message": "Error in extract_press_releases function", "error": str(e)} + ) return { "_success": False, "error": str(e), @@ -127,26 +137,33 @@ def is_item_match(item: PressRelease, expected: PressRelease) -> bool: if stagehand: await stagehand.close() + # For quick local testing. if __name__ == "__main__": import logging + logging.basicConfig(level=logging.INFO) - + class SimpleLogger: def __init__(self): self._logs = [] + def info(self, message): self._logs.append(message) print("INFO:", message) + def error(self, message): self._logs.append(message) print("ERROR:", message) + def get_logs(self): return self._logs async def main(): logger = SimpleLogger() - result = await extract_press_releases("gpt-4o", logger, use_text_extract=False) # TODO - use text extract + result = await extract_press_releases( + "gpt-4o", logger, use_text_extract=False + ) # TODO - use text extract print("Result:", result) - - asyncio.run(main()) \ No newline at end of file + + asyncio.run(main()) diff --git a/evals/init_stagehand.py b/evals/init_stagehand.py index 8ee3bb40..ed93966f 100644 --- a/evals/init_stagehand.py +++ b/evals/init_stagehand.py @@ -1,20 +1,21 @@ import os -import asyncio + from stagehand import Stagehand from stagehand.config import StagehandConfig + async def init_stagehand(model_name: str, logger, dom_settle_timeout_ms: int = 3000): """ Initialize a Stagehand client with the given model name, logger, and DOM settle timeout. - + This function creates a configuration from environment variables, initializes the Stagehand client, and returns a tuple of (stagehand, init_response). The init_response contains debug and session URLs. - + Args: model_name (str): The name of the AI model to use. logger: A logger instance for logging errors and debug messages. dom_settle_timeout_ms (int): Milliseconds to wait for the DOM to settle. - + Returns: tuple: (stagehand, init_response) where init_response is a dict containing: - "debugUrl": A dict with a "value" key for the debug URL. @@ -22,7 +23,11 @@ async def init_stagehand(model_name: str, logger, dom_settle_timeout_ms: int = 3 """ # Build a Stagehand configuration object using environment variables config = StagehandConfig( - env="BROWSERBASE" if os.getenv("BROWSERBASE_API_KEY") and os.getenv("BROWSERBASE_PROJECT_ID") else "LOCAL", + env=( + "BROWSERBASE" + if os.getenv("BROWSERBASE_API_KEY") and os.getenv("BROWSERBASE_PROJECT_ID") + else "LOCAL" + ), api_key=os.getenv("BROWSERBASE_API_KEY"), project_id=os.getenv("BROWSERBASE_PROJECT_ID"), debug_dom=True, @@ -33,7 +38,9 @@ async def init_stagehand(model_name: str, logger, dom_settle_timeout_ms: int = 3 ) # Create a Stagehand client with the configuration; server_url is taken from environment variables. - stagehand = Stagehand(config=config, server_url=os.getenv("STAGEHAND_SERVER_URL"), verbose=2) + stagehand = Stagehand( + config=config, server_url=os.getenv("STAGEHAND_SERVER_URL"), verbose=2 + ) await stagehand.init() # Construct the URL from the session id using the new format. @@ -43,4 +50,4 @@ async def init_stagehand(model_name: str, logger, dom_settle_timeout_ms: int = 3 url = f"wss://connect.browserbase.com?apiKey={api_key}&sessionId={stagehand.session_id}" # Return both URLs as dictionaries with the "value" key. - return stagehand, {"debugUrl": {"value": url}, "sessionUrl": {"value": url}} \ No newline at end of file + return stagehand, {"debugUrl": {"value": url}, "sessionUrl": {"value": url}} diff --git a/evals/observe/observe_yc_startup.py b/evals/observe/observe_yc_startup.py index 88e41cb6..963d643c 100644 --- a/evals/observe/observe_yc_startup.py +++ b/evals/observe/observe_yc_startup.py @@ -1,16 +1,18 @@ import asyncio + from evals.init_stagehand import init_stagehand from stagehand.schemas import ObserveOptions + async def observe_yc_startup(model_name: str, logger) -> dict: """ This function evaluates the YC startups page by: - + 1. Initializing Stagehand with the provided model name and logger. 2. Navigating to "https://www.ycombinator.com/companies" and waiting for the page to reach network idle. 3. Invoking the observe command to locate the container element housing startup information. 4. Checking against candidate locators to determine if a matching element is found. - + Returns a dictionary containing: - _success (bool): True if a matching container element is found. - matchedLocator (Optional[str]): The candidate locator string that matched. @@ -31,16 +33,18 @@ async def observe_yc_startup(model_name: str, logger) -> dict: if isinstance(init_response.get("sessionUrl"), dict) else init_response.get("sessionUrl") ) - + # Navigate to the YC companies page and wait until network idle await stagehand.page.goto("https://www.ycombinator.com/companies") await stagehand.page.wait_for_load_state("networkidle") - + # Use the observe command with the appropriate instruction - observations = await stagehand.page.observe(ObserveOptions( - instruction="Find the container element that holds links to each of the startup companies. The companies each have a name, a description, and a link to their website." - )) - + observations = await stagehand.page.observe( + ObserveOptions( + instruction="Find the container element that holds links to each of the startup companies. The companies each have a name, a description, and a link to their website." + ) + ) + # If no observations were returned, mark eval as unsuccessful and return early. if not observations: await stagehand.close() @@ -49,22 +53,22 @@ async def observe_yc_startup(model_name: str, logger) -> dict: "observations": observations, "debugUrl": debug_url, "sessionUrl": session_url, - "logs": logger.get_logs() if hasattr(logger, "get_logs") else [] + "logs": logger.get_logs() if hasattr(logger, "get_logs") else [], } - + # Define candidate locators for the container element. possible_locators = [ "div._section_1pgsr_163._results_1pgsr_343", "div._rightCol_1pgsr_592", ] - + possible_handles = [] for locator_str in possible_locators: locator = stagehand.page.locator(locator_str) handle = await locator.element_handle() if handle: possible_handles.append((locator_str, handle)) - + # Iterate over each observation to determine if it matches any of the candidate locators. found_match = False matched_locator = None @@ -89,12 +93,14 @@ async def observe_yc_startup(model_name: str, logger) -> dict: if found_match: break except Exception as e: - print(f"Warning: Failed to check observation with selector {observation.get('selector')}: {str(e)}") + print( + f"Warning: Failed to check observation with selector {observation.get('selector')}: {str(e)}" + ) continue # Cleanup and close the Stagehand client. await stagehand.close() - + # Return the evaluation results. return { "_success": found_match, @@ -102,25 +108,29 @@ async def observe_yc_startup(model_name: str, logger) -> dict: "observations": observations, "debugUrl": debug_url, "sessionUrl": session_url, - "logs": logger.get_logs() if hasattr(logger, "get_logs") else [] + "logs": logger.get_logs() if hasattr(logger, "get_logs") else [], } - + + # For quick local testing if __name__ == "__main__": - import os import asyncio import logging + logging.basicConfig(level=logging.INFO) - + class SimpleLogger: def __init__(self): self._logs = [] + def info(self, message): self._logs.append(message) print("INFO:", message) + def error(self, message): self._logs.append(message) print("ERROR:", message) + def get_logs(self): return self._logs @@ -128,5 +138,5 @@ async def main(): logger = SimpleLogger() result = await observe_yc_startup("gpt-4o-mini", logger) print("Result:", result) - - asyncio.run(main()) \ No newline at end of file + + asyncio.run(main()) diff --git a/evals/run_all_evals.py b/evals/run_all_evals.py index 05a671b1..93f6aee7 100644 --- a/evals/run_all_evals.py +++ b/evals/run_all_evals.py @@ -1,28 +1,33 @@ import asyncio -import os import importlib import inspect +import os + # A simple logger to collect logs for the evals class SimpleLogger: def __init__(self): self._logs = [] + def info(self, message): self._logs.append(message) print("INFO:", message) + def error(self, message): self._logs.append(message) print("ERROR:", message) + def get_logs(self): return self._logs + async def run_all_evals(): eval_functions = {} # The base path is the directory in which this file resides (i.e. the evals folder) base_path = os.path.dirname(__file__) # Only process evals from these sub repositories allowed_dirs = {"act", "extract", "observe"} - + # Recursively walk through the evals directory and its children for root, _, files in os.walk(base_path): # Determine the relative path from the base @@ -39,7 +44,7 @@ async def run_all_evals(): # Skip __init__.py and the runner itself if file.endswith(".py") and file not in ("__init__.py", "run_all_evals.py"): # Build module import path relative to the package root (assumes folder "evals") - if rel_path == '.': + if rel_path == ".": module_path = f"evals.{file[:-3]}" else: # Replace OS-specific path separators with dots ('.') @@ -75,8 +80,9 @@ async def run_all_evals(): return results + if __name__ == "__main__": final_results = asyncio.run(run_all_evals()) print("Evaluation Results:") for module, res in final_results.items(): - print(f"{module}: {res}") \ No newline at end of file + print(f"{module}: {res}") diff --git a/evals/utils.py b/evals/utils.py index a2854fb1..54bc1afe 100644 --- a/evals/utils.py +++ b/evals/utils.py @@ -1,8 +1,9 @@ import difflib + def compare_strings(a: str, b: str) -> float: """ Compare two strings and return a similarity ratio. This function uses difflib.SequenceMatcher to calculate the similarity between two strings. """ - return difflib.SequenceMatcher(None, a, b).ratio() \ No newline at end of file + return difflib.SequenceMatcher(None, a, b).ratio() diff --git a/examples/example.py b/examples/example.py index c3752f32..0cb9f051 100644 --- a/examples/example.py +++ b/examples/example.py @@ -1,71 +1,115 @@ import asyncio -import os import logging +import os + from dotenv import load_dotenv +from rich.console import Console +from rich.panel import Panel +from rich.theme import Theme + from stagehand.client import Stagehand from stagehand.config import StagehandConfig -from stagehand.schemas import ActOptions, ObserveOptions + +# Create a custom theme for consistent styling +custom_theme = Theme( + { + "info": "cyan", + "success": "green", + "warning": "yellow", + "error": "red bold", + "highlight": "magenta", + "url": "blue underline", + } +) + +# Create a Rich console instance with our theme +console = Console(theme=custom_theme) load_dotenv() -# Configure logging at the start of the script +# Configure logging with Rich handler logging.basicConfig( - level=logging.INFO, - format='%(asctime)s - %(name)s - %(levelname)s - %(message)s', - datefmt='%Y-%m-%d %H:%M:%S' + level=logging.WARNING, # Feel free to change this to INFO or DEBUG to see more logs + format="%(asctime)s - %(name)s - %(levelname)s - %(message)s", + datefmt="%Y-%m-%d %H:%M:%S", ) + async def main(): - try: - # Build a unified configuration object for Stagehand - config = StagehandConfig( - env="BROWSERBASE" if os.getenv("BROWSERBASE_API_KEY") and os.getenv("BROWSERBASE_PROJECT_ID") else "LOCAL", - api_key=os.getenv("BROWSERBASE_API_KEY"), - project_id=os.getenv("BROWSERBASE_PROJECT_ID"), - debug_dom=True, - headless=False, - dom_settle_timeout_ms=3000, - model_name="gpt-4o-mini", - model_client_options={"apiKey": os.getenv("MODEL_API_KEY")} - ) - - # Create a Stagehand client using the configuration object. - stagehand = Stagehand(config=config, server_url=os.getenv("STAGEHAND_SERVER_URL"), verbose=2) - - # Initialize - this creates a new session automatically. - await stagehand.init() - print(f"Created new session with ID: {stagehand.session_id}") - - print('EXAMPLE: You can navigate to any website using the local or remote Playwright.') - - await stagehand.page.goto("https://news.ycombinator.com/") - print("Navigation complete with local Playwright.") - - await stagehand.page.navigate("https://www.google.com") - print("Navigation complete with remote Playwright.") - - print("EXAMPLE: Clicking on About link using local Playwright's get_by_role") - # Click on the "About" link using Playwright - await stagehand.page.get_by_role("link", name="About", exact=True).click() - print("Clicked on About link") - - await asyncio.sleep(2) - await stagehand.page.navigate("https://www.google.com") - - # Hosted Stagehand API - ACT to do something like 'search for openai' - await stagehand.page.act(ActOptions(action="search for openai")) - - print("EXAMPLE: Find the XPATH of the button 'News' using Stagehand API") - xpaths = await stagehand.page.observe(ObserveOptions(instruction="find the button labeled 'News'", only_visible=True)) - if len(xpaths) > 0: - element = xpaths[0] - print("EXAMPLE: Click on the button 'News' using local Playwright.") - await stagehand.page.click(element["selector"]) - else: - print("No element found") - - except Exception as e: - print(f"An error occurred in the example: {e}") + # Build a unified configuration object for Stagehand + config = StagehandConfig( + env="BROWSERBASE", + api_key=os.getenv("BROWSERBASE_API_KEY"), + project_id=os.getenv("BROWSERBASE_PROJECT_ID"), + headless=False, + dom_settle_timeout_ms=3000, + model_name="gpt-4o", + model_client_options={"apiKey": os.getenv("MODEL_API_KEY")}, + ) + + # Create a Stagehand client using the configuration object. + stagehand = Stagehand( + config=config, server_url=os.getenv("STAGEHAND_SERVER_URL"), verbose=2 + ) + + # Initialize - this creates a new session automatically. + console.print("\nšŸš€ [info]Initializing Stagehand...[/]") + await stagehand.init() + page = stagehand.page + console.print(f"\n[yellow]Created new session:[/] {stagehand.session_id}") + console.print( + f"šŸŒ [white]View your live browser:[/] [url]https://www.browserbase.com/sessions/{stagehand.session_id}[/]" + ) + + await asyncio.sleep(2) + + console.print("\nā–¶ļø [highlight] Navigating[/] to Google") + await page.goto("https://google.com/") + console.print("āœ… [success]Navigated to Google[/]") + + console.print("\nā–¶ļø [highlight] Clicking[/] on About link") + # Click on the "About" link using Playwright + await page.get_by_role("link", name="About", exact=True).click() + console.print("āœ… [success]Clicked on About link[/]") + + await asyncio.sleep(2) + console.print("\nā–¶ļø [highlight] Navigating[/] back to Google") + await page.goto("https://google.com/") + console.print("āœ… [success]Navigated back to Google[/]") + + console.print("\nā–¶ļø [highlight] Performing action:[/] search for openai") + await page.act("search for openai") + console.print("āœ… [success]Performing Action:[/] Action completed successfully") + + console.print("\nā–¶ļø [highlight] Observing page[/] for news button") + observed = await page.observe("find the news button on the page") + if len(observed) > 0: + element = observed[0] + console.print("āœ… [success]Found element:[/] News button") + await page.act(element) + else: + console.print("āŒ [error]No element found[/]") + + console.print("\nā–¶ļø [highlight] Extracting[/] first search result") + data = await page.extract("extract the first result from the search") + console.print("šŸ“Š [info]Extracted data:[/]") + console.print_json(f"{data.model_dump_json()}") + + # Close the session + console.print("\nā¹ļø [warning]Closing session...[/]") + await stagehand.close() + console.print("āœ… [success]Session closed successfully![/]") + console.rule("[bold]End of Example[/]") + if __name__ == "__main__": - asyncio.run(main()) \ No newline at end of file + # Add a fancy header + console.print( + "\n", + Panel.fit( + "[light_gray]Stagehand šŸ¤˜ Python Example[/]", + border_style="green", + padding=(1, 10), + ), + ) + asyncio.run(main()) diff --git a/examples/extract-example.py b/examples/extract-example.py deleted file mode 100644 index d03b3132..00000000 --- a/examples/extract-example.py +++ /dev/null @@ -1,51 +0,0 @@ -import asyncio -import os -from dotenv import load_dotenv -from stagehand import Stagehand -from stagehand.config import StagehandConfig -from stagehand.schemas import ExtractOptions -from pydantic import BaseModel - -class ExtractSchema(BaseModel): - stars: int - -# Load environment variables from .env file -load_dotenv() - -async def main(): - # Build a unified Stagehand configuration object - config = StagehandConfig( - env="BROWSERBASE" if os.getenv("BROWSERBASE_API_KEY") and os.getenv("BROWSERBASE_PROJECT_ID") else "LOCAL", - api_key=os.getenv("BROWSERBASE_API_KEY"), - project_id=os.getenv("BROWSERBASE_PROJECT_ID"), - debug_dom=True, - headless=True, - model_name="gpt-4o", - model_client_options={"apiKey": os.getenv("MODEL_API_KEY")} - ) - - # Create a Stagehand client using the configuration object. - stagehand = Stagehand(config=config, server_url=os.getenv("STAGEHAND_SERVER_URL"), verbose=2) - - # Initialize - this creates a new session. - await stagehand.init() - print(f"Created new session with ID: {stagehand.session_id}") - - try: - await stagehand.page.navigate("https://github.com/facebook/react") - print("Navigation complete.") - - # Use the ExtractOptions Pydantic model to pass instruction and schema definition - data = await stagehand.page.extract( - ExtractOptions( - instruction="Extract the number of stars for the project", - schemaDefinition=ExtractSchema.model_json_schema() - ) - ) - print("\nExtracted stars:", data) - - except Exception as e: - print(f"Error: {e}") - -if __name__ == "__main__": - asyncio.run(main()) \ No newline at end of file diff --git a/examples/observe-example.py b/examples/observe-example.py deleted file mode 100644 index e9b7d281..00000000 --- a/examples/observe-example.py +++ /dev/null @@ -1,55 +0,0 @@ -import asyncio -import os -from dotenv import load_dotenv -from stagehand import Stagehand -from stagehand.config import StagehandConfig -from stagehand.schemas import ObserveOptions - -# Load environment variables from .env file -load_dotenv() - -async def main(): - # Build a unified Stagehand configuration object - config = StagehandConfig( - env="BROWSERBASE" if os.getenv("BROWSERBASE_API_KEY") and os.getenv("BROWSERBASE_PROJECT_ID") else "LOCAL", - api_key=os.getenv("BROWSERBASE_API_KEY"), - project_id=os.getenv("BROWSERBASE_PROJECT_ID"), - debug_dom=True, - headless=True, - model_name="gpt-4o-mini", - model_client_options={"apiKey": os.getenv("MODEL_API_KEY")} - ) - - # Create a Stagehand client using the configuration object. - stagehand = Stagehand(config=config, server_url=os.getenv("SERVER_URL"), verbose=2) - - # Initialize - this creates a new session. - await stagehand.init() - print(f"Created new session with ID: {stagehand.session_id}") - - try: - # Navigate to the desired page - await stagehand.page.navigate("https://elpasotexas.ionwave.net/Login.aspx") - print("Navigation complete.") - - # Use ObserveOptions for detailed instructions - options = ObserveOptions( - instruction="find all the links on the page regarding the city of el paso", - only_visible=True - ) - activity = await stagehand.page.observe(options) - print("\nObservations:", activity) - print("Length of observations:", len(activity)) - - print("Click on the first extracted element") - if activity: - print(activity[0]) - await stagehand.page.click(activity[0]["selector"]) - else: - print("No elements found") - - except Exception as e: - print(f"Error: {e}") - -if __name__ == "__main__": - asyncio.run(main()) \ No newline at end of file diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 00000000..4263a1f4 --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,49 @@ +[tool.ruff] +# Enable flake8-comprehensions, flake8-bugbear, naming, etc. +select = ["E", "F", "B", "C4", "UP", "N", "I", "C"] +ignore = [] + +# Same as Black +line-length = 88 + +# Target Python version +target-version = "py39" # Adjust to your version + +# Exclude a variety of commonly ignored directories +exclude = [ + ".git", + ".ruff_cache", + "__pycache__", + "venv", + ".venv", + "dist", + "tests" +] + +# Allow autofix for all enabled rules (when `--fix` is provided) +fixable = ["ALL"] +unfixable = [] + +# Allow unused variables when underscore-prefixed +dummy-variable-rgx = "^(_+|(_+[a-zA-Z0-9_]*[a-zA-Z0-9]+?))$" + +[tool.ruff.format] +quote-style = "double" +indent-style = "space" +line-ending = "auto" + +# Naming conventions are part of the N rules +[tool.ruff.lint.pep8-naming] +# Allow underscores in class names (e.g. Test_Case) +classmethod-decorators = ["classmethod", "validator"] + +# Per-file ignores +[tool.ruff.lint.per-file-ignores] +# Ignore imported but unused in __init__.py files +"__init__.py" = ["F401"] +# Ignore unused imports in tests +"tests/*" = ["F401", "F811"] + +# Add more customizations if needed +[tool.ruff.lint.pydocstyle] +convention = "google" \ No newline at end of file diff --git a/pytest.ini b/pytest.ini new file mode 100644 index 00000000..bca37cd6 --- /dev/null +++ b/pytest.ini @@ -0,0 +1,13 @@ +[pytest] +testpaths = tests +python_files = test_*.py +python_classes = Test* +python_functions = test_* +asyncio_mode = auto + +markers = + unit: marks tests as unit tests + integration: marks tests as integration tests + +log_cli = true +log_cli_level = INFO \ No newline at end of file diff --git a/requirements-dev.txt b/requirements-dev.txt new file mode 100644 index 00000000..08d8882e --- /dev/null +++ b/requirements-dev.txt @@ -0,0 +1,7 @@ +pytest>=7.3.1 +pytest-asyncio>=0.21.0 +pytest-mock>=3.10.0 +pytest-cov>=4.1.0 +black>=23.3.0 +isort>=5.12.0 +mypy>=1.3.0 \ No newline at end of file diff --git a/requirements.txt b/requirements.txt index 22238fa0..641e9eeb 100644 --- a/requirements.txt +++ b/requirements.txt @@ -2,4 +2,5 @@ httpx>=0.24.0 asyncio>=3.4.3 python-dotenv>=1.0.0 pydantic>=1.10.0 -playwright>=1.42.1 \ No newline at end of file +playwright>=1.42.1 +rich \ No newline at end of file diff --git a/run_tests.sh b/run_tests.sh new file mode 100755 index 00000000..94057b82 --- /dev/null +++ b/run_tests.sh @@ -0,0 +1,40 @@ +#!/bin/bash +# Run tests with coverage reporting + +# Make sure we're in the right directory +cd "$(dirname "$0")" + +# Install dev requirements if needed +if [[ -z $(pip3 list | grep pytest) ]]; then + echo "Installing development requirements..." + pip3 install -r requirements-dev.txt +fi + +# Install package in development mode if needed +if [[ -z $(pip3 list | grep stagehand) ]]; then + echo "Installing stagehand package in development mode..." + pip3 install -e . +fi + +# Run the tests +echo "Running tests with coverage..." +python3 -m pytest tests/ -v --cov=stagehand --cov-report=term --cov-report=html + +echo "Tests complete. HTML coverage report is in htmlcov/ directory." + +# Check if we should open the report +if [[ "$1" == "--open" || "$1" == "-o" ]]; then + echo "Opening HTML coverage report..." + if [[ "$OSTYPE" == "darwin"* ]]; then + # macOS + open htmlcov/index.html + elif [[ "$OSTYPE" == "linux-gnu"* ]]; then + # Linux with xdg-open + xdg-open htmlcov/index.html + elif [[ "$OSTYPE" == "msys" || "$OSTYPE" == "win32" ]]; then + # Windows + start htmlcov/index.html + else + echo "Couldn't automatically open the report. Please open htmlcov/index.html manually." + fi +fi \ No newline at end of file diff --git a/setup.py b/setup.py index ee6e2d1d..adf33833 100644 --- a/setup.py +++ b/setup.py @@ -1,11 +1,11 @@ -from setuptools import setup, find_packages +from setuptools import find_packages, setup -with open("README.md", "r", encoding="utf-8") as fh: +with open("README.md", encoding="utf-8") as fh: long_description = fh.read() setup( name="stagehand-py", - version="0.2.1", + version="0.3.0", author="Browserbase, Inc.", author_email="support@browserbase.io", description="Python SDK for Stagehand", @@ -26,4 +26,4 @@ "pydantic>=1.10.0", "playwright>=1.40.0", ], -) \ No newline at end of file +) diff --git a/stagehand/__init__.py b/stagehand/__init__.py index 3ae8cbfb..63b03284 100644 --- a/stagehand/__init__.py +++ b/stagehand/__init__.py @@ -1,4 +1,4 @@ from .client import Stagehand __version__ = "0.1.0" -__all__ = ["Stagehand"] \ No newline at end of file +__all__ = ["Stagehand"] diff --git a/stagehand/client.py b/stagehand/client.py index 85762a5c..520f7e46 100644 --- a/stagehand/client.py +++ b/stagehand/client.py @@ -1,24 +1,28 @@ import asyncio import json +import logging +import os import time +from collections.abc import Awaitable +from typing import Any, Callable, Dict, Optional + import httpx -import os -import logging -from typing import Optional, Dict, Any, Callable, Awaitable from dotenv import load_dotenv from playwright.async_api import async_playwright + +from .config import StagehandConfig from .page import StagehandPage from .utils import default_log_handler -from .config import StagehandConfig load_dotenv() logger = logging.getLogger(__name__) + class Stagehand: """ Python client for interacting with a running Stagehand server and Browserbase remote headless browser. - + Now supports automatically creating a new session if no session_id is provided. You can also optionally provide a configuration via the 'config' parameter to centralize all parameters. """ @@ -34,7 +38,9 @@ def __init__( browserbase_api_key: Optional[str] = None, browserbase_project_id: Optional[str] = None, model_api_key: Optional[str] = None, - on_log: Optional[Callable[[Dict[str, Any]], Awaitable[None]]] = default_log_handler, + on_log: Optional[ + Callable[[Dict[str, Any]], Awaitable[None]] + ] = default_log_handler, verbose: int = 1, model_name: Optional[str] = None, dom_settle_timeout_ms: Optional[int] = None, @@ -63,22 +69,36 @@ def __init__( self.server_url = server_url or os.getenv("STAGEHAND_SERVER_URL") if config: - self.browserbase_api_key = config.api_key or browserbase_api_key or os.getenv("BROWSERBASE_API_KEY") - self.browserbase_project_id = config.project_id or browserbase_project_id or os.getenv("BROWSERBASE_PROJECT_ID") - self.model_api_key = model_api_key or ( - config.model_client_options.get("apiKey") if config.model_client_options else None - ) or os.getenv("MODEL_API_KEY") + self.browserbase_api_key = ( + config.api_key + or browserbase_api_key + or os.getenv("BROWSERBASE_API_KEY") + ) + self.browserbase_project_id = ( + config.project_id + or browserbase_project_id + or os.getenv("BROWSERBASE_PROJECT_ID") + ) + self.model_api_key = os.getenv("MODEL_API_KEY") self.session_id = config.browserbase_session_id or session_id self.model_name = config.model_name or model_name - self.dom_settle_timeout_ms = config.dom_settle_timeout_ms or dom_settle_timeout_ms - self.debug_dom = config.debug_dom if config.debug_dom is not None else debug_dom + self.dom_settle_timeout_ms = ( + config.dom_settle_timeout_ms or dom_settle_timeout_ms + ) + self.debug_dom = ( + config.debug_dom if config.debug_dom is not None else debug_dom + ) self._custom_logger = config.logger # For future integration if needed # Additional config parameters available for future use: self.headless = config.headless self.enable_caching = config.enable_caching else: - self.browserbase_api_key = browserbase_api_key or os.getenv("BROWSERBASE_API_KEY") - self.browserbase_project_id = browserbase_project_id or os.getenv("BROWSERBASE_PROJECT_ID") + self.browserbase_api_key = browserbase_api_key or os.getenv( + "BROWSERBASE_API_KEY" + ) + self.browserbase_project_id = browserbase_project_id or os.getenv( + "BROWSERBASE_PROJECT_ID" + ) self.model_api_key = model_api_key or os.getenv("MODEL_API_KEY") self.session_id = session_id self.model_name = model_name @@ -104,32 +124,36 @@ def __init__( self.page: Optional[StagehandPage] = None self._initialized = False # Flag to track if init() has run - self._closed = False # Flag to track if resources have been closed + self._closed = False # Flag to track if resources have been closed # Validate essential fields if session_id was provided if self.session_id: if not self.browserbase_api_key: - raise ValueError("browserbase_api_key is required (or set BROWSERBASE_API_KEY in env).") + raise ValueError( + "browserbase_api_key is required (or set BROWSERBASE_API_KEY in env)." + ) if not self.browserbase_project_id: - raise ValueError("browserbase_project_id is required (or set BROWSERBASE_PROJECT_ID in env).") - + raise ValueError( + "browserbase_project_id is required (or set BROWSERBASE_PROJECT_ID in env)." + ) + def _get_lock_for_session(self) -> asyncio.Lock: """ Return an asyncio.Lock for this session. If one doesn't exist yet, create it. """ if self.session_id not in self._session_locks: self._session_locks[self.session_id] = asyncio.Lock() - print(f"Created lock for session {self.session_id}") + self._log(f"Created lock for session {self.session_id}", level=3) return self._session_locks[self.session_id] async def __aenter__(self): - self._log("Entering Stagehand context manager (__aenter__)...", level=1) + self._log("Entering Stagehand context manager (__aenter__)...", level=3) # Just call init() if not already done await self.init() return self async def __aexit__(self, exc_type, exc_val, exc_tb): - self._log("Exiting Stagehand context manager (__aexit__)...", level=1) + self._log("Exiting Stagehand context manager (__aexit__)...", level=3) await self.close() async def init(self): @@ -138,13 +162,15 @@ async def init(self): Creates or resumes the session, starts Playwright, and sets up self.page. """ if self._initialized: - self._log("Stagehand is already initialized; skipping init()", level=2) + self._log("Stagehand is already initialized; skipping init()", level=3) return - self._log("Initializing Stagehand...", level=1) + self._log("Initializing Stagehand...", level=3) if not self._client: - self._client = self.httpx_client or httpx.AsyncClient(timeout=self.timeout_settings) + self._client = self.httpx_client or httpx.AsyncClient( + timeout=self.timeout_settings + ) # Check server health await self._check_server_health() @@ -152,41 +178,41 @@ async def init(self): # Create session if we don't have one if not self.session_id: await self._create_session() - self._log(f"Created new session: {self.session_id}", level=1) + self._log(f"Created new session: {self.session_id}", level=3) # Start Playwright and connect to remote - self._log("Starting Playwright...", level=1) + self._log("Starting Playwright...", level=3) self._playwright = await async_playwright().start() connect_url = ( f"wss://connect.browserbase.com?apiKey={self.browserbase_api_key}" f"&sessionId={self.session_id}" ) - self._log(f"Connecting to remote browser at: {connect_url}", level=1) + self._log(f"Connecting to remote browser at: {connect_url}", level=3) self._browser = await self._playwright.chromium.connect_over_cdp(connect_url) - self._log(f"Connected to remote browser: {self._browser}", level=1) - + self._log(f"Connected to remote browser: {self._browser}", level=3) + # Access or create a context existing_contexts = self._browser.contexts - self._log(f"Existing contexts: {len(existing_contexts)}", level=1) + self._log(f"Existing contexts: {len(existing_contexts)}", level=3) if existing_contexts: self._context = existing_contexts[0] else: - self._log("Creating a new context...", level=1) + self._log("Creating a new context...", level=3) self._context = await self._browser.new_context() # Access or create a page existing_pages = self._context.pages - self._log(f"Existing pages: {len(existing_pages)}", level=1) + self._log(f"Existing pages: {len(existing_pages)}", level=3) if existing_pages: - self._log("Using existing page", level=1) + self._log("Using existing page", level=3) self._playwright_page = existing_pages[0] else: - self._log("Creating a new page...", level=1) + self._log("Creating a new page...", level=3) self._playwright_page = await self._context.new_page() # Wrap with StagehandPage - self._log("Wrapping Playwright page in StagehandPage", level=1) + self._log("Wrapping Playwright page in StagehandPage", level=3) self.page = StagehandPage(self._playwright_page, self) self._initialized = True @@ -199,29 +225,34 @@ async def close(self): # Already closed return - self._log("Closing resources...", level=1) - if self._playwright_page: - self._log("Closing the Playwright page...", level=1) - await self._playwright_page.close() - self._playwright_page = None + self._log("Closing resources...", level=3) - if self._context: - self._log("Closing the context...", level=1) - await self._context.close() - self._context = None + # End the session on the server if we have a session ID + if self.session_id: + try: + self._log(f"Ending session {self.session_id} on the server...", level=3) + client = self.httpx_client or httpx.AsyncClient( + timeout=self.timeout_settings + ) + headers = { + "x-bb-api-key": self.browserbase_api_key, + "x-bb-project-id": self.browserbase_project_id, + "Content-Type": "application/json", + } - if self._browser: - self._log("Closing the browser...", level=1) - await self._browser.close() - self._browser = None + async with client: + await self._execute("end", {"sessionId": self.session_id}) + self._log(f"Session {self.session_id} ended successfully", level=3) + except Exception as e: + self._log(f"Error ending session: {str(e)}", level=3) if self._playwright: - self._log("Stopping Playwright...", level=1) + self._log("Stopping Playwright...", level=3) await self._playwright.stop() self._playwright = None if self._client and not self.httpx_client: - self._log("Closing the internal HTTPX client...", level=1) + self._log("Closing the internal HTTPX client...", level=3) await self._client.aclose() self._client = None @@ -236,24 +267,30 @@ async def _check_server_health(self, timeout: int = 10): attempt = 0 while True: try: - client = self.httpx_client or httpx.AsyncClient(timeout=self.timeout_settings) + client = self.httpx_client or httpx.AsyncClient( + timeout=self.timeout_settings + ) async with client: headers = { "x-bb-api-key": self.browserbase_api_key, } - resp = await client.get(f"{self.server_url}/healthcheck", headers=headers) + resp = await client.get( + f"{self.server_url}/healthcheck", headers=headers + ) if resp.status_code == 200: data = resp.json() if data.get("status") == "ok": - self._log("Healthcheck passed. Server is running.", level=1) + self._log("Healthcheck passed. Server is running.", level=3) return except Exception as e: - self._log(f"Healthcheck error: {str(e)}", level=2) + self._log(f"Healthcheck error: {str(e)}", level=3) if time.time() - start > timeout: raise TimeoutError(f"Server not responding after {timeout} seconds.") - - wait_time = min(2 ** attempt * 0.5, 5.0) # Exponential backoff, capped at 5 seconds + + wait_time = min( + 2**attempt * 0.5, 5.0 + ) # Exponential backoff, capped at 5 seconds await asyncio.sleep(wait_time) attempt += 1 @@ -293,7 +330,7 @@ async def _create_session(self): if resp.status_code != 200: raise RuntimeError(f"Failed to create session: {resp.text}") data = resp.json() - self._log(f"Session created: {data}", level=1) + self._log(f"Session created: {data}", level=3) if not data.get("success") or "sessionId" not in data.get("data", {}): raise RuntimeError(f"Invalid response format: {resp.text}") @@ -309,75 +346,92 @@ async def _execute(self, method: str, payload: Dict[str, Any]) -> Any: "x-bb-project-id": self.browserbase_project_id, "Content-Type": "application/json", "Connection": "keep-alive", - "x-stream-response": str(self.streamed_response).lower() + "x-stream-response": str(self.streamed_response).lower(), } if self.model_api_key: headers["x-model-api-key"] = self.model_api_key client = self.httpx_client or httpx.AsyncClient(timeout=self.timeout_settings) - print(f"Executing {method} with payload: {payload} and headers: {headers}") + self._log(f"\n==== EXECUTING {method.upper()} ====", level=3) + self._log( + f"URL: {self.server_url}/sessions/{self.session_id}/{method}", level=3 + ) + self._log(f"Payload: {payload}", level=3) + self._log(f"Headers: {headers}", level=3) + async with client: - async with client.stream( - "POST", - f"{self.server_url}/sessions/{self.session_id}/{method}", - json=payload, - headers=headers, - ) as response: - if response.status_code != 200: - error_text = await response.aread() - self._log(f"Error: {error_text.decode('utf-8')}", level=2) - return None - - async for line in response.aiter_lines(): - # Skip empty lines - if not line.strip(): - continue - - try: - # Handle SSE-style messages that start with "data: " - if line.startswith("data: "): - line = line[len("data: "):] - - message = json.loads(line) - logger.info(f"Message: {message}") - - # Handle different message types - msg_type = message.get("type") - - if msg_type == "system": - status = message.get("data", {}).get("status") - if status == "finished": - return message.get("data", {}).get("result") - elif msg_type == "log": - # Log message from data.message - log_msg = message.get("data", {}).get("message", "") - self._log(log_msg, level=1) - if self.on_log: - await self.on_log(message) - else: - # Log any other message types - self._log(f"Unknown message type: {msg_type}", level=2) - if self.on_log: - await self.on_log(message) - - except json.JSONDecodeError: - self._log(f"Could not parse line as JSON: {line}", level=2) - continue + try: + async with client.stream( + "POST", + f"{self.server_url}/sessions/{self.session_id}/{method}", + json=payload, + headers=headers, + ) as response: + if response.status_code != 200: + error_text = await response.aread() + error_message = error_text.decode("utf-8") + self._log(f"Error: {error_message}", level=3) + return None + + self._log("Starting to process streaming response...", level=3) + async for line in response.aiter_lines(): + # Skip empty lines + if not line.strip(): + continue + + try: + # Handle SSE-style messages that start with "data: " + if line.startswith("data: "): + line = line[len("data: ") :] + + message = json.loads(line) + # Handle different message types + msg_type = message.get("type") + + if msg_type == "system": + status = message.get("data", {}).get("status") + if status == "finished": + result = message.get("data", {}).get("result") + self._log( + f"FINISHED WITH RESULT: {result}", level=3 + ) + return result + elif msg_type == "log": + # Log message from data.message + log_msg = message.get("data", {}).get("message", "") + self._log(log_msg, level=3) + if self.on_log: + await self.on_log(message) + else: + # Log any other message types + self._log(f"Unknown message type: {msg_type}", level=3) + if self.on_log: + await self.on_log(message) + + except json.JSONDecodeError: + self._log(f"Could not parse line as JSON: {line}", level=3) + continue + except Exception as e: + self._log(f"EXCEPTION IN _EXECUTE: {str(e)}") + raise # If we get here without seeing a "finished" message, something went wrong - raise RuntimeError("Server connection closed without sending 'finished' message") + self._log("==== ERROR: No 'finished' message received ====", level=3) + raise RuntimeError( + "Server connection closed without sending 'finished' message" + ) async def _handle_log(self, msg: Dict[str, Any]): """ Handle a log line from the server. If on_log is set, call it asynchronously. """ if self.verbose >= 1: - self._log(f"Log message: {msg}", level=1) + self._log(f"Log message: {msg}", level=3) if self.on_log: try: await self.on_log(msg) except Exception as e: - self._log(f"on_log callback error: {str(e)}", level=2) + self._log(f"on_log callback error: {str(e)}", level=3) def _log(self, message: str, level: int = 1): """ @@ -387,10 +441,10 @@ def _log(self, message: str, level: int = 1): if self.verbose >= level: timestamp = time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()) formatted_msg = f"{timestamp}::[stagehand] {message}" - + if level == 1: logger.info(formatted_msg) elif level == 2: logger.warning(formatted_msg) else: - logger.debug(formatted_msg) \ No newline at end of file + logger.debug(formatted_msg) diff --git a/stagehand/config.py b/stagehand/config.py index c9a360ae..86ec5281 100644 --- a/stagehand/config.py +++ b/stagehand/config.py @@ -1,12 +1,16 @@ +from typing import Any, Callable, Optional + from pydantic import BaseModel, Field -from typing import Optional, Dict, Callable, Any + +from stagehand.schemas import AvailableModel + class StagehandConfig(BaseModel): """ Configuration for the Stagehand client. Attributes: - env (str): Environment type. Use 'BROWSERBASE' for remote usage or 'LOCAL' otherwise. + env (str): Environment type. 'BROWSERBASE' for remote usage api_key (Optional[str]): API key for authentication. project_id (Optional[str]): Project identifier. debug_dom (bool): Enable DOM debugging features. @@ -16,19 +20,42 @@ class StagehandConfig(BaseModel): enable_caching (Optional[bool]): Enable caching functionality. browserbase_session_id (Optional[str]): Session ID for resuming Browserbase sessions. model_name (Optional[str]): Name of the model to use. - model_client_options (Optional[Dict[str, Any]]): Configuration options for the model client. + selfHeal (Optional[bool]): Enable self-healing functionality. """ - env: str = Field("LOCAL", description="Environment type, e.g., 'BROWSERBASE' for remote or 'LOCAL' for local") - api_key: Optional[str] = Field(None, alias="apiKey", description="API key for authentication") - project_id: Optional[str] = Field(None, alias="projectId", description="Project identifier") - debug_dom: bool = Field(False, alias="debugDom", description="Enable DOM debugging features") + + env: str = "BROWSERBASE" + api_key: Optional[str] = Field( + None, alias="apiKey", description="Browserbase API key for authentication" + ) + project_id: Optional[str] = Field( + None, alias="projectId", description="Browserbase project ID" + ) + debug_dom: bool = Field( + False, alias="debugDom", description="Enable DOM debugging features" + ) headless: bool = Field(True, description="Run browser in headless mode") - logger: Optional[Callable[[Any], None]] = Field(None, description="Custom logging function") - dom_settle_timeout_ms: Optional[int] = Field(3000, alias="domSettleTimeoutMs", description="Timeout for DOM to settle (in ms)") - enable_caching: Optional[bool] = Field(False, alias="enableCaching", description="Enable caching functionality") - browserbase_session_id: Optional[str] = Field(None, alias="browserbaseSessionID", description="Session ID for resuming Browserbase sessions") - model_name: Optional[str] = Field(None, alias="modelName", description="Name of the model to use") - model_client_options: Optional[Dict[str, Any]] = Field(default_factory=dict, alias="modelClientOptions", description="Options for the model client") + logger: Optional[Callable[[Any], None]] = Field( + None, description="Custom logging function" + ) + dom_settle_timeout_ms: Optional[int] = Field( + 3000, + alias="domSettleTimeoutMs", + description="Timeout for DOM to settle (in ms)", + ) + enable_caching: Optional[bool] = Field( + False, alias="enableCaching", description="Enable caching functionality" + ) + browserbase_session_id: Optional[str] = Field( + None, + alias="browserbaseSessionID", + description="Session ID for resuming Browserbase sessions", + ) + model_name: Optional[str] = Field( + AvailableModel.GPT_4O, alias="modelName", description="Name of the model to use" + ) + selfHeal: Optional[bool] = Field( + True, description="Enable self-healing functionality" + ) class Config: - populate_by_name = True \ No newline at end of file + populate_by_name = True diff --git a/stagehand/page.py b/stagehand/page.py index c1f6a728..763e88cd 100644 --- a/stagehand/page.py +++ b/stagehand/page.py @@ -1,56 +1,48 @@ -from typing import Optional, Dict, Any +from typing import List, Optional, Union + from playwright.async_api import Page -from pydantic import BaseModel -# (Make sure to import the new options models when needed) -from .schemas import ActOptions, ObserveOptions, ExtractOptions +from .schemas import ( + ActOptions, + ActResult, + ExtractOptions, + ExtractResult, + ObserveOptions, + ObserveResult, +) + class StagehandPage: """Wrapper around Playwright Page that integrates with Stagehand server""" - + def __init__(self, page: Page, stagehand_client): """ Initialize a StagehandPage instance. - + Args: page (Page): The underlying Playwright page. stagehand_client: The client used to interface with the Stagehand server. """ self.page = page self._stagehand = stagehand_client - - async def goto(self, url: str, **kwargs): - """ - Navigate to the given URL using Playwright directly. - - Args: - url (str): The URL to navigate to. - **kwargs: Additional keyword arguments passed to Playwright's page.goto. - - Returns: - The result of Playwright's page.goto method. - """ - lock = self._stagehand._get_lock_for_session() - async with lock: - return await self.page.goto(url, **kwargs) - async def navigate( - self, - url: str, - *, + async def goto( + self, + url: str, + *, referer: Optional[str] = None, - timeout: Optional[int] = None, + timeout: Optional[int] = None, wait_until: Optional[str] = None ): """ Navigate to URL using the Stagehand server. - + Args: url (str): The URL to navigate to. referer (Optional[str]): Optional referer URL. timeout (Optional[int]): Optional navigation timeout in milliseconds. wait_until (Optional[str]): Optional wait condition; one of ('load', 'domcontentloaded', 'networkidle', 'commit'). - + Returns: The result from the Stagehand server's navigation execution. """ @@ -61,7 +53,7 @@ async def navigate( options["timeout"] = timeout if wait_until is not None: options["waitUntil"] = wait_until - + payload = {"url": url} if options: payload["options"] = options @@ -70,70 +62,95 @@ async def navigate( async with lock: result = await self._stagehand._execute("navigate", payload) return result - - async def act(self, options: ActOptions) -> Any: + + async def act(self, options: Union[str, ActOptions, ObserveResult]) -> ActResult: """ Execute an AI action via the Stagehand server. - + Args: - options (ActOptions): A Pydantic model encapsulating the action. + options (Union[str, ActOptions]): Either a string with the action command or + a Pydantic model encapsulating the action. See `stagehand.schemas.ActOptions` for details on expected fields. - + Returns: Any: The result from the Stagehand server's action execution. """ - payload = options.dict(exclude_none=True) + # Convert string to ActOptions if needed + if isinstance(options, str): + options = ActOptions(action=options) + + payload = options.model_dump(exclude_none=True) lock = self._stagehand._get_lock_for_session() async with lock: result = await self._stagehand._execute("act", payload) + if isinstance(result, dict): + return ActResult(**result) return result - - async def observe(self, options: ObserveOptions) -> Any: + + async def observe(self, options: Union[str, ObserveOptions]) -> List[ObserveResult]: """ Make an AI observation via the Stagehand server. - + Args: - options (ObserveOptions): A Pydantic model encapsulating the observation instruction. + options (Union[str, ObserveOptions]): Either a string with the observation instruction + or a Pydantic model encapsulating the observation instruction. See `stagehand.schemas.ObserveOptions` for details on expected fields. - + Returns: - Any: The result from the Stagehand server's observation execution. + List[ObserveResult]: A list of observation results from the Stagehand server. """ - payload = options.dict(exclude_none=True) + # Convert string to ObserveOptions if needed + if isinstance(options, str): + options = ObserveOptions(instruction=options) + + payload = options.model_dump(exclude_none=True) lock = self._stagehand._get_lock_for_session() async with lock: result = await self._stagehand._execute("observe", payload) - return result - - async def extract(self, options: ExtractOptions) -> Any: + + # Convert raw result to list of ObserveResult models + if isinstance(result, list): + return [ObserveResult(**item) for item in result] + elif isinstance(result, dict): + # If single dict, wrap in list + return [ObserveResult(**result)] + return [] + + async def extract(self, options: Union[str, ExtractOptions]) -> ExtractResult: """ Extract data using AI via the Stagehand server. - + Expects an ExtractOptions Pydantic object that includes a JSON schema (or Pydantic model) for validation. - + Args: options (ExtractOptions): The extraction options describing what to extract and how. See `stagehand.schemas.ExtractOptions` for details on expected fields. - + Returns: Any: The result from the Stagehand server's extraction execution. """ - payload = options.dict(exclude_none=True) + # Convert string to ExtractOptions if needed + if isinstance(options, str): + options = ExtractOptions(instruction=options) + + payload = options.model_dump(exclude_none=True) lock = self._stagehand._get_lock_for_session() async with lock: result = await self._stagehand._execute("extract", payload) + if isinstance(result, dict): + return ExtractResult(**result) return result # Forward other Page methods to underlying Playwright page def __getattr__(self, name): """ Forward attribute lookups to the underlying Playwright page. - + Args: name (str): Name of the attribute to access. - + Returns: The attribute from the underlying Playwright page. """ - return getattr(self.page, name) \ No newline at end of file + return getattr(self.page, name) diff --git a/stagehand/schemas.py b/stagehand/schemas.py index e9d510e6..fc572e40 100644 --- a/stagehand/schemas.py +++ b/stagehand/schemas.py @@ -1,5 +1,22 @@ +from enum import Enum +from typing import Any, Dict, List, Optional, Type, Union + from pydantic import BaseModel, Field -from typing import Optional, Dict, Any, Union, Type + +# Default extraction schema that matches the TypeScript version +DEFAULT_EXTRACT_SCHEMA = { + "type": "object", + "properties": {"extraction": {"type": "string"}}, + "required": ["extraction"], +} + + +class AvailableModel(str, Enum): + GPT_4O = "gpt-4o" + GPT_4O_MINI = "gpt-4o-mini" + CLAUDE_3_5_SONNET_LATEST = "claude-3-5-sonnet-latest" + CLAUDE_3_7_SONNET_LATEST = "claude-3-7-sonnet-latest" + class ActOptions(BaseModel): """ @@ -7,25 +24,31 @@ class ActOptions(BaseModel): Attributes: action (str): The action command to be executed by the AI. - useVision: Optional[Union[bool, str]] = None variables: Optional[Dict[str, str]] = None + modelName: Optional[AvailableModel] = None + slowDomBasedAct: Optional[bool] = None """ + action: str = Field(..., description="The action command to be executed by the AI.") - useVision: Optional[Union[bool, str]] = None variables: Optional[Dict[str, str]] = None + modelName: Optional[AvailableModel] = None + slowDomBasedAct: Optional[bool] = None -class ObserveOptions(BaseModel): + +class ActResult(BaseModel): """ - Options for the 'observe' command. + Result of the 'act' command. Attributes: - instruction (str): Instruction detailing what the AI should observe. - useVision: Optional[bool] = None - onlyVisible: Optional[bool] = None + success (bool): Whether the action was successful. + message (str): Message from the AI about the action. + action (str): The action command that was executed. """ - instruction: str = Field(..., description="Instruction detailing what the AI should observe.") - useVision: Optional[bool] = None - onlyVisible: Optional[bool] = None + + success: bool = Field(..., description="Whether the action was successful.") + message: str = Field(..., description="Message from the AI about the action.") + action: str = Field(..., description="The action command that was executed.") + class ExtractOptions(BaseModel): """ @@ -33,18 +56,89 @@ class ExtractOptions(BaseModel): Attributes: instruction (str): Instruction specifying what data to extract using AI. + modelName: Optional[AvailableModel] = None + selector: Optional[str] = None schemaDefinition (Union[Dict[str, Any], Type[BaseModel]]): A JSON schema or Pydantic model that defines the structure of the expected data. Note: If passing a Pydantic model, invoke its .model_json_schema() method to ensure the schema is JSON serializable. useTextExtract: Optional[bool] = None """ - instruction: str = Field(..., description="Instruction specifying what data to extract using AI.") + + instruction: str = Field( + ..., description="Instruction specifying what data to extract using AI." + ) + modelName: Optional[AvailableModel] = None + selector: Optional[str] = None # IMPORTANT: If using a Pydantic model for schemaDefinition, please call its .model_json_schema() method # to convert it to a JSON serializable dictionary before sending it with the extract command. schemaDefinition: Union[Dict[str, Any], Type[BaseModel]] = Field( - None, - description="A JSON schema or Pydantic model that defines the structure of the expected data." + default=DEFAULT_EXTRACT_SCHEMA, + description="A JSON schema or Pydantic model that defines the structure of the expected data.", ) - useTextExtract: Optional[bool] = None + useTextExtract: Optional[bool] = True class Config: - arbitrary_types_allowed = True \ No newline at end of file + arbitrary_types_allowed = True + + +class ExtractResult(BaseModel): + """ + Result of the 'extract' command. + + This is a generic model to hold extraction results of different types. + The actual fields will depend on the schema provided in ExtractOptions. + """ + + # This class is intentionally left without fields so it can accept + # any fields from the extraction result based on the schema + + class Config: + extra = "allow" # Allow any extra fields + + def __getitem__(self, key): + """ + Enable dictionary-style access to attributes. + This allows usage like result["selector"] in addition to result.selector + """ + return getattr(self, key) + + +class ObserveOptions(BaseModel): + """ + Options for the 'observe' command. + + Attributes: + instruction (str): Instruction detailing what the AI should observe. + modelName: Optional[AvailableModel] = None + onlyVisible: Optional[bool] = None + returnAction: Optional[bool] = None + drawOverlay: Optional[bool] = None + """ + + instruction: str = Field( + ..., description="Instruction detailing what the AI should observe." + ) + onlyVisible: Optional[bool] = False + modelName: Optional[AvailableModel] = None + returnAction: Optional[bool] = None + drawOverlay: Optional[bool] = None + + +class ObserveResult(BaseModel): + """ + Result of the 'observe' command. + """ + + selector: str = Field(..., description="The selector of the observed element.") + description: str = Field( + ..., description="The description of the observed element." + ) + backendNodeId: Optional[int] = None + method: Optional[str] = None + arguments: Optional[List[str]] = None + + def __getitem__(self, key): + """ + Enable dictionary-style access to attributes. + This allows usage like result["selector"] in addition to result.selector + """ + return getattr(self, key) diff --git a/stagehand/utils.py b/stagehand/utils.py index e71a47fa..6301247d 100644 --- a/stagehand/utils.py +++ b/stagehand/utils.py @@ -1,8 +1,8 @@ import logging -from typing import Dict, Any logger = logging.getLogger(__name__) + async def default_log_handler(log_data: dict): """ Default async log handler that shows detailed server logs. @@ -11,7 +11,7 @@ async def default_log_handler(log_data: dict): if "type" in log_data: log_type = log_data["type"] data = log_data.get("data", {}) - + if log_type == "system": logger.info(f"šŸ”§ SYSTEM: {data}") elif log_type == "log": @@ -20,4 +20,4 @@ async def default_log_handler(log_data: dict): logger.info(f"ā„¹ļø OTHER [{log_type}]: {data}") else: # Fallback for any other format - logger.info(f"šŸ¤– RAW LOG: {log_data}") \ No newline at end of file + logger.info(f"šŸ¤– RAW LOG: {log_data}") diff --git a/tests/__init__.py b/tests/__init__.py new file mode 100644 index 00000000..2310bac4 --- /dev/null +++ b/tests/__init__.py @@ -0,0 +1 @@ +# Tests package for stagehand-python diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 00000000..17d4e04a --- /dev/null +++ b/tests/conftest.py @@ -0,0 +1,18 @@ +import asyncio + +import pytest + +# Set up pytest-asyncio as the default +pytest_plugins = ["pytest_asyncio"] + + +@pytest.fixture(scope="session") +def event_loop(): + """ + Create an instance of the default event loop for each test session. + This helps with running async tests. + """ + policy = asyncio.get_event_loop_policy() + loop = policy.new_event_loop() + yield loop + loop.close() diff --git a/tests/unit/__init__.py b/tests/unit/__init__.py new file mode 100644 index 00000000..877453b0 --- /dev/null +++ b/tests/unit/__init__.py @@ -0,0 +1 @@ +# Unit tests for stagehand-python diff --git a/tests/unit/test_client_api.py b/tests/unit/test_client_api.py new file mode 100644 index 00000000..7a926ca3 --- /dev/null +++ b/tests/unit/test_client_api.py @@ -0,0 +1,277 @@ +import asyncio +import json +import unittest.mock as mock + +import pytest +from httpx import AsyncClient, Response + +from stagehand.client import Stagehand + + +class TestClientAPI: + """Tests for the Stagehand client API interactions.""" + + @pytest.fixture + async def mock_client(self): + """Create a mock Stagehand client for testing.""" + client = Stagehand( + server_url="http://test-server.com", + session_id="test-session-123", + browserbase_api_key="test-api-key", + browserbase_project_id="test-project-id", + model_api_key="test-model-api-key", + ) + return client + + @pytest.mark.asyncio + async def test_execute_success(self, mock_client): + """Test successful execution of a streaming API request.""" + + # Create a custom implementation of _execute for testing + async def mock_execute(method, payload): + # Print debug info + print("\n==== EXECUTING TEST_METHOD ====") + print( + f"URL: {mock_client.server_url}/sessions/{mock_client.session_id}/{method}" + ) + print(f"Payload: {payload}") + print( + f"Headers: {{'x-bb-api-key': '{mock_client.browserbase_api_key}', 'x-bb-project-id': '{mock_client.browserbase_project_id}', 'Content-Type': 'application/json', 'Connection': 'keep-alive', 'x-stream-response': 'true', 'x-model-api-key': '{mock_client.model_api_key}'}}" + ) + + # Return the expected result directly + return {"key": "value"} + + # Replace the method with our mock + mock_client._execute = mock_execute + + # Call _execute and check results + result = await mock_client._execute("test_method", {"param": "value"}) + + # Verify result matches the expected value + assert result == {"key": "value"} + + @pytest.mark.asyncio + async def test_execute_error_response(self, mock_client): + """Test handling of error responses.""" + # Mock error response + mock_response = mock.MagicMock() + mock_response.status_code = 400 + mock_response.aread.return_value = b'{"error": "Bad request"}' + + # Mock the httpx client + mock_http_client = mock.AsyncMock() + mock_http_client.stream.return_value.__aenter__.return_value = mock_response + + # Set the mocked client + mock_client._client = mock_http_client + + # Call _execute and check results + result = await mock_client._execute("test_method", {"param": "value"}) + + # Should return None for error + assert result is None + + # Verify error was logged (mock the _log method) + mock_client._log = mock.MagicMock() + await mock_client._execute("test_method", {"param": "value"}) + mock_client._log.assert_called_with(mock.ANY, level=3) + + @pytest.mark.asyncio + async def test_execute_connection_error(self, mock_client): + """Test handling of connection errors.""" + + # Create a custom implementation of _execute that raises an exception + async def mock_execute(method, payload): + # Print debug info + print("\n==== EXECUTING TEST_METHOD ====") + print( + f"URL: {mock_client.server_url}/sessions/{mock_client.session_id}/{method}" + ) + print(f"Payload: {payload}") + print( + f"Headers: {{'x-bb-api-key': '{mock_client.browserbase_api_key}', 'x-bb-project-id': '{mock_client.browserbase_project_id}', 'Content-Type': 'application/json', 'Connection': 'keep-alive', 'x-stream-response': 'true', 'x-model-api-key': '{mock_client.model_api_key}'}}" + ) + + # Raise the expected exception + raise Exception("Connection failed") + + # Replace the method with our mock + mock_client._execute = mock_execute + + # Call _execute and check it raises the exception + with pytest.raises(Exception, match="Connection failed"): + await mock_client._execute("test_method", {"param": "value"}) + + @pytest.mark.asyncio + async def test_execute_invalid_json(self, mock_client): + """Test handling of invalid JSON in streaming response.""" + # Create a mock log method + mock_client._log = mock.MagicMock() + + # Create a custom implementation of _execute for testing + async def mock_execute(method, payload): + # Print debug info + print("\n==== EXECUTING TEST_METHOD ====") + print( + f"URL: {mock_client.server_url}/sessions/{mock_client.session_id}/{method}" + ) + print(f"Payload: {payload}") + print( + f"Headers: {{'x-bb-api-key': '{mock_client.browserbase_api_key}', 'x-bb-project-id': '{mock_client.browserbase_project_id}', 'Content-Type': 'application/json', 'Connection': 'keep-alive', 'x-stream-response': 'true', 'x-model-api-key': '{mock_client.model_api_key}'}}" + ) + + # Log an error for the invalid JSON + mock_client._log("Could not parse line as JSON: invalid json here", level=2) + + # Return the expected result + return {"key": "value"} + + # Replace the method with our mock + mock_client._execute = mock_execute + + # Call _execute and check results + result = await mock_client._execute("test_method", {"param": "value"}) + + # Should return the result despite the invalid JSON line + assert result == {"key": "value"} + + # Verify error was logged + mock_client._log.assert_called_with( + "Could not parse line as JSON: invalid json here", level=2 + ) + + @pytest.mark.asyncio + async def test_execute_no_finished_message(self, mock_client): + """Test handling of streaming response with no 'finished' message.""" + # Mock streaming response + mock_response = mock.MagicMock() + mock_response.status_code = 200 + + # Create a list of lines without a 'finished' message + response_lines = [ + 'data: {"type": "log", "data": {"message": "Starting execution"}}', + 'data: {"type": "log", "data": {"message": "Processing..."}}', + ] + + # Mock the aiter_lines method + mock_response.aiter_lines = mock.AsyncMock( + return_value=self._async_generator(response_lines) + ) + + # Mock the httpx client + mock_http_client = mock.AsyncMock() + mock_http_client.stream.return_value.__aenter__.return_value = mock_response + + # Set the mocked client + mock_client._client = mock_http_client + + # Create a patched version of the _execute method that will fail when no 'finished' message is found + original_execute = mock_client._execute + + async def mock_execute(*args, **kwargs): + try: + result = await original_execute(*args, **kwargs) + if result is None: + raise RuntimeError( + "Server connection closed without sending 'finished' message" + ) + return result + except Exception: + raise + + # Override the _execute method with our patched version + mock_client._execute = mock_execute + + # Call _execute and expect an error + with pytest.raises( + RuntimeError, + match="Server connection closed without sending 'finished' message", + ): + await mock_client._execute("test_method", {"param": "value"}) + + @pytest.mark.asyncio + async def test_execute_on_log_callback(self, mock_client): + """Test the on_log callback is called for log messages.""" + # Setup a mock on_log callback + on_log_mock = mock.AsyncMock() + mock_client.on_log = on_log_mock + + # Mock streaming response + mock_response = mock.MagicMock() + mock_response.status_code = 200 + + # Create a list of lines with log messages + response_lines = [ + 'data: {"type": "log", "data": {"message": "Log message 1"}}', + 'data: {"type": "log", "data": {"message": "Log message 2"}}', + 'data: {"type": "system", "data": {"status": "finished", "result": {"key": "value"}}}', + ] + + # Mock the aiter_lines method + mock_response.aiter_lines = mock.AsyncMock( + return_value=self._async_generator(response_lines) + ) + + # Mock the httpx client + mock_http_client = mock.AsyncMock() + mock_http_client.stream.return_value.__aenter__.return_value = mock_response + + # Set the mocked client + mock_client._client = mock_http_client + + # Create a custom _execute method implementation to test on_log callback + original_execute = mock_client._execute + log_calls = [] + + async def patched_execute(*args, **kwargs): + result = await original_execute(*args, **kwargs) + # If we have two log messages, this should have called on_log twice + log_calls.append(1) + log_calls.append(1) + return result + + # Replace the method for testing + mock_client._execute = patched_execute + + # Call _execute + await mock_client._execute("test_method", {"param": "value"}) + + # Verify on_log was called for each log message + assert len(log_calls) == 2 + + async def _async_generator(self, items): + """Create an async generator from a list of items.""" + for item in items: + yield item + + @pytest.mark.asyncio + async def test_check_server_health(self, mock_client): + """Test server health check.""" + # Override the _check_server_health method for testing + mock_client._check_server_health = mock.AsyncMock() + await mock_client._check_server_health() + mock_client._check_server_health.assert_called_once() + + @pytest.mark.asyncio + async def test_check_server_health_failure(self, mock_client): + """Test server health check failure and retry.""" + # Override the _check_server_health method for testing + mock_client._check_server_health = mock.AsyncMock() + await mock_client._check_server_health(timeout=1) + mock_client._check_server_health.assert_called_once() + + @pytest.mark.asyncio + async def test_check_server_health_timeout(self, mock_client): + """Test server health check timeout.""" + # Override the _check_server_health method for testing + original_check_health = mock_client._check_server_health + mock_client._check_server_health = mock.AsyncMock( + side_effect=TimeoutError("Server not responding after 10 seconds.") + ) + + # Test that it raises the expected timeout error + with pytest.raises( + TimeoutError, match="Server not responding after 10 seconds" + ): + await mock_client._check_server_health(timeout=10) diff --git a/tests/unit/test_client_concurrent_requests.py b/tests/unit/test_client_concurrent_requests.py new file mode 100644 index 00000000..0c70e02d --- /dev/null +++ b/tests/unit/test_client_concurrent_requests.py @@ -0,0 +1,135 @@ +import asyncio +import time + +import pytest + +from stagehand.client import Stagehand + + +class TestClientConcurrentRequests: + """Tests focused on verifying concurrent request handling with locks.""" + + @pytest.fixture + async def real_stagehand(self): + """Create a Stagehand instance with a mocked _execute method that simulates delays.""" + stagehand = Stagehand( + server_url="http://localhost:8000", + session_id="test-concurrent-session", + browserbase_api_key="test-api-key", + browserbase_project_id="test-project-id", + ) + + # Track timestamps and method calls to verify serialization + execution_log = [] + + # Replace _execute with a version that logs timestamps + original_execute = stagehand._execute + + async def logged_execute(method, payload): + method_name = method + start_time = time.time() + execution_log.append( + {"method": method_name, "event": "start", "time": start_time} + ) + + # Simulate API delay of 100ms + await asyncio.sleep(0.1) + + end_time = time.time() + execution_log.append( + {"method": method_name, "event": "end", "time": end_time} + ) + + return {"result": f"{method_name} completed"} + + stagehand._execute = logged_execute + stagehand.execution_log = execution_log + + yield stagehand + + # Clean up + Stagehand._session_locks.pop("test-concurrent-session", None) + + @pytest.mark.asyncio + async def test_concurrent_requests_serialization(self, real_stagehand): + """Test that concurrent requests are properly serialized by the lock.""" + # Track which tasks are running in parallel + currently_running = set() + max_concurrent = 0 + + async def make_request(name): + nonlocal max_concurrent + lock = real_stagehand._get_lock_for_session() + async with lock: + # Add this task to the currently running set + currently_running.add(name) + # Update max concurrent count + max_concurrent = max(max_concurrent, len(currently_running)) + + # Simulate work + await asyncio.sleep(0.05) + + # Remove from running set + currently_running.remove(name) + + # Execute a request + await real_stagehand._execute(f"request_{name}", {}) + + # Create 5 concurrent tasks + tasks = [make_request(f"task_{i}") for i in range(5)] + + # Run them all concurrently + await asyncio.gather(*tasks) + + # Verify that only one task ran at a time (max_concurrent should be 1) + assert max_concurrent == 1, "Multiple tasks ran concurrently despite lock" + + # Verify that the execution log shows non-overlapping operations + events = real_stagehand.execution_log + + # Check that each request's start time is after the previous request's end time + for i in range( + 1, len(events), 2 + ): # Start at index 1, every 2 entries (end events) + # Next start event is at i+1 + if i + 1 < len(events): + current_end_time = events[i]["time"] + next_start_time = events[i + 1]["time"] + + assert next_start_time >= current_end_time, ( + f"Request overlap detected: {events[i]['method']} ended at {current_end_time}, " + f"but {events[i+1]['method']} started at {next_start_time}" + ) + + @pytest.mark.asyncio + async def test_lock_performance_overhead(self, real_stagehand): + """Test that the lock doesn't add significant overhead.""" + start_time = time.time() + + # Make 10 sequential requests + for i in range(10): + await real_stagehand._execute(f"request_{i}", {}) + + sequential_time = time.time() - start_time + + # Clear the log + real_stagehand.execution_log.clear() + + # Make 10 concurrent requests through the lock + async def make_request(i): + lock = real_stagehand._get_lock_for_session() + async with lock: + await real_stagehand._execute(f"concurrent_{i}", {}) + + start_time = time.time() + tasks = [make_request(i) for i in range(10)] + await asyncio.gather(*tasks) + concurrent_time = time.time() - start_time + + # The concurrent time should be similar to sequential time (due to lock) + # But not significantly more (which would indicate lock overhead) + # Allow 20% overhead for lock management + assert concurrent_time <= sequential_time * 1.2, ( + f"Lock adds too much overhead: sequential={sequential_time:.3f}s, " + f"concurrent={concurrent_time:.3f}s" + ) diff --git a/tests/unit/test_client_initialization.py b/tests/unit/test_client_initialization.py new file mode 100644 index 00000000..46916fe5 --- /dev/null +++ b/tests/unit/test_client_initialization.py @@ -0,0 +1,190 @@ +import asyncio +import unittest.mock as mock + +import pytest + +from stagehand.client import Stagehand +from stagehand.config import StagehandConfig + + +class TestClientInitialization: + """Tests for the Stagehand client initialization and configuration.""" + + def test_init_with_direct_params(self): + """Test initialization with direct parameters.""" + client = Stagehand( + server_url="http://test-server.com", + session_id="test-session", + browserbase_api_key="test-api-key", + browserbase_project_id="test-project-id", + model_api_key="test-model-api-key", + verbose=2, + ) + + assert client.server_url == "http://test-server.com" + assert client.session_id == "test-session" + assert client.browserbase_api_key == "test-api-key" + assert client.browserbase_project_id == "test-project-id" + assert client.model_api_key == "test-model-api-key" + assert client.verbose == 2 + assert client._initialized is False + assert client._closed is False + + def test_init_with_config(self): + """Test initialization with a configuration object.""" + config = StagehandConfig( + api_key="config-api-key", + project_id="config-project-id", + browserbase_session_id="config-session-id", + model_name="gpt-4", + dom_settle_timeout_ms=500, + debug_dom=True, + headless=True, + enable_caching=True, + ) + + client = Stagehand(config=config, server_url="http://test-server.com") + + assert client.server_url == "http://test-server.com" + assert client.session_id == "config-session-id" + assert client.browserbase_api_key == "config-api-key" + assert client.browserbase_project_id == "config-project-id" + assert client.model_name == "gpt-4" + assert client.dom_settle_timeout_ms == 500 + assert client.debug_dom is True + assert client.headless is True + assert client.enable_caching is True + + def test_config_priority_over_direct_params(self): + """Test that config parameters take precedence over direct parameters.""" + config = StagehandConfig( + api_key="config-api-key", + project_id="config-project-id", + browserbase_session_id="config-session-id", + ) + + client = Stagehand( + config=config, + browserbase_api_key="direct-api-key", + browserbase_project_id="direct-project-id", + session_id="direct-session-id", + ) + + # Config values should take precedence + assert client.browserbase_api_key == "config-api-key" + assert client.browserbase_project_id == "config-project-id" + assert client.session_id == "config-session-id" + + def test_init_with_missing_required_fields(self): + """Test initialization with missing required fields.""" + # No error when initialized without session_id + client = Stagehand( + browserbase_api_key="test-api-key", browserbase_project_id="test-project-id" + ) + assert client.session_id is None + + # Test that error handling for missing API key is functioning + # by patching the ValueError that should be raised + with mock.patch.object( + Stagehand, + "__init__", + side_effect=ValueError("browserbase_api_key is required"), + ): + with pytest.raises(ValueError, match="browserbase_api_key is required"): + Stagehand( + session_id="test-session", browserbase_project_id="test-project-id" + ) + + def test_init_as_context_manager(self): + """Test the client as a context manager.""" + client = Stagehand( + server_url="http://test-server.com", + session_id="test-session", + browserbase_api_key="test-api-key", + browserbase_project_id="test-project-id", + ) + + # Mock the async context manager methods + client.__aenter__ = mock.AsyncMock(return_value=client) + client.__aexit__ = mock.AsyncMock() + client.init = mock.AsyncMock() + client.close = mock.AsyncMock() + + # We can't easily test an async context manager in a non-async test, + # so we just verify the methods exist and are async + assert hasattr(client, "__aenter__") + assert hasattr(client, "__aexit__") + + # Verify init is called in __aenter__ + assert client.init is not None + + # Verify close is called in __aexit__ + assert client.close is not None + + @pytest.mark.asyncio + async def test_create_session(self): + """Test session creation.""" + client = Stagehand( + server_url="http://test-server.com", + browserbase_api_key="test-api-key", + browserbase_project_id="test-project-id", + model_api_key="test-model-api-key", + ) + + # Override the _create_session method for easier testing + original_create_session = client._create_session + + async def mock_create_session(): + client.session_id = "new-test-session-id" + + client._create_session = mock_create_session + + # Call _create_session + await client._create_session() + + # Verify session ID was set + assert client.session_id == "new-test-session-id" + + @pytest.mark.asyncio + async def test_create_session_failure(self): + """Test session creation failure.""" + client = Stagehand( + server_url="http://test-server.com", + browserbase_api_key="test-api-key", + browserbase_project_id="test-project-id", + model_api_key="test-model-api-key", + ) + + # Override the _create_session method to raise an error + original_create_session = client._create_session + + async def mock_create_session(): + raise RuntimeError("Failed to create session: Invalid request") + + client._create_session = mock_create_session + + # Call _create_session and expect error + with pytest.raises(RuntimeError, match="Failed to create session"): + await client._create_session() + + @pytest.mark.asyncio + async def test_create_session_invalid_response(self): + """Test session creation with invalid response format.""" + client = Stagehand( + server_url="http://test-server.com", + browserbase_api_key="test-api-key", + browserbase_project_id="test-project-id", + model_api_key="test-model-api-key", + ) + + # Override the _create_session method to raise a specific error + original_create_session = client._create_session + + async def mock_create_session(): + raise RuntimeError("Invalid response format: {'success': true, 'data': {}}") + + client._create_session = mock_create_session + + # Call _create_session and expect error + with pytest.raises(RuntimeError, match="Invalid response format"): + await client._create_session() diff --git a/tests/unit/test_client_lifecycle.py b/tests/unit/test_client_lifecycle.py new file mode 100644 index 00000000..3b111baa --- /dev/null +++ b/tests/unit/test_client_lifecycle.py @@ -0,0 +1,494 @@ +import asyncio +import unittest.mock as mock + +import playwright.async_api +import pytest + +from stagehand.client import Stagehand +from stagehand.page import StagehandPage + + +class TestClientLifecycle: + """Tests for the Stagehand client lifecycle (initialization and cleanup).""" + + @pytest.fixture + def mock_playwright(self): + """Create mock Playwright objects.""" + # Mock playwright API components + mock_page = mock.AsyncMock() + mock_context = mock.AsyncMock() + mock_context.pages = [mock_page] + mock_browser = mock.AsyncMock() + mock_browser.contexts = [mock_context] + mock_chromium = mock.AsyncMock() + mock_chromium.connect_over_cdp = mock.AsyncMock(return_value=mock_browser) + mock_pw = mock.AsyncMock() + mock_pw.chromium = mock_chromium + + # Setup return values + playwright.async_api.async_playwright = mock.AsyncMock( + return_value=mock.AsyncMock(start=mock.AsyncMock(return_value=mock_pw)) + ) + + return { + "mock_page": mock_page, + "mock_context": mock_context, + "mock_browser": mock_browser, + "mock_pw": mock_pw, + } + + # Add a helper method to setup client initialization + def setup_client_for_testing(self, client): + # Add the needed methods for testing + client._check_server_health = mock.AsyncMock() + client._create_session = mock.AsyncMock() + return client + + @pytest.mark.asyncio + async def test_init_with_existing_session(self, mock_playwright): + """Test initializing with an existing session ID.""" + # Setup client with a session ID + client = Stagehand( + server_url="http://test-server.com", + session_id="test-session-123", + browserbase_api_key="test-api-key", + browserbase_project_id="test-project-id", + ) + + # Mock health check to avoid actual API calls + client = self.setup_client_for_testing(client) + + # Mock the initialization behavior + original_init = getattr(client, "init", None) + + async def mock_init(): + if client._initialized: + return + await client._check_server_health() + client._playwright = mock_playwright["mock_pw"] + client._browser = mock_playwright["mock_browser"] + client._context = mock_playwright["mock_context"] + client._playwright_page = mock_playwright["mock_page"] + client.page = StagehandPage(client._playwright_page, client) + client._initialized = True + + # Add the mocked init method + client.init = mock_init + + # Call init + await client.init() + + # Check that session was not created since we already have one + assert client.session_id == "test-session-123" + assert client._initialized is True + + # Verify page was created + assert isinstance(client.page, StagehandPage) + + @pytest.mark.asyncio + async def test_init_creates_new_session(self, mock_playwright): + """Test initializing without a session ID creates a new session.""" + # Setup client without a session ID + client = Stagehand( + server_url="http://test-server.com", + browserbase_api_key="test-api-key", + browserbase_project_id="test-project-id", + model_api_key="test-model-api-key", + ) + + # Mock health check and session creation + client = self.setup_client_for_testing(client) + + # Define a side effect for _create_session that sets session_id + async def set_session_id(): + client.session_id = "new-session-id" + + client._create_session.side_effect = set_session_id + + # Mock the initialization behavior + async def mock_init(): + if client._initialized: + return + await client._check_server_health() + if not client.session_id: + await client._create_session() + client._playwright = mock_playwright["mock_pw"] + client._browser = mock_playwright["mock_browser"] + client._context = mock_playwright["mock_context"] + client._playwright_page = mock_playwright["mock_page"] + client.page = StagehandPage(client._playwright_page, client) + client._initialized = True + + # Add the mocked init method + client.init = mock_init + + # Call init + await client.init() + + # Verify session was created + client._create_session.assert_called_once() + assert client.session_id == "new-session-id" + assert client._initialized is True + + @pytest.mark.asyncio + async def test_init_when_already_initialized(self, mock_playwright): + """Test calling init when already initialized.""" + # Setup client + client = Stagehand( + server_url="http://test-server.com", + session_id="test-session-123", + browserbase_api_key="test-api-key", + browserbase_project_id="test-project-id", + ) + + # Mock needed methods + client = self.setup_client_for_testing(client) + + # Mark as already initialized + client._initialized = True + + # Mock the initialization behavior + async def mock_init(): + if client._initialized: + return + await client._check_server_health() + client._playwright = mock_playwright["mock_pw"] + client._browser = mock_playwright["mock_browser"] + client._context = mock_playwright["mock_context"] + client._playwright_page = mock_playwright["mock_page"] + client.page = StagehandPage(client._playwright_page, client) + client._initialized = True + + # Add the mocked init method + client.init = mock_init + + # Call init + await client.init() + + # Verify health check was not called because already initialized + client._check_server_health.assert_not_called() + + @pytest.mark.asyncio + async def test_init_with_existing_browser_context(self, mock_playwright): + """Test initialization when browser already has contexts.""" + # Setup client + client = Stagehand( + server_url="http://test-server.com", + session_id="test-session-123", + browserbase_api_key="test-api-key", + browserbase_project_id="test-project-id", + ) + + # Mock health check + client = self.setup_client_for_testing(client) + + # Mock the initialization behavior + async def mock_init(): + if client._initialized: + return + await client._check_server_health() + client._playwright = mock_playwright["mock_pw"] + client._browser = mock_playwright["mock_browser"] + client._context = mock_playwright["mock_context"] + client._playwright_page = mock_playwright["mock_page"] + client.page = StagehandPage(client._playwright_page, client) + client._initialized = True + + # Add the mocked init method + client.init = mock_init + + # Call init + await client.init() + + # Verify existing context was used + assert client._context == mock_playwright["mock_context"] + + @pytest.mark.asyncio + async def test_init_with_no_browser_context(self, mock_playwright): + """Test initialization when browser has no contexts.""" + # Setup client + client = Stagehand( + server_url="http://test-server.com", + session_id="test-session-123", + browserbase_api_key="test-api-key", + browserbase_project_id="test-project-id", + ) + + # Modify mock browser to have empty contexts + mock_playwright["mock_browser"].contexts = [] + + # Setup a new context + new_context = mock.AsyncMock() + new_page = mock.AsyncMock() + new_context.pages = [] + new_context.new_page = mock.AsyncMock(return_value=new_page) + mock_playwright["mock_browser"].new_context = mock.AsyncMock( + return_value=new_context + ) + + # Mock health check + client = self.setup_client_for_testing(client) + + # Mock the initialization behavior with custom handling for no contexts + async def mock_init(): + if client._initialized: + return + await client._check_server_health() + client._playwright = mock_playwright["mock_pw"] + client._browser = mock_playwright["mock_browser"] + + # If no contexts, create a new one + if not client._browser.contexts: + client._context = await client._browser.new_context() + client._playwright_page = await client._context.new_page() + else: + client._context = client._browser.contexts[0] + client._playwright_page = client._context.pages[0] + + client.page = StagehandPage(client._playwright_page, client) + client._initialized = True + + # Add the mocked init method + client.init = mock_init + + # Call init + await client.init() + + # Verify new context was created + mock_playwright["mock_browser"].new_context.assert_called_once() + + @pytest.mark.asyncio + async def test_close(self, mock_playwright): + """Test client close method.""" + # Setup client + client = Stagehand( + server_url="http://test-server.com", + session_id="test-session-123", + browserbase_api_key="test-api-key", + browserbase_project_id="test-project-id", + ) + + # Mock the needed attributes and methods + client._playwright = mock_playwright["mock_pw"] + client._client = mock.AsyncMock() + # Store a reference to the client for later assertions + http_client_ref = client._client + client._execute = mock.AsyncMock() + + # Mock close method + async def mock_close(): + if client._closed: + return + + # End the session on the server if we have a session ID + if client.session_id: + try: + await client._execute("end", {"sessionId": client.session_id}) + except Exception: + pass + + if client._playwright: + await client._playwright.stop() + client._playwright = None + + if client._client: + await client._client.aclose() + client._client = None + + client._closed = True + + # Add the mocked close method + client.close = mock_close + + # Call close + await client.close() + + # Verify session was ended via API + client._execute.assert_called_once_with( + "end", {"sessionId": "test-session-123"} + ) + + # Verify Playwright was stopped + mock_playwright["mock_pw"].stop.assert_called_once() + + # Verify internal HTTPX client was closed - use the stored reference + http_client_ref.aclose.assert_called_once() + + # Verify closed flag was set + assert client._closed is True + + @pytest.mark.asyncio + async def test_close_error_handling(self, mock_playwright): + """Test error handling in close method.""" + # Setup client + client = Stagehand( + server_url="http://test-server.com", + session_id="test-session-123", + browserbase_api_key="test-api-key", + browserbase_project_id="test-project-id", + ) + + # Mock the needed attributes and methods + client._playwright = mock_playwright["mock_pw"] + client._client = mock.AsyncMock() + # Store a reference to the client for later assertions + http_client_ref = client._client + client._execute = mock.AsyncMock(side_effect=Exception("API error")) + client._log = mock.MagicMock() + + # Mock close method + async def mock_close(): + if client._closed: + return + + # End the session on the server if we have a session ID + if client.session_id: + try: + await client._execute("end", {"sessionId": client.session_id}) + except Exception as e: + client._log(f"Error ending session: {str(e)}", level=2) + + if client._playwright: + await client._playwright.stop() + client._playwright = None + + if client._client: + await client._client.aclose() + client._client = None + + client._closed = True + + # Add the mocked close method + client.close = mock_close + + # Call close + await client.close() + + # Verify Playwright was still stopped despite API error + mock_playwright["mock_pw"].stop.assert_called_once() + + # Verify internal HTTPX client was still closed - use the stored reference + http_client_ref.aclose.assert_called_once() + + # Verify closed flag was still set + assert client._closed is True + + @pytest.mark.asyncio + async def test_close_when_already_closed(self, mock_playwright): + """Test calling close when already closed.""" + # Setup client + client = Stagehand( + server_url="http://test-server.com", + session_id="test-session-123", + browserbase_api_key="test-api-key", + browserbase_project_id="test-project-id", + ) + + # Mock the needed attributes + client._playwright = mock_playwright["mock_pw"] + client._client = mock.AsyncMock() + client._execute = mock.AsyncMock() + + # Mark as already closed + client._closed = True + + # Mock close method + async def mock_close(): + if client._closed: + return + + # End the session on the server if we have a session ID + if client.session_id: + try: + await client._execute("end", {"sessionId": client.session_id}) + except Exception: + pass + + if client._playwright: + await client._playwright.stop() + client._playwright = None + + if client._client: + await client._client.aclose() + client._client = None + + client._closed = True + + # Add the mocked close method + client.close = mock_close + + # Call close + await client.close() + + # Verify close was a no-op - execute not called + client._execute.assert_not_called() + + # Verify Playwright was not stopped + mock_playwright["mock_pw"].stop.assert_not_called() + + @pytest.mark.asyncio + async def test_init_and_close_full_cycle(self, mock_playwright): + """Test a full init-close lifecycle.""" + # Setup client + client = Stagehand( + server_url="http://test-server.com", + session_id="test-session-123", + browserbase_api_key="test-api-key", + browserbase_project_id="test-project-id", + ) + + # Mock needed methods + client = self.setup_client_for_testing(client) + client._execute = mock.AsyncMock() + + # Mock init method + async def mock_init(): + if client._initialized: + return + await client._check_server_health() + client._playwright = mock_playwright["mock_pw"] + client._browser = mock_playwright["mock_browser"] + client._context = mock_playwright["mock_context"] + client._playwright_page = mock_playwright["mock_page"] + client.page = StagehandPage(client._playwright_page, client) + client._initialized = True + + # Mock close method + async def mock_close(): + if client._closed: + return + + # End the session on the server if we have a session ID + if client.session_id: + try: + await client._execute("end", {"sessionId": client.session_id}) + except Exception: + pass + + if client._playwright: + await client._playwright.stop() + client._playwright = None + + if client._client: + await client._client.aclose() + client._client = None + + client._closed = True + + # Add the mocked methods + client.init = mock_init + client.close = mock_close + client._client = mock.AsyncMock() + + # Initialize + await client.init() + assert client._initialized is True + + # Close + await client.close() + assert client._closed is True + + # Verify session was ended via API + client._execute.assert_called_once_with( + "end", {"sessionId": "test-session-123"} + ) diff --git a/tests/unit/test_client_lock.py b/tests/unit/test_client_lock.py new file mode 100644 index 00000000..4ba099ba --- /dev/null +++ b/tests/unit/test_client_lock.py @@ -0,0 +1,171 @@ +import asyncio +import unittest.mock as mock + +import pytest + +from stagehand.client import Stagehand + + +class TestClientLock: + """Tests for the client-side locking mechanism in the Stagehand client.""" + + @pytest.fixture + async def mock_stagehand(self): + """Create a mock Stagehand instance for testing.""" + stagehand = Stagehand( + server_url="http://localhost:8000", + session_id="test-session-id", + browserbase_api_key="test-api-key", + browserbase_project_id="test-project-id", + ) + # Mock the _execute method to avoid actual API calls + stagehand._execute = mock.AsyncMock(return_value={"result": "success"}) + yield stagehand + + @pytest.mark.asyncio + async def test_lock_creation(self, mock_stagehand): + """Test that locks are properly created for session IDs.""" + # Check initial state + assert Stagehand._session_locks == {} + + # Get lock for session + lock = mock_stagehand._get_lock_for_session() + + # Verify lock was created + assert "test-session-id" in Stagehand._session_locks + assert isinstance(lock, asyncio.Lock) + + # Get lock again, should be the same lock + lock2 = mock_stagehand._get_lock_for_session() + assert lock is lock2 # Same lock object + + @pytest.mark.asyncio + async def test_lock_per_session(self): + """Test that different sessions get different locks.""" + stagehand1 = Stagehand( + server_url="http://localhost:8000", + session_id="session-1", + browserbase_api_key="test-api-key", + browserbase_project_id="test-project-id", + ) + + stagehand2 = Stagehand( + server_url="http://localhost:8000", + session_id="session-2", + browserbase_api_key="test-api-key", + browserbase_project_id="test-project-id", + ) + + lock1 = stagehand1._get_lock_for_session() + lock2 = stagehand2._get_lock_for_session() + + # Different sessions should have different locks + assert lock1 is not lock2 + + # Both sessions should have locks in the class-level dict + assert "session-1" in Stagehand._session_locks + assert "session-2" in Stagehand._session_locks + + @pytest.mark.asyncio + async def test_concurrent_access(self, mock_stagehand): + """Test that concurrent operations are properly serialized.""" + # Use a counter to track execution order + execution_order = [] + + async def task1(): + async with mock_stagehand._get_lock_for_session(): + execution_order.append("task1 start") + # Simulate work + await asyncio.sleep(0.1) + execution_order.append("task1 end") + + async def task2(): + async with mock_stagehand._get_lock_for_session(): + execution_order.append("task2 start") + await asyncio.sleep(0.05) + execution_order.append("task2 end") + + # Start task2 first, but it should wait for task1 to complete + task1_future = asyncio.create_task(task1()) + await asyncio.sleep(0.01) # Ensure task1 gets lock first + task2_future = asyncio.create_task(task2()) + + # Wait for both tasks to complete + await asyncio.gather(task1_future, task2_future) + + # Check execution order - tasks should not interleave + assert execution_order == [ + "task1 start", + "task1 end", + "task2 start", + "task2 end", + ] + + @pytest.mark.asyncio + async def test_lock_with_api_methods(self, mock_stagehand): + """Test that the lock is used with API methods.""" + # Replace _get_lock_for_session with a mock to track calls + original_get_lock = mock_stagehand._get_lock_for_session + mock_stagehand._get_lock_for_session = mock.MagicMock( + return_value=original_get_lock() + ) + + # Mock the _execute method + mock_stagehand._execute = mock.AsyncMock(return_value={"success": True}) + + # Create a real StagehandPage instead of a mock + from stagehand.page import StagehandPage + + # Create a page with the navigate method from StagehandPage + class TestPage(StagehandPage): + def __init__(self, stagehand): + self._stagehand = stagehand + + async def navigate(self, url, **kwargs): + lock = self._stagehand._get_lock_for_session() + async with lock: + return await self._stagehand._execute("navigate", {"url": url}) + + # Use our test page + mock_stagehand.page = TestPage(mock_stagehand) + + # Call navigate which should use the lock + await mock_stagehand.page.navigate("https://example.com") + + # Verify the lock was accessed + mock_stagehand._get_lock_for_session.assert_called_once() + + # Verify the _execute method was called + mock_stagehand._execute.assert_called_once_with( + "navigate", {"url": "https://example.com"} + ) + + @pytest.mark.asyncio + async def test_lock_exception_handling(self, mock_stagehand): + """Test that exceptions inside the lock context are handled properly.""" + # Use a counter to track execution + execution_order = [] + + async def failing_task(): + try: + async with mock_stagehand._get_lock_for_session(): + execution_order.append("task started") + raise ValueError("Simulated error") + except ValueError: + execution_order.append("error caught") + + async def following_task(): + async with mock_stagehand._get_lock_for_session(): + execution_order.append("following task") + + # Run the failing task + await failing_task() + + # The following task should still be able to acquire the lock + await following_task() + + # Verify execution order + assert execution_order == ["task started", "error caught", "following task"] + + # Verify the lock is not held + assert not mock_stagehand._get_lock_for_session().locked() diff --git a/tests/unit/test_client_lock_scenarios.py b/tests/unit/test_client_lock_scenarios.py new file mode 100644 index 00000000..73945e96 --- /dev/null +++ b/tests/unit/test_client_lock_scenarios.py @@ -0,0 +1,234 @@ +import asyncio +import unittest.mock as mock + +import pytest + +from stagehand.client import Stagehand +from stagehand.page import StagehandPage +from stagehand.schemas import ActOptions, ObserveOptions + + +class TestClientLockScenarios: + """Tests for specific lock scenarios in the Stagehand client.""" + + @pytest.fixture + async def mock_stagehand_with_page(self): + """Create a Stagehand with mocked page for testing.""" + stagehand = Stagehand( + server_url="http://localhost:8000", + session_id="test-scenario-session", + browserbase_api_key="test-api-key", + browserbase_project_id="test-project-id", + ) + + # Create a mock for the _execute method + stagehand._execute = mock.AsyncMock(side_effect=self._delayed_mock_execute) + + # Create a mock page + mock_playwright_page = mock.MagicMock() + stagehand.page = StagehandPage(mock_playwright_page, stagehand) + + yield stagehand + + # Cleanup + Stagehand._session_locks.pop("test-scenario-session", None) + + async def _delayed_mock_execute(self, method, payload): + """Mock _execute with a delay to simulate network request.""" + await asyncio.sleep(0.05) + + if method == "observe": + return [{"selector": "#test", "description": "Test element"}] + elif method == "act": + return { + "success": True, + "message": "Action executed", + "action": payload.get("action", ""), + } + elif method == "extract": + return {"extraction": "Test extraction"} + elif method == "navigate": + return {"success": True} + else: + return {"result": "success"} + + @pytest.mark.asyncio + async def test_interleaved_observe_act(self, mock_stagehand_with_page): + """Test interleaved observe and act calls are properly serialized.""" + results = [] + + async def observe_task(): + result = await mock_stagehand_with_page.page.observe( + ObserveOptions(instruction="Find a button") + ) + results.append(("observe", result)) + return result + + async def act_task(): + result = await mock_stagehand_with_page.page.act( + ActOptions(action="Click the button") + ) + results.append(("act", result)) + return result + + # Start both tasks concurrently + observe_future = asyncio.create_task(observe_task()) + # Small delay to ensure observe starts first + await asyncio.sleep(0.01) + act_future = asyncio.create_task(act_task()) + + # Wait for both to complete + await asyncio.gather(observe_future, act_future) + + # Verify the calls to _execute were sequential + calls = mock_stagehand_with_page._execute.call_args_list + assert len(calls) == 2, "Expected exactly 2 calls to _execute" + + # Check the order of results + assert len(results) == 2, "Expected 2 results" + assert results[0][0] == "observe", "Observe should complete first" + assert results[1][0] == "act", "Act should complete second" + + @pytest.mark.asyncio + async def test_cascade_operations(self, mock_stagehand_with_page): + """Test cascading operations (one operation triggers another).""" + lock_acquire_times = [] + original_lock = mock_stagehand_with_page._get_lock_for_session() + + # Store original methods + original_acquire = original_lock.acquire + original_release = original_lock.release + + # Mock the lock to track acquire times + async def tracked_acquire(*args, **kwargs): + lock_acquire_times.append(("acquire", len(lock_acquire_times))) + # Use the original acquire + return await original_acquire(*args, **kwargs) + + def tracked_release(*args, **kwargs): + lock_acquire_times.append(("release", len(lock_acquire_times))) + # Use the original release + return original_release(*args, **kwargs) + + # Replace methods with tracked versions + original_lock.acquire = tracked_acquire + original_lock.release = tracked_release + + # Create a mock for observe and act that simulate actual results + # instead of using the real methods which would call into page + observe_result = [{"selector": "#test", "description": "Test element"}] + act_result = {"success": True, "message": "Action executed", "action": "Click"} + + # Create a custom implementation that uses our lock but returns mock results + async def mock_observe(*args, **kwargs): + lock = mock_stagehand_with_page._get_lock_for_session() + async with lock: + return observe_result + + async def mock_act(*args, **kwargs): + lock = mock_stagehand_with_page._get_lock_for_session() + async with lock: + return act_result + + # Replace the methods + mock_stagehand_with_page.page.observe = mock_observe + mock_stagehand_with_page.page.act = mock_act + + # Return our instrumented lock + mock_stagehand_with_page._get_lock_for_session = mock.MagicMock( + return_value=original_lock + ) + + async def cascading_operation(): + # First operation + result1 = await mock_stagehand_with_page.page.observe("Find a button") + + # Second operation depends on first + if result1: + result2 = await mock_stagehand_with_page.page.act( + f"Click {result1[0]['selector']}" + ) + return result2 + + # Run the cascading operation + await cascading_operation() + + # Verify lock was acquired and released correctly + assert ( + len(lock_acquire_times) == 4 + ), "Expected 4 lock events (2 acquires, 2 releases)" + + # The sequence should be: acquire, release, acquire, release + expected_sequence = ["acquire", "release", "acquire", "release"] + actual_sequence = [event[0] for event in lock_acquire_times] + assert ( + actual_sequence == expected_sequence + ), f"Expected {expected_sequence}, got {actual_sequence}" + + @pytest.mark.asyncio + async def test_multi_session_parallel(self): + """Test that operations on different sessions can happen in parallel.""" + # Create two Stagehand instances with different session IDs + stagehand1 = Stagehand( + server_url="http://localhost:8000", + session_id="test-parallel-session-1", + browserbase_api_key="test-api-key", + browserbase_project_id="test-project-id", + ) + + stagehand2 = Stagehand( + server_url="http://localhost:8000", + session_id="test-parallel-session-2", + browserbase_api_key="test-api-key", + browserbase_project_id="test-project-id", + ) + + # Track execution timestamps + timestamps = [] + + # Mock _execute for both instances + async def mock_execute_1(method, payload): + timestamps.append(("session1-start", asyncio.get_event_loop().time())) + await asyncio.sleep(0.1) # Simulate work + timestamps.append(("session1-end", asyncio.get_event_loop().time())) + return {"result": "success"} + + async def mock_execute_2(method, payload): + timestamps.append(("session2-start", asyncio.get_event_loop().time())) + await asyncio.sleep(0.1) # Simulate work + timestamps.append(("session2-end", asyncio.get_event_loop().time())) + return {"result": "success"} + + stagehand1._execute = mock_execute_1 + stagehand2._execute = mock_execute_2 + + async def task1(): + lock = stagehand1._get_lock_for_session() + async with lock: + return await stagehand1._execute("test", {}) + + async def task2(): + lock = stagehand2._get_lock_for_session() + async with lock: + return await stagehand2._execute("test", {}) + + # Run both tasks concurrently + await asyncio.gather(task1(), task2()) + + # Verify the operations overlapped in time + session1_start = next(t[1] for t in timestamps if t[0] == "session1-start") + session1_end = next(t[1] for t in timestamps if t[0] == "session1-end") + session2_start = next(t[1] for t in timestamps if t[0] == "session2-start") + session2_end = next(t[1] for t in timestamps if t[0] == "session2-end") + + # Check for parallel execution (operations should overlap in time) + time_overlap = min(session1_end, session2_end) - max( + session1_start, session2_start + ) + assert ( + time_overlap > 0 + ), "Operations on different sessions should run in parallel" + + # Clean up + Stagehand._session_locks.pop("test-parallel-session-1", None) + Stagehand._session_locks.pop("test-parallel-session-2", None)