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

optexity/schema/automation.py

@@ -0,0 +1,432 @@
+import logging
+from typing import Annotated, Any, ForwardRef, Literal
+
+from pydantic import BaseModel, Field, model_validator
+
+from optexity.schema.actions.assertion_action import AssertionAction
+from optexity.schema.actions.extraction_action import ExtractionAction
+from optexity.schema.actions.interaction_action import InteractionAction
+from optexity.schema.actions.misc_action import PythonScriptAction
+from optexity.schema.actions.two_fa_action import TwoFAAction
+from optexity.utils.utils import get_onepassword_value, get_totp_code
+
+logger = logging.getLogger(__name__)
+
+IfElseNodeRef = ForwardRef("IfElseNode")
+ForLoopNodeRef = ForwardRef("ForLoopNode")
+
+
+class OnePasswordParameter(BaseModel):
+    vault_name: str
+    item_name: str
+    field_name: str
+    type: Literal["raw", "totp_secret"] = "raw"
+    digits: int | None = None
+
+    @model_validator(mode="after")
+    def validate_onepassword_parameter(self):
+        if self.type == "totp_secret":
+            assert self.digits is not None, "digits must be provided for totp_secret"
+        else:
+            assert self.digits is None, "digits must not be provided for raw"
+        return self
+
+
+class AmazonSecretsManagerParameter(BaseModel):
+    pass
+
+    @model_validator(mode="after")
+    def validate_amazon_secrets_manager_parameter(
+        cls, model: "AmazonSecretsManagerParameter"
+    ):
+        raise NotImplementedError("Amazon Secrets Manager is not implemented yet")
+
+
+class TOTPParameter(BaseModel):
+    totp_secret: str
+    digits: int = 6
+
+
+class SecureParameter(BaseModel):
+    onepassword: OnePasswordParameter | None = None
+    amazon_secrets_manager: AmazonSecretsManagerParameter | None = None
+    totp: TOTPParameter | None = None
+
+    @model_validator(mode="after")
+    def validate_secure_parameter(self):
+        non_null = [k for k, v in self.model_dump().items() if v is not None]
+        if len(non_null) != 1:
+            raise ValueError(
+                "Exactly one of onepassword or amazon_secrets_manager or totp must be provided"
+            )
+        return self
+
+
+class ActionNode(BaseModel):
+    type: Literal["action_node"]
+    interaction_action: InteractionAction | None = None
+    assertion_action: AssertionAction | None = None
+    extraction_action: ExtractionAction | None = None
+    python_script_action: PythonScriptAction | None = None
+    two_fa_action: TwoFAAction | None = None
+    before_sleep_time: float = 0.0
+    end_sleep_time: float = 5.0
+    expect_new_tab: bool = False
+    max_new_tab_wait_time: float = 0.0
+    localized_axtree_string: str | None = None
+
+    @model_validator(mode="after")
+    def validate_one_node(cls, model: "ActionNode"):
+        """Ensure exactly one of the node types is set and matches the type."""
+        provided = {
+            "interaction_action": model.interaction_action,
+            "assertion_action": model.assertion_action,
+            "extraction_action": model.extraction_action,
+            "python_script_action": model.python_script_action,
+            "two_fa_action": model.two_fa_action,
+        }
+        non_null = [k for k, v in provided.items() if v is not None]
+
+        if len(non_null) != 1:
+            raise ValueError(
+                "Exactly one of interaction_action, assertion_action, extraction_action, python_script_action, or two_fa_action must be provided"
+            )
+
+        assert (
+            model.end_sleep_time >= 0 and model.end_sleep_time <= 30
+        ), "end_sleep_time must be greater than 0 and less than 30"
+        assert (
+            model.max_new_tab_wait_time >= 0 and model.max_new_tab_wait_time <= 30
+        ), "max_new_tab_wait_time must be greater than 0 and less than 30"
+
+        # --- Adjust defaults only if user didn't override them ---
+        # We detect user-provided fields using model.__pydantic_fields_set__
+        user_set = model.__pydantic_fields_set__
+
+        if "end_sleep_time" not in user_set:
+            if model.assertion_action or model.extraction_action or model.two_fa_action:
+                model.end_sleep_time = 0.0
+
+        if "before_sleep_time" not in user_set:
+            model.before_sleep_time = 3.0 if model.extraction_action else 0.0
+
+        if model.expect_new_tab:
+            assert (
+                model.interaction_action is not None
+            ), "expect_new_tab is only allowed for interaction actions"
+            model.max_new_tab_wait_time = 10.0
+        else:
+            model.max_new_tab_wait_time = 0.0
+
+        return model
+
+    def replace(self, pattern: str, replacement: str | int | float | bool | None):
+        replacement = str(replacement)
+        if self.interaction_action:
+            self.interaction_action.replace(pattern, replacement)
+        if self.assertion_action:
+            self.assertion_action.replace(pattern, replacement)
+        if self.extraction_action:
+            self.extraction_action.replace(pattern, replacement)
+        if self.python_script_action:
+            pass
+        if self.two_fa_action:
+            pass
+
+        return self
+
+    async def replace_variables(
+        self, variables: dict[str, list[str | SecureParameter]]
+    ):
+        for key, values in variables.items():
+
+            for index, value in enumerate(values):
+                pattern = f"{{{key}[{index}]}}"
+
+                if isinstance(value, SecureParameter):
+                    if value.onepassword:
+                        str_value = await get_onepassword_value(
+                            value.onepassword.vault_name,
+                            value.onepassword.item_name,
+                            value.onepassword.field_name,
+                        )
+                        if value.onepassword.type == "totp_secret":
+                            str_value = get_totp_code(
+                                str_value, value.onepassword.digits
+                            )
+
+                    elif value.amazon_secrets_manager:
+                        raise NotImplementedError(
+                            "Amazon Secrets Manager is not implemented yet"
+                        )
+                    elif value.totp:
+                        str_value = get_totp_code(
+                            value.totp.totp_secret, value.totp.digits
+                        )
+
+                elif (
+                    isinstance(value, str)
+                    or isinstance(value, int)
+                    or isinstance(value, float)
+                    or isinstance(value, bool)
+                ):
+                    str_value = str(value)
+                else:
+                    raise ValueError(f"Invalid value type for {key}: {type(value)}")
+
+                self.replace(pattern, str_value)
+
+        return self
+
+
+class ForLoopNode(BaseModel):
+    # Loops through range of values of {variable_name[index]}
+    type: Literal["for_loop_node"]
+    variable_name: str
+    nodes: list[Annotated[ActionNode | IfElseNodeRef, Field(discriminator="type")]]
+    reset_nodes: list[
+        Annotated[ActionNode | IfElseNodeRef, Field(discriminator="type")]
+    ] = []
+    on_error_in_loop: Literal["continue", "break", "raise"] = "raise"
+
+    @model_validator(mode="before")
+    def migrate_old_nodes(cls, data: dict[str, Any]):
+        for key in ["nodes"]:
+            raw_nodes = data.get(key, [])
+            new_nodes = []
+            used_old_format = False
+
+            for item in raw_nodes:
+                if (
+                    isinstance(item, ActionNode)
+                    or isinstance(item, ForLoopNode)
+                    or isinstance(item, IfElseNode)
+                ):
+                    new_nodes.append(item)
+                    continue
+
+                # --- new format: already has a type ---
+                if isinstance(item, dict) and "type" in item:
+                    new_nodes.append(item)
+                    continue
+
+                # --- old format cases ---
+                used_old_format = True
+
+                if isinstance(item, dict) and "condition" in item:
+                    new_nodes.append({"type": "if_else_node", **item})
+                    continue
+
+                if isinstance(item, dict) and "variable_name" in item:
+                    new_nodes.append({"type": "for_loop_node", **item})
+                    continue
+
+                new_nodes.append({"type": "action_node", **item})
+
+            if used_old_format:
+                logger.warning(
+                    "Old node format without 'type' is deprecated. "
+                    "Use the new format: {'type': 'action_node'|'for_loop_node'|'if_else_node', ...}"
+                )
+
+            data[key] = new_nodes
+        return data
+
+
+class IfElseNode(BaseModel):
+    type: Literal["if_else_node"]
+    condition: str
+    if_nodes: list[ActionNode | IfElseNodeRef | ForLoopNodeRef]
+    else_nodes: list[ActionNode | IfElseNodeRef | ForLoopNodeRef] = []
+
+    @model_validator(mode="before")
+    def migrate_old_nodes(cls, data: dict[str, Any]):
+        for key in ["if_nodes", "else_nodes"]:
+            raw_nodes = data.get(key, [])
+            new_nodes = []
+            used_old_format = False
+
+            for item in raw_nodes:
+                if (
+                    isinstance(item, ActionNode)
+                    or isinstance(item, ForLoopNode)
+                    or isinstance(item, IfElseNode)
+                ):
+                    new_nodes.append(item)
+                    continue
+
+                # --- new format: already has a type ---
+                if isinstance(item, dict) and "type" in item:
+                    new_nodes.append(item)
+                    continue
+
+                # --- old format cases ---
+                used_old_format = True
+
+                if isinstance(item, dict) and "condition" in item:
+                    new_nodes.append({"type": "if_else_node", **item})
+                    continue
+
+                if isinstance(item, dict) and "variable_name" in item:
+                    new_nodes.append({"type": "for_loop_node", **item})
+                    continue
+
+                new_nodes.append({"type": "action_node", **item})
+
+            if used_old_format:
+                logger.warning(
+                    "Old node format without 'type' is deprecated. "
+                    "Use the new format: {'type': 'action_node'|'for_loop_node'|'if_else_node', ...}"
+                )
+
+            data[key] = new_nodes
+        return data
+
+
+class Parameters(BaseModel):
+    input_parameters: dict[str, list[str | int | float | bool]]
+    secure_parameters: dict[str, list[SecureParameter]] = Field(default_factory=dict)
+    generated_parameters: dict[str, list[str | int | float | bool | None]]
+
+    @model_validator(mode="after")
+    def validate_parameters(self):
+        reserved_parameter_names = set(["current_page_url"])
+
+        for d in [
+            self.input_parameters,
+            self.generated_parameters,
+            self.secure_parameters,
+        ]:
+            for key in d.keys():
+                if key in reserved_parameter_names:
+                    raise ValueError(f"Parameter name {key} is reserved")
+                if not key.isidentifier():
+                    raise ValueError(
+                        f"Parameter name {key} is not a valid variable name"
+                    )
+        return self
+
+
+## TODO: fix expected downloads for ForLoop
+class Automation(BaseModel):
+    browser_channel: Literal["chromium", "chrome"] = "chromium"
+    expected_downloads: int = 0
+    url: str
+    parameters: Parameters
+    nodes: list[
+        Annotated[ActionNode | ForLoopNode | IfElseNode, Field(discriminator="type")]
+    ]
+    automation_description: str | None = None
+    automation_endpoint: str | None = None
+
+    @model_validator(mode="before")
+    def migrate_old_nodes(cls, data: dict[str, Any]):
+        raw_nodes = data.get("nodes", [])
+        new_nodes = []
+        used_old_format = False
+
+        for item in raw_nodes:
+            if (
+                isinstance(item, ActionNode)
+                or isinstance(item, ForLoopNode)
+                or isinstance(item, IfElseNode)
+            ):
+                new_nodes.append(item)
+                continue
+
+            # --- new format: already has a type ---
+            if isinstance(item, dict) and "type" in item:
+                new_nodes.append(item)
+                continue
+
+            # --- old format cases ---
+            used_old_format = True
+
+            if isinstance(item, dict) and "condition" in item:
+                new_nodes.append({"type": "if_else_node", **item})
+                continue
+
+            if isinstance(item, dict) and "variable_name" in item:
+                new_nodes.append({"type": "for_loop_node", **item})
+                continue
+
+            new_nodes.append({"type": "action_node", **item})
+
+        if used_old_format:
+            logger.warning(
+                "Old node format without 'type' is deprecated. "
+                "Use the new format: {'type': 'action_node'|'for_loop_node'|'if_else_node', ...}"
+            )
+
+        data["nodes"] = new_nodes
+        return data
+
+    @model_validator(mode="after")
+    def validate_parameters_with_examples(cls, model: "Automation"):
+        ## TODO: static check that all parameters with examples are used in the nodes
+        return model
+
+    def model_dump(self, *, sort_params_by_nodes: bool = False, **kwargs):
+        """
+        Extended model_dump with option to sort parameters by node order
+
+        Args:
+            sort_params_by_nodes: If True, sort input_parameters by their
+                                 appearance order in nodes. Fails gracefully
+                                 if sorting encounters any errors.
+            **kwargs: All standard Pydantic model_dump arguments (exclude,
+                     exclude_none, exclude_defaults, etc.)
+        """
+        data = super().model_dump(**kwargs)
+
+        if sort_params_by_nodes:
+            data = self._sort_parameters_by_node_order(data)
+
+        return data
+
+    def _sort_parameters_by_node_order(self, data: dict) -> dict:
+        """
+        Sort input_parameters based on their first appearance in nodes.
+        Returns data unchanged if any error occurs.
+
+        This method searches for parameter references in the format {param_name[index]}
+        throughout the entire nodes array and reorders input_parameters accordingly.
+        Parameters that don't appear in nodes are placed at the end.
+        """
+        try:
+            import json
+            import re
+
+            # Convert nodes to string to search for all parameter references
+            nodes_str = json.dumps(data.get("nodes", []))
+            # Extract all {param_name[index]} references
+            pattern = r"\{(\w+)\[\d+\]\}"
+            matches = re.findall(pattern, nodes_str)
+            # Preserve order of first occurrence
+            param_order = []
+            seen = set()
+            for param in matches:
+                if param not in seen:
+                    param_order.append(param)
+                    seen.add(param)
+            # Reorder input_parameters if they exist
+            if "parameters" in data and "input_parameters" in data["parameters"]:
+                old_params = data["parameters"]["input_parameters"]
+                sorted_params = {}
+                # Add params in order they appear in nodes
+                for param_name in param_order:
+                    if param_name in old_params:
+                        sorted_params[param_name] = old_params[param_name]
+                # Add remaining params that don't appear in nodes (at the end)
+                for param_name, param_value in old_params.items():
+                    if param_name not in sorted_params:
+                        sorted_params[param_name] = param_value
+                data["parameters"]["input_parameters"] = sorted_params
+            return data
+
+        except Exception as e:
+            # Log the error if logging is available
+            logger.warning(f"Failed to sort parameters by node order: {e}")
+
+            # Return original data unchanged
+            return data

