@prometheus-ai/memory 0.5.0

package/src/core/embeddings.ts

@@ -0,0 +1,423 @@
+import { mkdirSync } from "node:fs";
+import {
+	$env,
+	$flag,
+	APP_DISPLAY_NAME,
+	extractHttpStatusFromError,
+	fetchWithRetry,
+	getFastembedCacheDir,
+	logger,
+} from "@prometheus-ai/utils";
+import type { EmbeddingModel } from "fastembed";
+import { LRUCache } from "lru-cache/raw";
+import packageJson from "../../package.json" with { type: "json" };
+import { type EmbeddingOutput, getMnemopiRuntimeOptions, resolveEmbeddingProvider } from "./runtime-options";
+
+export type { EmbeddingOutput } from "./runtime-options";
+export { cosineSimilarity } from "./vector-math";
+
+export type Vector = Float32Array;
+export type EmbeddingMatrix = Vector[];
+
+export interface EmbeddingProvider {
+	embed(texts: readonly string[]): EmbeddingOutput | Promise<EmbeddingOutput>;
+	available?(): boolean | Promise<boolean>;
+}
+
+type StandardEmbeddingModel = Exclude<EmbeddingModel, EmbeddingModel.CUSTOM>;
+
+interface LocalEmbeddingModel {
+	embed(texts: string[], batchSize?: number): EmbeddingOutput;
+	queryEmbed?(query: string): Promise<number[]>;
+}
+
+type LocalModelInitOptions = {
+	model: StandardEmbeddingModel;
+	cacheDir?: string;
+	showDownloadProgress?: boolean;
+};
+type LocalModelInitializer = (options: LocalModelInitOptions) => Promise<LocalEmbeddingModel>;
+
+const QUERY_CACHE_MAX = 512;
+
+let providerOverride: EmbeddingProvider | null = null;
+let localModelPromise: Promise<LocalEmbeddingModel> | null = null;
+let localModelInitializer: LocalModelInitializer = defaultLocalModelInitializer;
+let apiCallCount = 0;
+const queryCache = new LRUCache<string, Vector>({ max: QUERY_CACHE_MAX });
+
+// Provider identity table for the cache key. Each unique `provider` object/function
+// (configured via `withMnemopiRuntimeOptions`) gets a stable integer id so the cache
+// scope reflects the runtime's actual embedding source. Two Mnemopi instances in the
+// same process using different providers/models hash to disjoint keys and never
+// collide on the same query text. `0` is the sentinel for "env-default fallback".
+const providerIds = new WeakMap<object, number>();
+let nextProviderId = 1;
+
+async function defaultLocalModelInitializer(options: LocalModelInitOptions): Promise<LocalEmbeddingModel> {
+	// Preload ORT 1.24 before fastembed's bundled ORT 1.21 — only on Windows,
+	// where loading the older binding first triggers a DLL-reuse crash. The 1.24
+	// line also has no darwin/x64 prebuilt, so importing it unconditionally breaks
+	// the darwin-x64 `bun build --compile` (Bun folds process.platform/arch and
+	// fails to resolve a binding that doesn't ship). The `win32` literal guard is
+	// statically foldable, so Bun dead-code-eliminates this import on every
+	// non-Windows target; fastembed loads its own ORT 1.21 binding there.
+	if (process.platform === "win32") {
+		await import("onnxruntime-node");
+	}
+	const { FlagEmbedding } = await import("fastembed");
+	return FlagEmbedding.init(options);
+}
+
+function activeEmbeddingOptions() {
+	return getMnemopiRuntimeOptions()?.embeddings;
+}
+
+/**
+ * Compose the per-query cache key. Includes the active provider's identity, the
+ * resolved model name, and the API base URL so two `Mnemopi` instances in the same
+ * process that point at different providers/models never share a cached query
+ * vector. Provider identity comes from `providerIds` (WeakMap-assigned integer);
+ * `0` is the sentinel for "no provider configured, fall back to env defaults".
+ */
+function queryCacheKey(text: string): string {
+	const active = activeEmbeddingOptions();
+	const provider = active?.provider as object | undefined;
+	let providerId = 0;
+	if (provider !== undefined) {
+		const existing = providerIds.get(provider);
+		if (existing === undefined) {
+			providerId = nextProviderId++;
+			providerIds.set(provider, providerId);
+		} else {
+			providerId = existing;
+		}
+	}
+	const model = defaultModel();
+	const apiUrl = active?.apiUrl ?? "";
+	return `${providerId}::${model}::${apiUrl}::${text}`;
+}
+
+function inTestRuntime(): boolean {
+	return $env.NODE_ENV === "test" || $env.BUN_ENV === "test";
+}
+
+function embeddingsDisabled(): boolean {
+	const active = activeEmbeddingOptions();
+	if (active?.disabled !== undefined) {
+		return active.disabled;
+	}
+	return $flag("PROMETHEUS_MEMORY_NO_EMBEDDINGS");
+}
+
+function embeddingApiKey(): string {
+	const active = activeEmbeddingOptions();
+	if (active?.apiKey !== undefined) {
+		return active.apiKey;
+	}
+	return $env.PROMETHEUS_MEMORY_EMBEDDING_API_KEY || $env.OPENROUTER_API_KEY || $env.OPENAI_API_KEY || "";
+}
+
+function embeddingBaseUrl(): string {
+	const active = activeEmbeddingOptions();
+	if (active?.apiUrl !== undefined) {
+		return active.apiUrl;
+	}
+	return $env.PROMETHEUS_MEMORY_EMBEDDING_API_URL || $env.OPENROUTER_BASE_URL || "https://openrouter.ai/api/v1";
+}
+
+function defaultModel(): string {
+	const active = activeEmbeddingOptions();
+	if (active?.model !== undefined) {
+		return active.model;
+	}
+	return $env.PROMETHEUS_MEMORY_EMBEDDING_MODEL || "BAAI/bge-small-en-v1.5";
+}
+
+/**
+ * Resolve the embedding model name for the currently active runtime scope.
+ *
+ * Reads (in order): the active provider's `model` from `withMnemopiRuntimeOptions`,
+ * the `PROMETHEUS_MEMORY_EMBEDDING_MODEL` env var, then the bundled fastembed default. Stored
+ * alongside each row in `memory_embeddings.model` so migrations can re-embed when
+ * the active model changes.
+ */
+export function currentEmbeddingModel(): string {
+	return defaultModel();
+}
+
+export function isApiModel(modelName: string): boolean {
+	if (
+		modelName.startsWith("openai/") ||
+		modelName.includes("text-embedding") ||
+		modelName.startsWith("text-embedding")
+	) {
+		return true;
+	}
+	const active = activeEmbeddingOptions();
+	const baseUrl = active?.apiUrl ?? ($env.PROMETHEUS_MEMORY_EMBEDDING_API_URL || $env.OPENROUTER_BASE_URL);
+	if (baseUrl !== undefined && baseUrl !== "" && !baseUrl.includes("openrouter.ai")) {
+		return true;
+	}
+	return $flag("PROMETHEUS_MEMORY_EMBEDDINGS_VIA_API");
+}
+
+const MODEL_DIMS: Record<string, number> = {
+	"BAAI/bge-small-en-v1.5": 384,
+	"BAAI/bge-base-en-v1.5": 768,
+	"BAAI/bge-large-en-v1.5": 1024,
+	"BAAI/bge-small-zh-v1.5": 512,
+	"BAAI/bge-base-zh-v1.5": 768,
+	"BAAI/bge-large-zh-v1.5": 1024,
+	"intfloat/multilingual-e5-small": 384,
+	"intfloat/multilingual-e5-base": 768,
+	"intfloat/multilingual-e5-large": 1024,
+	"BAAI/bge-m3": 1024,
+	"BAAI/bge-multilingual-gemma2": 3584,
+	"openai/text-embedding-3-small": 1536,
+	"openai/text-embedding-3-large": 3072,
+	"text-embedding-3-small": 1536,
+	"text-embedding-3-large": 3072,
+	"jina-embeddings-v5-omni-nano": 768,
+	"jina-embeddings-v5-omni-small": 1024,
+};
+export function embeddingDimFor(modelName: string): number {
+	const override = Number.parseInt($env.PROMETHEUS_MEMORY_EMBEDDING_DIM ?? "", 10);
+	if (Number.isFinite(override)) {
+		return override;
+	}
+	return MODEL_DIMS[modelName] ?? 384;
+}
+
+/** Drain an embedding stream (a custom provider or fastembed) into a `Float32Array` matrix. */
+async function collectMatrix(batches: EmbeddingOutput): Promise<EmbeddingMatrix> {
+	const rows: Vector[] = [];
+	for await (const batch of batches) {
+		for (const row of batch) {
+			rows.push(new Float32Array(row));
+		}
+	}
+	return rows;
+}
+
+const KNOWN_MODEL_NAMES: Record<string, string> = {
+	"BAAI/bge-small-en-v1.5": "fast-bge-small-en-v1.5",
+	"BAAI/bge-base-en-v1.5": "fast-bge-base-en-v1.5",
+	"BAAI/bge-small-en": "fast-bge-small-en",
+	"BAAI/bge-base-en": "fast-bge-base-en",
+	"BAAI/bge-small-zh-v1.5": "fast-bge-small-zh-v1.5",
+	"intfloat/multilingual-e5-large": "fast-multilingual-e5-large",
+	"sentence-transformers/all-MiniLM-L6-v2": "fast-all-MiniLM-L6-v2",
+};
+function fastembedModelName(modelName: string): StandardEmbeddingModel | null {
+	// Fastembed `EmbeddingModel` enum string values, inlined so resolving a model name
+	// (and `available()`) never imports `fastembed` — its module eagerly loads the
+	// `onnxruntime-node` native addon, which segfaults in some runtimes.
+	const id = KNOWN_MODEL_NAMES[modelName];
+	return id === undefined ? null : (id as StandardEmbeddingModel);
+}
+
+async function getLocalModel(): Promise<LocalEmbeddingModel | null> {
+	if (isApiModel(defaultModel()) || embeddingsDisabled() || inTestRuntime()) {
+		return null;
+	}
+	if (localModelPromise !== null) {
+		return localModelPromise;
+	}
+
+	const modelName = fastembedModelName(defaultModel());
+	if (modelName === null) {
+		return null;
+	}
+	const cacheDir = getFastembedCacheDir();
+	mkdirSync(cacheDir, { recursive: true });
+	const loading = localModelInitializer({
+		model: modelName,
+		cacheDir,
+		showDownloadProgress: false,
+	});
+	localModelPromise = loading;
+	try {
+		return await loading;
+	} catch {
+		if (localModelPromise === loading) localModelPromise = null;
+		return null;
+	}
+}
+
+async function embedApi(texts: readonly string[]): Promise<EmbeddingMatrix | null> {
+	const baseUrl = embeddingBaseUrl();
+	const isCustom = !baseUrl.includes("openrouter.ai");
+	const apiKey = embeddingApiKey();
+	if (!isCustom && apiKey === "") {
+		return null;
+	}
+
+	const headers: Record<string, string> = {
+		"Content-Type": "application/json",
+		"User-Agent": `${APP_DISPLAY_NAME}/${packageJson.version}`,
+		"HTTP-Referer": "https://prometheus.trivlab.com/",
+		"X-OpenRouter-Title": APP_DISPLAY_NAME,
+		"X-OpenRouter-Categories": "cli-agent",
+	};
+	if (apiKey !== "") {
+		headers.Authorization = `Bearer ${apiKey}`;
+	}
+
+	try {
+		const response = await fetchWithRetry(`${baseUrl.replace(/\/+$/, "")}/embeddings`, {
+			method: "POST",
+			headers,
+			body: JSON.stringify({ model: defaultModel(), input: texts }),
+			signal: AbortSignal.timeout(30000),
+			maxAttempts: 3,
+			defaultDelayMs: attempt => 2 ** attempt * 1000,
+		});
+		if (!response.ok) {
+			return null;
+		}
+		const { data: rows } = (await response.json()) as { data?: Array<{ embedding: number[] }> };
+		if (rows === undefined) {
+			return null;
+		}
+		apiCallCount += 1;
+		return rows.map(row => new Float32Array(row.embedding));
+	} catch (error) {
+		logger.debug("mnemopi embedding request failed", { status: extractHttpStatusFromError(error) });
+		return null;
+	}
+}
+
+async function providerAvailable(provider: EmbeddingProvider): Promise<boolean> {
+	if (provider.available === undefined) {
+		return true;
+	}
+	try {
+		return await provider.available();
+	} catch {
+		return false;
+	}
+}
+
+export function setEmbeddingProviderForTests(provider: EmbeddingProvider | null | undefined): void {
+	providerOverride = provider ?? null;
+	queryCache.clear();
+}
+
+export const setEmbeddingProvider = setEmbeddingProviderForTests;
+
+export function setLocalModelInitializerForTests(initializer: LocalModelInitializer | null | undefined): void {
+	localModelInitializer = initializer ?? defaultLocalModelInitializer;
+	localModelPromise = null;
+	queryCache.clear();
+}
+
+export function resetEmbeddingProviderForTests(): void {
+	providerOverride = null;
+	localModelPromise = null;
+	localModelInitializer = defaultLocalModelInitializer;
+	apiCallCount = 0;
+	queryCache.clear();
+}
+
+export const resetEmbeddingStateForTests = resetEmbeddingProviderForTests;
+
+export async function available(): Promise<boolean> {
+	if (embeddingsDisabled()) {
+		return false;
+	}
+	const active = activeEmbeddingOptions();
+	const activeProvider = resolveEmbeddingProvider(active?.provider);
+	if (activeProvider !== undefined) {
+		return providerAvailable(activeProvider);
+	}
+	if (providerOverride !== null) {
+		return providerAvailable(providerOverride);
+	}
+	if (isApiModel(defaultModel())) {
+		const baseUrl = active?.apiUrl ?? ($env.PROMETHEUS_MEMORY_EMBEDDING_API_URL || $env.OPENROUTER_BASE_URL);
+		if (baseUrl !== undefined && baseUrl !== "" && !baseUrl.includes("openrouter.ai")) {
+			return true;
+		}
+		return embeddingApiKey() !== "";
+	}
+	if (inTestRuntime()) {
+		return false;
+	}
+	return fastembedModelName(defaultModel()) !== null;
+}
+
+export function availableApi(): boolean {
+	return embeddingApiKey() !== "";
+}
+
+export async function embedQuery(text: string): Promise<Vector | null> {
+	if (text === "" || embeddingsDisabled()) {
+		return null;
+	}
+	const key = queryCacheKey(text);
+	const cached = queryCache.get(key);
+	if (cached !== undefined) {
+		return cached;
+	}
+	const vectors = await embed([text]);
+	const vector = vectors?.[0] ?? null;
+	if (vector !== null) {
+		queryCache.set(key, vector);
+	}
+	return vector;
+}
+
+export async function embed(texts: readonly string[]): Promise<EmbeddingMatrix | null> {
+	if (texts.length === 0 || embeddingsDisabled()) {
+		return null;
+	}
+	const activeProvider = resolveEmbeddingProvider(activeEmbeddingOptions()?.provider);
+	if (activeProvider !== undefined) {
+		try {
+			return await collectMatrix(await activeProvider.embed(texts));
+		} catch {
+			return null;
+		}
+	}
+	if (providerOverride !== null) {
+		try {
+			return await collectMatrix(await providerOverride.embed(texts));
+		} catch {
+			return null;
+		}
+	}
+	if (isApiModel(defaultModel())) {
+		return embedApi(texts);
+	}
+	if (texts.length === 1) {
+		const key = queryCacheKey(texts[0] ?? "");
+		const cached = queryCache.get(key);
+		if (cached !== undefined) {
+			return [cached];
+		}
+	}
+	const model = await getLocalModel();
+	if (model === null) {
+		return null;
+	}
+	try {
+		const vectors = await collectMatrix(model.embed([...texts]));
+		if (vectors.length === 1) {
+			const vector = vectors[0];
+			if (vector !== undefined) {
+				queryCache.set(queryCacheKey(texts[0] ?? ""), vector);
+			}
+		}
+		return vectors;
+	} catch {
+		return null;
+	}
+}
+
+export function getEmbeddingApiCallCountForTests(): number {
+	return apiCallCount;
+}
+
+export const DEFAULT_MODEL = defaultModel();
+export const EMBEDDING_DIM = embeddingDimFor(DEFAULT_MODEL);

