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

optexity/inference/models/llm_model.py

@@ -0,0 +1,210 @@
+import ast
+import logging
+import re
+import time
+from enum import Enum, unique
+from typing import Optional
+
+import tokencost.costs
+from pydantic import BaseModel, ValidationError
+
+from optexity.schema.token_usage import TokenUsage
+
+logger = logging.getLogger(__name__)
+
+
+@unique
+class HumanModels(Enum):
+    TERMINAL_INPUT = "terminal-input"
+
+
+@unique
+class GeminiModels(Enum):
+    GEMINI_1_5_FLASH = "gemini-1.5-flash"
+    GEMINI_2_0_FLASH = "gemini-2.0-flash"
+    GEMINI_2_5_FLASH = "gemini-2.5-flash"
+    GEMINI_2_5_FLASH_LITE = "gemini-2.5-flash-lite-preview-06-17"
+    GEMINI_2_5_PRO = "gemini-2.5-pro"
+
+
+@unique
+class OpenAIModels(Enum):
+    GPT_4O = "gpt-4o"
+    GPT_4O_MINI = "gpt-4o-mini"
+    GPT_4_1 = "gpt-4.1"
+    GPT_4_1_MINI = "gpt-4.1-mini"
+
+
+class LLMModel:
+    def __init__(
+        self,
+        model_name: GeminiModels | HumanModels | OpenAIModels,
+        use_structured_output: bool,
+    ):
+
+        self.model_name = model_name
+        self.use_structured_output = use_structured_output
+
+    def _get_model_response(
+        self, prompt: str, system_instruction: Optional[str] = None
+    ) -> tuple[str, TokenUsage]:
+        raise NotImplementedError("This method should be implemented by subclasses.")
+
+    def _get_model_response_with_structured_output(
+        self,
+        prompt: str,
+        response_schema: BaseModel,
+        screenshot: Optional[str] = None,
+        pdf_url: Optional[str] = None,
+        system_instruction: Optional[str] = None,
+    ) -> tuple[BaseModel, TokenUsage]:
+        raise NotImplementedError("This method should be implemented by subclasses.")
+
+    def get_model_response(
+        self, prompt: str, system_instruction: Optional[str] = None
+    ) -> tuple[str, TokenUsage]:
+
+        max_retries = 3
+        for i in range(max_retries):
+            try:
+                return self._get_model_response(prompt, system_instruction)
+            except Exception as e:
+                logger.error(f"LLM Error during inference: {e}")
+                if i < max_retries - 1:
+                    logger.info(f"Retrying... {i + 1}/{max_retries}")
+                    time.sleep(5)
+                continue
+        raise Exception("Max retries exceeded for LLM")
+
+    def get_model_response_with_structured_output(
+        self,
+        prompt: str,
+        response_schema: BaseModel,
+        screenshot: Optional[str] = None,
+        pdf_url: Optional[str] = None,
+        system_instruction: Optional[str] = None,
+    ) -> tuple[BaseModel, TokenUsage]:
+
+        total_token_usage = TokenUsage()
+        max_retries = 3
+        last_exception = ""
+        for i in range(max_retries):
+            try:
+                # raise Exception("Test error")
+                parsed_response, token_usage = (
+                    self._get_model_response_with_structured_output(
+                        prompt=prompt,
+                        response_schema=response_schema,
+                        screenshot=screenshot,
+                        pdf_url=pdf_url,
+                        system_instruction=system_instruction,
+                    )
+                )
+                total_token_usage += token_usage
+                if parsed_response is not None:
+                    return parsed_response, total_token_usage
+            except Exception as e:
+                logger.error(f"LLM with structured output Error during inference: {e}")
+                if i < max_retries - 1:
+                    logger.info(f"Retrying... {i + 1}/{max_retries}")
+                    time.sleep(20)
+                last_exception = str(e)
+
+        raise Exception(
+            "Max retries exceeded for LLM with structured output"
+            + "\n"
+            + last_exception
+        )
+
+    def extract_json_objects(self, text):
+        stack = []  # Stack to track `{` positions
+        json_candidates = []  # Potential JSON substrings
+
+        # Iterate through the text to find balanced { }
+        for i, char in enumerate(text):
+            if char == "{":
+                stack.append(i)  # Store index of '{'
+            elif char == "}" and stack:
+                start = stack.pop()  # Get the last unmatched '{'
+                json_candidates.append(text[start : i + 1])  # Extract substring
+
+        return json_candidates
+
+    def parse_from_completion(
+        self, content: str, response_schema: BaseModel
+    ) -> BaseModel:
+        patterns = [r"```json\n(.*?)\n```"]
+        json_blocks = []
+        for pattern in patterns:
+            json_blocks += re.findall(pattern, content, re.DOTALL)
+        json_blocks += self.extract_json_objects(content)
+        for block in json_blocks:
+            block = block.strip()
+            try:
+                response = response_schema.model_validate_json(block)
+                return response
+            except Exception as e:
+                try:
+                    block_dict = ast.literal_eval(block)
+                    response = response_schema.model_validate(block_dict)
+                    return response
+                except Exception as e:
+                    continue
+
+        raise ValidationError("Could not parse response from completion.")
+
+    def get_token_usage(
+        self,
+        input_tokens: int | None = None,
+        output_tokens: int | None = None,
+        tool_use_tokens: int | None = None,
+        thoughts_tokens: int | None = None,
+        total_tokens: Optional[int] = None,
+    ) -> TokenUsage:
+        if input_tokens is None:
+            input_tokens = 0
+        if output_tokens is None:
+            output_tokens = 0
+        if tool_use_tokens is None:
+            tool_use_tokens = 0
+        if thoughts_tokens is None:
+            thoughts_tokens = 0
+        if total_tokens is None:
+            total_tokens = 0
+        input_cost = tokencost.costs.calculate_cost_by_tokens(
+            model=self.model_name.value,
+            num_tokens=input_tokens,
+            token_type="input",
+        )
+        output_cost = tokencost.costs.calculate_cost_by_tokens(
+            model=self.model_name.value,
+            num_tokens=output_tokens,
+            token_type="output",
+        )
+        tool_use_cost = tokencost.costs.calculate_cost_by_tokens(
+            model=self.model_name.value,
+            num_tokens=tool_use_tokens,
+            token_type="output",
+        )
+        thoughts_cost = tokencost.costs.calculate_cost_by_tokens(
+            model=self.model_name.value,
+            num_tokens=thoughts_tokens,
+            token_type="output",
+        )
+        calculated_total_tokens = (
+            input_tokens + output_tokens + tool_use_tokens + thoughts_tokens
+        )
+        total_cost = input_cost + output_cost + tool_use_cost + thoughts_cost
+        return TokenUsage(
+            input_tokens=input_tokens,
+            output_tokens=output_tokens,
+            tool_use_tokens=tool_use_tokens,
+            thoughts_tokens=thoughts_tokens,
+            total_tokens=total_tokens,
+            calculated_total_tokens=calculated_total_tokens,
+            input_cost=input_cost,
+            output_cost=output_cost,
+            tool_use_cost=tool_use_cost,
+            thoughts_cost=thoughts_cost,
+            total_cost=total_cost,
+        )