optexity/schema/callback.py

@@ -0,0 +1,16 @@
+from typing import Literal
+
+from pydantic import BaseModel
+
+
+class CallbackResponse(BaseModel):
+    task_id: str
+    recording_id: str
+    output_data: list[dict | str] | None
+    status: Literal["queued", "allocated", "running", "success", "failed", "cancelled"]
+    error: str | None
+    final_screenshot: str | None = None
+    endpoint_name: str
+    downloads: list[dict] | None = None
+    input_parameters: dict[str, list[str | int | float | bool]] | None = None
+    unique_parameter_names: list[str] | None = None

optexity/schema/inference.py

@@ -0,0 +1,87 @@
+from datetime import datetime
+
+from pydantic import BaseModel, Field, model_validator
+
+from optexity.schema.automation import SecureParameter
+
+
+class InferenceRequest(BaseModel):
+    endpoint_name: str
+    input_parameters: dict[str, list[str | int | float | bool]]
+    unique_parameter_names: list[str] = Field(default_factory=list)
+    secure_parameters: dict[str, list[SecureParameter]] = Field(default_factory=dict)
+    use_proxy: bool = False
+
+    @model_validator(mode="after")
+    def validate_unique_parameter_names(self):
+        for unique_parameter_name in self.unique_parameter_names:
+            if unique_parameter_name not in self.input_parameters and (
+                self.secure_parameters is None
+                or unique_parameter_name not in self.secure_parameters
+            ):
+                raise ValueError(
+                    f"unique_parameter_name {unique_parameter_name} not found in input_parameters or secure_parameters"
+                )
+        return self
+
+
+class FetchEmailMessagesRequest(BaseModel):
+    receiver_email_address: str  # receiver's email address
+    sender_email_address: str  # sender's email address
+    start_2fa_time: datetime
+    end_2fa_time: datetime
+
+    @model_validator(mode="after")
+    def validate_time_parameters(self):
+        assert (
+            self.start_2fa_time.tzinfo is not None
+        ), "start_2fa_time must be timezone-aware"
+        assert (
+            self.end_2fa_time.tzinfo is not None
+        ), "end_2fa_time must be timezone-aware"
+        assert (
+            self.start_2fa_time < self.end_2fa_time
+        ), "start_2fa_time must be before end_2fa_time"
+        return self
+
+    class Config:
+        json_encoders = {datetime: lambda v: v.isoformat() if v is not None else None}
+
+
+class FetchSlackMessagesRequest(BaseModel):
+    slack_workspace_domain: str
+    channel_name: str
+    sender_name: str
+    start_2fa_time: datetime
+    end_2fa_time: datetime
+
+    @model_validator(mode="after")
+    def validate_time_parameters(self):
+        assert (
+            self.start_2fa_time.tzinfo is not None
+        ), "start_2fa_time must be timezone-aware"
+        assert (
+            self.end_2fa_time.tzinfo is not None
+        ), "end_2fa_time must be timezone-aware"
+        assert (
+            self.start_2fa_time < self.end_2fa_time
+        ), "start_2fa_time must be before end_2fa_time"
+        return self
+
+    class Config:
+        json_encoders = {datetime: lambda v: v.isoformat() if v is not None else None}
+
+
+class Message(BaseModel):
+    message_id: str | None = None
+    message_text: str
+    timestamp: datetime
+
+    @model_validator(mode="after")
+    def validate_timestamp(self):
+        assert self.timestamp.tzinfo is not None, "timestamp must be timezone-aware"
+        return self
+
+
+class FetchMessagesResponse(BaseModel):
+    messages: list[Message]

