- 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
101 lines
3.0 KiB
TypeScript
101 lines
3.0 KiB
TypeScript
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
|
|
.toLowerCase()
|
|
.replace(/[^a-z0-9\s]/g, '')
|
|
.split(/\s+/)
|
|
.filter((w) => w.length > 2);
|
|
}
|
|
|
|
export function rankByRelevance(query: string, entries: MemoryEntry[]): MemoryEntry[] {
|
|
const keywords = extractKeywords(query);
|
|
if (keywords.length === 0) return entries.slice(0, 5);
|
|
|
|
const scored = entries.map((entry) => {
|
|
let score = 0;
|
|
const searchText = `${entry.title} ${entry.content} ${entry.tags.join(' ')}`.toLowerCase();
|
|
for (const kw of keywords) {
|
|
if (entry.title.toLowerCase().includes(kw)) score += 3;
|
|
if (entry.tags.some((t) => t.toLowerCase().includes(kw))) score += 2;
|
|
if (entry.content.toLowerCase().includes(kw)) score += 1;
|
|
}
|
|
return { entry, score };
|
|
});
|
|
|
|
return scored
|
|
.filter((s) => s.score > 0)
|
|
.sort((a, b) => b.score - a.score)
|
|
.slice(0, 10)
|
|
.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,
|
|
query?: string,
|
|
): Promise<string[]> {
|
|
const entries = await scanProjectMemory(projectRoot);
|
|
if (entries.length === 0) return [];
|
|
|
|
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';
|