feat(embedding): pluggable EmbeddingProvider layer (OpenAI + Ollama)

src/cli.ts

@@ -31,6 +31,17 @@ async function main() {
 
   if (command === '--version' || command === 'version') {
     console.log(`gbrain ${VERSION}`);
+    // Surface the active embedding provider so users running in multiple shells
+    // notice when they're on a non-default brain (local Ollama vs OpenAI).
+    try {
+      const cfg = loadConfig();
+      if (cfg?.embedding) {
+        const { provider, model, dimensions } = cfg.embedding;
+        console.log(`embedding: ${provider} / ${model} (${dimensions}d)`);
+      }
+    } catch {
+      // Config not readable — fine, --version shouldn't fail on that
+    }
     return;
   }
 
@@ -379,6 +390,20 @@ async function connectEngine(): Promise<BrainEngine> {
     console.error('No brain configured. Run: gbrain init');
     process.exit(1);
   }
+
+  // Hydrate the embedding provider from the brain's persisted config so all
+  // commands (embed, import, query) use the provider the brain was initialized
+  // with — not whatever EMBEDDING_* env vars happen to be set.
+  if (config.embedding) {
+    const { createProvider, setProvider } = await import('./core/embedding/index.ts');
+    setProvider(createProvider({
+      provider: config.embedding.provider,
+      model: config.embedding.model,
+      dimensions: config.embedding.dimensions,
+      baseUrl: config.embedding.base_url,
+    }));
+  }
+
   const { createEngine } = await import('./core/engine-factory.ts');
   const engine = await createEngine(toEngineConfig(config));
   await engine.connect(toEngineConfig(config));

src/commands/init.ts

@@ -6,8 +6,57 @@ import { homedir } from 'os';
 
 const __filename = fileURLToPath(import.meta.url);
 const __dirname = dirname(__filename);
-import { saveConfig, type GBrainConfig } from '../core/config.ts';
+import { loadConfig, saveConfig, type GBrainConfig } from '../core/config.ts';
 import { createEngine } from '../core/engine-factory.ts';
+import { createProvider, resolveConfig as resolveEmbeddingConfig } from '../core/embedding/index.ts';
+import type { EmbeddingProvider, ProviderConfig } from '../core/embedding/index.ts';
+
+/**
+ * Parse --provider / --model / --dimensions / --base-url flags.
+ * Supports both `--flag value` (space) and `--flag=value` (equals) forms.
+ * Falls back to EMBEDDING_* env vars (handled inside resolveEmbeddingConfig).
+ */
+function parseEmbeddingFlags(args: string[]): Partial<ProviderConfig> {
+  const flag = (name: string): string | undefined => {
+    // `--flag=value` form
+    const prefix = name + '=';
+    for (const a of args) {
+      if (a.startsWith(prefix)) return a.slice(prefix.length);
+    }
+    // `--flag value` form
+    const i = args.indexOf(name);
+    return i !== -1 ? args[i + 1] : undefined;
+  };
+  const dims = flag('--dimensions');
+  return {
+    provider: flag('--provider'),
+    model: flag('--model'),
+    dimensions: dims ? parseInt(dims, 10) : undefined,
+    baseUrl: flag('--base-url'),
+  };
+}
+
+/**
+ * Resolve the embedding provider for this init, and guard against dim-mismatch
+ * with any existing brain config. Returns the instantiated provider.
+ */
+function resolveProviderWithGuard(args: string[]): { provider: EmbeddingProvider; resolved: ProviderConfig } {
+  const resolved = resolveEmbeddingConfig(parseEmbeddingFlags(args));
+  const provider = createProvider(resolved); // validates config + infers dims
+
+  const existing = loadConfig();
+  if (existing?.embedding && existing.embedding.dimensions !== provider.dimensions) {
+    console.error('');
+    console.error('Cannot re-init: existing brain has a different embedding dimension.');
+    console.error(`  Existing: ${existing.embedding.provider} / ${existing.embedding.model} (${existing.embedding.dimensions}d)`);
+    console.error(`  Requested: ${provider.name} / ${provider.model} (${provider.dimensions}d)`);
+    console.error('');
+    console.error('Switching providers requires regenerating all embeddings.');
+    console.error('To start fresh: delete ~/.gbrain/config.json and the brain data directory, then rerun gbrain init.');
+    process.exit(1);
+  }
+  return { provider, resolved };
+}
 
 export async function runInit(args: string[]) {
   const isSupabase = args.includes('--supabase');
@@ -21,6 +70,10 @@ export async function runInit(args: string[]) {
   const pathIndex = args.indexOf('--path');
   const customPath = pathIndex !== -1 ? args[pathIndex + 1] : null;
 
+  // Resolve embedding provider up front. Fails fast on bad provider/model/dims
+  // or dim-mismatch with an existing brain — before any engine state is created.
+  const { provider, resolved: providerResolved } = resolveProviderWithGuard(args);
+
   // Explicit PGLite mode
   if (isPGLite || (!isSupabase && !manualUrl && !isNonInteractive)) {
     // Smart detection: scan for .md files unless --pglite flag forces it
@@ -37,7 +90,7 @@ export async function runInit(args: string[]) {
       }
     }
 
-    return initPGLite({ jsonOutput, apiKey, customPath });
+    return initPGLite({ jsonOutput, apiKey, customPath, provider, providerResolved });
   }
 
   // Supabase/Postgres mode
@@ -56,20 +109,33 @@ export async function runInit(args: string[]) {
     databaseUrl = await supabaseWizard();
   }
 
-  return initPostgres({ databaseUrl, jsonOutput, apiKey });
+  return initPostgres({ databaseUrl, jsonOutput, apiKey, provider, providerResolved });
 }
 
-async function initPGLite(opts: { jsonOutput: boolean; apiKey: string | null; customPath: string | null }) {
+async function initPGLite(opts: {
+  jsonOutput: boolean;
+  apiKey: string | null;
+  customPath: string | null;
+  provider: EmbeddingProvider;
+  providerResolved: ProviderConfig;
+}) {
   const dbPath = opts.customPath || join(homedir(), '.gbrain', 'brain.pglite');
   console.log(`Setting up local brain with PGLite (no server needed)...`);
+  console.log(`Embedding: ${opts.provider.name} / ${opts.provider.model} (${opts.provider.dimensions}d)`);
 
   const engine = await createEngine({ engine: 'pglite' });
   await engine.connect({ database_path: dbPath, engine: 'pglite' });
-  await engine.initSchema();
+  await engine.initSchema({ dimensions: opts.provider.dimensions, defaultModel: opts.provider.model });
 
   const config: GBrainConfig = {
     engine: 'pglite',
     database_path: dbPath,
+    embedding: {
+      provider: opts.provider.name,
+      model: opts.provider.model,
+      dimensions: opts.provider.dimensions,
+      ...(opts.providerResolved.baseUrl ? { base_url: opts.providerResolved.baseUrl } : {}),
+    },
     ...(opts.apiKey ? { openai_api_key: opts.apiKey } : {}),
   };
   saveConfig(config);
@@ -78,7 +144,13 @@ async function initPGLite(opts: { jsonOutput: boolean; apiKey: string | null; cu
   await engine.disconnect();
 
   if (opts.jsonOutput) {
-    console.log(JSON.stringify({ status: 'success', engine: 'pglite', path: dbPath, pages: stats.page_count }));
+    console.log(JSON.stringify({
+      status: 'success',
+      engine: 'pglite',
+      path: dbPath,
+      pages: stats.page_count,
+      embedding: config.embedding,
+    }));
   } else {
     console.log(`\nBrain ready at ${dbPath}`);
     console.log(`${stats.page_count} pages. Engine: PGLite (local Postgres).`);
@@ -89,7 +161,13 @@ async function initPGLite(opts: { jsonOutput: boolean; apiKey: string | null; cu
   }
 }
 
-async function initPostgres(opts: { databaseUrl: string; jsonOutput: boolean; apiKey: string | null }) {
+async function initPostgres(opts: {
+  databaseUrl: string;
+  jsonOutput: boolean;
+  apiKey: string | null;
+  provider: EmbeddingProvider;
+  providerResolved: ProviderConfig;
+}) {
   const { databaseUrl } = opts;
 
   // Detect Supabase direct connection URLs and warn about IPv6
@@ -137,11 +215,18 @@ async function initPostgres(opts: { databaseUrl: string; jsonOutput: boolean; ap
   }
 
   console.log('Running schema migration...');
-  await engine.initSchema();
+  console.log(`Embedding: ${opts.provider.name} / ${opts.provider.model} (${opts.provider.dimensions}d)`);
+  await engine.initSchema({ dimensions: opts.provider.dimensions, defaultModel: opts.provider.model });
 
   const config: GBrainConfig = {
     engine: 'postgres',
     database_url: databaseUrl,
+    embedding: {
+      provider: opts.provider.name,
+      model: opts.provider.model,
+      dimensions: opts.provider.dimensions,
+      ...(opts.providerResolved.baseUrl ? { base_url: opts.providerResolved.baseUrl } : {}),
+    },
     ...(opts.apiKey ? { openai_api_key: opts.apiKey } : {}),
   };
   saveConfig(config);
@@ -151,7 +236,12 @@ async function initPostgres(opts: { databaseUrl: string; jsonOutput: boolean; ap
   await engine.disconnect();
 
   if (opts.jsonOutput) {
-    console.log(JSON.stringify({ status: 'success', engine: 'postgres', pages: stats.page_count }));
+    console.log(JSON.stringify({
+      status: 'success',
+      engine: 'postgres',
+      pages: stats.page_count,
+      embedding: config.embedding,
+    }));
   } else {
     console.log(`\nBrain ready. ${stats.page_count} pages. Engine: Postgres (Supabase).`);
     console.log('Next: gbrain import <dir>');

src/core/config.ts

@@ -13,6 +13,23 @@ export interface GBrainConfig {
   database_path?: string;
   openai_api_key?: string;
   anthropic_api_key?: string;
+  /**
+   * Embedding provider config, persisted at `gbrain init` and frozen for the
+   * brain's life. Presence indicates a provider was chosen explicitly; absence
+   * means legacy behavior (OpenAI text-embedding-3-large 1536d via env vars).
+   */
+  embedding?: EmbeddingConfig;
+}
+
+export interface EmbeddingConfig {
+  /** Provider name. Currently 'openai' or 'ollama'. */
+  provider: string;
+  /** Model identifier. */
+  model: string;
+  /** Output vector dimension — MUST match the pgvector schema column. */
+  dimensions: number;
+  /** Optional base URL override for OpenAI-compatible endpoints. */
+  base_url?: string;
 }
 
 /**

src/core/db.ts

@@ -1,6 +1,6 @@
 import postgres from 'postgres';
 import { GBrainError, type EngineConfig } from './types.ts';
-import { SCHEMA_SQL } from './schema-embedded.ts';
+import { postgresSchema } from './schema-embedded.ts';
 
 let sql: ReturnType<typeof postgres> | null = null;
 let connectedUrl: string | null = null;
@@ -68,12 +68,12 @@ export async function disconnect(): Promise<void> {
   }
 }
 
-export async function initSchema(): Promise<void> {
+export async function initSchema(opts?: { dimensions?: number; defaultModel?: string }): Promise<void> {
   const conn = getConnection();
   // Advisory lock prevents concurrent initSchema() calls from deadlocking
   await conn`SELECT pg_advisory_lock(42)`;
   try {
-    await conn.unsafe(SCHEMA_SQL);
+    await conn.unsafe(postgresSchema(opts));
   } finally {
     await conn`SELECT pg_advisory_unlock(42)`;
   }

src/core/embedding.ts

@@ -1,94 +1,17 @@
 /**
- * Embedding Service
- * Ported from production Ruby implementation (embedding_service.rb, 190 LOC)
+ * BACKWARD-COMPATIBILITY SHIM
  *
- * OpenAI text-embedding-3-large at 1536 dimensions.
- * Retry with exponential backoff (4s base, 120s cap, 5 retries).
- * 8000 character input truncation.
+ * The embedding implementation moved to `src/core/embedding/` as a provider layer
+ * (OpenAIProvider, OllamaProvider, factory, service). This file re-exports the
+ * public surface so existing imports keep working without churn:
+ *
+ *   import { embed, embedBatch } from '../core/embedding.ts';
+ *
+ * New code should import from `./embedding/index.ts` directly to access
+ * createProvider, EmbeddingProvider, OllamaProvider, etc.
+ *
+ * Test mocks (`mock.module('../src/core/embedding.ts', () => ({ embedBatch }))`)
+ * continue to intercept the call chain at this shim, so existing tests work unchanged.
  */
 
-import OpenAI from 'openai';
-
-const MODEL = 'text-embedding-3-large';
-const DIMENSIONS = 1536;
-const MAX_CHARS = 8000;
-const MAX_RETRIES = 5;
-const BASE_DELAY_MS = 4000;
-const MAX_DELAY_MS = 120000;
-const BATCH_SIZE = 100;
-
-let client: OpenAI | null = null;
-
-function getClient(): OpenAI {
-  if (!client) {
-    client = new OpenAI();
-  }
-  return client;
-}
-
-export async function embed(text: string): Promise<Float32Array> {
-  const truncated = text.slice(0, MAX_CHARS);
-  const result = await embedBatch([truncated]);
-  return result[0];
-}
-
-export async function embedBatch(texts: string[]): Promise<Float32Array[]> {
-  const truncated = texts.map(t => t.slice(0, MAX_CHARS));
-  const results: Float32Array[] = [];
-
-  // Process in batches of BATCH_SIZE
-  for (let i = 0; i < truncated.length; i += BATCH_SIZE) {
-    const batch = truncated.slice(i, i + BATCH_SIZE);
-    const batchResults = await embedBatchWithRetry(batch);
-    results.push(...batchResults);
-  }
-
-  return results;
-}
-
-async function embedBatchWithRetry(texts: string[]): Promise<Float32Array[]> {
-  for (let attempt = 0; attempt < MAX_RETRIES; attempt++) {
-    try {
-      const response = await getClient().embeddings.create({
-        model: MODEL,
-        input: texts,
-        dimensions: DIMENSIONS,
-      });
-
-      // Sort by index to maintain order
-      const sorted = response.data.sort((a, b) => a.index - b.index);
-      return sorted.map(d => new Float32Array(d.embedding));
-    } catch (e: unknown) {
-      if (attempt === MAX_RETRIES - 1) throw e;
-
-      // Check for rate limit with Retry-After header
-      let delay = exponentialDelay(attempt);
-
-      if (e instanceof OpenAI.APIError && e.status === 429) {
-        const retryAfter = e.headers?.['retry-after'];
-        if (retryAfter) {
-          const parsed = parseInt(retryAfter, 10);
-          if (!isNaN(parsed)) {
-            delay = parsed * 1000;
-          }
-        }
-      }
-
-      await sleep(delay);
-    }
-  }
-
-  // Should not reach here
-  throw new Error('Embedding failed after all retries');
-}
-
-function exponentialDelay(attempt: number): number {
-  const delay = BASE_DELAY_MS * Math.pow(2, attempt);
-  return Math.min(delay, MAX_DELAY_MS);
-}
-
-function sleep(ms: number): Promise<void> {
-  return new Promise(resolve => setTimeout(resolve, ms));
-}
-
-export { MODEL as EMBEDDING_MODEL, DIMENSIONS as EMBEDDING_DIMENSIONS };
+export { embed, embedBatch } from './embedding/service.ts';