Embeddings Engine Skill

Role

You are an elite embeddings engineer. You know every major embedding model,

their trade-offs, optimal batch sizes, caching strategies, and how to choose

the right model for each task. Embeddings are the foundation of all semantic AI.

Part 1: Model Selection Guide

IMI Recommendation:

• Production with API budget: text-embedding-3-small (speed+cost balance)
• High-precision client deliverables: text-embedding-3-large
• Fully local/private: bge-large-en-v1.5 via sentence-transformers
• Long documents: voyage-large-2-instruct

Part 2: OpenAI Embeddings

from openai import OpenAI
import numpy as np
from typing import Union
import time, logging

logger = logging.getLogger(__name__)
client = OpenAI()  # Uses OPENAI_API_KEY env var

def embed_openai(
    texts: Union[str, list[str]],
    model: str = "text-embedding-3-small",
    dimensions: int = None,  # Only for text-embedding-3-*
    batch_size: int = 100,
    max_retries: int = 3
) -> list[list[float]]:
    """
    Generate embeddings using OpenAI API with batching and retry.
    Returns list of embedding vectors.
    """
    if isinstance(texts, str):
        texts = [texts]

    all_embeddings = []

    for i in range(0, len(texts), batch_size):
        batch = texts[i:i + batch_size]

        for attempt in range(max_retries):
            try:
                kwargs = {"input": batch, "model": model}
                if dimensions:
                    kwargs["dimensions"] = dimensions

                response = client.embeddings.create(**kwargs)
                batch_embeddings = [r.embedding for r in response.data]
                all_embeddings.extend(batch_embeddings)
                break

            except Exception as e:
                if attempt == max_retries - 1:
                    raise
                wait = 2 ** attempt
                logger.warning(f"Embed attempt {attempt+1} failed: {e}. Retrying in {wait}s")
                time.sleep(wait)

    return all_embeddings

Part 3: Local Embeddings (sentence-transformers)

# pip install sentence-transformers
from sentence_transformers import SentenceTransformer
import numpy as np
from functools import lru_cache

class LocalEmbedder:
    """
    Local embeddings using sentence-transformers.
    No API calls, no cost, fully private.
    Best models: bge-large-en-v1.5, all-mpnet-base-v2, nomic-embed-text
    """
    _instances = {}  # Model cache

    def __init__(self, model_name: str = "BAAI/bge-large-en-v1.5",
                  device: str = "cpu"):
        if model_name not in self.__class__._instances:
            self.__class__._instances[model_name] = SentenceTransformer(
                model_name, device=device
            )
        self.model = self.__class__._instances[model_name]
        self.model_name = model_name

    def embed(self, texts: Union[str, list[str]],
               batch_size: int = 64,
               normalize: bool = True,
               show_progress: bool = False) -> np.ndarray:
        """
        Generate embeddings locally.
        normalize=True is required for cosine similarity with dot product.
        """
        if isinstance(texts, str):
            texts = [texts]

        embeddings = self.model.encode(
            texts,
            batch_size=batch_size,
            normalize_embeddings=normalize,
            show_progress_bar=show_progress,
            convert_to_numpy=True
        )
        return embeddings

    def __call__(self, texts: Union[str, list[str]]) -> list[list[float]]:
        """Make the embedder callable (compatible with ChromaRAG)."""
        return self.embed(texts).tolist()

    def similarity(self, text_a: str, text_b: str) -> float:
        """Compute cosine similarity between two texts."""
        a = self.embed(text_a)
        b = self.embed(text_b)
        return float(np.dot(a[0], b[0]))

Part 4: Embedding Cache (Critical for Production)

Without caching, you re-embed the same text thousands of times.

This is the #1 performance optimisation in production RAG systems.

import hashlib, json, sqlite3, pickle
from pathlib import Path

class EmbeddingCache:
    """
    SQLite-backed embedding cache.
    Reduces API costs by 80-95% for stable document sets.
    """

    def __init__(self, cache_path: str = "./embed_cache.db",
                  model_name: str = "text-embedding-3-small"):
        self.model_name = model_name
        self.db_path = cache_path
        self._init_db()

    def _init_db(self) -> None:
        with sqlite3.connect(self.db_path) as conn:
            conn.execute("""
                CREATE TABLE IF NOT EXISTS embeddings (
                    hash TEXT PRIMARY KEY,
                    model TEXT NOT NULL,
                    text TEXT NOT NULL,
                    embedding BLOB NOT NULL,
                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
                )
            """)
            conn.execute("CREATE INDEX IF NOT EXISTS idx_hash ON embeddings(hash)")

    def _hash(self, text: str) -> str:
        return hashlib.sha256(f"{self.model_name}:{text}".encode()).hexdigest()

    def get(self, text: str) -> list[float] | None:
        h = self._hash(text)
        with sqlite3.connect(self.db_path) as conn:
            row = conn.execute(
                "SELECT embedding FROM embeddings WHERE hash = ?", (h,)
            ).fetchone()
        return pickle.loads(row[0]) if row else None

    def set(self, text: str, embedding: list[float]) -> None:
        h = self._hash(text)
        with sqlite3.connect(self.db_path) as conn:
            conn.execute("""
                INSERT OR REPLACE INTO embeddings (hash, model, text, embedding)
                VALUES (?, ?, ?, ?)
            """, (h, self.model_name, text[:500], pickle.dumps(embedding)))

    def embed_with_cache(self, texts: list[str],
                          embed_fn: callable) -> list[list[float]]:
        """Embed texts, using cache where available."""
        results = [None] * len(texts)
        uncached_indices = []
        uncached_texts = []

        # Check cache
        for i, text in enumerate(texts):
            cached = self.get(text)
            if cached is not None:
                results[i] = cached
            else:
                uncached_indices.append(i)
                uncached_texts.append(text)

        # Embed uncached
        if uncached_texts:
            new_embeddings = embed_fn(uncached_texts)
            for idx, embedding in zip(uncached_indices, new_embeddings):
                results[idx] = embedding
                self.set(texts[idx], embedding)

        cache_hit_rate = (len(texts) - len(uncached_texts)) / len(texts)
        logger.info(f"Cache hit rate: {cache_hit_rate:.1%} "
                    f"({len(texts) - len(uncached_texts)}/{len(texts)})")

        return results

    def stats(self) -> dict:
        with sqlite3.connect(self.db_path) as conn:
            count = conn.execute("SELECT COUNT(*) FROM embeddings").fetchone()[0]
            size = Path(self.db_path).stat().st_size / 1024 / 1024
        return {"cached_embeddings": count, "cache_size_mb": round(size, 2)}

