@remit/search-service 0.0.1

package/package.json

@@ -0,0 +1,66 @@
+{
+	"name": "@remit/search-service",
+	"version": "0.0.1",
+	"type": "module",
+	"main": "src/index.ts",
+	"types": "src/index.ts",
+	"exports": {
+		".": {
+			"types": "./src/index.ts",
+			"default": "./src/index.ts"
+		},
+		"./from-env": {
+			"types": "./src/from-env.ts",
+			"default": "./src/from-env.ts"
+		},
+		"./s3-vectors": {
+			"types": "./src/s3-vectors.ts",
+			"default": "./src/s3-vectors.ts"
+		},
+		"./bedrock": {
+			"types": "./src/bedrock.ts",
+			"default": "./src/bedrock.ts"
+		},
+		"./sqlite-vec": {
+			"types": "./src/sqlite-vec.ts",
+			"default": "./src/sqlite-vec.ts"
+		},
+		"./pgvector": {
+			"types": "./src/pgvector.ts",
+			"default": "./src/pgvector.ts"
+		}
+	},
+	"scripts": {
+		"test:typecheck": "tsgo --noEmit",
+		"test:run": "node --env-file=../../localhost-test-unit.env --import tsx --test 'src/**/*.test.ts'",
+		"test:integ": "RUN_INTEG_TESTS=1 node --env-file=../../localhost-test-unit.env --import tsx --test 'src/**/*.integ.test.ts'",
+		"test": "npm run test:typecheck && npm run test:run"
+	},
+	"dependencies": {
+		"@huggingface/transformers": "^4.2.0",
+		"@remit/domain-enums": "*",
+		"@remit/api-openapi-types": "*",
+		"better-sqlite3": "^12.11.1",
+		"p-limit": "^6.2.0",
+		"pg": "^8.0.0",
+		"sqlite-vec": "^0.1.9",
+		"zod": "*"
+	},
+	"devDependencies": {
+		"@aws-sdk/client-bedrock-runtime": "*",
+		"@aws-sdk/client-s3vectors": "*",
+		"@remit/storage-service": "*",
+		"@types/better-sqlite3": "^7.6.13",
+		"@types/pg": "^8.20.0",
+		"aws-sdk-client-mock": "*"
+	},
+	"license": "MIT",
+	"publishConfig": {
+		"access": "public"
+	},
+	"repository": {
+		"type": "git",
+		"url": "git+https://github.com/remit-mail/remit.git",
+		"directory": "packages/search-service"
+	}
+}

package/src/anchor.test.ts

