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

openadapt_ml/retrieval/README.md

@@ -0,0 +1,226 @@
+# Demo Retrieval Module
+
+This module provides functionality to index and retrieve similar demonstrations for few-shot prompting in GUI automation.
+
+## Overview
+
+The retrieval module consists of three main components:
+
+1. **TextEmbedder** (`embeddings.py`) - Simple TF-IDF based text embeddings
+2. **DemoIndex** (`index.py`) - Stores episodes with metadata and embeddings
+3. **DemoRetriever** (`retriever.py`) - Retrieves top-K similar demos
+
+## Quick Start
+
+```python
+from openadapt_ml.retrieval import DemoIndex, DemoRetriever
+from openadapt_ml.schema import Episode
+
+# 1. Create index and add episodes
+index = DemoIndex()
+index.add_many(episodes)  # episodes is a list of Episode objects
+index.build()  # Compute embeddings
+
+# 2. Create retriever
+retriever = DemoRetriever(index, domain_bonus=0.2)
+
+# 3. Retrieve similar demos
+task = "Turn off Night Shift on macOS"
+app_context = "System Settings"
+similar_demos = retriever.retrieve(task, app_context, top_k=3)
+
+# 4. Use with prompt formatting
+from openadapt_ml.experiments.demo_prompt.format_demo import format_episode_as_demo
+formatted_demo = format_episode_as_demo(similar_demos[0])
+```
+
+## Features
+
+### Text Similarity
+- Uses TF-IDF with cosine similarity for v1
+- No external ML libraries required
+- Can be upgraded to sentence-transformers later
+
+### Domain Matching
+- Auto-extracts app name from observations
+- Auto-extracts domain from URLs
+- Applies bonus score for domain/app matches
+
+### Metadata Support
+- Stores arbitrary metadata with each demo
+- Tracks app name, domain, and custom fields
+- Efficient filtering by app/domain
+
+## API Reference
+
+### DemoIndex
+
+```python
+index = DemoIndex()
+
+# Add episodes
+index.add(episode, app_name="Chrome", domain="github.com")
+index.add_many(episodes)
+
+# Build index (required before retrieval)
+index.build()
+
+# Query index
+index.get_apps()      # List of unique app names
+index.get_domains()   # List of unique domains
+len(index)            # Number of demos
+index.is_fitted()     # Check if built
+```
+
+### DemoRetriever
+
+```python
+retriever = DemoRetriever(
+    index,
+    domain_bonus=0.2,  # Bonus score for domain match
+)
+
+# Retrieve episodes
+episodes = retriever.retrieve(
+    task="Description of task",
+    app_context="Chrome",  # Optional
+    top_k=3,
+)
+
+# Retrieve with scores (for debugging)
+results = retriever.retrieve_with_scores(task, app_context, top_k=3)
+for result in results:
+    print(f"Score: {result.score}")
+    print(f"  Text similarity: {result.text_score}")
+    print(f"  Domain bonus: {result.domain_bonus}")
+    print(f"  Goal: {result.demo.episode.goal}")
+```
+
+### TextEmbedder
+
+```python
+from openadapt_ml.retrieval.embeddings import TextEmbedder
+
+embedder = TextEmbedder()
+
+# Fit on corpus
+documents = ["task 1", "task 2", "task 3"]
+embedder.fit(documents)
+
+# Embed text
+vec1 = embedder.embed("new task")
+vec2 = embedder.embed("another task")
+
+# Compute similarity
+similarity = embedder.cosine_similarity(vec1, vec2)
+```
+
+## Scoring
+
+The retrieval score combines text similarity and domain matching:
+
+```
+total_score = text_similarity + domain_bonus
+```
+
+- **Text similarity**: TF-IDF cosine similarity between task descriptions (0-1)
+- **Domain bonus**: Fixed bonus if app_context matches demo's app or domain (default: 0.2)
+
+### Example Scores
+
+```
+Query: "Search GitHub for ML papers"
+App context: "github.com"
+
+Demo 1: "Search for machine learning papers on GitHub"
+  - Text similarity: 0.678
+  - Domain bonus: 0.200 (github.com match)
+  - Total: 0.878 ⭐ Best match
+
+Demo 2: "Create a new GitHub repository"
+  - Text similarity: 0.111
+  - Domain bonus: 0.200 (github.com match)
+  - Total: 0.311
+
+Demo 3: "Search for Python documentation on Google"
+  - Text similarity: 0.232
+  - Domain bonus: 0.000 (no match)
+  - Total: 0.232
+```
+
+## Loading Real Episodes
+
+```python
+from openadapt_ml.ingest.capture import load_capture
+from openadapt_ml.retrieval import DemoIndex, DemoRetriever
+
+# Load from capture directory
+capture_path = "/path/to/capture"
+episodes = load_capture(capture_path)
+
+# Build index
+index = DemoIndex()
+index.add_many(episodes)
+index.build()
+
+# Retrieve
+retriever = DemoRetriever(index)
+demos = retriever.retrieve("New task description", top_k=3)
+```
+
+## Integration with Prompting
+
+```python
+from openadapt_ml.experiments.demo_prompt.format_demo import format_episode_as_demo
+
+# Retrieve demo
+demos = retriever.retrieve(task, app_context, top_k=1)
+
+# Format for prompt
+demo_text = format_episode_as_demo(demos[0], max_steps=10)
+
+# Inject into prompt
+prompt = f"""Here is a demonstration of a similar task:
+
+{demo_text}
+
+Now perform this task:
+Task: {task}
+"""
+```
+
+## Examples
+
+See `examples/demo_retrieval_example.py` for a complete working example.
+
+Run it with:
+```bash
+uv run python examples/demo_retrieval_example.py
+```
+
+## Future Improvements
+
+### v2: Better Embeddings
+Replace TF-IDF with sentence-transformers:
+```python
+from sentence_transformers import SentenceTransformer
+model = SentenceTransformer('all-MiniLM-L6-v2')
+```
+
+### v3: Semantic Search
+- Use FAISS or Qdrant for large-scale retrieval
+- Add metadata filtering before similarity search
+- Support multi-modal embeddings (text + screenshots)
+
+### v4: Learning to Rank
+- Train a ranking model using success/failure data
+- Incorporate user feedback
+- Personalized retrieval based on agent history
+
+## Design Principles
+
+1. **Start simple** - v1 uses no ML models, just text matching
+2. **Functional over optimal** - Works out of the box, can be improved later
+3. **Clear API** - Simple retrieve() interface, complex details hidden
+4. **Composable** - Each component can be used independently
+5. **Schema-first** - Works with Episode schema, no custom data structures

