feat(memory): add path-based weights for QMD search

This commit is contained in:
Frad LEE 2026-02-19 00:54:31 +08:00
parent f08fe02a1b
commit 328cf1db6a
3 changed files with 47 additions and 7 deletions

View File

@ -20,6 +20,7 @@ export type MemoryQmdConfig = {
update?: MemoryQmdUpdateConfig;
limits?: MemoryQmdLimitsConfig;
scope?: SessionSendPolicyConfig;
weights?: Record<string, number>;
};
export type MemoryQmdMcporterConfig = {

View File

@ -67,6 +67,7 @@ export type ResolvedQmdConfig = {
limits: ResolvedQmdLimitsConfig;
includeDefaultMemory: boolean;
scope?: SessionSendPolicyConfig;
weights?: Record<string, number>;
};
const DEFAULT_BACKEND: MemoryBackend = "builtin";
@ -344,6 +345,7 @@ export function resolveMemoryBackendConfig(params: {
},
limits: resolveLimits(qmdCfg?.limits),
scope: qmdCfg?.scope ?? DEFAULT_QMD_SCOPE,
weights: qmdCfg?.weights,
};
return {

View File

@ -622,6 +622,9 @@ export class QmdMemoryManager implements MemorySearchManager {
this.qmd.limits.maxResults,
opts?.maxResults ?? this.qmd.limits.maxResults,
);
// Fetch more results to allow for client-side re-ranking (weighting)
const fetchLimit = this.qmd.weights ? limit * 5 : limit;
const collectionNames = this.listManagedCollectionNames();
if (collectionNames.length === 0) {
log.warn("qmd query skipped: no managed collections configured");
@ -645,7 +648,7 @@ export class QmdMemoryManager implements MemorySearchManager {
return await this.runMcporterAcrossCollections({
tool,
query: trimmed,
limit,
limit: fetchLimit,
minScore,
collectionNames,
});
@ -654,7 +657,7 @@ export class QmdMemoryManager implements MemorySearchManager {
mcporter: this.qmd.mcporter,
tool,
query: trimmed,
limit,
limit: fetchLimit,
minScore,
collection: collectionNames[0],
timeoutMs: this.qmd.limits.timeoutMs,
@ -663,12 +666,12 @@ export class QmdMemoryManager implements MemorySearchManager {
if (collectionNames.length > 1) {
return await this.runQueryAcrossCollections(
trimmed,
limit,
fetchLimit,
collectionNames,
qmdSearchCommand,
);
}
const args = this.buildSearchArgs(qmdSearchCommand, trimmed, limit);
const args = this.buildSearchArgs(qmdSearchCommand, trimmed, fetchLimit);
args.push(...this.buildCollectionFilterArgs(collectionNames));
// Always scope to managed collections (default + custom). Even for `search`/`vsearch`,
// pass collection filters; if a given QMD build rejects these flags, we fall back to `query`.
@ -688,9 +691,9 @@ export class QmdMemoryManager implements MemorySearchManager {
);
try {
if (collectionNames.length > 1) {
return await this.runQueryAcrossCollections(trimmed, limit, collectionNames, "query");
return await this.runQueryAcrossCollections(trimmed, fetchLimit, collectionNames, "query");
}
const fallbackArgs = this.buildSearchArgs("query", trimmed, limit);
const fallbackArgs = this.buildSearchArgs("query", trimmed, fetchLimit);
fallbackArgs.push(...this.buildCollectionFilterArgs(collectionNames));
const fallback = await this.runQmd(fallbackArgs, {
timeoutMs: this.qmd.limits.timeoutMs,
@ -717,6 +720,9 @@ export class QmdMemoryManager implements MemorySearchManager {
parsed = await runSearchAttempt(false);
}
const results: MemorySearchResult[] = [];
const weights = this.qmd.weights || {};
const hasWeights = Object.keys(weights).length > 0;
for (const entry of parsed) {
const doc = await this.resolveDocLocation(entry.docid, {
preferredCollection: entry.collection,
@ -727,7 +733,17 @@ export class QmdMemoryManager implements MemorySearchManager {
}
const snippet = entry.snippet?.slice(0, this.qmd.limits.maxSnippetChars) ?? "";
const lines = this.extractSnippetLines(snippet);
const score = typeof entry.score === "number" ? entry.score : 0;
let score = typeof entry.score === "number" ? entry.score : 0;
// Apply weights
if (hasWeights) {
for (const [pattern, weight] of Object.entries(weights)) {
if (this.matchesPath(doc.rel, pattern)) {
score *= weight;
}
}
}
const minScore = opts?.minScore ?? 0;
if (score < minScore) {
continue;
@ -741,9 +757,30 @@ export class QmdMemoryManager implements MemorySearchManager {
source: doc.source,
});
}
// Re-sort if weights were applied
if (hasWeights) {
results.sort((a, b) => b.score - a.score);
}
return this.clampResultsByInjectedChars(this.diversifyResultsBySource(results, limit));
}
private matchesPath(target: string, pattern: string): boolean {
// Simple glob matching
// Handle "dir/**"
if (pattern.endsWith("/**")) {
const prefix = pattern.slice(0, -3);
return target.startsWith(prefix + "/") || target === prefix;
}
// Handle "*.md"
if (pattern.startsWith("*")) {
return target.endsWith(pattern.slice(1));
}
// Exact match
return target === pattern;
}
async sync(params?: {
reason?: string;
force?: boolean;