@@ -0,0 +1,164 @@
+import assert from "node:assert/strict";
+import { describe, it } from "node:test";
+import {
+	buildAnchorSourceText,
+	buildMessageAnchor,
+	poolChunkVectors,
+} from "./anchor.js";
+import { createMemoryVectorStore } from "./backends/memory.js";
+import { createDeterministicEmbeddingService } from "./embeddings.js";
+import type { ChunkMetadata, VectorRecord } from "./types.js";
+
+const l2Norm = (vector: number[]): number =>
+	Math.sqrt(vector.reduce((sum, value) => sum + value * value, 0));
+
+const baseMetadata = (
+	overrides: Partial<ChunkMetadata> = {},
+): ChunkMetadata => ({
+	messageId: "msg-1",
+	threadId: "thread-1",
+	accountConfigId: "acct-1",
+	mailboxIds: ["mbox-1"],
+	chunkType: "body",
+	sentDate: 1_700_000_000,
+	isRead: false,
+	hasAttachment: false,
+	hasStars: false,
+	...overrides,
+});
+
+describe("poolChunkVectors", () => {
+	it("mean-pools chunk vectors and L2-normalizes the result", () => {
+		const pooled = poolChunkVectors([
+			[2, 0, 0],
+			[0, 2, 0],
+		]);
+		// Mean is [1, 1, 0]; normalized to unit length.
+		assert.ok(Math.abs(l2Norm(pooled) - 1) < 1e-9);
+		assert.ok(Math.abs(pooled[0] - pooled[1]) < 1e-9);
+		assert.ok(Math.abs(pooled[2]) < 1e-9);
+	});
+
+	it("returns the single vector's direction when pooling one chunk", () => {
+		const pooled = poolChunkVectors([[3, 4]]);
+		assert.ok(Math.abs(pooled[0] - 0.6) < 1e-9);
+		assert.ok(Math.abs(pooled[1] - 0.8) < 1e-9);
+	});
+
+	it("returns a zero vector unchanged rather than dividing by a zero norm", () => {
+		assert.deepEqual(poolChunkVectors([[0, 0, 0]]), [0, 0, 0]);
+	});
+
+	it("throws on an empty set", () => {
+		assert.throws(() => poolChunkVectors([]), /empty set/);
+	});
+
+	it("throws on a dimension mismatch", () => {
+		assert.throws(
+			() =>
+				poolChunkVectors([
+					[1, 2, 3],
+					[1, 2],
+				]),
+			/dimension mismatch/,
+		);
+	});
+});
+
+describe("buildAnchorSourceText", () => {
+	it("prefers subject then body previews", () => {
+		const text = buildAnchorSourceText([
+			{ chunkType: "body", textPreview: "the body" },
+			{ chunkType: "subject", textPreview: "the subject" },
+			{ chunkType: "sender", textPreview: "someone@example.com" },
+		]);
+		assert.equal(text, "the subject\nthe body");
+	});
+
+	it("falls back to every preview when there is no subject or body", () => {
+		const text = buildAnchorSourceText([
+			{ chunkType: "sender", textPreview: "someone@example.com" },
+			{ chunkType: "entities", textPreview: "ACME Corp" },
+		]);
+		assert.equal(text, "someone@example.com\nACME Corp");
+	});
+
+	it("never exceeds the 512-char preview bound", () => {
+		const long = "x".repeat(5000);
+		const text = buildAnchorSourceText([
+			{ chunkType: "body", textPreview: long },
+		]);
+		assert.ok(text.length <= 512);
+	});
+});
+
+describe("buildMessageAnchor", () => {
+	const embedder = createDeterministicEmbeddingService({ dimensions: 8 });
+
+	const putRecord = async (
+		store: ReturnType<typeof createMemoryVectorStore>,
+		record: VectorRecord,
+	): Promise<void> => store.upsert([record]);
+
+	it("pools a message's chunk vectors and derives the anchor payload", async () => {
+		const store = createMemoryVectorStore();
+		await putRecord(store, {
+			chunkId: "msg-1::subject",
+			vector: [1, 0, 0, 0, 0, 0, 0, 0],
+			metadata: baseMetadata({
+				chunkType: "subject",
+				textPreview: "booking confirmed",
+			}),
+		});
+		await putRecord(store, {
+			chunkId: "msg-1::body-0",
+			vector: [0, 1, 0, 0, 0, 0, 0, 0],
+			metadata: baseMetadata({
+				chunkType: "body",
+				textPreview: "your trip is booked",
+			}),
+		});
+
+		const anchor = await buildMessageAnchor(
+			{ store, embedder },
+			{ accountConfigId: "acct-1", anchorMessageId: "msg-1" },
+		);
+
+		assert.ok(anchor);
+		assert.equal(anchor.anchorEmbeddingId, embedder.embeddingId);
+		assert.equal(anchor.anchorEmbedding.length, 8);
+		assert.ok(Math.abs(l2Norm(anchor.anchorEmbedding) - 1) < 1e-9);
+		assert.equal(
+			anchor.anchorSourceText,
+			"booking confirmed\nyour trip is booked",
+		);
+	});
+
+	it("returns null when the message has no indexed chunks", async () => {
+		const store = createMemoryVectorStore();
+		const anchor = await buildMessageAnchor(
+			{ store, embedder },
+			{ accountConfigId: "acct-1", anchorMessageId: "absent" },
+		);
+		assert.equal(anchor, null);
+	});
+
+	it("ignores chunks belonging to another account", async () => {
+		const store = createMemoryVectorStore();
+		await putRecord(store, {
+			chunkId: "msg-1::subject",
+			vector: [1, 0, 0, 0, 0, 0, 0, 0],
+			metadata: baseMetadata({
+				chunkType: "subject",
+				accountConfigId: "other-acct",
+				textPreview: "not mine",
+			}),
+		});
+
+		const anchor = await buildMessageAnchor(
+			{ store, embedder },
+			{ accountConfigId: "acct-1", anchorMessageId: "msg-1" },
+		);
+		assert.equal(anchor, null);
+	});
+});

