How Vector Embeddings Are Integrated Into the ai-memory Retrieval Pipeline

The ai-memory retrieval pipeline generates dense vector embeddings for user queries, computes cosine similarity against stored page embeddings, and fuses the vector stream with FTS5, entity-match, and graph signals via Reciprocal Rank Fusion to produce semantically ranked results.

The akitaonrails/ai-memory project implements a hybrid search architecture that combines traditional full-text search with neural vector embeddings. This article explains how the pipeline embeds queries, persists and retrieves vectors, scores semantic similarity, and merges multiple ranking signals into a unified result set.

Query Embedding Generation

Every search begins by transforming the user's text into a dense vector representation.

The memory_query command (and the corresponding web API at /api/v1/search) delegates embedding generation to the EmbeddingClient trait defined in crates/ai-memory-llm/src/embedding.rs. The system supports multiple providers:

• OpenAI (text-embedding-3-small, text-embedding-3-large)
• Ollama (local embedding models)
• Local MiniLM (on-device inference)

The embedder returns a float vector together with metadata: provider, model, and dim. This triple ensures that vectors remain comparable throughout the pipeline.

// From ai-memory-llm/src/embedding.rs
let embedder = EmbeddingClient::new_openai("text-embedding-3-small");
let query_vec = embedder.embed("reset project").await?;
// Returns: Embedding { vector: Vec<f32>, provider: "openai", model: "text-embedding-3-small", dim: 1536 }

Embedding Storage and Metadata Tracking

Page embeddings are persisted in SQLite with strict versioning to prevent model mismatch.

The write path is handled by store_embedding in crates/ai-memory-store/src/writer.rs (line 1162). Each row in the page_embeddings table stores:

• The raw vector bytes (as a BLOB)
• The provider/model/dim triple (as TEXT/INTEGER columns)
• A timestamp for staleness detection

// From ai-memory-store/src/writer.rs
pub async fn store_embedding(
    &self,
    page_id: i64,
    embedding: &Embedding,
) -> Result<(), StoreError> {
    let bytes = f32_slice_to_bytes(&embedding.vector);
    sqlx::query(
        "INSERT INTO page_embeddings (page_id, vector, provider, model, dim, created_at)
         VALUES (?1, ?2, ?3, ?4, ?5, ?6)
         ON CONFLICT(page_id) DO UPDATE SET
           vector=excluded.vector,
           provider=excluded.provider,
           model=excluded.model,
           dim=excluded.dim,
           created_at=excluded.created_at"
    )
    .bind(page_id)
    .bind(bytes)
    .bind(&embedding.provider)
    .bind(&embedding.model)
    .bind(embedding.dim as i64)
    .bind(Utc::now())
    .execute(&self.pool)
    .await?;
    Ok(())
}

The batch variant store_embeddings optimizes back-fill operations for large projects.

Loading Compatible Embeddings for Retrieval

Before scoring, the system filters embeddings to ensure vector compatibility.

The function load_embeddings in crates/ai-memory-store/src/reader.rs (line 3063) executes:

SELECT page_id, vector, provider, model, dim
FROM page_embeddings
WHERE provider = ? AND model = ? AND dim = ?
ORDER BY created_at DESC

This guarantees that only embeddings from the same model configuration as the query embedder are loaded. Attempting to compare vectors from different dimensions or models would produce meaningless similarity scores.

Cosine Similarity Scoring

The core ranking computation happens in dot_embedding_bytes, used inside top_embedding_hits_for_project (line 3202 in reader.rs).

The function:

1. Deserializes the stored BLOB into a Vec<f32>
2. Computes the dot product between query and document vectors
3. Normalizes to cosine similarity: cosine = dot / (‖q‖ × ‖d‖)

// Conceptual implementation from ai-memory-store/src/reader.rs
fn dot_embedding_bytes(query: &[f32], stored: &[u8]) -> f32 {
    let doc: Vec<f32> = bytes_to_f32_slice(stored);
    let dot: f32 = query.iter().zip(&doc).map(|(a, b)| a * b).sum();
    let query_norm = query.iter().map(|x| x * x).sum::<f32>().sqrt();
    let doc_norm = doc.iter().map(|x| x * x).sum::<f32>().sqrt();
    dot / (query_norm * doc_norm + 1e-8) // epsilon for numerical stability
}

