diff --git a/README.md b/README.md index d1c36bb..ede27aa 100644 --- a/README.md +++ b/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 | diff --git a/src/db.ts b/src/db.ts index 5fe7ab4..3e72de9 100644 --- a/src/db.ts +++ b/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 any>(fn: T): T; loadExtension(path: string): void; close(): void; } diff --git a/src/llm.ts b/src/llm.ts index bd70e80..5022d27 100644 --- a/src/llm.ts +++ b/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 { 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)"); diff --git a/src/store.ts b/src/store.ts index 5289917..c9d1115 100644 --- a/src/store.ts +++ b/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++) { diff --git a/test/llm.test.ts b/test/llm.test.ts index 42bce0e..b1dd64f 100644 --- a/test/llm.test.ts +++ b/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);