optexity 0.1.2__py3-none-any.whl → 0.1.3__py3-none-any.whl

optexity/inference/core/run_extraction.py

@@ -0,0 +1,240 @@
+import logging
+import traceback
+
+import aiofiles
+import httpx
+
+from optexity.inference.infra.browser import Browser
+from optexity.inference.models import GeminiModels, get_llm_model
+from optexity.schema.actions.extraction_action import (
+    ExtractionAction,
+    LLMExtraction,
+    NetworkCallExtraction,
+    ScreenshotExtraction,
+    StateExtraction,
+)
+from optexity.schema.memory import (
+    BrowserState,
+    Memory,
+    NetworkRequest,
+    NetworkResponse,
+    OutputData,
+    ScreenshotData,
+)
+from optexity.schema.task import Task
+
+logger = logging.getLogger(__name__)
+
+llm_model = get_llm_model(GeminiModels.GEMINI_2_5_FLASH, True)
+
+
+async def run_extraction_action(
+    extraction_action: ExtractionAction, memory: Memory, browser: Browser, task: Task
+):
+    logger.debug(
+        f"---------Running extraction action {extraction_action.model_dump_json()}---------"
+    )
+
+    if extraction_action.llm:
+        await handle_llm_extraction(
+            extraction_action.llm, memory, browser, extraction_action.unique_identifier
+        )
+    elif extraction_action.network_call:
+        await handle_network_call_extraction(
+            extraction_action.network_call,
+            memory,
+            browser,
+            task,
+            extraction_action.unique_identifier,
+        )
+    elif extraction_action.screenshot:
+        await handle_screenshot_extraction(
+            extraction_action.screenshot,
+            memory,
+            browser,
+            extraction_action.unique_identifier,
+        )
+    elif extraction_action.state:
+        await handle_state_extraction(
+            extraction_action.state,
+            memory,
+            browser,
+            extraction_action.unique_identifier,
+        )
+
+
+async def handle_state_extraction(
+    state_extraction: StateExtraction,
+    memory: Memory,
+    browser: Browser,
+    unique_identifier: str | None = None,
+):
+    page = await browser.get_current_page()
+    if page is None:
+        return
+
+    memory.variables.output_data.append(
+        OutputData(
+            unique_identifier=unique_identifier,
+            json_data={"page_url": page.url, "page_title": await page.title()},
+        )
+    )
+
+
+async def handle_screenshot_extraction(
+    screenshot_extraction: ScreenshotExtraction,
+    memory: Memory,
+    browser: Browser,
+    unique_identifier: str | None = None,
+):
+
+    screenshot_base64 = await browser.get_screenshot(
+        full_page=screenshot_extraction.full_page
+    )
+    if screenshot_base64 is None:
+        return
+
+    memory.variables.output_data.append(
+        OutputData(
+            unique_identifier=unique_identifier,
+            screenshot=ScreenshotData(
+                filename=screenshot_extraction.filename, base64=screenshot_base64
+            ),
+        )
+    )
+
+
+async def handle_llm_extraction(
+    llm_extraction: LLMExtraction,
+    memory: Memory,
+    browser: Browser,
+    unique_identifier: str | None = None,
+):
+    browser_state_summary = await browser.get_browser_state_summary()
+    memory.browser_states[-1] = BrowserState(
+        url=browser_state_summary.url,
+        screenshot=browser_state_summary.screenshot,
+        title=browser_state_summary.title,
+        axtree=browser_state_summary.dom_state.llm_representation(),
+    )
+
+    # TODO: fix this double calling of screenshot and axtree
+    if "axtree" in llm_extraction.source:
+        axtree = memory.browser_states[-1].axtree
+    else:
+        axtree = None
+
+    if "screenshot" in llm_extraction.source:
+        screenshot = memory.browser_states[-1].screenshot
+    else:
+        screenshot = None
+
+    system_instruction = f"""
+    You are an expert in extracting information from a website. You will be given an axtree of a webpage.
+    Your task is to extract the information from the webpage and return it in the format specified by the instructions.
+    {llm_extraction.extraction_instructions}
+    """
+
+    prompt = f"""
+    [INPUT]
+    Axtree: {axtree}
+    [/INPUT]
+    """
+
+    if llm_extraction.llm_provider == "gemini":
+        model_name = GeminiModels(llm_extraction.llm_model_name)
+        llm_model.model_name = model_name
+    else:
+        raise ValueError(f"Invalid LLM provider: {llm_extraction.llm_provider}")
+
+    response, token_usage = llm_model.get_model_response_with_structured_output(
+        prompt=prompt,
+        response_schema=llm_extraction.build_model(),
+        screenshot=screenshot,
+        system_instruction=system_instruction,
+    )
+    response_dict = response.model_dump()
+    output_data = OutputData(
+        unique_identifier=unique_identifier, json_data=response_dict
+    )
+
+    logger.debug(f"Response: {response_dict}")
+
+    memory.token_usage += token_usage
+    memory.variables.output_data.append(output_data)
+
+    if llm_extraction.output_variable_names is not None:
+        for output_variable_name in llm_extraction.output_variable_names:
+            v = response_dict[output_variable_name]
+            if isinstance(v, list):
+                memory.variables.generated_variables[output_variable_name] = v
+            elif (
+                isinstance(v, str)
+                or isinstance(v, int)
+                or isinstance(v, float)
+                or isinstance(v, bool)
+            ):
+                memory.variables.generated_variables[output_variable_name] = [v]
+            else:
+                raise ValueError(
+                    f"Output variable {output_variable_name} must be a string, int, float, bool, or a list of strings, ints, floats, or bools. Extracted values: {response_dict[output_variable_name]}"
+                )
+    return output_data
+
+
+async def handle_network_call_extraction(
+    network_call_extraction: NetworkCallExtraction,
+    memory: Memory,
+    browser: Browser,
+    task: Task,
+    unique_identifier: str | None = None,
+):
+
+    for network_call in browser.network_calls:
+        if network_call_extraction.url_pattern not in network_call.url:
+            continue
+
+        if network_call_extraction.download_from == "request" and isinstance(
+            network_call, NetworkRequest
+        ):
+            await download_request(
+                network_call, network_call_extraction.download_filename, task, memory
+            )
+
+        if (
+            network_call_extraction.extract_from == "request"
+            and isinstance(network_call, NetworkRequest)
+        ) or (
+            network_call_extraction.extract_from == "response"
+            and isinstance(network_call, NetworkResponse)
+        ):
+            memory.variables.output_data.append(
+                OutputData(
+                    unique_identifier=unique_identifier,
+                    json_data=network_call.model_dump(include={"body"}),
+                )
+            )
+
+
+async def download_request(
+    network_call: NetworkRequest, download_filename: str, task: Task, memory: Memory
+):
+    try:
+        async with httpx.AsyncClient(follow_redirects=True) as client:
+            response = await client.request(
+                network_call.method,
+                network_call.url,
+                headers=network_call.headers,
+                content=network_call.body,  # not data=
+            )
+
+            response.raise_for_status()
+
+        # Save raw response to PDF
+        download_path = task.downloads_directory / download_filename
+        async with aiofiles.open(download_path, "wb") as f:
+            await f.write(response.content)
+
+        memory.downloads.append(download_path)
+    except Exception as e:
+        logger.error(f"Failed to download request: {e}, {traceback.format_exc()}")