optexity/inference/run_local.py

@@ -0,0 +1,200 @@
+import asyncio
+import logging
+import os
+import uuid
+from datetime import datetime, timezone
+
+from dotenv import load_dotenv
+
+from optexity.examples.fadv import fadv_test
+from optexity.examples.i94 import automation
+from optexity.examples.pshpgeorgia_medicaid import (
+    pshpgeorgia_login_test,
+    pshpgeorgia_medicaid_test,
+)
+from optexity.examples.shein import shein_test
+from optexity.examples.supabase_login import supabase_login_test
+from optexity.inference.core.run_automation import run_automation
+from optexity.inference.infra.browser import Browser
+from optexity.schema.memory import Memory, Variables
+from optexity.schema.task import Task
+
+load_dotenv()
+
+
+logger = logging.getLogger(__name__)
+logging.getLogger(__name__).setLevel(logging.DEBUG)
+
+
+async def run_supabase_login_test():
+    logger.debug("Starting Supabase login test")
+    browser = Browser()
+    memory = Memory(
+        variables=Variables(
+            input_variables={
+                "username": ["test@test.com"],
+                "password": ["password"],
+            }
+        )
+    )
+
+    await browser.start()
+    logger.info("Browser started")
+    logger.info("Navigating to Supabase")
+    await browser.go_to_url("https://supabase.com")
+    logger.info("Navigated to Supabase")
+    logger.info("Sleeping for 5 seconds")
+    await asyncio.sleep(2)
+
+    logger.info("Running automation")
+    await run_automation(supabase_login_test, memory, browser)
+    logger.info("Automation finished")
+    await asyncio.sleep(5)
+
+    await browser.stop()
+
+
+async def run_pshpgeorgia_test():
+    try:
+        logger.debug("Starting PSHP Georgia test")
+        browser = Browser()
+        memory = Memory(
+            variables=Variables(
+                input_variables={
+                    "username": [os.environ.get("USERNAME")],
+                    "password": [os.environ.get("PASSWORD")],
+                    "plan_type": [os.environ.get("PLAN_TYPE")],
+                    "member_id": [os.environ.get("MEMBER_ID")],
+                    "dob": [os.environ.get("DOB")],
+                }
+            )
+        )
+
+        await browser.start()
+        logger.debug("Browser started")
+        logger.debug("Navigating to PSHP Georgia")
+        await browser.go_to_url(
+            "https://sso.entrykeyid.com/as/authorization.oauth2?response_type=code&client_id=f6a6219c-be42-421b-b86c-e4fc509e2e87&scope=openid%20profile&state=_igWklSsnrkO5DQfjBMMuN41ksMJePZQ_SM_61wTJlA%3D&redirect_uri=https://provider.pshpgeorgia.com/careconnect/login/oauth2/code/pingcloud&code_challenge_method=S256&nonce=xG41TJjco_x7Vs_MQgcS3bw5njLiJsXCqvO-V8THmY0&code_challenge=ZTaVHaZCNFTejXNJo51RlJ3Kv9dH0tMODPTqO7hiP3A&app_origin=https://provider.pshpgeorgia.com/careconnect/login/oauth2/code/pingcloud&brand=pshpgeorgia"
+        )
+        logger.debug("Navigated to PSHP Georgia")
+
+        logger.debug("Running login test")
+        await run_automation(pshpgeorgia_login_test, memory, browser)
+        logger.debug("Login test finished")
+
+        logger.debug("Running Medicaid test")
+        await run_automation(pshpgeorgia_medicaid_test, memory, browser)
+        logger.debug("Medicaid test finished")
+
+        await asyncio.sleep(5)
+        await browser.stop()
+    except Exception as e:
+        logger.error(f"Error running PSHP Georgia test: {e}")
+        raise e
+    finally:
+        await browser.stop()
+
+
+async def run_i94_test():
+    try:
+        logger.debug("Starting I-94 test")
+        browser = Browser(stealth=True)
+        memory = Memory(
+            variables=Variables(
+                input_variables={
+                    "last_name": [os.environ.get("LAST_NAME")],
+                    "first_name": [os.environ.get("FIRST_NAME")],
+                    "nationality": [os.environ.get("NATIONALITY")],
+                    "date_of_birth": [os.environ.get("DATE_OF_BIRTH")],
+                    "document_number": [os.environ.get("DOCUMENT_NUMBER")],
+                }
+            )
+        )
+
+        await browser.start()
+        logger.debug("Browser started")
+        logger.debug("Navigating to I-94")
+        await browser.go_to_url(automation.url)
+        logger.debug("Navigated to I-94")
+
+        logger.debug("Running I-94 test")
+        await asyncio.sleep(5)
+        await run_automation(automation, memory, browser)
+        logger.debug("I-94 test finished")
+
+        await asyncio.sleep(5)
+        await browser.stop()
+    except Exception as e:
+        logger.error(f"Error running I-94 test: {e}")
+        raise e
+    finally:
+        await browser.stop()
+
+
+async def run_shein_test():
+
+    try:
+        logger.debug("Starting Shein test")
+        task = Task(
+            task_id=str(uuid.uuid4()),
+            user_id=str(uuid.uuid4()),
+            recording_id=str(uuid.uuid4()),
+            automation=shein_test,
+            input_parameters={},
+            unique_parameter_names=[],
+            created_at=datetime.now(timezone.utc),
+            status="queued",
+        )
+        await run_automation(task, 0)
+    except Exception as e:
+        logger.error(f"Error running Shein test: {e}")
+        raise e
+    finally:
+
+        logger.debug("Remaining tasks:")
+        for task in asyncio.all_tasks():
+            if task is not asyncio.current_task():
+                logger.debug(f"Remaining task: {task.get_coro()}")
+
+    logger.debug("Shein test finished")
+
+
+async def run_fadv_test():
+    try:
+        logger.debug("Starting FADV test task")
+        task = Task(
+            task_id=str(uuid.uuid4()),
+            user_id=str(uuid.uuid4()),
+            recording_id=str(uuid.uuid4()),
+            automation=fadv_test,
+            input_parameters={
+                "client_id": [os.environ.get("client_id")],
+                "user_id": [os.environ.get("user_id")],
+                "password": [os.environ.get("password")],
+                "secret_answer": [os.environ.get("secret_answer")],
+                "start_date": [os.environ.get("start_date")],
+            },
+            unique_parameter_names=[],
+            created_at=datetime.now(timezone.utc),
+            status="queued",
+        )
+        await run_automation(task, 0)
+        await asyncio.sleep(5)
+    except Exception as e:
+        logger.error(f"Error running FADV test: {e}")
+        raise e
+    finally:
+        logger.debug("Remaining tasks:")
+        for task in asyncio.all_tasks():
+            if task is not asyncio.current_task():
+                logger.debug(f"Remaining task: {task.get_coro()}")
+    logger.debug("FADV test finished")
+
+
+if __name__ == "__main__":
+
+    # asyncio.run(run_supabase_login_test())
+    # asyncio.run(run_pshpgeorgia_test())
+    # asyncio.run(run_i94_test())
+    asyncio.run(run_fadv_test())
+    # asyncio.run(run_shein_test())

