openadapt-ml 0.1.0__py3-none-any.whl → 0.2.1__py3-none-any.whl

openadapt_ml/retrieval/index.py

@@ -0,0 +1,194 @@
+"""Demo index for storing and retrieving demonstrations."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Any, Dict, List, Optional
+
+from openadapt_ml.retrieval.embeddings import TextEmbedder
+from openadapt_ml.schema import Episode
+
+
+@dataclass
+class DemoMetadata:
+    """Metadata for a single demonstration.
+
+    Stores both the episode and computed features for retrieval.
+    """
+
+    episode: Episode
+    app_name: Optional[str] = None
+    domain: Optional[str] = None
+    metadata: Dict[str, Any] = field(default_factory=dict)
+
+    # Computed at index time
+    text_embedding: Dict[str, float] = field(default_factory=dict)
+
+
+class DemoIndex:
+    """Index for demonstrations.
+
+    Stores episodes with their metadata and embeddings for efficient retrieval.
+    """
+
+    def __init__(self) -> None:
+        """Initialize the demo index."""
+        self.demos: List[DemoMetadata] = []
+        self.embedder = TextEmbedder()
+        self._is_fitted = False
+
+    def _extract_app_name(self, episode: Episode) -> Optional[str]:
+        """Extract app name from episode steps.
+
+        Args:
+            episode: Episode to extract from.
+
+        Returns:
+            App name if found, None otherwise.
+        """
+        # Look through observations to find app_name
+        for step in episode.steps:
+            if step.observation and step.observation.app_name:
+                return step.observation.app_name
+        return None
+
+    def _extract_domain(self, episode: Episode) -> Optional[str]:
+        """Extract domain from episode metadata or URL.
+
+        Args:
+            episode: Episode to extract from.
+
+        Returns:
+            Domain if found, None otherwise.
+        """
+        # Try to extract from URL in observations
+        for step in episode.steps:
+            if step.observation and step.observation.url:
+                url = step.observation.url
+                # Simple domain extraction (e.g., "github.com" from "https://github.com/...")
+                if "://" in url:
+                    domain = url.split("://")[1].split("/")[0]
+                    # Remove www. prefix
+                    if domain.startswith("www."):
+                        domain = domain[4:]
+                    return domain
+
+        return None
+
+    def add(
+        self,
+        episode: Episode,
+        app_name: Optional[str] = None,
+        domain: Optional[str] = None,
+        metadata: Optional[Dict[str, Any]] = None,
+    ) -> None:
+        """Add an episode to the index.
+
+        Args:
+            episode: Episode to add.
+            app_name: Optional app name (auto-extracted if not provided).
+            domain: Optional domain (auto-extracted if not provided).
+            metadata: Additional metadata for the episode.
+        """
+        # Auto-extract app_name and domain if not provided
+        if app_name is None:
+            app_name = self._extract_app_name(episode)
+        if domain is None:
+            domain = self._extract_domain(episode)
+
+        demo_meta = DemoMetadata(
+            episode=episode,
+            app_name=app_name,
+            domain=domain,
+            metadata=metadata or {},
+        )
+
+        self.demos.append(demo_meta)
+        # Mark as not fitted since we added new data
+        self._is_fitted = False
+
+    def add_many(self, episodes: List[Episode]) -> None:
+        """Add multiple episodes to the index.
+
+        Args:
+            episodes: List of episodes to add.
+        """
+        for episode in episodes:
+            self.add(episode)
+
+    def build(self) -> None:
+        """Build the index by computing embeddings.
+
+        This must be called after adding all demos and before retrieval.
+        """
+        if not self.demos:
+            return
+
+        # Fit embedder on all instruction texts
+        instruction_texts = [demo.episode.instruction for demo in self.demos]
+        self.embedder.fit(instruction_texts)
+
+        # Compute embeddings for each demo
+        for demo in self.demos:
+            demo.text_embedding = self.embedder.embed(demo.episode.instruction)
+
+        self._is_fitted = True
+
+    def is_empty(self) -> bool:
+        """Check if the index is empty.
+
+        Returns:
+            True if no demos have been added.
+        """
+        return len(self.demos) == 0
+
+    def is_fitted(self) -> bool:
+        """Check if the index has been built.
+
+        Returns:
+            True if build() has been called.
+        """
+        return self._is_fitted
+
+    def get_all_demos(self) -> List[DemoMetadata]:
+        """Get all demos in the index.
+
+        Returns:
+            List of all DemoMetadata objects.
+        """
+        return self.demos
+
+    def get_apps(self) -> List[str]:
+        """Get list of unique app names in the index.
+
+        Returns:
+            List of app names (excluding None).
+        """
+        apps = {demo.app_name for demo in self.demos if demo.app_name is not None}
+        return sorted(apps)
+
+    def get_domains(self) -> List[str]:
+        """Get list of unique domains in the index.
+
+        Returns:
+            List of domains (excluding None).
+        """
+        domains = {demo.domain for demo in self.demos if demo.domain is not None}
+        return sorted(domains)
+
+    def __len__(self) -> int:
+        """Return number of demos in the index.
+
+        Returns:
+            Number of demos.
+        """
+        return len(self.demos)
+
+    def __repr__(self) -> str:
+        """String representation of the index.
+
+        Returns:
+            String representation.
+        """
+        status = "fitted" if self._is_fitted else "not fitted"
+        return f"DemoIndex({len(self.demos)} demos, {status})"

openadapt_ml/retrieval/retriever.py

@@ -0,0 +1,162 @@
+"""Demo retriever for finding similar demonstrations."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import List, Optional
+
+from openadapt_ml.retrieval.index import DemoIndex, DemoMetadata
+from openadapt_ml.schema import Episode
+
+
+@dataclass
+class RetrievalResult:
+    """A single retrieval result with score.
+
+    Attributes:
+        demo: The demo metadata.
+        score: Retrieval score (higher is better).
+        text_score: Text similarity component.
+        domain_bonus: Domain match bonus applied.
+    """
+
+    demo: DemoMetadata
+    score: float
+    text_score: float
+    domain_bonus: float
+
+
+class DemoRetriever:
+    """Retrieves top-K similar demonstrations from an index.
+
+    Uses text similarity (TF-IDF cosine) with optional domain match bonus.
+    """
+
+    def __init__(
+        self,
+        index: DemoIndex,
+        domain_bonus: float = 0.2,
+    ) -> None:
+        """Initialize the retriever.
+
+        Args:
+            index: DemoIndex to retrieve from.
+            domain_bonus: Bonus score for domain match (default: 0.2).
+
+        Raises:
+            ValueError: If index is empty or not fitted.
+        """
+        if index.is_empty():
+            raise ValueError("Cannot create retriever from empty index")
+        if not index.is_fitted():
+            raise ValueError(
+                "Index must be built before retrieval (call index.build())"
+            )
+
+        self.index = index
+        self.domain_bonus = domain_bonus
+
+    def _compute_score(
+        self,
+        task: str,
+        demo: DemoMetadata,
+        app_context: Optional[str] = None,
+    ) -> RetrievalResult:
+        """Compute retrieval score for a demo.
+
+        Args:
+            task: Task description to match against.
+            demo: Demo metadata to score.
+            app_context: Optional app/domain context for bonus.
+
+        Returns:
+            RetrievalResult with computed scores.
+        """
+        # Text similarity using TF-IDF
+        query_embedding = self.index.embedder.embed(task)
+        text_score = self.index.embedder.cosine_similarity(
+            query_embedding,
+            demo.text_embedding,
+        )
+
+        # Domain match bonus
+        bonus = 0.0
+        if app_context is not None:
+            # Check if app_context matches app_name or domain
+            app_match = demo.app_name and app_context.lower() in demo.app_name.lower()
+            domain_match = demo.domain and app_context.lower() in demo.domain.lower()
+
+            if app_match or domain_match:
+                bonus = self.domain_bonus
+
+        # Final score is text similarity + bonus
+        total_score = text_score + bonus
+
+        return RetrievalResult(
+            demo=demo,
+            score=total_score,
+            text_score=text_score,
+            domain_bonus=bonus,
+        )
+
+    def retrieve(
+        self,
+        task: str,
+        app_context: Optional[str] = None,
+        top_k: int = 3,
+    ) -> List[Episode]:
+        """Retrieve top-K most similar demos.
+
+        Args:
+            task: Task description to find demos for.
+            app_context: Optional app/domain context (e.g., "Chrome", "github.com").
+            top_k: Number of demos to retrieve.
+
+        Returns:
+            List of Episode objects, ordered by relevance (most similar first).
+        """
+        if self.index.is_empty():
+            return []
+
+        # Score all demos
+        results = [
+            self._compute_score(task, demo, app_context)
+            for demo in self.index.get_all_demos()
+        ]
+
+        # Sort by score (descending)
+        results.sort(key=lambda r: r.score, reverse=True)
+
+        # Return top-K episodes
+        top_results = results[:top_k]
+        return [r.demo.episode for r in top_results]
+
+    def retrieve_with_scores(
+        self,
+        task: str,
+        app_context: Optional[str] = None,
+        top_k: int = 3,
+    ) -> List[RetrievalResult]:
+        """Retrieve top-K demos with their scores.
+
+        Args:
+            task: Task description to find demos for.
+            app_context: Optional app/domain context.
+            top_k: Number of demos to retrieve.
+
+        Returns:
+            List of RetrievalResult objects with scores.
+        """
+        if self.index.is_empty():
+            return []
+
+        # Score all demos
+        results = [
+            self._compute_score(task, demo, app_context)
+            for demo in self.index.get_all_demos()
+        ]
+
+        # Sort by score (descending)
+        results.sort(key=lambda r: r.score, reverse=True)
+
+        return results[:top_k]

