From 4121977bb41d064676d9d7ebd204c60d0dae9b44 Mon Sep 17 00:00:00 2001 From: Kaspre Date: Tue, 2 Jun 2026 11:20:39 -0400 Subject: [PATCH 1/6] feat(llm): OpenAI-compatible remote embedding, reranking & query expansion Add a RemoteLLM backend that talks to any OpenAI-compatible HTTP API (vLLM, TEI, Ollama, llama.cpp --server, LiteLLM, OpenAI, ...), composed with the local LlamaCpp via a HybridLLM that routes each operation independently: - RemoteLLM: /v1/embeddings, /v1/rerank, /v1/chat/completions; per-endpoint circuit breakers; bearer auth; char-based token approximation so chunking works without a local tokenizer. - HybridLLM: embed/embedBatch/rerank/expandQuery -> remote, generate/tokenize/detokenize -> local, with per-operation local fallback. - Widened LLM interface (embedBatch, tokenize/detokenize/countTokens, isRemote, embedModelName, rerankModelName, generateModelName, usesRemoteEmbedding, supportsRerank/supportsExpand) plus getDefaultLLM/setDefaultLLM alongside the existing getDefaultLlamaCpp/setDefaultLlamaCpp (no breaking change). - Sigmoid normalization of log-odds rerank scores; RemoteLLM.expandQuery via chat completions; startup pre-flight embed probe; HybridLLM.rerank local fallback (symmetric with the expandQuery fallback). Opt-in via models.*_api_url / QMD_*_API_* env vars; with nothing configured the local-only path is byte-for-byte unchanged. Builds on @georgelichen's remote-LLM work in #629. Co-Authored-By: Claude Opus 4.8 --- CHANGELOG.md | 10 +- README.md | 35 ++ src/cli/qmd.ts | 58 ++- src/collections.ts | 12 + src/hybrid-llm.ts | 96 +++++ src/llm.ts | 126 ++++-- src/remote-llm.ts | 546 ++++++++++++++++++++++++++ src/store.ts | 83 +++- test/remote-llm-integration.test.ts | 412 ++++++++++++++++++++ test/remote-llm.test.ts | 580 ++++++++++++++++++++++++++++ 10 files changed, 1900 insertions(+), 58 deletions(-) create mode 100644 src/hybrid-llm.ts create mode 100644 src/remote-llm.ts create mode 100644 test/remote-llm-integration.test.ts create mode 100644 test/remote-llm.test.ts diff --git a/CHANGELOG.md b/CHANGELOG.md index b367f4080..d0fb7c470 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,15 @@ ## [Unreleased] +### Added + +- Remote embedding, reranking, and query expansion via OpenAI-compatible API + (vLLM, Ollama, OpenAI, etc.). Set `QMD_EMBED_API_URL` / `QMD_EMBED_API_MODEL` + (and optionally `QMD_RERANK_API_*` / `QMD_EXPAND_API_*`) env vars or add + the equivalent keys to `models:` in `index.yml`. Local generation and + tokenization are preserved via a hybrid routing layer. Includes circuit + breakers, dimension validation, and batch splitting. + ### Fixed - Filesystem paths with special characters (`#`, `&`, spaces, `[]`, `()`, etc.) @@ -108,7 +117,6 @@ - Launcher: Rewrite `bin/qmd` as a Node-based shebang polyglot to fix global npm installation execution failures on Windows (#668 / #452), while supporting seamless fallback to Bun in Node-less environments. - ## [2.5.1] - 2026-05-20 ### Changes diff --git a/README.md b/README.md index 7eadb9347..3a7861553 100644 --- a/README.md +++ b/README.md @@ -942,6 +942,41 @@ Uses node-llama-cpp's `createRankingContext()` and `rankAndSort()` API for cross Used for generating query variations via `LlamaChatSession`. +### Remote Embedding & Reranking + +QMD can offload embedding and reranking to a remote OpenAI-compatible server (vLLM, Ollama, LM Studio, OpenAI, etc.) while keeping query expansion local. + +**Environment variables** (presence of `QMD_EMBED_API_URL` activates remote mode): + +| Variable | Required | Description | +|----------|----------|-------------| +| `QMD_EMBED_API_URL` | Yes | Base URL, e.g. `http://gpu-host:8000/v1` | +| `QMD_EMBED_API_MODEL` | Yes | Model name, e.g. `BAAI/bge-m3` | +| `QMD_EMBED_API_KEY` | No | Bearer token for auth | +| `QMD_RERANK_API_URL` | No | Rerank endpoint (defaults to embed URL) | +| `QMD_RERANK_API_MODEL` | No | Rerank model name | +| `QMD_RERANK_API_KEY` | No | Rerank auth (defaults to embed key) | + +**YAML config** (`~/.config/qmd/index.yml`): +```yaml +models: + embed_api_url: "http://gpu-host:8000/v1" + embed_api_model: "BAAI/bge-m3" + rerank_api_model: "BAAI/bge-reranker-v2-m3" +``` + +**Example with vLLM:** +```sh +# Start vLLM with an embedding model +vllm serve BAAI/bge-m3 --task embed + +# Point QMD at it +export QMD_EMBED_API_URL=http://localhost:8000/v1 +export QMD_EMBED_API_MODEL=BAAI/bge-m3 +qmd embed +qmd query "your search query" +``` + ## License MIT diff --git a/src/cli/qmd.ts b/src/cli/qmd.ts index 105506d37..ad159748b 100755 --- a/src/cli/qmd.ts +++ b/src/cli/qmd.ts @@ -81,7 +81,9 @@ import { type ReindexResult, type ChunkStrategy, } from "../store.js"; -import { disposeDefaultLlamaCpp, getDefaultLlamaCpp, setDefaultLlamaCpp, LlamaCpp, withLLMSession, pullModels, DEFAULT_MODEL_CACHE_DIR, resolveEmbedModel, resolveGenerateModel, resolveRerankModel, resolveModels, inspectGgufFile, isDarwinMetalMitigationActive } from "../llm.js"; +import { disposeDefaultLlamaCpp, getDefaultLLM, setDefaultLLM, getDefaultLlamaCpp, setDefaultLlamaCpp, LlamaCpp, withLLMSession, pullModels, DEFAULT_MODEL_CACHE_DIR, resolveEmbedModel, resolveGenerateModel, resolveRerankModel, resolveModels, inspectGgufFile, isDarwinMetalMitigationActive } from "../llm.js"; +import { RemoteLLM, remoteConfigFromEnv } from "../remote-llm.js"; +import { HybridLLM } from "../hybrid-llm.js"; import { formatSearchResults, formatDocuments, @@ -134,11 +136,18 @@ function getStore(): ReturnType { const activeModels = ensureModelsConfiguredForCli(); const config = loadConfig(); syncConfigToDb(store.db, config); - setDefaultLlamaCpp(new LlamaCpp({ + const localLlm = new LlamaCpp({ embedModel: activeModels.embed, generateModel: activeModels.generate, rerankModel: activeModels.rerank, - })); + }); + // Remote embedding/rerank: env vars (QMD_EMBED_API_URL etc) take precedence over YAML models.*_api_* + const remoteConfig = remoteConfigFromEnv(config.models); + if (remoteConfig) { + setDefaultLLM(new HybridLLM(new RemoteLLM(remoteConfig), localLlm)); + } else { + setDefaultLLM(localLlm); + } } catch { // Config may not exist yet — that's fine, DB works without it } @@ -1984,6 +1993,39 @@ async function vectorIndex( const storeInstance = getStore(); const db = storeInstance.db; + // PR #629 follow-up: when HybridLLM/RemoteLLM is configured, the actual embedding + // model name comes from the LLM (e.g. "nomic-embed-text") rather than the local + // GGUF URI we got from CLI defaults. Override here so generateEmbeddings, fingerprint + // computation, and the pending-doc lookup all use the remote model identifier. + const llm = getDefaultLLM(); + model = llm.embedModelName; + + // Pre-flight probe: when remote embedding is configured, verify the remote + // endpoint actually works BEFORE writing any vectors. This prevents the + // silent-fallback failure mode where a misconfigured/unreachable remote + // backend would otherwise let qmd appear to "succeed" while writing + // local-model-tagged vectors (or no vectors at all). Belt-and-suspenders + // on top of RemoteLLM's call-time errors — fail fast at startup with a + // clear message rather than mid-batch. + if (llm.usesRemoteEmbedding) { + try { + const probe = await llm.embed("preflight probe", { model }); + if (!probe || !Array.isArray(probe.embedding) || probe.embedding.length === 0) { + throw new Error( + `Pre-flight probe to remote embedder returned no embedding ` + + `(model=${model}). Aborting embed run to avoid silent fallback. ` + + `Check llama-server / remote endpoint health.` + ); + } + } catch (err) { + const msg = err instanceof Error ? err.message : String(err); + throw new Error( + `Pre-flight probe to remote embedder FAILED (model=${model}): ${msg}. ` + + `Aborting embed run to avoid silent fallback.` + ); + } + } + if (force) { console.log(`${c.yellow}Force re-indexing: clearing all vectors...${c.reset}`); } @@ -3868,7 +3910,15 @@ async function runDoctorDeviceChecks(nextSteps: string[]): Promise { } try { - const device = await getDefaultLlamaCpp().getDeviceInfo({ allowBuild: false }); + const llm = getDefaultLLM(); + if (!(llm instanceof LlamaCpp)) { + if (process.stdout.isTTY) { + process.stdout.write(`\r${" ".repeat(crashHint.length)}\r`); + } + console.log(` ${c.dim}device probe unavailable for non-local LLM backend${c.reset}`); + return; + } + const device = await llm.getDeviceInfo({ allowBuild: false }); if (process.stdout.isTTY) { process.stdout.write(`\r${" ".repeat(crashHint.length)}\r`); } diff --git a/src/collections.ts b/src/collections.ts index 6950493d1..768a7eeb7 100644 --- a/src/collections.ts +++ b/src/collections.ts @@ -40,6 +40,18 @@ export interface ModelsConfig { embed?: string; rerank?: string; generate?: string; + /** Remote embedding API base URL (e.g. http://gpu-host:8000/v1) */ + embed_api_url?: string; + /** Remote embedding model name (e.g. BAAI/bge-m3) */ + embed_api_model?: string; + /** Bearer token for remote embedding API */ + embed_api_key?: string; + /** Remote rerank API base URL (defaults to embed_api_url) */ + rerank_api_url?: string; + /** Remote rerank model name */ + rerank_api_model?: string; + /** Bearer token for remote rerank API */ + rerank_api_key?: string; } /** diff --git a/src/hybrid-llm.ts b/src/hybrid-llm.ts new file mode 100644 index 000000000..be37bf656 --- /dev/null +++ b/src/hybrid-llm.ts @@ -0,0 +1,96 @@ +/** + * hybrid-llm.ts - Compositor that routes LLM operations between remote and local backends + * + * Embed/rerank → remote (GPU-heavy, benefits from offloading) + * Generate/expandQuery → local LlamaCpp (QMD's fine-tuned query expansion model) + * tokenize/countTokens → local LlamaCpp (CPU-cheap, needed for chunking) + */ + +import type { + LLM, + EmbedOptions, + EmbeddingResult, + GenerateOptions, + GenerateResult, + ModelInfo, + Queryable, + RerankDocument, + RerankOptions, + RerankResult, +} from "./llm.js"; +import type { Token as LlamaToken } from "node-llama-cpp"; +import { RemoteLLM } from "./remote-llm.js"; + +export class HybridLLM implements LLM { + constructor( + private readonly remote: LLM, + private readonly local: LLM, + ) {} + + get embedModelName(): string { + return this.remote.embedModelName; + } + + get generateModelName(): string { + return this.local.generateModelName; + } + + get rerankModelName(): string { + if (this.remote instanceof RemoteLLM && !this.remote.supportsRerank) { + return this.local.rerankModelName; + } + return this.remote.rerankModelName; + } + + get usesRemoteEmbedding(): boolean { + return this.remote.usesRemoteEmbedding === true; + } + + // Route to remote + embed(text: string, options?: EmbedOptions): Promise { + return this.remote.embed(text, options); + } + + embedBatch(texts: string[], options?: EmbedOptions): Promise<(EmbeddingResult | null)[]> { + return this.remote.embedBatch(texts, options); + } + + rerank(query: string, documents: RerankDocument[], options?: RerankOptions): Promise { + // When remote is a RemoteLLM without a rerank model configured, fall back to local rerank + // (same fallback shape as expandQuery → local). + if (this.remote instanceof RemoteLLM && !this.remote.supportsRerank) { + return this.local.rerank(query, documents, options); + } + return this.remote.rerank(query, documents, options); + } + + // Route to local + generate(prompt: string, options?: GenerateOptions): Promise { + return this.local.generate(prompt, options); + } + + tokenize(text: string): Promise { + return this.local.tokenize(text); + } + + detokenize(tokens: readonly LlamaToken[]): Promise { + return this.local.detokenize(tokens); + } + + expandQuery(query: string, options?: { context?: string; includeLexical?: boolean; intent?: string }): Promise { + // Route to remote when configured for it; otherwise local (same fallback + // shape as rerank → local when remote doesn't support rerank). + if (this.remote instanceof RemoteLLM && this.remote.supportsExpand) { + return this.remote.expandQuery(query, options); + } + return this.local.expandQuery(query, options); + } + + modelExists(model: string): Promise { + return this.local.modelExists(model); + } + + async dispose(): Promise { + await Promise.all([this.remote.dispose(), this.local.dispose()]); + } +} diff --git a/src/llm.ts b/src/llm.ts index 7c2d464cb..597d0e976 100644 --- a/src/llm.ts +++ b/src/llm.ts @@ -86,13 +86,24 @@ export function isQwen3EmbeddingModel(modelUri: string): boolean { return /qwen.*embed/i.test(modelUri) || /embed.*qwen/i.test(modelUri); } +/** + * Detect if a model URI refers to a remote API model (not a local GGUF model). + * Remote models handle their own prompt formatting, so no prefixes should be added. + */ +export function isRemoteModel(modelUri: string): boolean { + // Local models use hf: URIs or local file paths ending in .gguf + return !modelUri.startsWith("hf:") && !modelUri.endsWith(".gguf"); +} + /** * Format a query for embedding. * Uses nomic-style task prefix format for embeddinggemma (default). * Uses Qwen3-Embedding instruct format when a Qwen embedding model is active. + * Remote models receive raw text (they handle their own formatting). */ export function formatQueryForEmbedding(query: string, modelUri?: string): string { const uri = modelUri ?? resolveEmbedModel(); + if (isRemoteModel(uri)) return query; if (isQwen3EmbeddingModel(uri)) { return `Instruct: Retrieve relevant documents for the given query\nQuery: ${query}`; } @@ -103,9 +114,11 @@ export function formatQueryForEmbedding(query: string, modelUri?: string): strin * Format a document for embedding. * Uses nomic-style format with title and text fields (default). * Qwen3-Embedding encodes documents as raw text without special prefixes. + * Remote models receive raw text (they handle their own formatting). */ export function formatDocForEmbedding(text: string, title?: string, modelUri?: string): string { const uri = modelUri ?? resolveEmbedModel(); + if (isRemoteModel(uri)) return title ? `${title}\n${text}` : text; if (isQwen3EmbeddingModel(uri)) { // Qwen3-Embedding: documents are raw text, no task prefix return title ? `${title}\n${text}` : text; @@ -524,6 +537,43 @@ export interface LLM { */ embed(text: string, options?: EmbedOptions): Promise; + /** + * Batch embed multiple texts + */ + embedBatch(texts: string[], options?: EmbedOptions): Promise<(EmbeddingResult | null)[]>; + + /** + * The embedding model name/URI + */ + readonly embedModelName: string; + + /** + * The generation model name/URI (used by query expansion). + */ + readonly generateModelName: string; + + /** + * The reranking model name/URI. + */ + readonly rerankModelName: string; + + /** + * True when embedding requests are served by a remote backend and callers may + * want to avoid local tokenizer/model initialization for preprocessing. + */ + readonly usesRemoteEmbedding?: boolean; + + /** + * Tokenize text using the embedding model's tokenizer (used for chunking). + * For remote-only backends this may throw or be unavailable. + */ + tokenize(text: string): Promise; + + /** + * Detokenize tokens back to text using the embedding model's tokenizer. + */ + detokenize(tokens: readonly LlamaToken[]): Promise; + /** * Generate text completion */ @@ -538,7 +588,7 @@ export interface LLM { * Expand a search query into multiple variations for different backends. * Returns a list of Queryable objects. */ - expandQuery(query: string, options?: { context?: string, includeLexical?: boolean }): Promise; + expandQuery(query: string, options?: { context?: string; includeLexical?: boolean; intent?: string }): Promise; /** * Rerank documents by relevance to a query @@ -1714,11 +1764,11 @@ export class LlamaCpp implements LLM { * Coordinates with LlamaCpp idle timeout to prevent disposal during active sessions. */ class LLMSessionManager { - private llm: LlamaCpp; + private llm: LLM; private _activeSessionCount = 0; private _inFlightOperations = 0; - constructor(llm: LlamaCpp) { + constructor(llm: LLM) { this.llm = llm; } @@ -1754,7 +1804,7 @@ class LLMSessionManager { this._inFlightOperations = Math.max(0, this._inFlightOperations - 1); } - getLlamaCpp(): LlamaCpp { + getLLM(): LLM { return this.llm; } } @@ -1857,18 +1907,18 @@ class LLMSession implements ILLMSession { } async embed(text: string, options?: EmbedOptions): Promise { - return this.withOperation(() => this.manager.getLlamaCpp().embed(text, options)); + return this.withOperation(() => this.manager.getLLM().embed(text, options)); } async embedBatch(texts: string[], options?: EmbedOptions): Promise<(EmbeddingResult | null)[]> { - return this.withOperation(() => this.manager.getLlamaCpp().embedBatch(texts, options)); + return this.withOperation(() => this.manager.getLLM().embedBatch(texts, options)); } async expandQuery( query: string, options?: { context?: string; includeLexical?: boolean } ): Promise { - return this.withOperation(() => this.manager.getLlamaCpp().expandQuery(query, options)); + return this.withOperation(() => this.manager.getLLM().expandQuery(query, options)); } async rerank( @@ -1876,19 +1926,19 @@ class LLMSession implements ILLMSession { documents: RerankDocument[], options?: RerankOptions ): Promise { - return this.withOperation(() => this.manager.getLlamaCpp().rerank(query, documents, options)); + return this.withOperation(() => this.manager.getLLM().rerank(query, documents, options)); } } -// Session manager for the default LlamaCpp instance +// Session manager for the default LLM instance let defaultSessionManager: LLMSessionManager | null = null; /** - * Get the session manager for the default LlamaCpp instance. + * Get the session manager for the default LLM instance. */ function getSessionManager(): LLMSessionManager { - const llm = getDefaultLlamaCpp(); - if (!defaultSessionManager || defaultSessionManager.getLlamaCpp() !== llm) { + const llm = getDefaultLLM(); + if (!defaultSessionManager || defaultSessionManager.getLLM() !== llm) { defaultSessionManager = new LLMSessionManager(llm); } return defaultSessionManager; @@ -1923,11 +1973,11 @@ export async function withLLMSession( } /** - * Execute a function with a scoped LLM session using a specific LlamaCpp instance. + * Execute a function with a scoped LLM session using a specific LLM instance. * Unlike withLLMSession, this does not use the global singleton. */ export async function withLLMSessionForLlm( - llm: LlamaCpp, + llm: LLM, fn: (session: ILLMSession) => Promise, options?: LLMSessionOptions ): Promise { @@ -2011,49 +2061,61 @@ export function isDarwinExitGuardInstalled(): boolean { } // ============================================================================= -// Singleton for default LlamaCpp instance +// Singleton for default LLM instance // ============================================================================= -let defaultLlamaCpp: LlamaCpp | null = null; +let defaultLLMInstance: LLM | null = null; /** - * Get the default LlamaCpp instance (creates one if needed). The LlamaCpp - * constructor installs the darwin exit guard, so any code path that obtains - * the singleton is protected. + * Get the default LLM instance (creates a LlamaCpp if none set). */ -export function getDefaultLlamaCpp(): LlamaCpp { - if (!defaultLlamaCpp) { - defaultLlamaCpp = new LlamaCpp(); +export function getDefaultLLM(): LLM { + if (!defaultLLMInstance) { + defaultLLMInstance = new LlamaCpp(); } - return defaultLlamaCpp; + return defaultLLMInstance; +} + +/** + * Set the default LLM instance + */ +export function setDefaultLLM(llm: LLM | null): void { + defaultLLMInstance = llm; +} + +/** @deprecated Use getDefaultLLM() */ +export function getDefaultLlamaCpp(): LLM { + return getDefaultLLM(); } /** - * Set a custom default LlamaCpp instance (useful for testing). Setting a + * Set a custom default LLM instance (useful for testing). Setting a * non-null instance also ensures the darwin exit guard is installed — keeps * the invariant intact for test doubles that didn't go through the real * constructor. + * + * @deprecated Use setDefaultLLM() */ -export function setDefaultLlamaCpp(llm: LlamaCpp | null): void { +export function setDefaultLlamaCpp(llm: LLM | null): void { if (llm !== null) installDarwinExitGuard(); - defaultLlamaCpp = llm; + setDefaultLLM(llm); } /** - * Peek at the default LlamaCpp instance without instantiating one. Used by + * Peek at the default LLM instance without instantiating one. Used by * doctor and lifecycle diagnostics. */ export function hasDefaultLlamaCpp(): boolean { - return defaultLlamaCpp !== null; + return defaultLLMInstance !== null; } /** - * Dispose the default LlamaCpp instance if it exists. + * Dispose the default LLM instance if it exists. * Call this before process exit to prevent NAPI crashes. */ export async function disposeDefaultLlamaCpp(): Promise { - if (defaultLlamaCpp) { - await defaultLlamaCpp.dispose(); - defaultLlamaCpp = null; + if (defaultLLMInstance) { + await defaultLLMInstance.dispose(); + defaultLLMInstance = null; } } diff --git a/src/remote-llm.ts b/src/remote-llm.ts new file mode 100644 index 000000000..d9b1554d6 --- /dev/null +++ b/src/remote-llm.ts @@ -0,0 +1,546 @@ +/** + * remote-llm.ts - OpenAI-compatible remote embedding & reranking backend + * + * Implements the LLM interface by calling HTTP endpoints (vLLM, Ollama, OpenAI, etc.). + * Only supports embed/rerank operations — generate/expandQuery throw. + */ + +import type { + LLM, + EmbedOptions, + EmbeddingResult, + GenerateOptions, + GenerateResult, + ModelInfo, + Queryable, + QueryType, + RerankDocument, + RerankOptions, + RerankResult, +} from "./llm.js"; + +// ============================================================================= +// Configuration +// ============================================================================= + +export type RemoteLLMConfig = { + /** Base URL for embedding endpoint (e.g. http://gpu-host:8000/v1) */ + embedApiUrl: string; + /** Model name for embedding (e.g. BAAI/bge-m3) */ + embedApiModel: string; + /** Optional bearer token for embedding endpoint */ + embedApiKey?: string; + /** Base URL for rerank endpoint (defaults to embedApiUrl) */ + rerankApiUrl?: string; + /** Model name for reranking */ + rerankApiModel?: string; + /** Optional bearer token for rerank endpoint */ + rerankApiKey?: string; + /** Base URL for query-expansion endpoint (defaults to embedApiUrl). Hits POST /chat/completions. */ + expandApiUrl?: string; + /** Model name for query expansion (any chat-completion model). */ + expandApiModel?: string; + /** Optional bearer token for expand endpoint. */ + expandApiKey?: string; + /** Connect timeout in ms (default: 5000) */ + connectTimeoutMs?: number; + /** Read timeout for embedding in ms (default: 30000) */ + embedReadTimeoutMs?: number; + /** Read timeout for reranking in ms (default: 60000) */ + rerankReadTimeoutMs?: number; + /** Read timeout for query expansion in ms (default: 30000) */ + expandReadTimeoutMs?: number; + /** Max texts per embed HTTP request (default: 32) */ + maxBatchSize?: number; +}; + +// ============================================================================= +// Circuit Breaker +// ============================================================================= + +type CircuitState = "closed" | "open" | "half-open"; + +class CircuitBreaker { + private state: CircuitState = "closed"; + private failures = 0; + private lastFailureTime = 0; + private readonly maxFailures: number; + private readonly cooldownMs: number; + + constructor(maxFailures = 3, cooldownMs = 10 * 60 * 1000) { + this.maxFailures = maxFailures; + this.cooldownMs = cooldownMs; + } + + canAttempt(): boolean { + if (this.state === "closed") return true; + if (this.state === "open") { + if (Date.now() - this.lastFailureTime >= this.cooldownMs) { + this.state = "half-open"; + return true; + } + return false; + } + // half-open: allow one attempt + return true; + } + + onSuccess(): void { + this.state = "closed"; + this.failures = 0; + } + + onFailure(): void { + this.failures++; + this.lastFailureTime = Date.now(); + if (this.state === "half-open" || this.failures >= this.maxFailures) { + this.state = "open"; + } + } + + getState(): CircuitState { + return this.state; + } +} + +// ============================================================================= +// RemoteLLM +// ============================================================================= + +export class RemoteLLM implements LLM { + private readonly config: Required< + Pick + > & RemoteLLMConfig; + + private readonly embedBreaker = new CircuitBreaker(); + private readonly rerankBreaker = new CircuitBreaker(); + private expectedDimensions: number | null = null; + + constructor(config: RemoteLLMConfig) { + this.config = { + connectTimeoutMs: 5000, + embedReadTimeoutMs: 30000, + rerankReadTimeoutMs: 60000, + maxBatchSize: 32, + ...config, + }; + } + + get embedModelName(): string { + return this.config.embedApiModel; + } + + /** Rerank model — defaults to the embedding model when no separate rerank model is configured. */ + get rerankModelName(): string { + return this.config.rerankApiModel || this.config.embedApiModel; + } + + /** Remote backend exposes no local generation model; use the embed model as a placeholder identifier. */ + get generateModelName(): string { + return this.config.embedApiModel; + } + + get usesRemoteEmbedding(): boolean { + return true; + } + + /** True when expandApiModel is configured and remote query expansion is available. */ + get supportsExpand(): boolean { + return !!this.config.expandApiModel; + } + + /** True when rerankApiModel is configured and remote reranking is available. */ + get supportsRerank(): boolean { + return !!this.config.rerankApiModel; + } + + /** + * Remote backends have no local tokenizer. HybridLLM proxies tokenize/detokenize + * to the local LlamaCpp; bare RemoteLLM use will throw. + */ + async tokenize(_text: string): Promise { + throw new Error("RemoteLLM.tokenize is unavailable; use HybridLLM to access the local tokenizer."); + } + + async detokenize(_tokens: readonly never[]): Promise { + throw new Error("RemoteLLM.detokenize is unavailable; use HybridLLM to access the local tokenizer."); + } + + // --------------------------------------------------------------------------- + // Embedding + // --------------------------------------------------------------------------- + + async embed(text: string, options?: EmbedOptions): Promise { + const results = await this.embedBatch([text], options); + return results[0] ?? null; + } + + async embedBatch(texts: string[], _options?: EmbedOptions): Promise<(EmbeddingResult | null)[]> { + if (texts.length === 0) return []; + + if (!this.embedBreaker.canAttempt()) { + throw new Error( + `Remote embedding circuit breaker is open — endpoint ${this.config.embedApiUrl} is unavailable. ` + + `Will retry after cooldown.` + ); + } + + const batchSize = this.config.maxBatchSize; + const results: (EmbeddingResult | null)[] = []; + + for (let i = 0; i < texts.length; i += batchSize) { + const batch = texts.slice(i, i + batchSize); + const batchResults = await this.embedBatchRequest(batch); + results.push(...batchResults); + } + + return results; + } + + private async embedBatchRequest(texts: string[]): Promise<(EmbeddingResult | null)[]> { + const url = normalizeUrl(this.config.embedApiUrl, "/embeddings"); + const headers: Record = { "Content-Type": "application/json" }; + if (this.config.embedApiKey) { + headers["Authorization"] = `Bearer ${this.config.embedApiKey}`; + } + + const body = JSON.stringify({ + model: this.config.embedApiModel, + input: texts, + }); + + try { + const response = await fetchWithTimeout(url, { + method: "POST", + headers, + body, + }, this.config.embedReadTimeoutMs); + + if (!response.ok) { + const errText = await response.text().catch(() => ""); + throw new Error(`Embedding API returned ${response.status}: ${errText}`); + } + + const json = await response.json() as { + data: { embedding: number[]; index: number }[]; + }; + + // Validate dimensions consistency + if (json.data.length > 0) { + const dim = json.data[0]!.embedding.length; + if (this.expectedDimensions === null) { + this.expectedDimensions = dim; + } else if (dim !== this.expectedDimensions) { + throw new Error( + `Embedding dimension mismatch: expected ${this.expectedDimensions}, got ${dim}. ` + + `This usually means the remote model changed.` + ); + } + } + + // Sort by index to match input order + const sorted = [...json.data].sort((a, b) => a.index - b.index); + const results: (EmbeddingResult | null)[] = sorted.map(item => ({ + embedding: item.embedding, + model: this.config.embedApiModel, + })); + + this.embedBreaker.onSuccess(); + return results; + } catch (err) { + this.embedBreaker.onFailure(); + throw err; + } + } + + // --------------------------------------------------------------------------- + // Reranking + // --------------------------------------------------------------------------- + + async rerank(query: string, documents: RerankDocument[], _options?: RerankOptions): Promise { + const rerankUrl = this.config.rerankApiUrl || this.config.embedApiUrl; + const rerankModel = this.config.rerankApiModel; + const rerankKey = this.config.rerankApiKey || this.config.embedApiKey; + + if (!rerankModel) { + throw new Error("Remote reranking requires rerankApiModel to be configured"); + } + + if (!this.rerankBreaker.canAttempt()) { + throw new Error( + `Remote rerank circuit breaker is open — endpoint ${rerankUrl} is unavailable. ` + + `Will retry after cooldown.` + ); + } + + const url = normalizeUrl(rerankUrl, "/rerank"); + const headers: Record = { "Content-Type": "application/json" }; + if (rerankKey) { + headers["Authorization"] = `Bearer ${rerankKey}`; + } + + const body = JSON.stringify({ + model: rerankModel, + query, + documents: documents.map(d => d.text), + }); + + try { + const response = await fetchWithTimeout(url, { + method: "POST", + headers, + body, + }, this.config.rerankReadTimeoutMs); + + if (!response.ok) { + const errText = await response.text().catch(() => ""); + throw new Error(`Rerank API returned ${response.status}: ${errText}`); + } + + const json = await response.json() as { + results: { index: number; relevance_score: number }[]; + }; + + // Normalize relevance_score via sigmoid σ(x) = 1/(1+e^-x). + // + // Many cross-encoder rerankers exposed through llama.cpp's /v1/rerank + // endpoint (notably bge-reranker-v2-m3, BAAI/bge-reranker-large, jina- + // reranker-v2) emit log-odds (range roughly -10..+10, negative = poor + // match, positive = good match). The qmd consumer (store.ts blend + // formula at line ~4767 and the --min-score default of 0.3) assumes a + // 0..1 probability range. Without normalization, every blended score + // ends up negative and the default min-score filter drops all results. + // + // Sigmoid is monotonic so ordering is preserved. For rerankers that + // already emit 0..1 scores (some Voyage / Cohere endpoints), sigmoid + // is a no-op-ish squash that keeps them in [0,1]. The conversion is + // safe in both directions. + const results = json.results.map(r => ({ + file: documents[r.index]!.file, + score: 1 / (1 + Math.exp(-r.relevance_score)), + index: r.index, + })); + + this.rerankBreaker.onSuccess(); + return { results, model: rerankModel }; + } catch (err) { + this.rerankBreaker.onFailure(); + throw err; + } + } + + // --------------------------------------------------------------------------- + // Unsupported operations (these require local models) + // --------------------------------------------------------------------------- + + async generate(_prompt: string, _options?: GenerateOptions): Promise { + throw new Error("RemoteLLM does not support text generation — use HybridLLM to route generation to a local backend"); + } + + async modelExists(_model: string): Promise { + return { name: this.config.embedApiModel, exists: true }; + } + + async expandQuery( + query: string, + options?: { context?: string; includeLexical?: boolean; intent?: string }, + ): Promise { + const expandUrl = this.config.expandApiUrl || this.config.embedApiUrl; + const expandModel = this.config.expandApiModel; + const expandKey = this.config.expandApiKey || this.config.embedApiKey; + const includeLexical = options?.includeLexical ?? true; + const intent = options?.intent; + + // Shared fallback shape — used whenever the remote call fails OR returns + // nothing parseable. Mirrors LocalLLM.expandQuery's fallback exactly so + // downstream code doesn't see a behavior difference. + const defaultFallback = (): Queryable[] => { + const triple: Queryable[] = [ + { type: "hyde", text: `Information about ${query}` }, + { type: "lex", text: query }, + { type: "vec", text: query }, + ]; + return includeLexical ? triple : triple.filter(q => q.type !== "lex"); + }; + + if (!expandModel) { + // Configured to use remote but no expand model set → safe default. + return defaultFallback(); + } + + // Prompt the chat model to emit the lex/vec/hyde format that + // LocalLLM.expandQuery produces via grammar-constrained sampling. + // Without llama.cpp grammar we have to ask politely; parsing below is + // tolerant of variations and falls back if the model goes off-script. + const systemPrompt = + "You expand search queries for a hybrid retrieval system. " + + "Output 3 to 6 query variants, one per line, each prefixed with its type. " + + "Types:\n" + + " lex - keyword/BM25-friendly phrasing (extract distinctive terms)\n" + + " vec - semantic embedding-friendly phrasing (paraphrase intent)\n" + + " hyde - a hypothetical answer or document passage that would match the query\n" + + "\n" + + "Format strictly as:\n" + + ": \n" + + "\n" + + "No preamble, no explanation, no blank lines. Every line MUST start with " + + "exactly 'lex: ', 'vec: ', or 'hyde: '. Each variant should contain at " + + "least one term from the original query."; + + const userPrompt = intent + ? `Expand this search query: ${query}\nQuery intent: ${intent}` + : `Expand this search query: ${query}`; + + const url = normalizeUrl(expandUrl, "/chat/completions"); + const headers: Record = { "Content-Type": "application/json" }; + if (expandKey) headers["Authorization"] = `Bearer ${expandKey}`; + + const body = JSON.stringify({ + model: expandModel, + messages: [ + { role: "system", content: systemPrompt }, + { role: "user", content: userPrompt }, + ], + temperature: 0.7, + top_p: 0.8, + max_tokens: 600, + }); + + let content = ""; + try { + const response = await fetchWithTimeout( + url, + { method: "POST", headers, body }, + this.config.expandReadTimeoutMs ?? 30000, + ); + if (!response.ok) { + const errText = await response.text().catch(() => ""); + throw new Error(`Expand API returned ${response.status}: ${errText}`); + } + const json = (await response.json()) as { + choices?: { message?: { content?: string } }[]; + }; + content = json.choices?.[0]?.message?.content ?? ""; + } catch (err) { + // Network error, timeout, or non-2xx — fall back to the default triple. + // Don't escalate to caller; LocalLLM also masks failures behind the + // fallback to keep search resilient. + console.error("Remote query expansion failed:", err); + return defaultFallback(); + } + + // Parse — mirror LocalLLM's parsing exactly so downstream sees consistent + // shapes regardless of which backend produced the expansion. + const lines = content.trim().split("\n"); + const queryLower = query.toLowerCase(); + const queryTerms = queryLower + .replace(/[^a-z0-9\s]/g, " ") + .split(/\s+/) + .filter(Boolean); + + const hasQueryTerm = (text: string): boolean => { + const lower = text.toLowerCase(); + if (queryTerms.length === 0) return true; + return queryTerms.some(term => lower.includes(term)); + }; + + const queryables: Queryable[] = lines + .map((line): Queryable | null => { + const colonIdx = line.indexOf(":"); + if (colonIdx === -1) return null; + const type = line.slice(0, colonIdx).trim().toLowerCase(); + if (type !== "lex" && type !== "vec" && type !== "hyde") return null; + const text = line.slice(colonIdx + 1).trim(); + if (!text) return null; + if (!hasQueryTerm(text)) return null; + return { type: type as QueryType, text }; + }) + .filter((q): q is Queryable => q !== null); + + const filtered = includeLexical + ? queryables + : queryables.filter(q => q.type !== "lex"); + + if (filtered.length > 0) return filtered; + return defaultFallback(); + } + + async dispose(): Promise { + // Nothing to dispose for HTTP client + } +} + +// ============================================================================= +// Helpers +// ============================================================================= + +/** + * Normalize a base URL and append a path, handling trailing slashes. + */ +function normalizeUrl(baseUrl: string, path: string): string { + const base = baseUrl.replace(/\/+$/, ""); + return `${base}${path}`; +} + +/** + * Fetch with a timeout using AbortSignal.timeout(). + */ +async function fetchWithTimeout( + url: string, + init: RequestInit, + timeoutMs: number +): Promise { + return fetch(url, { + ...init, + signal: AbortSignal.timeout(timeoutMs), + }); +} + +// ============================================================================= +// Configuration from environment +// ============================================================================= + +/** + * Create a RemoteLLMConfig from environment variables and optional YAML config. + * Returns null if remote embedding is not configured. + */ +export function remoteConfigFromEnv(yamlModels?: { + embed_api_url?: string; + embed_api_model?: string; + embed_api_key?: string; + rerank_api_url?: string; + rerank_api_model?: string; + rerank_api_key?: string; + expand_api_url?: string; + expand_api_model?: string; + expand_api_key?: string; +}): RemoteLLMConfig | null { + const embedApiUrl = process.env.QMD_EMBED_API_URL || yamlModels?.embed_api_url; + const embedApiModel = process.env.QMD_EMBED_API_MODEL || yamlModels?.embed_api_model; + + if (!embedApiUrl || !embedApiModel) return null; + + return { + embedApiUrl, + embedApiModel, + embedApiKey: process.env.QMD_EMBED_API_KEY || yamlModels?.embed_api_key, + rerankApiUrl: process.env.QMD_RERANK_API_URL || yamlModels?.rerank_api_url, + rerankApiModel: process.env.QMD_RERANK_API_MODEL || yamlModels?.rerank_api_model, + rerankApiKey: process.env.QMD_RERANK_API_KEY || yamlModels?.rerank_api_key, + expandApiUrl: process.env.QMD_EXPAND_API_URL || yamlModels?.expand_api_url, + expandApiModel: process.env.QMD_EXPAND_API_MODEL || yamlModels?.expand_api_model, + expandApiKey: process.env.QMD_EXPAND_API_KEY || yamlModels?.expand_api_key, + connectTimeoutMs: parseEnvInt("QMD_REMOTE_CONNECT_TIMEOUT", 5000), + embedReadTimeoutMs: parseEnvInt("QMD_REMOTE_READ_TIMEOUT", 30000), + rerankReadTimeoutMs: parseEnvInt("QMD_REMOTE_RERANK_TIMEOUT", 60000), + expandReadTimeoutMs: parseEnvInt("QMD_REMOTE_EXPAND_TIMEOUT", 30000), + maxBatchSize: parseEnvInt("QMD_REMOTE_BATCH_SIZE", 32), + }; +} + +function parseEnvInt(name: string, defaultValue: number): number { + const val = process.env[name]; + if (!val) return defaultValue; + const parsed = parseInt(val, 10); + return Number.isFinite(parsed) && parsed > 0 ? parsed : defaultValue; +} diff --git a/src/store.ts b/src/store.ts index 99e36b861..4ebbd7000 100644 --- a/src/store.ts +++ b/src/store.ts @@ -20,14 +20,14 @@ import { readFileSync, realpathSync, statSync, mkdirSync } from "node:fs"; import fastGlob from "fast-glob"; import { qmdHomedir } from "./paths.js"; import { - LlamaCpp, - getDefaultLlamaCpp, + getDefaultLLM, formatQueryForEmbedding, formatDocForEmbedding, withLLMSessionForLlm, DEFAULT_EMBED_MODEL_URI, DEFAULT_RERANK_MODEL_URI, DEFAULT_GENERATE_MODEL_URI, + type LLM, type RerankDocument, type ILLMSession, } from "./llm.js"; @@ -77,11 +77,11 @@ export function getEmbeddingFingerprint(model: string = DEFAULT_EMBED_MODEL): st } /** - * Get the LlamaCpp instance for a store — prefers the store's own instance, + * Get the LLM instance for a store — prefers the store's own instance, * falls back to the global singleton. */ -function getLlm(store: Store): LlamaCpp { - return store.llm ?? getDefaultLlamaCpp(); +function getLlm(store: Store): LLM { + return store.llm ?? getDefaultLLM(); } // ============================================================================= @@ -1176,8 +1176,8 @@ function ensureVecTableInternal(db: Database, dimensions: number): void { export type Store = { db: Database; dbPath: string; - /** Optional LlamaCpp instance for this store (overrides the global singleton) */ - llm?: LlamaCpp; + /** Optional LLM instance for this store (overrides the global singleton) */ + llm?: LLM; close: () => void; ensureVecTable: (dimensions: number) => void; @@ -1604,6 +1604,7 @@ export async function generateEmbeddings( // Use store's LlamaCpp or global singleton, wrapped in a session const embedModelUri = model; + const usesRemoteEmbedding = llm.usesRemoteEmbedding === true; // Create a session manager for this llm instance const result = await withLLMSessionForLlm(llm, async (session) => { @@ -1701,13 +1702,21 @@ export async function generateEmbeddings( if (!doc.body.trim()) continue; const title = extractTitle(doc.body, doc.path); - const chunks = await chunkDocumentByTokens( - doc.body, - undefined, undefined, undefined, - doc.path, - options?.chunkStrategy, - session.signal, - ); + const chunks = usesRemoteEmbedding + ? await chunkDocumentByApproxTokens( + doc.body, + undefined, undefined, undefined, + doc.path, + options?.chunkStrategy, + session.signal, + ) + : await chunkDocumentByTokens( + doc.body, + undefined, undefined, undefined, + doc.path, + options?.chunkStrategy, + session.signal, + ); for (let seq = 0; seq < chunks.length; seq++) { batchChunks.push({ @@ -2658,7 +2667,10 @@ export async function chunkDocumentByTokens( chunkStrategy: ChunkStrategy = "regex", signal?: AbortSignal ): Promise<{ text: string; pos: number; tokens: number }[]> { - const llm = getDefaultLlamaCpp(); + const llm = getDefaultLLM(); + + // Check if the LLM supports tokenization (LlamaCpp does, RemoteLLM doesn't) + const canTokenize = typeof (llm as any).tokenize === "function"; // Use moderate chars/token estimate (prose ~4, code ~2, mixed ~3) // If chunks exceed limit, they'll be re-split with actual ratio @@ -2739,6 +2751,35 @@ export async function chunkDocumentByTokens( return results; } +/** + * Chunk a document using only character-space heuristics. + * Used when the active embedding backend is remote and we want to avoid + * initializing a local tokenizer/model just for preprocessing. + */ +export async function chunkDocumentByApproxTokens( + content: string, + maxTokens: number = CHUNK_SIZE_TOKENS, + overlapTokens: number = CHUNK_OVERLAP_TOKENS, + windowTokens: number = CHUNK_WINDOW_TOKENS, + filepath?: string, + chunkStrategy: ChunkStrategy = "regex", + signal?: AbortSignal +): Promise<{ text: string; pos: number; tokens: number }[]> { + if (signal?.aborted) return []; + + const avgCharsPerToken = 3; + const maxChars = maxTokens * avgCharsPerToken; + const overlapChars = overlapTokens * avgCharsPerToken; + const windowChars = windowTokens * avgCharsPerToken; + const charChunks = await chunkDocumentAsync(content, maxChars, overlapChars, windowChars, filepath, chunkStrategy); + + return charChunks.map((chunk) => ({ + text: chunk.text, + pos: chunk.pos, + tokens: Math.max(1, Math.ceil(chunk.text.length / avgCharsPerToken)), + })); +} + // ============================================================================= // Fuzzy matching // ============================================================================= @@ -3618,12 +3659,12 @@ export async function searchVec(db: Database, query: string, model: string, limi // Embeddings // ============================================================================= -async function getEmbedding(text: string, model: string, isQuery: boolean, session?: ILLMSession, llmOverride?: LlamaCpp): Promise { +async function getEmbedding(text: string, model: string, isQuery: boolean, session?: ILLMSession, llmOverride?: LLM): Promise { // Format text using the appropriate prompt template const formattedText = isQuery ? formatQueryForEmbedding(text, model) : formatDocForEmbedding(text, undefined, model); const result = session ? await session.embed(formattedText, { model, isQuery }) - : await (llmOverride ?? getDefaultLlamaCpp()).embed(formattedText, { model, isQuery }); + : await (llmOverride ?? getDefaultLLM()).embed(formattedText, { model, isQuery }); return result?.embedding || null; } @@ -3778,7 +3819,7 @@ function removeIncompleteEmbeddings(db: Database, expectedChunksByHash: Map { +export async function expandQuery(query: string, model: string = DEFAULT_QUERY_MODEL, db: Database, intent?: string, llmOverride?: LLM): Promise { // Check cache first — stored as JSON preserving types const cacheKey = getCacheKey("expandQuery", { query, model, ...(intent && { intent }) }); const cached = getCachedResult(db, cacheKey); @@ -3798,7 +3839,7 @@ export async function expandQuery(query: string, model: string = DEFAULT_QUERY_M } } - const llm = llmOverride ?? getDefaultLlamaCpp(); + const llm = llmOverride ?? getDefaultLLM(); // Note: LlamaCpp uses hardcoded model, model parameter is ignored const results = await llm.expandQuery(query, { intent }); @@ -3819,7 +3860,7 @@ export async function expandQuery(query: string, model: string = DEFAULT_QUERY_M // Reranking // ============================================================================= -export async function rerank(query: string, documents: { file: string; text: string }[], model: string = DEFAULT_RERANK_MODEL, db: Database, intent?: string, llmOverride?: LlamaCpp): Promise<{ file: string; score: number }[]> { +export async function rerank(query: string, documents: { file: string; text: string }[], model: string = DEFAULT_RERANK_MODEL, db: Database, intent?: string, llmOverride?: LLM): Promise<{ file: string; score: number }[]> { // Prepend intent to rerank query so the reranker scores with domain context const rerankQuery = intent ? `${intent}\n\n${query}` : query; @@ -3844,7 +3885,7 @@ export async function rerank(query: string, documents: { file: string; text: str // Rerank uncached documents using LlamaCpp if (uncachedDocsByChunk.size > 0) { - const llm = llmOverride ?? getDefaultLlamaCpp(); + const llm = llmOverride ?? getDefaultLLM(); const uncachedDocs = [...uncachedDocsByChunk.values()]; const rerankResult = await llm.rerank(rerankQuery, uncachedDocs, { model }); diff --git a/test/remote-llm-integration.test.ts b/test/remote-llm-integration.test.ts new file mode 100644 index 000000000..7f28b1116 --- /dev/null +++ b/test/remote-llm-integration.test.ts @@ -0,0 +1,412 @@ +/** + * Integration tests for RemoteLLM against live vLLM servers. + * + * Requires environment variables: + * VLLM_EMBED_URL - e.g. http://gpu-host:8002/v1 + * VLLM_EMBED_MODEL - e.g. Qwen/Qwen3-Embedding-0.6B + * VLLM_RERANK_URL - e.g. http://gpu-host:8001/v1 + * VLLM_RERANK_MODEL - e.g. qwen3-reranker-4b + * + * Skip these tests when no server is available (all tests guard on EMBED_URL). + */ + +import { describe, it, expect, beforeAll, beforeEach } from "vitest"; +import { RemoteLLM } from "../src/remote-llm.js"; +import { HybridLLM } from "../src/hybrid-llm.js"; +import { formatQueryForEmbedding, formatDocForEmbedding } from "../src/llm.js"; +import type { LLM } from "../src/llm.js"; + +const EMBED_URL = process.env.VLLM_EMBED_URL ?? ""; +const EMBED_MODEL = process.env.VLLM_EMBED_MODEL ?? ""; +const RERANK_URL = process.env.VLLM_RERANK_URL ?? ""; +const RERANK_MODEL = process.env.VLLM_RERANK_MODEL ?? ""; + +const SKIP = !EMBED_URL || !EMBED_MODEL; + +let remoteLlm: RemoteLLM; + +beforeAll(() => { + if (SKIP) return; + remoteLlm = new RemoteLLM({ + embedApiUrl: EMBED_URL, + embedApiModel: EMBED_MODEL, + rerankApiUrl: RERANK_URL, + rerankApiModel: RERANK_MODEL, + }); +}); + +// ============================================================================= +// Connectivity +// ============================================================================= + +describe.skipIf(SKIP)("Server connectivity", () => { + it("can reach the embedding server", async () => { + const res = await fetch(`${EMBED_URL}/models`); + expect(res.ok).toBe(true); + const json = await res.json() as any; + expect(json.data.length).toBeGreaterThan(0); + }); + + it("can reach the reranking server", async () => { + const res = await fetch(`${RERANK_URL}/models`); + expect(res.ok).toBe(true); + }); +}); + +// ============================================================================= +// Single embedding +// ============================================================================= + +describe.skipIf(SKIP)("Single embedding", () => { + it("returns a non-empty embedding vector", async () => { + const result = await remoteLlm.embed("The quick brown fox jumps over the lazy dog"); + expect(result).not.toBeNull(); + expect(result!.embedding.length).toBeGreaterThan(0); + expect(result!.model).toBe(EMBED_MODEL); + }); + + it("embedding values are finite numbers", async () => { + const result = await remoteLlm.embed("test embedding quality"); + expect(result).not.toBeNull(); + for (const val of result!.embedding) { + expect(Number.isFinite(val)).toBe(true); + } + }); + + it("embedding is normalized (L2 norm ≈ 1.0)", async () => { + const result = await remoteLlm.embed("normalization check"); + expect(result).not.toBeNull(); + const norm = Math.sqrt(result!.embedding.reduce((sum, v) => sum + v * v, 0)); + expect(norm).toBeCloseTo(1.0, 1); // within 0.1 + }); + + it("different texts produce different embeddings", async () => { + const [a, b] = await Promise.all([ + remoteLlm.embed("cats are wonderful pets"), + remoteLlm.embed("quantum computing research paper"), + ]); + expect(a).not.toBeNull(); + expect(b).not.toBeNull(); + // Cosine similarity should be < 1 (they are different) + const dot = a!.embedding.reduce((sum, v, i) => sum + v * b!.embedding[i]!, 0); + expect(dot).toBeLessThan(0.95); + }); + + it("similar texts produce similar embeddings", async () => { + const [a, b] = await Promise.all([ + remoteLlm.embed("how to train a puppy"), + remoteLlm.embed("puppy training tips"), + ]); + expect(a).not.toBeNull(); + expect(b).not.toBeNull(); + const dot = a!.embedding.reduce((sum, v, i) => sum + v * b!.embedding[i]!, 0); + expect(dot).toBeGreaterThan(0.7); + }); +}); + +// ============================================================================= +// Dimension consistency +// ============================================================================= + +describe.skipIf(SKIP)("Dimension consistency", () => { + it("all embeddings have the same dimension", async () => { + const texts = [ + "short", + "a medium length sentence about embedding dimensions", + "a much longer piece of text that goes on and on to test whether the embedding dimension stays consistent regardless of input length, which it absolutely should because the model always projects to a fixed-size output vector", + ]; + const results = await Promise.all(texts.map(t => remoteLlm.embed(t))); + const dims = results.map(r => r!.embedding.length); + expect(new Set(dims).size).toBe(1); + console.log(` Embedding dimension: ${dims[0]}`); + }); +}); + +// ============================================================================= +// Batch embedding +// ============================================================================= + +describe.skipIf(SKIP)("Batch embedding", () => { + it("embeds a batch of texts", async () => { + const texts = [ + "document one about machine learning", + "document two about cooking recipes", + "document three about space exploration", + ]; + const results = await remoteLlm.embedBatch(texts); + expect(results).toHaveLength(3); + for (const r of results) { + expect(r).not.toBeNull(); + expect(r!.embedding.length).toBeGreaterThan(0); + } + }); + + it("batch results match individual results", async () => { + const texts = ["alpha text", "beta text"]; + const [batchResults, individual1, individual2] = await Promise.all([ + remoteLlm.embedBatch(texts), + remoteLlm.embed("alpha text"), + remoteLlm.embed("beta text"), + ]); + + // Compare batch[0] with individual1 + expect(batchResults[0]!.embedding.length).toBe(individual1!.embedding.length); + // Embeddings should be very close (may not be exactly identical due to batching) + const dot = batchResults[0]!.embedding.reduce( + (sum, v, i) => sum + v * individual1!.embedding[i]!, 0 + ); + expect(dot).toBeGreaterThan(0.99); + }); + + it("handles empty batch", async () => { + const results = await remoteLlm.embedBatch([]); + expect(results).toEqual([]); + }); + + it("handles large batch (>32 texts, triggers splitting)", async () => { + const texts = Array.from({ length: 50 }, (_, i) => `document number ${i} about topic ${i % 5}`); + const results = await remoteLlm.embedBatch(texts); + expect(results).toHaveLength(50); + for (const r of results) { + expect(r).not.toBeNull(); + } + }); +}); + +// ============================================================================= +// Edge cases +// ============================================================================= + +describe.skipIf(SKIP)("Edge cases", () => { + it("handles very short text", async () => { + const result = await remoteLlm.embed("a"); + expect(result).not.toBeNull(); + expect(result!.embedding.length).toBeGreaterThan(0); + }); + + it("handles text with special characters", async () => { + const result = await remoteLlm.embed("café résumé naïve 日本語 中文 🎉 "); + expect(result).not.toBeNull(); + expect(result!.embedding.length).toBeGreaterThan(0); + }); + + it("handles multi-paragraph text", async () => { + const text = `# Introduction + +This is a long document with multiple paragraphs and markdown formatting. + +## Section 1 + +Some content here with **bold** and *italic* text. + +## Section 2 + +More content with a list: +- item one +- item two +- item three + +\`\`\`python +def hello(): + print("hello world") +\`\`\` +`; + const result = await remoteLlm.embed(text); + expect(result).not.toBeNull(); + expect(result!.embedding.length).toBeGreaterThan(0); + }); +}); + +// ============================================================================= +// Reranking +// ============================================================================= + +describe.skipIf(SKIP)("Reranking", () => { + it("reranks documents by relevance", async () => { + const query = "how to bake chocolate chip cookies"; + const documents = [ + { file: "space.md", text: "The Mars rover collected soil samples from the crater rim." }, + { file: "cookies.md", text: "Preheat oven to 375°F. Mix flour, butter, sugar and chocolate chips. Bake for 12 minutes." }, + { file: "quantum.md", text: "Quantum entanglement allows particles to be correlated over large distances." }, + { file: "baking.md", text: "Cookie recipes require precise measurements of ingredients like flour and sugar." }, + ]; + + const result = await remoteLlm.rerank(query, documents); + expect(result.model).toBe(RERANK_MODEL); + expect(result.results).toHaveLength(4); + + // The cookie/baking docs should rank higher than space/quantum + const scores = new Map(result.results.map(r => [r.file, r.score])); + console.log(" Rerank scores:", Object.fromEntries(scores)); + + expect(scores.get("cookies.md")!).toBeGreaterThan(scores.get("space.md")!); + expect(scores.get("cookies.md")!).toBeGreaterThan(scores.get("quantum.md")!); + }); + + it("scores are between 0 and 1", async () => { + const result = await remoteLlm.rerank("test query", [ + { file: "a.md", text: "relevant document about testing" }, + { file: "b.md", text: "unrelated document about gardening" }, + ]); + for (const r of result.results) { + expect(r.score).toBeGreaterThanOrEqual(0); + expect(r.score).toBeLessThanOrEqual(1); + } + }); + + it("preserves file mapping through index", async () => { + const documents = [ + { file: "first.md", text: "First document" }, + { file: "second.md", text: "Second document" }, + { file: "third.md", text: "Third document" }, + ]; + const result = await remoteLlm.rerank("query", documents); + const files = new Set(result.results.map(r => r.file)); + expect(files).toEqual(new Set(["first.md", "second.md", "third.md"])); + }); + + it("handles single document", async () => { + const result = await remoteLlm.rerank("test", [ + { file: "only.md", text: "The only document" }, + ]); + expect(result.results).toHaveLength(1); + expect(result.results[0]!.file).toBe("only.md"); + }); + + it("handles many documents", async () => { + const documents = Array.from({ length: 20 }, (_, i) => ({ + file: `doc${i}.md`, + text: `Document ${i} contains some text about topic ${i % 4}`, + })); + const result = await remoteLlm.rerank("topic about topic 2", documents); + expect(result.results).toHaveLength(20); + }); +}); + +// ============================================================================= +// Embedding format (remote models skip prefixes) +// ============================================================================= + +describe.skipIf(SKIP)("Embedding format for remote models", () => { + it("formatQueryForEmbedding returns raw text for remote model name", () => { + const formatted = formatQueryForEmbedding("search query", EMBED_MODEL); + expect(formatted).toBe("search query"); + }); + + it("formatDocForEmbedding returns raw text for remote model name", () => { + const formatted = formatDocForEmbedding("doc content", undefined, EMBED_MODEL); + expect(formatted).toBe("doc content"); + }); + + it("formatDocForEmbedding with title prepends title for remote model", () => { + const formatted = formatDocForEmbedding("doc content", "Title", EMBED_MODEL); + expect(formatted).toBe("Title\ndoc content"); + }); +}); + +// ============================================================================= +// HybridLLM integration +// ============================================================================= + +describe.skipIf(SKIP)("HybridLLM with real remote backend", () => { + // Mock local LLM for generate/expandQuery + function createMockLocal(): LLM { + return { + embedModelName: "local-embed-model", + embed: async () => ({ embedding: [0.5], model: "local" }), + embedBatch: async (texts) => texts.map(() => ({ embedding: [0.5], model: "local" })), + generate: async () => ({ text: "generated text", model: "local", done: true }), + modelExists: async (model) => ({ name: model, exists: true }), + expandQuery: async () => [{ type: "lex" as const, text: "expanded" }], + rerank: async () => ({ results: [], model: "local" }), + dispose: async () => {}, + }; + } + + it("routes embed through remote, returning real embeddings", async () => { + const hybrid = new HybridLLM(remoteLlm, createMockLocal()); + const result = await hybrid.embed("testing hybrid embedding"); + expect(result).not.toBeNull(); + // Real embedding has many dimensions, not just [0.5] + expect(result!.embedding.length).toBeGreaterThan(10); + expect(result!.model).toBe(EMBED_MODEL); + }); + + it("routes embedBatch through remote", async () => { + const hybrid = new HybridLLM(remoteLlm, createMockLocal()); + const results = await hybrid.embedBatch(["text one", "text two"]); + expect(results).toHaveLength(2); + expect(results[0]!.embedding.length).toBeGreaterThan(10); + }); + + it("routes rerank through remote", async () => { + const hybrid = new HybridLLM(remoteLlm, createMockLocal()); + const result = await hybrid.rerank("cookies", [ + { file: "a.md", text: "baking cookies at 350 degrees" }, + { file: "b.md", text: "orbiting space station" }, + ]); + expect(result.model).toBe(RERANK_MODEL); + expect(result.results).toHaveLength(2); + }); + + it("routes generate through local mock", async () => { + const hybrid = new HybridLLM(remoteLlm, createMockLocal()); + const result = await hybrid.generate("prompt"); + expect(result!.text).toBe("generated text"); + expect(result!.model).toBe("local"); + }); + + it("routes expandQuery through local mock", async () => { + const hybrid = new HybridLLM(remoteLlm, createMockLocal()); + const result = await hybrid.expandQuery("query"); + expect(result[0]!.text).toBe("expanded"); + }); + + it("embedModelName comes from remote", async () => { + const hybrid = new HybridLLM(remoteLlm, createMockLocal()); + expect(hybrid.embedModelName).toBe(EMBED_MODEL); + }); +}); + +// ============================================================================= +// End-to-end: embed → cosine similarity search +// ============================================================================= + +describe.skipIf(SKIP)("End-to-end embed + search simulation", () => { + it("finds the most relevant document via cosine similarity", async () => { + // Index some "documents" + const docs = [ + { file: "git.md", text: "Git is a distributed version control system for tracking changes in source code" }, + { file: "cooking.md", text: "To make pasta, boil water, add salt, cook noodles for 8 minutes" }, + { file: "docker.md", text: "Docker containers package applications with their dependencies for consistent deployment" }, + { file: "gardening.md", text: "Tomatoes need full sun and regular watering to produce fruit" }, + { file: "typescript.md", text: "TypeScript adds static type checking to JavaScript for safer code" }, + ]; + + // Embed all documents + const docEmbeddings = await remoteLlm.embedBatch(docs.map(d => d.text)); + + // Embed a query + const queryResult = await remoteLlm.embed("how to use version control for my code"); + expect(queryResult).not.toBeNull(); + + // Compute cosine similarities + const similarities = docEmbeddings.map((docEmb, i) => { + const dot = queryResult!.embedding.reduce((sum, v, j) => sum + v * docEmb!.embedding[j]!, 0); + return { file: docs[i]!.file, similarity: dot }; + }); + + similarities.sort((a, b) => b.similarity - a.similarity); + console.log(" Similarity ranking:"); + for (const s of similarities) { + console.log(` ${s.file}: ${s.similarity.toFixed(4)}`); + } + + // git.md should be the top result for a version control query + expect(similarities[0]!.file).toBe("git.md"); + // cooking/gardening should be near the bottom + const cookingRank = similarities.findIndex(s => s.file === "cooking.md"); + const gitRank = similarities.findIndex(s => s.file === "git.md"); + expect(gitRank).toBeLessThan(cookingRank); + }); +}); diff --git a/test/remote-llm.test.ts b/test/remote-llm.test.ts new file mode 100644 index 000000000..19f05f1bb --- /dev/null +++ b/test/remote-llm.test.ts @@ -0,0 +1,580 @@ +/** + * Tests for RemoteLLM and HybridLLM + * + * Uses a local HTTP server to mock OpenAI-compatible endpoints. + */ + +import { describe, it, expect, beforeAll, afterAll, beforeEach, afterEach } from "vitest"; +import { createServer, type Server, type IncomingMessage, type ServerResponse } from "http"; +import { RemoteLLM, remoteConfigFromEnv, type RemoteLLMConfig } from "../src/remote-llm.js"; +import { HybridLLM } from "../src/hybrid-llm.js"; +import { isRemoteModel, formatQueryForEmbedding, formatDocForEmbedding, getDefaultLLM, setDefaultLLM, LlamaCpp } from "../src/llm.js"; +import type { LLM, EmbeddingResult, RerankResult, Queryable, GenerateResult, ModelInfo } from "../src/llm.js"; + +// ============================================================================= +// Mock HTTP server +// ============================================================================= + +type MockHandler = (req: IncomingMessage, body: string) => { status: number; body: any }; + +let server: Server; +let serverPort: number; +let mockHandler: MockHandler; + +function setMockHandler(handler: MockHandler) { + mockHandler = handler; +} + +beforeAll(async () => { + server = createServer(async (req: IncomingMessage, res: ServerResponse) => { + const chunks: Buffer[] = []; + for await (const chunk of req) { + chunks.push(chunk as Buffer); + } + const body = Buffer.concat(chunks).toString(); + + try { + const result = mockHandler(req, body); + res.writeHead(result.status, { "Content-Type": "application/json" }); + res.end(JSON.stringify(result.body)); + } catch (err: any) { + res.writeHead(500, { "Content-Type": "application/json" }); + res.end(JSON.stringify({ error: err.message })); + } + }); + + await new Promise((resolve) => { + server.listen(0, "127.0.0.1", () => { + const addr = server.address(); + if (typeof addr === "object" && addr) { + serverPort = addr.port; + } + resolve(); + }); + }); +}); + +afterAll(() => { + server.close(); +}); + +function baseUrl(): string { + return `http://127.0.0.1:${serverPort}/v1`; +} + +function createRemoteLLM(overrides?: Partial): RemoteLLM { + return new RemoteLLM({ + embedApiUrl: baseUrl(), + embedApiModel: "test-model", + ...overrides, + }); +} + +// ============================================================================= +// RemoteLLM Tests +// ============================================================================= + +describe("RemoteLLM", () => { + describe("embed", () => { + it("should embed a single text", async () => { + setMockHandler((req, body) => { + const parsed = JSON.parse(body); + expect(parsed.model).toBe("test-model"); + expect(parsed.input).toEqual(["hello world"]); + return { + status: 200, + body: { + data: [{ embedding: [0.1, 0.2, 0.3], index: 0 }], + }, + }; + }); + + const llm = createRemoteLLM(); + const result = await llm.embed("hello world"); + expect(result).not.toBeNull(); + expect(result!.embedding).toEqual([0.1, 0.2, 0.3]); + expect(result!.model).toBe("test-model"); + }); + + it("should embed a batch of texts", async () => { + setMockHandler((_req, body) => { + const parsed = JSON.parse(body); + return { + status: 200, + body: { + data: parsed.input.map((text: string, i: number) => ({ + embedding: [i * 0.1, i * 0.2], + index: i, + })), + }, + }; + }); + + const llm = createRemoteLLM(); + const results = await llm.embedBatch(["text1", "text2", "text3"]); + expect(results).toHaveLength(3); + expect(results[0]!.embedding).toEqual([0, 0]); + expect(results[2]!.embedding).toEqual([0.2, 0.4]); + }); + + it("should return empty array for empty input", async () => { + const llm = createRemoteLLM(); + const results = await llm.embedBatch([]); + expect(results).toEqual([]); + }); + + it("should split large batches", async () => { + const requestBodies: string[][] = []; + setMockHandler((_req, body) => { + const parsed = JSON.parse(body); + requestBodies.push(parsed.input); + return { + status: 200, + body: { + data: parsed.input.map((_: string, i: number) => ({ + embedding: [1.0], + index: i, + })), + }, + }; + }); + + const llm = createRemoteLLM({ maxBatchSize: 2 }); + const texts = ["a", "b", "c", "d", "e"]; + const results = await llm.embedBatch(texts); + + expect(results).toHaveLength(5); + // Should have made 3 requests: [a,b], [c,d], [e] + expect(requestBodies).toHaveLength(3); + expect(requestBodies[0]).toEqual(["a", "b"]); + expect(requestBodies[1]).toEqual(["c", "d"]); + expect(requestBodies[2]).toEqual(["e"]); + }); + + it("should sort response by index", async () => { + setMockHandler(() => ({ + status: 200, + body: { + // Return in reverse order + data: [ + { embedding: [0.3], index: 2 }, + { embedding: [0.1], index: 0 }, + { embedding: [0.2], index: 1 }, + ], + }, + })); + + const llm = createRemoteLLM(); + const results = await llm.embedBatch(["a", "b", "c"]); + expect(results[0]!.embedding).toEqual([0.1]); + expect(results[1]!.embedding).toEqual([0.2]); + expect(results[2]!.embedding).toEqual([0.3]); + }); + }); + + describe("auth", () => { + it("should send Authorization header when key is set", async () => { + let authHeader: string | undefined; + setMockHandler((req) => { + authHeader = req.headers["authorization"] as string; + return { + status: 200, + body: { data: [{ embedding: [1.0], index: 0 }] }, + }; + }); + + const llm = createRemoteLLM({ embedApiKey: "test-key-123" }); + await llm.embed("test"); + expect(authHeader).toBe("Bearer test-key-123"); + }); + + it("should not send Authorization header when no key", async () => { + let authHeader: string | undefined; + setMockHandler((req) => { + authHeader = req.headers["authorization"] as string; + return { + status: 200, + body: { data: [{ embedding: [1.0], index: 0 }] }, + }; + }); + + const llm = createRemoteLLM(); + await llm.embed("test"); + expect(authHeader).toBeUndefined(); + }); + }); + + describe("dimension validation", () => { + it("should reject dimension mismatch after first response", async () => { + let callCount = 0; + setMockHandler(() => { + callCount++; + const dim = callCount === 1 ? [1.0, 2.0, 3.0] : [1.0, 2.0]; + return { + status: 200, + body: { data: [{ embedding: dim, index: 0 }] }, + }; + }); + + const llm = createRemoteLLM(); + // First call succeeds and locks dimensions to 3 + await llm.embed("first"); + // Second call should fail because dimensions changed + await expect(llm.embed("second")).rejects.toThrow("dimension mismatch"); + }); + }); + + describe("error handling", () => { + it("should throw on HTTP error", async () => { + setMockHandler(() => ({ + status: 500, + body: { error: "Internal server error" }, + })); + + const llm = createRemoteLLM(); + await expect(llm.embed("test")).rejects.toThrow("500"); + }); + + it("should open circuit breaker after failures", async () => { + setMockHandler(() => ({ + status: 500, + body: { error: "down" }, + })); + + const llm = createRemoteLLM(); + // Fail 3 times to trip the breaker + for (let i = 0; i < 3; i++) { + await expect(llm.embed("test")).rejects.toThrow(); + } + // Next call should fail immediately with circuit breaker message + await expect(llm.embed("test")).rejects.toThrow("circuit breaker"); + }); + }); + + describe("rerank", () => { + it("should rerank documents", async () => { + setMockHandler((_req, body) => { + const parsed = JSON.parse(body); + expect(parsed.model).toBe("rerank-model"); + expect(parsed.query).toBe("test query"); + expect(parsed.documents).toEqual(["doc A text", "doc B text"]); + return { + status: 200, + body: { + results: [ + { index: 1, relevance_score: 0.9 }, + { index: 0, relevance_score: 0.3 }, + ], + }, + }; + }); + + const llm = createRemoteLLM({ + rerankApiModel: "rerank-model", + }); + const result = await llm.rerank( + "test query", + [ + { file: "a.md", text: "doc A text" }, + { file: "b.md", text: "doc B text" }, + ] + ); + + expect(result.model).toBe("rerank-model"); + expect(result.results).toHaveLength(2); + // RemoteLLM.rerank sigmoid-normalizes the raw relevance_score (log-odds) + // into a 0..1 probability: σ(x) = 1/(1+e^-x). σ(0.9) ≈ 0.7109, + // σ(0.3) ≈ 0.5744. Ordering (monotonic) is preserved. + const bScore = result.results.find(r => r.file === "b.md")!.score; + const aScore = result.results.find(r => r.file === "a.md")!.score; + expect(bScore).toBeCloseTo(1 / (1 + Math.exp(-0.9)), 10); + expect(aScore).toBeCloseTo(1 / (1 + Math.exp(-0.3)), 10); + expect(bScore).toBeGreaterThan(aScore); + }); + + it("should throw when rerankApiModel not configured", async () => { + const llm = createRemoteLLM(); + await expect( + llm.rerank("query", [{ file: "a.md", text: "text" }]) + ).rejects.toThrow("rerankApiModel"); + }); + }); + + describe("unsupported operations", () => { + it("should throw on generate", async () => { + const llm = createRemoteLLM(); + await expect(llm.generate("prompt")).rejects.toThrow("does not support text generation"); + }); + + it("should fall back to the default expansion triple when no expand model is configured", async () => { + // RemoteLLM.expandQuery now implements query expansion via chat + // completions. When no expand model is configured it returns the same + // lex/vec/hyde fallback that LocalLLM produces, rather than throwing. + const llm = createRemoteLLM(); + const result = await llm.expandQuery("query"); + expect(result).toEqual([ + { type: "hyde", text: "Information about query" }, + { type: "lex", text: "query" }, + { type: "vec", text: "query" }, + ]); + }); + }); +}); + +// ============================================================================= +// HybridLLM Tests +// ============================================================================= + +describe("HybridLLM", () => { + // Simple mock local LLM + function createMockLocalLLM(): LLM { + return { + embedModelName: "local-model", + embed: async () => ({ embedding: [0.5], model: "local-model" }), + embedBatch: async (texts) => texts.map(() => ({ embedding: [0.5], model: "local-model" })), + generate: async () => ({ text: "expanded", model: "local-model", done: true }), + modelExists: async (model) => ({ name: model, exists: true }), + expandQuery: async () => [{ type: "lex" as const, text: "expanded query" }], + rerank: async () => ({ results: [], model: "local-model" }), + dispose: async () => {}, + }; + } + + it("should route embed to remote", async () => { + setMockHandler(() => ({ + status: 200, + body: { data: [{ embedding: [0.9], index: 0 }] }, + })); + + const remote = createRemoteLLM(); + const local = createMockLocalLLM(); + const hybrid = new HybridLLM(remote, local); + + const result = await hybrid.embed("test"); + // Should come from remote (0.9), not local (0.5) + expect(result!.embedding).toEqual([0.9]); + }); + + it("should route embedBatch to remote", async () => { + setMockHandler((_req, body) => { + const parsed = JSON.parse(body); + return { + status: 200, + body: { + data: parsed.input.map((_: string, i: number) => ({ + embedding: [0.9 + i * 0.01], + index: i, + })), + }, + }; + }); + + const remote = createRemoteLLM(); + const local = createMockLocalLLM(); + const hybrid = new HybridLLM(remote, local); + + const results = await hybrid.embedBatch(["a", "b"]); + expect(results[0]!.embedding).toEqual([0.9]); + expect(results[1]!.embedding).toEqual([0.91]); + }); + + it("should route generate to local", async () => { + const local = createMockLocalLLM(); + const remote = createRemoteLLM(); + const hybrid = new HybridLLM(remote, local); + + const result = await hybrid.generate("prompt"); + expect(result!.text).toBe("expanded"); + expect(result!.model).toBe("local-model"); + }); + + it("should route expandQuery to local", async () => { + const local = createMockLocalLLM(); + const remote = createRemoteLLM(); + const hybrid = new HybridLLM(remote, local); + + const result = await hybrid.expandQuery("test query"); + expect(result[0]!.text).toBe("expanded query"); + }); + + it("should use remote embedModelName", async () => { + const remote = createRemoteLLM({ embedApiModel: "BAAI/bge-m3" }); + const local = createMockLocalLLM(); + const hybrid = new HybridLLM(remote, local); + + expect(hybrid.embedModelName).toBe("BAAI/bge-m3"); + }); +}); + +// ============================================================================= +// Config Tests +// ============================================================================= + +describe("remoteConfigFromEnv", () => { + const origEnv = { ...process.env }; + + beforeEach(() => { + // Clear any QMD_ env vars + for (const key of Object.keys(process.env)) { + if (key.startsWith("QMD_") && key.includes("API")) { + delete process.env[key]; + } + } + }); + + afterAll(() => { + // Restore original env + for (const key of Object.keys(process.env)) { + if (key.startsWith("QMD_") && key.includes("API")) { + delete process.env[key]; + } + } + Object.assign(process.env, origEnv); + }); + + it("should return null when no config", () => { + expect(remoteConfigFromEnv()).toBeNull(); + }); + + it("should parse env vars", () => { + process.env.QMD_EMBED_API_URL = "http://gpu:8000/v1"; + process.env.QMD_EMBED_API_MODEL = "bge-m3"; + process.env.QMD_EMBED_API_KEY = "secret"; + + const config = remoteConfigFromEnv(); + expect(config).not.toBeNull(); + expect(config!.embedApiUrl).toBe("http://gpu:8000/v1"); + expect(config!.embedApiModel).toBe("bge-m3"); + expect(config!.embedApiKey).toBe("secret"); + }); + + it("should use YAML config as fallback", () => { + const config = remoteConfigFromEnv({ + embed_api_url: "http://yaml:8000/v1", + embed_api_model: "yaml-model", + }); + expect(config).not.toBeNull(); + expect(config!.embedApiUrl).toBe("http://yaml:8000/v1"); + }); + + it("should prefer env vars over YAML", () => { + process.env.QMD_EMBED_API_URL = "http://env:8000/v1"; + process.env.QMD_EMBED_API_MODEL = "env-model"; + + const config = remoteConfigFromEnv({ + embed_api_url: "http://yaml:8000/v1", + embed_api_model: "yaml-model", + }); + expect(config!.embedApiUrl).toBe("http://env:8000/v1"); + expect(config!.embedApiModel).toBe("env-model"); + }); +}); + +// ============================================================================= +// Embedding format tests +// ============================================================================= + +describe("isRemoteModel", () => { + it("should detect remote models", () => { + expect(isRemoteModel("BAAI/bge-m3")).toBe(true); + expect(isRemoteModel("intfloat/multilingual-e5-large")).toBe(true); + expect(isRemoteModel("text-embedding-ada-002")).toBe(true); + }); + + it("should detect local models", () => { + expect(isRemoteModel("hf:ggml-org/embeddinggemma-300M-GGUF/embeddinggemma-300M-Q8_0.gguf")).toBe(false); + expect(isRemoteModel("/path/to/model.gguf")).toBe(false); + }); +}); + +describe("formatQueryForEmbedding with remote models", () => { + it("should return raw query for remote models", () => { + expect(formatQueryForEmbedding("test query", "BAAI/bge-m3")).toBe("test query"); + }); + + it("should add prefix for local nomic models", () => { + expect(formatQueryForEmbedding("test query", "hf:ggml-org/embeddinggemma-300M-GGUF/embeddinggemma-300M-Q8_0.gguf")).toContain("task:"); + }); +}); + +describe("formatDocForEmbedding with remote models", () => { + it("should return raw text for remote models", () => { + expect(formatDocForEmbedding("doc text", undefined, "BAAI/bge-m3")).toBe("doc text"); + }); + + it("should include title when provided for remote models", () => { + expect(formatDocForEmbedding("doc text", "My Title", "BAAI/bge-m3")).toBe("My Title\ndoc text"); + }); +}); + +// ============================================================================= +// Local-only path (no remote config) +// ============================================================================= + +describe("Local-only LlamaCpp path", () => { + afterEach(() => { + // Reset to default so other tests aren't affected + setDefaultLLM(null); + }); + + it("getDefaultLLM() returns a LlamaCpp instance when nothing is configured", () => { + setDefaultLLM(null); + const llm = getDefaultLLM(); + expect(llm).toBeInstanceOf(LlamaCpp); + }); + + it("LlamaCpp instance satisfies the LLM interface", () => { + const llm = new LlamaCpp(); + // All LLM interface methods exist + expect(typeof llm.embed).toBe("function"); + expect(typeof llm.embedBatch).toBe("function"); + expect(typeof llm.generate).toBe("function"); + expect(typeof llm.modelExists).toBe("function"); + expect(typeof llm.expandQuery).toBe("function"); + expect(typeof llm.rerank).toBe("function"); + expect(typeof llm.dispose).toBe("function"); + expect(typeof llm.embedModelName).toBe("string"); + }); + + it("LlamaCpp has tokenize method (used by chunkDocumentByTokens duck-typing)", () => { + const llm = new LlamaCpp(); + expect(typeof llm.tokenize).toBe("function"); + }); + + it("setDefaultLLM with LlamaCpp is retrievable via getDefaultLLM", () => { + const llm = new LlamaCpp(); + setDefaultLLM(llm); + expect(getDefaultLLM()).toBe(llm); + }); + + it("remoteConfigFromEnv returns null when no env vars or YAML set", () => { + // Clear any remote env vars + const saved: Record = {}; + for (const key of ["QMD_EMBED_API_URL", "QMD_EMBED_API_MODEL"]) { + saved[key] = process.env[key]; + delete process.env[key]; + } + try { + expect(remoteConfigFromEnv()).toBeNull(); + expect(remoteConfigFromEnv({})).toBeNull(); + expect(remoteConfigFromEnv({ embed_api_url: undefined })).toBeNull(); + } finally { + for (const [key, val] of Object.entries(saved)) { + if (val !== undefined) process.env[key] = val; + } + } + }); + + it("formatQueryForEmbedding adds nomic prefix for default local model", () => { + // Default model is embeddinggemma (hf: URI), should get task prefix + const formatted = formatQueryForEmbedding("hello"); + expect(formatted).toContain("task:"); + expect(formatted).toContain("hello"); + }); + + it("formatDocForEmbedding adds nomic prefix for default local model", () => { + const formatted = formatDocForEmbedding("doc content", "My Doc"); + expect(formatted).toContain("title: My Doc"); + expect(formatted).toContain("text: doc content"); + }); +}); From 0b2f41e1a90f30ab148459d956e0c09b7686600b Mon Sep 17 00:00:00 2001 From: Kaspre Date: Tue, 2 Jun 2026 11:20:39 -0400 Subject: [PATCH 2/6] feat(rerank): recover from oversized rerank payloads by splitting/truncating When a remote /v1/rerank request is rejected as too large (HTTP 413 / "too large to process" / context length), RemoteLLM.rerank recursively bisects the batch and halve-truncates a single oversized document down to a 32-char floor, remapping response indices to the originals and re-sorting by score. Non-oversized errors still propagate so the circuit breaker / HybridLLM local fallback can react. Adapted from the rerank recovery in #619 (@loopyd). Co-Authored-By: Claude Opus 4.8 --- src/remote-llm.ts | 102 +++++++++++++++++++++++++++++++++------- test/remote-llm.test.ts | 57 ++++++++++++++++++++++ 2 files changed, 141 insertions(+), 18 deletions(-) diff --git a/src/remote-llm.ts b/src/remote-llm.ts index d9b1554d6..2cd9bd387 100644 --- a/src/remote-llm.ts +++ b/src/remote-llm.ts @@ -15,6 +15,7 @@ import type { Queryable, QueryType, RerankDocument, + RerankDocumentResult, RerankOptions, RerankResult, } from "./llm.js"; @@ -103,6 +104,26 @@ class CircuitBreaker { } } +/** Floor for halve-truncating a single oversized document during rerank recovery. */ +const RERANK_MIN_DOC_CHARS = 32; + +/** + * True when a rerank request failed because the payload exceeded what the + * server/model can process (the whole batch, or a single very long document). + * Drives batch-splitting recovery; non-oversized errors propagate instead so + * the circuit breaker / HybridLLM local fallback can react. + */ +function isOversizedRerankError(err: unknown): boolean { + if (!(err instanceof Error)) return false; + const m = err.message.toLowerCase(); + return m.includes("too large to process") + || m.includes("payload too large") + || m.includes(" 413") + || m.includes("context length") + || m.includes("maximum context") + || m.includes("too long"); +} + // ============================================================================= // RemoteLLM // ============================================================================= @@ -279,17 +300,22 @@ export class RemoteLLM implements LLM { headers["Authorization"] = `Bearer ${rerankKey}`; } - const body = JSON.stringify({ - model: rerankModel, - query, - documents: documents.map(d => d.text), - }); - - try { + // One rerank request for a slice of `documents` starting at `start`. + // Response indices are local to the submitted slice, so they are remapped + // to original document positions via `start`. relevance_score is + // sigmoid-normalized (see note below). + const rerankBatch = async ( + batch: RerankDocument[], + start: number, + ): Promise => { const response = await fetchWithTimeout(url, { method: "POST", headers, - body, + body: JSON.stringify({ + model: rerankModel, + query, + documents: batch.map(d => d.text), + }), }, this.config.rerankReadTimeoutMs); if (!response.ok) { @@ -307,20 +333,60 @@ export class RemoteLLM implements LLM { // endpoint (notably bge-reranker-v2-m3, BAAI/bge-reranker-large, jina- // reranker-v2) emit log-odds (range roughly -10..+10, negative = poor // match, positive = good match). The qmd consumer (store.ts blend - // formula at line ~4767 and the --min-score default of 0.3) assumes a - // 0..1 probability range. Without normalization, every blended score - // ends up negative and the default min-score filter drops all results. + // formula and the --min-score default of 0.3) assumes a 0..1 probability + // range. Without normalization, every blended score ends up negative and + // the default min-score filter drops all results. // // Sigmoid is monotonic so ordering is preserved. For rerankers that // already emit 0..1 scores (some Voyage / Cohere endpoints), sigmoid - // is a no-op-ish squash that keeps them in [0,1]. The conversion is - // safe in both directions. - const results = json.results.map(r => ({ - file: documents[r.index]!.file, - score: 1 / (1 + Math.exp(-r.relevance_score)), - index: r.index, - })); + // is a no-op-ish squash that keeps them in [0,1]. Safe in both directions. + return json.results.map(r => { + const documentIndex = start + r.index; + return { + file: documents[documentIndex]!.file, + score: 1 / (1 + Math.exp(-r.relevance_score)), + index: documentIndex, + }; + }); + }; + // Recovery for servers that reject an over-large rerank payload (the batch + // or a single very long document exceeds the model's context). On an + // oversized error we bisect the batch; a single oversized document is + // halve-truncated down to a floor and re-scored on the truncated text. + // Adapted from the recovery in tobi/qmd#619 (loopyd). Non-oversized errors + // propagate so the circuit breaker / HybridLLM local fallback can react. + const scoreBatch = async ( + batch: RerankDocument[], + start: number, + ): Promise => { + try { + return await rerankBatch(batch, start); + } catch (err) { + if (!isOversizedRerankError(err) || batch.length === 0) { + throw err; + } + if (batch.length === 1) { + const doc = batch[0]!; + const nextLength = Math.max(RERANK_MIN_DOC_CHARS, Math.floor(doc.text.length / 2)); + if (doc.text.length <= RERANK_MIN_DOC_CHARS || nextLength >= doc.text.length) { + throw err; + } + return scoreBatch([{ ...doc, text: doc.text.slice(0, nextLength) }], start); + } + const mid = Math.ceil(batch.length / 2); + const left = await scoreBatch(batch.slice(0, mid), start); + const right = await scoreBatch(batch.slice(mid), start + mid); + return [...left, ...right]; + } + }; + + try { + const results = await scoreBatch(documents, 0); + // Sub-batch recovery merges results out of global rank order; sort by the + // (monotonic) normalized score for a coherent final ranking. No-op on the + // common single-request path, where the server already returns ranked. + results.sort((a, b) => b.score - a.score); this.rerankBreaker.onSuccess(); return { results, model: rerankModel }; } catch (err) { diff --git a/test/remote-llm.test.ts b/test/remote-llm.test.ts index 19f05f1bb..87d1cdcc5 100644 --- a/test/remote-llm.test.ts +++ b/test/remote-llm.test.ts @@ -292,6 +292,63 @@ describe("RemoteLLM", () => { expect(bScore).toBeGreaterThan(aScore); }); + it("recovers from an oversized rerank batch by splitting", async () => { + // log-odds keyed by document text; the server rejects any multi-document + // batch as "too large", forcing recursive bisection down to single docs. + const scores: Record = { + "doc a": 0.5, "doc b": 1.5, "doc c": -0.5, "doc d": 2.5, + }; + setMockHandler((_req, body) => { + const parsed = JSON.parse(body); + if (parsed.documents.length > 1) { + return { status: 413, body: { error: "input is too large to process" } }; + } + return { + status: 200, + body: { results: [{ index: 0, relevance_score: scores[parsed.documents[0]] }] }, + }; + }); + + const llm = createRemoteLLM({ rerankApiModel: "rerank-model" }); + const result = await llm.rerank("q", [ + { file: "a.md", text: "doc a" }, + { file: "b.md", text: "doc b" }, + { file: "c.md", text: "doc c" }, + { file: "d.md", text: "doc d" }, + ]); + + // Every document is scored, response indices remap back to the original + // positions, and results come back sorted by normalized score descending. + expect(result.results).toHaveLength(4); + expect(result.results.map(r => r.file)).toEqual(["d.md", "b.md", "a.md", "c.md"]); + const byFile = Object.fromEntries(result.results.map(r => [r.file, r.index])); + expect(byFile["a.md"]).toBe(0); + expect(byFile["d.md"]).toBe(3); + expect(result.results[0].score).toBeCloseTo(1 / (1 + Math.exp(-2.5)), 10); + }); + + it("recovers from a single oversized document by truncating", async () => { + // The server rejects any request whose document exceeds 64 chars, forcing + // halve-truncation of the single oversized doc until it fits. + setMockHandler((_req, body) => { + const parsed = JSON.parse(body); + const longest = Math.max(...parsed.documents.map((d: string) => d.length)); + if (longest > 64) { + return { status: 400, body: { error: "input too long" } }; + } + return { + status: 200, + body: { results: parsed.documents.map((_: string, i: number) => ({ index: i, relevance_score: 1.0 })) }, + }; + }); + + const llm = createRemoteLLM({ rerankApiModel: "rerank-model" }); + const result = await llm.rerank("q", [{ file: "big.md", text: "x".repeat(500) }]); + expect(result.results).toHaveLength(1); + expect(result.results[0].file).toBe("big.md"); + expect(result.results[0].score).toBeCloseTo(1 / (1 + Math.exp(-1.0)), 10); + }); + it("should throw when rerankApiModel not configured", async () => { const llm = createRemoteLLM(); await expect( From df6950eff513bd3a168805764b4d86e63f24ca8f Mon Sep 17 00:00:00 2001 From: Kaspre Date: Tue, 2 Jun 2026 11:20:39 -0400 Subject: [PATCH 3/6] fix(remote): harden config + rerank score normalization - remoteConfigFromEnv: throw on a half-configured remote backend (embed_api_url set without embed_api_model, or vice-versa) instead of silently falling back to the local backend and skipping the remote pre-flight probe. - RemoteLLM.rerank: normalize scores once over the full (possibly recovery-split) result set, applying sigmoid only when logit-range values are present (any score < 0 or > 1). Rerankers that already return [0,1] probabilities (Cohere/Voyage-style) pass through unchanged, avoiding distortion of the blend and --min-score filtering. Co-Authored-By: Claude Opus 4.8 --- src/remote-llm.ts | 64 +++++++++++++++++++++++++++-------------- test/remote-llm.test.ts | 62 ++++++++++++++++++++++++++++++--------- 2 files changed, 92 insertions(+), 34 deletions(-) diff --git a/src/remote-llm.ts b/src/remote-llm.ts index 2cd9bd387..504c1c5a3 100644 --- a/src/remote-llm.ts +++ b/src/remote-llm.ts @@ -302,8 +302,9 @@ export class RemoteLLM implements LLM { // One rerank request for a slice of `documents` starting at `start`. // Response indices are local to the submitted slice, so they are remapped - // to original document positions via `start`. relevance_score is - // sigmoid-normalized (see note below). + // to original document positions via `start`. The RAW relevance_score is + // returned; score normalization happens once in the caller (see below) so + // the logit-vs-probability decision is consistent across recovered batches. const rerankBatch = async ( batch: RerankDocument[], start: number, @@ -327,24 +328,11 @@ export class RemoteLLM implements LLM { results: { index: number; relevance_score: number }[]; }; - // Normalize relevance_score via sigmoid σ(x) = 1/(1+e^-x). - // - // Many cross-encoder rerankers exposed through llama.cpp's /v1/rerank - // endpoint (notably bge-reranker-v2-m3, BAAI/bge-reranker-large, jina- - // reranker-v2) emit log-odds (range roughly -10..+10, negative = poor - // match, positive = good match). The qmd consumer (store.ts blend - // formula and the --min-score default of 0.3) assumes a 0..1 probability - // range. Without normalization, every blended score ends up negative and - // the default min-score filter drops all results. - // - // Sigmoid is monotonic so ordering is preserved. For rerankers that - // already emit 0..1 scores (some Voyage / Cohere endpoints), sigmoid - // is a no-op-ish squash that keeps them in [0,1]. Safe in both directions. return json.results.map(r => { const documentIndex = start + r.index; return { file: documents[documentIndex]!.file, - score: 1 / (1 + Math.exp(-r.relevance_score)), + score: r.relevance_score, // raw; normalized once in the caller index: documentIndex, }; }); @@ -382,10 +370,31 @@ export class RemoteLLM implements LLM { }; try { - const results = await scoreBatch(documents, 0); - // Sub-batch recovery merges results out of global rank order; sort by the - // (monotonic) normalized score for a coherent final ranking. No-op on the - // common single-request path, where the server already returns ranked. + const merged = await scoreBatch(documents, 0); + + // Normalize once, over the full merged set, so the decision is consistent + // even when oversized-recovery split the request into sub-batches. + // + // Cross-encoder rerankers exposed via llama.cpp's /v1/rerank (notably + // bge-reranker-v2-m3, BAAI/bge-reranker-large, jina-reranker-v2) emit + // log-odds (~−10..+10), which the qmd consumer (store.ts blend formula + + // the --min-score 0.3 default) would otherwise read as sub-zero + // probabilities and drop. We map those into [0,1] with sigmoid + // σ(x)=1/(1+e^-x) (monotonic, so ordering is preserved). + // + // But Cohere/Voyage-style endpoints already return probabilities in [0,1]; + // sigmoid would distort those (0.9→0.71, 0.01→0.50) and skew min-score + // filtering. So normalize ONLY when logit-range values are actually present + // (any score < 0 or > 1). (A logit reranker whose scores for a query all + // land within [0,1] is left as-is; an explicit per-endpoint normalization + // config could remove that ambiguity as a follow-up.) + const needsSigmoid = merged.some(r => r.score < 0 || r.score > 1); + const results = needsSigmoid + ? merged.map(r => ({ ...r, score: 1 / (1 + Math.exp(-r.score)) })) + : merged; + + // Recovery merges sub-batches out of global rank order; sort for a coherent + // final ranking (no-op on the common single-request path). results.sort((a, b) => b.score - a.score); this.rerankBreaker.onSuccess(); return { results, model: rerankModel }; @@ -584,7 +593,20 @@ export function remoteConfigFromEnv(yamlModels?: { const embedApiUrl = process.env.QMD_EMBED_API_URL || yamlModels?.embed_api_url; const embedApiModel = process.env.QMD_EMBED_API_MODEL || yamlModels?.embed_api_model; - if (!embedApiUrl || !embedApiModel) return null; + // Neither set → remote mode not requested; caller uses the local backend. + if (!embedApiUrl && !embedApiModel) return null; + + // Exactly one set → a misconfiguration that would otherwise silently install + // the local backend (and skip the remote pre-flight probe), so indexing would + // quietly run on the wrong embeddings. Fail fast so the operator notices. + if (!embedApiUrl || !embedApiModel) { + const present = embedApiUrl ? "embed_api_url (QMD_EMBED_API_URL)" : "embed_api_model (QMD_EMBED_API_MODEL)"; + const missing = embedApiUrl ? "embed_api_model (QMD_EMBED_API_MODEL)" : "embed_api_url (QMD_EMBED_API_URL)"; + throw new Error( + `Incomplete remote embedding configuration: ${present} is set but ${missing} is missing. ` + + `Set both to use a remote backend, or neither to use the local model.` + ); + } return { embedApiUrl, diff --git a/test/remote-llm.test.ts b/test/remote-llm.test.ts index 87d1cdcc5..cdd49335f 100644 --- a/test/remote-llm.test.ts +++ b/test/remote-llm.test.ts @@ -252,7 +252,7 @@ describe("RemoteLLM", () => { }); describe("rerank", () => { - it("should rerank documents", async () => { + it("sigmoid-normalizes logit-style rerank scores", async () => { setMockHandler((_req, body) => { const parsed = JSON.parse(body); expect(parsed.model).toBe("rerank-model"); @@ -262,16 +262,14 @@ describe("RemoteLLM", () => { status: 200, body: { results: [ - { index: 1, relevance_score: 0.9 }, - { index: 0, relevance_score: 0.3 }, + { index: 1, relevance_score: 2.0 }, // log-odds (outside [0,1]) + { index: 0, relevance_score: -1.0 }, ], }, }; }); - const llm = createRemoteLLM({ - rerankApiModel: "rerank-model", - }); + const llm = createRemoteLLM({ rerankApiModel: "rerank-model" }); const result = await llm.rerank( "test query", [ @@ -280,18 +278,44 @@ describe("RemoteLLM", () => { ] ); + // Scores fall outside [0,1] (log-odds), so they're mapped to a 0..1 + // probability via σ(x)=1/(1+e^-x). σ(2.0)≈0.881, σ(-1.0)≈0.269. Ordering + // (monotonic) is preserved. expect(result.model).toBe("rerank-model"); expect(result.results).toHaveLength(2); - // RemoteLLM.rerank sigmoid-normalizes the raw relevance_score (log-odds) - // into a 0..1 probability: σ(x) = 1/(1+e^-x). σ(0.9) ≈ 0.7109, - // σ(0.3) ≈ 0.5744. Ordering (monotonic) is preserved. const bScore = result.results.find(r => r.file === "b.md")!.score; const aScore = result.results.find(r => r.file === "a.md")!.score; - expect(bScore).toBeCloseTo(1 / (1 + Math.exp(-0.9)), 10); - expect(aScore).toBeCloseTo(1 / (1 + Math.exp(-0.3)), 10); + expect(bScore).toBeCloseTo(1 / (1 + Math.exp(-2.0)), 10); + expect(aScore).toBeCloseTo(1 / (1 + Math.exp(1.0)), 10); // σ(-1.0) expect(bScore).toBeGreaterThan(aScore); }); + it("preserves rerank scores already in [0,1] (no sigmoid distortion)", async () => { + // Cohere/Voyage-style rerankers return probabilities in [0,1] already; + // applying sigmoid would compress them (0.9→0.71) and skew min-score + // filtering, so in-range scores must pass through unchanged. + setMockHandler((_req, body) => { + const parsed = JSON.parse(body); + return { + status: 200, + body: { + results: [ + { index: 0, relevance_score: 0.9 }, + { index: 1, relevance_score: 0.2 }, + ], + }, + }; + }); + + const llm = createRemoteLLM({ rerankApiModel: "rerank-model" }); + const result = await llm.rerank("q", [ + { file: "a.md", text: "doc a" }, + { file: "b.md", text: "doc b" }, + ]); + expect(result.results.find(r => r.file === "a.md")!.score).toBe(0.9); + expect(result.results.find(r => r.file === "b.md")!.score).toBe(0.2); + }); + it("recovers from an oversized rerank batch by splitting", async () => { // log-odds keyed by document text; the server rejects any multi-document // batch as "too large", forcing recursive bisection down to single docs. @@ -338,7 +362,7 @@ describe("RemoteLLM", () => { } return { status: 200, - body: { results: parsed.documents.map((_: string, i: number) => ({ index: i, relevance_score: 1.0 })) }, + body: { results: parsed.documents.map((_: string, i: number) => ({ index: i, relevance_score: 2.0 })) }, }; }); @@ -346,7 +370,7 @@ describe("RemoteLLM", () => { const result = await llm.rerank("q", [{ file: "big.md", text: "x".repeat(500) }]); expect(result.results).toHaveLength(1); expect(result.results[0].file).toBe("big.md"); - expect(result.results[0].score).toBeCloseTo(1 / (1 + Math.exp(-1.0)), 10); + expect(result.results[0].score).toBeCloseTo(1 / (1 + Math.exp(-2.0)), 10); // σ(2.0) }); it("should throw when rerankApiModel not configured", async () => { @@ -525,6 +549,18 @@ describe("remoteConfigFromEnv", () => { expect(config!.embedApiUrl).toBe("http://env:8000/v1"); expect(config!.embedApiModel).toBe("env-model"); }); + + it("throws on incomplete remote config (url without model)", () => { + // Half-configured remote would otherwise silently install the local backend + // and skip the remote pre-flight probe — fail fast instead. + expect(() => remoteConfigFromEnv({ embed_api_url: "http://gpu:8000/v1" })) + .toThrow(/incomplete remote embedding/i); + }); + + it("throws on incomplete remote config (model without url)", () => { + expect(() => remoteConfigFromEnv({ embed_api_model: "bge-m3" })) + .toThrow(/incomplete remote embedding/i); + }); }); // ============================================================================= From 6c0793887165795c1004ebfcaae316b7111b9998 Mon Sep 17 00:00:00 2001 From: Kaspre Date: Tue, 2 Jun 2026 11:56:48 -0400 Subject: [PATCH 4/6] fix(remote): align hybrid backend wiring --- README.md | 9 ++++- src/cli/qmd.ts | 46 +++++++++++---------- src/collections.ts | 6 +++ src/configured-llm.ts | 21 ++++++++++ src/hybrid-llm.ts | 17 ++++++-- src/index.ts | 14 +++---- src/remote-llm.ts | 8 ++-- src/store.ts | 5 ++- test/cli.test.ts | 26 ++++++++++++ test/remote-llm.test.ts | 88 +++++++++++++++++++++++++++++++++++++++- test/sdk.test.ts | 90 +++++++++++++++++++++++++++++++++++++++++ 11 files changed, 287 insertions(+), 43 deletions(-) create mode 100644 src/configured-llm.ts diff --git a/README.md b/README.md index 3a7861553..61542434f 100644 --- a/README.md +++ b/README.md @@ -942,9 +942,9 @@ Uses node-llama-cpp's `createRankingContext()` and `rankAndSort()` API for cross Used for generating query variations via `LlamaChatSession`. -### Remote Embedding & Reranking +### Remote Embedding, Reranking & Query Expansion -QMD can offload embedding and reranking to a remote OpenAI-compatible server (vLLM, Ollama, LM Studio, OpenAI, etc.) while keeping query expansion local. +QMD can offload embedding, reranking, and query expansion to remote OpenAI-compatible servers (vLLM, Ollama, LM Studio, OpenAI, etc.) while keeping local generation and tokenization available for hybrid fallback. **Environment variables** (presence of `QMD_EMBED_API_URL` activates remote mode): @@ -956,6 +956,9 @@ QMD can offload embedding and reranking to a remote OpenAI-compatible server (vL | `QMD_RERANK_API_URL` | No | Rerank endpoint (defaults to embed URL) | | `QMD_RERANK_API_MODEL` | No | Rerank model name | | `QMD_RERANK_API_KEY` | No | Rerank auth (defaults to embed key) | +| `QMD_EXPAND_API_URL` | No | Query expansion chat endpoint (defaults to embed URL) | +| `QMD_EXPAND_API_MODEL` | No | Chat model for query expansion | +| `QMD_EXPAND_API_KEY` | No | Query expansion auth (defaults to embed key) | **YAML config** (`~/.config/qmd/index.yml`): ```yaml @@ -963,6 +966,8 @@ models: embed_api_url: "http://gpu-host:8000/v1" embed_api_model: "BAAI/bge-m3" rerank_api_model: "BAAI/bge-reranker-v2-m3" + expand_api_url: "https://chat-host/v1" + expand_api_model: "qwen3-4b" ``` **Example with vLLM:** diff --git a/src/cli/qmd.ts b/src/cli/qmd.ts index ad159748b..f426512b9 100755 --- a/src/cli/qmd.ts +++ b/src/cli/qmd.ts @@ -33,6 +33,7 @@ import { formatDocForEmbedding, getEmbeddingFingerprint, chunkDocumentByTokens, + chunkDocumentByApproxTokens, clearCache, getCacheKey, getCachedResult, @@ -82,8 +83,7 @@ import { type ChunkStrategy, } from "../store.js"; import { disposeDefaultLlamaCpp, getDefaultLLM, setDefaultLLM, getDefaultLlamaCpp, setDefaultLlamaCpp, LlamaCpp, withLLMSession, pullModels, DEFAULT_MODEL_CACHE_DIR, resolveEmbedModel, resolveGenerateModel, resolveRerankModel, resolveModels, inspectGgufFile, isDarwinMetalMitigationActive } from "../llm.js"; -import { RemoteLLM, remoteConfigFromEnv } from "../remote-llm.js"; -import { HybridLLM } from "../hybrid-llm.js"; +import { createConfiguredLLM } from "../configured-llm.js"; import { formatSearchResults, formatDocuments, @@ -132,25 +132,19 @@ function getStore(): ReturnType { if (!store) { store = createStore(storeDbPathOverride); // Sync YAML config into SQLite store_collections so store.ts reads from DB + const activeModels = ensureModelsConfiguredForCli(); + let config: CollectionConfig | undefined; try { - const activeModels = ensureModelsConfiguredForCli(); - const config = loadConfig(); + config = loadConfig(); syncConfigToDb(store.db, config); - const localLlm = new LlamaCpp({ - embedModel: activeModels.embed, - generateModel: activeModels.generate, - rerankModel: activeModels.rerank, - }); - // Remote embedding/rerank: env vars (QMD_EMBED_API_URL etc) take precedence over YAML models.*_api_* - const remoteConfig = remoteConfigFromEnv(config.models); - if (remoteConfig) { - setDefaultLLM(new HybridLLM(new RemoteLLM(remoteConfig), localLlm)); - } else { - setDefaultLLM(localLlm); - } } catch { // Config may not exist yet — that's fine, DB works without it } + setDefaultLLM(createConfiguredLLM(config?.models, { + embedModel: activeModels.embed, + generateModel: activeModels.generate, + rerankModel: activeModels.rerank, + })); } return store; } @@ -320,7 +314,11 @@ function formatETA(seconds: number): string { // Check index health and print warnings/tips -function checkIndexHealth(db: Database, model: string = resolveEmbedModelForCli()): void { +function getActiveEmbedModelForCli(): string { + return getDefaultLLM().embedModelName; +} + +function checkIndexHealth(db: Database, model: string = getActiveEmbedModelForCli()): void { const { needsEmbedding, totalDocs, daysStale } = getIndexHealth(db, model); // Warn if many docs need embedding @@ -493,7 +491,7 @@ async function showStatus(): Promise { // Overall stats const totalDocs = db.prepare(`SELECT COUNT(*) as count FROM documents WHERE active = 1`).get() as { count: number }; const vectorCount = db.prepare(`SELECT COUNT(*) as count FROM content_vectors`).get() as { count: number }; - const statusEmbedModel = resolveEmbedModelForCli(); + const statusEmbedModel = getActiveEmbedModelForCli(); const needsEmbedding = getHashesNeedingEmbedding(db, undefined, statusEmbedModel); // Most recent update across all collections @@ -750,7 +748,7 @@ async function updateCollections(): Promise { } // Check if any documents need embedding (show once at end) - const needsEmbedding = getHashesNeedingEmbedding(db); + const needsEmbedding = getHashesNeedingEmbedding(db, undefined, getActiveEmbedModelForCli()); closeDb(); console.log(`${c.green}✓ All collections updated.${c.reset}`); @@ -1909,7 +1907,7 @@ async function indexFiles(pwd?: string, globPattern: string = DEFAULT_GLOB, coll const orphanedContent = cleanupOrphanedContent(db); // Check if vector index needs updating - const needsEmbedding = getHashesNeedingEmbedding(db); + const needsEmbedding = getHashesNeedingEmbedding(db, undefined, getActiveEmbedModelForCli()); progress.clear(); console.log(`\nIndexed: ${indexed} new, ${updated} updated, ${unchanged} unchanged, ${removed} removed`); @@ -3795,11 +3793,15 @@ async function checkEmbeddingVectorSamples(db: Database, model: string, fingerpr const threshold = 0.0001; const mismatches: string[] = []; + const llm = getDefaultLLM(); + const usesRemoteEmbedding = llm.usesRemoteEmbedding === true; await withLLMSession(async (session) => { for (const sample of samples) { const hashSeq = `${sample.hash}_${sample.seq}`; - const chunks = await chunkDocumentByTokens(sample.body, undefined, undefined, undefined, sample.path, undefined, session.signal); + const chunks = usesRemoteEmbedding + ? await chunkDocumentByApproxTokens(sample.body, undefined, undefined, undefined, sample.path, undefined, session.signal) + : await chunkDocumentByTokens(sample.body, undefined, undefined, undefined, sample.path, undefined, session.signal); const chunk = chunks[sample.seq]; if (!chunk) { mismatches.push(`${shortHashSeq(hashSeq)}: chunk no longer exists`); @@ -3984,7 +3986,7 @@ async function showDoctor(): Promise { const db = storeInstance.db; const pkg = readPackageJson(); const activeModels = resolveModelsForCli(); - const embedModel = activeModels.embed; + const embedModel = getActiveEmbedModelForCli(); const fingerprint = getEmbeddingFingerprint(embedModel); const nextSteps: string[] = []; diff --git a/src/collections.ts b/src/collections.ts index 768a7eeb7..2cdc9496e 100644 --- a/src/collections.ts +++ b/src/collections.ts @@ -52,6 +52,12 @@ export interface ModelsConfig { rerank_api_model?: string; /** Bearer token for remote rerank API */ rerank_api_key?: string; + /** Remote query expansion API base URL */ + expand_api_url?: string; + /** Remote query expansion chat model name */ + expand_api_model?: string; + /** Bearer token for remote query expansion API */ + expand_api_key?: string; } /** diff --git a/src/configured-llm.ts b/src/configured-llm.ts new file mode 100644 index 000000000..44f1d72de --- /dev/null +++ b/src/configured-llm.ts @@ -0,0 +1,21 @@ +import type { ModelsConfig } from "./collections.js"; +import { HybridLLM } from "./hybrid-llm.js"; +import { LlamaCpp, type LLM, type LlamaCppConfig } from "./llm.js"; +import { RemoteLLM, remoteConfigFromEnv } from "./remote-llm.js"; + +/** + * Build the LLM backend implied by config/env. + * + * Remote embedding is opt-in via remoteConfigFromEnv(). When configured, a + * HybridLLM keeps local generation/tokenization/fallback behavior while routing + * remote-capable operations through the OpenAI-compatible API. + */ +export function createConfiguredLLM( + models?: ModelsConfig, + localConfig: LlamaCppConfig = {}, +): LLM { + const remoteConfig = remoteConfigFromEnv(models); + const local = new LlamaCpp(localConfig); + if (!remoteConfig) return local; + return new HybridLLM(new RemoteLLM(remoteConfig), local); +} diff --git a/src/hybrid-llm.ts b/src/hybrid-llm.ts index be37bf656..c9efb5ee2 100644 --- a/src/hybrid-llm.ts +++ b/src/hybrid-llm.ts @@ -55,13 +55,18 @@ export class HybridLLM implements LLM { return this.remote.embedBatch(texts, options); } - rerank(query: string, documents: RerankDocument[], options?: RerankOptions): Promise { + async rerank(query: string, documents: RerankDocument[], options?: RerankOptions): Promise { // When remote is a RemoteLLM without a rerank model configured, fall back to local rerank // (same fallback shape as expandQuery → local). if (this.remote instanceof RemoteLLM && !this.remote.supportsRerank) { return this.local.rerank(query, documents, options); } - return this.remote.rerank(query, documents, options); + try { + return await this.remote.rerank(query, documents, options); + } catch (error) { + console.error("Remote rerank failed; falling back to local rerank:", error); + return this.local.rerank(query, documents, options); + } } // Route to local @@ -77,11 +82,15 @@ export class HybridLLM implements LLM { return this.local.detokenize(tokens); } - expandQuery(query: string, options?: { context?: string; includeLexical?: boolean; intent?: string }): Promise { + async expandQuery(query: string, options?: { context?: string; includeLexical?: boolean; intent?: string }): Promise { // Route to remote when configured for it; otherwise local (same fallback // shape as rerank → local when remote doesn't support rerank). if (this.remote instanceof RemoteLLM && this.remote.supportsExpand) { - return this.remote.expandQuery(query, options); + try { + return await this.remote.expandQuery(query, options); + } catch (error) { + console.error("Remote query expansion failed; falling back to local expansion:", error); + } } return this.local.expandQuery(query, options); } diff --git a/src/index.ts b/src/index.ts index f853a974c..5bae86109 100644 --- a/src/index.ts +++ b/src/index.ts @@ -63,9 +63,7 @@ import { type EmbedResult, type ChunkStrategy, } from "./store.js"; -import { - LlamaCpp, -} from "./llm.js"; +import { createConfiguredLLM } from "./configured-llm.js"; import { setConfigSource, loadConfig, @@ -211,8 +209,8 @@ export interface StoreOptions { * The QMD SDK store — provides search, retrieval, collection management, * context management, and indexing operations. * - * All methods are async. The store manages its own LlamaCpp instance - * (lazy-loaded, auto-unloaded after inactivity) — no global singletons. + * All methods are async. The store manages its own LLM instance + * (lazy-loaded, auto-unloaded after inactivity for local models) — no global singletons. */ export interface QMDStore { /** The underlying internal store (for advanced use) */ @@ -368,9 +366,9 @@ export async function createStore(options: StoreOptions): Promise { } // else: DB-only mode — no external config, use existing store_collections - // Create a per-store LlamaCpp instance — lazy-loads models on first use, - // auto-unloads after 5 min inactivity to free VRAM. - const llm = new LlamaCpp({ + // Create a per-store LLM instance — local-only by default, HybridLLM when + // remote embedding is configured. + const llm = createConfiguredLLM(config?.models, { embedModel: config?.models?.embed, generateModel: config?.models?.generate, rerankModel: config?.models?.rerank, diff --git a/src/remote-llm.ts b/src/remote-llm.ts index 504c1c5a3..e83af1fd4 100644 --- a/src/remote-llm.ts +++ b/src/remote-llm.ts @@ -497,11 +497,9 @@ export class RemoteLLM implements LLM { }; content = json.choices?.[0]?.message?.content ?? ""; } catch (err) { - // Network error, timeout, or non-2xx — fall back to the default triple. - // Don't escalate to caller; LocalLLM also masks failures behind the - // fallback to keep search resilient. - console.error("Remote query expansion failed:", err); - return defaultFallback(); + // Network error, timeout, or non-2xx. Let HybridLLM fall back to local + // query expansion when available; bare RemoteLLM callers see the failure. + throw err; } // Parse — mirror LocalLLM's parsing exactly so downstream sees consistent diff --git a/src/store.ts b/src/store.ts index 4ebbd7000..5c932a5ad 100644 --- a/src/store.ts +++ b/src/store.ts @@ -2195,9 +2195,12 @@ export async function maybeAdoptLegacyEmbeddingFingerprint(store: Store, model: const expectedHashSeq = `${sample.hash}_${sample.seq}`; const title = extractTitle(sample.body, sample.path); const llm = getLlm(store); + const usesRemoteEmbedding = llm.usesRemoteEmbedding === true; return await withLLMSessionForLlm(llm, async (session) => { - const chunks = await chunkDocumentByTokens(sample.body, undefined, undefined, undefined, sample.path, undefined, session.signal); + const chunks = usesRemoteEmbedding + ? await chunkDocumentByApproxTokens(sample.body, undefined, undefined, undefined, sample.path, undefined, session.signal) + : await chunkDocumentByTokens(sample.body, undefined, undefined, undefined, sample.path, undefined, session.signal); const chunk = chunks[sample.seq]; if (!chunk) { return { checked: true, adopted: 0, reason: `sample chunk ${expectedHashSeq} no longer exists` }; diff --git a/test/cli.test.ts b/test/cli.test.ts index 5f4e13827..f10f7cfff 100644 --- a/test/cli.test.ts +++ b/test/cli.test.ts @@ -553,6 +553,32 @@ describe("CLI Status Command", () => { expect(stdout).toContain("fix the YAML"); }, 20000); + test("qmd status surfaces incomplete remote config instead of falling back locally", async () => { + const env = await createIsolatedTestEnv("status-incomplete-remote-config"); + await writeFile(join(env.configDir, "index.yml"), [ + "collections: {}", + "models:", + " embed_api_url: http://remote.example/v1", + "", + ].join("\n")); + + const { stderr, exitCode } = await runQmd(["status"], { + dbPath: env.dbPath, + configDir: env.configDir, + env: { + QMD_EMBED_API_URL: "", + QMD_EMBED_API_MODEL: "", + QMD_RERANK_API_URL: "", + QMD_RERANK_API_MODEL: "", + QMD_EXPAND_API_URL: "", + QMD_EXPAND_API_MODEL: "", + }, + }); + + expect(exitCode).not.toBe(0); + expect(stderr).toMatch(/incomplete remote embedding/i); + }, 20000); + test("qmd doctor warns when configured models differ from code defaults", async () => { const env = await createIsolatedTestEnv("doctor-custom-models"); await writeFile(join(env.configDir, "index.yml"), `collections: {}\nmodels:\n embed: hf:example/custom-embed/custom.gguf\n generate: ${DEFAULT_GENERATE_MODEL_URI}\n rerank: ${DEFAULT_RERANK_MODEL_URI}\n`); diff --git a/test/remote-llm.test.ts b/test/remote-llm.test.ts index cdd49335f..bab942e71 100644 --- a/test/remote-llm.test.ts +++ b/test/remote-llm.test.ts @@ -6,6 +6,7 @@ import { describe, it, expect, beforeAll, afterAll, beforeEach, afterEach } from "vitest"; import { createServer, type Server, type IncomingMessage, type ServerResponse } from "http"; +import { createConfiguredLLM } from "../src/configured-llm.js"; import { RemoteLLM, remoteConfigFromEnv, type RemoteLLMConfig } from "../src/remote-llm.js"; import { HybridLLM } from "../src/hybrid-llm.js"; import { isRemoteModel, formatQueryForEmbedding, formatDocForEmbedding, getDefaultLLM, setDefaultLLM, LlamaCpp } from "../src/llm.js"; @@ -70,6 +71,23 @@ function createRemoteLLM(overrides?: Partial): RemoteLLM { }); } +function withRemoteEnvCleared(fn: () => T): T { + const saved: Record = {}; + for (const key of Object.keys(process.env)) { + if (key.startsWith("QMD_") && key.includes("API")) { + saved[key] = process.env[key]; + delete process.env[key]; + } + } + try { + return fn(); + } finally { + for (const [key, value] of Object.entries(saved)) { + if (value !== undefined) process.env[key] = value; + } + } +} + // ============================================================================= // RemoteLLM Tests // ============================================================================= @@ -409,14 +427,22 @@ describe("RemoteLLM", () => { describe("HybridLLM", () => { // Simple mock local LLM function createMockLocalLLM(): LLM { + const localTokens = [1, 2, 3] as any; return { embedModelName: "local-model", + generateModelName: "local-generate-model", + rerankModelName: "local-rerank-model", embed: async () => ({ embedding: [0.5], model: "local-model" }), embedBatch: async (texts) => texts.map(() => ({ embedding: [0.5], model: "local-model" })), generate: async () => ({ text: "expanded", model: "local-model", done: true }), modelExists: async (model) => ({ name: model, exists: true }), expandQuery: async () => [{ type: "lex" as const, text: "expanded query" }], - rerank: async () => ({ results: [], model: "local-model" }), + rerank: async (_query, documents) => ({ + results: documents.map((doc, index) => ({ file: doc.file, score: 0.42, index })), + model: "local-rerank-model", + }), + tokenize: async () => localTokens, + detokenize: async () => "detokenized", dispose: async () => {}, }; } @@ -478,6 +504,35 @@ describe("HybridLLM", () => { expect(result[0]!.text).toBe("expanded query"); }); + it("falls back to local rerank when configured remote rerank fails", async () => { + setMockHandler(() => ({ + status: 500, + body: { error: "rerank down" }, + })); + + const remote = createRemoteLLM({ rerankApiModel: "remote-rerank" }); + const local = createMockLocalLLM(); + const hybrid = new HybridLLM(remote, local); + + const result = await hybrid.rerank("query", [{ file: "doc.md", text: "doc text" }]); + expect(result.model).toBe("local-rerank-model"); + expect(result.results).toEqual([{ file: "doc.md", score: 0.42, index: 0 }]); + }); + + it("falls back to local expansion when configured remote expansion fails", async () => { + setMockHandler(() => ({ + status: 500, + body: { error: "chat down" }, + })); + + const remote = createRemoteLLM({ expandApiModel: "remote-chat" }); + const local = createMockLocalLLM(); + const hybrid = new HybridLLM(remote, local); + + const result = await hybrid.expandQuery("test query"); + expect(result).toEqual([{ type: "lex", text: "expanded query" }]); + }); + it("should use remote embedModelName", async () => { const remote = createRemoteLLM({ embedApiModel: "BAAI/bge-m3" }); const local = createMockLocalLLM(); @@ -487,6 +542,37 @@ describe("HybridLLM", () => { }); }); +describe("createConfiguredLLM", () => { + it("returns local LlamaCpp when remote config is absent", () => { + const llm = withRemoteEnvCleared(() => createConfiguredLLM(undefined, { + embedModel: "hf:local/embed.gguf", + generateModel: "hf:local/generate.gguf", + rerankModel: "hf:local/rerank.gguf", + })); + + expect(llm).toBeInstanceOf(LlamaCpp); + expect(llm.embedModelName).toBe("hf:local/embed.gguf"); + }); + + it("returns HybridLLM when remote embedding is configured", () => { + const llm = withRemoteEnvCleared(() => createConfiguredLLM({ + embed_api_url: baseUrl(), + embed_api_model: "remote-embed", + rerank_api_model: "remote-rerank", + expand_api_model: "remote-chat", + })); + + expect(llm).toBeInstanceOf(HybridLLM); + expect(llm.embedModelName).toBe("remote-embed"); + expect(llm.usesRemoteEmbedding).toBe(true); + }); + + it("throws instead of falling back locally when remote config is incomplete", () => { + expect(() => withRemoteEnvCleared(() => createConfiguredLLM({ embed_api_url: baseUrl() }))) + .toThrow(/incomplete remote embedding/i); + }); +}); + // ============================================================================= // Config Tests // ============================================================================= diff --git a/test/sdk.test.ts b/test/sdk.test.ts index 53764c560..fc28736a4 100644 --- a/test/sdk.test.ts +++ b/test/sdk.test.ts @@ -61,6 +61,24 @@ function freshDbPath(): string { return join(testDir, `test-${Date.now()}-${Math.random().toString(36).slice(2)}.sqlite`); } +async function withRemoteEnvCleared(fn: () => Promise): Promise { + const saved: Record = {}; + for (const key of Object.keys(process.env)) { + if (key.startsWith("QMD_") && key.includes("API")) { + saved[key] = process.env[key]; + delete process.env[key]; + } + } + try { + return await fn(); + } finally { + for (const [key, value] of Object.entries(saved)) { + if (value === undefined) delete process.env[key]; + else process.env[key] = value; + } + } +} + // ============================================================================= // Constructor Tests // ============================================================================= @@ -82,6 +100,78 @@ describe("createStore", () => { await store.close(); }); + test("creates hybrid remote LLM from inline config", async () => { + await withRemoteEnvCleared(async () => { + const store = await createStore({ + dbPath: freshDbPath(), + config: { + collections: {}, + models: { + embed_api_url: "http://remote.example/v1", + embed_api_model: "remote-embed", + rerank_api_model: "remote-rerank", + expand_api_model: "remote-chat", + }, + }, + }); + + try { + expect(store.internal.llm?.embedModelName).toBe("remote-embed"); + expect(store.internal.llm?.usesRemoteEmbedding).toBe(true); + expect(store.internal.llm?.rerankModelName).toBe("remote-rerank"); + } finally { + await store.close(); + } + }); + }); + + test("creates hybrid remote LLM from YAML config path used by MCP startup", async () => { + await withRemoteEnvCleared(async () => { + const configPath = join(testDir, "test-remote-mcp-config.yml"); + const config: CollectionConfig = { + collections: {}, + models: { + embed_api_url: "http://remote.example/v1", + embed_api_model: "remote-embed", + rerank_api_model: "remote-rerank", + expand_api_model: "remote-chat", + }, + }; + writeFileSync(configPath, YAML.stringify(config)); + + const store = await createStore({ + dbPath: freshDbPath(), + configPath, + }); + + try { + expect(store.internal.llm?.embedModelName).toBe("remote-embed"); + expect(store.internal.llm?.usesRemoteEmbedding).toBe(true); + expect(store.internal.llm?.rerankModelName).toBe("remote-rerank"); + } finally { + await store.close(); + } + }); + }); + + test("rejects incomplete remote YAML config instead of falling back locally", async () => { + await withRemoteEnvCleared(async () => { + const configPath = join(testDir, "test-incomplete-remote-config.yml"); + const config: CollectionConfig = { + collections: {}, + models: { + embed_api_url: "http://remote.example/v1", + }, + }; + writeFileSync(configPath, YAML.stringify(config)); + + await expect(createStore({ + dbPath: freshDbPath(), + configPath, + })).rejects.toThrow(/incomplete remote embedding/i); + }); + }); + test("creates store with YAML config file", async () => { const configPath = join(testDir, "test-config.yml"); const config: CollectionConfig = { From 46539d812f9aa8299d9ab7e7718c4a4f28e5ecbc Mon Sep 17 00:00:00 2001 From: Kaspre Date: Tue, 2 Jun 2026 13:06:37 -0400 Subject: [PATCH 5/6] fix(remote): harden expansion cache and breaker --- src/hybrid-llm.ts | 13 ++++- src/llm.ts | 27 +++++++-- src/remote-llm.ts | 51 +++++++++++++---- src/store.ts | 15 ++++- test/remote-llm.test.ts | 118 ++++++++++++++++++++++++++++++++++++++++ test/store.test.ts | 74 ++++++++++++++++++++++++- 6 files changed, 278 insertions(+), 20 deletions(-) diff --git a/src/hybrid-llm.ts b/src/hybrid-llm.ts index c9efb5ee2..1134cfd2c 100644 --- a/src/hybrid-llm.ts +++ b/src/hybrid-llm.ts @@ -2,7 +2,8 @@ * hybrid-llm.ts - Compositor that routes LLM operations between remote and local backends * * Embed/rerank → remote (GPU-heavy, benefits from offloading) - * Generate/expandQuery → local LlamaCpp (QMD's fine-tuned query expansion model) + * Generate → local LlamaCpp + * expandQuery → remote when configured, otherwise local * tokenize/countTokens → local LlamaCpp (CPU-cheap, needed for chunking) */ @@ -12,6 +13,7 @@ import type { EmbeddingResult, GenerateOptions, GenerateResult, + LLMExpandQueryOptions, ModelInfo, Queryable, RerankDocument, @@ -35,6 +37,13 @@ export class HybridLLM implements LLM { return this.local.generateModelName; } + get expandModelName(): string { + if (this.remote instanceof RemoteLLM && this.remote.supportsExpand) { + return this.remote.expandModelName ?? this.remote.generateModelName; + } + return this.local.expandModelName ?? this.local.generateModelName; + } + get rerankModelName(): string { if (this.remote instanceof RemoteLLM && !this.remote.supportsRerank) { return this.local.rerankModelName; @@ -82,7 +91,7 @@ export class HybridLLM implements LLM { return this.local.detokenize(tokens); } - async expandQuery(query: string, options?: { context?: string; includeLexical?: boolean; intent?: string }): Promise { + async expandQuery(query: string, options?: LLMExpandQueryOptions): Promise { // Route to remote when configured for it; otherwise local (same fallback // shape as rerank → local when remote doesn't support rerank). if (this.remote instanceof RemoteLLM && this.remote.supportsExpand) { diff --git a/src/llm.ts b/src/llm.ts index 597d0e976..e479023b9 100644 --- a/src/llm.ts +++ b/src/llm.ts @@ -225,7 +225,7 @@ export type LLMSessionOptions = { export interface ILLMSession { embed(text: string, options?: EmbedOptions): Promise; embedBatch(texts: string[], options?: EmbedOptions): Promise<(EmbeddingResult | null)[]>; - expandQuery(query: string, options?: { context?: string; includeLexical?: boolean }): Promise; + expandQuery(query: string, options?: LLMExpandQueryOptions): Promise; rerank(query: string, documents: RerankDocument[], options?: RerankOptions): Promise; /** Whether this session is still valid (not released or aborted) */ readonly isValid: boolean; @@ -246,6 +246,14 @@ export type Queryable = { text: string; }; +export type LLMExpandQueryOptions = { + context?: string; + includeLexical?: boolean; + intent?: string; + /** Reports the model that actually produced the returned expansion. */ + onModelUsed?: (model: string) => void; +}; + /** * Document to rerank */ @@ -552,6 +560,12 @@ export interface LLM { */ readonly generateModelName: string; + /** + * The query-expansion model name/URI. Defaults to generateModelName for + * local backends, but remote chat expansion can use an independent model. + */ + readonly expandModelName?: string; + /** * The reranking model name/URI. */ @@ -588,7 +602,7 @@ export interface LLM { * Expand a search query into multiple variations for different backends. * Returns a list of Queryable objects. */ - expandQuery(query: string, options?: { context?: string; includeLexical?: boolean; intent?: string }): Promise; + expandQuery(query: string, options?: LLMExpandQueryOptions): Promise; /** * Rerank documents by relevance to a query @@ -793,6 +807,10 @@ export class LlamaCpp implements LLM { return this.generateModelUri; } + get expandModelName(): string { + return this.generateModelUri; + } + get rerankModelName(): string { return this.rerankModelUri; } @@ -1481,8 +1499,9 @@ export class LlamaCpp implements LLM { // High-level abstractions // ========================================================================== - async expandQuery(query: string, options: { context?: string, includeLexical?: boolean, intent?: string } = {}): Promise { + async expandQuery(query: string, options: LLMExpandQueryOptions = {}): Promise { if (this._ciMode) throw new Error("LLM operations are disabled in CI (set CI=true)"); + options.onModelUsed?.(this.expandModelName); // Ping activity at start to keep models alive during this operation this.touchActivity(); @@ -1916,7 +1935,7 @@ class LLMSession implements ILLMSession { async expandQuery( query: string, - options?: { context?: string; includeLexical?: boolean } + options?: LLMExpandQueryOptions ): Promise { return this.withOperation(() => this.manager.getLLM().expandQuery(query, options)); } diff --git a/src/remote-llm.ts b/src/remote-llm.ts index e83af1fd4..7265bb727 100644 --- a/src/remote-llm.ts +++ b/src/remote-llm.ts @@ -2,7 +2,8 @@ * remote-llm.ts - OpenAI-compatible remote embedding & reranking backend * * Implements the LLM interface by calling HTTP endpoints (vLLM, Ollama, OpenAI, etc.). - * Only supports embed/rerank operations — generate/expandQuery throw. + * Supports embedding, optional reranking, and optional chat-based query expansion. + * Text generation remains local-only through HybridLLM. */ import type { @@ -12,6 +13,7 @@ import type { GenerateOptions, GenerateResult, ModelInfo, + LLMExpandQueryOptions, Queryable, QueryType, RerankDocument, @@ -65,6 +67,7 @@ class CircuitBreaker { private state: CircuitState = "closed"; private failures = 0; private lastFailureTime = 0; + private halfOpenProbeInFlight = false; private readonly maxFailures: number; private readonly cooldownMs: number; @@ -78,22 +81,29 @@ class CircuitBreaker { if (this.state === "open") { if (Date.now() - this.lastFailureTime >= this.cooldownMs) { this.state = "half-open"; + this.halfOpenProbeInFlight = true; return true; } return false; } - // half-open: allow one attempt - return true; + // half-open: allow one probe at a time + if (!this.halfOpenProbeInFlight) { + this.halfOpenProbeInFlight = true; + return true; + } + return false; } onSuccess(): void { this.state = "closed"; this.failures = 0; + this.halfOpenProbeInFlight = false; } onFailure(): void { this.failures++; this.lastFailureTime = Date.now(); + this.halfOpenProbeInFlight = false; if (this.state === "half-open" || this.failures >= this.maxFailures) { this.state = "open"; } @@ -135,6 +145,7 @@ export class RemoteLLM implements LLM { private readonly embedBreaker = new CircuitBreaker(); private readonly rerankBreaker = new CircuitBreaker(); + private readonly expandBreaker = new CircuitBreaker(); private expectedDimensions: number | null = null; constructor(config: RemoteLLMConfig) { @@ -161,6 +172,10 @@ export class RemoteLLM implements LLM { return this.config.embedApiModel; } + get expandModelName(): string { + return this.config.expandApiModel || this.config.embedApiModel; + } + get usesRemoteEmbedding(): boolean { return true; } @@ -418,7 +433,7 @@ export class RemoteLLM implements LLM { async expandQuery( query: string, - options?: { context?: string; includeLexical?: boolean; intent?: string }, + options?: LLMExpandQueryOptions, ): Promise { const expandUrl = this.config.expandApiUrl || this.config.embedApiUrl; const expandModel = this.config.expandApiModel; @@ -426,9 +441,9 @@ export class RemoteLLM implements LLM { const includeLexical = options?.includeLexical ?? true; const intent = options?.intent; - // Shared fallback shape — used whenever the remote call fails OR returns - // nothing parseable. Mirrors LocalLLM.expandQuery's fallback exactly so - // downstream code doesn't see a behavior difference. + // Shared fallback shape for RemoteLLM without a configured expand model. + // When remote expansion is configured but unavailable or unusable, throw + // so HybridLLM can fall back to local expansion instead. const defaultFallback = (): Queryable[] => { const triple: Queryable[] = [ { type: "hyde", text: `Information about ${query}` }, @@ -440,13 +455,22 @@ export class RemoteLLM implements LLM { if (!expandModel) { // Configured to use remote but no expand model set → safe default. + options?.onModelUsed?.(this.expandModelName); return defaultFallback(); } + if (!this.expandBreaker.canAttempt()) { + throw new Error( + `Remote expand circuit breaker is open — endpoint ${expandUrl} is unavailable. ` + + `Will retry after cooldown.` + ); + } + // Prompt the chat model to emit the lex/vec/hyde format that // LocalLLM.expandQuery produces via grammar-constrained sampling. // Without llama.cpp grammar we have to ask politely; parsing below is - // tolerant of variations and falls back if the model goes off-script. + // tolerant of small variations but rejects unusable output so HybridLLM + // can fall back to local expansion. const systemPrompt = "You expand search queries for a hybrid retrieval system. " + "Output 3 to 6 query variants, one per line, each prefixed with its type. " + @@ -497,6 +521,7 @@ export class RemoteLLM implements LLM { }; content = json.choices?.[0]?.message?.content ?? ""; } catch (err) { + this.expandBreaker.onFailure(); // Network error, timeout, or non-2xx. Let HybridLLM fall back to local // query expansion when available; bare RemoteLLM callers see the failure. throw err; @@ -534,8 +559,14 @@ export class RemoteLLM implements LLM { ? queryables : queryables.filter(q => q.type !== "lex"); - if (filtered.length > 0) return filtered; - return defaultFallback(); + if (filtered.length > 0) { + this.expandBreaker.onSuccess(); + options?.onModelUsed?.(this.expandModelName); + return filtered; + } + + this.expandBreaker.onFailure(); + throw new Error("Expand API returned no parseable query expansions"); } async dispose(): Promise { diff --git a/src/store.ts b/src/store.ts index 5c932a5ad..e4a3cac67 100644 --- a/src/store.ts +++ b/src/store.ts @@ -1902,7 +1902,7 @@ export function createStore(dbPath?: string): Store { searchVec: (query: string, model: string, limit?: number, collectionName?: string, session?: ILLMSession, precomputedEmbedding?: number[]) => searchVec(db, query, model, limit, collectionName, session, precomputedEmbedding), // Query expansion & reranking - expandQuery: (query: string, model?: string, intent?: string) => expandQuery(query, model ?? store.llm?.generateModelName ?? DEFAULT_QUERY_MODEL, db, intent, store.llm), + expandQuery: (query: string, model?: string, intent?: string) => expandQuery(query, model ?? store.llm?.expandModelName ?? store.llm?.generateModelName ?? DEFAULT_QUERY_MODEL, db, intent, store.llm), rerank: (query: string, documents: { file: string; text: string }[], model?: string, intent?: string) => rerank(query, documents, model ?? store.llm?.rerankModelName ?? DEFAULT_RERANK_MODEL, db, intent, store.llm), // Document retrieval @@ -3843,8 +3843,14 @@ export async function expandQuery(query: string, model: string = DEFAULT_QUERY_M } const llm = llmOverride ?? getDefaultLLM(); + let modelUsed = model; // Note: LlamaCpp uses hardcoded model, model parameter is ignored - const results = await llm.expandQuery(query, { intent }); + const results = await llm.expandQuery(query, { + intent, + onModelUsed: (usedModel) => { + modelUsed = usedModel; + }, + }); // Map Queryable[] → ExpandedQuery[] (same shape, decoupled from llm.ts internals). // Filter out entries that duplicate the original query text. @@ -3853,7 +3859,10 @@ export async function expandQuery(query: string, model: string = DEFAULT_QUERY_M .map(r => ({ type: r.type, query: r.text })); if (expanded.length > 0) { - setCachedResult(db, cacheKey, JSON.stringify(expanded)); + const writeCacheKey = modelUsed === model + ? cacheKey + : getCacheKey("expandQuery", { query, model: modelUsed, ...(intent && { intent }) }); + setCachedResult(db, writeCacheKey, JSON.stringify(expanded)); } return expanded; diff --git a/test/remote-llm.test.ts b/test/remote-llm.test.ts index bab942e71..17716dcb3 100644 --- a/test/remote-llm.test.ts +++ b/test/remote-llm.test.ts @@ -405,6 +405,91 @@ describe("RemoteLLM", () => { await expect(llm.generate("prompt")).rejects.toThrow("does not support text generation"); }); + it("should expand queries through chat completions when configured", async () => { + setMockHandler((_req, body) => { + const parsed = JSON.parse(body); + expect(parsed.model).toBe("remote-chat"); + expect(parsed.messages.at(-1).content).toContain("search docs"); + return { + status: 200, + body: { + choices: [{ + message: { + content: [ + "lex: search docs keywords", + "vec: semantic search docs", + "hyde: Information about search docs", + ].join("\n"), + }, + }], + }, + }; + }); + + const llm = createRemoteLLM({ expandApiModel: "remote-chat" }); + const result = await llm.expandQuery("search docs"); + expect(result).toEqual([ + { type: "lex", text: "search docs keywords" }, + { type: "vec", text: "semantic search docs" }, + { type: "hyde", text: "Information about search docs" }, + ]); + }); + + it("opens an independent circuit breaker after repeated expansion failures", async () => { + let requestCount = 0; + setMockHandler(() => { + requestCount++; + return { + status: 500, + body: { error: "chat down" }, + }; + }); + + const llm = createRemoteLLM({ expandApiModel: "remote-chat" }); + for (let i = 0; i < 3; i++) { + await expect(llm.expandQuery("test query")).rejects.toThrow("500"); + } + + await expect(llm.expandQuery("test query")).rejects.toThrow("circuit breaker"); + expect(requestCount).toBe(3); + + setMockHandler((req) => { + expect(req.url).toContain("/embeddings"); + return { + status: 200, + body: { data: [{ embedding: [0.4], index: 0 }] }, + }; + }); + + const embedResult = await llm.embed("embed still works"); + expect(embedResult?.embedding).toEqual([0.4]); + }); + + it("treats unparseable expansion responses as failures", async () => { + let requestCount = 0; + setMockHandler(() => { + requestCount++; + return { + status: 200, + body: { + choices: [{ + message: { + content: "I cannot help with that.", + }, + }], + }, + }; + }); + + const llm = createRemoteLLM({ expandApiModel: "remote-chat" }); + for (let i = 0; i < 3; i++) { + await expect(llm.expandQuery("test query")).rejects.toThrow("no parseable"); + } + + await expect(llm.expandQuery("test query")).rejects.toThrow("circuit breaker"); + expect(requestCount).toBe(3); + }); + it("should fall back to the default expansion triple when no expand model is configured", async () => { // RemoteLLM.expandQuery now implements query expansion via chat // completions. When no expand model is configured it returns the same @@ -533,6 +618,26 @@ describe("HybridLLM", () => { expect(result).toEqual([{ type: "lex", text: "expanded query" }]); }); + it("falls back to local expansion when configured remote expansion is unparseable", async () => { + setMockHandler(() => ({ + status: 200, + body: { + choices: [{ + message: { + content: "Here are some ideas without the required prefixes.", + }, + }], + }, + })); + + const remote = createRemoteLLM({ expandApiModel: "remote-chat" }); + const local = createMockLocalLLM(); + const hybrid = new HybridLLM(remote, local); + + const result = await hybrid.expandQuery("test query"); + expect(result).toEqual([{ type: "lex", text: "expanded query" }]); + }); + it("should use remote embedModelName", async () => { const remote = createRemoteLLM({ embedApiModel: "BAAI/bge-m3" }); const local = createMockLocalLLM(); @@ -540,6 +645,19 @@ describe("HybridLLM", () => { expect(hybrid.embedModelName).toBe("BAAI/bge-m3"); }); + + it("should use remote expandModelName only when remote expansion is configured", async () => { + const local = createMockLocalLLM(); + + const embedOnlyHybrid = new HybridLLM(createRemoteLLM(), local); + expect(embedOnlyHybrid.expandModelName).toBe("local-generate-model"); + + const remoteExpandHybrid = new HybridLLM( + createRemoteLLM({ expandApiModel: "remote-chat" }), + local, + ); + expect(remoteExpandHybrid.expandModelName).toBe("remote-chat"); + }); }); describe("createConfiguredLLM", () => { diff --git a/test/store.test.ts b/test/store.test.ts index b080fc625..d7ed81da0 100644 --- a/test/store.test.ts +++ b/test/store.test.ts @@ -14,7 +14,7 @@ import { tmpdir } from "node:os"; import { join } from "node:path"; import YAML from "yaml"; import * as llmModule from "../src/llm.js"; -import { disposeDefaultLlamaCpp, setDefaultLlamaCpp } from "../src/llm.js"; +import { disposeDefaultLlamaCpp, setDefaultLlamaCpp, type LLMExpandQueryOptions } from "../src/llm.js"; import { createStore, verifySqliteVecLoaded, @@ -1151,6 +1151,78 @@ describe("Caching", () => { await cleanupTestDb(store); }); + + test("expandQuery cache separates local and remote expansion models", async () => { + const store = await createTestStore(); + const query = "same expansion query"; + const calls: string[] = []; + + const makeLlm = (label: string, expandModelName?: string, modelUsed = expandModelName ?? "local-generate") => ({ + embedModelName: "remote-embed", + generateModelName: "local-generate", + ...(expandModelName ? { expandModelName } : {}), + rerankModelName: "remote-rerank", + embed: async () => ({ embedding: [0.1], model: "remote-embed" }), + embedBatch: async (texts: string[]) => texts.map(() => ({ embedding: [0.1], model: "remote-embed" })), + generate: async () => ({ text: label, model: "local-generate", done: true }), + modelExists: async (model: string) => ({ name: model, exists: true }), + expandQuery: async (_query: string, options?: LLMExpandQueryOptions) => { + calls.push(label); + options?.onModelUsed?.(modelUsed); + return [{ type: "lex" as const, text: `${query} ${label}` }]; + }, + rerank: async () => ({ results: [], model: "remote-rerank" }), + tokenize: async () => [] as any, + detokenize: async () => "", + dispose: async () => {}, + }); + + // Simulate an existing local-generation cache entry keyed by generateModelName. + store.llm = makeLlm("local") as any; + const local = await store.expandQuery(query); + expect(calls).toEqual(["local"]); + + // LLMs without expandModelName remain compatible by falling back to + // generateModelName for the cache identity. + store.llm = makeLlm("local-again") as any; + const localAgain = await store.expandQuery(query); + expect(calls).toEqual(["local"]); + expect(localAgain).toEqual(local); + + // Remote expansion keeps local generation for generate(), but must use a + // distinct expansion model identity so it does not reuse the local cache. + store.llm = makeLlm("remote-v1", "remote-chat-v1") as any; + const remoteV1 = await store.expandQuery(query); + expect(calls).toEqual(["local", "remote-v1"]); + expect(remoteV1).not.toEqual(local); + + // Same remote expansion model should hit cache. + const remoteV1Again = await store.expandQuery(query); + expect(calls).toEqual(["local", "remote-v1"]); + expect(remoteV1Again).toEqual(remoteV1); + + // Changing only the remote chat model must miss the old remote cache. + store.llm = makeLlm("remote-v2", "remote-chat-v2") as any; + const remoteV2 = await store.expandQuery(query); + expect(calls).toEqual(["local", "remote-v1", "remote-v2"]); + expect(remoteV2).not.toEqual(remoteV1); + + store.clearCache(); + calls.length = 0; + + // If HybridLLM falls back locally after a remote expansion failure, Store + // must not cache that local output under the remote chat model key. + store.llm = makeLlm("fallback-local", "remote-chat", "local-generate") as any; + const fallbackLocal = await store.expandQuery(query); + expect(calls).toEqual(["fallback-local"]); + + store.llm = makeLlm("remote-after-recovery", "remote-chat") as any; + const recoveredRemote = await store.expandQuery(query); + expect(calls).toEqual(["fallback-local", "remote-after-recovery"]); + expect(recoveredRemote).not.toEqual(fallbackLocal); + + await cleanupTestDb(store); + }); }); // ============================================================================= From db5a32e7cad9717eeafa9cf5eec590ef3de7868a Mon Sep 17 00:00:00 2001 From: Kaspre Date: Wed, 3 Jun 2026 17:05:42 -0400 Subject: [PATCH 6/6] fix(remote): preserve producing LLM identity --- src/index.ts | 15 ++++++++- src/store.ts | 9 +++--- test/sdk.test.ts | 79 ++++++++++++++++++++++++++++++++++++++++++++++ test/store.test.ts | 45 ++++++++++++++++++++++++++ 4 files changed, 143 insertions(+), 5 deletions(-) diff --git a/src/index.ts b/src/index.ts index 5bae86109..b1ab19e88 100644 --- a/src/index.ts +++ b/src/index.ts @@ -288,6 +288,7 @@ export interface QMDStore { /** Generate vector embeddings for documents that need them */ embed(options?: { force?: boolean; + /** Local embedding model override; remote embedding rejects mismatches with the configured remote model. */ model?: string; /** Restrict embedding to documents in one collection. */ collection?: string; @@ -516,9 +517,21 @@ export async function createStore(options: StoreOptions): Promise { }, embed: async (embedOpts) => { + const activeEmbedModel = internal.llm?.embedModelName; + if ( + internal.llm?.usesRemoteEmbedding === true + && embedOpts?.model + && activeEmbedModel + && embedOpts.model !== activeEmbedModel + ) { + throw new Error( + `Remote embedding is configured for model '${activeEmbedModel}'; ` + + `store.embed({ model }) cannot override it.` + ); + } return generateEmbeddings(internal, { force: embedOpts?.force, - model: embedOpts?.model, + model: internal.llm?.usesRemoteEmbedding === true ? activeEmbedModel : embedOpts?.model, collection: embedOpts?.collection, maxDocsPerBatch: embedOpts?.maxDocsPerBatch, maxBatchBytes: embedOpts?.maxBatchBytes, diff --git a/src/store.ts b/src/store.ts index e4a3cac67..93fe17df3 100644 --- a/src/store.ts +++ b/src/store.ts @@ -1899,7 +1899,7 @@ export function createStore(dbPath?: string): Store { // Search searchFTS: (query: string, limit?: number, collectionName?: string) => searchFTS(db, query, limit, collectionName), - searchVec: (query: string, model: string, limit?: number, collectionName?: string, session?: ILLMSession, precomputedEmbedding?: number[]) => searchVec(db, query, model, limit, collectionName, session, precomputedEmbedding), + searchVec: (query: string, model: string, limit?: number, collectionName?: string, session?: ILLMSession, precomputedEmbedding?: number[]) => searchVec(db, query, model, limit, collectionName, session, precomputedEmbedding, store.llm), // Query expansion & reranking expandQuery: (query: string, model?: string, intent?: string) => expandQuery(query, model ?? store.llm?.expandModelName ?? store.llm?.generateModelName ?? DEFAULT_QUERY_MODEL, db, intent, store.llm), @@ -3572,11 +3572,11 @@ export function searchFTS(db: Database, query: string, limit: number = 20, colle // Vector Search // ============================================================================= -export async function searchVec(db: Database, query: string, model: string, limit: number = 20, collectionName?: string, session?: ILLMSession, precomputedEmbedding?: number[]): Promise { +export async function searchVec(db: Database, query: string, model: string, limit: number = 20, collectionName?: string, session?: ILLMSession, precomputedEmbedding?: number[], llmOverride?: LLM): Promise { const tableExists = db.prepare(`SELECT name FROM sqlite_master WHERE type='table' AND name='vectors_vec'`).get(); if (!tableExists) return []; - const embedding = precomputedEmbedding ?? await getEmbedding(query, model, true, session); + const embedding = precomputedEmbedding ?? await getEmbedding(query, model, true, session, llmOverride); if (!embedding) return []; // IMPORTANT: We use a two-step query approach here because sqlite-vec virtual tables @@ -3905,7 +3905,8 @@ export async function rerank(query: string, documents: { file: string; text: str const textByFile = new Map(uncachedDocs.map(d => [d.file, d.text])); for (const result of rerankResult.results) { const chunk = textByFile.get(result.file) || ""; - const cacheKey = getCacheKey("rerank", { query: rerankQuery, model, chunk }); + const cacheModel = rerankResult.model || model; + const cacheKey = getCacheKey("rerank", { query: rerankQuery, model: cacheModel, chunk }); setCachedResult(db, cacheKey, result.score.toString()); cachedResults.set(chunk, result.score); } diff --git a/test/sdk.test.ts b/test/sdk.test.ts index fc28736a4..fa4793dd9 100644 --- a/test/sdk.test.ts +++ b/test/sdk.test.ts @@ -665,6 +665,61 @@ describe("searchLex (BM25)", () => { }); }); +describe("searchVector", () => { + test("uses the per-store LLM for query embeddings with remote SDK config", async () => { + await withRemoteEnvCleared(async () => { + const store = await createStore({ + dbPath: freshDbPath(), + config: { + collections: {}, + models: { + embed_api_url: "http://remote.example/v1", + embed_api_model: "remote-embed", + }, + }, + }); + const embedCalls: Array<{ text: string; options?: { model?: string; isQuery?: boolean } }> = []; + + setDefaultLlamaCpp({ + embed: async () => { + throw new Error("searchVector used the global default LLM"); + }, + } as any); + + store.internal.llm = { + embedModelName: "remote-embed", + generateModelName: "local-generate", + rerankModelName: "local-rerank", + usesRemoteEmbedding: true, + embed: async (text: string, options?: { model?: string; isQuery?: boolean }) => { + embedCalls.push({ text, options }); + return { embedding: [0.1, 0.2, 0.3], model: "remote-embed" }; + }, + embedBatch: async (texts: string[]) => texts.map(() => ({ embedding: [0.1, 0.2, 0.3], model: "remote-embed" })), + generate: async () => ({ text: "", model: "local-generate", done: true }), + modelExists: async (model: string) => ({ name: model, exists: true }), + expandQuery: async () => [], + rerank: async () => ({ results: [], model: "local-rerank" }), + tokenize: async () => [] as any, + detokenize: async () => "", + dispose: async () => {}, + } as any; + + try { + store.internal.ensureVecTable(3); + await expect(store.searchVector("remote query", { limit: 1 })).resolves.toEqual([]); + expect(embedCalls).toEqual([{ + text: "remote query", + options: { model: "remote-embed", isQuery: true }, + }]); + } finally { + setDefaultLlamaCpp(null); + await store.close(); + } + }); + }); +}); + // ============================================================================= // Unified search() API Tests // ============================================================================= @@ -1086,6 +1141,30 @@ describe("embed", () => { } }); + test("store.embed rejects model overrides when remote embedding is configured", async () => { + await withRemoteEnvCleared(async () => { + const store = await createStore({ + dbPath: freshDbPath(), + config: { + collections: {}, + models: { + embed_api_url: "http://remote.example/v1", + embed_api_model: "remote-embed", + }, + }, + }); + + try { + await expect(store.embed({ model: "hf:local/embed-model.gguf" })) + .rejects.toThrow(/Remote embedding.*cannot override/i); + await expect(store.embed({ model: "remote-embed" })) + .resolves.toMatchObject({ docsProcessed: 0, chunksEmbedded: 0 }); + } finally { + await store.close(); + } + }); + }); + test("store.embed scopes pending documents to the requested collection", async () => { const store = await createStore({ dbPath: freshDbPath(), diff --git a/test/store.test.ts b/test/store.test.ts index d7ed81da0..402f8d016 100644 --- a/test/store.test.ts +++ b/test/store.test.ts @@ -1223,6 +1223,51 @@ describe("Caching", () => { await cleanupTestDb(store); }); + + test("rerank cache uses the model that actually produced fallback scores", async () => { + const store = await createTestStore(); + const query = "same rerank query"; + const docs = [{ file: "doc.md", text: "same chunk text" }]; + const calls: string[] = []; + + const makeLlm = (label: string, resultModel: string, score: number) => ({ + embedModelName: "remote-embed", + generateModelName: "local-generate", + rerankModelName: "remote-rerank", + embed: async () => ({ embedding: [0.1], model: "remote-embed" }), + embedBatch: async (texts: string[]) => texts.map(() => ({ embedding: [0.1], model: "remote-embed" })), + generate: async () => ({ text: label, model: "local-generate", done: true }), + modelExists: async (model: string) => ({ name: model, exists: true }), + expandQuery: async () => [{ type: "lex" as const, text: "expanded query" }], + rerank: async (_query: string, documents: { file: string; text: string }[]) => { + calls.push(label); + return { + model: resultModel, + results: documents.map((doc, index) => ({ file: doc.file, score, index })), + }; + }, + tokenize: async () => [] as any, + detokenize: async () => "", + dispose: async () => {}, + }); + + store.llm = makeLlm("fallback-local", "local-rerank", 0.42) as any; + const fallback = await store.rerank(query, docs); + expect(calls).toEqual(["fallback-local"]); + expect(fallback).toEqual([{ file: "doc.md", score: 0.42 }]); + + const remoteCacheKey = getCacheKey("rerank", { query, model: "remote-rerank", chunk: docs[0]!.text }); + const localCacheKey = getCacheKey("rerank", { query, model: "local-rerank", chunk: docs[0]!.text }); + expect(store.getCachedResult(remoteCacheKey)).toBeNull(); + expect(store.getCachedResult(localCacheKey)).toBe("0.42"); + + store.llm = makeLlm("remote-after-recovery", "remote-rerank", 0.91) as any; + const recoveredRemote = await store.rerank(query, docs); + expect(calls).toEqual(["fallback-local", "remote-after-recovery"]); + expect(recoveredRemote).toEqual([{ file: "doc.md", score: 0.91 }]); + + await cleanupTestDb(store); + }); }); // =============================================================================