> ## Documentation Index
> Fetch the complete documentation index at: https://docs.orpheus.run/llms.txt
> Use this file to discover all available pages before exploring further.

# Embedding Cache Agent

> High-performance agent that keeps an embedding model loaded in memory for fast vector generation.

# Embedding Cache Agent

This agent demonstrates **memory-intensive workloads** and **warm-start performance**. It loads a heavy embedding model (e.g., SentenceTransformers) into memory once during startup and reuses it for all subsequent requests.

## How it Works

1. **Startup:** Loads the model (simulated 2GB model) into global memory. This happens only once when the worker starts.
2. **Request:** Receives text input.
3. **Processing:** Generates vector embeddings instantly using the pre-loaded model.
4. **Output:** Returns the vector array.

## Key Features

* **High Memory Config:** Uses `memory: 4096` (4GB) to accommodate the model.
* **Zero Cold Starts:** Subsequent requests are millisecond-fast because the model stays in memory.
* **Singleton Pattern:** Uses a global variable to cache the model across requests.

## Usage

```bash theme={null}
orpheus run embedding-cache '{"text": "Hello world"}'
```

## Source Code

<Tabs>
  <Tab title="agent.yaml">
    ```yaml theme={null}
    name: embedding-cache
    runtime: python3
    module: agent
    entrypoint: handler

    memory: 512  # CONFIGURABLE for different OOM tests
    timeout: 300

    # ServiceManager integration
    model: tinyllama
    engine: ollama  # Using Ollama for CPU testing (no rate limits)

    env:
      - COLLECTIONS_DIR=/agent/data/collections
      - MAX_COLLECTIONS=2

    # Telemetry configuration - custom labels for Prometheus filtering
    telemetry:
      enabled: true
      labels:
        team: ml_platform
        tier: memory_intensive
        use_case: embedding

    scaling:
      min_workers: 1
      max_workers: 5
      target_utilization: 1.5
      scale_up_threshold: 2.0
      scale_down_threshold: 0.3

    ```
  </Tab>

  <Tab title="agent.py">
    ```python theme={null}
    #!/usr/bin/env python3
    """
    Embedding Cache Agent for Orpheus: Memory-Intensive ML Workload (OOM Testing).

    This agent demonstrates Orpheus's OOM detection and recovery:

    1. Real ML Memory Patterns:
       - Loads sentence-transformers embedding model (~400MB)
       - Loads multiple FAISS indices (~150MB per 10K docs)
       - Tests actual memory pressure from ML libraries

    2. OOM Detection (exit 137):
       - ServiceManager detects OOM kills
       - Applies exponential backoff (60s+)
       - Auto-replaces worker after 3 failures

    3. Memory Tracking:
       - Uses psutil to track real RSS memory
       - Logs memory before/after each load
       - Returns actual memory stats

    The agent supports actions:
    - load: Load collections (triggers OOM if over limit)
    - search: Search loaded collections
    - status: Return memory usage and loaded collections
    """

    import os
    import sys
    import time
    import psutil
    from pathlib import Path
    from typing import Dict, List

    # RAG dependencies (same as rag-agent)
    from langchain_huggingface import HuggingFaceEmbeddings
    from langchain_community.vectorstores import FAISS
    from langchain_ollama import OllamaLLM
    from langchain_core.prompts import PromptTemplate

    # Configuration
    COLLECTIONS_DIR = os.getenv("COLLECTIONS_DIR", "./data/collections")
    MAX_COLLECTIONS = int(os.getenv("MAX_COLLECTIONS", "2"))
    EMBEDDING_MODEL = "sentence-transformers/all-MiniLM-L6-v2"

    # Model endpoint - ServiceManager injects MODEL_URL automatically
    MODEL_URL = os.getenv("MODEL_URL") or os.getenv("OLLAMA_BASE_URL", "http://localhost:11434")
    MODEL_NAME = os.getenv("OLLAMA_MODEL", "mistral")


    class EmbeddingCache:
        """Manages multiple FAISS indices with memory tracking."""

        def __init__(self):
            self.embedding_model = None  # HuggingFaceEmbeddings
            self.collections = {}  # Dict[str, FAISS]
            self.max_collections = MAX_COLLECTIONS
            self.llm = None
            self._load_times = {}

        def get_memory_usage_mb(self) -> float:
            """Get current process memory usage in MB."""
            process = psutil.Process()
            return process.memory_info().rss / 1024 / 1024

        def load_embedding_model(self):
            """Load sentence-transformers embedding model (~400MB)."""
            if self.embedding_model is not None:
                return  # Already loaded

            print(f"[cache] Loading embedding model: {EMBEDDING_MODEL}", file=sys.stderr, flush=True)
            print(f"[cache] This will use approximately 400MB of memory...", file=sys.stderr, flush=True)

            mem_before = self.get_memory_usage_mb()
            start_time = time.time()

            self.embedding_model = HuggingFaceEmbeddings(
                model_name=EMBEDDING_MODEL,
                model_kwargs={"device": "cpu"},
                encode_kwargs={"normalize_embeddings": True},
            )

            elapsed = time.time() - start_time
            mem_after = self.get_memory_usage_mb()
            mem_delta = mem_after - mem_before

            print(
                f"[cache] Embedding model loaded: {mem_delta:.1f}MB in {elapsed:.2f}s (total: {mem_after:.1f}MB)",
                file=sys.stderr,
                flush=True,
            )

            self._load_times["embedding_model"] = elapsed

        def load_collection(self, name: str) -> bool:
            """Load a FAISS index (~150MB per 10K docs)."""
            if name in self.collections:
                print(f"[cache] Collection '{name}' already loaded", file=sys.stderr, flush=True)
                return True

            # Ensure embedding model is loaded first
            if self.embedding_model is None:
                self.load_embedding_model()

            # Check if we're at max collections
            if len(self.collections) >= self.max_collections:
                print(
                    f"[cache] Max collections reached ({self.max_collections}), cannot load '{name}'",
                    file=sys.stderr,
                    flush=True,
                )
                return False

            # Build path to FAISS index
            collection_path = Path(COLLECTIONS_DIR) / name / "faiss_index"
            if not collection_path.exists():
                print(f"[cache] Collection '{name}' not found at {collection_path}", file=sys.stderr, flush=True)
                return False

            print(f"[cache] Loading collection '{name}' from {collection_path}", file=sys.stderr, flush=True)
            print(f"[cache] This will use approximately 150MB of memory...", file=sys.stderr, flush=True)

            mem_before = self.get_memory_usage_mb()
            start_time = time.time()

            vectorstore = FAISS.load_local(
                str(collection_path),
                self.embedding_model,
                allow_dangerous_deserialization=True,
            )

            elapsed = time.time() - start_time
            mem_after = self.get_memory_usage_mb()
            mem_delta = mem_after - mem_before

            self.collections[name] = vectorstore
            self._load_times[f"collection_{name}"] = elapsed

            print(
                f"[cache] Collection '{name}' loaded: {mem_delta:.1f}MB in {elapsed:.2f}s (total: {mem_after:.1f}MB)",
                file=sys.stderr,
                flush=True,
            )

            return True

        def load_all_available(self) -> List[str]:
            """Load all available collections up to MAX_COLLECTIONS limit."""
            collections_dir = Path(COLLECTIONS_DIR)
            if not collections_dir.exists():
                print(f"[cache] Collections directory not found: {collections_dir}", file=sys.stderr, flush=True)
                return []

            # Find all collection directories
            available = sorted([d.name for d in collections_dir.iterdir() if d.is_dir()])
            print(f"[cache] Found {len(available)} collections: {available}", file=sys.stderr, flush=True)

            loaded = []
            for name in available[: self.max_collections]:
                if self.load_collection(name):
                    loaded.append(name)

            return loaded

        def search_collection(self, collection_name: str, query: str, top_k: int = 3) -> dict:
            """Search a loaded collection."""
            if collection_name not in self.collections:
                return {
                    "status": "error",
                    "error": f"Collection '{collection_name}' not loaded. Available: {list(self.collections.keys())}",
                }

            vectorstore = self.collections[collection_name]

            start_time = time.time()
            docs = vectorstore.similarity_search(query, k=top_k)
            elapsed = time.time() - start_time

            return {
                "status": "success",
                "collection": collection_name,
                "query": query,
                "results": [
                    {"content": doc.page_content[:200] + "...", "metadata": doc.metadata} for doc in docs
                ],
                "retrieval_ms": int(elapsed * 1000),
            }

        def get_status(self) -> dict:
            """Get current cache status and memory usage."""
            mem_current = self.get_memory_usage_mb()

            return {
                "status": "success",
                "memory_mb": round(mem_current, 1),
                "max_collections": self.max_collections,
                "collections_loaded": list(self.collections.keys()),
                "embedding_model_loaded": self.embedding_model is not None,
                "load_times": {k: round(v, 2) for k, v in self._load_times.items()},
            }


    # Global cache instance (persists across invocations)
    _cache = None


    def get_cache() -> EmbeddingCache:
        """Get or create the global cache instance."""
        global _cache
        if _cache is None:
            _cache = EmbeddingCache()
        return _cache


    def handler(input_data: dict) -> dict:
        """
        Orpheus agent handler for embedding cache operations.

        Input:
            {
                "action": "load" | "search" | "status",
                "collection": "collection_1",  # for search
                "query": "search query",       # for search
                "top_k": 3                      # for search (optional)
            }

        Output:
            {
                "status": "success" | "error",
                "action": "load" | "search" | "status",
                "collections_loaded": [...],
                "memory_mb": 687.3,
                ...
            }
        """
        try:
            cache = get_cache()
            action = input_data.get("action", "load")

            print(f"[cache] Action: {action}", file=sys.stderr, flush=True)
            print(f"[cache] Current memory: {cache.get_memory_usage_mb():.1f}MB", file=sys.stderr, flush=True)

            if action == "load":
                # Load collections - this is where OOM can occur
                start_time = time.time()
                loaded = cache.load_all_available()
                elapsed = time.time() - start_time

                status = cache.get_status()
                status.update(
                    {
                        "action": "load",
                        "load_time_sec": round(elapsed, 2),
                        "collections_loaded": loaded,
                    }
                )

                print(f"[cache] Load complete: {len(loaded)} collections, {status['memory_mb']}MB", file=sys.stderr, flush=True)
                return status

            elif action == "search":
                collection = input_data.get("collection")
                query = input_data.get("query")
                top_k = input_data.get("top_k", 3)

                if not collection or not query:
                    return {
                        "status": "error",
                        "action": "search",
                        "error": "Missing 'collection' or 'query' field",
                    }

                result = cache.search_collection(collection, query, top_k)
                result["action"] = "search"
                result["memory_mb"] = round(cache.get_memory_usage_mb(), 1)

                return result

            elif action == "status":
                status = cache.get_status()
                status["action"] = "status"
                return status

            else:
                return {
                    "status": "error",
                    "error": f"Unknown action: {action}. Use 'load', 'search', or 'status'",
                }

        except Exception as e:
            print(f"[cache] ERROR: {e}", file=sys.stderr, flush=True)
            return {
                "status": "error",
                "action": input_data.get("action", "unknown"),
                "error": str(e),
            }

    ```
  </Tab>
</Tabs>