Unify QMD model resolution
This commit is contained in:
Tobi Lütke
parent
1f757379e2
commit
da184e58e9
@ -8,6 +8,7 @@
|
||||
|
||||
### Fixes
|
||||
|
||||
- Model config: centralize embedding/rerank/generation model resolution so `qmd embed`, `status`, `query`, `vsearch`, `pull`, SDK vector search, and `bench` use the same active `.qmd/index.yaml` model hints and environment fallbacks.
|
||||
- GPU/status: `qmd status` now uses the same embedding model identity as `qmd embed` when computing pending embeddings, so URI-backed embeddings are not incorrectly reported as pending under the legacy `embeddinggemma` alias.
|
||||
- GPU status: `qmd status` now always shows GPU mode/configuration without unsafe native probing, and CPU-fallback warnings point to `QMD_STATUS_DEVICE_PROBE=1 qmd status` for an actual backend probe. The no-GPU warning is emitted once per process instead of once per LLM instance during benchmarks.
|
||||
- GPU: add `QMD_FORCE_CPU=1` / `--no-gpu` to bypass CUDA/Vulkan/Metal probing entirely, and route native llama.cpp stdout noise to stderr so JSON output stays parseable during search/query commands.
|
||||
|
||||
@ -78,7 +78,7 @@ import {
|
||||
type ReindexResult,
|
||||
type ChunkStrategy,
|
||||
} from "../store.js";
|
||||
import { disposeDefaultLlamaCpp, getDefaultLlamaCpp, setDefaultLlamaCpp, LlamaCpp, withLLMSession, pullModels, DEFAULT_EMBED_MODEL_URI, DEFAULT_GENERATE_MODEL_URI, DEFAULT_RERANK_MODEL_URI, DEFAULT_MODEL_CACHE_DIR } from "../llm.js";
|
||||
import { disposeDefaultLlamaCpp, getDefaultLlamaCpp, setDefaultLlamaCpp, LlamaCpp, withLLMSession, pullModels, DEFAULT_MODEL_CACHE_DIR, resolveEmbedModel, resolveGenerateModel, resolveRerankModel, resolveModels } from "../llm.js";
|
||||
import {
|
||||
formatSearchResults,
|
||||
formatDocuments,
|
||||
@ -537,10 +537,11 @@ async function showStatus(): Promise<void> {
|
||||
const match = uri.match(/^hf:([^/]+\/[^/]+)\//);
|
||||
return match ? `https://huggingface.co/${match[1]}` : uri;
|
||||
};
|
||||
const activeModels = resolveModelsForCli();
|
||||
console.log(`\n${c.bold}Models${c.reset}`);
|
||||
console.log(` Embedding: ${hfLink(DEFAULT_EMBED_MODEL_URI)}`);
|
||||
console.log(` Reranking: ${hfLink(DEFAULT_RERANK_MODEL_URI)}`);
|
||||
console.log(` Generation: ${hfLink(DEFAULT_GENERATE_MODEL_URI)}`);
|
||||
console.log(` Embedding: ${hfLink(activeModels.embed)}`);
|
||||
console.log(` Reranking: ${hfLink(activeModels.rerank)}`);
|
||||
console.log(` Generation: ${hfLink(activeModels.generate)}`);
|
||||
}
|
||||
|
||||
// Device / GPU info
|
||||
@ -1802,7 +1803,35 @@ function parseChunkStrategy(value: unknown): ChunkStrategy | undefined {
|
||||
}
|
||||
|
||||
export function resolveEmbedModelForCli(): string {
|
||||
return process.env.QMD_EMBED_MODEL ?? DEFAULT_EMBED_MODEL_URI;
|
||||
try {
|
||||
return resolveEmbedModel(loadConfig().models);
|
||||
} catch {
|
||||
return resolveEmbedModel();
|
||||
}
|
||||
}
|
||||
|
||||
export function resolveGenerateModelForCli(): string {
|
||||
try {
|
||||
return resolveGenerateModel(loadConfig().models);
|
||||
} catch {
|
||||
return resolveGenerateModel();
|
||||
}
|
||||
}
|
||||
|
||||
export function resolveRerankModelForCli(): string {
|
||||
try {
|
||||
return resolveRerankModel(loadConfig().models);
|
||||
} catch {
|
||||
return resolveRerankModel();
|
||||
}
|
||||
}
|
||||
|
||||
function resolveModelsForCli(): { embed: string; generate: string; rerank: string } {
|
||||
try {
|
||||
return resolveModels(loadConfig().models);
|
||||
} catch {
|
||||
return resolveModels();
|
||||
}
|
||||
}
|
||||
|
||||
async function vectorIndex(
|
||||
@ -3539,10 +3568,11 @@ if (isMain) {
|
||||
|
||||
case "pull": {
|
||||
const refresh = cli.values.refresh === undefined ? false : Boolean(cli.values.refresh);
|
||||
const activeModels = resolveModelsForCli();
|
||||
const models = [
|
||||
DEFAULT_EMBED_MODEL_URI,
|
||||
DEFAULT_GENERATE_MODEL_URI,
|
||||
DEFAULT_RERANK_MODEL_URI,
|
||||
activeModels.embed,
|
||||
activeModels.generate,
|
||||
activeModels.rerank,
|
||||
];
|
||||
console.log(`${c.bold}Pulling models${c.reset}`);
|
||||
const results = await pullModels(models, {
|
||||
|
||||
@ -23,7 +23,6 @@ import {
|
||||
structuredSearch,
|
||||
extractSnippet,
|
||||
addLineNumbers,
|
||||
DEFAULT_EMBED_MODEL,
|
||||
DEFAULT_MULTI_GET_MAX_BYTES,
|
||||
reindexCollection,
|
||||
generateEmbeddings,
|
||||
@ -423,7 +422,7 @@ export async function createStore(options: StoreOptions): Promise<QMDStore> {
|
||||
});
|
||||
},
|
||||
searchLex: async (q, opts) => internal.searchFTS(q, opts?.limit, opts?.collection),
|
||||
searchVector: async (q, opts) => internal.searchVec(q, DEFAULT_EMBED_MODEL, opts?.limit, opts?.collection),
|
||||
searchVector: async (q, opts) => internal.searchVec(q, llm.embedModelName, opts?.limit, opts?.collection),
|
||||
expandQuery: async (q, opts) => internal.expandQuery(q, undefined, opts?.intent),
|
||||
get: async (pathOrDocid, opts) => internal.findDocument(pathOrDocid, opts),
|
||||
getDocumentBody: async (pathOrDocid, opts) => {
|
||||
|
||||
44
src/llm.ts
44
src/llm.ts
@ -84,7 +84,7 @@ export function isQwen3EmbeddingModel(modelUri: string): boolean {
|
||||
* Uses Qwen3-Embedding instruct format when a Qwen embedding model is active.
|
||||
*/
|
||||
export function formatQueryForEmbedding(query: string, modelUri?: string): string {
|
||||
const uri = modelUri ?? process.env.QMD_EMBED_MODEL ?? DEFAULT_EMBED_MODEL;
|
||||
const uri = modelUri ?? resolveEmbedModel();
|
||||
if (isQwen3EmbeddingModel(uri)) {
|
||||
return `Instruct: Retrieve relevant documents for the given query\nQuery: ${query}`;
|
||||
}
|
||||
@ -97,7 +97,7 @@ export function formatQueryForEmbedding(query: string, modelUri?: string): strin
|
||||
* Qwen3-Embedding encodes documents as raw text without special prefixes.
|
||||
*/
|
||||
export function formatDocForEmbedding(text: string, title?: string, modelUri?: string): string {
|
||||
const uri = modelUri ?? process.env.QMD_EMBED_MODEL ?? DEFAULT_EMBED_MODEL;
|
||||
const uri = modelUri ?? resolveEmbedModel();
|
||||
if (isQwen3EmbeddingModel(uri)) {
|
||||
// Qwen3-Embedding: documents are raw text, no task prefix
|
||||
return title ? `${title}\n${text}` : text;
|
||||
@ -256,6 +256,32 @@ export const DEFAULT_EMBED_MODEL_URI = DEFAULT_EMBED_MODEL;
|
||||
export const DEFAULT_RERANK_MODEL_URI = DEFAULT_RERANK_MODEL;
|
||||
export const DEFAULT_GENERATE_MODEL_URI = DEFAULT_GENERATE_MODEL;
|
||||
|
||||
export type ModelResolutionConfig = {
|
||||
embed?: string;
|
||||
generate?: string;
|
||||
rerank?: string;
|
||||
};
|
||||
|
||||
export function resolveEmbedModel(config?: ModelResolutionConfig): string {
|
||||
return config?.embed || process.env.QMD_EMBED_MODEL || DEFAULT_EMBED_MODEL;
|
||||
}
|
||||
|
||||
export function resolveGenerateModel(config?: ModelResolutionConfig): string {
|
||||
return config?.generate || process.env.QMD_GENERATE_MODEL || DEFAULT_GENERATE_MODEL;
|
||||
}
|
||||
|
||||
export function resolveRerankModel(config?: ModelResolutionConfig): string {
|
||||
return config?.rerank || process.env.QMD_RERANK_MODEL || DEFAULT_RERANK_MODEL;
|
||||
}
|
||||
|
||||
export function resolveModels(config?: ModelResolutionConfig): Required<ModelResolutionConfig> {
|
||||
return {
|
||||
embed: resolveEmbedModel(config),
|
||||
generate: resolveGenerateModel(config),
|
||||
rerank: resolveRerankModel(config),
|
||||
};
|
||||
}
|
||||
|
||||
// Local model cache directory
|
||||
const MODEL_CACHE_DIR = process.env.XDG_CACHE_HOME
|
||||
? join(process.env.XDG_CACHE_HOME, "qmd", "models")
|
||||
@ -612,9 +638,9 @@ export class LlamaCpp implements LLM {
|
||||
|
||||
|
||||
constructor(config: LlamaCppConfig = {}) {
|
||||
this.embedModelUri = config.embedModel || process.env.QMD_EMBED_MODEL || DEFAULT_EMBED_MODEL;
|
||||
this.generateModelUri = config.generateModel || process.env.QMD_GENERATE_MODEL || DEFAULT_GENERATE_MODEL;
|
||||
this.rerankModelUri = config.rerankModel || process.env.QMD_RERANK_MODEL || DEFAULT_RERANK_MODEL;
|
||||
this.embedModelUri = resolveEmbedModel({ embed: config.embedModel });
|
||||
this.generateModelUri = resolveGenerateModel({ generate: config.generateModel });
|
||||
this.rerankModelUri = resolveRerankModel({ rerank: config.rerankModel });
|
||||
this.modelCacheDir = config.modelCacheDir || MODEL_CACHE_DIR;
|
||||
this.expandContextSize = resolveExpandContextSize(config.expandContextSize);
|
||||
this.inactivityTimeoutMs = config.inactivityTimeoutMs ?? DEFAULT_INACTIVITY_TIMEOUT_MS;
|
||||
@ -625,6 +651,14 @@ export class LlamaCpp implements LLM {
|
||||
return this.embedModelUri;
|
||||
}
|
||||
|
||||
get generateModelName(): string {
|
||||
return this.generateModelUri;
|
||||
}
|
||||
|
||||
get rerankModelName(): string {
|
||||
return this.rerankModelUri;
|
||||
}
|
||||
|
||||
/**
|
||||
* Reset the inactivity timer. Called after each model operation.
|
||||
* When timer fires, models are unloaded to free memory (if no active sessions).
|
||||
|
||||
13
src/store.ts
13
src/store.ts
@ -25,6 +25,9 @@ import {
|
||||
formatQueryForEmbedding,
|
||||
formatDocForEmbedding,
|
||||
withLLMSessionForLlm,
|
||||
DEFAULT_EMBED_MODEL_URI,
|
||||
DEFAULT_RERANK_MODEL_URI,
|
||||
DEFAULT_GENERATE_MODEL_URI,
|
||||
type RerankDocument,
|
||||
type ILLMSession,
|
||||
} from "./llm.js";
|
||||
@ -39,9 +42,9 @@ import type {
|
||||
// Configuration
|
||||
// =============================================================================
|
||||
|
||||
export const DEFAULT_EMBED_MODEL = "embeddinggemma";
|
||||
export const DEFAULT_RERANK_MODEL = "ExpedientFalcon/qwen3-reranker:0.6b-q8_0";
|
||||
export const DEFAULT_QUERY_MODEL = "Qwen/Qwen3-1.7B";
|
||||
export const DEFAULT_EMBED_MODEL = DEFAULT_EMBED_MODEL_URI;
|
||||
export const DEFAULT_RERANK_MODEL = DEFAULT_RERANK_MODEL_URI;
|
||||
export const DEFAULT_QUERY_MODEL = DEFAULT_GENERATE_MODEL_URI;
|
||||
export const DEFAULT_GLOB = "**/*.md";
|
||||
export const DEFAULT_MULTI_GET_MAX_BYTES = 10 * 1024; // 10KB
|
||||
export const DEFAULT_EMBED_MAX_DOCS_PER_BATCH = 64;
|
||||
@ -1749,8 +1752,8 @@ 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, db, intent, store.llm),
|
||||
rerank: (query: string, documents: { file: string; text: string }[], model?: string, intent?: string) => rerank(query, documents, model, db, intent, store.llm),
|
||||
expandQuery: (query: string, model?: string, intent?: string) => expandQuery(query, model ?? 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
|
||||
findDocument: (filename: string, options?: { includeBody?: boolean }) => findDocument(db, filename, options),
|
||||
|
||||
@ -17,6 +17,10 @@ import {
|
||||
withNativeStdoutRedirectedToStderr,
|
||||
resolveParallelismOverride,
|
||||
resolveSafeParallelism,
|
||||
resolveEmbedModel,
|
||||
resolveGenerateModel,
|
||||
resolveRerankModel,
|
||||
resolveModels,
|
||||
withLLMSession,
|
||||
canUnloadLLM,
|
||||
SessionReleasedError,
|
||||
@ -24,6 +28,63 @@ import {
|
||||
type ILLMSession,
|
||||
} from "../src/llm.js";
|
||||
|
||||
describe("model name resolution", () => {
|
||||
function withModelEnv(env: Record<string, string | undefined>, fn: () => void): void {
|
||||
const previous = {
|
||||
QMD_EMBED_MODEL: process.env.QMD_EMBED_MODEL,
|
||||
QMD_GENERATE_MODEL: process.env.QMD_GENERATE_MODEL,
|
||||
QMD_RERANK_MODEL: process.env.QMD_RERANK_MODEL,
|
||||
};
|
||||
try {
|
||||
for (const [key, value] of Object.entries(env)) {
|
||||
if (value === undefined) delete process.env[key];
|
||||
else process.env[key] = value;
|
||||
}
|
||||
fn();
|
||||
} finally {
|
||||
for (const [key, value] of Object.entries(previous)) {
|
||||
if (value === undefined) delete process.env[key];
|
||||
else process.env[key] = value;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
test("all model roles resolve config hints before env fallbacks", () => {
|
||||
withModelEnv({
|
||||
QMD_EMBED_MODEL: "env-embed",
|
||||
QMD_GENERATE_MODEL: "env-generate",
|
||||
QMD_RERANK_MODEL: "env-rerank",
|
||||
}, () => {
|
||||
const config = {
|
||||
embed: "config-embed",
|
||||
generate: "config-generate",
|
||||
rerank: "config-rerank",
|
||||
};
|
||||
expect(resolveEmbedModel(config)).toBe("config-embed");
|
||||
expect(resolveGenerateModel(config)).toBe("config-generate");
|
||||
expect(resolveRerankModel(config)).toBe("config-rerank");
|
||||
expect(resolveModels(config)).toEqual(config);
|
||||
});
|
||||
});
|
||||
|
||||
test("LlamaCpp constructor uses the same resolver as status/embed/query helpers", () => {
|
||||
withModelEnv({
|
||||
QMD_EMBED_MODEL: "env-embed",
|
||||
QMD_GENERATE_MODEL: "env-generate",
|
||||
QMD_RERANK_MODEL: "env-rerank",
|
||||
}, () => {
|
||||
const llm = new LlamaCpp({
|
||||
embedModel: "config-embed",
|
||||
generateModel: "config-generate",
|
||||
rerankModel: "config-rerank",
|
||||
});
|
||||
expect(llm.embedModelName).toBe(resolveEmbedModel({ embed: "config-embed" }));
|
||||
expect(llm.generateModelName).toBe(resolveGenerateModel({ generate: "config-generate" }));
|
||||
expect(llm.rerankModelName).toBe(resolveRerankModel({ rerank: "config-rerank" }));
|
||||
});
|
||||
});
|
||||
});
|
||||
|
||||
// =============================================================================
|
||||
// Singleton Tests (no model loading required)
|
||||
// =============================================================================
|
||||
|
||||
@ -56,7 +56,7 @@ describe("local .qmd project config", () => {
|
||||
mkdirSync(join(root, ".qmd"), { recursive: true });
|
||||
mkdirSync(join(root, "docs"), { recursive: true });
|
||||
writeFileSync(join(root, "docs", "a.md"), "# A\n\nLocal test document.\n");
|
||||
writeFileSync(join(root, ".qmd", "index.yaml"), `collections:\n docs:\n path: ${JSON.stringify(join(root, "docs"))}\n pattern: "**/*.md"\n context:\n /: Local test docs\n`);
|
||||
writeFileSync(join(root, ".qmd", "index.yaml"), `collections:\n docs:\n path: ${JSON.stringify(join(root, "docs"))}\n pattern: "**/*.md"\n context:\n /: Local test docs\nmodels:\n embed: local-embed-model\n rerank: local-rerank-model\n generate: local-generate-model\n`);
|
||||
|
||||
const home = join(root, "home");
|
||||
const tsxBin = join(process.cwd(), "node_modules", ".bin", "tsx");
|
||||
@ -69,12 +69,19 @@ describe("local .qmd project config", () => {
|
||||
HOME: home,
|
||||
XDG_CONFIG_HOME: join(home, ".config"),
|
||||
XDG_CACHE_HOME: join(home, ".cache"),
|
||||
QMD_EMBED_MODEL: "env-embed-model",
|
||||
QMD_RERANK_MODEL: "env-rerank-model",
|
||||
QMD_GENERATE_MODEL: "env-generate-model",
|
||||
},
|
||||
});
|
||||
|
||||
const localIndex = join(root, ".qmd", "index.sqlite");
|
||||
expect(output).toContain(`Index: ${realpathSync(localIndex)}`);
|
||||
expect(output).toContain("docs (qmd://docs/)");
|
||||
expect(output).toContain("Embedding: local-embed-model");
|
||||
expect(output).toContain("Reranking: local-rerank-model");
|
||||
expect(output).toContain("Generation: local-generate-model");
|
||||
expect(output).not.toContain("env-embed-model");
|
||||
expect(existsSync(localIndex)).toBe(true);
|
||||
expect(existsSync(join(home, ".cache", "qmd", "index.sqlite"))).toBe(false);
|
||||
});
|
||||
|
||||
Loading…
Reference in New Issue
Block a user