optexity/inference/core/run_interaction.py

@@ -0,0 +1,254 @@
+import asyncio
+import logging
+from datetime import datetime, timezone
+
+import aiofiles
+
+from optexity.exceptions import AssertLocatorPresenceException
+from optexity.inference.agents.error_handler.error_handler import ErrorHandlerAgent
+from optexity.inference.core.interaction.handle_agentic_task import handle_agentic_task
+from optexity.inference.core.interaction.handle_check import (
+    handle_check_element,
+    handle_uncheck_element,
+)
+from optexity.inference.core.interaction.handle_click import handle_click_element
+from optexity.inference.core.interaction.handle_input import handle_input_text
+from optexity.inference.core.interaction.handle_keypress import handle_key_press
+from optexity.inference.core.interaction.handle_select import handle_select_option
+from optexity.inference.core.interaction.handle_upload import handle_upload_file
+from optexity.inference.infra.browser import Browser
+from optexity.schema.actions.interaction_action import (
+    CloseOverlayPopupAction,
+    CloseTabsUntil,
+    DownloadUrlAsPdfAction,
+    GoBackAction,
+    GoToUrlAction,
+    InteractionAction,
+)
+from optexity.schema.memory import BrowserState, Memory, OutputData
+from optexity.schema.task import Task
+
+error_handler_agent = ErrorHandlerAgent()
+
+
+logger = logging.getLogger(__name__)
+
+
+async def run_interaction_action(
+    interaction_action: InteractionAction,
+    task: Task,
+    memory: Memory,
+    browser: Browser,
+    retries_left: int,
+):
+    if retries_left <= 0:
+        return
+
+    logger.debug(
+        f"---------Running interaction action {interaction_action.model_dump_json(exclude_none=True)}---------"
+    )
+
+    try:
+        memory.automation_state.start_2fa_time = datetime.now(timezone.utc)
+        if interaction_action.click_element:
+            await handle_click_element(
+                interaction_action.click_element,
+                task,
+                memory,
+                browser,
+                interaction_action.max_timeout_seconds_per_try,
+                interaction_action.max_tries,
+            )
+        elif interaction_action.input_text:
+            await handle_input_text(
+                interaction_action.input_text,
+                task,
+                memory,
+                browser,
+                interaction_action.max_timeout_seconds_per_try,
+                interaction_action.max_tries,
+            )
+        elif interaction_action.select_option:
+            await handle_select_option(
+                interaction_action.select_option,
+                task,
+                memory,
+                browser,
+                interaction_action.max_timeout_seconds_per_try,
+                interaction_action.max_tries,
+            )
+        elif interaction_action.check:
+            await handle_check_element(
+                interaction_action.check,
+                task,
+                memory,
+                browser,
+                interaction_action.max_timeout_seconds_per_try,
+                interaction_action.max_tries,
+            )
+        elif interaction_action.uncheck:
+            await handle_uncheck_element(
+                interaction_action.uncheck,
+                task,
+                memory,
+                browser,
+                interaction_action.max_timeout_seconds_per_try,
+                interaction_action.max_tries,
+            )
+        elif interaction_action.go_back:
+            await handle_go_back(interaction_action.go_back, memory, browser)
+        elif interaction_action.download_url_as_pdf:
+            await handle_download_url_as_pdf(
+                interaction_action.download_url_as_pdf, task, memory, browser
+            )
+        elif interaction_action.agentic_task:
+            await handle_agentic_task(
+                interaction_action.agentic_task, task, memory, browser
+            )
+        elif interaction_action.close_overlay_popup:
+            await handle_agentic_task(
+                interaction_action.close_overlay_popup, task, memory, browser
+            )
+        elif interaction_action.go_to_url:
+            await handle_go_to_url(interaction_action.go_to_url, task, memory, browser)
+        elif interaction_action.upload_file:
+            await handle_upload_file(
+                interaction_action.upload_file,
+                task,
+                memory,
+                browser,
+                interaction_action.max_timeout_seconds_per_try,
+                interaction_action.max_tries,
+            )
+        elif interaction_action.close_current_tab:
+            await browser.close_current_tab()
+        elif interaction_action.switch_tab:
+            await browser.switch_tab(interaction_action.switch_tab.tab_index)
+        elif interaction_action.close_tabs_until:
+            await handle_close_tabs_until(
+                interaction_action.close_tabs_until, task, memory, browser
+            )
+        elif interaction_action.key_press:
+            await handle_key_press(interaction_action.key_press, memory, browser)
+    except AssertLocatorPresenceException as e:
+        await handle_assert_locator_presence_error(
+            e, interaction_action, task, memory, browser, retries_left
+        )
+
+
+async def handle_close_tabs_until(
+    close_tabs_until_action: CloseTabsUntil,
+    task: Task,
+    memory: Memory,
+    browser: Browser,
+):
+
+    while True:
+        page = await browser.get_current_page()
+        if page is None:
+            return
+
+        if close_tabs_until_action.matching_url is not None:
+            if close_tabs_until_action.matching_url in page.url:
+                break
+        elif close_tabs_until_action.tab_index is not None:
+            if len(browser.context.pages) == close_tabs_until_action.tab_index + 1:
+                break
+
+        await browser.close_current_tab()
+
+
+async def handle_go_to_url(
+    go_to_url_action: GoToUrlAction, task: Task, memory: Memory, browser: Browser
+):
+    await browser.go_to_url(go_to_url_action.url)
+
+
+async def handle_go_back(
+    go_back_action: GoBackAction, memory: Memory, browser: Browser
+):
+    page = await browser.get_current_page()
+    if page is None:
+        return
+    await page.go_back()
+
+
+async def handle_download_url_as_pdf(
+    download_url_as_pdf_action: DownloadUrlAsPdfAction,
+    task: Task,
+    memory: Memory,
+    browser: Browser,
+):
+    if download_url_as_pdf_action.url is not None:
+        pdf_url = download_url_as_pdf_action.url
+    else:
+        pdf_url = await browser.get_current_page_url()
+
+    if pdf_url is None:
+        logger.error("No PDF URL found for current page")
+        return
+    download_path = (
+        task.downloads_directory / download_url_as_pdf_action.download_filename
+    )
+
+    resp = await browser.context.request.get(pdf_url)
+
+    if not resp.ok:
+        logger.error(f"Failed to download PDF: {resp.status}")
+        return
+
+    content = await resp.body()
+    async with aiofiles.open(download_path, "wb") as f:
+        await f.write(content)
+
+    memory.downloads.append(download_path)
+
+
+async def handle_assert_locator_presence_error(
+    error: AssertLocatorPresenceException,
+    interaction_action: InteractionAction,
+    task: Task,
+    memory: Memory,
+    browser: Browser,
+    retries_left: int,
+):
+    logger.debug(f"Handling assert locator presence error: {error.command}")
+    if retries_left > 1:
+        browser_state_summary = await browser.get_browser_state_summary()
+        memory.browser_states[-1] = BrowserState(
+            url=browser_state_summary.url,
+            screenshot=browser_state_summary.screenshot,
+            title=browser_state_summary.title,
+            axtree=browser_state_summary.dom_state.llm_representation(),
+        )
+        final_prompt, response, token_usage = error_handler_agent.classify_error(
+            error.command, memory.browser_states[-1].screenshot
+        )
+        memory.token_usage += token_usage
+
+        if response.error_type == "website_not_loaded":
+            await asyncio.sleep(5)
+            await run_interaction_action(
+                interaction_action, task, memory, browser, retries_left - 1
+            )
+        elif response.error_type == "overlay_popup_blocking":
+            close_overlay_popup_action = CloseOverlayPopupAction()
+            await handle_agentic_task(close_overlay_popup_action, task, memory, browser)
+            await run_interaction_action(
+                interaction_action, task, memory, browser, retries_left - 1
+            )
+        elif response.error_type == "fatal_error":
+            logger.error(
+                f"Fatal error running node {memory.automation_state.step_index} after {retries_left} retries: {error.original_error}. Error: {response.detailed_reason}"
+            )
+            memory.variables.output_data.append(
+                OutputData(unique_identifier="error", text=response.detailed_reason)
+            )
+            raise Exception(
+                f"Fatal error running node {memory.automation_state.step_index} after {retries_left} retries: {error.original_error}. Final reason: {response.detailed_reason}"
+            )
+    else:
+        logger.error(
+            f"Error running node {memory.automation_state.step_index} after {retries_left} retries: {error.original_error}"
+        )
+        raise error

