When Your Firestore × Gemini Embeddings RAG Quietly Degrades — Designing for Re-Embedding

// embed.js — a wrapper that produces Gemini embeddings "with provenance"
import { GoogleGenAI } from "@google/genai";
import crypto from "node:crypto";
 
const ai = new GoogleGenAI({ apiKey: process.env.GEMINI_API_KEY });
 
// Pin the embedding config in one place. Swaps happen only here.
export const EMBED_CONFIG = {
  model: "gemini-embedding-001",
  dimensions: 1536,        // truncated from the 3072 default; must match the Firestore index
  version: "v2",           // internal version; bump whenever the config changes
};
 
function sha256(text) {
  return crypto.createHash("sha256").update(text).digest("hex");
}
 
// pass RETRIEVAL_DOCUMENT on write, RETRIEVAL_QUERY on search
export async function embed(text, taskType) {
  const res = await ai.models.embedContent({
    model: EMBED_CONFIG.model,
    contents: text,
    config: {
      taskType,
      outputDimensionality: EMBED_CONFIG.dimensions,
    },
  });
  const values = res.embeddings[0].values;
  return {
    values,
    meta: {
      embedModel: EMBED_CONFIG.model,
      embedDims: EMBED_CONFIG.dimensions,
      embedVersion: EMBED_CONFIG.version,
      taskType,
      textHash: sha256(text),
    },
  };
}

// store.js — save the vector and its metadata together
import { getFirestore, FieldValue } from "firebase-admin/firestore";
import { embed, EMBED_CONFIG } from "./embed.js";
 
const db = getFirestore();
 
export async function upsertDoc(docId, text, extraFields = {}) {
  const { values, meta } = await embed(text, "RETRIEVAL_DOCUMENT");
  await db.collection("kb").doc(docId).set(
    {
      text,
      ...extraFields,
      embedding: FieldValue.vector(values),  // stored as the native vector type
      embedVersion: meta.embedVersion,       // keep top-level; we filter on it
      embedModel: meta.embedModel,
      embedDims: meta.embedDims,
      textHash: meta.textHash,
      updatedAt: FieldValue.serverTimestamp(),
    },
    { merge: true },
  );
}

// drift-check.js — count how many stale vectors remain
import { getFirestore } from "firebase-admin/firestore";
import { EMBED_CONFIG } from "./embed.js";
 
const db = getFirestore();
 
export async function countStaleVectors() {
  const snap = await db
    .collection("kb")
    .where("embedVersion", "!=", EMBED_CONFIG.version)
    .count()
    .get();
  const stale = snap.data().count;
  if (stale > 0) {
    console.warn(`[drift] ${stale} vectors are still on an old version`);
  }
  return stale;
}

// golden-eval.js — measure Recall@5 over fixed queries
import { search } from "./search.js";
 
const GOLDEN = [
  { query: "where do I check the refund policy", expectId: "faq-refund" },
  { query: "how to cancel my subscription", expectId: "guide-cancel-subscription" },
  // ... 20-30 pairs
];
 
export async function evalRecallAt5() {
  let hit = 0;
  for (const g of GOLDEN) {
    const results = await search(g.query, { limit: 5 });
    if (results.some((r) => r.id === g.expectId)) hit++;
  }
  const recall = hit / GOLDEN.length;
  console.log(`Recall@5 = ${recall.toFixed(3)}`);
  return recall;
}

// reembed.js — backfill into a new field, then flip the read target on completion
import { getFirestore, FieldValue } from "firebase-admin/firestore";
import { embed, EMBED_CONFIG } from "./embed.js";
 
const db = getFirestore();
const NEW_FIELD = "embedding_v2";
 
// skip the API call when text is unchanged (the key to cost control)
export async function backfill(batchSize = 100) {
  let last = null;
  let processed = 0;
  for (;;) {
    let q = db.collection("kb").orderBy("__name__").limit(batchSize);
    if (last) q = q.startAfter(last);
    const snap = await q.get();
    if (snap.empty) break;
 
    for (const doc of snap.docs) {
      const d = doc.data();
      // already embedded with the latest config? skip it
      if (d.embedVersionV2 === EMBED_CONFIG.version) continue;
      const { values, meta } = await embed(d.text, "RETRIEVAL_DOCUMENT");
      await doc.ref.update({
        [NEW_FIELD]: FieldValue.vector(values),
        embedVersionV2: meta.embedVersion,
        embedDimsV2: meta.embedDims,
      });
      processed++;
    }
    last = snap.docs[snap.docs.length - 1];
  }
  console.log(`backfill done: re-embedded ${processed} documents`);
}

// search.js — vector search with a distance threshold
import { getFirestore } from "firebase-admin/firestore";
import { embed } from "./embed.js";
 
const db = getFirestore();
const VECTOR_FIELD = "embedding_v2";     // after migration, flip just this one line
const MAX_DISTANCE = 0.55;               // COSINE distance; tune against the golden set
 
export async function search(query, { limit = 8 } = {}) {
  const { values } = await embed(query, "RETRIEVAL_QUERY");
  const snap = await db
    .collection("kb")
    .where("embedVersionV2", "==", "v2")  // physically exclude un-migrated rows
    .findNearest({
      vectorField: VECTOR_FIELD,
      queryVector: values,
      limit,
      distanceMeasure: "COSINE",
      distanceResultField: "_distance",
    })
    .get();
 
  return snap.docs
    .map((d) => ({ id: d.id, ...d.data(), distance: d.get("_distance") }))
    .filter((r) => r.distance <= MAX_DISTANCE);  // discard candidates that are too far
}