The output becomes the vector_rank for each hit, ranging from -1.0 to 1.0 (though typically 0.3–0.9 for semantically related content).

RRF Fusion: Combining Vector, FTS5, Entity, and Graph Signals

The vector stream does not standalone—it merges with three other retrieval signals:

Stream	Source	Rank Field
FTS5	SQLite full-text index	fts_rank
Entity	Named entity mention matches	entity_rank
Graph	Page graph neighborhood	graph_rank
Vector	Cosine similarity (this article)	vector_rank

The Reciprocal Rank Fusion (RRF) step in crates/ai-memory-store/src/reader.rs (line 682) computes:

// From RrfContributions
let fused_score = rrf_weight * (
    1.0 / (k + fts_rank) +
    1.0 / (k + entity_rank) +
    1.0 / (k + graph_rank) +
    1.0 / (k + vector_rank)
);

Where k = 60 (standard RRF constant) and rrf_weight is configurable. The fused score is stored in SearchExplain::fused, and the raw cosine similarity is preserved in SearchExplain::cosine for API consumers.

Result Presentation and API Exposure

Final PageHit objects include the hybrid ranking. The web API endpoint in crates/ai-memory-web/src/routes/api.rs (line 287) returns:

{
  "hits": [
    {
      "page_id": 42,
      "title": "Resetting Project State",
      "rank": 0.847,
      "explain": {
        "fts_rank": 3,
        "vector_rank": 2,
        "graph_rank": 5,
        "entity_rank": null,
        "fused": 0.847,
        "cosine": 0.912
      }
    }
  ]
}

The optional cosine field allows callers to debug semantic relevance independently of the fusion layer.

Embedding Back-fill and Model Migration

When the configured embedder changes (new model, new provider, or dimension shift), the system automatically re-embeds affected pages.

The run_backfill function in crates/ai-memory-consolidate/src/embed.rs identifies pages where:

• No embedding exists
• The stored triple differs from current configuration
• The embedding exceeds a staleness threshold

use ai_memory_consolidate::embed::run_backfill;

// Re-embed all incompatible or missing pages
run_backfill(&writer, &embedder, &scope).await?;

This ensures the vector retrieval pipeline never encounters dimension mismatches or stale semantic representations.

Summary

• Embedding generation in ai-memory-llm/src/embedding.rs produces query vectors with provider/model/dim metadata
• Storage via store_embedding in writer.rs persists vectors with versioning for compatibility checks
• Retrieval via load_embeddings in reader.rs filters by the embedder triple to ensure comparable vectors
• Scoring via dot_embedding_bytes computes cosine similarity for semantic ranking
• Fusion via RrfContributions merges vector, FTS5, entity, and graph streams into a unified score
• Back-fill via run_backfill maintains embedding freshness across model changes

Frequently Asked Questions

How does ai-memory prevent comparing embeddings from different models?

The system stores the (provider, model, dim) triple alongside every vector. The load_embeddings function in reader.rs filters the SQL query to only rows matching the current embedder's configuration, ensuring dimensional and semantic compatibility.

Can I use a local embedding model instead of OpenAI?

Yes. The EmbeddingClient trait in ai-memory-llm/src/embedding.rs supports multiple backends including Ollama and local MiniLM. Configure your preferred provider in the project settings; the pipeline handles the rest transparently.

What happens to existing embeddings when I switch models?

The back-fill system in ai-memory-consolidate/src/embed.rs automatically detects embeddings with mismatched triples and re-computes them using the new model. This occurs during the auto-improve scheduler runs or can be triggered manually.

Why does the API return both fused and cosine scores?

The fused score incorporates all retrieval signals (vector, FTS5, entity, graph) through RRF and represents the final ranking. The cosine score exposes the raw semantic similarity for debugging, allowing you to verify whether vector relevance aligns with perceived result quality.

{
  "mcpServers": {
    "instagit": {
      "command": "npx",
      "args": ["-y", "instagit@latest"]
    }
  }
}