openadapt-ml 0.1.0__py3-none-any.whl → 0.2.0__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,160 @@
+"""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/policy.py

@@ -8,11 +8,11 @@ from typing import Any, Dict, List, 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*\)")
@@ -119,7 +119,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 +127,14 @@ 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 +145,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 +153,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

openadapt_ml/schema/__init__.py

@@ -0,0 +1,104 @@
+"""
+Episode Schema - Canonical format for GUI trajectory data.
+
+A standardized contract for representing GUI automation episodes, enabling
+interoperability across training pipelines, benchmarks, and recording tools.
+
+Installation:
+    pip install openadapt-ml
+    # or: uv add openadapt-ml
+
+Basic Usage:
+    from openadapt_ml.schema import Episode, Step, Action, Observation, ActionType
+
+    # Create an episode
+    episode = Episode(
+        episode_id="demo_001",
+        instruction="Open Notepad and type Hello World",
+        steps=[
+            Step(
+                step_index=0,
+                observation=Observation(screenshot_path="step_0.png"),
+                action=Action(type=ActionType.CLICK, coordinates={"x": 100, "y": 200}),
+            ),
+            Step(
+                step_index=1,
+                observation=Observation(screenshot_path="step_1.png"),
+                action=Action(type=ActionType.TYPE, text="Hello World"),
+            ),
+        ],
+        success=True,
+    )
+
+    # Save/load JSON
+    save_episode(episode, "episode.json")
+    episode = load_episode("episode.json")
+
+    # Validate external data
+    is_valid, error = validate_episode({"episode_id": "x", ...})
+
+Coordinate Systems:
+    # Pixel coordinates (absolute)
+    Action(type=ActionType.CLICK, coordinates={"x": 512, "y": 384})
+
+    # Normalized coordinates (0.0-1.0, resolution-independent)
+    Action(type=ActionType.CLICK, normalized_coordinates=(0.5, 0.375))
+
+    # Both can coexist - use whichever fits your pipeline
+
+Converting from Other Formats:
+    from openadapt_ml.schema.converters import from_waa_trajectory
+
+    # Convert Windows Agent Arena format
+    episode = from_waa_trajectory(trajectory_list, task_info_dict)
+
+    # Convert back
+    trajectory, task_info = to_waa_trajectory(episode)
+
+JSON Schema Export:
+    # For external validation tools (e.g., JSON Schema validators, TypeScript codegen)
+    export_json_schema("episode.schema.json")
+
+See Also:
+    - docs/schema/episode.schema.json - Full JSON Schema
+    - openadapt_ml.schema.episode - Model definitions
+    - openadapt_ml.schema.converters - Format converters
+"""
+
+from openadapt_ml.schema.episode import (
+    SCHEMA_VERSION,
+    Episode,
+    Step,
+    Action,
+    Observation,
+    ActionType,
+    BenchmarkSource,
+    Coordinates,
+    BoundingBox,
+    UIElement,
+    validate_episode,
+    load_episode,
+    save_episode,
+    export_json_schema,
+)
+
+__all__ = [
+    # Version
+    "SCHEMA_VERSION",
+    # Core models
+    "Episode",
+    "Step",
+    "Action",
+    "Observation",
+    # Supporting models
+    "ActionType",
+    "BenchmarkSource",
+    "Coordinates",
+    "BoundingBox",
+    "UIElement",
+    # Utilities
+    "validate_episode",
+    "load_episode",
+    "save_episode",
+    "export_json_schema",
+]