Embedding cache (KV)

Source: platform/apps/app/src/lib/rag/embedding-cache.ts · rendered from main on every deploy — edit in the repo, not here

platform/apps/app/src/lib/rag/embedding-cache.ts

/**
 * Embedding cache — KV in front of bge-m3 (FinOps + latency).
 *
 * Every retrieval path embeds the user's query before hitting Vectorize.
 * The bge-m3 call costs ~$0.0001 and adds ~50ms; for repeated FAQ
 * queries that's pure waste. KV is the right fit:
 *
 *   - bge-m3 is deterministic — same query → same vector. Cache key
 *     can be a simple SHA-256 of (model, canonical query).
 *   - Vectors are tiny (1024 floats = 4KB). KV holds them comfortably
 *     and serves them in ~10-15ms — faster than the AI call's ~50ms.
 *   - Embeddings are tenant-agnostic (the model has no idea which
 *     tenant asked) so the cache is shared globally. A single popular
 *     phrase across many tenants pays the AI cost once.
 *
 * 7-day TTL is generous — embedding model versions change rarely, and
 * we version the key prefix (`v1`) so a model bump invalidates without
 * a manual flush.
 *
 * KV economics check:
 *   - KV read: ~$0.50/1M ($5e-7 per call)
 *   - bge-m3 call: ~$1e-4 per call
 *   - Even at 1% hit rate, savings dominate the read cost by 200x.
 *   - 4KB × 1M cached vectors = 4GB ≈ $2/month storage. Negligible.
 */

const CACHE_KEY_PREFIX = 'embed:bge-m3:v1:';
const CACHE_TTL_SECONDS = 7 * 24 * 60 * 60;
/** Skip the cache for unusually long queries — the SHA-256 keys would
 * still be fixed-size, but caching a 5KB-prompt embedding is unlikely
 * to repeat. Keeps the hot set lean. */
const MAX_QUERY_LEN_FOR_CACHE = 1000;

/** Build a stable KV key. Canonicalizes whitespace so trailing spaces
 * etc. don't fragment the cache. */
export async function buildEmbeddingCacheKey(query: string): Promise<string | null> {
  const canonical = query.trim();
  if (canonical.length === 0 || canonical.length > MAX_QUERY_LEN_FOR_CACHE) return null;
  const buf = await crypto.subtle.digest('SHA-256', new TextEncoder().encode(canonical));
  const hex = Array.from(new Uint8Array(buf))
    .map((b) => b.toString(16).padStart(2, '0'))
    .join('');
  return CACHE_KEY_PREFIX + hex;
}

/** Minimal duck-typed view of KVNamespace.get + put — keeps the helper
 * runtime-agnostic for tests. */
export interface EmbeddingCacheLike {
  get(key: string, type: 'arrayBuffer'): Promise<ArrayBuffer | null>;
  put(key: string, value: ArrayBuffer, options?: { expirationTtl?: number }): Promise<void>;
}

/** Returns the cached vector or null on miss / error. Cache must never
 * break the request path. */
export async function readEmbedding(
  cache: EmbeddingCacheLike,
  key: string,
): Promise<number[] | null> {
  try {
    const buf = await cache.get(key, 'arrayBuffer');
    if (!buf) return null;
    // Workers AI returns number[] from .run(); persist & return the same
    // shape so downstream Vectorize.query() doesn't have to branch.
    return Array.from(new Float32Array(buf));
  } catch {
    return null;
  }
}

/** Schedule a write without blocking the response. */
export function writeEmbedding(
  cache: EmbeddingCacheLike,
  key: string,
  vector: number[],
  waitUntil: (p: Promise<unknown>) => void,
): void {
  try {
    const f32 = new Float32Array(vector);
    waitUntil(cache.put(key, f32.buffer, { expirationTtl: CACHE_TTL_SECONDS }));
  } catch {
    // Cache write failed — non-fatal.
  }
}