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

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}")

optexity/inference/core/run_assertion.py

@@ -0,0 +1,55 @@
+import logging
+from copy import deepcopy
+
+from optexity.inference.core.run_extraction import handle_llm_extraction
+from optexity.inference.infra.browser import Browser
+from optexity.inference.models import GeminiModels, get_llm_model
+from optexity.schema.actions.assertion_action import AssertionAction, LLMAssertion
+from optexity.schema.memory import Memory
+
+logger = logging.getLogger(__name__)
+
+llm_model = get_llm_model(GeminiModels.GEMINI_2_5_FLASH, True)
+
+
+async def run_assertion_action(
+    assertion_action: AssertionAction, memory: Memory, browser: Browser
+):
+    logger.debug(
+        f"---------Running assertion action {assertion_action.model_dump_json()}---------"
+    )
+
+    if assertion_action.llm:
+        await handle_llm_assertion(assertion_action.llm, memory, browser)
+    elif assertion_action.network_call:
+        raise ValueError("Network call assertions are not supported yet")
+        # await handle_network_call_assertion(
+        #     assertion_action.network_call, memory, browser
+        # )
+    elif assertion_action.python_script:
+        raise ValueError("Python script assertions are not supported yet")
+        # await handle_python_script_assertion(
+        #     assertion_action.python_script, memory, browser
+        # )
+
+
+async def handle_llm_assertion(
+    llm_assertion: LLMAssertion, memory: Memory, browser: Browser
+):
+    extra_instruction = """You are a helpful assistant that verifies if the condition is met. 
+        Use the info supplied below to verify the condition.
+        The assertion_reason should be a short explanation of why the condition was met or not met.
+        The assertion_result should be True if the condition is met, False otherwise.
+        """
+    llm_assertion_new = deepcopy(llm_assertion)
+    llm_assertion_new.extraction_instructions = (
+        extra_instruction + "\n" + llm_assertion_new.extraction_instructions
+    )
+    output_data = await handle_llm_extraction(llm_assertion_new, memory, browser)
+
+    if output_data.json_data["assertion_result"]:
+        return True
+    else:
+        raise AssertionError(
+            f"Assertion failed on node {memory.automation_state.step_index}: {output_data.json_data['assertion_reason']}"
+        )