optexity/schema/memory.py

@@ -0,0 +1,100 @@
+import asyncio
+from datetime import datetime
+from pathlib import Path
+from typing import Any, Literal
+
+from playwright.async_api import Download
+from pydantic import BaseModel, Field, model_validator
+
+from optexity.schema.token_usage import TokenUsage
+
+
+class NetworkRequest(BaseModel):
+    url: str
+    method: str
+    headers: dict
+    body: str | bytes | None | dict | Any
+
+
+class NetworkError(BaseModel):
+    url: str = Field(...)
+    message: str = Field(...)
+    stack_trace: str = Field(...)
+
+
+class NetworkResponse(BaseModel):
+    url: str = Field(...)
+    status: int = Field(...)
+    headers: dict = Field(...)
+    body: dict | str | None | bytes | Any = Field(default=None)
+    method: str = Field(...)
+    content_length: int = Field(...)
+
+
+class AutomationState(BaseModel):
+    step_index: int = Field(default_factory=lambda: -1)
+    try_index: int = Field(default_factory=lambda: -1)
+    start_2fa_time: datetime | None = Field(default=None)
+
+    @model_validator(mode="after")
+    def validate_start_2fa_time(self):
+        if self.start_2fa_time is not None:
+            assert (
+                self.start_2fa_time.tzinfo is not None
+            ), "start_2fa_time must be timezone-aware"
+        return self
+
+
+class BrowserState(BaseModel):
+    url: str = Field(...)
+    title: str | None = Field(default=None)
+    screenshot: str | None = Field(default=None)
+    html: str | None = Field(default=None)
+    axtree: str | None = Field(default=None)
+    final_prompt: str | None = Field(default=None)
+    llm_response: str | dict | None = Field(default=None)
+
+
+class ScreenshotData(BaseModel):
+    filename: str = Field(...)
+    base64: str = Field(...)
+
+
+class OutputData(BaseModel):
+    unique_identifier: str | None = None
+    json_data: dict | None = Field(default=None)
+    screenshot: ScreenshotData = Field(default=None)
+    text: str | None = Field(default=None)
+
+
+class ForLoopStatus(BaseModel):
+    variable_name: str
+    index: int
+    value: str | int | float | bool
+    error: str | None = None
+    status: Literal["success", "error", "skipped"]
+
+
+class Variables(BaseModel):
+    output_data: list[OutputData] = Field(default_factory=list)
+    for_loop_status: list[list[ForLoopStatus]] = Field(default_factory=list)
+    generated_variables: dict = Field(default_factory=dict)
+
+
+class Memory(BaseModel):
+    variables: Variables = Field(default_factory=Variables)
+    automation_state: AutomationState = Field(default_factory=AutomationState)
+    browser_states: list[BrowserState] = Field(default_factory=list)
+    token_usage: TokenUsage = Field(default_factory=TokenUsage)
+    download_lock: asyncio.Lock = Field(default_factory=asyncio.Lock)
+    raw_downloads: dict[Path, tuple[bool, Download | None]] = Field(
+        default_factory=dict
+    )
+    urls_to_downloads: list[tuple[str, str]] = Field(default_factory=list)
+    downloads: list[Path] = Field(default_factory=list)
+    final_screenshot: str | None = Field(default=None)
+
+    model_config = {
+        "arbitrary_types_allowed": True,
+        "exclude": {"download_lock"},
+    }