feat: Phase 2 — embeddings + hybrid memory.search

Adds a small embedder sidecar (Xenova/bge-small-en-v1.5, ONNX, CPU-only)
that the web app calls inline on memory.write and memory.update, and on
demand from the new memory.search tool.

memory.search performs three candidate fetches in parallel — pgvector
cosine similarity, Postgres full-text via plainto_tsquery + ts_rank_cd,
and tag-set overlap — then fuses them with Reciprocal Rank Fusion
(k=60). Each result carries its per-source rank so the model can see
*why* a memory surfaced.

The migrator boot step gained an idempotent embedding backfill: any row
with embedding IS NULL is batched (32 at a time) through the embedder
after SQL migrations apply. Safe to run on every boot.

New tool memory.update fixes the missing edit path; centralises the
re-embed-on-content-change rule alongside write.

Stack additions:
- apps/embedder/ — Fastify server, persistent /data/models volume so the
  ~30 MB model only downloads once
- apps/web/lib/embedder.ts — typed HTTP client with batched embed +
  health probe
- packages/schemas — MemoryUpdateInput, MemorySearchInput
- docker-compose — embedder service, healthcheck, app + migrator both
  depend_on it healthy; EMBEDDER_URL promoted to a required env var

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
2026-05-15 09:04:44 -07:00
co-authored by Claude Opus 4.7
parent e1fa1197f4
commit 9a8b504f51
12 changed files with 1523 additions and 19 deletions
+62
View File
@@ -0,0 +1,62 @@
# syntax=docker/dockerfile:1.7
# -----------------------------------------------------------------------------
# Embedder sidecar.
#
# Builds from the repo root: docker build -f apps/embedder/Dockerfile .
# -----------------------------------------------------------------------------
FROM node:20-alpine AS base
RUN corepack enable
WORKDIR /app
# ---------- deps ----------
FROM base AS deps
COPY package.json pnpm-workspace.yaml pnpm-lock.yaml .npmrc ./
COPY apps/embedder/package.json ./apps/embedder/
# Other workspace package.json files needed so pnpm install can resolve the
# workspace before --filter narrows things down.
COPY apps/web/package.json ./apps/web/
COPY packages/schemas/package.json ./packages/schemas/
RUN --mount=type=cache,id=pnpm,target=/root/.local/share/pnpm/store \
pnpm install --frozen-lockfile --filter @shared-memory/embedder...
# ---------- builder ----------
FROM base AS builder
COPY --from=deps /app/node_modules ./node_modules
COPY --from=deps /app/apps/embedder/node_modules ./apps/embedder/node_modules
COPY . .
# Compile TS to JS.
RUN cd apps/embedder \
&& pnpm exec tsc -p tsconfig.json --noEmit false --outDir dist
# Prune devDependencies so the runtime image only ships production deps.
RUN cd apps/embedder \
&& pnpm install --prod --frozen-lockfile --filter @shared-memory/embedder...
# ---------- runner ----------
FROM node:20-alpine AS runner
WORKDIR /app
ENV NODE_ENV=production \
PORT=8080 \
HOST=0.0.0.0 \
MODEL_CACHE_DIR=/data/models
RUN apk add --no-cache wget \
&& addgroup --system --gid 1001 nodejs \
&& adduser --system --uid 1001 --ingroup nodejs node-embedder \
&& mkdir -p /data/models \
&& chown -R node-embedder:nodejs /data
COPY --from=builder --chown=node-embedder:nodejs /app/apps/embedder/dist ./dist
COPY --from=builder --chown=node-embedder:nodejs /app/apps/embedder/node_modules ./node_modules
COPY --from=builder --chown=node-embedder:nodejs /app/apps/embedder/package.json ./package.json
USER node-embedder
EXPOSE 8080
VOLUME ["/data/models"]
HEALTHCHECK --interval=15s --timeout=5s --start-period=120s --retries=5 \
CMD wget -q -O - http://127.0.0.1:8080/health | grep -q '"ready":true' || exit 1
CMD ["node", "--enable-source-maps", "dist/index.js"]
+21
View File
@@ -0,0 +1,21 @@
{
"name": "@shared-memory/embedder",
"version": "0.1.0",
"private": true,
"type": "module",
"scripts": {
"dev": "tsx watch src/index.ts",
"build": "tsc --noEmit",
"start": "node --enable-source-maps dist/index.js",
"typecheck": "tsc --noEmit"
},
"dependencies": {
"@xenova/transformers": "^2.17.2",
"fastify": "^5.2.0"
},
"devDependencies": {
"@types/node": "^22.10.2",
"tsx": "^4.19.2",
"typescript": "^5.7.2"
}
}
+118
View File
@@ -0,0 +1,118 @@
/**
* Embedder sidecar — loads a small ONNX model once at boot and serves
* mean-pooled, L2-normalized sentence embeddings over HTTP.
*
* Endpoints:
* GET /health → { status, ready, model, dim }
* POST /embed → { vectors: number[][] } given { texts: string[] }
*
* Used by the web app's memory.write / memory.update / memory.search and
* by the migrator's one-shot backfill step.
*/
import Fastify from "fastify";
import { pipeline, env as txEnv } from "@xenova/transformers";
// Persist the downloaded model on a named docker volume so subsequent
// boots don't re-fetch ~30 MB.
txEnv.cacheDir = process.env.MODEL_CACHE_DIR ?? "/data/models";
txEnv.allowLocalModels = true;
txEnv.allowRemoteModels = true;
const MODEL_NAME = process.env.EMBEDDING_MODEL ?? "Xenova/bge-small-en-v1.5";
const EXPECTED_DIM = Number.parseInt(process.env.EMBEDDING_DIM ?? "384", 10);
const PORT = Number.parseInt(process.env.PORT ?? "8080", 10);
const HOST = process.env.HOST ?? "0.0.0.0";
// The pipeline()'s return type is a giant union covering every task; we
// only use feature-extraction, so a narrower call signature is much easier
// to work with than the upstream typing.
interface FeatureExtractor {
(
texts: string[],
options: { pooling: "mean" | "cls"; normalize: boolean },
): Promise<{ tolist: () => number[] | number[][] }>;
}
let extractor: FeatureExtractor | null = null;
async function loadModel() {
const start = Date.now();
console.log(`[embedder] loading ${MODEL_NAME}`);
// Quantized=true is the @xenova default and is fast enough; flip via env if
// we ever need the full-precision model.
extractor = (await pipeline("feature-extraction", MODEL_NAME, {
quantized: process.env.EMBEDDER_QUANTIZED !== "false",
})) as unknown as FeatureExtractor;
console.log(`[embedder] model ready in ${Date.now() - start}ms`);
}
const app = Fastify({
logger: { level: process.env.LOG_LEVEL ?? "info" },
bodyLimit: 5 * 1024 * 1024, // 5 MB — generous for batched embeds
});
app.get("/health", async () => ({
status: "ok",
ready: extractor !== null,
model: MODEL_NAME,
dim: EXPECTED_DIM,
}));
interface EmbedRequest {
texts: string[];
}
app.post("/embed", async (req, reply) => {
if (!extractor) {
return reply.code(503).send({ error: "model not loaded yet" });
}
const body = req.body as EmbedRequest | null;
if (!body || !Array.isArray(body.texts)) {
return reply.code(400).send({ error: "body must be { texts: string[] }" });
}
if (body.texts.length === 0) {
return { vectors: [] };
}
if (body.texts.length > 256) {
return reply.code(400).send({ error: "max 256 texts per request" });
}
if (body.texts.some((t) => typeof t !== "string")) {
return reply.code(400).send({ error: "every entry in texts must be a string" });
}
// Mean-pool the per-token hidden states and L2-normalize so cosine sim
// matches the inner-product distance we'll feed into pgvector.
const output = await extractor(body.texts, {
pooling: "mean",
normalize: true,
});
// Transformers.js returns a Tensor; .tolist() gives nested JS arrays.
// For batches the shape is [batch, dim]; for a single input the wrapper
// may collapse to [dim] — defensively re-wrap.
const raw = output.tolist();
const vectors: number[][] = Array.isArray(raw[0])
? (raw as number[][])
: [raw as number[]];
// Sanity-check the dimension once at runtime — catches a model swap that
// wasn't accompanied by an EMBEDDING_DIM bump.
if (vectors[0] && vectors[0].length !== EXPECTED_DIM) {
return reply.code(500).send({
error: `model produced dim=${vectors[0].length}, expected ${EXPECTED_DIM}`,
});
}
return { vectors };
});
async function start() {
await loadModel();
await app.listen({ host: HOST, port: PORT });
console.log(`[embedder] listening on http://${HOST}:${PORT}`);
}
start().catch((err) => {
console.error("[embedder] startup failed:", err);
process.exit(1);
});
+13
View File
@@ -0,0 +1,13 @@
{
"extends": "../../tsconfig.base.json",
"compilerOptions": {
"rootDir": "./src",
"outDir": "./dist",
"noEmit": false,
"declaration": false,
"module": "ESNext",
"moduleResolution": "Bundler",
"lib": ["ES2022"]
},
"include": ["src/**/*.ts"]
}
+65
View File
@@ -0,0 +1,65 @@
import { env } from "@/lib/env";
/**
* Thin HTTP client for the embedder sidecar. Used by memory.write /
* memory.update / memory.search and by the migrator's backfill step.
*
* Calls are blocking on purpose — write-path latency is a worthwhile
* trade for "the memory I just wrote is searchable now."
*/
export class EmbedderError extends Error {
constructor(message: string, public readonly status?: number) {
super(message);
this.name = "EmbedderError";
}
}
function url(): string {
const u = env().EMBEDDER_URL;
if (!u) throw new EmbedderError("EMBEDDER_URL is not configured");
return u.replace(/\/$/, "");
}
/** Embed a batch of texts. Returns one vector per input. */
export async function embedTexts(texts: string[]): Promise<number[][]> {
if (texts.length === 0) return [];
const res = await fetch(`${url()}/embed`, {
method: "POST",
headers: { "content-type": "application/json" },
body: JSON.stringify({ texts }),
});
if (!res.ok) {
const detail = await res.text().catch(() => "");
throw new EmbedderError(
`embedder returned ${res.status}: ${detail.slice(0, 200)}`,
res.status,
);
}
const body = (await res.json()) as { vectors: number[][] };
if (!Array.isArray(body.vectors) || body.vectors.length !== texts.length) {
throw new EmbedderError("embedder response shape mismatch");
}
return body.vectors;
}
/** Embed a single text — convenience for one-off calls. */
export async function embedText(text: string): Promise<number[]> {
const [vec] = await embedTexts([text]);
if (!vec) throw new EmbedderError("embedder returned no vector");
return vec;
}
/** Quick check used by the migrator before backfilling. */
export async function embedderReady(): Promise<boolean> {
try {
const res = await fetch(`${url()}/health`);
if (!res.ok) return false;
const body = (await res.json()) as { ready?: boolean };
return body.ready === true;
} catch {
return false;
}
}
+3 -4
View File
@@ -21,9 +21,8 @@ const envSchema = z.object({
// Database
DATABASE_URL: z.string().url(),
// Embedder (used in Phase 2; present-but-empty allowed in Phase 1)
EMBEDDER_URL: z
.preprocess((v) => (v === "" ? undefined : v), z.string().url().optional()),
// Embedder sidecar — required in Phase 2 since memory.write embeds inline.
EMBEDDER_URL: z.string().url(),
EMBEDDING_MODEL: z.string().default("Xenova/bge-small-en-v1.5"),
EMBEDDING_DIM: z.coerce.number().int().positive().default(384),
@@ -72,7 +71,7 @@ function buildPhaseStub(): Env {
OIDC_CLIENT_ID_MCP: "build",
OIDC_AUDIENCE: "build",
DATABASE_URL: "postgres://build:build@build-phase.invalid:5432/build",
EMBEDDER_URL: undefined,
EMBEDDER_URL: "http://embedder.invalid:8080",
EMBEDDING_MODEL: "Xenova/bge-small-en-v1.5",
EMBEDDING_DIM: 384,
NEXTAUTH_SECRET: "build-phase-secret-not-used-at-runtime-xxxxxxxx",
+234 -5
View File
@@ -1,22 +1,23 @@
import { and, desc, eq, isNull, sql } from "drizzle-orm";
import { db } from "@/lib/db/client";
import { and, desc, eq, inArray, isNull, sql } from "drizzle-orm";
import { db, pg } from "@/lib/db/client";
import { memories, projects, auditLog } from "@/lib/db/schema";
import {
MemoryIdInput,
MemoryListInput,
MemorySearchInput,
MemoryUpdateInput,
MemoryWriteInput,
ProjectIdentifyInput,
} from "@shared-memory/schemas";
import { embedText } from "@/lib/embedder";
import type { UserContext } from "./context";
/**
* MCP tool definitions for v1 (Phase 1). Each tool has:
* MCP tool definitions. Each tool has:
* - name: dotted identifier exposed to clients
* - description: shown to the model
* - inputSchema: JSON Schema for the arguments object
* - handler: async function that runs the tool
*
* Search (memory.search) and snippets come in later phases.
*/
export interface ToolResult {
@@ -61,6 +62,11 @@ async function resolveProjectId(
return row[0]?.id ?? null;
}
/** pgvector accepts vectors as text literals like "[0.1,0.2,...]". */
function toVectorLiteral(v: number[]): string {
return `[${v.join(",")}]`;
}
// ---------- tools ----------
const projectIdentify: ToolDef = {
@@ -151,6 +157,12 @@ const memoryWrite: ToolDef = {
}
}
// Embed inline so the new memory is searchable immediately. Slower
// writes (~50150 ms) are an acceptable price for that guarantee; if
// embedder pressure ever forces an async path, only this section
// needs to change.
const embedding = await embedText(parsed.data.content);
const inserted = await db
.insert(memories)
.values({
@@ -159,6 +171,7 @@ const memoryWrite: ToolDef = {
scope,
content: parsed.data.content,
tags: parsed.data.tags ?? [],
embedding,
})
.returning({ id: memories.id, createdAt: memories.createdAt });
@@ -297,11 +310,227 @@ const memoryDelete: ToolDef = {
},
};
const memoryUpdate: ToolDef = {
name: "memory.update",
description:
"Edit an existing memory. Provide id and any of content or tags. If content changes, the embedding is re-computed automatically. Useful for fixing a typo without re-creating the row.",
inputSchema: {
type: "object",
properties: {
id: { type: "string", format: "uuid" },
content: { type: "string", description: "Replacement content (164,000 chars)." },
tags: { type: "array", items: { type: "string" }, description: "Replacement tag list." },
},
required: ["id"],
},
async handler(args, ctx) {
const parsed = MemoryUpdateInput.safeParse(args);
if (!parsed.success) return err(parsed.error.message);
const existing = await db
.select({ id: memories.id, content: memories.content })
.from(memories)
.where(
and(
eq(memories.id, parsed.data.id),
eq(memories.userId, ctx.userId),
isNull(memories.deletedAt),
),
)
.limit(1);
if (!existing[0]) return err("not found");
const update: Record<string, unknown> = { updatedAt: new Date() };
if (parsed.data.tags !== undefined) update.tags = parsed.data.tags;
if (parsed.data.content !== undefined && parsed.data.content !== existing[0].content) {
update.content = parsed.data.content;
update.embedding = await embedText(parsed.data.content);
}
const updated = await db
.update(memories)
.set(update)
.where(eq(memories.id, parsed.data.id))
.returning({ id: memories.id, updatedAt: memories.updatedAt });
await db.insert(auditLog).values({
userId: ctx.userId,
actor: "mcp",
action: "memory.update",
entityType: "memory",
entityId: updated[0]!.id,
payload: {
fields: Object.keys(update).filter((k) => k !== "updatedAt"),
},
});
return ok(updated[0]!, `updated memory ${updated[0]!.id}`);
},
};
interface RankAccumulator {
vectorRank?: number;
ftsRank?: number;
tagRank?: number;
rrfScore: number;
}
const memorySearch: ToolDef = {
name: "memory.search",
description:
"Hybrid search across this user's memories. Combines three signals — vector similarity (semantic), Postgres full-text rank (keyword), and tag overlap — via reciprocal rank fusion. Returns top results with per-source ranks visible so the model can judge confidence.",
inputSchema: {
type: "object",
properties: {
query: { type: "string", description: "Natural-language query." },
project: { type: "string", description: "Restrict to a single project key." },
scope: { type: "string", enum: ["project", "user"] },
tags: { type: "array", items: { type: "string" }, description: "Boost results with these tags." },
limit: { type: "integer", minimum: 1, maximum: 50, default: 10 },
},
required: ["query"],
},
async handler(args, ctx) {
const parsed = MemorySearchInput.safeParse(args);
if (!parsed.success) return err(parsed.error.message);
const { query, scope, tags, limit } = parsed.data;
const projectId = parsed.data.project
? await resolveProjectId(ctx, parsed.data.project)
: null;
if (parsed.data.project && !projectId) {
return ok({ items: [], _ranks: {} }, "0 results (unknown project)");
}
const queryVec = await embedText(query);
const vecLit = toVectorLiteral(queryVec);
const CANDIDATES = 50;
const RRF_K = 60;
// Run the three candidate-fetch queries in parallel. The filter is
// expressed via pg's tagged-template binding so values are safely
// interpolated.
const userId = ctx.userId;
const vecPromise = pg<{ id: string }[]>`
SELECT id
FROM memories
WHERE user_id = ${userId}
AND deleted_at IS NULL
AND embedding IS NOT NULL
${scope ? pg`AND scope = ${scope}` : pg``}
${projectId ? pg`AND project_id = ${projectId}` : pg``}
ORDER BY embedding <=> ${vecLit}::vector ASC
LIMIT ${CANDIDATES}
`;
const ftsPromise = pg<{ id: string }[]>`
SELECT id
FROM memories, plainto_tsquery('english', ${query}) AS q
WHERE user_id = ${userId}
AND deleted_at IS NULL
AND content_tsv @@ q
${scope ? pg`AND scope = ${scope}` : pg``}
${projectId ? pg`AND project_id = ${projectId}` : pg``}
ORDER BY ts_rank_cd(content_tsv, q) DESC
LIMIT ${CANDIDATES}
`;
const tagPromise =
tags && tags.length > 0
? pg<{ id: string }[]>`
SELECT id
FROM memories
WHERE user_id = ${userId}
AND deleted_at IS NULL
AND tags && ${tags}::text[]
${scope ? pg`AND scope = ${scope}` : pg``}
${projectId ? pg`AND project_id = ${projectId}` : pg``}
ORDER BY cardinality(
ARRAY(SELECT unnest(tags) INTERSECT SELECT unnest(${tags}::text[]))
) DESC
LIMIT ${CANDIDATES}
`
: Promise.resolve([] as { id: string }[]);
const [vecHits, ftsHits, tagHits] = await Promise.all([
vecPromise,
ftsPromise,
tagPromise,
]);
// Fuse via RRF: score(d) = Σ_r 1/(k + rank_r(d))
const scores = new Map<string, RankAccumulator>();
const accum = (id: string, rank: number, key: "vectorRank" | "ftsRank" | "tagRank") => {
const e = scores.get(id) ?? { rrfScore: 0 };
e[key] = rank;
e.rrfScore += 1 / (RRF_K + rank);
scores.set(id, e);
};
vecHits.forEach((h, i) => accum(h.id, i + 1, "vectorRank"));
ftsHits.forEach((h, i) => accum(h.id, i + 1, "ftsRank"));
tagHits.forEach((h, i) => accum(h.id, i + 1, "tagRank"));
if (scores.size === 0) {
return ok({ items: [], debug: { vec: 0, fts: 0, tag: 0 } }, "0 results");
}
const sorted = [...scores.entries()]
.sort(([, a], [, b]) => b.rrfScore - a.rrfScore)
.slice(0, limit);
const topIds = sorted.map(([id]) => id);
const rows = await db
.select({
id: memories.id,
scope: memories.scope,
projectId: memories.projectId,
content: memories.content,
tags: memories.tags,
createdAt: memories.createdAt,
updatedAt: memories.updatedAt,
})
.from(memories)
.where(inArray(memories.id, topIds));
const byId = new Map(rows.map((r) => [r.id, r]));
const items = sorted.flatMap(([id, rank]) => {
const row = byId.get(id);
if (!row) return [];
return [
{
...row,
_rank: {
rrfScore: Number(rank.rrfScore.toFixed(6)),
vectorRank: rank.vectorRank ?? null,
ftsRank: rank.ftsRank ?? null,
tagRank: rank.tagRank ?? null,
},
},
];
});
return ok(
{
items,
debug: {
vec: vecHits.length,
fts: ftsHits.length,
tag: tagHits.length,
},
},
`${items.length} result(s)`,
);
},
};
export const tools: ToolDef[] = [
projectIdentify,
memoryWrite,
memoryUpdate,
memoryList,
memoryGet,
memorySearch,
memoryDelete,
];
+77
View File
@@ -72,11 +72,88 @@ async function main() {
}
console.log("Migrations complete.");
if (process.env.EMBEDDER_URL) {
await backfillEmbeddings(sql);
} else {
console.log("EMBEDDER_URL not set — skipping embedding backfill.");
}
} finally {
await sql.end({ timeout: 5 });
}
}
/**
* Backfill embedding column for any memory written before embeddings were
* online. Idempotent: only touches rows where embedding IS NULL. Runs on
* every migrator boot, so deploying Phase 2 — or recovering from an
* embedder outage that left fresh rows unembedded — needs no manual step.
*/
async function backfillEmbeddings(sql: ReturnType<typeof postgres>) {
const embedderUrl = process.env.EMBEDDER_URL!.replace(/\/$/, "");
const BATCH = 32;
// Wait for the embedder to report ready — its first boot has to download
// and load the model, which can take 3060s on a cold container.
const waitDeadline = Date.now() + 180_000;
for (;;) {
try {
const res = await fetch(`${embedderUrl}/health`);
if (res.ok) {
const body = (await res.json()) as { ready?: boolean };
if (body.ready) break;
}
} catch {
/* embedder not up yet */
}
if (Date.now() > waitDeadline) {
throw new Error("embedder did not become ready within 180s");
}
await new Promise((r) => setTimeout(r, 2000));
}
let total = 0;
for (;;) {
const rows = await sql<{ id: string; content: string }[]>`
SELECT id, content FROM memories
WHERE embedding IS NULL AND deleted_at IS NULL
ORDER BY created_at
LIMIT ${BATCH}
`;
if (rows.length === 0) break;
const res = await fetch(`${embedderUrl}/embed`, {
method: "POST",
headers: { "content-type": "application/json" },
body: JSON.stringify({ texts: rows.map((r) => r.content) }),
});
if (!res.ok) {
const detail = await res.text().catch(() => "");
throw new Error(`embedder error ${res.status}: ${detail.slice(0, 200)}`);
}
const { vectors } = (await res.json()) as { vectors: number[][] };
await sql.begin(async (tx) => {
for (let i = 0; i < rows.length; i++) {
const id = rows[i]!.id;
const vec = vectors[i];
if (!vec) continue;
const literal = `[${vec.join(",")}]`;
await tx`UPDATE memories SET embedding = ${literal}::vector WHERE id = ${id}`;
}
});
total += rows.length;
console.log(` embedded ${rows.length} memories (total: ${total})`);
}
if (total === 0) {
console.log("Embedding backfill: nothing to do.");
} else {
console.log(`Embedding backfill complete: ${total} memories embedded.`);
}
}
main().catch((err) => {
console.error("Migration failed:", err);
process.exit(1);