optexity/schema/actions/assertion_action.py

@@ -0,0 +1,66 @@
+from typing import Literal, Optional
+
+from pydantic import BaseModel, field_validator, model_validator
+
+from optexity.schema.actions.extraction_action import LLMExtraction
+
+
+class LLMAssertion(LLMExtraction):
+    source: list[Literal["axtree", "screenshot"]] = ["screenshot"]
+    extraction_format: dict = {"assertion_result": "bool", "assertion_reason": "str"}
+
+    @model_validator(mode="after")
+    def validate_output_var_in_format(self):
+        if "screenshot" not in self.source:
+            self.source.append("screenshot")
+
+        return self
+
+
+class NetworkCallAssertion(BaseModel):
+    url_pattern: Optional[str] = None
+    header_filter: Optional[dict[str, str]] = None
+
+
+class PythonScriptAssertion(BaseModel):
+    script: str
+    ## TODO: add output to memory variables
+
+    @field_validator("script")
+    @classmethod
+    def validate_script(cls, v: str):
+        if not v.strip():
+            raise ValueError("Script cannot be empty")
+        return v
+
+
+class AssertionAction(BaseModel):
+    network_call: Optional[NetworkCallAssertion] = None
+    llm: Optional[LLMAssertion] = None
+    python_script: Optional[PythonScriptAssertion] = None
+
+    @model_validator(mode="after")
+    def validate_one_assertion(cls, model: "AssertionAction"):
+        """Ensure exactly one of the extraction types is set and matches the type."""
+        provided = {
+            "llm": model.llm,
+            "network_call": model.network_call,
+            "python_script": model.python_script,
+        }
+        non_null = [k for k, v in provided.items() if v is not None]
+
+        if len(non_null) != 1:
+            raise ValueError(
+                "Exactly one of llm, networkcall, or python must be provided"
+            )
+
+        return model
+
+    def replace(self, pattern: str, replacement: str):
+        if self.network_call:
+            pass
+        if self.llm:
+            self.llm.replace(pattern, replacement)
+        if self.python_script:
+            pass
+        return self