openadapt_ml/runtime/__init__.py

@@ -0,0 +1,50 @@
+"""
+Runtime module for GUI automation agents.
+
+This module provides:
+- AgentPolicy: Runtime policy wrapper for VLM-based action prediction
+- SafetyGate: Deterministic safety checks for action validation
+
+Example usage:
+    from openadapt_ml.runtime import AgentPolicy, SafetyGate, SafetyConfig, SafetyDecision
+
+    # Create policy and safety gate
+    policy = AgentPolicy(adapter)
+    gate = SafetyGate(SafetyConfig(confidence_threshold=0.8))
+
+    # In agent loop
+    action, thought, state, raw = policy.predict_action_from_sample(sample)
+    assessment = gate.assess(action, observation)
+
+    if assessment.decision == SafetyDecision.ALLOW:
+        execute(action)
+    elif assessment.decision == SafetyDecision.REQUIRE_CONFIRMATION:
+        if user_confirms():
+            execute(action)
+    else:  # BLOCK
+        log_blocked(assessment.reason)
+"""
+
+from openadapt_ml.runtime.policy import (
+    AgentPolicy,
+    PolicyOutput,
+    parse_thought_state_action,
+)
+from openadapt_ml.runtime.safety_gate import (
+    SafetyAssessment,
+    SafetyConfig,
+    SafetyDecision,
+    SafetyGate,
+)
+
+__all__ = [
+    # Policy
+    "AgentPolicy",
+    "PolicyOutput",
+    "parse_thought_state_action",
+    # Safety Gate
+    "SafetyGate",
+    "SafetyConfig",
+    "SafetyDecision",
+    "SafetyAssessment",
+]

