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

optexity/inference/core/interaction/handle_select_utils.py

@@ -0,0 +1,132 @@
+import logging
+import re
+
+from pydantic import BaseModel
+
+from optexity.inference.agents.select_value_prediction.select_value_prediction import (
+    SelectValuePredictionAgent,
+)
+from optexity.schema.actions.interaction_action import Locator
+from optexity.schema.memory import Memory
+
+logger = logging.getLogger(__name__)
+select_value_prediction_agent = SelectValuePredictionAgent()
+
+
+class SelectOptionValue(BaseModel):
+    value: str
+    label: str
+
+
+def llm_select_match(
+    options: list[SelectOptionValue], patterns: list[str], memory: Memory
+) -> list[str]:
+    final_prompt, response, token_usage = (
+        select_value_prediction_agent.predict_select_value(
+            [o.model_dump() for o in options], patterns
+        )
+    )
+    memory.token_usage += token_usage
+    memory.browser_states[-1].final_prompt = final_prompt
+    memory.browser_states[-1].llm_response = response.model_dump()
+
+    matched_values = response.matched_values
+
+    all_values = [o.value for o in options]
+
+    final_matched_values = []
+    for value in matched_values:
+        if value in all_values:
+            final_matched_values.append(value)
+
+    return final_matched_values
+
+
+def score_match(pat: str, val: str) -> int:
+    # higher is better
+    if pat == val:
+        return 100
+    if val.startswith(pat):
+        return 80
+    if pat in val:
+        return 60
+    return 0
+
+
+async def smart_select(
+    options: list[SelectOptionValue], patterns: list[str], memory: Memory
+):
+    # Get all options from the <select>
+
+    matched_values = []
+
+    for p in patterns:
+        # If pattern contains regex characters, treat as regex
+        is_regex = p.startswith("^") or p.endswith("$") or ".*" in p
+
+        ## Check if reggex pattern and then try finding the option by value and label
+        if is_regex:
+            regex = re.compile(p)
+            for opt in options:
+                if regex.search(opt.value) or regex.search(opt.label):
+                    matched_values.append(opt.value)
+        else:
+            # try exact match
+            for opt in options:
+                if opt.value == p or opt.label == p:
+                    matched_values.append(opt.value)
+
+    if len(matched_values) == 0:
+        ## If no matches, check if all values are unique and try score matching of values
+
+        processed_values = [
+            (v.value.lower().replace(" ", ""), v.value) for v in options
+        ]
+
+        if len(processed_values) == len(set(processed_values)):
+            for p in patterns:
+                processed_pattern = p.lower().replace(" ", "")
+
+                best_score = 0
+                best_value = None
+
+                for processed_value, value in processed_values:
+                    score = score_match(processed_pattern, processed_value)
+                    if score > best_score:
+                        best_score = score
+                        best_value = value
+
+                if best_value is not None and best_score > 0:
+                    matched_values.append(best_value)
+
+    if len(matched_values) == 0:
+        processed_labels = [
+            (v.label.lower().replace(" ", ""), v.label) for v in options
+        ]
+
+        if len(processed_labels) == len(set(processed_labels)):
+            for p in patterns:
+                processed_pattern = p.lower().replace(" ", "")
+
+                best_score = 0
+                best_label = None
+                best_value = None
+
+                for opt in options:
+                    processed_label = opt.label.lower().replace(" ", "")
+                    score = score_match(processed_pattern, processed_label)
+                    if score > best_score:
+                        best_score = score
+                        best_label = opt.label
+                        best_value = opt.value
+
+                if best_label is not None and best_score > 0:
+                    matched_values.append(best_value)
+
+    if len(matched_values) == 0:
+        matched_values = llm_select_match(options, patterns, memory)
+
+    if len(matched_values) == 0:
+        matched_values = patterns
+
+    return matched_values

optexity/inference/core/interaction/handle_upload.py

