feat(server): memory v2 — BM25 + local embedding hybrid search
- Bm25Ranker: Okapi BM25 scoring (pure TS, no deps) - Embedding module: ONNX-based local embeddings via onnxruntime-node - Hybrid recall: BM25 (30%) + cosine similarity (70%) weighted merge - Falls back to keyword-only via MEMORY_SEARCH=keyword env var - extract_memory agent tool for persisting memory entries
This commit is contained in:
@@ -1,5 +1,9 @@
|
||||
import type { MemoryEntry } from './entries.js';
|
||||
import { scanProjectMemory } from './scan.js';
|
||||
import { Bm25Ranker } from './bm25.js';
|
||||
import { embed, isEmbeddingAvailable } from './embeddings.js';
|
||||
|
||||
const SEARCH_MODE = process.env['MEMORY_SEARCH'] ?? 'hybrid';
|
||||
|
||||
function extractKeywords(query: string): string[] {
|
||||
return query
|
||||
@@ -31,6 +35,51 @@ export function rankByRelevance(query: string, entries: MemoryEntry[]): MemoryEn
|
||||
.map((s) => s.entry);
|
||||
}
|
||||
|
||||
export async function rankByHybrid(
|
||||
query: string,
|
||||
entries: MemoryEntry[],
|
||||
): Promise<MemoryEntry[]> {
|
||||
if (entries.length === 0) return [];
|
||||
const texts = entries.map((e) => `${e.title} ${e.content} ${e.tags.join(' ')}`);
|
||||
|
||||
const bm25 = new Bm25Ranker();
|
||||
bm25.fit(texts);
|
||||
const bm25Scores = texts.map((_, i) => bm25.score(query, i));
|
||||
const maxBm25 = Math.max(...bm25Scores, 1);
|
||||
const normBm25 = bm25Scores.map((s) => s / maxBm25);
|
||||
|
||||
let cosineScores: number[] = [];
|
||||
if (isEmbeddingAvailable()) {
|
||||
const vectors = await embed([query, ...texts]);
|
||||
if (vectors) {
|
||||
const queryVec = vectors[0]!;
|
||||
cosineScores = texts.map((_, i) => {
|
||||
const vec = vectors[i + 1];
|
||||
if (!vec) return 0;
|
||||
let dot = 0, nA = 0, nB = 0;
|
||||
for (let j = 0; j < queryVec.length; j++) {
|
||||
dot += queryVec[j]! * vec[j]!;
|
||||
nA += queryVec[j]! * queryVec[j]!;
|
||||
nB += vec[j]! * vec[j]!;
|
||||
}
|
||||
const denom = Math.sqrt(nA) * Math.sqrt(nB);
|
||||
return denom === 0 ? 0 : dot / denom;
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
const scored = entries.map((entry, i) => {
|
||||
const combined = (normBm25[i]! * 0.3) + ((cosineScores[i] ?? 0) * 0.7);
|
||||
return { entry, score: combined };
|
||||
});
|
||||
|
||||
return scored
|
||||
.filter((s) => s.score >= 0.15)
|
||||
.sort((a, b) => b.score - a.score)
|
||||
.slice(0, 10)
|
||||
.map((s) => s.entry);
|
||||
}
|
||||
|
||||
export async function loadMemoryForSession(
|
||||
projectRoot: string,
|
||||
_sessionId?: string,
|
||||
@@ -39,6 +88,13 @@ export async function loadMemoryForSession(
|
||||
const entries = await scanProjectMemory(projectRoot);
|
||||
if (entries.length === 0) return [];
|
||||
|
||||
const relevant = query ? rankByRelevance(query, entries) : entries.slice(0, 5);
|
||||
const relevant = query
|
||||
? SEARCH_MODE === 'keyword'
|
||||
? rankByRelevance(query, entries)
|
||||
: await rankByHybrid(query, entries)
|
||||
: entries.slice(0, 5);
|
||||
|
||||
return relevant.map((e) => `[${e.topic}] ${e.title}: ${e.content}`);
|
||||
}
|
||||
|
||||
export { initEmbeddings } from './embeddings.js';
|
||||
|
||||
Reference in New Issue
Block a user