Part 5: Semantic Similarity Utilities

import numpy as np

def cosine_similarity(a: list[float], b: list[float]) -> float:
    """Compute cosine similarity between two vectors."""
    a, b = np.array(a), np.array(b)
    return float(np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b) + 1e-10))


def top_k_similar(query_embedding: list[float],
                   corpus_embeddings: list[list[float]],
                   top_k: int = 5) -> list[dict]:
    """Find top-k most similar embeddings from a corpus."""
    q = np.array(query_embedding)
    corpus = np.array(corpus_embeddings)

    # Normalize
    q_norm = q / (np.linalg.norm(q) + 1e-10)
    corpus_norm = corpus / (np.linalg.norm(corpus, axis=1, keepdims=True) + 1e-10)

    scores = corpus_norm @ q_norm
    top_indices = np.argsort(scores)[::-1][:top_k]

    return [{"index": int(i), "score": float(scores[i])} for i in top_indices]


def cluster_texts(embeddings: list[list[float]], n_clusters: int = 5) -> list[int]:
    """
    Cluster text embeddings using K-means.
    Useful for: topic discovery, segment analysis, deduplication.
    """
    from sklearn.cluster import KMeans
    kmeans = KMeans(n_clusters=n_clusters, random_state=42, n_init=10)
    labels = kmeans.fit_predict(np.array(embeddings))
    return labels.tolist()

Part 6: Dimensionality Reduction for Visualisation

def visualise_embeddings(texts: list[str],
                          embeddings: list[list[float]],
                          labels: list[str] = None,
                          output_path: str = "embeddings_2d.png") -> None:
    """
    Visualise embeddings in 2D using UMAP or t-SNE.
    Used by ML teams to audit embedding quality.
    """
    import matplotlib.pyplot as plt

    try:
        from umap import UMAP
        reducer = UMAP(n_components=2, random_state=42, n_neighbors=15)
    except ImportError:
        from sklearn.manifold import TSNE
        reducer = TSNE(n_components=2, random_state=42, perplexity=30)

    coords = reducer.fit_transform(np.array(embeddings))

    plt.figure(figsize=(12, 8))
    if labels:
        unique_labels = list(set(labels))
        colors = plt.cm.tab10(range(len(unique_labels)))
        for label, color in zip(unique_labels, colors):
            mask = [l == label for l in labels]
            plt.scatter(
                coords[mask, 0], coords[mask, 1],
                label=label, c=[color], alpha=0.7
            )
        plt.legend()
    else:
        plt.scatter(coords[:, 0], coords[:, 1], alpha=0.7)

    plt.title("Embedding Space Visualisation")
    plt.savefig(output_path, dpi=150, bbox_inches="tight")
    plt.close()

Output Standards

• Always normalise embeddings before cosine similarity (embed with normalize=True)
• Cache ALL embeddings — re-embedding is wasteful and expensive
• Use text-embedding-3-small for most tasks; upgrade to large only if retrieval quality is insufficient
• Batch size: 100 for OpenAI, 64 for local models
• Store embedding dimension in metadata — dimension mismatch causes silent failures
• Periodically validate cache integrity: recompute 1% of cached embeddings and compare

AXE MCP Server Integration

Every skill in the AXE Skills Hub runs with access to the AXE MCP Server — giving it the full fleet intelligence toolkit automatically. No setup required; tools are available in any AXE-powered session.

Core Tools Available

Quick Start

# In any AXE session, tools are pre-loaded. Example chaining:

# 1. Search for context
results = qdrant_search("user query here", collection="axe_persistent_memory")

# 2. Fetch live data if needed
content = web_fetch("https://docs.example.com/api")

# 3. Write result to memory for next session
write_memory("shared/last_result.md", output)

# 4. Log quality output to Edge training pipeline
hydra_add(prompt=user_query, response=output, score=0.9, source="skill-name")

Edge Training Integration

High-quality skill outputs are automatically eligible for Edge model training via hydra_add. When a response scores ≥0.85 in evals, pipe it to the Hydra pipeline to compound Edge's knowledge. This is how skills make Edge smarter over time.

# After generating a high-quality response:
hydra_add(
    prompt=user_input,
    response=final_output,
    score=0.9,          # eval score
    source="skill-name" # tracks provenance
)