add embeddings and vector search modules

Embeddings module wraps the embeddinggemma pipeline for query, batch,
and document chunk embedding using Transformers.js feature-extraction.
Vector search implements cosine similarity ranking over embedded chunks.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

src/pipeline/embeddings.ts

@@ -0,0 +1,41 @@
+import { getEmbeddingPipeline } from "./models";
+import { EMBED_QUERY_TEMPLATE, EMBED_DOC_TEMPLATE } from "../constants";
+
+/**
+ * Embed a single query string using the embeddinggemma model.
+ * Uses the query template: "task: search result | query: {query}"
+ */
+export async function embedQuery(query: string): Promise<Float32Array> {
+  const pipe = getEmbeddingPipeline();
+  if (!pipe) throw new Error("Embedding model not loaded");
+  const text = EMBED_QUERY_TEMPLATE(query);
+  const output = await pipe(text, { pooling: "mean", normalize: true });
+  return new Float32Array(output.tolist()[0]);
+}
+
+/**
+ * Embed multiple texts in batch.
+ */
+export async function embedBatch(
+  texts: string[],
+): Promise<Float32Array[]> {
+  const pipe = getEmbeddingPipeline();
+  if (!pipe) throw new Error("Embedding model not loaded");
+  const output = await pipe(texts, { pooling: "mean", normalize: true });
+  return output.tolist().map((arr: number[]) => new Float32Array(arr));
+}
+
+/**
+ * Embed a document chunk using the doc template.
+ * Uses the doc template: "title: {title} | text: {body}"
+ */
+export async function embedDocChunk(
+  title: string,
+  chunkText: string,
+): Promise<Float32Array> {
+  const pipe = getEmbeddingPipeline();
+  if (!pipe) throw new Error("Embedding model not loaded");
+  const text = EMBED_DOC_TEMPLATE(title, chunkText);
+  const output = await pipe(text, { pooling: "mean", normalize: true });
+  return new Float32Array(output.tolist()[0]);
+}

src/pipeline/vectorSearch.test.ts