package/src/anchor.ts

@@ -0,0 +1,127 @@
+import type { VectorStoreService } from "./backends/memory.js";
+import type { EmbeddingService } from "./embeddings.js";
+import { buildTextPreview } from "./search.js";
+import type { ChunkType, VectorRecord } from "./types.js";
+
+/**
+ * The persisted anchor of a semantic filter (RFC 034 Decision 2.1). A single
+ * mean-pooled snapshot of the anchor message's chunk vectors, plus the bounded
+ * source text needed to re-embed it after a model migration (Decision 2.4).
+ * Written once onto the sibling `FilterAnchor` row at filter-save time; never
+ * re-derived per match.
+ */
+export interface AnchorPayload {
+	anchorEmbedding: number[];
+	anchorEmbeddingId: string;
+	anchorSourceText: string;
+}
+
+/**
+ * Mean-pool chunk vectors into one, then L2-normalize so the anchor is unit
+ * length like the per-chunk vectors it summarizes (cosine is scale-invariant,
+ * but a normalized anchor keeps the stored vector consistent with the chunk
+ * vectors it was pooled from). Throws on an empty set or a dimension mismatch —
+ * both are programmer errors the caller cannot recover from (let it crash).
+ */
+export const poolChunkVectors = (vectors: number[][]): number[] => {
+	if (vectors.length === 0) {
+		throw new Error("Cannot pool an empty set of chunk vectors");
+	}
+	const dimensions = vectors[0].length;
+	const sum = new Array<number>(dimensions).fill(0);
+	for (const vector of vectors) {
+		if (vector.length !== dimensions) {
+			throw new Error(
+				`Chunk vector dimension mismatch: ${vector.length} vs ${dimensions}`,
+			);
+		}
+		for (let i = 0; i < dimensions; i++) sum[i] += vector[i];
+	}
+	for (let i = 0; i < dimensions; i++) sum[i] /= vectors.length;
+
+	let norm = 0;
+	for (const value of sum) norm += value * value;
+	norm = Math.sqrt(norm);
+	if (norm === 0) return sum;
+	return sum.map((value) => value / norm);
+};
+
+// The chunk types whose text carries the semantic meaning of "messages like
+// this" — subject and body. Structured chunks (sender, recipient, attachment,
+// entities) are excluded from the re-embeddable source text; they add no signal
+// a user's plain-sentence anchor is about.
+const SOURCE_TEXT_CHUNK_TYPES: readonly ChunkType[] = ["subject", "body"];
+
+/**
+ * Assemble the anchor's re-embeddable source text (RFC 034 Decision 2.4) from
+ * the message's chunk previews, preferring subject then body. Falls back to
+ * every available preview when the message has neither. Bounded by the same
+ * `buildTextPreview` char/byte budget the chunk vectors already pay, so it never
+ * exceeds the 512-char (and S3 Vectors byte) cap.
+ */
+export const buildAnchorSourceText = (
+	chunks: Array<{ chunkType: ChunkType; textPreview?: string }>,
+): string => {
+	const previews: string[] = [];
+	for (const type of SOURCE_TEXT_CHUNK_TYPES) {
+		for (const chunk of chunks) {
+			if (chunk.chunkType === type && chunk.textPreview) {
+				previews.push(chunk.textPreview);
+			}
+		}
+	}
+	if (previews.length === 0) {
+		for (const chunk of chunks) {
+			if (chunk.textPreview) previews.push(chunk.textPreview);
+		}
+	}
+	return buildTextPreview(previews.join("\n"));
+};
+
+export interface AnchorBuildDeps {
+	store: Pick<VectorStoreService, "getByMessage">;
+	embedder: Pick<EmbeddingService, "embeddingId">;
+}
+
+export interface AnchorBuildParams {
+	accountConfigId: string;
+	anchorMessageId: string;
+}
+
+/**
+ * Build a filter's persisted anchor from the anchor message's already-indexed
+ * chunk vectors (RFC 034 Decision 2.1). Reads the message's chunk vectors,
+ * pools them into one, and derives the bounded source text — never embeds
+ * anything new here; the vectors already exist from index time. Returns `null`
+ * when the message has no indexed chunks, so the caller can decline to write a
+ * `FilterAnchor` row (and leave `Filter.hasAnchor` false) rather than persist an
+ * empty anchor.
+ *
+ * `anchorEmbeddingId` is the current embedder's identifier — the model the
+ * chunk vectors were embedded under, which the indexing pipeline keeps current.
+ */
+export const buildMessageAnchor = async (
+	deps: AnchorBuildDeps,
+	params: AnchorBuildParams,
+): Promise<AnchorPayload | null> => {
+	const records = (await deps.store.getByMessage(params.anchorMessageId))
+		.filter(
+			(record: VectorRecord) =>
+				record.metadata.accountConfigId === params.accountConfigId,
+		)
+		.sort((a: VectorRecord, b: VectorRecord) =>
+			a.chunkId.localeCompare(b.chunkId),
+		);
+	if (records.length === 0) return null;
+
+	return {
+		anchorEmbedding: poolChunkVectors(records.map((r) => r.vector)),
+		anchorEmbeddingId: deps.embedder.embeddingId,
+		anchorSourceText: buildAnchorSourceText(
+			records.map((r) => ({
+				chunkType: r.metadata.chunkType,
+				textPreview: r.metadata.textPreview,
+			})),
+		),
+	};
+};

