@winci/local-rag 0.2.1
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- package/.claude-plugin/plugin.json +24 -0
- package/.mcp.json +11 -0
- package/LICENSE +21 -0
- package/README.md +567 -0
- package/hooks/hooks.json +25 -0
- package/hooks/scripts/reindex-file.sh +19 -0
- package/hooks/scripts/session-start.sh +11 -0
- package/package.json +52 -0
- package/skills/local-rag/SKILL.md +42 -0
- package/src/cli/commands/analytics.ts +58 -0
- package/src/cli/commands/benchmark.ts +30 -0
- package/src/cli/commands/checkpoint.ts +85 -0
- package/src/cli/commands/conversation.ts +102 -0
- package/src/cli/commands/demo.ts +119 -0
- package/src/cli/commands/eval.ts +31 -0
- package/src/cli/commands/index-cmd.ts +26 -0
- package/src/cli/commands/init.ts +35 -0
- package/src/cli/commands/map.ts +21 -0
- package/src/cli/commands/remove.ts +15 -0
- package/src/cli/commands/search-cmd.ts +59 -0
- package/src/cli/commands/serve.ts +5 -0
- package/src/cli/commands/status.ts +13 -0
- package/src/cli/index.ts +117 -0
- package/src/cli/progress.ts +21 -0
- package/src/cli/setup.ts +192 -0
- package/src/config/index.ts +101 -0
- package/src/conversation/indexer.ts +147 -0
- package/src/conversation/parser.ts +323 -0
- package/src/db/analytics.ts +116 -0
- package/src/db/annotations.ts +161 -0
- package/src/db/checkpoints.ts +166 -0
- package/src/db/conversation.ts +241 -0
- package/src/db/files.ts +146 -0
- package/src/db/graph.ts +250 -0
- package/src/db/index.ts +468 -0
- package/src/db/search.ts +244 -0
- package/src/db/types.ts +85 -0
- package/src/embeddings/embed.ts +73 -0
- package/src/graph/resolver.ts +305 -0
- package/src/indexing/chunker.ts +523 -0
- package/src/indexing/indexer.ts +263 -0
- package/src/indexing/parse.ts +99 -0
- package/src/indexing/watcher.ts +84 -0
- package/src/main.ts +8 -0
- package/src/search/benchmark.ts +139 -0
- package/src/search/eval.ts +171 -0
- package/src/search/hybrid.ts +194 -0
- package/src/search/reranker.ts +99 -0
- package/src/search/usages.ts +27 -0
- package/src/server/index.ts +126 -0
- package/src/tools/analytics-tools.ts +58 -0
- package/src/tools/annotation-tools.ts +89 -0
- package/src/tools/checkpoint-tools.ts +147 -0
- package/src/tools/conversation-tools.ts +86 -0
- package/src/tools/git-tools.ts +103 -0
- package/src/tools/graph-tools.ts +163 -0
- package/src/tools/index-tools.ts +91 -0
- package/src/tools/index.ts +33 -0
- package/src/tools/search.ts +238 -0
- package/src/types.ts +9 -0
- package/src/utils/log.ts +39 -0
|
@@ -0,0 +1,171 @@
|
|
|
1
|
+
import { readFile, writeFile } from "fs/promises";
|
|
2
|
+
import { resolve, basename } from "path";
|
|
3
|
+
import { RagDB } from "../db";
|
|
4
|
+
import { search, type DedupedResult } from "./hybrid";
|
|
5
|
+
import { loadConfig } from "../config";
|
|
6
|
+
|
|
7
|
+
export interface EvalTask {
|
|
8
|
+
task: string;
|
|
9
|
+
grading: string; // human-readable criteria for what a good answer looks like
|
|
10
|
+
expectedFiles?: string[]; // optional: files the agent should reference
|
|
11
|
+
}
|
|
12
|
+
|
|
13
|
+
export interface EvalTrace {
|
|
14
|
+
task: string;
|
|
15
|
+
grading: string;
|
|
16
|
+
condition: "with-rag" | "without-rag";
|
|
17
|
+
searchResults: DedupedResult[];
|
|
18
|
+
filesReferenced: string[];
|
|
19
|
+
searchCount: number;
|
|
20
|
+
durationMs: number;
|
|
21
|
+
}
|
|
22
|
+
|
|
23
|
+
export interface EvalSummary {
|
|
24
|
+
totalTasks: number;
|
|
25
|
+
withRag: {
|
|
26
|
+
avgSearchResults: number;
|
|
27
|
+
avgFilesReferenced: number;
|
|
28
|
+
avgDurationMs: number;
|
|
29
|
+
fileHitRate: number; // % of tasks where expected files were found
|
|
30
|
+
};
|
|
31
|
+
withoutRag: {
|
|
32
|
+
avgSearchResults: number;
|
|
33
|
+
avgFilesReferenced: number;
|
|
34
|
+
avgDurationMs: number;
|
|
35
|
+
fileHitRate: number;
|
|
36
|
+
};
|
|
37
|
+
traces: EvalTrace[];
|
|
38
|
+
}
|
|
39
|
+
|
|
40
|
+
export async function loadEvalTasks(path: string): Promise<EvalTask[]> {
|
|
41
|
+
const raw = await readFile(path, "utf-8");
|
|
42
|
+
const parsed = JSON.parse(raw);
|
|
43
|
+
|
|
44
|
+
if (!Array.isArray(parsed)) {
|
|
45
|
+
throw new Error("Eval file must be a JSON array of { task, grading } objects");
|
|
46
|
+
}
|
|
47
|
+
|
|
48
|
+
for (const entry of parsed) {
|
|
49
|
+
if (!entry.task || !entry.grading) {
|
|
50
|
+
throw new Error(`Invalid eval entry: ${JSON.stringify(entry)}. Each entry needs "task" (string) and "grading" (string)`);
|
|
51
|
+
}
|
|
52
|
+
}
|
|
53
|
+
|
|
54
|
+
return parsed;
|
|
55
|
+
}
|
|
56
|
+
|
|
57
|
+
/**
|
|
58
|
+
* Simulate an agent's search behavior for a task.
|
|
59
|
+
* "with-rag": runs semantic search on the task description and returns what was found.
|
|
60
|
+
* "without-rag": returns empty results (simulating an agent with no RAG).
|
|
61
|
+
*/
|
|
62
|
+
export async function runEvalTask(
|
|
63
|
+
task: EvalTask,
|
|
64
|
+
db: RagDB,
|
|
65
|
+
projectDir: string,
|
|
66
|
+
condition: "with-rag" | "without-rag",
|
|
67
|
+
topK: number = 5
|
|
68
|
+
): Promise<EvalTrace> {
|
|
69
|
+
const start = performance.now();
|
|
70
|
+
let searchResults: DedupedResult[] = [];
|
|
71
|
+
let searchCount = 0;
|
|
72
|
+
|
|
73
|
+
if (condition === "with-rag") {
|
|
74
|
+
const config = await loadConfig(projectDir);
|
|
75
|
+
searchResults = await search(task.task, db, topK, 0, config.hybridWeight);
|
|
76
|
+
searchCount = 1;
|
|
77
|
+
}
|
|
78
|
+
|
|
79
|
+
const durationMs = Math.round(performance.now() - start);
|
|
80
|
+
const filesReferenced = searchResults.map((r) => r.path);
|
|
81
|
+
|
|
82
|
+
return {
|
|
83
|
+
task: task.task,
|
|
84
|
+
grading: task.grading,
|
|
85
|
+
condition,
|
|
86
|
+
searchResults,
|
|
87
|
+
filesReferenced,
|
|
88
|
+
searchCount,
|
|
89
|
+
durationMs,
|
|
90
|
+
};
|
|
91
|
+
}
|
|
92
|
+
|
|
93
|
+
export async function runEval(
|
|
94
|
+
tasks: EvalTask[],
|
|
95
|
+
db: RagDB,
|
|
96
|
+
projectDir: string,
|
|
97
|
+
topK: number = 5
|
|
98
|
+
): Promise<EvalSummary> {
|
|
99
|
+
const traces: EvalTrace[] = [];
|
|
100
|
+
|
|
101
|
+
for (const task of tasks) {
|
|
102
|
+
const withRag = await runEvalTask(task, db, projectDir, "with-rag", topK);
|
|
103
|
+
const withoutRag = await runEvalTask(task, db, projectDir, "without-rag", topK);
|
|
104
|
+
traces.push(withRag, withoutRag);
|
|
105
|
+
}
|
|
106
|
+
|
|
107
|
+
const withRagTraces = traces.filter((t) => t.condition === "with-rag");
|
|
108
|
+
const withoutRagTraces = traces.filter((t) => t.condition === "without-rag");
|
|
109
|
+
|
|
110
|
+
function computeStats(traceSet: EvalTrace[], tasks: EvalTask[]) {
|
|
111
|
+
const n = traceSet.length || 1;
|
|
112
|
+
const avgSearchResults = traceSet.reduce((s, t) => s + t.searchResults.length, 0) / n;
|
|
113
|
+
const avgFilesReferenced = traceSet.reduce((s, t) => s + t.filesReferenced.length, 0) / n;
|
|
114
|
+
const avgDurationMs = traceSet.reduce((s, t) => s + t.durationMs, 0) / n;
|
|
115
|
+
|
|
116
|
+
// File hit rate: % of tasks with expectedFiles where at least one was found
|
|
117
|
+
let hits = 0;
|
|
118
|
+
let withExpected = 0;
|
|
119
|
+
for (let i = 0; i < tasks.length; i++) {
|
|
120
|
+
if (tasks[i].expectedFiles && tasks[i].expectedFiles!.length > 0) {
|
|
121
|
+
withExpected++;
|
|
122
|
+
const expected = tasks[i].expectedFiles!;
|
|
123
|
+
const found = traceSet[i].filesReferenced;
|
|
124
|
+
const hasHit = expected.some((e) =>
|
|
125
|
+
found.some((f) => f === e || f.endsWith(e) || e.endsWith(f))
|
|
126
|
+
);
|
|
127
|
+
if (hasHit) hits++;
|
|
128
|
+
}
|
|
129
|
+
}
|
|
130
|
+
const fileHitRate = withExpected > 0 ? hits / withExpected : 0;
|
|
131
|
+
|
|
132
|
+
return { avgSearchResults, avgFilesReferenced, avgDurationMs, fileHitRate };
|
|
133
|
+
}
|
|
134
|
+
|
|
135
|
+
return {
|
|
136
|
+
totalTasks: tasks.length,
|
|
137
|
+
withRag: computeStats(withRagTraces, tasks),
|
|
138
|
+
withoutRag: computeStats(withoutRagTraces, tasks),
|
|
139
|
+
traces,
|
|
140
|
+
};
|
|
141
|
+
}
|
|
142
|
+
|
|
143
|
+
export function formatEvalReport(summary: EvalSummary): string {
|
|
144
|
+
const lines: string[] = [];
|
|
145
|
+
|
|
146
|
+
lines.push(`A/B Eval results (${summary.totalTasks} tasks):`);
|
|
147
|
+
lines.push("");
|
|
148
|
+
lines.push(" With RAG Without RAG");
|
|
149
|
+
lines.push(` Avg results: ${summary.withRag.avgSearchResults.toFixed(1).padStart(8)} ${summary.withoutRag.avgSearchResults.toFixed(1).padStart(11)}`);
|
|
150
|
+
lines.push(` Avg files found: ${summary.withRag.avgFilesReferenced.toFixed(1).padStart(8)} ${summary.withoutRag.avgFilesReferenced.toFixed(1).padStart(11)}`);
|
|
151
|
+
lines.push(` File hit rate: ${(summary.withRag.fileHitRate * 100).toFixed(0).padStart(7)}% ${(summary.withoutRag.fileHitRate * 100).toFixed(0).padStart(10)}%`);
|
|
152
|
+
lines.push(` Avg latency: ${summary.withRag.avgDurationMs.toFixed(0).padStart(7)}ms ${summary.withoutRag.avgDurationMs.toFixed(0).padStart(10)}ms`);
|
|
153
|
+
|
|
154
|
+
// Per-task breakdown
|
|
155
|
+
lines.push("\nPer-task breakdown:");
|
|
156
|
+
const withRagTraces = summary.traces.filter((t) => t.condition === "with-rag");
|
|
157
|
+
for (const trace of withRagTraces) {
|
|
158
|
+
const files = trace.filesReferenced.length > 0
|
|
159
|
+
? trace.filesReferenced.map((f) => basename(f)).join(", ")
|
|
160
|
+
: "(none)";
|
|
161
|
+
lines.push(` "${trace.task}"`);
|
|
162
|
+
lines.push(` files found: ${files}`);
|
|
163
|
+
lines.push(` grading: ${trace.grading}`);
|
|
164
|
+
}
|
|
165
|
+
|
|
166
|
+
return lines.join("\n");
|
|
167
|
+
}
|
|
168
|
+
|
|
169
|
+
export async function saveEvalTraces(traces: EvalTrace[], outputPath: string): Promise<void> {
|
|
170
|
+
await writeFile(outputPath, JSON.stringify(traces, null, 2));
|
|
171
|
+
}
|
|
@@ -0,0 +1,194 @@
|
|
|
1
|
+
import { embed } from "../embeddings/embed";
|
|
2
|
+
import { RagDB, type SearchResult, type ChunkSearchResult } from "../db";
|
|
3
|
+
import { rerank } from "./reranker";
|
|
4
|
+
import { log } from "../utils/log";
|
|
5
|
+
|
|
6
|
+
export interface DedupedResult {
|
|
7
|
+
path: string;
|
|
8
|
+
score: number;
|
|
9
|
+
snippets: string[];
|
|
10
|
+
}
|
|
11
|
+
|
|
12
|
+
export interface ChunkResult {
|
|
13
|
+
path: string;
|
|
14
|
+
score: number;
|
|
15
|
+
content: string;
|
|
16
|
+
chunkIndex: number;
|
|
17
|
+
entityName: string | null;
|
|
18
|
+
chunkType: string | null;
|
|
19
|
+
startLine: number | null;
|
|
20
|
+
endLine: number | null;
|
|
21
|
+
}
|
|
22
|
+
|
|
23
|
+
// Default: 70% vector, 30% BM25
|
|
24
|
+
const DEFAULT_HYBRID_WEIGHT = 0.7;
|
|
25
|
+
|
|
26
|
+
/**
|
|
27
|
+
* Merge vector and text search results using hybrid scoring.
|
|
28
|
+
* Each result must have `score`, `path`, and `chunkIndex` at minimum.
|
|
29
|
+
* Extra fields from the vector results are preserved on the merged output.
|
|
30
|
+
*/
|
|
31
|
+
export function mergeHybridScores<T extends { score: number; path: string; chunkIndex: number }>(
|
|
32
|
+
vectorResults: T[],
|
|
33
|
+
textResults: T[],
|
|
34
|
+
hybridWeight: number
|
|
35
|
+
): T[] {
|
|
36
|
+
const scoreMap = new Map<string, { item: T; vectorScore: number; textScore: number }>();
|
|
37
|
+
|
|
38
|
+
for (const r of vectorResults) {
|
|
39
|
+
const key = `${r.path}:${r.chunkIndex}`;
|
|
40
|
+
scoreMap.set(key, { item: r, vectorScore: r.score, textScore: 0 });
|
|
41
|
+
}
|
|
42
|
+
|
|
43
|
+
for (const r of textResults) {
|
|
44
|
+
const key = `${r.path}:${r.chunkIndex}`;
|
|
45
|
+
const existing = scoreMap.get(key);
|
|
46
|
+
if (existing) {
|
|
47
|
+
existing.textScore = r.score;
|
|
48
|
+
} else {
|
|
49
|
+
scoreMap.set(key, { item: r, vectorScore: 0, textScore: r.score });
|
|
50
|
+
}
|
|
51
|
+
}
|
|
52
|
+
|
|
53
|
+
return Array.from(scoreMap.values()).map((entry) => ({
|
|
54
|
+
...entry.item,
|
|
55
|
+
score: hybridWeight * entry.vectorScore + (1 - hybridWeight) * entry.textScore,
|
|
56
|
+
}));
|
|
57
|
+
}
|
|
58
|
+
|
|
59
|
+
export async function search(
|
|
60
|
+
query: string,
|
|
61
|
+
db: RagDB,
|
|
62
|
+
topK: number = 5,
|
|
63
|
+
threshold: number = 0,
|
|
64
|
+
hybridWeight: number = DEFAULT_HYBRID_WEIGHT,
|
|
65
|
+
enableReranking: boolean = false
|
|
66
|
+
): Promise<DedupedResult[]> {
|
|
67
|
+
const start = performance.now();
|
|
68
|
+
const queryEmbedding = await embed(query);
|
|
69
|
+
|
|
70
|
+
// Fetch more than topK to allow deduplication
|
|
71
|
+
const vectorResults = db.search(queryEmbedding, topK * 3);
|
|
72
|
+
|
|
73
|
+
// BM25 text search for keyword matching
|
|
74
|
+
let textResults: typeof vectorResults = [];
|
|
75
|
+
try {
|
|
76
|
+
textResults = db.textSearch(query, topK * 3);
|
|
77
|
+
} catch (err) {
|
|
78
|
+
log.debug(`FTS query failed, falling back to vector-only: ${err instanceof Error ? err.message : err}`, "search");
|
|
79
|
+
}
|
|
80
|
+
|
|
81
|
+
const merged = mergeHybridScores(vectorResults, textResults, hybridWeight);
|
|
82
|
+
|
|
83
|
+
// Deduplicate by file path, keeping the best score per file
|
|
84
|
+
const byFile = new Map<string, DedupedResult>();
|
|
85
|
+
|
|
86
|
+
for (const result of merged) {
|
|
87
|
+
if (threshold > 0 && result.score < threshold) continue;
|
|
88
|
+
|
|
89
|
+
const existing = byFile.get(result.path);
|
|
90
|
+
if (existing) {
|
|
91
|
+
if (result.score > existing.score) {
|
|
92
|
+
existing.score = result.score;
|
|
93
|
+
}
|
|
94
|
+
if (!existing.snippets.includes(result.snippet)) {
|
|
95
|
+
existing.snippets.push(result.snippet);
|
|
96
|
+
}
|
|
97
|
+
} else {
|
|
98
|
+
byFile.set(result.path, {
|
|
99
|
+
path: result.path,
|
|
100
|
+
score: result.score,
|
|
101
|
+
snippets: [result.snippet],
|
|
102
|
+
});
|
|
103
|
+
}
|
|
104
|
+
}
|
|
105
|
+
|
|
106
|
+
// Sort by score descending, take candidates for reranking
|
|
107
|
+
let results = Array.from(byFile.values())
|
|
108
|
+
.sort((a, b) => b.score - a.score)
|
|
109
|
+
.slice(0, enableReranking ? topK * 2 : topK);
|
|
110
|
+
|
|
111
|
+
// Cross-encoder reranking: re-score top candidates for precision
|
|
112
|
+
if (enableReranking && results.length > 0) {
|
|
113
|
+
try {
|
|
114
|
+
const passages = results.map((r) => r.snippets[0] ?? "");
|
|
115
|
+
const rerankScores = await rerank(query, passages);
|
|
116
|
+
results = results
|
|
117
|
+
.map((r, i) => ({ ...r, score: rerankScores[i] }))
|
|
118
|
+
.sort((a, b) => b.score - a.score)
|
|
119
|
+
.slice(0, topK);
|
|
120
|
+
} catch (err) {
|
|
121
|
+
log.warn(`Reranking failed, using hybrid scores: ${err instanceof Error ? err.message : err}`, "search");
|
|
122
|
+
results = results.slice(0, topK);
|
|
123
|
+
}
|
|
124
|
+
}
|
|
125
|
+
|
|
126
|
+
// Log query for analytics
|
|
127
|
+
const durationMs = Math.round(performance.now() - start);
|
|
128
|
+
db.logQuery(
|
|
129
|
+
query,
|
|
130
|
+
results.length,
|
|
131
|
+
results[0]?.score ?? null,
|
|
132
|
+
results[0]?.path ?? null,
|
|
133
|
+
durationMs
|
|
134
|
+
);
|
|
135
|
+
|
|
136
|
+
return results;
|
|
137
|
+
}
|
|
138
|
+
|
|
139
|
+
/**
|
|
140
|
+
* Chunk-level search: returns individual semantic chunks ranked by relevance.
|
|
141
|
+
* No file deduplication — two chunks from the same file can both appear.
|
|
142
|
+
*/
|
|
143
|
+
export async function searchChunks(
|
|
144
|
+
query: string,
|
|
145
|
+
db: RagDB,
|
|
146
|
+
topK: number = 8,
|
|
147
|
+
threshold: number = 0.3,
|
|
148
|
+
hybridWeight: number = DEFAULT_HYBRID_WEIGHT,
|
|
149
|
+
enableReranking: boolean = false
|
|
150
|
+
): Promise<ChunkResult[]> {
|
|
151
|
+
const start = performance.now();
|
|
152
|
+
const queryEmbedding = await embed(query);
|
|
153
|
+
|
|
154
|
+
const vectorResults = db.searchChunks(queryEmbedding, topK * 3);
|
|
155
|
+
|
|
156
|
+
let textResults: ChunkSearchResult[] = [];
|
|
157
|
+
try {
|
|
158
|
+
textResults = db.textSearchChunks(query, topK * 3);
|
|
159
|
+
} catch (err) {
|
|
160
|
+
log.debug(`FTS chunk query failed, falling back to vector-only: ${err instanceof Error ? err.message : err}`, "search");
|
|
161
|
+
}
|
|
162
|
+
|
|
163
|
+
let results = mergeHybridScores(vectorResults, textResults, hybridWeight)
|
|
164
|
+
.filter((r) => r.score >= threshold)
|
|
165
|
+
.sort((a, b) => b.score - a.score)
|
|
166
|
+
.slice(0, enableReranking ? topK * 2 : topK);
|
|
167
|
+
|
|
168
|
+
// Cross-encoder reranking: re-score top candidates for precision
|
|
169
|
+
if (enableReranking && results.length > 0) {
|
|
170
|
+
try {
|
|
171
|
+
const passages = results.map((r) => r.content);
|
|
172
|
+
const rerankScores = await rerank(query, passages);
|
|
173
|
+
results = results
|
|
174
|
+
.map((r, i) => ({ ...r, score: rerankScores[i] }))
|
|
175
|
+
.sort((a, b) => b.score - a.score)
|
|
176
|
+
.slice(0, topK);
|
|
177
|
+
} catch (err) {
|
|
178
|
+
log.warn(`Reranking failed, using hybrid scores: ${err instanceof Error ? err.message : err}`, "search");
|
|
179
|
+
results = results.slice(0, topK);
|
|
180
|
+
}
|
|
181
|
+
}
|
|
182
|
+
|
|
183
|
+
// Log query for analytics
|
|
184
|
+
const durationMs = Math.round(performance.now() - start);
|
|
185
|
+
db.logQuery(
|
|
186
|
+
query,
|
|
187
|
+
results.length,
|
|
188
|
+
results[0]?.score ?? null,
|
|
189
|
+
results[0]?.path ?? null,
|
|
190
|
+
durationMs
|
|
191
|
+
);
|
|
192
|
+
|
|
193
|
+
return results;
|
|
194
|
+
}
|
|
@@ -0,0 +1,99 @@
|
|
|
1
|
+
import { env, AutoTokenizer, AutoModelForSequenceClassification, type PreTrainedTokenizer, type PreTrainedModel } from "@huggingface/transformers";
|
|
2
|
+
import { join } from "node:path";
|
|
3
|
+
import { homedir } from "node:os";
|
|
4
|
+
import { rmSync } from "node:fs";
|
|
5
|
+
import { log } from "../utils/log";
|
|
6
|
+
|
|
7
|
+
// Share the same cache directory as embeddings
|
|
8
|
+
const CACHE_DIR = join(homedir(), ".cache", "local-rag", "models");
|
|
9
|
+
env.cacheDir = CACHE_DIR;
|
|
10
|
+
|
|
11
|
+
const RERANKER_MODEL_ID = "Xenova/ms-marco-MiniLM-L-6-v2";
|
|
12
|
+
|
|
13
|
+
let tokenizer: PreTrainedTokenizer | null = null;
|
|
14
|
+
let model: PreTrainedModel | null = null;
|
|
15
|
+
let loadingPromise: Promise<void> | null = null;
|
|
16
|
+
|
|
17
|
+
async function loadReranker(): Promise<void> {
|
|
18
|
+
if (tokenizer && model) return;
|
|
19
|
+
if (loadingPromise) return loadingPromise;
|
|
20
|
+
|
|
21
|
+
loadingPromise = (async () => {
|
|
22
|
+
const start = performance.now();
|
|
23
|
+
try {
|
|
24
|
+
tokenizer = await AutoTokenizer.from_pretrained(RERANKER_MODEL_ID, {
|
|
25
|
+
cache_dir: CACHE_DIR,
|
|
26
|
+
});
|
|
27
|
+
model = await AutoModelForSequenceClassification.from_pretrained(RERANKER_MODEL_ID, {
|
|
28
|
+
dtype: "fp32",
|
|
29
|
+
cache_dir: CACHE_DIR,
|
|
30
|
+
});
|
|
31
|
+
const elapsed = Math.round(performance.now() - start);
|
|
32
|
+
log.debug(`Reranker loaded in ${elapsed}ms`, "reranker");
|
|
33
|
+
} catch (err) {
|
|
34
|
+
// If the cached model is corrupted, delete it and retry once
|
|
35
|
+
const msg = (err as Error).message || "";
|
|
36
|
+
if (msg.includes("Protobuf parsing failed") || msg.includes("Load model")) {
|
|
37
|
+
const modelDir = join(CACHE_DIR, ...RERANKER_MODEL_ID.split("/"));
|
|
38
|
+
rmSync(modelDir, { recursive: true, force: true });
|
|
39
|
+
tokenizer = await AutoTokenizer.from_pretrained(RERANKER_MODEL_ID, {
|
|
40
|
+
cache_dir: CACHE_DIR,
|
|
41
|
+
});
|
|
42
|
+
model = await AutoModelForSequenceClassification.from_pretrained(RERANKER_MODEL_ID, {
|
|
43
|
+
dtype: "fp32",
|
|
44
|
+
cache_dir: CACHE_DIR,
|
|
45
|
+
});
|
|
46
|
+
} else {
|
|
47
|
+
loadingPromise = null;
|
|
48
|
+
throw err;
|
|
49
|
+
}
|
|
50
|
+
}
|
|
51
|
+
})();
|
|
52
|
+
|
|
53
|
+
return loadingPromise;
|
|
54
|
+
}
|
|
55
|
+
|
|
56
|
+
/**
|
|
57
|
+
* Score a list of (query, passage) pairs using a cross-encoder reranker.
|
|
58
|
+
* Returns relevance scores (higher = more relevant).
|
|
59
|
+
* Scores are sigmoid-normalized to [0, 1].
|
|
60
|
+
*/
|
|
61
|
+
export async function rerank(
|
|
62
|
+
query: string,
|
|
63
|
+
passages: string[],
|
|
64
|
+
): Promise<number[]> {
|
|
65
|
+
if (passages.length === 0) return [];
|
|
66
|
+
|
|
67
|
+
await loadReranker();
|
|
68
|
+
if (!tokenizer || !model) throw new Error("Reranker failed to load");
|
|
69
|
+
|
|
70
|
+
const scores: number[] = [];
|
|
71
|
+
|
|
72
|
+
// Process one at a time to avoid OOM on large result sets
|
|
73
|
+
for (const passage of passages) {
|
|
74
|
+
const inputs = tokenizer(query, {
|
|
75
|
+
text_pair: passage,
|
|
76
|
+
padding: true,
|
|
77
|
+
truncation: true,
|
|
78
|
+
max_length: 512,
|
|
79
|
+
});
|
|
80
|
+
|
|
81
|
+
const output = await model(inputs);
|
|
82
|
+
// Cross-encoder output is a single logit — apply sigmoid
|
|
83
|
+
const logit = output.logits.data[0] as number;
|
|
84
|
+
scores.push(sigmoid(logit));
|
|
85
|
+
}
|
|
86
|
+
|
|
87
|
+
return scores;
|
|
88
|
+
}
|
|
89
|
+
|
|
90
|
+
function sigmoid(x: number): number {
|
|
91
|
+
return 1 / (1 + Math.exp(-x));
|
|
92
|
+
}
|
|
93
|
+
|
|
94
|
+
/** Reset the singleton — only for testing */
|
|
95
|
+
export function resetReranker(): void {
|
|
96
|
+
tokenizer = null;
|
|
97
|
+
model = null;
|
|
98
|
+
loadingPromise = null;
|
|
99
|
+
}
|
|
@@ -0,0 +1,27 @@
|
|
|
1
|
+
/**
|
|
2
|
+
* Utilities for the find_usages feature.
|
|
3
|
+
*
|
|
4
|
+
* find_usages works at query time rather than pre-indexing call sites:
|
|
5
|
+
* 1. FTS search finds chunks containing the symbol name.
|
|
6
|
+
* 2. Defining files are excluded via the file_exports table.
|
|
7
|
+
* 3. Within each matching chunk, a word-boundary regex locates the exact line.
|
|
8
|
+
* 4. Absolute line numbers are computed from the chunk's stored start_line.
|
|
9
|
+
*
|
|
10
|
+
* This avoids a full re-index pass and handles all supported file types since
|
|
11
|
+
* all chunks are in the FTS index regardless of language.
|
|
12
|
+
*/
|
|
13
|
+
|
|
14
|
+
export function escapeRegex(s: string): string {
|
|
15
|
+
return s.replace(/[.*+?^${}()|[\]\\]/g, "\\$&");
|
|
16
|
+
}
|
|
17
|
+
|
|
18
|
+
/**
|
|
19
|
+
* Sanitize a user query for FTS5 MATCH by quoting each token.
|
|
20
|
+
* FTS5 treats bare +, -, *, AND, OR, NOT, NEAR, ( ) as operators.
|
|
21
|
+
* Wrapping each token in double quotes forces literal matching.
|
|
22
|
+
*/
|
|
23
|
+
export function sanitizeFTS(query: string): string {
|
|
24
|
+
const tokens = query.split(/\s+/).filter(Boolean);
|
|
25
|
+
if (tokens.length === 0) return '""';
|
|
26
|
+
return tokens.map(t => `"${t.replace(/"/g, '""')}"`).join(" ");
|
|
27
|
+
}
|
|
@@ -0,0 +1,126 @@
|
|
|
1
|
+
import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
|
|
2
|
+
import { StdioServerTransport } from "@modelcontextprotocol/sdk/server/stdio.js";
|
|
3
|
+
import { resolve } from "path";
|
|
4
|
+
import { homedir } from "os";
|
|
5
|
+
import { RagDB } from "../db";
|
|
6
|
+
import { loadConfig } from "../config";
|
|
7
|
+
import { indexDirectory } from "../indexing/indexer";
|
|
8
|
+
import { startWatcher, type Watcher } from "../indexing/watcher";
|
|
9
|
+
import { discoverSessions } from "../conversation/parser";
|
|
10
|
+
import { indexConversation, startConversationTail } from "../conversation/indexer";
|
|
11
|
+
import { registerAllTools } from "../tools";
|
|
12
|
+
import { log } from "../utils/log";
|
|
13
|
+
|
|
14
|
+
// Read version from package.json at module load time
|
|
15
|
+
const { version } = await import("../../package.json");
|
|
16
|
+
|
|
17
|
+
// Lazy-init DB per project directory — keep all open to avoid
|
|
18
|
+
// closing a DB that background tasks (auto-index, watcher) still use.
|
|
19
|
+
const dbMap = new Map<string, RagDB>();
|
|
20
|
+
|
|
21
|
+
function getDB(projectDir: string): RagDB {
|
|
22
|
+
const resolved = resolve(projectDir);
|
|
23
|
+
let db = dbMap.get(resolved);
|
|
24
|
+
if (db) return db;
|
|
25
|
+
db = new RagDB(resolved);
|
|
26
|
+
dbMap.set(resolved, db);
|
|
27
|
+
return db;
|
|
28
|
+
}
|
|
29
|
+
|
|
30
|
+
export async function startServer() {
|
|
31
|
+
const server = new McpServer({
|
|
32
|
+
name: "local-rag",
|
|
33
|
+
version,
|
|
34
|
+
});
|
|
35
|
+
|
|
36
|
+
// Register all MCP tools
|
|
37
|
+
registerAllTools(server, getDB);
|
|
38
|
+
|
|
39
|
+
// Auto-index on startup + start file watcher
|
|
40
|
+
const startupDir = process.env.RAG_PROJECT_DIR || process.cwd();
|
|
41
|
+
|
|
42
|
+
const isHomeDirTrap = resolve(startupDir) === homedir();
|
|
43
|
+
if (isHomeDirTrap) {
|
|
44
|
+
process.stderr.write(
|
|
45
|
+
`[local-rag] WARNING: project directory is your home folder (${startupDir}).\n` +
|
|
46
|
+
`[local-rag] Skipping auto-index and file watcher. Set RAG_PROJECT_DIR to your project path.\n` +
|
|
47
|
+
`[local-rag] Example: "env": { "RAG_PROJECT_DIR": "/path/to/your/project" }\n`
|
|
48
|
+
);
|
|
49
|
+
}
|
|
50
|
+
const startupDb = getDB(startupDir);
|
|
51
|
+
const startupConfig = await loadConfig(startupDir);
|
|
52
|
+
|
|
53
|
+
let watcher: Watcher | null = null;
|
|
54
|
+
let convWatcher: Watcher | null = null;
|
|
55
|
+
|
|
56
|
+
if (!isHomeDirTrap) {
|
|
57
|
+
// Index in background — don't block server startup
|
|
58
|
+
indexDirectory(startupDir, startupDb, startupConfig, (msg) => {
|
|
59
|
+
process.stderr.write(`[local-rag] ${msg}\n`);
|
|
60
|
+
}).then((result) => {
|
|
61
|
+
process.stderr.write(
|
|
62
|
+
`[local-rag] Startup index: ${result.indexed} indexed, ${result.skipped} skipped, ${result.pruned} pruned\n`
|
|
63
|
+
);
|
|
64
|
+
|
|
65
|
+
// Start watching after initial index completes
|
|
66
|
+
watcher = startWatcher(startupDir, startupDb, startupConfig, (msg) => {
|
|
67
|
+
process.stderr.write(`[local-rag] ${msg}\n`);
|
|
68
|
+
});
|
|
69
|
+
}).catch((err) => {
|
|
70
|
+
log.warn(`Startup indexing failed: ${err instanceof Error ? err.message : err}`, "server");
|
|
71
|
+
});
|
|
72
|
+
}
|
|
73
|
+
|
|
74
|
+
// Start conversation tailing — find and tail the current session's JSONL
|
|
75
|
+
const sessions = discoverSessions(startupDir);
|
|
76
|
+
if (sessions.length > 0) {
|
|
77
|
+
// Tail the most recent session (likely the current one)
|
|
78
|
+
const currentSession = sessions[0];
|
|
79
|
+
process.stderr.write(`[local-rag] Indexing conversation: ${currentSession.sessionId.slice(0, 8)}...\n`);
|
|
80
|
+
|
|
81
|
+
convWatcher = startConversationTail(
|
|
82
|
+
currentSession.jsonlPath,
|
|
83
|
+
currentSession.sessionId,
|
|
84
|
+
startupDb,
|
|
85
|
+
(msg) => process.stderr.write(`[local-rag] ${msg}\n`)
|
|
86
|
+
);
|
|
87
|
+
|
|
88
|
+
// Also index any older sessions that haven't been indexed yet
|
|
89
|
+
for (const session of sessions.slice(1)) {
|
|
90
|
+
const existing = startupDb.getSession(session.sessionId);
|
|
91
|
+
if (!existing || existing.mtime < session.mtime) {
|
|
92
|
+
indexConversation(
|
|
93
|
+
session.jsonlPath,
|
|
94
|
+
session.sessionId,
|
|
95
|
+
startupDb
|
|
96
|
+
).then((result) => {
|
|
97
|
+
if (result.turnsIndexed > 0) {
|
|
98
|
+
process.stderr.write(
|
|
99
|
+
`[local-rag] Indexed past session ${session.sessionId.slice(0, 8)}...: ${result.turnsIndexed} turns\n`
|
|
100
|
+
);
|
|
101
|
+
}
|
|
102
|
+
}).catch((err) => {
|
|
103
|
+
log.warn(`Failed to index session ${session.sessionId.slice(0, 8)}: ${err instanceof Error ? err.message : err}`, "conversation");
|
|
104
|
+
});
|
|
105
|
+
}
|
|
106
|
+
}
|
|
107
|
+
}
|
|
108
|
+
|
|
109
|
+
// Graceful shutdown
|
|
110
|
+
function cleanup() {
|
|
111
|
+
process.stderr.write("[local-rag] Shutting down...\n");
|
|
112
|
+
if (watcher) watcher.close();
|
|
113
|
+
if (convWatcher) convWatcher.close();
|
|
114
|
+
for (const d of dbMap.values()) d.close();
|
|
115
|
+
dbMap.clear();
|
|
116
|
+
process.exit(0);
|
|
117
|
+
}
|
|
118
|
+
|
|
119
|
+
process.on("SIGINT", cleanup);
|
|
120
|
+
process.on("SIGTERM", cleanup);
|
|
121
|
+
process.on("SIGHUP", cleanup);
|
|
122
|
+
|
|
123
|
+
// Start server
|
|
124
|
+
const transport = new StdioServerTransport();
|
|
125
|
+
await server.connect(transport);
|
|
126
|
+
}
|
|
@@ -0,0 +1,58 @@
|
|
|
1
|
+
import { z } from "zod";
|
|
2
|
+
import type { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
|
|
3
|
+
import { type GetDB, resolveProject } from "./index";
|
|
4
|
+
|
|
5
|
+
export function registerAnalyticsTools(server: McpServer, getDB: GetDB) {
|
|
6
|
+
server.tool(
|
|
7
|
+
"search_analytics",
|
|
8
|
+
"Show search usage analytics: query counts, zero-result queries, low-relevance queries, top searched terms.",
|
|
9
|
+
{
|
|
10
|
+
directory: z
|
|
11
|
+
.string()
|
|
12
|
+
.optional()
|
|
13
|
+
.describe("Project directory. Defaults to RAG_PROJECT_DIR env or cwd"),
|
|
14
|
+
days: z
|
|
15
|
+
.number()
|
|
16
|
+
.optional()
|
|
17
|
+
.default(30)
|
|
18
|
+
.describe("Number of days to look back (default: 30)"),
|
|
19
|
+
},
|
|
20
|
+
async ({ directory, days }) => {
|
|
21
|
+
const { db: ragDb } = await resolveProject(directory, getDB);
|
|
22
|
+
const analytics = ragDb.getAnalytics(days);
|
|
23
|
+
|
|
24
|
+
const lines: string[] = [
|
|
25
|
+
`Search analytics (last ${days} days):`,
|
|
26
|
+
` Total queries: ${analytics.totalQueries}`,
|
|
27
|
+
` Avg results: ${analytics.avgResultCount.toFixed(1)}`,
|
|
28
|
+
` Avg top score: ${analytics.avgTopScore?.toFixed(2) ?? "n/a"}`,
|
|
29
|
+
` Zero-result rate: ${analytics.totalQueries > 0 ? ((analytics.zeroResultQueries.reduce((s, q) => s + q.count, 0) / analytics.totalQueries) * 100).toFixed(0) : 0}%`,
|
|
30
|
+
];
|
|
31
|
+
|
|
32
|
+
if (analytics.topSearchedTerms.length > 0) {
|
|
33
|
+
lines.push("", "Top searches:");
|
|
34
|
+
for (const t of analytics.topSearchedTerms) {
|
|
35
|
+
lines.push(` - "${t.query}" (${t.count}×)`);
|
|
36
|
+
}
|
|
37
|
+
}
|
|
38
|
+
|
|
39
|
+
if (analytics.zeroResultQueries.length > 0) {
|
|
40
|
+
lines.push("", "Zero-result queries (consider indexing these topics):");
|
|
41
|
+
for (const q of analytics.zeroResultQueries) {
|
|
42
|
+
lines.push(` - "${q.query}" (${q.count}×)`);
|
|
43
|
+
}
|
|
44
|
+
}
|
|
45
|
+
|
|
46
|
+
if (analytics.lowScoreQueries.length > 0) {
|
|
47
|
+
lines.push("", "Low-relevance queries (top score < 0.3):");
|
|
48
|
+
for (const q of analytics.lowScoreQueries) {
|
|
49
|
+
lines.push(` - "${q.query}" (score: ${q.topScore.toFixed(2)})`);
|
|
50
|
+
}
|
|
51
|
+
}
|
|
52
|
+
|
|
53
|
+
return {
|
|
54
|
+
content: [{ type: "text" as const, text: lines.join("\n") }],
|
|
55
|
+
};
|
|
56
|
+
}
|
|
57
|
+
);
|
|
58
|
+
}
|