@@ -0,0 +1,128 @@
+import { describe, it, expect } from "vitest";
+import { cosineSimilarity, vectorSearch } from "./vectorSearch";
+import type { EmbeddedChunk } from "../types";
+
+// ---------------------------------------------------------------------------
+// Helper to create an EmbeddedChunk with a given embedding
+// ---------------------------------------------------------------------------
+function makeEmbeddedChunk(
+  embedding: number[],
+  docId = "doc1",
+  chunkIndex = 0,
+): EmbeddedChunk {
+  return {
+    docId,
+    chunkIndex,
+    text: `chunk ${chunkIndex} of ${docId}`,
+    startChar: 0,
+    title: "Test",
+    embedding: new Float32Array(embedding),
+  };
+}
+
+// ---------------------------------------------------------------------------
+// cosineSimilarity
+// ---------------------------------------------------------------------------
+describe("cosineSimilarity", () => {
+  it("returns 1 for identical vectors", () => {
+    const v = new Float32Array([1, 2, 3]);
+    expect(cosineSimilarity(v, v)).toBeCloseTo(1.0, 5);
+  });
+
+  it("returns -1 for opposite vectors", () => {
+    const a = new Float32Array([1, 0, 0]);
+    const b = new Float32Array([-1, 0, 0]);
+    expect(cosineSimilarity(a, b)).toBeCloseTo(-1.0, 5);
+  });
+
+  it("returns 0 for orthogonal vectors", () => {
+    const a = new Float32Array([1, 0, 0]);
+    const b = new Float32Array([0, 1, 0]);
+    expect(cosineSimilarity(a, b)).toBeCloseTo(0.0, 5);
+  });
+
+  it("is symmetric", () => {
+    const a = new Float32Array([1, 2, 3]);
+    const b = new Float32Array([4, 5, 6]);
+    expect(cosineSimilarity(a, b)).toBeCloseTo(cosineSimilarity(b, a), 10);
+  });
+
+  it("is scale-invariant", () => {
+    const a = new Float32Array([1, 2, 3]);
+    const b = new Float32Array([2, 4, 6]); // 2x scale of a
+    expect(cosineSimilarity(a, b)).toBeCloseTo(1.0, 5);
+  });
+
+  it("computes correct value for known vectors", () => {
+    // a = [3, 4], b = [4, 3]
+    // dot = 12 + 12 = 24, |a| = 5, |b| = 5 → cos = 24/25 = 0.96
+    const a = new Float32Array([3, 4]);
+    const b = new Float32Array([4, 3]);
+    expect(cosineSimilarity(a, b)).toBeCloseTo(0.96, 2);
+  });
+});
+
+// ---------------------------------------------------------------------------
+// vectorSearch
+// ---------------------------------------------------------------------------
+describe("vectorSearch", () => {
+  // Create chunks with embeddings at known angles from the query
+  const query = new Float32Array([1, 0, 0]);
+
+  const chunks: EmbeddedChunk[] = [
+    makeEmbeddedChunk([0, 1, 0], "orthogonal", 0), // cos = 0
+    makeEmbeddedChunk([1, 0, 0], "identical", 0), // cos = 1
+    makeEmbeddedChunk([0.7, 0.7, 0], "similar", 0), // cos ≈ 0.707
+    makeEmbeddedChunk([-1, 0, 0], "opposite", 0), // cos = -1
+    makeEmbeddedChunk([0.9, 0.1, 0], "very-similar", 0), // cos ≈ 0.994
+  ];
+
+  it("returns results sorted by descending score", () => {
+    const results = vectorSearch(query, chunks);
+    for (let i = 1; i < results.length; i++) {
+      expect(results[i].score).toBeLessThanOrEqual(results[i - 1].score);
+    }
+  });
+
+  it("ranks identical vector highest", () => {
+    const results = vectorSearch(query, chunks);
+    expect(results[0].chunk.docId).toBe("identical");
+  });
+
+  it("ranks opposite vector lowest", () => {
+    const results = vectorSearch(query, chunks);
+    expect(results[results.length - 1].chunk.docId).toBe("opposite");
+  });
+
+  it("all results have source 'vector'", () => {
+    const results = vectorSearch(query, chunks);
+    for (const r of results) {
+      expect(r.source).toBe("vector");
+    }
+  });
+
+  it("respects topK parameter", () => {
+    const results = vectorSearch(query, chunks, 2);
+    expect(results.length).toBe(2);
+    expect(results[0].chunk.docId).toBe("identical");
+    expect(results[1].chunk.docId).toBe("very-similar");
+  });
+
+  it("returns all when topK exceeds chunk count", () => {
+    const results = vectorSearch(query, chunks, 100);
+    expect(results.length).toBe(chunks.length);
+  });
+
+  it("handles empty chunks array", () => {
+    const results = vectorSearch(query, []);
+    expect(results).toEqual([]);
+  });
+
+  it("handles single chunk", () => {
+    const single = [makeEmbeddedChunk([0.5, 0.5, 0], "only", 0)];
+    const results = vectorSearch(query, single);
+    expect(results.length).toBe(1);
+    expect(results[0].source).toBe("vector");
+    expect(results[0].score).toBeGreaterThan(0);
+  });
+});

src/pipeline/vectorSearch.ts

@@ -0,0 +1,34 @@
+import type { EmbeddedChunk, ScoredChunk } from "../types";
+
+/**
+ * Cosine similarity between two vectors.
+ */
+export function cosineSimilarity(a: Float32Array, b: Float32Array): number {
+  let dot = 0;
+  let normA = 0;
+  let normB = 0;
+  for (let i = 0; i < a.length; i++) {
+    dot += a[i] * b[i];
+    normA += a[i] * a[i];
+    normB += b[i] * b[i];
+  }
+  return dot / (Math.sqrt(normA) * Math.sqrt(normB));
+}
+
+/**
+ * Search embedded chunks by cosine similarity to query embedding.
+ * Returns the top-K results sorted by descending similarity score.
+ */
+export function vectorSearch(
+  queryEmbedding: Float32Array,
+  chunks: EmbeddedChunk[],
+  topK: number = 20,
+): ScoredChunk[] {
+  const scored = chunks.map((chunk) => ({
+    chunk,
+    score: cosineSimilarity(queryEmbedding, chunk.embedding),
+    source: "vector" as const,
+  }));
+  scored.sort((a, b) => b.score - a.score);
+  return scored.slice(0, topK);
+}