package/src/backends/bedrock.test.ts

@@ -0,0 +1,148 @@
+import assert from "node:assert/strict";
+import { describe, it } from "node:test";
+import {
+	BedrockRuntimeClient,
+	InvokeModelCommand,
+} from "@aws-sdk/client-bedrock-runtime";
+import { mockClient } from "aws-sdk-client-mock";
+import { BedrockEmbeddingService } from "./bedrock.js";
+
+const encodeTitanResponse = (dimensions: number): Uint8Array =>
+	new TextEncoder().encode(
+		JSON.stringify({ embedding: new Array(dimensions).fill(0) }),
+	);
+
+describe("BedrockEmbeddingService", () => {
+	it("caps in-flight InvokeModel calls to the configured concurrency", async () => {
+		const concurrency = 6;
+		const total = 20;
+		let inFlight = 0;
+		let peak = 0;
+
+		const bedrockMock = mockClient(BedrockRuntimeClient);
+		bedrockMock.on(InvokeModelCommand).callsFake(async () => {
+			inFlight += 1;
+			peak = Math.max(peak, inFlight);
+			await new Promise((resolve) => setTimeout(resolve, 10));
+			inFlight -= 1;
+			return { body: encodeTitanResponse(1024) };
+		});
+
+		const service = new BedrockEmbeddingService({
+			client: new BedrockRuntimeClient({}),
+			concurrency,
+		});
+
+		const texts = Array.from({ length: total }, (_, i) => `text-${i}`);
+		const results = await service.embed(texts);
+
+		assert.equal(results.length, total);
+		assert.ok(
+			peak <= concurrency,
+			`peak in-flight ${peak} exceeded concurrency ${concurrency}`,
+		);
+		assert.equal(bedrockMock.commandCalls(InvokeModelCommand).length, total);
+	});
+
+	it("truncates over-budget input before sending it to Bedrock", async () => {
+		const bedrockMock = mockClient(BedrockRuntimeClient);
+		let sentLength = 0;
+		bedrockMock.on(InvokeModelCommand).callsFake((input) => {
+			const body = JSON.parse(input.body as string) as { inputText: string };
+			sentLength = body.inputText.length;
+			return { body: encodeTitanResponse(1024) };
+		});
+
+		const service = new BedrockEmbeddingService({
+			client: new BedrockRuntimeClient({}),
+		});
+
+		await service.embed(["x".repeat(50000)]);
+
+		assert.ok(
+			sentLength <= 6000,
+			`inputText length ${sentLength} exceeds the 6000-char budget`,
+		);
+	});
+
+	it("shrinks and retries when Bedrock rejects a dense input as over-budget", async () => {
+		// A 6000-char chunk of dense/non-Latin text can still exceed Titan's
+		// 8192-token limit even though it's within the char budget. Rather than
+		// dead-letter the message forever, the embedder halves the input and
+		// retries until Bedrock accepts it (#910).
+		const bedrockMock = mockClient(BedrockRuntimeClient);
+		const sentLengths: number[] = [];
+		bedrockMock.on(InvokeModelCommand).callsFake((input) => {
+			const body = JSON.parse(input.body as string) as { inputText: string };
+			sentLengths.push(body.inputText.length);
+			if (body.inputText.length > 3000) {
+				throw new Error(
+					"ValidationException: 400 Bad Request: Too many input tokens. Max input tokens: 8192, request input token count: 14022",
+				);
+			}
+			return { body: encodeTitanResponse(1024) };
+		});
+
+		const service = new BedrockEmbeddingService({
+			client: new BedrockRuntimeClient({}),
+		});
+
+		const [vector] = await service.embed(["字".repeat(50000)]);
+
+		assert.equal(vector.length, 1024);
+		assert.equal(sentLengths[0], 6000, "first attempt uses the full budget");
+		assert.ok(
+			// biome-ignore lint/style/noNonNullAssertion: test assertion, value is guaranteed by test setup
+			sentLengths.at(-1) !== undefined && sentLengths.at(-1)! <= 3000,
+			"a later attempt shrinks under the token limit",
+		);
+	});
+
+	it("surfaces a non-token error without shrinking and retrying", async () => {
+		const bedrockMock = mockClient(BedrockRuntimeClient);
+		let calls = 0;
+		bedrockMock.on(InvokeModelCommand).callsFake(() => {
+			calls += 1;
+			throw new Error("AccessDeniedException: not authorized");
+		});
+
+		const service = new BedrockEmbeddingService({
+			client: new BedrockRuntimeClient({}),
+		});
+
+		await assert.rejects(
+			service.embed(["hello"]),
+			/AccessDeniedException/,
+			"non-token errors must propagate",
+		);
+		assert.equal(calls, 1, "no shrink-retry for non-token errors");
+	});
+
+	it("returns embeddings in the same order as the input texts", async () => {
+		const bedrockMock = mockClient(BedrockRuntimeClient);
+		bedrockMock.on(InvokeModelCommand).callsFake(async (input) => {
+			const body = JSON.parse(input.body as string) as { inputText: string };
+			const tag = Number.parseInt(body.inputText.split("-")[1] ?? "0", 10);
+			await new Promise((resolve) => setTimeout(resolve, (5 - (tag % 5)) * 2));
+			return {
+				body: new TextEncoder().encode(
+					JSON.stringify({ embedding: [tag, 0, 0, 0] }),
+				),
+			};
+		});
+
+		const service = new BedrockEmbeddingService({
+			client: new BedrockRuntimeClient({}),
+			dimensions: 4,
+			concurrency: 3,
+		});
+
+		const texts = ["text-0", "text-1", "text-2", "text-3", "text-4"];
+		const results = await service.embed(texts);
+
+		assert.deepEqual(
+			results.map((v) => v[0]),
+			[0, 1, 2, 3, 4],
+		);
+	});
+});