openadapt_ml/retrieval/USAGE.md

@@ -0,0 +1,391 @@
+# Demo Retrieval - Usage Guide
+
+## Quick Reference
+
+```python
+# 1. Build index
+from openadapt_ml.retrieval import DemoIndex, DemoRetriever
+index = DemoIndex()
+index.add_many(episodes)
+index.build()
+
+# 2. Retrieve
+retriever = DemoRetriever(index)
+similar_demos = retriever.retrieve("Turn off Night Shift", top_k=3)
+```
+
+## Complete Examples
+
+### Example 1: Basic Usage with Synthetic Data
+
+```python
+from openadapt_ml.retrieval import DemoIndex, DemoRetriever
+from openadapt_ml.schema import Action, ActionType, Episode, Observation, Step
+
+# Create test episodes
+def create_episode(instruction, app_name=None):
+    obs = Observation(app_name=app_name)
+    action = Action(type=ActionType.CLICK, normalized_coordinates=(0.5, 0.5))
+    step = Step(step_index=0, observation=obs, action=action)
+    return Episode(episode_id=f"ep_{instruction[:10]}", instruction=instruction, steps=[step])
+
+episodes = [
+    create_episode("Turn off Night Shift", app_name="System Settings"),
+    create_episode("Search GitHub", app_name="Chrome"),
+    create_episode("Open calculator", app_name="Calculator"),
+]
+
+# Build index
+index = DemoIndex()
+index.add_many(episodes)
+index.build()
+
+# Retrieve
+retriever = DemoRetriever(index, domain_bonus=0.2)
+results = retriever.retrieve("Disable Night Shift", top_k=2)
+
+print(f"Found {len(results)} similar demos:")
+for ep in results:
+    print(f"- {ep.goal}")
+```
+
+### Example 2: Loading from Capture
+
+```python
+from openadapt_ml.ingest.capture import capture_to_episode
+from openadapt_ml.retrieval import DemoIndex, DemoRetriever
+
+# Load multiple captures
+capture_paths = [
+    "/path/to/capture1",
+    "/path/to/capture2",
+    "/path/to/capture3",
+]
+
+episodes = [
+    capture_to_episode(path, include_moves=False)
+    for path in capture_paths
+]
+
+# Build index
+index = DemoIndex()
+index.add_many(episodes)
+index.build()
+
+# Retrieve for new task
+retriever = DemoRetriever(index)
+task = "Turn on dark mode"
+app = "System Settings"
+demos = retriever.retrieve(task, app_context=app, top_k=3)
+```
+
+### Example 3: Integration with Prompting
+
+```python
+from openadapt_ml.experiments.demo_prompt.format_demo import format_episode_as_demo
+from openadapt_ml.retrieval import DemoIndex, DemoRetriever
+
+# Build index (assume episodes already loaded)
+index = DemoIndex()
+index.add_many(episodes)
+index.build()
+
+# Retrieve for new task
+retriever = DemoRetriever(index)
+task = "Turn off Night Shift"
+demos = retriever.retrieve(task, top_k=1)
+
+# Format for prompt
+if demos:
+    demo_text = format_episode_as_demo(demos[0], max_steps=10)
+
+    # Create few-shot prompt
+    prompt = f"""You are a GUI automation agent.
+
+DEMONSTRATION OF SIMILAR TASK:
+{demo_text}
+
+NEW TASK:
+{task}
+
+What is your first action?"""
+
+    print(prompt)
+```
+
+### Example 4: Retrieval with Scores (Debugging)
+
+```python
+from openadapt_ml.retrieval import DemoRetriever
+
+retriever = DemoRetriever(index, domain_bonus=0.3)
+
+# Retrieve with scores for analysis
+results = retriever.retrieve_with_scores(
+    task="Search for Python docs",
+    app_context="github.com",
+    top_k=5,
+)
+
+# Analyze scores
+for i, result in enumerate(results, 1):
+    print(f"\n{i}. {result.demo.episode.goal}")
+    print(f"   Total score: {result.score:.3f}")
+    print(f"   Text similarity: {result.text_score:.3f}")
+    print(f"   Domain bonus: {result.domain_bonus:.3f}")
+
+    if result.demo.app_name:
+        print(f"   App: {result.demo.app_name}")
+    if result.demo.domain:
+        print(f"   Domain: {result.demo.domain}")
+```
+
+### Example 5: Custom Metadata
+
+```python
+from openadapt_ml.retrieval import DemoIndex
+
+index = DemoIndex()
+
+# Add episodes with custom metadata
+for episode in episodes:
+    metadata = {
+        "difficulty": "easy",
+        "success_rate": 0.95,
+        "duration_seconds": 30,
+        "tags": ["settings", "macOS"],
+    }
+
+    index.add(
+        episode,
+        app_name="System Settings",
+        domain=None,
+        metadata=metadata,
+    )
+
+index.build()
+
+# Access metadata after retrieval
+retriever = DemoRetriever(index)
+results = retriever.retrieve_with_scores("Turn off Night Shift", top_k=1)
+
+if results:
+    demo = results[0].demo
+    print(f"Difficulty: {demo.metadata.get('difficulty')}")
+    print(f"Tags: {demo.metadata.get('tags')}")
+```
+
+## CLI Examples
+
+Run the provided example scripts:
+
+```bash
+# Basic demo with synthetic data
+uv run python examples/demo_retrieval_example.py
+
+# Test with real capture
+uv run python examples/retrieval_with_capture.py /path/to/capture
+
+# With custom task
+uv run python examples/retrieval_with_capture.py /path/to/capture "Turn off dark mode"
+```
+
+## Common Patterns
+
+### Pattern 1: Multi-Domain Index
+
+```python
+# Build index with episodes from multiple domains
+web_episodes = load_web_captures()
+desktop_episodes = load_desktop_captures()
+
+index = DemoIndex()
+index.add_many(web_episodes)
+index.add_many(desktop_episodes)
+index.build()
+
+# Retrieve with domain filtering via app_context
+retriever = DemoRetriever(index, domain_bonus=0.5)
+
+# This will prefer github.com demos
+web_demos = retriever.retrieve("Search code", app_context="github.com", top_k=3)
+
+# This will prefer System Settings demos
+desktop_demos = retriever.retrieve("Change settings", app_context="System Settings", top_k=3)
+```
+
+### Pattern 2: Incremental Index Updates
+
+```python
+# Build initial index
+index = DemoIndex()
+index.add_many(initial_episodes)
+index.build()
+
+# Add new episodes
+index.add(new_episode)
+
+# Rebuild required after adding
+index.build()
+
+# Now retriever will use updated index
+retriever = DemoRetriever(index)
+```
+
+### Pattern 3: Batch Retrieval
+
+```python
+# Retrieve for multiple tasks
+tasks = [
+    "Turn off Night Shift",
+    "Enable dark mode",
+    "Adjust brightness",
+]
+
+retriever = DemoRetriever(index)
+
+for task in tasks:
+    demos = retriever.retrieve(task, top_k=3)
+    print(f"\nTask: {task}")
+    for demo in demos:
+        print(f"  - {demo.goal}")
+```
+
+## Tuning Parameters
+
+### Domain Bonus
+
+Controls how much to favor domain/app matches:
+
+```python
+# No domain bonus - pure text similarity
+retriever = DemoRetriever(index, domain_bonus=0.0)
+
+# Small bonus (default)
+retriever = DemoRetriever(index, domain_bonus=0.2)
+
+# Large bonus - heavily favor same domain
+retriever = DemoRetriever(index, domain_bonus=0.5)
+```
+
+**Rule of thumb:**
+- `0.0-0.1`: When task text is very specific and domain doesn't matter much
+- `0.2-0.3`: Good default for most cases
+- `0.4-0.5`: When domain matching is critical (e.g., domain-specific workflows)
+
+### Top-K
+
+Number of demos to retrieve:
+
+```python
+# Single best match
+demos = retriever.retrieve(task, top_k=1)
+
+# Few-shot with 3 examples
+demos = retriever.retrieve(task, top_k=3)
+
+# Retrieve more for analysis/selection
+demos = retriever.retrieve(task, top_k=10)
+```
+
+**Rule of thumb:**
+- `top_k=1`: When prompt length is constrained
+- `top_k=3`: Good default for few-shot learning
+- `top_k=5+`: For ensemble methods or human selection
+
+## Performance Tips
+
+### 1. Build Once, Retrieve Many
+
+```python
+# Good: Build once
+index.build()
+retriever = DemoRetriever(index)
+for task in many_tasks:
+    retriever.retrieve(task)
+
+# Bad: Build repeatedly
+for task in many_tasks:
+    index.build()  # Wasteful!
+    retriever = DemoRetriever(index)
+    retriever.retrieve(task)
+```
+
+### 2. Pre-extract Metadata
+
+```python
+# Good: Extract once when adding
+index.add(episode, app_name="Chrome", domain="github.com")
+
+# Less efficient: Let auto-extraction scan every episode
+index.add(episode)  # Will scan steps for app_name and domain
+```
+
+### 3. Filter Before Retrieval
+
+```python
+# If you have a large index but know the domain, create a filtered index
+web_demos = [d for d in index.get_all_demos() if d.domain]
+web_index = DemoIndex()
+for demo in web_demos:
+    web_index.add(demo.episode)
+web_index.build()
+```
+
+## Troubleshooting
+
+### Issue: All scores are 0.0
+
+**Cause:** Only one episode in index, so IDF is undefined.
+
+**Solution:** Add more episodes or use a larger demo library.
+
+```python
+# Need at least 2-3 episodes for meaningful scores
+assert len(index) >= 3, "Add more demos to the index"
+```
+
+### Issue: Domain bonus not applied
+
+**Cause:** app_context doesn't match app_name or domain.
+
+**Debug:**
+```python
+results = retriever.retrieve_with_scores(task, app_context, top_k=5)
+for r in results:
+    print(f"App: {r.demo.app_name}, Domain: {r.demo.domain}, Bonus: {r.domain_bonus}")
+```
+
+**Solution:** Check exact string matching (case-insensitive contains).
+
+### Issue: Poor retrieval quality
+
+**Causes:**
+1. Task descriptions too generic
+2. Demo library too small
+3. TF-IDF limitations
+
+**Solutions:**
+1. Use more specific task descriptions
+2. Add more diverse demos to index
+3. Upgrade to sentence-transformers (see README.md § Future Improvements)
+
+## Testing
+
+Run unit tests:
+```bash
+uv run pytest tests/test_retrieval.py -v
+```
+
+Run integration tests:
+```bash
+uv run python test_retrieval.py
+```
+
+## Next Steps
+
+1. **Integrate with training**: Use retrieval in data augmentation
+2. **Experiment with prompting**: Test different demo counts and formats
+3. **Upgrade embeddings**: Try sentence-transformers for better similarity
+4. **Add filtering**: Support domain/app filtering before similarity scoring
+5. **Evaluate impact**: Measure action accuracy with/without retrieval