optexity/schema/actions/extraction_action.py

@@ -0,0 +1,143 @@
+from typing import Any, List, Literal, Optional
+from uuid import uuid4
+
+from pydantic import BaseModel, Field, field_validator, model_validator
+
+from optexity.utils.utils import build_model
+
+
+class LLMExtraction(BaseModel):
+    source: list[Literal["axtree", "screenshot"]] = ["axtree"]
+    extraction_format: dict
+    extraction_instructions: str
+    output_variable_names: list[str] | None = None
+    llm_provider: Literal["gemini"] = "gemini"
+    llm_model_name: str = "gemini-2.5-flash"
+
+    def build_model(self):
+        return build_model(self.extraction_format)
+
+    @field_validator("extraction_format")
+    def validate_extraction_format(cls, v):
+        if isinstance(v, dict):
+            try:
+                build_model(v)
+            except Exception as e:
+                raise ValueError(f"Invalid extraction_format dict: {e}")
+            return v
+        raise ValueError("extraction_format must be either a string or a dict")
+
+    @model_validator(mode="after")
+    def validate_output_var_in_format(self):
+
+        if self.output_variable_names is not None:
+            for key in self.output_variable_names:
+                if key not in self.extraction_format:
+                    raise ValueError(
+                        f"Output variable {key} not found in extraction_format"
+                    )
+                ## TODO: fix this
+                # if eval(self.extraction_format[key]) not in [
+                #     int,
+                #     float,
+                #     bool,
+                #     str,
+                #     None,
+                #     list[str | int | float | bool | None],
+                #     List[str | int | float | bool | None],
+                # ]:
+                #     raise ValueError(
+                #         f"Output variable {key} must be a string, int, float, bool, or a list of strings, ints, floats, or bools"
+                #     )
+
+        return self
+
+    def replace(self, pattern: str, replacement: str):
+        return self
+
+
+class NetworkCallExtraction(BaseModel):
+    url_pattern: Optional[str] = None
+    extract_from: None | Literal["request", "response"] = None
+    download_from: None | Literal["request", "response"] = None
+    download_filename: str | None = None
+
+    @model_validator(mode="before")
+    def download_filename_if_download_from_is_set(cls, data: dict[str, Any]):
+        if (
+            "downlowd_from" in data
+            and data["download_from"] is not None
+            and ("download_filename" not in data or data["download_filename"] is None)
+        ):
+            data["download_filename"] = str(uuid4())
+
+        return data
+
+    def replace(self, pattern: str, replacement: str):
+        return self
+
+
+class PythonScriptExtraction(BaseModel):
+    script: str
+    ## TODO: add output to memory variables
+
+    @field_validator("script")
+    @classmethod
+    def validate_script(cls, v: str):
+        if not v.strip():
+            raise ValueError("Script cannot be empty")
+        return v
+
+    def replace(self, pattern: str, replacement: str):
+        self.script = self.script.replace(pattern, replacement)
+        return self
+
+
+class ScreenshotExtraction(BaseModel):
+    filename: str
+    full_page: bool = True
+
+
+class StateExtraction(BaseModel):
+    pass
+
+
+class ExtractionAction(BaseModel):
+    unique_identifier: str | None = None
+    network_call: Optional[NetworkCallExtraction] = None
+    llm: Optional[LLMExtraction] = None
+    python_script: Optional[PythonScriptExtraction] = None
+    screenshot: Optional[ScreenshotExtraction] = None
+    state: Optional[StateExtraction] = None
+
+    @model_validator(mode="after")
+    def validate_one_extraction(cls, model: "ExtractionAction"):
+        """Ensure exactly one of the extraction types is set and matches the type."""
+        provided = {
+            "llm": model.llm,
+            "network_call": model.network_call,
+            "python_script": model.python_script,
+            "screenshot": model.screenshot,
+            "state": model.state,
+        }
+        non_null = [k for k, v in provided.items() if v is not None]
+
+        if len(non_null) != 1:
+            raise ValueError(
+                "Exactly one of llm, networkcall, python_script, or screenshot must be provided"
+            )
+
+        return model
+
+    def replace(self, pattern: str, replacement: str):
+        if self.network_call:
+            self.network_call.replace(pattern, replacement)
+        if self.llm:
+            self.llm.replace(pattern, replacement)
+        if self.python_script:
+            self.python_script.replace(pattern, replacement)
+        if self.unique_identifier:
+            self.unique_identifier = self.unique_identifier.replace(
+                pattern, replacement
+            )
+        return self