@@ -0,0 +1,59 @@
+import logging
+
+from optexity.inference.core.interaction.handle_command import (
+    command_based_action_with_retry,
+)
+from optexity.inference.core.interaction.utils import get_index_from_prompt
+from optexity.inference.infra.browser import Browser
+from optexity.schema.actions.interaction_action import UploadFileAction
+from optexity.schema.memory import Memory
+from optexity.schema.task import Task
+
+logger = logging.getLogger(__name__)
+
+
+async def handle_upload_file(
+    upload_file_action: UploadFileAction,
+    task: Task,
+    memory: Memory,
+    browser: Browser,
+    max_timeout_seconds_per_try: float,
+    max_tries: int,
+):
+    if upload_file_action.command and not upload_file_action.skip_command:
+        last_error = await command_based_action_with_retry(
+            upload_file_action,
+            browser,
+            memory,
+            task,
+            max_tries,
+            max_timeout_seconds_per_try,
+        )
+        if last_error is None:
+            return
+
+    if not upload_file_action.skip_prompt:
+        logger.debug(
+            f"Executing prompt-based action: {upload_file_action.__class__.__name__}"
+        )
+        await upload_file_index(upload_file_action, browser, memory)
+
+
+async def upload_file_index(
+    upload_file_action: UploadFileAction, browser: Browser, memory: Memory
+):
+
+    try:
+        index = await get_index_from_prompt(
+            memory, upload_file_action.prompt_instructions, browser
+        )
+        if index is None:
+            return
+
+        action_model = browser.backend_agent.ActionModel(
+            **{"upload_file": {"index": index, "path": upload_file_action.file_path}}
+        )
+        await browser.backend_agent.multi_act([action_model])
+    except Exception as e:
+        logger.error(f"Error in upload_file_index: {e}")
+        return

optexity/inference/core/interaction/utils.py

@@ -0,0 +1,81 @@
+import logging
+from pathlib import Path
+from typing import Callable
+
+import aiofiles
+
+from optexity.inference.agents.index_prediction.action_prediction_locator_axtree import (
+    ActionPredictionLocatorAxtree,
+)
+from optexity.inference.infra.browser import Browser
+from optexity.schema.memory import BrowserState, Memory
+from optexity.schema.task import Task
+
+logger = logging.getLogger(__name__)
+
+
+index_prediction_agent = ActionPredictionLocatorAxtree()
+
+
+async def get_index_from_prompt(
+    memory: Memory, prompt_instructions: str, browser: Browser
+):
+    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(),
+    )
+
+    try:
+        final_prompt, response, token_usage = index_prediction_agent.predict_action(
+            prompt_instructions, memory.browser_states[-1].axtree
+        )
+        memory.token_usage += token_usage
+        memory.browser_states[-1].final_prompt = final_prompt
+        memory.browser_states[-1].llm_response = response.model_dump()
+
+        return response.index
+    except Exception as e:
+        logger.error(f"Error in get_index_from_prompt: {e}")
+
+
+async def handle_download(
+    func: Callable, memory: Memory, browser: Browser, task: Task, download_filename: str
+):
+    page = await browser.get_current_page()
+    if page is None:
+        logger.error("No page found for current page")
+        return
+    download_path: Path = task.downloads_directory / download_filename
+    async with page.expect_download() as download_info:
+        await func()
+        download = await download_info.value
+
+        if download:
+            temp_path = await download.path()
+            async with memory.download_lock:
+                memory.raw_downloads[temp_path] = (True, None)
+
+            await download.save_as(download_path)
+            memory.downloads.append(download_path)
+            await clean_download(download_path)
+        else:
+            logger.error("No download found")
+
+
+async def clean_download(download_path: Path):
+
+    if download_path.suffix == ".csv":
+        # Read full file
+        async with aiofiles.open(download_path, "r", encoding="utf-8") as f:
+            content = await f.read()
+        # Remove everything between <script>...</script> (multiline safe)
+
+        if "</script>" in content:
+            clean_content = content.split("</script>")[-1]
+
+            # Write cleaned CSV back
+            async with aiofiles.open(download_path, "w", encoding="utf-8") as f:
+                await f.write(clean_content)

optexity/inference/core/logging.py