openadapt_ml/retrieval/__init__.py

@@ -0,0 +1,91 @@
+"""Demo retrieval module for finding similar demonstrations.
+
+This module provides functionality for indexing and retrieving demonstrations
+based on semantic similarity of task descriptions.
+
+Main Components:
+- DemoRetriever: Main class for indexing and retrieving demos
+- Embedders: TFIDFEmbedder, SentenceTransformerEmbedder, OpenAIEmbedder
+- DemoIndex: Legacy index class (use DemoRetriever instead)
+
+Quick Start:
+    from openadapt_ml.retrieval import DemoRetriever
+    from openadapt_ml.schema import Episode
+
+    # Create retriever (TF-IDF is default, no external dependencies)
+    retriever = DemoRetriever()
+
+    # Or use sentence-transformers for better semantic matching
+    retriever = DemoRetriever(embedding_method="sentence_transformers")
+
+    # Add demos
+    retriever.add_demo(episode1)
+    retriever.add_demo(episode2, app_name="Chrome", domain="github.com")
+
+    # Build index (required before retrieval)
+    retriever.build_index()
+
+    # Retrieve similar demos
+    results = retriever.retrieve("Turn off Night Shift", top_k=3)
+
+    # Format for inclusion in a prompt
+    prompt_text = retriever.format_for_prompt(results)
+
+Embedding Methods:
+    # TF-IDF (default, no dependencies)
+    retriever = DemoRetriever(embedding_method="tfidf")
+
+    # Sentence Transformers (recommended, requires: pip install sentence-transformers)
+    retriever = DemoRetriever(
+        embedding_method="sentence_transformers",
+        embedding_model="all-MiniLM-L6-v2",  # Fast, 22MB
+    )
+
+    # OpenAI (requires: pip install openai, OPENAI_API_KEY env var)
+    retriever = DemoRetriever(
+        embedding_method="openai",
+        embedding_model="text-embedding-3-small",
+    )
+
+See Also:
+    - docs/demo_retrieval_design.md - Full design document
+    - openadapt_ml/experiments/demo_prompt/ - Demo-conditioned prompting
+"""
+
+# Main retrieval class (recommended)
+from openadapt_ml.retrieval.demo_retriever import (
+    DemoRetriever,
+    DemoMetadata,
+    RetrievalResult,
+)
+
+# Embedders
+from openadapt_ml.retrieval.embeddings import (
+    BaseEmbedder,
+    TFIDFEmbedder,
+    TextEmbedder,  # Alias for TFIDFEmbedder (backward compat)
+    SentenceTransformerEmbedder,
+    OpenAIEmbedder,
+    create_embedder,
+)
+
+# Legacy classes (for backward compatibility)
+from openadapt_ml.retrieval.index import DemoIndex
+from openadapt_ml.retrieval.retriever import DemoRetriever as LegacyDemoRetriever
+
+__all__ = [
+    # Main classes
+    "DemoRetriever",
+    "DemoMetadata",
+    "RetrievalResult",
+    # Embedders
+    "BaseEmbedder",
+    "TFIDFEmbedder",
+    "TextEmbedder",
+    "SentenceTransformerEmbedder",
+    "OpenAIEmbedder",
+    "create_embedder",
+    # Legacy (backward compat)
+    "DemoIndex",
+    "LegacyDemoRetriever",
+]