Vector Embedding Integration with OpenAI, Voyage, and Gemini in ai-memory

ai-memory unifies vector embedding integration with OpenAI, Voyage, and Gemini behind a single async Embedder trait, enabling provider-agnostic semantic search with SQLite-backed hybrid retrieval.

The akitaonrails/ai-memory project implements a provider-agnostic vector embedding integration with OpenAI, Voyage, and Gemini through a shared Rust trait. This design allows the MCP server to swap embedding backends without changing storage or search logic. The embedding layer lives in crates/ai-memory-llm and feeds directly into the hybrid search system.

Provider-Agnostic Embedder Trait

All embedder implementations implement the same async interface defined in crates/ai-memory-llm/src/embedding.rs.

Core Interface

The Embedder trait requires every provider to expose its metadata and embedding methods:

#[async_trait]
pub trait Embedder: Send + Sync {
    fn provider(&self) -> &'static str;            // e.g. "voyage" or "google"
    fn model(&self) -> &str;
    fn dim(&self) -> u32;
    async fn embed(&self, text: &str) -> LlmResult<Vec<f32>>;
    async fn embed_document(&self, text: &str) -> LlmResult<Vec<f32>> { … }
    async fn embed_query(&self, text: &str) -> LlmResult<Vec<f32>> { … }
}

This contract ensures that OpenAI, Voyage, and Gemini embedders can be treated interchangeably as Arc<dyn Embedder> throughout the server.

Voyage Embedder: OpenAI-Compatible Requests

The Voyage implementation sends HTTP POST requests to https://api.voyageai.com/v1/embeddings. It reuses OpenAI-style endpoint normalization via normalize_openai_base, making its request shape identical to OpenAI's embedding endpoint.

Key implementation details from crates/ai-memory-llm/src/embedding.rs:

• Constructed with an API key, model name, and expected vector dimensionality.
• Builds the URL through normalize_openai_base(&self.base_url, "embeddings").
• Transmits a JSON body shaped as { input: [text], model: model }.
• Validates response dimensionality and unit-normalises the returned vector.

Because Voyage mirrors the OpenAI request format, the same JSON serialization and response parsing logic can service both providers.

Gemini (Google) Embedder

The Gemini embedder targets https://generativelanguage.googleapis.com/v1/models/{model}:embedContent in crates/ai-memory-llm/src/google.rs.

Task-Type Handling and Retry Logic

The GoogleEmbedder struct detects whether the model uses the newer embedding_v2 format to select the correct task-type header. Its embed_with_task method formats text differently for documents versus queries when embedding_v2 is enabled. The implementation also performs exponential back-off on HTTP 429 responses and validates vector length before normalising.

Source references:

• crates/ai-memory-llm/src/google.rs lines 22–52 for the struct and constructor.
• crates/ai-memory-llm/src/google.rs lines 62–84 for embed_with_task.

Wiring Embeddings into the MCP Server

The MCP server holds an optional embedder inside ServerBuilder:

pub struct ServerBuilder {
    // …
    embedder: Option<Arc<dyn Embedder>>,
}

When a search query arrives, the server checks for an embedder at crates/ai-memory-mcp/src/server.rs lines 1520–1528. If present, it calls embed_query; otherwise it falls back to pure FTS5, entity, and graph search.

Hybrid Search Pipeline

During retrieval, the server extracts provider metadata to blend vector similarity with textual results:

let (provider, model, dim) = match (&self.embedder, options.query_vec) {
    (Some(e), Some(v)) => (e.provider(), e.model(), e.dim()),
    // …
};

This vector stream merges with FTS5 and graph streams via reciprocal rank fusion (RRF).

Persisting Vectors

The writer actor in crates/ai-memory-store/src/writer.rs receives StoreEmbedding commands at lines 251–257 and persists vectors alongside page IDs in the SQLite embeddings table.

Configuring the Embedder Provider

Both providers are instantiated from EmbedderConfig via crates/ai-memory-llm/src/factory.rs. A typical ai_memory.toml entry looks like:

[embedder]
provider = "voyage"          # or "google"

api_key = "YOUR_KEY"
model = "voyage-2"           # or "gemini-embedding-001"

dim = 1536

The factory selects the concrete VoyageEmbedder or GoogleEmbedder and returns an Arc<dyn Embedder> that is injected into the server builder.

Practical Code Examples

Creating a Voyage Embedder

use ai_memory_llm::{VoyageEmbedder, Embedder};
use secrecy::SecretString;

let voyage = VoyageEmbedder::new(
    SecretString::from("YOUR_VOYAGE_API_KEY"),
    "voyage-2",
    1536,
)?;

let embedder = std::sync::Arc::new(voyage);

Reference: crates/ai-memory-llm/src/embedding.rs lines 60–78.

Creating a Gemini Embedder

use ai_memory_llm::{GoogleEmbedder, Embedder};
use secrecy::SecretString;

let gemini = GoogleEmbedder::new(
    SecretString::from("YOUR_GEMINI_API_KEY"),
    "gemini-embedding-001",
    768,
)?;

let gemini = gemini.with_base_url("http://localhost:8080");
let embedder = std::sync::Arc::new(gemini);

Reference: crates/ai-memory-llm/src/google.rs lines 34–46.

Attaching the Embedder to the Server

use ai_memory_mcp::ServerBuilder;

let server = ServerBuilder::default()
    .with_embedder(embedder)
    .build()?;

Reference: crates/ai-memory-mcp/src/server.rs lines 1326–1328.

Performing Semantic Search

use reqwest::Client;

let client = Client::new();
let resp = client
    .get("http://127.0.0.1:49374/api/v1/search")
    .query(&[("q", "rust async traits")])
    .send()
    .await?
    .json::<serde_json::Value>()
    .await?;

If the server has a Voyage or Gemini embedder configured, the query text is embedded via embed_query and combined with FTS5 results for ranked output.

Summary

• Provider-agnostic design: The Embedder trait in embedding.rs abstracts OpenAI, Voyage, and Gemini behind a single async interface.
• OpenAI compatibility: Voyage uses normalize_openai_base to match OpenAI's request shape, while Gemini uses a distinct embedContent API with task-type headers.
• Server integration: The MCP server accepts any Arc<dyn Embedder> through ServerBuilder::with_embedder, enabling optional hybrid vector search.
• Storage: Vectors are persisted via StoreEmbedding commands in the SQLite-backed store writer.
• Configuration: The factory in factory.rs maps TOML config to concrete embedder instances.

Frequently Asked Questions

How does ai-memory keep embedding providers interchangeable?

All providers implement the Embedder trait defined in crates/ai-memory-llm/src/embedding.rs. The server only interacts with Arc<dyn Embedder>, so swapping Voyage for Gemini requires no changes to search or storage logic.

Why does the Voyage embedder use OpenAI-style request normalization?

Voyage's endpoint at /v1/embeddings accepts the same JSON shape as OpenAI. The normalize_openai_base utility in embedding.rs lets the Voyage client reuse OpenAI-compatible URL building and request formatting logic.

What happens if no embedder is configured in the server?

If ServerBuilder has no embedder, the query handler at server.rs lines 1520–1528 skips vector generation and falls back to pure FTS5, entity extraction, and graph traversal without RRF vector scoring.

How does the Gemini embedder handle rate limits?

The GoogleEmbedder in google.rs implements exponential back-off when it encounters HTTP 429 responses, ensuring temporary quota errors do not terminate the embedding pipeline.

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