package/src/core/entities.ts

@@ -0,0 +1,259 @@
+const ENTITY_EXTRACTION_STOP_WORD_VALUES = [
+	"the",
+	"a",
+	"an",
+	"and",
+	"or",
+	"but",
+	"in",
+	"on",
+	"at",
+	"to",
+	"for",
+	"of",
+	"with",
+	"by",
+	"from",
+	"as",
+	"is",
+	"was",
+	"are",
+	"were",
+	"be",
+	"been",
+	"being",
+	"have",
+	"has",
+	"had",
+	"do",
+	"does",
+	"did",
+	"will",
+	"would",
+	"could",
+	"should",
+	"may",
+	"might",
+	"can",
+	"shall",
+	"i",
+	"you",
+	"he",
+	"she",
+	"it",
+	"we",
+	"they",
+	"me",
+	"him",
+	"her",
+	"us",
+	"them",
+	"my",
+	"your",
+	"his",
+	"its",
+	"our",
+	"their",
+	"this",
+	"that",
+	"these",
+	"those",
+	"here",
+	"there",
+	"where",
+	"when",
+	"what",
+	"which",
+	"who",
+	"whom",
+	"whose",
+	"how",
+	"why",
+	"assistant",
+	"user",
+	"skill",
+	"review",
+	"target",
+	"class",
+	"level",
+	"signals",
+	"phase",
+	"api",
+	"pi",
+	"summary",
+	"added",
+	"active",
+	"not",
+	"whether",
+	"all",
+	"no",
+	"replying",
+	"ai",
+	"memory",
+	"conversation",
+	"fact",
+	"false",
+	"true",
+	"none",
+	"null",
+	"signal",
+	"hermes",
+	"agent",
+	"model",
+	"system",
+	"note",
+	"task",
+	"project",
+	"result",
+	"output",
+	"input",
+	"data",
+	"step",
+	"process",
+	"point",
+	"way",
+	"thing",
+	"time",
+	"work",
+] as const;
+
+export const ENTITY_EXTRACTION_STOP_WORDS: ReadonlySet<string> = new Set(ENTITY_EXTRACTION_STOP_WORD_VALUES);
+const ENTITY_PATTERNS: readonly RegExp[] = [
+	/@(\w{2,30})/g,
+	/#(\w{2,30})/g,
+	/"([^"]{2,50})"/g,
+	/'([^']{2,50})'/g,
+	/\b([A-Z][a-zA-Z]*(?:\s+[A-Z][a-zA-Z]*){1,4})\b/g,
+	/\b([A-Z][a-zA-Z]{1,20})\b/g,
+];
+
+function chars(value: string): string[] {
+	return Array.from(value);
+}
+
+export function levenshteinDistance(s1: string, s2: string): number {
+	let left = chars(s1);
+	let right = chars(s2);
+	if (left.length < right.length) {
+		const tmp = left;
+		left = right;
+		right = tmp;
+	}
+	if (right.length === 0) return left.length;
+
+	let previousRow = new Array<number>(right.length + 1);
+	let currentRow = new Array<number>(right.length + 1).fill(0);
+	for (let i = 0; i <= right.length; i++) previousRow[i] = i;
+
+	for (let i = 0; i < left.length; i++) {
+		currentRow[0] = i + 1;
+		const c1 = left[i];
+		for (let j = 0; j < right.length; j++) {
+			const insertions = (previousRow[j + 1] ?? 0) + 1;
+			const deletions = (currentRow[j] ?? 0) + 1;
+			const substitutions = (previousRow[j] ?? 0) + (c1 === right[j] ? 0 : 1);
+			currentRow[j + 1] = Math.min(insertions, deletions, substitutions);
+		}
+		const tmp = previousRow;
+		previousRow = currentRow;
+		currentRow = tmp;
+	}
+	return previousRow[right.length] ?? 0;
+}
+export function similarity(s1: string, s2: string): number {
+	const s1Lower = s1.toLowerCase().trim();
+	const s2Lower = s2.toLowerCase().trim();
+	if (s1Lower === s2Lower) return 1.0;
+
+	const maxLen = Math.max(chars(s1Lower).length, chars(s2Lower).length);
+	if (maxLen === 0) return 1.0;
+
+	if (s1Lower.startsWith(s2Lower) || s2Lower.startsWith(s1Lower)) {
+		const longer = Math.max(chars(s1Lower).length, chars(s2Lower).length);
+		const shorter = Math.min(chars(s1Lower).length, chars(s2Lower).length);
+		if (shorter / longer < 0.3) return 0.0;
+		return 0.7 + (shorter / longer) * 0.3;
+	}
+
+	if (s1Lower.includes(s2Lower) || s2Lower.includes(s1Lower)) {
+		const longer = Math.max(chars(s1Lower).length, chars(s2Lower).length);
+		const shorter = Math.min(chars(s1Lower).length, chars(s2Lower).length);
+		return 0.5 + (shorter / longer) * 0.3;
+	}
+
+	const dist = levenshteinDistance(s1Lower, s2Lower);
+	return 1.0 - dist / maxLen;
+}
+
+function isPureNumber(entity: string): boolean {
+	const normalized = entity.replaceAll(".", "").replaceAll(",", "");
+	return normalized.length > 0 && /^\d+$/.test(normalized);
+}
+
+export function extractEntitiesRegex(text: string): string[] {
+	if (typeof text !== "string" || text.length === 0) return [];
+
+	const entities = new Set<string>();
+	for (const sourcePattern of ENTITY_PATTERNS) {
+		const pattern = new RegExp(sourcePattern.source, sourcePattern.flags);
+		for (const match of text.matchAll(pattern)) {
+			const captured = match[1];
+			if (captured === undefined) continue;
+			const entity = captured.trim();
+			if (entity.length < 2) continue;
+
+			const words = entity.split(/\s+/).filter(word => word.length > 0);
+			if (words.length === 1 && ENTITY_EXTRACTION_STOP_WORDS.has(entity.toLowerCase())) continue;
+			if (words.some(word => ENTITY_EXTRACTION_STOP_WORDS.has(word.toLowerCase()))) continue;
+			if (isPureNumber(entity)) continue;
+
+			const first = entity[0];
+			if (words.length === 1 && first !== undefined && first >= "a" && first <= "z") {
+				const groupStart = match.index + match[0].indexOf(captured);
+				const prefixChar = groupStart > 0 ? text[groupStart - 1] : undefined;
+				if (prefixChar !== "@" && prefixChar !== "#") continue;
+			}
+
+			entities.add(entity);
+		}
+	}
+
+	const result = Array.from(entities).sort();
+	const filtered = new Set<string>();
+	for (const entity of result) {
+		let isSubstring = false;
+		for (const other of result) {
+			if (other === entity || !other.includes(entity)) continue;
+			if (entity.startsWith("@") || entity.startsWith("#")) continue;
+			if (other.startsWith("@") || other.startsWith("#")) continue;
+			isSubstring = true;
+			break;
+		}
+		if (!isSubstring) filtered.add(entity);
+	}
+	return Array.from(filtered).sort();
+}
+export type SimilarEntity = readonly [entity: string, score: number];
+
+export function findSimilarEntities(
+	entity: string,
+	knownEntities: readonly string[],
+	threshold = 0.8,
+): SimilarEntity[] {
+	const matches: SimilarEntity[] = [];
+	for (const known of knownEntities) {
+		if (known === entity) {
+			matches.push([known, 1.0]);
+			continue;
+		}
+		const score = similarity(entity, known);
+		if (score >= threshold) matches.push([known, score]);
+	}
+	matches.sort((left, right) => right[1] - left[1]);
+	return matches;
+}
+export function entityExtractionPerformance(text: string, iterations = 1000): number {
+	const start = performance.now();
+	for (let i = 0; i < iterations; i++) extractEntitiesRegex(text);
+	return (performance.now() - start) / iterations;
+}