brainbank 0.1.0-beta.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/LICENSE +21 -0
- package/README.md +155 -0
- package/assets/architecture.png +0 -0
- package/bin/brainbank +18 -0
- package/bin/brainbank-mcp +19 -0
- package/dist/chunk-3YBCD6DI.js +117 -0
- package/dist/chunk-3YBCD6DI.js.map +1 -0
- package/dist/chunk-63GBCDS5.js +3249 -0
- package/dist/chunk-63GBCDS5.js.map +1 -0
- package/dist/chunk-DMFMTOHF.js +123 -0
- package/dist/chunk-DMFMTOHF.js.map +1 -0
- package/dist/chunk-FQYKWB2Q.js +136 -0
- package/dist/chunk-FQYKWB2Q.js.map +1 -0
- package/dist/chunk-IMJJ2VEM.js +74 -0
- package/dist/chunk-IMJJ2VEM.js.map +1 -0
- package/dist/chunk-M744PCJQ.js +43 -0
- package/dist/chunk-M744PCJQ.js.map +1 -0
- package/dist/chunk-O3J6ZIXK.js +82 -0
- package/dist/chunk-O3J6ZIXK.js.map +1 -0
- package/dist/chunk-OPH7GZ7U.js +124 -0
- package/dist/chunk-OPH7GZ7U.js.map +1 -0
- package/dist/chunk-PXEWQMN7.js +89 -0
- package/dist/chunk-PXEWQMN7.js.map +1 -0
- package/dist/chunk-RDQYDLYZ.js +69 -0
- package/dist/chunk-RDQYDLYZ.js.map +1 -0
- package/dist/chunk-VIIHPCC4.js +254 -0
- package/dist/chunk-VIIHPCC4.js.map +1 -0
- package/dist/chunk-WCQVDF3K.js +14 -0
- package/dist/chunk-WCQVDF3K.js.map +1 -0
- package/dist/cli.d.ts +1 -0
- package/dist/cli.js +3076 -0
- package/dist/cli.js.map +1 -0
- package/dist/haiku-expander-YRSIPGKP.js +8 -0
- package/dist/haiku-expander-YRSIPGKP.js.map +1 -0
- package/dist/haiku-pruner-SHAXUPY6.js +8 -0
- package/dist/haiku-pruner-SHAXUPY6.js.map +1 -0
- package/dist/http-server-QUXHLWUM.js +9 -0
- package/dist/http-server-QUXHLWUM.js.map +1 -0
- package/dist/index.d.ts +2161 -0
- package/dist/index.js +357 -0
- package/dist/index.js.map +1 -0
- package/dist/local-embedding-NZQTILGV.js +8 -0
- package/dist/local-embedding-NZQTILGV.js.map +1 -0
- package/dist/mcp.d.ts +2 -0
- package/dist/mcp.js +334 -0
- package/dist/mcp.js.map +1 -0
- package/dist/openai-embedding-ZP5TSUJG.js +8 -0
- package/dist/openai-embedding-ZP5TSUJG.js.map +1 -0
- package/dist/perplexity-context-embedding-GI5PHE6X.js +9 -0
- package/dist/perplexity-context-embedding-GI5PHE6X.js.map +1 -0
- package/dist/perplexity-embedding-KZRYGJRC.js +10 -0
- package/dist/perplexity-embedding-KZRYGJRC.js.map +1 -0
- package/dist/plugin-IKQ6IRSJ.js +32 -0
- package/dist/plugin-IKQ6IRSJ.js.map +1 -0
- package/dist/resolve-ASGLBNUC.js +10 -0
- package/dist/resolve-ASGLBNUC.js.map +1 -0
- package/dist/stats-tui-ZY2NQSEA.js +1904 -0
- package/dist/stats-tui-ZY2NQSEA.js.map +1 -0
- package/package.json +96 -0
- package/src/brainbank.ts +617 -0
- package/src/cli/commands/collection.ts +77 -0
- package/src/cli/commands/context.ts +179 -0
- package/src/cli/commands/daemon.ts +100 -0
- package/src/cli/commands/docs.ts +71 -0
- package/src/cli/commands/files.ts +69 -0
- package/src/cli/commands/help.ts +77 -0
- package/src/cli/commands/index.ts +482 -0
- package/src/cli/commands/kv.ts +140 -0
- package/src/cli/commands/mcp-export.ts +273 -0
- package/src/cli/commands/mcp.ts +6 -0
- package/src/cli/commands/reembed.ts +30 -0
- package/src/cli/commands/scan.ts +336 -0
- package/src/cli/commands/search.ts +203 -0
- package/src/cli/commands/stats.ts +68 -0
- package/src/cli/commands/status.ts +47 -0
- package/src/cli/commands/watch.ts +47 -0
- package/src/cli/factory/brain-context.ts +43 -0
- package/src/cli/factory/builtin-registration.ts +87 -0
- package/src/cli/factory/config-loader.ts +77 -0
- package/src/cli/factory/index.ts +69 -0
- package/src/cli/factory/plugin-loader.ts +325 -0
- package/src/cli/index.ts +71 -0
- package/src/cli/server-client.ts +178 -0
- package/src/cli/tui/index-tui.tsx +667 -0
- package/src/cli/tui/stats-data.ts +523 -0
- package/src/cli/tui/stats-search.ts +262 -0
- package/src/cli/tui/stats-tui.tsx +1465 -0
- package/src/cli/tui/tree-scanner.ts +650 -0
- package/src/cli/utils.ts +137 -0
- package/src/config.ts +49 -0
- package/src/constants.ts +21 -0
- package/src/db/adapter.ts +112 -0
- package/src/db/metadata.ts +130 -0
- package/src/db/migrations.ts +66 -0
- package/src/db/sqlite-adapter.ts +218 -0
- package/src/db/tracker.ts +91 -0
- package/src/engine/index-api.ts +81 -0
- package/src/engine/reembed.ts +206 -0
- package/src/engine/search-api.ts +218 -0
- package/src/index.ts +154 -0
- package/src/lib/fts.ts +57 -0
- package/src/lib/languages.ts +180 -0
- package/src/lib/logger.ts +126 -0
- package/src/lib/math.ts +87 -0
- package/src/lib/provider-key.ts +20 -0
- package/src/lib/prune.ts +71 -0
- package/src/lib/rrf.ts +133 -0
- package/src/lib/write-lock.ts +108 -0
- package/src/mcp/mcp-server.ts +195 -0
- package/src/mcp/workspace-factory.ts +68 -0
- package/src/mcp/workspace-pool.ts +224 -0
- package/src/plugin.ts +381 -0
- package/src/providers/embeddings/embedding-worker-thread.ts +95 -0
- package/src/providers/embeddings/embedding-worker.ts +141 -0
- package/src/providers/embeddings/local-embedding.ts +115 -0
- package/src/providers/embeddings/openai-embedding.ts +167 -0
- package/src/providers/embeddings/perplexity-context-embedding.ts +195 -0
- package/src/providers/embeddings/perplexity-embedding.ts +165 -0
- package/src/providers/embeddings/resolve.ts +34 -0
- package/src/providers/pruners/haiku-expander.ts +166 -0
- package/src/providers/pruners/haiku-pruner.ts +112 -0
- package/src/providers/vector/hnsw-index.ts +174 -0
- package/src/providers/vector/hnsw-loader.ts +129 -0
- package/src/search/bm25-boost.ts +69 -0
- package/src/search/context-builder.ts +251 -0
- package/src/search/keyword/composite-bm25-search.ts +47 -0
- package/src/search/types.ts +37 -0
- package/src/search/vector/composite-vector-search.ts +61 -0
- package/src/search/vector/mmr.ts +64 -0
- package/src/services/collection.ts +384 -0
- package/src/services/daemon.ts +87 -0
- package/src/services/http-server.ts +336 -0
- package/src/services/kv-service.ts +64 -0
- package/src/services/plugin-registry.ts +77 -0
- package/src/services/watch.ts +340 -0
- package/src/services/webhook-server.ts +100 -0
- package/src/types.ts +493 -0
|
@@ -0,0 +1,166 @@
|
|
|
1
|
+
/**
|
|
2
|
+
* BrainBank — Haiku Expander
|
|
3
|
+
*
|
|
4
|
+
* LLM-powered context expansion using Anthropic's Haiku 4.5 model.
|
|
5
|
+
* After search + pruning, reviews a manifest of available chunks
|
|
6
|
+
* and requests additional IDs to include.
|
|
7
|
+
*
|
|
8
|
+
* Flow:
|
|
9
|
+
* 1. Receives lightweight manifest (~20 chars per chunk)
|
|
10
|
+
* 2. Haiku selects additional chunk IDs (just numbers, fast)
|
|
11
|
+
* 3. Caller fetches those chunks from DB and splices into results
|
|
12
|
+
*
|
|
13
|
+
* Designed for minimal token usage:
|
|
14
|
+
* - Input: ~2,000-3,000 tokens (manifest)
|
|
15
|
+
* - Output: ~50-100 tokens (ID array)
|
|
16
|
+
* - Cost: ~$0.001 per call
|
|
17
|
+
* - Latency: ~300-600ms
|
|
18
|
+
*
|
|
19
|
+
* Fail-open: any error returns empty array (no expansion).
|
|
20
|
+
*/
|
|
21
|
+
|
|
22
|
+
import type { Expander, ExpanderManifestItem, ExpanderResult } from '@/types.ts';
|
|
23
|
+
|
|
24
|
+
const DEFAULT_MODEL = 'claude-haiku-4-5-20251001';
|
|
25
|
+
|
|
26
|
+
export interface HaikuExpanderOptions {
|
|
27
|
+
/** Anthropic API key. Falls back to ANTHROPIC_API_KEY env var. */
|
|
28
|
+
apiKey?: string;
|
|
29
|
+
/** Model to use. Default: claude-haiku-4-5-20251001 */
|
|
30
|
+
model?: string;
|
|
31
|
+
}
|
|
32
|
+
|
|
33
|
+
export class HaikuExpander implements Expander {
|
|
34
|
+
private readonly _apiKey: string;
|
|
35
|
+
private readonly _model: string;
|
|
36
|
+
|
|
37
|
+
constructor(options: HaikuExpanderOptions = {}) {
|
|
38
|
+
this._apiKey = options.apiKey ?? process.env.ANTHROPIC_API_KEY ?? '';
|
|
39
|
+
this._model = options.model ?? DEFAULT_MODEL;
|
|
40
|
+
|
|
41
|
+
if (!this._apiKey) {
|
|
42
|
+
throw new Error(
|
|
43
|
+
'HaikuExpander: No API key provided. Set ANTHROPIC_API_KEY env var or pass apiKey option.',
|
|
44
|
+
);
|
|
45
|
+
}
|
|
46
|
+
}
|
|
47
|
+
|
|
48
|
+
async expand(
|
|
49
|
+
query: string,
|
|
50
|
+
currentIds: number[],
|
|
51
|
+
manifest: ExpanderManifestItem[],
|
|
52
|
+
): Promise<ExpanderResult> {
|
|
53
|
+
if (manifest.length === 0) return { ids: [] };
|
|
54
|
+
|
|
55
|
+
// Filter out chunks already in results
|
|
56
|
+
const currentSet = new Set(currentIds);
|
|
57
|
+
const available = manifest.filter(m => !currentSet.has(m.id));
|
|
58
|
+
if (available.length === 0) return { ids: [] };
|
|
59
|
+
|
|
60
|
+
// Split manifest into priority (import-graph neighbors) and general
|
|
61
|
+
const priorityChunks = available.filter(m => m.priority);
|
|
62
|
+
const otherChunks = available.filter(m => !m.priority);
|
|
63
|
+
|
|
64
|
+
const currentSummary = manifest
|
|
65
|
+
.filter(m => currentSet.has(m.id))
|
|
66
|
+
.map(m => `#${m.id} ${m.filePath} | ${m.chunkType} ${m.name}`)
|
|
67
|
+
.join('\n');
|
|
68
|
+
|
|
69
|
+
// Build manifest sections
|
|
70
|
+
let manifestSection = '';
|
|
71
|
+
if (priorityChunks.length > 0) {
|
|
72
|
+
const prioLines = priorityChunks.map(m =>
|
|
73
|
+
`#${m.id} ${m.filePath} | ${m.chunkType} ${m.name} ${m.lines}`
|
|
74
|
+
).join('\n');
|
|
75
|
+
manifestSection += `DEPENDENCY chunks (imported by or importing the search result files):\n${prioLines}\n\n`;
|
|
76
|
+
}
|
|
77
|
+
if (otherChunks.length > 0) {
|
|
78
|
+
const otherLines = otherChunks.map(m =>
|
|
79
|
+
`#${m.id} ${m.filePath} | ${m.chunkType} ${m.name} ${m.lines}`
|
|
80
|
+
).join('\n');
|
|
81
|
+
manifestSection += `Other available chunks:\n${otherLines}`;
|
|
82
|
+
}
|
|
83
|
+
|
|
84
|
+
const prompt =
|
|
85
|
+
`Task: "${query}"\n\n` +
|
|
86
|
+
`Already included chunks:\n${currentSummary}\n\n` +
|
|
87
|
+
`${manifestSection}\n\n` +
|
|
88
|
+
`You are a code context expander. The search already found the "included" chunks above.\n` +
|
|
89
|
+
`Review the available chunks and select any that would help an AI agent complete the task.\n\n` +
|
|
90
|
+
`Rules:\n` +
|
|
91
|
+
`- STRONGLY PREFER dependency chunks — they are structurally connected to the search results via imports\n` +
|
|
92
|
+
`- Select type definitions, interfaces, models, or configs needed to understand included code\n` +
|
|
93
|
+
`- Select initialization or setup code if the task involves debugging or modifying a feature\n` +
|
|
94
|
+
`- Do NOT select test files, documentation, or unrelated utilities\n` +
|
|
95
|
+
`- Be selective: only include chunks that fill clear gaps. Quality over quantity.\n` +
|
|
96
|
+
`- If nothing useful is available, return an empty ids array\n\n` +
|
|
97
|
+
`Respond with ONLY a JSON object:\n` +
|
|
98
|
+
`{ "ids": [42, 17, 89], "note": "Brief 1-2 sentence observation about the codebase relevant to the task" }\n\n` +
|
|
99
|
+
`The "note" is optional — use it to mention things like missing files, architectural patterns, ` +
|
|
100
|
+
`deprecated modules, or important relationships you noticed. If nothing notable, omit it.\n` +
|
|
101
|
+
`If nothing to add: { "ids": [] }`;
|
|
102
|
+
|
|
103
|
+
try {
|
|
104
|
+
const response = await fetch('https://api.anthropic.com/v1/messages', {
|
|
105
|
+
method: 'POST',
|
|
106
|
+
headers: {
|
|
107
|
+
'Content-Type': 'application/json',
|
|
108
|
+
'x-api-key': this._apiKey,
|
|
109
|
+
'anthropic-version': '2023-06-01',
|
|
110
|
+
},
|
|
111
|
+
body: JSON.stringify({
|
|
112
|
+
model: this._model,
|
|
113
|
+
max_tokens: 512,
|
|
114
|
+
messages: [{
|
|
115
|
+
role: 'user',
|
|
116
|
+
content: prompt,
|
|
117
|
+
}],
|
|
118
|
+
}),
|
|
119
|
+
});
|
|
120
|
+
|
|
121
|
+
if (!response.ok) {
|
|
122
|
+
return { ids: [] };
|
|
123
|
+
}
|
|
124
|
+
|
|
125
|
+
const data = await response.json() as {
|
|
126
|
+
content: { type: string; text: string }[];
|
|
127
|
+
};
|
|
128
|
+
|
|
129
|
+
const text = data.content?.[0]?.text ?? '';
|
|
130
|
+
return this._parseResponse(text, available);
|
|
131
|
+
} catch {
|
|
132
|
+
return { ids: [] };
|
|
133
|
+
}
|
|
134
|
+
}
|
|
135
|
+
|
|
136
|
+
/** Parse Haiku response — handles both `{ ids, note }` and bare `[...]` formats. */
|
|
137
|
+
private _parseResponse(text: string, available: ExpanderManifestItem[]): ExpanderResult {
|
|
138
|
+
const validIds = new Set(available.map(m => m.id));
|
|
139
|
+
|
|
140
|
+
// Try JSON object first: { "ids": [...], "note": "..." }
|
|
141
|
+
const objMatch = text.match(/\{[\s\S]*"ids"\s*:\s*\[[\d\s,]*\][\s\S]*\}/);
|
|
142
|
+
if (objMatch) {
|
|
143
|
+
try {
|
|
144
|
+
const parsed = JSON.parse(objMatch[0]) as { ids: number[]; note?: string };
|
|
145
|
+
const ids = parsed.ids.filter(id => validIds.has(id));
|
|
146
|
+
const note = parsed.note?.trim() || undefined;
|
|
147
|
+
return { ids, note };
|
|
148
|
+
} catch {
|
|
149
|
+
// Fall through to array parsing
|
|
150
|
+
}
|
|
151
|
+
}
|
|
152
|
+
|
|
153
|
+
// Fallback: bare array [42, 17, 89]
|
|
154
|
+
const arrMatch = text.match(/\[[\d\s,]*\]/);
|
|
155
|
+
if (arrMatch) {
|
|
156
|
+
const ids = (JSON.parse(arrMatch[0]) as number[]).filter(id => validIds.has(id));
|
|
157
|
+
return { ids };
|
|
158
|
+
}
|
|
159
|
+
|
|
160
|
+
return { ids: [] };
|
|
161
|
+
}
|
|
162
|
+
|
|
163
|
+
async close(): Promise<void> {
|
|
164
|
+
// No resources to release (stateless HTTP)
|
|
165
|
+
}
|
|
166
|
+
}
|
|
@@ -0,0 +1,112 @@
|
|
|
1
|
+
/**
|
|
2
|
+
* BrainBank — Haiku Pruner
|
|
3
|
+
*
|
|
4
|
+
* LLM-based noise filter using Anthropic's Haiku 4.5 model.
|
|
5
|
+
* Binary classification: for each search result, Haiku decides
|
|
6
|
+
* "relevant" or "noise" based on filePath, metadata, and full
|
|
7
|
+
* file content (capped at ~8K chars per item by prune.ts).
|
|
8
|
+
*
|
|
9
|
+
* Latency: ~300-600ms.
|
|
10
|
+
*/
|
|
11
|
+
|
|
12
|
+
import type { Pruner, PrunerItem } from '@/types.ts';
|
|
13
|
+
|
|
14
|
+
const DEFAULT_MODEL = 'claude-haiku-4-5-20251001';
|
|
15
|
+
|
|
16
|
+
export interface HaikuPrunerOptions {
|
|
17
|
+
/** Anthropic API key. Falls back to ANTHROPIC_API_KEY env var. */
|
|
18
|
+
apiKey?: string;
|
|
19
|
+
/** Model to use. Default: claude-haiku-4-5-20251001 */
|
|
20
|
+
model?: string;
|
|
21
|
+
}
|
|
22
|
+
|
|
23
|
+
export class HaikuPruner implements Pruner {
|
|
24
|
+
private readonly _apiKey: string;
|
|
25
|
+
private readonly _model: string;
|
|
26
|
+
|
|
27
|
+
constructor(options: HaikuPrunerOptions = {}) {
|
|
28
|
+
this._apiKey = options.apiKey ?? process.env.ANTHROPIC_API_KEY ?? '';
|
|
29
|
+
this._model = options.model ?? DEFAULT_MODEL;
|
|
30
|
+
|
|
31
|
+
if (!this._apiKey) {
|
|
32
|
+
throw new Error(
|
|
33
|
+
'HaikuPruner: No API key provided. Set ANTHROPIC_API_KEY env var or pass apiKey option.',
|
|
34
|
+
);
|
|
35
|
+
}
|
|
36
|
+
}
|
|
37
|
+
|
|
38
|
+
async prune(query: string, items: PrunerItem[]): Promise<number[]> {
|
|
39
|
+
if (items.length === 0) return [];
|
|
40
|
+
if (items.length === 1) return [items[0].id];
|
|
41
|
+
|
|
42
|
+
const itemLines = items.map(item => {
|
|
43
|
+
// Only show useful metadata fields (skip raw scores, IDs, large arrays)
|
|
44
|
+
const SKIP_KEYS = new Set(['id', 'chunkIds', 'rrfScore', 'filePath']);
|
|
45
|
+
const meta = Object.entries(item.metadata)
|
|
46
|
+
.filter(([k, v]) => v !== undefined && v !== null && !SKIP_KEYS.has(k))
|
|
47
|
+
.map(([k, v]) => `${k}=${v}`)
|
|
48
|
+
.join(' | ');
|
|
49
|
+
return `#${item.id} ${item.filePath} | ${meta}\n${item.preview}`;
|
|
50
|
+
}).join('\n---\n');
|
|
51
|
+
|
|
52
|
+
const prompt =
|
|
53
|
+
`Query: "${query}"\n\nSearch results (full file content):\n${itemLines}\n\n` +
|
|
54
|
+
`You are a precision search filter and ranker. You have the FULL source code of each file.\n` +
|
|
55
|
+
`Return a JSON array of #IDs to KEEP, ordered by relevance (most relevant FIRST).\n\n` +
|
|
56
|
+
`Rules:\n` +
|
|
57
|
+
`- Understand the SPECIFIC system/feature the query targets. Don't match on shared vocabulary alone.\n` +
|
|
58
|
+
` Example: "snackbar toast notification" targets the toast popup system, NOT a notification center/bell icon.\n` +
|
|
59
|
+
`- KEEP files that directly implement, define types for, or configure the queried system.\n` +
|
|
60
|
+
`- KEEP files where the queried system is mounted, initialized, or composed into a workflow.\n` +
|
|
61
|
+
`- DROP files that only CONSUME the system (e.g. a component that calls showNotification once but has 400 lines of unrelated logic).\n` +
|
|
62
|
+
`- DROP files that implement a DIFFERENT system sharing similar vocabulary.\n` +
|
|
63
|
+
`- DROP infrastructure/boilerplate (font loaders, theme definitions, CSS-only layout shells) unless they directly configure the queried feature.\n` +
|
|
64
|
+
`- Aim for 40-70% keep rate. Returning fewer, highly relevant files is BETTER than returning many tangential ones.\n` +
|
|
65
|
+
`- ORDER: core implementation → types/config → mount points → peripheral.\n\n` +
|
|
66
|
+
`Respond with ONLY the JSON array. Example: [3, 0, 5, 1]`;
|
|
67
|
+
|
|
68
|
+
try {
|
|
69
|
+
const response = await fetch('https://api.anthropic.com/v1/messages', {
|
|
70
|
+
method: 'POST',
|
|
71
|
+
headers: {
|
|
72
|
+
'Content-Type': 'application/json',
|
|
73
|
+
'x-api-key': this._apiKey,
|
|
74
|
+
'anthropic-version': '2023-06-01',
|
|
75
|
+
},
|
|
76
|
+
body: JSON.stringify({
|
|
77
|
+
model: this._model,
|
|
78
|
+
max_tokens: 512,
|
|
79
|
+
messages: [{
|
|
80
|
+
role: 'user',
|
|
81
|
+
content: prompt,
|
|
82
|
+
}],
|
|
83
|
+
}),
|
|
84
|
+
});
|
|
85
|
+
|
|
86
|
+
if (!response.ok) {
|
|
87
|
+
// API error → fail-open, return all
|
|
88
|
+
return items.map(i => i.id);
|
|
89
|
+
}
|
|
90
|
+
|
|
91
|
+
const data = await response.json() as {
|
|
92
|
+
content: { type: string; text: string }[];
|
|
93
|
+
};
|
|
94
|
+
|
|
95
|
+
const text = data.content?.[0]?.text ?? '';
|
|
96
|
+
// Haiku may wrap in ```json ... ``` — extract any JSON array
|
|
97
|
+
const match = text.match(/\[[\d\s,]+\]/);
|
|
98
|
+
if (!match) return items.map(i => i.id);
|
|
99
|
+
|
|
100
|
+
const keepIds = JSON.parse(match[0]) as number[];
|
|
101
|
+
const validIds = new Set(items.map(i => i.id));
|
|
102
|
+
return keepIds.filter(id => validIds.has(id));
|
|
103
|
+
} catch {
|
|
104
|
+
// Network error → fail-open, return all
|
|
105
|
+
return items.map(i => i.id);
|
|
106
|
+
}
|
|
107
|
+
}
|
|
108
|
+
|
|
109
|
+
async close(): Promise<void> {
|
|
110
|
+
// No resources to release (stateless HTTP)
|
|
111
|
+
}
|
|
112
|
+
}
|
|
@@ -0,0 +1,174 @@
|
|
|
1
|
+
/**
|
|
2
|
+
* BrainBank — HNSW Vector Index
|
|
3
|
+
*
|
|
4
|
+
* Wraps hnswlib-node for O(log n) approximate nearest neighbor search.
|
|
5
|
+
* M=16 connections, ef=200 construction, ef=50 search by default.
|
|
6
|
+
* 150x faster than brute force at 1M vectors.
|
|
7
|
+
*
|
|
8
|
+
* Supports disk persistence: save(path) / tryLoad(path, count)
|
|
9
|
+
* to skip costly vector-by-vector rebuild on startup.
|
|
10
|
+
*/
|
|
11
|
+
|
|
12
|
+
import type { VectorIndex, SearchHit } from '@/types.ts';
|
|
13
|
+
|
|
14
|
+
import { existsSync } from 'node:fs';
|
|
15
|
+
|
|
16
|
+
/** Shape of the HNSW index from hnswlib-node. */
|
|
17
|
+
interface HnswlibIndex {
|
|
18
|
+
initIndex(maxElements: number, M: number, efConstruction: number): void;
|
|
19
|
+
setEf(ef: number): void;
|
|
20
|
+
addPoint(vector: number[], id: number): void;
|
|
21
|
+
markDelete(id: number): void;
|
|
22
|
+
searchKnn(vector: number[], k: number): { neighbors: number[]; distances: number[] };
|
|
23
|
+
writeIndexSync(path: string): void;
|
|
24
|
+
readIndexSync(path: string): void;
|
|
25
|
+
getCurrentCount(): number;
|
|
26
|
+
getIdsList(): number[];
|
|
27
|
+
}
|
|
28
|
+
|
|
29
|
+
/** Shape of the hnswlib-node module (supports default + named exports). */
|
|
30
|
+
interface HnswlibModule {
|
|
31
|
+
default?: { HierarchicalNSW: new (space: 'cosine' | 'l2' | 'ip', dims: number) => HnswlibIndex };
|
|
32
|
+
HierarchicalNSW?: new (space: 'cosine' | 'l2' | 'ip', dims: number) => HnswlibIndex;
|
|
33
|
+
}
|
|
34
|
+
|
|
35
|
+
export class HNSWIndex implements VectorIndex {
|
|
36
|
+
private _index: HnswlibIndex | null = null;
|
|
37
|
+
private _lib: HnswlibModule | null = null;
|
|
38
|
+
private _ids = new Set<number>();
|
|
39
|
+
|
|
40
|
+
constructor(
|
|
41
|
+
private _dims: number,
|
|
42
|
+
private _maxElements: number = 2_000_000,
|
|
43
|
+
private _M: number = 16,
|
|
44
|
+
private _efConstruction: number = 200,
|
|
45
|
+
private _efSearch: number = 50,
|
|
46
|
+
) {}
|
|
47
|
+
|
|
48
|
+
/**
|
|
49
|
+
* Initialize the HNSW index.
|
|
50
|
+
* Must be called before add/search.
|
|
51
|
+
*/
|
|
52
|
+
async init(): Promise<this> {
|
|
53
|
+
this._lib = await import('hnswlib-node');
|
|
54
|
+
this._createIndex();
|
|
55
|
+
return this;
|
|
56
|
+
}
|
|
57
|
+
|
|
58
|
+
/**
|
|
59
|
+
* Reinitialize the index in-place, clearing all vectors.
|
|
60
|
+
* Required after reembed or full re-index to avoid duplicate IDs.
|
|
61
|
+
* init() must have been called first.
|
|
62
|
+
*/
|
|
63
|
+
reinit(): void {
|
|
64
|
+
if (!this._lib) throw new Error('HNSW not initialized — call init() first');
|
|
65
|
+
this._createIndex();
|
|
66
|
+
}
|
|
67
|
+
|
|
68
|
+
private _createIndex(): void {
|
|
69
|
+
if (!this._lib) throw new Error('HNSW lib not loaded');
|
|
70
|
+
const HNSW = this._lib.default?.HierarchicalNSW ?? this._lib.HierarchicalNSW;
|
|
71
|
+
if (!HNSW) throw new Error('HierarchicalNSW not found in hnswlib-node module');
|
|
72
|
+
this._index = new HNSW('cosine', this._dims);
|
|
73
|
+
this._index.initIndex(this._maxElements, this._M, this._efConstruction);
|
|
74
|
+
this._index.setEf(this._efSearch);
|
|
75
|
+
this._ids = new Set();
|
|
76
|
+
}
|
|
77
|
+
|
|
78
|
+
/** Maximum capacity of this index. */
|
|
79
|
+
get maxElements(): number { return this._maxElements; }
|
|
80
|
+
|
|
81
|
+
/**
|
|
82
|
+
* Add a vector with an integer ID.
|
|
83
|
+
* The vector should be pre-normalized for cosine distance.
|
|
84
|
+
*/
|
|
85
|
+
add(vector: Float32Array, id: number): void {
|
|
86
|
+
if (!this._index) throw new Error('HNSW index not initialized — call init() first');
|
|
87
|
+
if (this._ids.has(id)) return; // idempotent: skip duplicates
|
|
88
|
+
if (this._ids.size >= this._maxElements) {
|
|
89
|
+
throw new Error(
|
|
90
|
+
`HNSW index full (${this._maxElements} elements). ` +
|
|
91
|
+
`Increase maxElements in config or prune old data.`
|
|
92
|
+
);
|
|
93
|
+
}
|
|
94
|
+
this._index.addPoint(Array.from(vector), id);
|
|
95
|
+
this._ids.add(id);
|
|
96
|
+
}
|
|
97
|
+
|
|
98
|
+
/**
|
|
99
|
+
* Mark a vector as deleted so it no longer appears in searches.
|
|
100
|
+
* Uses hnswlib-node markDelete under the hood.
|
|
101
|
+
* Safe to call with an ID that doesn't exist.
|
|
102
|
+
*/
|
|
103
|
+
remove(id: number): void {
|
|
104
|
+
if (!this._index || this._ids.size === 0) return;
|
|
105
|
+
if (!this._ids.has(id)) return;
|
|
106
|
+
try {
|
|
107
|
+
this._index.markDelete(id);
|
|
108
|
+
this._ids.delete(id);
|
|
109
|
+
} catch {
|
|
110
|
+
// ID not found — ignore silently
|
|
111
|
+
}
|
|
112
|
+
}
|
|
113
|
+
|
|
114
|
+
/**
|
|
115
|
+
* Search for the k nearest neighbors.
|
|
116
|
+
* Returns results sorted by score (highest first).
|
|
117
|
+
* Score is 1 - cosine_distance (1.0 = identical).
|
|
118
|
+
*/
|
|
119
|
+
search(query: Float32Array, k: number): SearchHit[] {
|
|
120
|
+
if (!this._index || this._ids.size === 0) return [];
|
|
121
|
+
|
|
122
|
+
const actualK = Math.min(k, this._ids.size);
|
|
123
|
+
const result = this._index.searchKnn(Array.from(query), actualK);
|
|
124
|
+
|
|
125
|
+
return result.neighbors.map((id: number, i: number) => ({
|
|
126
|
+
id,
|
|
127
|
+
score: 1 - result.distances[i],
|
|
128
|
+
}));
|
|
129
|
+
}
|
|
130
|
+
|
|
131
|
+
/** Number of vectors in the index. */
|
|
132
|
+
get size(): number {
|
|
133
|
+
return this._ids.size;
|
|
134
|
+
}
|
|
135
|
+
|
|
136
|
+
/**
|
|
137
|
+
* Save the HNSW graph to disk.
|
|
138
|
+
* The file can be loaded later with tryLoad() to skip vector-by-vector insertion.
|
|
139
|
+
*/
|
|
140
|
+
save(path: string): void {
|
|
141
|
+
if (!this._index || this._ids.size === 0) return;
|
|
142
|
+
this._index.writeIndexSync(path);
|
|
143
|
+
}
|
|
144
|
+
|
|
145
|
+
/**
|
|
146
|
+
* Try to load a previously saved HNSW index from disk.
|
|
147
|
+
* Returns true if loaded successfully, false if stale or missing.
|
|
148
|
+
* @param path File path to the saved index
|
|
149
|
+
* @param expectedCount Expected number of vectors (from SQLite) — used to detect staleness
|
|
150
|
+
*/
|
|
151
|
+
tryLoad(path: string, expectedCount: number): boolean {
|
|
152
|
+
if (!this._index || !existsSync(path)) return false;
|
|
153
|
+
|
|
154
|
+
try {
|
|
155
|
+
this._index.readIndexSync(path);
|
|
156
|
+
const loadedCount = this._index.getCurrentCount();
|
|
157
|
+
|
|
158
|
+
// Stale: vector count in DB differs from saved index
|
|
159
|
+
if (loadedCount !== expectedCount) {
|
|
160
|
+
this.reinit();
|
|
161
|
+
return false;
|
|
162
|
+
}
|
|
163
|
+
|
|
164
|
+
// Rebuild _ids set from the loaded index
|
|
165
|
+
const ids = this._index.getIdsList();
|
|
166
|
+
this._ids = new Set(ids);
|
|
167
|
+
this._index.setEf(this._efSearch);
|
|
168
|
+
return true;
|
|
169
|
+
} catch {
|
|
170
|
+
this.reinit();
|
|
171
|
+
return false;
|
|
172
|
+
}
|
|
173
|
+
}
|
|
174
|
+
}
|
|
@@ -0,0 +1,129 @@
|
|
|
1
|
+
/**
|
|
2
|
+
* BrainBank — HNSW Loader
|
|
3
|
+
*
|
|
4
|
+
* Utilities for persisting and loading HNSW indexes to/from disk.
|
|
5
|
+
* Used by BrainBank._runInitialize() and PluginContext.loadVectors().
|
|
6
|
+
*
|
|
7
|
+
* Includes cross-process write locking and hot-reload support
|
|
8
|
+
* for multi-process coordination.
|
|
9
|
+
*/
|
|
10
|
+
|
|
11
|
+
import type { DatabaseAdapter, CountRow } from '@/db/adapter.ts';
|
|
12
|
+
import type { HNSWIndex } from './hnsw-index.ts';
|
|
13
|
+
|
|
14
|
+
import { dirname, join } from 'node:path';
|
|
15
|
+
import { withLock } from '@/lib/write-lock.ts';
|
|
16
|
+
|
|
17
|
+
/** Derive the HNSW index file path from the DB path. */
|
|
18
|
+
export function hnswPath(dbPath: string, name: string): string {
|
|
19
|
+
return join(dirname(dbPath), `hnsw-${name}.index`);
|
|
20
|
+
}
|
|
21
|
+
|
|
22
|
+
/** Derive the lock directory from the DB path. */
|
|
23
|
+
export function lockDir(dbPath: string): string {
|
|
24
|
+
return dirname(dbPath);
|
|
25
|
+
}
|
|
26
|
+
|
|
27
|
+
/** Count rows in a vector table (fast, no data transfer). */
|
|
28
|
+
export function countRows(db: DatabaseAdapter, table: string): number {
|
|
29
|
+
const row = db.prepare(`SELECT COUNT(*) as c FROM ${table}`).get() as CountRow;
|
|
30
|
+
return row?.c ?? 0;
|
|
31
|
+
}
|
|
32
|
+
|
|
33
|
+
/**
|
|
34
|
+
* Save all HNSW indexes to disk with cross-process file locking.
|
|
35
|
+
* Prevents concurrent writes from corrupting `.index` files.
|
|
36
|
+
*/
|
|
37
|
+
export async function saveAllHnsw(
|
|
38
|
+
dbPath: string,
|
|
39
|
+
kvHnsw: HNSWIndex,
|
|
40
|
+
sharedHnsw: Map<string, { hnsw: HNSWIndex; vecCache: Map<number, Float32Array> }>,
|
|
41
|
+
privateHnsw: Map<string, HNSWIndex>,
|
|
42
|
+
): Promise<boolean> {
|
|
43
|
+
try {
|
|
44
|
+
await withLock(lockDir(dbPath), 'hnsw', () => {
|
|
45
|
+
kvHnsw.save(hnswPath(dbPath, 'kv'));
|
|
46
|
+
for (const [name, { hnsw }] of sharedHnsw) {
|
|
47
|
+
hnsw.save(hnswPath(dbPath, name));
|
|
48
|
+
}
|
|
49
|
+
for (const [name, hnsw] of privateHnsw) {
|
|
50
|
+
hnsw.save(hnswPath(dbPath, name));
|
|
51
|
+
}
|
|
52
|
+
});
|
|
53
|
+
return true;
|
|
54
|
+
} catch {
|
|
55
|
+
// Non-fatal: next startup rebuilds from SQLite (slower).
|
|
56
|
+
return false;
|
|
57
|
+
}
|
|
58
|
+
}
|
|
59
|
+
|
|
60
|
+
/** Load vectors from SQLite into HNSW + cache. */
|
|
61
|
+
export function loadVectors(
|
|
62
|
+
db: DatabaseAdapter,
|
|
63
|
+
table: string,
|
|
64
|
+
idCol: string,
|
|
65
|
+
hnsw: HNSWIndex,
|
|
66
|
+
cache: Map<number, Float32Array>,
|
|
67
|
+
): void {
|
|
68
|
+
const iter = db.prepare(`SELECT ${idCol}, embedding FROM ${table}`).iterate() as IterableIterator<{ embedding: Uint8Array; [key: string]: unknown }>;
|
|
69
|
+
for (const row of iter) {
|
|
70
|
+
const vec = new Float32Array(
|
|
71
|
+
row.embedding.buffer.slice(
|
|
72
|
+
row.embedding.byteOffset,
|
|
73
|
+
row.embedding.byteOffset + row.embedding.byteLength,
|
|
74
|
+
),
|
|
75
|
+
);
|
|
76
|
+
hnsw.add(vec, row[idCol] as number);
|
|
77
|
+
cache.set(row[idCol] as number, vec);
|
|
78
|
+
}
|
|
79
|
+
}
|
|
80
|
+
|
|
81
|
+
/** Populate only the vecCache from SQLite (HNSW already loaded from file). */
|
|
82
|
+
export function loadVecCache(
|
|
83
|
+
db: DatabaseAdapter,
|
|
84
|
+
table: string,
|
|
85
|
+
idCol: string,
|
|
86
|
+
cache: Map<number, Float32Array>,
|
|
87
|
+
): void {
|
|
88
|
+
const iter = db.prepare(`SELECT ${idCol}, embedding FROM ${table}`).iterate() as IterableIterator<{ embedding: Uint8Array; [key: string]: unknown }>;
|
|
89
|
+
for (const row of iter) {
|
|
90
|
+
const vec = new Float32Array(
|
|
91
|
+
row.embedding.buffer.slice(
|
|
92
|
+
row.embedding.byteOffset,
|
|
93
|
+
row.embedding.byteOffset + row.embedding.byteLength,
|
|
94
|
+
),
|
|
95
|
+
);
|
|
96
|
+
cache.set(row[idCol] as number, vec);
|
|
97
|
+
}
|
|
98
|
+
}
|
|
99
|
+
|
|
100
|
+
/** Deps for reloading a single HNSW index from disk. */
|
|
101
|
+
interface ReloadDeps {
|
|
102
|
+
dbPath: string;
|
|
103
|
+
db: DatabaseAdapter;
|
|
104
|
+
name: string;
|
|
105
|
+
hnsw: HNSWIndex;
|
|
106
|
+
vecCache: Map<number, Float32Array>;
|
|
107
|
+
vectorTable: string;
|
|
108
|
+
idCol: string;
|
|
109
|
+
}
|
|
110
|
+
|
|
111
|
+
/**
|
|
112
|
+
* Reload a single HNSW index from disk after detecting a stale version.
|
|
113
|
+
* Reinitializes the in-memory HNSW, loads the saved index file, and
|
|
114
|
+
* refreshes the vector cache from SQLite.
|
|
115
|
+
*/
|
|
116
|
+
export function reloadHnsw(deps: ReloadDeps): void {
|
|
117
|
+
const { dbPath, db, name, hnsw, vecCache, vectorTable, idCol } = deps;
|
|
118
|
+
const indexPath = hnswPath(dbPath, name);
|
|
119
|
+
const rowCount = countRows(db, vectorTable);
|
|
120
|
+
|
|
121
|
+
hnsw.reinit();
|
|
122
|
+
vecCache.clear();
|
|
123
|
+
|
|
124
|
+
if (hnsw.tryLoad(indexPath, rowCount)) {
|
|
125
|
+
loadVecCache(db, vectorTable, idCol, vecCache);
|
|
126
|
+
} else {
|
|
127
|
+
loadVectors(db, vectorTable, idCol, hnsw, vecCache);
|
|
128
|
+
}
|
|
129
|
+
}
|
|
@@ -0,0 +1,69 @@
|
|
|
1
|
+
/**
|
|
2
|
+
* BrainBank — BM25 Intersection Boost
|
|
3
|
+
*
|
|
4
|
+
* Pure functions for post-processing vector search results with BM25 keyword overlap.
|
|
5
|
+
* Extracted from ContextBuilder for single-responsibility and testability.
|
|
6
|
+
*/
|
|
7
|
+
|
|
8
|
+
import type { SearchResult } from '@/types.ts';
|
|
9
|
+
import type { SearchStrategy } from './types.ts';
|
|
10
|
+
|
|
11
|
+
/** BM25 boost factor applied to vector results that also match keywords. */
|
|
12
|
+
export const BM25_BOOST = 0.15;
|
|
13
|
+
|
|
14
|
+
/**
|
|
15
|
+
* Boost vector results that also appear in BM25 keyword results.
|
|
16
|
+
* Does NOT add new results — only re-scores and re-sorts existing vector hits.
|
|
17
|
+
* This promotes keyword-relevant files without introducing BM25-only noise.
|
|
18
|
+
*/
|
|
19
|
+
export async function boostWithBM25(
|
|
20
|
+
vectorResults: SearchResult[],
|
|
21
|
+
bm25: SearchStrategy,
|
|
22
|
+
query: string,
|
|
23
|
+
sources: Record<string, number>,
|
|
24
|
+
): Promise<SearchResult[]> {
|
|
25
|
+
if (vectorResults.length === 0) return vectorResults;
|
|
26
|
+
|
|
27
|
+
const bm25Results = await bm25.search(query, { sources });
|
|
28
|
+
if (bm25Results.length === 0) return vectorResults;
|
|
29
|
+
|
|
30
|
+
// Build a set of BM25 hit keys for fast lookup
|
|
31
|
+
const bm25Keys = new Set<string>();
|
|
32
|
+
for (const r of bm25Results) {
|
|
33
|
+
bm25Keys.add(resultKey(r));
|
|
34
|
+
}
|
|
35
|
+
|
|
36
|
+
// Boost scores of vector results that also appear in BM25
|
|
37
|
+
const boosted = vectorResults.map(r => {
|
|
38
|
+
const k = resultKey(r);
|
|
39
|
+
if (bm25Keys.has(k)) {
|
|
40
|
+
return { ...r, score: r.score + BM25_BOOST };
|
|
41
|
+
}
|
|
42
|
+
return r;
|
|
43
|
+
});
|
|
44
|
+
|
|
45
|
+
// Re-sort by boosted score
|
|
46
|
+
boosted.sort((a, b) => b.score - a.score);
|
|
47
|
+
return boosted;
|
|
48
|
+
}
|
|
49
|
+
|
|
50
|
+
/** Filter results whose filePath starts with any of the given prefixes. */
|
|
51
|
+
export function filterByPath(results: SearchResult[], prefix: string | string[] | undefined): SearchResult[] {
|
|
52
|
+
if (!prefix) return results;
|
|
53
|
+
const prefixes = Array.isArray(prefix) ? prefix : [prefix];
|
|
54
|
+
if (prefixes.length === 0) return results;
|
|
55
|
+
return results.filter(r => prefixes.some(p => r.filePath?.startsWith(p)));
|
|
56
|
+
}
|
|
57
|
+
|
|
58
|
+
/** Exclude results whose filePath starts with any of the given prefixes. */
|
|
59
|
+
export function filterByIgnore(results: SearchResult[], ignorePaths: string[] | undefined): SearchResult[] {
|
|
60
|
+
if (!ignorePaths || ignorePaths.length === 0) return results;
|
|
61
|
+
return results.filter(r => !r.filePath || !ignorePaths.some(p => r.filePath!.startsWith(p)));
|
|
62
|
+
}
|
|
63
|
+
|
|
64
|
+
/** Generate a dedup key for a search result (file:startLine:endLine). */
|
|
65
|
+
export function resultKey(r: SearchResult): string {
|
|
66
|
+
const sl = 'startLine' in r.metadata ? r.metadata.startLine : '';
|
|
67
|
+
const el = 'endLine' in r.metadata ? r.metadata.endLine : '';
|
|
68
|
+
return `${r.filePath ?? ''}:${sl}:${el}`;
|
|
69
|
+
}
|