From 36412100a418093e9d83478867aa95aaa1c8b1ca Mon Sep 17 00:00:00 2001
From: Basit Mustafa <basit.mustafa@gmail.com>
Date: Fri, 17 Apr 2026 07:53:10 -0700
Subject: [PATCH] fix: use config API key for embeddings

---
 src/core/embedding.ts         | 12 +++++--
 test/embedding-config.test.ts | 66 +++++++++++++++++++++++++++++++++++
 2 files changed, 76 insertions(+), 2 deletions(-)
 create mode 100644 test/embedding-config.test.ts

diff --git a/src/core/embedding.ts b/src/core/embedding.ts
index 4689ccd1..b138f8f8 100644
--- a/src/core/embedding.ts
+++ b/src/core/embedding.ts
@@ -8,6 +8,7 @@
  */
 
 import OpenAI from 'openai';
+import { loadConfig } from './config.ts';
 
 const MODEL = 'text-embedding-3-large';
 const DIMENSIONS = 1536;
@@ -18,10 +19,17 @@ const MAX_DELAY_MS = 120000;
 const BATCH_SIZE = 100;
 
 let client: OpenAI | null = null;
+let clientApiKey: string | undefined;
+
+export function resolveApiKey(): string | undefined {
+  return process.env.OPENAI_API_KEY || loadConfig()?.openai_api_key;
+}
 
 function getClient(): OpenAI {
-  if (!client) {
-    client = new OpenAI();
+  const apiKey = resolveApiKey();
+  if (!client || clientApiKey !== apiKey) {
+    client = apiKey ? new OpenAI({ apiKey }) : new OpenAI();
+    clientApiKey = apiKey;
   }
   return client;
 }
diff --git a/test/embedding-config.test.ts b/test/embedding-config.test.ts
new file mode 100644
index 00000000..bcc654f1
--- /dev/null
+++ b/test/embedding-config.test.ts
@@ -0,0 +1,66 @@
+import { describe, test, expect } from 'bun:test';
+import { mkdtempSync, mkdirSync, writeFileSync, rmSync } from 'fs';
+import { join } from 'path';
+import { tmpdir } from 'os';
+
+function runResolveApiKey(homeDir: string, env: Record<string, string | undefined> = {}) {
+  const proc = Bun.spawnSync({
+    cmd: [
+      process.execPath,
+      '--eval',
+      `import { resolveApiKey } from './src/core/embedding.ts'; console.log(resolveApiKey() ?? '');`,
+    ],
+    cwd: '/Users/basitmustafa/Documents/GitHub/upstream/gbrain',
+    env: {
+      ...process.env,
+      HOME: homeDir,
+      OPENAI_API_KEY: env.OPENAI_API_KEY,
+    },
+    stdout: 'pipe',
+    stderr: 'pipe',
+  });
+
+  return {
+    exitCode: proc.exitCode,
+    stdout: new TextDecoder().decode(proc.stdout).trim(),
+    stderr: new TextDecoder().decode(proc.stderr).trim(),
+  };
+}
+
+describe('embedding config', () => {
+  test('uses config openai_api_key when env var is absent', () => {
+    const homeDir = mkdtempSync(join(tmpdir(), 'gbrain-config-test-'));
+    try {
+      mkdirSync(join(homeDir, '.gbrain'), { recursive: true });
+      writeFileSync(
+        join(homeDir, '.gbrain', 'config.json'),
+        JSON.stringify({ engine: 'postgres', openai_api_key: 'config-key' }),
+      );
+
+      const result = runResolveApiKey(homeDir);
+      expect(result.exitCode).toBe(0);
+      expect(result.stderr).toBe('');
+      expect(result.stdout).toBe('config-key');
+    } finally {
+      rmSync(homeDir, { recursive: true, force: true });
+    }
+  });
+
+  test('prefers OPENAI_API_KEY env var over config openai_api_key', () => {
+    const homeDir = mkdtempSync(join(tmpdir(), 'gbrain-config-test-'));
+    try {
+      mkdirSync(join(homeDir, '.gbrain'), { recursive: true });
+      writeFileSync(
+        join(homeDir, '.gbrain', 'config.json'),
+        JSON.stringify({ engine: 'postgres', openai_api_key: 'config-key' }),
+      );
+
+      const result = runResolveApiKey(homeDir, { OPENAI_API_KEY: 'env-key' });
+      expect(result.exitCode).toBe(0);
+      expect(result.stderr).toBe('');
+      expect(result.stdout).toBe('env-key');
+    } finally {
+      rmSync(homeDir, { recursive: true, force: true });
+    }
+  });
+});