openadapt_ml/runtime/policy.py

@@ -3,16 +3,16 @@ from __future__ import annotations
 import json
 import re
 from dataclasses import dataclass
-from typing import Any, Dict, List, Optional, Tuple
+from typing import Any, Dict, Optional, Tuple
 
 from PIL import Image
 
 from openadapt_ml.models.base_adapter import BaseVLMAdapter
-from openadapt_ml.schemas.sessions import Action
+from openadapt_ml.schema import Action, ActionType, UIElement
 
 
 # Coordinate-based DSL patterns
-_CLICK_RE = re.compile(r"CLICK\(x=([0-9]*\.?[0-9]+),\s*y=([0-9]*\.?[0-9]+)\)")
+_CLICK_RE = re.compile(r"CLICK\(x=(-?[0-9]*\.?[0-9]+),\s*y=(-?[0-9]*\.?[0-9]+)\)")
 _TYPE_RE = re.compile(r'TYPE\(text="([^"\\]*(?:\\.[^"\\]*)*)"\)')
 _WAIT_RE = re.compile(r"\bWAIT\s*\(\s*\)")
 _DONE_RE = re.compile(r"\bDONE\s*\(\s*\)")
@@ -26,13 +26,16 @@ _TYPE_SOM_SIMPLE_RE = re.compile(r'TYPE\(["\']([^"\']*(?:\\.[^"\']*)*)["\']\)')
 @dataclass
 class PolicyOutput:
     """Result of a single policy step."""