optexity/inference/core/run_python_script.py

@@ -0,0 +1,20 @@
+import logging
+
+from optexity.inference.infra.browser import Browser
+from optexity.schema.actions.misc_action import PythonScriptAction
+from optexity.schema.memory import Memory
+
+logger = logging.getLogger(__name__)
+
+
+async def run_python_script_action(
+    python_script_action: PythonScriptAction, memory: Memory, browser: Browser
+):
+    local_vars = {}
+    exec(python_script_action.execution_code, {}, local_vars)
+
+    # Get the function
+    code_fn = local_vars["code_fn"]
+
+    page = await browser.get_current_page()
+    await code_fn(page)

optexity/inference/core/run_two_fa.py

@@ -0,0 +1,120 @@
+import asyncio
+import logging
+from datetime import timedelta
+from urllib.parse import urljoin
+
+import httpx
+
+from optexity.inference.agents.two_fa_extraction.two_fa_extraction import (
+    TwoFAExtraction,
+)
+from optexity.schema.actions.two_fa_action import (
+    EmailTwoFAAction,
+    SlackTwoFAAction,
+    TwoFAAction,
+)
+from optexity.schema.inference import (
+    FetchEmailMessagesRequest,
+    FetchMessagesResponse,
+    FetchSlackMessagesRequest,
+)
+from optexity.schema.memory import Memory
+from optexity.utils.settings import settings
+
+logger = logging.getLogger(__name__)
+
+two_fa_extraction_agent = TwoFAExtraction()
+
+
+async def run_two_fa_action(two_fa_action: TwoFAAction, memory: Memory):
+    logger.debug(
+        f"---------Running 2fa action {two_fa_action.model_dump_json()}---------"
+    )
+
+    elapsed = 0
+    messages = None
+
+    while elapsed < two_fa_action.max_wait_time:
+        messages = await fetch_messages(
+            two_fa_action.action, memory, two_fa_action.max_wait_time
+        )
+        if messages and len(messages) > 0:
+            final_prompt, response, token_usage = two_fa_extraction_agent.extract_code(
+                two_fa_action.instructions, messages
+            )
+            memory.token_usage += token_usage
+            code = None
+            if response.code is not None:
+                if isinstance(response.code, str):
+                    code = response.code
+                elif isinstance(response.code, list):
+                    if len(response.code) > 0:
+                        raise ValueError(f"Multiple 2FA codes found, {response.code}")
+                    else:
+                        code = response.code[0]
+
+            if code is not None:
+                logger.debug(
+                    f"2FA code {code} found after {elapsed} seconds from {messages}"
+                )
+                break
+            logger.debug(
+                f"No 2FA code found in messages, {messages}, waiting for {two_fa_action.check_interval} seconds"
+            )
+        else:
+            logger.debug(
+                f"No messages found for 2FA code after {elapsed} seconds, waiting for {two_fa_action.check_interval} seconds"
+            )
+
+        await asyncio.sleep(two_fa_action.check_interval)
+        elapsed += two_fa_action.check_interval
+
+    memory.automation_state.start_2fa_time = None
+    if code is None:
+        raise ValueError("2FA code not found")
+
+    memory.variables.generated_variables[two_fa_action.output_variable_name] = [code]
+
+    return code
+
+
+async def fetch_messages(
+    action: EmailTwoFAAction | SlackTwoFAAction,
+    memory: Memory,
+    max_wait_time: float,
+):
+
+    start_2fa_time = memory.automation_state.start_2fa_time
+    end_2fa_time = memory.automation_state.start_2fa_time + timedelta(
+        seconds=max_wait_time
+    )
+
+    headers = {"x-api-key": settings.API_KEY}
+
+    if isinstance(action, EmailTwoFAAction):
+        url = urljoin(settings.SERVER_URL, settings.FETCH_EMAIL_MESSAGES_ENDPOINT)
+        body = FetchEmailMessagesRequest(
+            receiver_email_address=action.receiver_email_address,
+            sender_email_address=action.sender_email_address,
+            start_2fa_time=start_2fa_time,
+            end_2fa_time=end_2fa_time,
+        )
+    elif isinstance(action, SlackTwoFAAction):
+        url = urljoin(settings.SERVER_URL, settings.FETCH_SLACK_MESSAGES_ENDPOINT)
+        body = FetchSlackMessagesRequest(
+            slack_workspace_domain=action.slack_workspace_domain,
+            channel_name=action.channel_name,
+            sender_name=action.sender_name,
+            start_2fa_time=start_2fa_time,
+            end_2fa_time=end_2fa_time,
+        )
+
+    async with httpx.AsyncClient(timeout=30.0) as client:
+
+        response = await client.post(
+            url, json=body.model_dump(mode="json"), headers=headers
+        )
+        response.raise_for_status()
+        response_data = FetchMessagesResponse.model_validate(response.json())
+
+        return response_data.messages