Skip to content

Handle embedding fallback dimension drift #29

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
AVIDS2 merged 2 commits into from
Mar 28, 2026
Merged

Handle embedding fallback dimension drift #29

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
f4aed38
Handle embedding fallback dimension drift
RaviTharuma Mar 28, 2026
06cd7ad
Fix reindex embedding test expectations
RaviTharuma Mar 28, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
199 changes: 148 additions & 51 deletions src/embedding/provider.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -42,6 +42,9 @@ export interface EmbeddingProvider {
let provider: EmbeddingProvider | null = null;
let initPromise: Promise<EmbeddingProvider | null> | null = null;

type EmbeddingMode = 'off' | 'fastembed' | 'transformers' | 'api' | 'auto';
type ProviderKind = 'api' | 'fastembed' | 'transformers' | 'unknown';

/**
* Tracks whether the last init attempt resulted in a temporary failure
* (mode != 'off' but provider returned null). When true, the next
Expand All @@ -53,7 +56,7 @@ let lastInitWasTemporaryFailure = false;
* Get configured embedding mode from environment.
* Default is 'off' to minimize resource usage.
*/
function getEmbeddingMode(): 'off' | 'fastembed' | 'transformers' | 'api' | 'auto' {
function getEmbeddingMode(): EmbeddingMode {
// Unified: env vars > config.json > 'off'
try {
const { getEmbeddingMode: cfgMode } = require('../config.js');
Expand Down Expand Up @@ -89,6 +92,127 @@ function hasAPIEmbeddingConfig(): boolean {
}
}

function getProviderKind(candidate: EmbeddingProvider): ProviderKind {
if (candidate.name.startsWith('api-')) return 'api';
if (candidate.name.startsWith('fastembed-')) return 'fastembed';
if (candidate.name.startsWith('transformers-')) return 'transformers';
return 'unknown';
}

function isTemporaryEmbeddingFailure(error: unknown): boolean {
if (!(error instanceof Error)) return false;
return (
/embedding api timeout/i.test(error.message) ||
/embedding api error \((402|429|5\d\d)\)/i.test(error.message) ||
/quota exceeded/i.test(error.message) ||
/account balance/i.test(error.message) ||
/fetch failed/i.test(error.message) ||
/econnreset/i.test(error.message) ||
/econnrefused/i.test(error.message) ||
/temporarily unavailable/i.test(error.message)
);
}

function markTemporaryFailure(reason: unknown): void {
provider = null;
initPromise = null;
lastInitWasTemporaryFailure = true;
lastFailureTimestamp = Date.now();
const message = reason instanceof Error ? reason.message : String(reason);
console.error(`[memorix] Embedding provider temporarily unavailable at runtime: ${message}`);
}

async function createFastEmbedProvider(): Promise<EmbeddingProvider | null> {
try {
const { FastEmbedProvider } = await import('./fastembed-provider.js');
return await FastEmbedProvider.create();
} catch (e) {
console.error(`[memorix] Failed to load fastembed: ${e instanceof Error ? e.message : e}`);
console.error('[memorix] Install with: npm install fastembed');
return null;
}
}

async function createTransformersProvider(): Promise<EmbeddingProvider | null> {
try {
const { TransformersProvider } = await import('./transformers-provider.js');
return await TransformersProvider.create();
} catch (e) {
console.error(`[memorix] Failed to load transformers: ${e instanceof Error ? e.message : e}`);
console.error('[memorix] Install with: npm install @huggingface/transformers');
return null;
}
}

async function createAPIProvider(): Promise<EmbeddingProvider | null> {
try {
const { APIEmbeddingProvider } = await import('./api-provider.js');
return await APIEmbeddingProvider.create();
} catch (e) {
console.error(`[memorix] Failed to init API embedding: ${e instanceof Error ? e.message : e}`);
return null;
}
}

async function createLocalFallbackProvider(): Promise<EmbeddingProvider | null> {
const fastembed = await createFastEmbedProvider();
if (fastembed) return fastembed;

const transformers = await createTransformersProvider();
if (transformers) return transformers;

return null;
}

function wrapProvider(candidate: EmbeddingProvider): EmbeddingProvider {
const kind = getProviderKind(candidate);

return {
name: candidate.name,
dimensions: candidate.dimensions,
async embed(text: string): Promise<number[]> {
try {
return await candidate.embed(text);
} catch (error) {
if (kind === 'api' && getEmbeddingMode() === 'auto' && isTemporaryEmbeddingFailure(error)) {
console.error('[memorix] API embedding temporarily unavailable — switching to local fallback provider');
const fallback = await createLocalFallbackProvider();
if (fallback) {
provider = wrapProvider(fallback);
console.error(`[memorix] Embedding fallback activated: ${provider.name} (${provider.dimensions}d)`);
return provider.embed(text);
}
}

if (isTemporaryEmbeddingFailure(error)) {
markTemporaryFailure(error);
}
throw error;
}
},
async embedBatch(texts: string[]): Promise<number[][]> {
try {
return await candidate.embedBatch(texts);
} catch (error) {
if (kind === 'api' && getEmbeddingMode() === 'auto' && isTemporaryEmbeddingFailure(error)) {
console.error('[memorix] API embedding temporarily unavailable — switching to local fallback provider');
const fallback = await createLocalFallbackProvider();
if (fallback) {
provider = wrapProvider(fallback);
console.error(`[memorix] Embedding fallback activated: ${provider.name} (${provider.dimensions}d)`);
return provider.embedBatch(texts);
}
}

if (isTemporaryEmbeddingFailure(error)) {
markTemporaryFailure(error);
}
throw error;
}
},
};
}

/** Minimum interval between retry attempts after a temporary failure (ms). */
const RETRY_COOLDOWN_MS = 30_000;
let lastFailureTimestamp = 0;
Expand Down Expand Up @@ -133,73 +257,46 @@ export async function getEmbeddingProvider(): Promise<EmbeddingProvider | null>

// Explicit fastembed
if (mode === 'fastembed') {
try {
const { FastEmbedProvider } = await import('./fastembed-provider.js');
provider = await FastEmbedProvider.create();
console.error(`[memorix] Embedding provider: ${provider!.name} (${provider!.dimensions}d)`);
return provider;
} catch (e) {
console.error(`[memorix] Failed to load fastembed: ${e instanceof Error ? e.message : e}`);
console.error('[memorix] Install with: npm install fastembed');
return null;
}
const initialized = await createFastEmbedProvider();
if (!initialized) return null;
provider = wrapProvider(initialized);
console.error(`[memorix] Embedding provider: ${provider.name} (${provider.dimensions}d)`);
return provider;
}

// Explicit transformers
if (mode === 'transformers') {
try {
const { TransformersProvider } = await import('./transformers-provider.js');
provider = await TransformersProvider.create();
console.error(`[memorix] Embedding provider: ${provider!.name} (${provider!.dimensions}d)`);
return provider;
} catch (e) {
console.error(`[memorix] Failed to load transformers: ${e instanceof Error ? e.message : e}`);
console.error('[memorix] Install with: npm install @huggingface/transformers');
return null;
}
const initialized = await createTransformersProvider();
if (!initialized) return null;
provider = wrapProvider(initialized);
console.error(`[memorix] Embedding provider: ${provider.name} (${provider.dimensions}d)`);
return provider;
}

// API mode: remote embedding via OpenAI-compatible endpoint
if (mode === 'api') {
try {
const { APIEmbeddingProvider } = await import('./api-provider.js');
provider = await APIEmbeddingProvider.create();
console.error(`[memorix] Embedding provider: ${provider!.name} (${provider!.dimensions}d)`);
return provider;
} catch (e) {
console.error(`[memorix] Failed to init API embedding: ${e instanceof Error ? e.message : e}`);
return null;
}
const initialized = await createAPIProvider();
if (!initialized) return null;
provider = wrapProvider(initialized);
console.error(`[memorix] Embedding provider: ${provider.name} (${provider.dimensions}d)`);
return provider;
}

// Auto mode: try configured API first, then local fallbacks
if (hasAPIEmbeddingConfig()) {
try {
const { APIEmbeddingProvider } = await import('./api-provider.js');
provider = await APIEmbeddingProvider.create();
console.error(`[memorix] Embedding provider: ${provider!.name} (${provider!.dimensions}d)`);
const initialized = await createAPIProvider();
if (initialized) {
provider = wrapProvider(initialized);
console.error(`[memorix] Embedding provider: ${provider.name} (${provider.dimensions}d)`);
return provider;
} catch (e) {
console.error(`[memorix] API embedding unavailable in auto mode: ${e instanceof Error ? e.message : e}`);
}
}

try {
const { FastEmbedProvider } = await import('./fastembed-provider.js');
provider = await FastEmbedProvider.create();
console.error(`[memorix] Embedding provider: ${provider!.name} (${provider!.dimensions}d)`);
return provider;
} catch {
// fastembed not installed — try next
}

try {
const { TransformersProvider } = await import('./transformers-provider.js');
provider = await TransformersProvider.create();
console.error(`[memorix] Embedding provider: ${provider!.name} (${provider!.dimensions}d)`);
const localFallback = await createLocalFallbackProvider();
if (localFallback) {
provider = wrapProvider(localFallback);
console.error(`[memorix] Embedding provider: ${provider.name} (${provider.dimensions}d)`);
return provider;
} catch {
// transformers not installed — degrade to fulltext
}

console.error('[memorix] No embedding provider available — using BM25 fulltext search');
Expand Down
54 changes: 48 additions & 6 deletions src/memory/observations.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,13 +16,14 @@ import {
resetDb,
generateEmbedding,
batchGenerateEmbeddings,
getVectorDimensions,
makeOramaObservationId,
} from '../store/orama-store.js';
import { saveObservationsJson, loadObservationsJson, saveIdCounter, loadIdCounter } from '../store/persistence.js';
import { withFileLock } from '../store/file-lock.js';
import { countTextTokens } from '../compact/token-budget.js';
import { extractEntities, enrichConcepts } from './entity-extractor.js';
import { isEmbeddingExplicitlyDisabled } from '../embedding/provider.js';
import { getEmbeddingProvider, isEmbeddingExplicitlyDisabled } from '../embedding/provider.js';

/** In-memory observation list (loaded from persistence on init) */
let observations: Observation[] = [];
Expand All @@ -35,6 +36,12 @@ let projectDir: string | null = null;
const vectorMissingIds = new Set<number>();
let vectorBackfillRunning = false;

function isVectorCompatibleWithCurrentIndex(embedding: number[] | null): boolean {
if (!embedding) return false;
const vectorDimensions = getVectorDimensions();
return vectorDimensions === null || embedding.length === vectorDimensions;
}

/**
* Initialize the observations manager with a project directory.
*/
Expand Down Expand Up @@ -236,6 +243,13 @@ export async function storeObservation(input: {
const searchableText = [input.title, input.narrative, ...(input.facts ?? [])].join(' ');
generateEmbedding(searchableText).then(async (embedding) => {
if (embedding) {
if (!isVectorCompatibleWithCurrentIndex(embedding)) {
const vectorDimensions = getVectorDimensions();
console.error(
`[memorix] Embedding dimension mismatch for obs-${obsId}: provider returned ${embedding.length}d, index expects ${vectorDimensions ?? 'unknown'}d (kept in backfill queue)`,
);
return;
}
try {
const { removeObservation: removeObs } = await import('../store/orama-store.js');
await removeObs(makeOramaObservationId(input.projectId, obsId));
Expand Down Expand Up @@ -560,24 +574,41 @@ export async function reindexObservations(): Promise<number> {

// Reset the Orama index to ensure clean reindex (idempotent)
await resetDb();
vectorMissingIds.clear();

// Batch-generate all embeddings at once (much faster than individual calls)
let embeddings: (number[] | null)[] = [];
try {
let embeddings: (number[] | null)[] = observations.map(() => null);
const provider = await getEmbeddingProvider();
const canBatchEmbedAtStartup = provider !== null && !provider.name.startsWith('api-');

if (provider && !canBatchEmbedAtStartup) {
console.error('[memorix] Startup reindex: skipping synchronous API embeddings; background backfill will hydrate vectors');
}

if (canBatchEmbedAtStartup) {
try {
const texts = observations.map(obs =>
[obs.title, obs.narrative, ...obs.facts].join(' '),
);
embeddings = await batchGenerateEmbeddings(texts);
// Batch embedding failed — fall back to no embeddings
} catch {
// Batch embedding failed; fall back to no embeddings.
} catch {
// Batch embedding failed; fall back to no embeddings.
}
}

let count = 0;
for (let i = 0; i < observations.length; i++) {
const obs = observations[i];
try {
const embedding = embeddings[i] ?? null;
const compatibleEmbedding = isVectorCompatibleWithCurrentIndex(embedding) ? embedding : null;
if (embedding && !compatibleEmbedding) {
const vectorDimensions = getVectorDimensions();
console.error(
`[memorix] Startup reindex embedding mismatch for obs-${obs.id}: provider returned ${embedding.length}d, index expects ${vectorDimensions ?? 'unknown'}d (queued for backfill)`,
);
}
const docId = makeOramaObservationId(obs.projectId, obs.id);
const doc: MemorixDocument = {
id: docId,
Expand All @@ -596,9 +627,12 @@ export async function reindexObservations(): Promise<number> {
lastAccessedAt: '',
status: obs.status ?? 'active',
source: obs.source ?? 'agent',
...(embedding ? { embedding } : {}),
...(compatibleEmbedding ? { embedding: compatibleEmbedding } : {}),
};
await insertObservation(doc);
if (!compatibleEmbedding && !isEmbeddingExplicitlyDisabled()) {
vectorMissingIds.add(obs.id);
}
count++;
} catch (err) {
console.error(`[memorix] Failed to reindex observation #${obs.id}: ${err}`);
Expand Down Expand Up @@ -668,6 +702,14 @@ export async function backfillVectorEmbeddings(): Promise<{
try {
const embedding = await generateEmbedding(text);
if (embedding) {
if (!isVectorCompatibleWithCurrentIndex(embedding)) {
const vectorDimensions = getVectorDimensions();
console.error(
`[memorix] Backfill embedding mismatch for obs-${id}: provider returned ${embedding.length}d, index expects ${vectorDimensions ?? 'unknown'}d (kept in queue)`,
);
failed++;
continue;
}
const oramaId = makeOramaObservationId(obs.projectId, obs.id);
try {
const { removeObservation: removeObs } = await import('../store/orama-store.js');
Expand Down
Loading
Loading