@@ -0,0 +1,406 @@
+import base64
+import io
+import json
+import logging
+import shutil
+import tarfile
+from datetime import datetime, timezone
+from pathlib import Path
+from urllib.parse import urljoin
+
+import aiofiles
+import httpx
+
+from optexity.schema.automation import ActionNode
+from optexity.schema.memory import Memory
+from optexity.schema.task import Task
+from optexity.schema.token_usage import TokenUsage
+from optexity.utils.settings import settings
+from optexity.utils.utils import save_screenshot
+
+logger = logging.getLogger(__name__)
+
+
+def create_tar_in_memory(directory: Path | str, name: str) -> io.BytesIO:
+    if isinstance(directory, str):
+        directory = Path(directory)
+    tar_bytes = io.BytesIO()
+    with tarfile.open(fileobj=tar_bytes, mode="w:gz") as tar:
+        tar.add(directory, arcname=name)
+    tar_bytes.seek(0)  # rewind to start
+    return tar_bytes
+
+
+async def start_task_in_server(task: Task):
+    try:
+        task.started_at = datetime.now(timezone.utc)
+        task.status = "running"
+
+        url = urljoin(settings.SERVER_URL, settings.START_TASK_ENDPOINT)
+        headers = {"x-api-key": task.api_key}
+        body = {
+            "task_id": task.task_id,
+            "started_at": task.started_at.isoformat(),
+        }
+        if task.allocated_at:
+            body["allocated_at"] = task.allocated_at.isoformat()
+        async with httpx.AsyncClient(timeout=30.0) as client:
+            response = await client.post(
+                url,
+                headers=headers,
+                json=body,
+            )
+
+            response.raise_for_status()
+            return response.json()
+    except httpx.HTTPStatusError as e:
+        raise ValueError(
+            f"Failed to start task in server: {e.response.status_code} - {e.response.text}"
+        )
+    except Exception as e:
+        raise ValueError(f"Failed to start task in server: {e}")
+
+
+async def complete_task_in_server(
+    task: Task, token_usage: TokenUsage, child_process_id: int
+):
+    try:
+        task.completed_at = datetime.now(timezone.utc)
+
+        url = urljoin(settings.SERVER_URL, settings.COMPLETE_TASK_ENDPOINT)
+        headers = {"x-api-key": task.api_key}
+        body = {
+            "task_id": task.task_id,
+            "child_process_id": child_process_id,
+            "completed_at": task.completed_at.isoformat(),
+            "status": task.status,
+            "error": task.error,
+            "token_usage": token_usage.model_dump(),
+        }
+        async with httpx.AsyncClient(timeout=30.0) as client:
+            response = await client.post(
+                url,
+                headers=headers,
+                json=body,
+            )
+
+            response.raise_for_status()
+            return response.json()
+    except httpx.HTTPStatusError as e:
+        logger.error(
+            f"Failed to complete task in server: {e.response.status_code} - {e.response.text}"
+        )
+
+    except Exception as e:
+        logger.error(f"Failed to complete task in server: {e}")
+
+
+async def save_output_data_in_server(task: Task, memory: Memory):
+    try:
+        if len(memory.variables.output_data) == 0 and memory.final_screenshot is None:
+            return
+
+        url = urljoin(settings.SERVER_URL, settings.SAVE_OUTPUT_DATA_ENDPOINT)
+        headers = {"x-api-key": task.api_key}
+
+        output_data = [
+            output_data.model_dump(exclude_none=True, exclude={"screenshot"})
+            for output_data in memory.variables.output_data
+        ]
+        output_data = [data for data in output_data if data and len(data.keys()) > 0]
+        body = {
+            "task_id": task.task_id,
+            "output_data": output_data,
+            "final_screenshot": memory.final_screenshot,
+        }
+
+        for_loop_status = []
+        for loop_status in memory.variables.for_loop_status:
+            loop_status = [item.model_dump(exclude_none=True) for item in loop_status]
+            for_loop_status.append(loop_status)
+
+        if len(for_loop_status) > 0:
+            body["for_loop_status"] = for_loop_status
+
+        async with httpx.AsyncClient(timeout=30.0) as client:
+            response = await client.post(
+                url,
+                headers=headers,
+                json=body,
+            )
+
+            response.raise_for_status()
+            return response.json()
+    except httpx.HTTPStatusError as e:
+        logger.error(
+            f"Failed to save output data in server: {e.response.status_code} - {e.response.text}"
+        )
+    except Exception as e:
+        logger.error(f"Failed to save output data in server: {e}")
+
+
+async def save_downloads_in_server(task: Task, memory: Memory):
+    try:
+        # if len(memory.downloads) == 0:
+        #     return
+
+        url = urljoin(settings.SERVER_URL, settings.SAVE_DOWNLOADS_ENDPOINT)
+        headers = {"x-api-key": task.api_key}
+
+        payload = {
+            "task_id": task.task_id,  # form field
+        }
+
+        files = []
+        downloads = [
+            download
+            for download in task.downloads_directory.iterdir()
+            if download.is_file()
+        ]
+        if len(downloads) > 0:
+            tar_bytes = create_tar_in_memory(task.downloads_directory, task.task_id)
+            # add tar.gz
+            files.append(
+                (
+                    "compressed_downloads",
+                    (f"{task.task_id}.tar.gz", tar_bytes, "application/gzip"),
+                )
+            )
+
+        # add screenshots
+        for data in memory.variables.output_data:
+            if data.screenshot:
+                files.append(
+                    (
+                        "screenshots",
+                        (
+                            data.screenshot.filename,
+                            base64.b64decode(data.screenshot.base64),
+                            "image/png",
+                        ),
+                    )
+                )
+
+        if memory.final_screenshot:
+            files.append(
+                (
+                    "screenshots",
+                    (
+                        "final_screenshot.png",
+                        base64.b64decode(memory.final_screenshot),
+                        "image/png",
+                    ),
+                )
+            )
+
+        if len(files) == 0:
+            return
+
+        async with httpx.AsyncClient(timeout=30.0) as client:
+
+            response = await client.post(
+                url, headers=headers, data=payload, files=files
+            )
+
+            response.raise_for_status()
+            return response.json()
+    except httpx.HTTPStatusError as e:
+        logger.error(
+            f"Failed to save downloads in server: {e.response.status_code} - {e.response.text}"
+        )
+    except Exception as e:
+        logger.error(f"Failed to save downloads in server: {e}")
+
+
+async def save_trajectory_in_server(task: Task, memory: Memory):
+    try:
+        url = urljoin(settings.SERVER_URL, settings.SAVE_TRAJECTORY_ENDPOINT)
+        headers = {"x-api-key": task.api_key}
+
+        data = {
+            "task_id": task.task_id,  # form field
+        }
+
+        tar_bytes = create_tar_in_memory(task.task_directory, task.task_id)
+        files = {
+            "compressed_trajectory": (
+                f"{task.task_id}.tar.gz",
+                tar_bytes,
+                "application/gzip",
+            )
+        }
+        async with httpx.AsyncClient(timeout=30.0) as client:
+
+            response = await client.post(url, headers=headers, data=data, files=files)
+
+            response.raise_for_status()
+            return response.json()
+    except httpx.HTTPStatusError as e:
+        logger.error(
+            f"Failed to save trajectory in server: {e.response.status_code} - {e.response.text}"
+        )
+    except Exception as e:
+        logger.error(f"Failed to save trajectory in server: {e}")
+
+
+async def initiate_callback(task: Task):
+
+    if settings.DEPLOYMENT == "dev" and settings.LOCAL_CALLBACK_URL is not None:
+        logger.info("initiating local callback")
+        callback_data = None
+        try:
+            url = urljoin(settings.SERVER_URL, settings.GET_CALLBACK_DATA_ENDPOINT)
+            headers = {"x-api-key": task.api_key}
+            data = {
+                "task_id": task.task_id,
+                "endpoint_name": task.endpoint_name,
+            }
+            async with httpx.AsyncClient(timeout=30.0) as client:
+                response = await client.post(url, headers=headers, json=data)
+                response.raise_for_status()
+                callback_data = response.json()["data"]
+        except Exception as e:
+            logger.error(f"Failed to get callback data: {e}")
+            return
+
+        if callback_data is None:
+            return
+
+        try:
+            async with httpx.AsyncClient(timeout=30.0) as client:
+                response = await client.post(
+                    settings.LOCAL_CALLBACK_URL, json=callback_data
+                )
+                response.raise_for_status()
+        except Exception as e:
+            logger.error(f"Failed to initiate local callback: {e}")
+            return
+
+        return
+
+    try:
+        logger.info("initiating callback")
+        if task.callback_url is None:
+            return
+
+        url = urljoin(settings.SERVER_URL, settings.INITIATE_CALLBACK_ENDPOINT)
+        headers = {"x-api-key": task.api_key}
+
+        data = {
+            "task_id": task.task_id,
+            "endpoint_name": task.endpoint_name,
+            "callback_url": task.callback_url.model_dump(),
+        }
+
+        async with httpx.AsyncClient(timeout=30.0) as client:
+
+            response = await client.post(url, headers=headers, json=data)
+
+            response.raise_for_status()
+            return response.json()
+    except httpx.HTTPStatusError as e:
+        logger.error(
+            f"Failed to save trajectory in server: {e.response.status_code} - {e.response.text}"
+        )
+    except Exception as e:
+        logger.error(f"Failed to save trajectory in server: {e}")
+
+
+async def save_latest_memory_state_locally(
+    task: Task, memory: Memory, node: ActionNode | None
+):
+
+    try:
+        browser_state = memory.browser_states[-1]
+        automation_state = memory.automation_state
+        step_directory = (
+            task.logs_directory / f"step_{str(automation_state.step_index)}"
+        )
+        step_directory.mkdir(parents=True, exist_ok=True)
+
+        if browser_state.screenshot:
+            save_screenshot(browser_state.screenshot, step_directory / "screenshot.png")
+        else:
+            logger.warning(
+                "No screenshot found for step %s", automation_state.step_index
+            )
+
+        state_dict = {
+            "title": browser_state.title,
+            "url": browser_state.url,
+            "step_index": automation_state.step_index,
+            "try_index": automation_state.try_index,
+            "downloaded_files": [
+                downloaded_file.name for downloaded_file in memory.downloads
+            ],
+            "token_usage": memory.token_usage.model_dump(),
+        }
+
+        async with aiofiles.open(step_directory / "state.json", "w") as f:
+            await f.write(json.dumps(state_dict, indent=4))
+
+        if browser_state.axtree:
+            async with aiofiles.open(step_directory / "axtree.txt", "w") as f:
+                await f.write(browser_state.axtree)
+
+        if browser_state.final_prompt:
+            async with aiofiles.open(step_directory / "final_prompt.txt", "w") as f:
+                await f.write(browser_state.final_prompt)
+
+        if browser_state.llm_response:
+            async with aiofiles.open(step_directory / "llm_response.json", "w") as f:
+                await f.write(json.dumps(browser_state.llm_response, indent=4))
+
+        if node:
+            async with aiofiles.open(step_directory / "action_node.json", "w") as f:
+                await f.write(
+                    json.dumps(
+                        node.model_dump(exclude_none=True, exclude_defaults=True),
+                        indent=4,
+                    )
+                )
+
+        async with aiofiles.open(step_directory / "input_parameters.json", "w") as f:
+            await f.write(json.dumps(task.input_parameters, indent=4))
+
+        async with aiofiles.open(step_directory / "secure_parameters.json", "w") as f:
+            await f.write(json.dumps(task.secure_parameters, indent=4))
+
+        async with aiofiles.open(step_directory / "generated_variables.json", "w") as f:
+            await f.write(json.dumps(memory.variables.generated_variables, indent=4))
+
+        async with aiofiles.open(step_directory / "output_data.json", "w") as f:
+            await f.write(
+                json.dumps(
+                    [
+                        output_data.model_dump(
+                            exclude_none=True,
+                            exclude={"screenshot"},
+                            exclude_defaults=True,
+                        )
+                        for output_data in memory.variables.output_data
+                    ],
+                    indent=4,
+                )
+            )
+
+        for output_data in memory.variables.output_data:
+            if output_data.screenshot:
+                async with aiofiles.open(
+                    step_directory
+                    / f"screenshot_{output_data.screenshot.filename}.png",
+                    "wb",
+                ) as f:
+                    await f.write(base64.b64decode(output_data.screenshot.base64))
+    except Exception as e:
+        logger.error(f"Failed to save latest memory state locally: {e}")
+
+
+async def delete_local_data(task: Task):
+    try:
+        if settings.DEPLOYMENT == "dev" or task.task_directory is None:
+            return
+
+        shutil.rmtree(task.task_directory, ignore_errors=True)
+    except Exception as e:
+        logger.error(f"Failed to delete local data: {e}")