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
@@ -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,
|
||||
];
|
||||
|
||||
|
||||
Reference in New Issue
Block a user