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:
Claude Opus 4.7
+2
-2
@@ -48,9 +48,9 @@ POSTGRES_DB=memory
|
||||
# DATABASE_URL=postgres://memory:...@db:5432/memory
|
||||
|
||||
# -----------------------------------------------------------------------------
|
||||
# Embedder sidecar (added in Phase 2; leave EMBEDDER_URL empty in Phase 1)
|
||||
# Embedder sidecar. Default points at the in-compose service.
|
||||
# -----------------------------------------------------------------------------
|
||||
EMBEDDER_URL=
|
||||
EMBEDDER_URL=http://embedder:8080
|
||||
EMBEDDING_MODEL=Xenova/bge-small-en-v1.5
|
||||
EMBEDDING_DIM=384
|
||||
|
||||
|
||||
@@ -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"]
|
||||
@@ -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"
|
||||
}
|
||||
}
|
||||
@@ -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);
|
||||
});
|
||||
@@ -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"]
|
||||
}
|
||||
@@ -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
@@ -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
@@ -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 (~50–150 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 (1–64,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,
|
||||
];
|
||||
|
||||
|
||||
@@ -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 30–60s 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);
|
||||
|
||||
+33
-3
@@ -41,8 +41,32 @@ services:
|
||||
networks:
|
||||
- internal
|
||||
|
||||
# One-shot migration runner. Exits 0 when migrations are up-to-date;
|
||||
# `app` waits on its successful completion before starting.
|
||||
# Embedding sidecar — loads bge-small-en-v1.5 once and serves /embed.
|
||||
# First boot downloads the model (~30 MB) into a named volume so future
|
||||
# boots are warm.
|
||||
embedder:
|
||||
image: ${EMBEDDER_IMAGE_REF:-shared-memory-embedder:local}
|
||||
build:
|
||||
context: .
|
||||
dockerfile: apps/embedder/Dockerfile
|
||||
restart: unless-stopped
|
||||
environment:
|
||||
EMBEDDING_MODEL: ${EMBEDDING_MODEL:-Xenova/bge-small-en-v1.5}
|
||||
EMBEDDING_DIM: ${EMBEDDING_DIM:-384}
|
||||
LOG_LEVEL: ${LOG_LEVEL:-info}
|
||||
volumes:
|
||||
- embedder_models:/data/models
|
||||
healthcheck:
|
||||
test: ["CMD-SHELL", "wget -q -O - http://127.0.0.1:8080/health | grep -q '\"ready\":true' || exit 1"]
|
||||
interval: 15s
|
||||
timeout: 5s
|
||||
retries: 5
|
||||
start_period: 180s
|
||||
networks:
|
||||
- internal
|
||||
|
||||
# One-shot migration runner + embedding backfill. Exits 0 when both are
|
||||
# up-to-date; `app` waits on its successful completion before starting.
|
||||
migrator:
|
||||
image: ${IMAGE_REF:-shared-memory-web:local}
|
||||
build:
|
||||
@@ -52,8 +76,11 @@ services:
|
||||
depends_on:
|
||||
db:
|
||||
condition: service_healthy
|
||||
embedder:
|
||||
condition: service_healthy
|
||||
environment:
|
||||
DATABASE_URL: postgres://${POSTGRES_USER}:${POSTGRES_PASSWORD}@db:5432/${POSTGRES_DB}
|
||||
EMBEDDER_URL: ${EMBEDDER_URL:-http://embedder:8080}
|
||||
command: ["node", "apps/web/migrate.mjs"]
|
||||
networks:
|
||||
- internal
|
||||
@@ -67,6 +94,8 @@ services:
|
||||
depends_on:
|
||||
db:
|
||||
condition: service_healthy
|
||||
embedder:
|
||||
condition: service_healthy
|
||||
migrator:
|
||||
condition: service_completed_successfully
|
||||
environment:
|
||||
@@ -87,7 +116,7 @@ services:
|
||||
|
||||
DATABASE_URL: postgres://${POSTGRES_USER}:${POSTGRES_PASSWORD}@db:5432/${POSTGRES_DB}
|
||||
|
||||
EMBEDDER_URL: ${EMBEDDER_URL:-}
|
||||
EMBEDDER_URL: ${EMBEDDER_URL:-http://embedder:8080}
|
||||
EMBEDDING_MODEL: ${EMBEDDING_MODEL:-Xenova/bge-small-en-v1.5}
|
||||
EMBEDDING_DIM: ${EMBEDDING_DIM:-384}
|
||||
|
||||
@@ -137,6 +166,7 @@ volumes:
|
||||
db_data:
|
||||
caddy_data:
|
||||
caddy_config:
|
||||
embedder_models:
|
||||
|
||||
networks:
|
||||
internal:
|
||||
|
||||
@@ -42,6 +42,24 @@ export const MemoryIdInput = z.object({
|
||||
});
|
||||
export type MemoryIdInput = z.infer<typeof MemoryIdInput>;
|
||||
|
||||
export const MemoryUpdateInput = z.object({
|
||||
id: z.string().uuid(),
|
||||
content: MemoryContent.optional(),
|
||||
tags: Tags.optional(),
|
||||
}).refine((v) => v.content !== undefined || v.tags !== undefined, {
|
||||
message: "memory.update requires content or tags",
|
||||
});
|
||||
export type MemoryUpdateInput = z.infer<typeof MemoryUpdateInput>;
|
||||
|
||||
export const MemorySearchInput = z.object({
|
||||
query: z.string().min(1).max(2000),
|
||||
project: ProjectKey.optional(),
|
||||
scope: MemoryScope.optional(),
|
||||
tags: z.array(z.string()).optional(),
|
||||
limit: z.number().int().min(1).max(50).default(10),
|
||||
});
|
||||
export type MemorySearchInput = z.infer<typeof MemorySearchInput>;
|
||||
|
||||
export const ProjectIdentifyInput = z.object({
|
||||
key: ProjectKey,
|
||||
display_name: z.string().min(1).max(200).optional(),
|
||||
|
||||
Generated
+877
-5
File diff suppressed because it is too large
Load Diff
Reference in New Issue
Block a user