package/src/backends/bedrock.ts

@@ -0,0 +1,105 @@
+import {
+	BedrockRuntimeClient,
+	InvokeModelCommand,
+} from "@aws-sdk/client-bedrock-runtime";
+import pLimit from "p-limit";
+import type { EmbeddingService } from "../embeddings.js";
+
+export interface BedrockEmbeddingConfig {
+	client?: BedrockRuntimeClient;
+	region?: string;
+	modelId?: string;
+	dimensions?: number;
+	concurrency?: number;
+}
+
+const DEFAULT_MODEL_ID = "amazon.titan-embed-text-v2:0";
+const DEFAULT_DIMENSIONS = 1024;
+const DEFAULT_CONCURRENCY = 6;
+
+/**
+ * Titan v2 rejects inputs over its 8192-token limit. Chunk builders already cap
+ * their output, but this is a final backstop so no chunk source can ever push
+ * an over-budget inputText to Bedrock. 6000 chars is safe for typical Latin text
+ * (~4 chars/token), but dense or non-Latin scripts (CJK, packed entity lists)
+ * can exceed 8192 tokens at far fewer chars — observed up to ~21k tokens from a
+ * 6000-char chunk. On a token-limit rejection we halve and retry rather than let
+ * a valid-but-dense message dead-letter forever (#910).
+ */
+const MAX_INPUT_CHARS = 6000;
+// Below this we stop shrinking and let the error surface: a chunk this small
+// that still overflows is not something a retry can fix.
+const MIN_INPUT_CHARS = 256;
+
+const isTokenLimitError = (error: unknown): boolean =>
+	error instanceof Error &&
+	/too many input tokens|input is too long|maximum.*token/i.test(error.message);
+
+const isNumberArray = (value: unknown): value is number[] =>
+	Array.isArray(value) && value.every((n) => typeof n === "number");
+
+const parseTitanResponse = (raw: Uint8Array): number[] => {
+	const text = new TextDecoder().decode(raw);
+	const parsed: unknown = JSON.parse(text);
+	if (
+		typeof parsed === "object" &&
+		parsed !== null &&
+		"embedding" in parsed &&
+		isNumberArray((parsed as { embedding: unknown }).embedding)
+	) {
+		return (parsed as { embedding: number[] }).embedding;
+	}
+	throw new Error("Bedrock Titan response missing 'embedding' array");
+};
+
+export class BedrockEmbeddingService implements EmbeddingService {
+	private client: BedrockRuntimeClient;
+	private modelId: string;
+	readonly dimensions: number;
+	readonly embeddingId: string;
+	private limit: ReturnType<typeof pLimit>;
+
+	constructor(config: BedrockEmbeddingConfig = {}) {
+		this.client =
+			config.client ?? new BedrockRuntimeClient({ region: config.region });
+		this.modelId = config.modelId ?? DEFAULT_MODEL_ID;
+		this.dimensions = config.dimensions ?? DEFAULT_DIMENSIONS;
+		this.embeddingId = `${this.modelId}@${this.dimensions}`;
+		this.limit = pLimit(config.concurrency ?? DEFAULT_CONCURRENCY);
+	}
+
+	embed = async (texts: string[]): Promise<number[][]> =>
+		Promise.all(texts.map((text) => this.limit(() => this.embedOne(text))));
+
+	private embedOne = async (text: string): Promise<number[]> => {
+		let charBudget = MAX_INPUT_CHARS;
+		while (true) {
+			try {
+				return await this.invoke(text.slice(0, charBudget));
+			} catch (error) {
+				if (!isTokenLimitError(error) || charBudget <= MIN_INPUT_CHARS) {
+					throw error;
+				}
+				charBudget = Math.max(MIN_INPUT_CHARS, Math.floor(charBudget / 2));
+			}
+		}
+	};
+
+	private invoke = async (inputText: string): Promise<number[]> => {
+		const cmd = new InvokeModelCommand({
+			modelId: this.modelId,
+			contentType: "application/json",
+			accept: "application/json",
+			body: JSON.stringify({
+				inputText,
+				dimensions: this.dimensions,
+				normalize: true,
+			}),
+		});
+		const response = await this.client.send(cmd);
+		if (!response.body) {
+			throw new Error("Bedrock InvokeModel returned empty body");
+		}
+		return parseTitanResponse(response.body);
+	};
+}