+
     action: Action
     thought: Optional[str] = None
     state: Optional[Dict[str, Any]] = None
     raw_text: str = ""
 
 
-def parse_thought_state_action(text: str) -> Tuple[Optional[str], Optional[Dict[str, Any]], str]:
+def parse_thought_state_action(
+    text: str,
+) -> Tuple[Optional[str], Optional[Dict[str, Any]], str]:
     """Parse Thought / State / Action blocks from model output.
 
     Expected format:
@@ -54,12 +57,18 @@ def parse_thought_state_action(text: str) -> Tuple[Optional[str], Optional[Dict[
     action_str: str = text.strip()
 
     # Extract Thought - find the LAST occurrence (model's response, not template)
-    thought_matches = list(re.finditer(r"Thought:\s*(.+?)(?=State:|Action:|$)", text, re.DOTALL | re.IGNORECASE))
+    thought_matches = list(
+        re.finditer(
+            r"Thought:\s*(.+?)(?=State:|Action:|$)", text, re.DOTALL | re.IGNORECASE
+        )
+    )
     if thought_matches:
         thought = thought_matches[-1].group(1).strip()
 
     # Extract State (JSON on same line or next line) - last occurrence
-    state_matches = list(re.finditer(r"State:\s*(\{.*?\})", text, re.DOTALL | re.IGNORECASE))
+    state_matches = list(
+        re.finditer(r"State:\s*(\{.*?\})", text, re.DOTALL | re.IGNORECASE)
+    )
     if state_matches:
         try:
             state = json.loads(state_matches[-1].group(1))
@@ -119,7 +128,7 @@ class AgentPolicy:
         m = _CLICK_SOM_RE.search(text)
         if m:
             idx = int(m.group(1))
-            return Action(type="click", element_index=idx)
+            return Action(type=ActionType.CLICK, element=UIElement(element_id=str(idx)))
 
         # TYPE([N], "text")
         m = _TYPE_SOM_RE.search(text)
@@ -127,14 +136,18 @@ class AgentPolicy:
             idx = int(m.group(1))
             raw_text = m.group(2)
             unescaped = raw_text.replace('\\"', '"').replace("\\\\", "\\")
-            return Action(type="type", text=unescaped, element_index=idx)
+            return Action(
+                type=ActionType.TYPE,
+                text=unescaped,
+                element=UIElement(element_id=str(idx)),
+            )
 
         # TYPE("text") - SoM style without index
         m = _TYPE_SOM_SIMPLE_RE.search(text)
         if m:
             raw_text = m.group(1)
             unescaped = raw_text.replace('\\"', '"').replace("\\\\", "\\")
-            return Action(type="type", text=unescaped)
+            return Action(type=ActionType.TYPE, text=unescaped)
 
         # Coordinate-based patterns
         # CLICK(x=..., y=...)
@@ -145,7 +158,7 @@ class AgentPolicy:
             # Clamp to [0, 1]
             x = max(0.0, min(1.0, x))
             y = max(0.0, min(1.0, y))
-            return Action(type="click", x=x, y=y)
+            return Action(type=ActionType.CLICK, normalized_coordinates=(x, y))
 
         # TYPE(text="...")
         m = _TYPE_RE.search(text)
@@ -153,18 +166,18 @@ class AgentPolicy:
             # Unescape the text content
             raw_text = m.group(1)
             unescaped = raw_text.replace('\\"', '"').replace("\\\\", "\\")
-            return Action(type="type", text=unescaped)
+            return Action(type=ActionType.TYPE, text=unescaped)
 
         # WAIT()
         if _WAIT_RE.search(text):
-            return Action(type="wait")
+            return Action(type=ActionType.WAIT)
 
         # DONE()
         if _DONE_RE.search(text):
-            return Action(type="done")
+            return Action(type=ActionType.DONE)
 
         # Fallback
-        return Action(type="failed", raw={"text": text})
+        return Action(type=ActionType.FAIL, raw={"text": text})
 
     def predict_action_from_sample(
         self, sample: Dict[str, Any], max_new_tokens: int = 150