feat: support NVIDIA embedding API

README.md

@@ -492,6 +492,16 @@ export QMD_EMBED_API_BASE_URL="https://api.openai.com/v1"
 export QMD_EMBED_MODEL="text-embedding-3-small"
 ```
 
+NVIDIA NIM's OpenAI-compatible endpoint can be used directly. QMD reads
+`NVIDIA_API_KEY` when `QMD_EMBED_API_KEY` is not set and sends NVIDIA's required
+`input_type` automatically (`passage` while indexing, `query` while searching):
+
+```sh
+export NVIDIA_API_KEY="..."
+export QMD_EMBED_API_BASE_URL="https://integrate.api.nvidia.com/v1"
+export QMD_EMBED_MODEL="nvidia/llama-3.2-nv-embedqa-1b-v2"
+```
+
 Reranking and query expansion still use local GGUF models via node-llama-cpp:
 
 | Model | Purpose | Size |

src/db.ts

@@ -68,6 +68,7 @@ export function openDatabase(path: string): Database {
 export interface Database {
   exec(sql: string): void;
   prepare(sql: string): Statement;
+  transaction<T extends (...args: any[]) => any>(fn: T): T;
   loadExtension(path: string): void;
   close(): void;
 }

src/llm.ts

@@ -513,7 +513,7 @@ export class LlamaCpp implements LLM {
   constructor(config: LlamaCppConfig = {}) {
     this.embedModelUri = config.embedModel || process.env.QMD_EMBED_MODEL || DEFAULT_EMBED_MODEL;
     this.embedApiBaseUrl = (config.embedApiBaseUrl || process.env.QMD_EMBED_API_BASE_URL || process.env.OPENAI_BASE_URL || DEFAULT_EMBED_API_BASE_URL).replace(/\/+$/, "");
-    this.embedApiKey = config.embedApiKey || process.env.QMD_EMBED_API_KEY || process.env.OPENAI_API_KEY;
+    this.embedApiKey = config.embedApiKey || process.env.QMD_EMBED_API_KEY || process.env.NVIDIA_API_KEY || process.env.OPENAI_API_KEY;
     this.generateModelUri = config.generateModel || process.env.QMD_GENERATE_MODEL || DEFAULT_GENERATE_MODEL;
     this.rerankModelUri = config.rerankModel || process.env.QMD_RERANK_MODEL || DEFAULT_RERANK_MODEL;
     this.modelCacheDir = config.modelCacheDir || MODEL_CACHE_DIR;
@@ -991,7 +991,11 @@ export class LlamaCpp implements LLM {
     return { text: truncatedText, truncated: true, limit: maxTokens };
   }
 
-  private async embedExternal(texts: string[], model: string): Promise<(EmbeddingResult | null)[]> {
+  private isNvidiaEmbedApi(): boolean {
+    return /(^|\.)nvidia\.com$/i.test(new URL(this.embedApiBaseUrl).hostname);
+  }
+
+  private async embedExternal(texts: string[], model: string, options: EmbedOptions = {}): Promise<(EmbeddingResult | null)[]> {
     if (texts.length === 0) return [];
     if (!this.embedApiKey) {
       throw new Error(
@@ -1006,7 +1010,11 @@ export class LlamaCpp implements LLM {
         "Authorization": `Bearer ${this.embedApiKey}`,
         "Content-Type": "application/json",
       },
-      body: JSON.stringify({ model, input: texts }),
+      body: JSON.stringify({
+        model,
+        input: texts,
+        ...(this.isNvidiaEmbedApi() ? { input_type: options.isQuery ? "query" : "passage" } : {}),
+      }),
     });
 
     if (!response.ok) {
@@ -1036,7 +1044,7 @@ export class LlamaCpp implements LLM {
   async embed(text: string, options: EmbedOptions = {}): Promise<EmbeddingResult | null> {
     const model = options.model ?? this.embedModelUri;
     if (!isLocalEmbeddingModel(model)) {
-      const results = await this.embedExternal([text], model);
+      const results = await this.embedExternal([text], model, options);
       return results[0] ?? null;
     }
 
@@ -1071,7 +1079,7 @@ export class LlamaCpp implements LLM {
   async embedBatch(texts: string[], options: EmbedOptions = {}): Promise<(EmbeddingResult | null)[]> {
     const model = options.model ?? this.embedModelUri;
     if (!isLocalEmbeddingModel(model)) {
-      return this.embedExternal(texts, model);
+      return this.embedExternal(texts, model, options);
     }
 
     if (this._ciMode) throw new Error("LLM operations are disabled in CI (set CI=true)");

src/store.ts

@@ -4092,7 +4092,7 @@ export async function hybridQuery(
     const textsToEmbed = vecQueries.map(q => formatQueryForEmbedding(q.text, llm.embedModelName));
     hooks?.onEmbedStart?.(textsToEmbed.length);
     const embedStart = Date.now();
-    const embeddings = await llm.embedBatch(textsToEmbed);
+    const embeddings = await llm.embedBatch(textsToEmbed, { isQuery: true });
     hooks?.onEmbedDone?.(Date.now() - embedStart);
 
     // Run sqlite-vec lookups with pre-computed embeddings
@@ -4475,7 +4475,7 @@ export async function structuredSearch(
       const textsToEmbed = vecSearches.map(s => formatQueryForEmbedding(s.query, llm.embedModelName));
       hooks?.onEmbedStart?.(textsToEmbed.length);
       const embedStart = Date.now();
-      const embeddings = await llm.embedBatch(textsToEmbed);
+      const embeddings = await llm.embedBatch(textsToEmbed, { isQuery: true });
       hooks?.onEmbedDone?.(Date.now() - embedStart);
 
       for (let i = 0; i < vecSearches.length; i++) {

test/llm.test.ts

@@ -220,6 +220,41 @@ describe("LlamaCpp model resolution (config > env > default)", () => {
     }
   });
 
+  test("NVIDIA embedding API uses NVIDIA_API_KEY and input_type", async () => {
+    const prevEmbedKey = process.env.QMD_EMBED_API_KEY;
+    const prevNvidiaKey = process.env.NVIDIA_API_KEY;
+    const prevBaseUrl = process.env.QMD_EMBED_API_BASE_URL;
+    delete process.env.QMD_EMBED_API_KEY;
+    process.env.NVIDIA_API_KEY = "nvidia-test-key";
+    process.env.QMD_EMBED_API_BASE_URL = "https://integrate.api.nvidia.com/v1";
+    const fetchMock = vi.spyOn(globalThis, "fetch").mockResolvedValue({
+      ok: true,
+      json: async () => ({
+        model: "nvidia/llama-3.2-nv-embedqa-1b-v2",
+        data: [{ index: 0, embedding: [0.1, 0.2, 0.3] }],
+      }),
+    } as Response);
+
+    try {
+      const llm = new LlamaCpp({ embedModel: "nvidia/llama-3.2-nv-embedqa-1b-v2" });
+      await llm.embed("hello", { isQuery: true });
+      const [, init] = fetchMock.mock.calls[0]!;
+      expect(JSON.parse((init as RequestInit).body as string)).toEqual({
+        model: "nvidia/llama-3.2-nv-embedqa-1b-v2",
+        input: ["hello"],
+        input_type: "query",
+      });
+    } finally {
+      fetchMock.mockRestore();
+      if (prevEmbedKey === undefined) delete process.env.QMD_EMBED_API_KEY;
+      else process.env.QMD_EMBED_API_KEY = prevEmbedKey;
+      if (prevNvidiaKey === undefined) delete process.env.NVIDIA_API_KEY;
+      else process.env.NVIDIA_API_KEY = prevNvidiaKey;
+      if (prevBaseUrl === undefined) delete process.env.QMD_EMBED_API_BASE_URL;
+      else process.env.QMD_EMBED_API_BASE_URL = prevBaseUrl;
+    }
+  });
+
   test("hf embedding model opts into local embedding", () => {
     const llm = new LlamaCpp({ embedModel: "hf:custom/embed.gguf" });
     expect(llm.usesLocalEmbedding).toBe(true);