@lotargo/memory_plugin 1.5.3 → 1.6.0
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/CHANGELOG.md +108 -0
- package/README.md +383 -352
- package/mcp-server/admin/auth.js +13 -4
- package/mcp-server/admin/snapshot.js +24 -7
- package/mcp-server/cli/direct_commands.js +334 -313
- package/mcp-server/cli/handlers/engine_actions.js +41 -0
- package/mcp-server/cli/handlers/storage_actions.js +58 -0
- package/mcp-server/cli/secret_input.js +44 -0
- package/mcp-server/cli/ui.js +564 -565
- package/mcp-server/cli.js +356 -324
- package/mcp-server/config/auth_store.js +74 -16
- package/mcp-server/config/config_manager.js +4 -0
- package/mcp-server/db/database.js +33 -14
- package/mcp-server/db/sync_queue.js +9 -19
- package/mcp-server/index.js +112 -42
- package/mcp-server/ingest/normalizer.js +116 -29
- package/mcp-server/ingest/pipeline.js +94 -6
- package/mcp-server/logger.js +49 -0
- package/mcp-server/memory.js +6 -9
- package/mcp-server/ml/gpu_monitor.js +169 -166
- package/mcp-server/ml/model_manager.js +17 -4
- package/mcp-server/retrieval/retriever.js +35 -15
- package/mcp-server/security/path_guard.js +67 -0
- package/mcp-server/setup.js +10 -2
- package/mcp-server/storage/blob_store.js +15 -2
- package/mcp-server/tools/core/memory_core.js +393 -0
- package/mcp-server/tools/helpers.js +59 -39
- package/mcp-server/tools/memory_tools.js +123 -516
- package/mcp-server/tools/rag_tools.js +49 -1
- package/opencode-plugin/index.js +94 -397
- package/package.json +7 -1
- package/skills/using-memory/SKILL.md +7 -2
|
@@ -5,9 +5,11 @@ import { getDatabase, BLOBS_DIR } from "../db/database.js";
|
|
|
5
5
|
import { saveBlob, deleteBlob } from "../storage/blob_store.js";
|
|
6
6
|
import { normalizeContent, fetchUrlContent } from "./normalizer.js";
|
|
7
7
|
import { buildTripleHierarchy } from "./chunker.js";
|
|
8
|
-
import {
|
|
8
|
+
import { embedBatch, vectorToBuffer } from "../ml/model_manager.js";
|
|
9
9
|
import { buildGraphEdges, saveGraphEdges } from "../graph/graph_extractor.js";
|
|
10
10
|
import { getConfig } from "../config/config_manager.js";
|
|
11
|
+
import { assertIngestPathAllowed } from "../security/path_guard.js";
|
|
12
|
+
import { logger } from "../logger.js";
|
|
11
13
|
|
|
12
14
|
export async function ingestDocument({
|
|
13
15
|
content,
|
|
@@ -34,9 +36,10 @@ export async function ingestDocument({
|
|
|
34
36
|
const filePath = effectivePath || content;
|
|
35
37
|
const needsRead = !content || content === filePath;
|
|
36
38
|
if (needsRead && filePath) {
|
|
37
|
-
const
|
|
39
|
+
const safePath = assertIngestPathAllowed(filePath);
|
|
40
|
+
const ext = extname(safePath).toLowerCase();
|
|
38
41
|
const isBinary = [".pdf", ".docx", ".xlsx", ".xls"].includes(ext);
|
|
39
|
-
content = await readFile(
|
|
42
|
+
content = await readFile(safePath, isBinary ? null : "utf-8");
|
|
40
43
|
effectivePath = filePath;
|
|
41
44
|
}
|
|
42
45
|
}
|
|
@@ -77,7 +80,7 @@ export async function ingestDocument({
|
|
|
77
80
|
}
|
|
78
81
|
} else {
|
|
79
82
|
for (const micro of hierarchy.microChunks) {
|
|
80
|
-
micro.vector = Buffer.alloc(
|
|
83
|
+
micro.vector = Buffer.alloc(0);
|
|
81
84
|
}
|
|
82
85
|
}
|
|
83
86
|
|
|
@@ -155,7 +158,7 @@ export async function ingestDocument({
|
|
|
155
158
|
const exportedData = await exportDocumentData(docId, db);
|
|
156
159
|
await enqueueSyncTask("ingest_document", docId, exportedData);
|
|
157
160
|
} catch (err) {
|
|
158
|
-
|
|
161
|
+
logger.error("Failed to queue document ingest sync task:", err.message);
|
|
159
162
|
}
|
|
160
163
|
}
|
|
161
164
|
|
|
@@ -226,9 +229,94 @@ export async function deleteDocument(docIdOrPath, customDb = null, customBlobDir
|
|
|
226
229
|
const { enqueueSyncTask } = await import("../db/sync_queue.js");
|
|
227
230
|
await enqueueSyncTask("delete_document", docIdOrPath);
|
|
228
231
|
} catch (err) {
|
|
229
|
-
|
|
232
|
+
logger.error("Failed to queue document delete sync task:", err.message);
|
|
230
233
|
}
|
|
231
234
|
}
|
|
232
235
|
|
|
233
236
|
return { deleted: true, docId: doc.id, title: doc.title, linksCleaned: true };
|
|
234
237
|
}
|
|
238
|
+
|
|
239
|
+
export async function reindexEmbeddings({
|
|
240
|
+
model = null,
|
|
241
|
+
dimension = null,
|
|
242
|
+
customDb = null,
|
|
243
|
+
embedFn = null,
|
|
244
|
+
progressCallback = null,
|
|
245
|
+
} = {}) {
|
|
246
|
+
const db = customDb || await getDatabase();
|
|
247
|
+
const config = getConfig();
|
|
248
|
+
const targetModel = model || config.embeddingModel || "Xenova/multilingual-e5-small";
|
|
249
|
+
const targetDim = dimension !== null && dimension !== undefined ? Number(dimension) : (config.vectorDimension || 0);
|
|
250
|
+
|
|
251
|
+
const countRow = await db.prepare("SELECT COUNT(*) as cnt FROM micro_chunks").get();
|
|
252
|
+
const total = countRow ? countRow.cnt : 0;
|
|
253
|
+
if (total === 0) return { reindexed: 0, documentsAffected: 0, model: targetModel, dimension: targetDim };
|
|
254
|
+
|
|
255
|
+
const rows = await db.prepare(`
|
|
256
|
+
SELECT m.id, m.doc_id, m.content, s.breadcrumbs, d.title as doc_title
|
|
257
|
+
FROM micro_chunks m
|
|
258
|
+
LEFT JOIN sections s ON m.section_id = s.id
|
|
259
|
+
LEFT JOIN documents d ON m.doc_id = d.id
|
|
260
|
+
ORDER BY m.doc_id, m.id
|
|
261
|
+
`).all();
|
|
262
|
+
|
|
263
|
+
const items = rows.map((r) => ({
|
|
264
|
+
id: r.id,
|
|
265
|
+
doc_id: r.doc_id,
|
|
266
|
+
text: r.breadcrumbs
|
|
267
|
+
? `${r.content}\n\nContext: ${r.doc_title || ""} > ${r.breadcrumbs}`
|
|
268
|
+
: `${r.content}\n\nContext: ${r.doc_title || ""}`,
|
|
269
|
+
}));
|
|
270
|
+
|
|
271
|
+
const defaultEmbed = async (texts) =>
|
|
272
|
+
embedBatch(texts, false, targetModel, progressCallback, null, {}, targetDim || null);
|
|
273
|
+
const embed = embedFn || defaultEmbed;
|
|
274
|
+
|
|
275
|
+
const BATCH_SIZE = config.batchSize || 12;
|
|
276
|
+
const vectors = [];
|
|
277
|
+
for (let i = 0; i < items.length; i += BATCH_SIZE) {
|
|
278
|
+
const batch = items.slice(i, i + BATCH_SIZE);
|
|
279
|
+
const batchVecs = await embed(batch.map((b) => b.text));
|
|
280
|
+
if (!batchVecs || batchVecs.length !== batch.length) {
|
|
281
|
+
throw new Error(`Embedding batch returned ${batchVecs ? batchVecs.length : 0} vectors, expected ${batch.length}`);
|
|
282
|
+
}
|
|
283
|
+
for (let j = 0; j < batch.length; j++) {
|
|
284
|
+
vectors.push({ id: batch[j].id, doc_id: batch[j].doc_id, vector: vectorToBuffer(batchVecs[j]) });
|
|
285
|
+
}
|
|
286
|
+
if (progressCallback) progressCallback({ done: vectors.length, total });
|
|
287
|
+
}
|
|
288
|
+
|
|
289
|
+
await db.exec("BEGIN IMMEDIATE;");
|
|
290
|
+
try {
|
|
291
|
+
const stmt = db.prepare("UPDATE micro_chunks SET vector = ? WHERE id = ?;");
|
|
292
|
+
for (const v of vectors) {
|
|
293
|
+
await stmt.run(v.vector, v.id);
|
|
294
|
+
}
|
|
295
|
+
await db.exec("COMMIT;");
|
|
296
|
+
} catch (err) {
|
|
297
|
+
await db.exec("ROLLBACK;");
|
|
298
|
+
throw new Error(`Re-index transaction failed: ${err.message}`);
|
|
299
|
+
}
|
|
300
|
+
|
|
301
|
+
const affectedDocIds = [...new Set(items.map((i) => i.doc_id).filter(Boolean))];
|
|
302
|
+
|
|
303
|
+
if (config.mode === "hybrid-sync") {
|
|
304
|
+
try {
|
|
305
|
+
const { enqueueSyncTask } = await import("../db/sync_queue.js");
|
|
306
|
+
const { exportDocumentData } = await import("./exporter.js");
|
|
307
|
+
for (const docId of affectedDocIds) {
|
|
308
|
+
const exportedData = await exportDocumentData(docId, db);
|
|
309
|
+
await enqueueSyncTask("ingest_document", docId, exportedData);
|
|
310
|
+
}
|
|
311
|
+
} catch (err) {
|
|
312
|
+
logger.error("Failed to queue document re-index sync tasks:", err.message);
|
|
313
|
+
}
|
|
314
|
+
}
|
|
315
|
+
|
|
316
|
+
return {
|
|
317
|
+
reindexed: vectors.length,
|
|
318
|
+
documentsAffected: affectedDocIds.length,
|
|
319
|
+
model: targetModel,
|
|
320
|
+
dimension: targetDim,
|
|
321
|
+
};
|
|
322
|
+
}
|
|
@@ -0,0 +1,49 @@
|
|
|
1
|
+
// Central logger. Everything goes to stderr: stdout is the MCP JSON-RPC
|
|
2
|
+
// channel and any stray write there corrupts the protocol stream.
|
|
3
|
+
//
|
|
4
|
+
// Level is taken from MEMORY_LOG_LEVEL (silent|error|warn|info|debug),
|
|
5
|
+
// defaulting to "warn". Hosts embedding the server can swap the sink with
|
|
6
|
+
// setLogSink() — e.g. to forward into the OpenCode client log.
|
|
7
|
+
|
|
8
|
+
const LEVELS = { silent: 0, error: 1, warn: 2, info: 3, debug: 4 };
|
|
9
|
+
|
|
10
|
+
function envLevel() {
|
|
11
|
+
const raw = String(process.env.MEMORY_LOG_LEVEL || "").toLowerCase();
|
|
12
|
+
return raw in LEVELS ? raw : "warn";
|
|
13
|
+
}
|
|
14
|
+
|
|
15
|
+
let currentLevel = envLevel();
|
|
16
|
+
let sink = (level, message) => {
|
|
17
|
+
process.stderr.write(`[memory:${level}] ${message}\n`);
|
|
18
|
+
};
|
|
19
|
+
|
|
20
|
+
export function setLogLevel(level) {
|
|
21
|
+
if (level in LEVELS) currentLevel = level;
|
|
22
|
+
}
|
|
23
|
+
|
|
24
|
+
export function getLogLevel() {
|
|
25
|
+
return currentLevel;
|
|
26
|
+
}
|
|
27
|
+
|
|
28
|
+
export function setLogSink(fn) {
|
|
29
|
+
sink = typeof fn === "function" ? fn : sink;
|
|
30
|
+
}
|
|
31
|
+
|
|
32
|
+
function emit(level, args) {
|
|
33
|
+
if (LEVELS[level] > LEVELS[currentLevel]) return;
|
|
34
|
+
const message = args
|
|
35
|
+
.map((a) => (a instanceof Error ? a.stack || a.message : typeof a === "string" ? a : JSON.stringify(a)))
|
|
36
|
+
.join(" ");
|
|
37
|
+
try {
|
|
38
|
+
sink(level, message);
|
|
39
|
+
} catch {}
|
|
40
|
+
}
|
|
41
|
+
|
|
42
|
+
export const logger = {
|
|
43
|
+
error: (...args) => emit("error", args),
|
|
44
|
+
warn: (...args) => emit("warn", args),
|
|
45
|
+
info: (...args) => emit("info", args),
|
|
46
|
+
debug: (...args) => emit("debug", args),
|
|
47
|
+
};
|
|
48
|
+
|
|
49
|
+
export default logger;
|
package/mcp-server/memory.js
CHANGED
|
@@ -3,6 +3,7 @@ import { existsSync, mkdirSync } from "fs";
|
|
|
3
3
|
import { join, basename, resolve } from "path";
|
|
4
4
|
import { homedir } from "os";
|
|
5
5
|
import { resolveProjectIdentity } from "./identity.js";
|
|
6
|
+
import { logger } from "./logger.js";
|
|
6
7
|
|
|
7
8
|
function resolveMemoryDir() {
|
|
8
9
|
if (process.env.MEMORY_DIR) return process.env.MEMORY_DIR;
|
|
@@ -93,10 +94,6 @@ export function parseMeta(content) {
|
|
|
93
94
|
return { key: m ? m[1].trim() : null };
|
|
94
95
|
}
|
|
95
96
|
|
|
96
|
-
function isSimpleKey(key) {
|
|
97
|
-
return /^[a-zA-Z0-9_-]+$/.test(key);
|
|
98
|
-
}
|
|
99
|
-
|
|
100
97
|
export async function readMemory(key) {
|
|
101
98
|
if (!key) return [];
|
|
102
99
|
const { getConfig } = await import("./config/config_manager.js");
|
|
@@ -110,7 +107,7 @@ export async function readMemory(key) {
|
|
|
110
107
|
return row.content.split("\n").filter((l) => l.startsWith("- ["));
|
|
111
108
|
}
|
|
112
109
|
} catch (err) {
|
|
113
|
-
|
|
110
|
+
logger.error("Failed to read memory from cloud database:", err.message);
|
|
114
111
|
}
|
|
115
112
|
return [];
|
|
116
113
|
}
|
|
@@ -121,7 +118,7 @@ export async function readMemory(key) {
|
|
|
121
118
|
const { ensureReverseSync } = await import("./db/sync_queue.js");
|
|
122
119
|
await ensureReverseSync();
|
|
123
120
|
} catch (err) {
|
|
124
|
-
|
|
121
|
+
logger.error("Failed to reverse-sync before read:", err.message);
|
|
125
122
|
}
|
|
126
123
|
}
|
|
127
124
|
if (existsSync(fp)) {
|
|
@@ -171,7 +168,7 @@ export async function writeMemory(key, entries) {
|
|
|
171
168
|
ON CONFLICT(key) DO UPDATE SET content = excluded.content, updated_at = excluded.updated_at;
|
|
172
169
|
`).run(key, content, Date.now());
|
|
173
170
|
} catch (err) {
|
|
174
|
-
|
|
171
|
+
logger.error("Failed to write memory to cloud database:", err.message);
|
|
175
172
|
}
|
|
176
173
|
return;
|
|
177
174
|
}
|
|
@@ -183,7 +180,7 @@ export async function writeMemory(key, entries) {
|
|
|
183
180
|
const { enqueueSyncTask } = await import("./db/sync_queue.js");
|
|
184
181
|
await enqueueSyncTask("write_memory", key, content);
|
|
185
182
|
} catch (err) {
|
|
186
|
-
|
|
183
|
+
logger.error("Failed to queue memory sync task:", err.message);
|
|
187
184
|
}
|
|
188
185
|
}
|
|
189
186
|
}
|
|
@@ -214,7 +211,7 @@ export async function listProjectStores() {
|
|
|
214
211
|
stores.sort((a, b) => a.basename.localeCompare(b.basename));
|
|
215
212
|
return stores;
|
|
216
213
|
} catch (err) {
|
|
217
|
-
|
|
214
|
+
logger.error("Failed to list memory stores from cloud database:", err.message);
|
|
218
215
|
}
|
|
219
216
|
return [];
|
|
220
217
|
}
|
|
@@ -1,166 +1,169 @@
|
|
|
1
|
-
import { execFile } from "node:child_process";
|
|
2
|
-
import { promisify } from "node:util";
|
|
3
|
-
|
|
4
|
-
const execFileAsync = promisify(execFile);
|
|
5
|
-
|
|
6
|
-
export async function getGpuUtilizationAsync() {
|
|
7
|
-
try {
|
|
8
|
-
const { stdout } = await execFileAsync("nvidia-smi",
|
|
9
|
-
["--query-gpu=utilization.gpu", "--format=csv,noheader,nounits"],
|
|
10
|
-
{ timeout: 1000 }
|
|
11
|
-
);
|
|
12
|
-
const val = parseInt(stdout.trim(), 10);
|
|
13
|
-
if (!isNaN(val)) return val;
|
|
14
|
-
} catch {}
|
|
15
|
-
|
|
16
|
-
if (process.platform === "win32") {
|
|
17
|
-
try {
|
|
18
|
-
const { stdout } = await execFileAsync("powershell",
|
|
19
|
-
["-NoProfile", "-Command", "(Get-CimInstance Win32_PerfFormattedData_GPUPerformanceCounters_GPUEngine | Measure-Object -Property UtilizationPercentage -Sum).Sum"],
|
|
20
|
-
{ timeout: 1500 }
|
|
21
|
-
);
|
|
22
|
-
const val = parseInt(stdout.trim(), 10);
|
|
23
|
-
if (!isNaN(val)) return Math.min(100, val);
|
|
24
|
-
} catch {}
|
|
25
|
-
}
|
|
26
|
-
|
|
27
|
-
return null;
|
|
28
|
-
}
|
|
29
|
-
|
|
30
|
-
export class GpuMonitor {
|
|
31
|
-
constructor(sampleIntervalMs = 100) {
|
|
32
|
-
this.intervalMs = sampleIntervalMs;
|
|
33
|
-
this.samples = [];
|
|
34
|
-
this.timer = null;
|
|
35
|
-
this.isMonitoring = false;
|
|
36
|
-
}
|
|
37
|
-
|
|
38
|
-
start() {
|
|
39
|
-
this.samples = [];
|
|
40
|
-
this.isMonitoring = true;
|
|
41
|
-
this.sample();
|
|
42
|
-
|
|
43
|
-
this.timer = setInterval(() => {
|
|
44
|
-
if (this.isMonitoring) {
|
|
45
|
-
this.sample();
|
|
46
|
-
}
|
|
47
|
-
}, this.intervalMs);
|
|
48
|
-
}
|
|
49
|
-
|
|
50
|
-
async sample() {
|
|
51
|
-
const util = await getGpuUtilizationAsync();
|
|
52
|
-
if (util !== null) {
|
|
53
|
-
this.samples.push(util);
|
|
54
|
-
}
|
|
55
|
-
}
|
|
56
|
-
|
|
57
|
-
stop() {
|
|
58
|
-
this.isMonitoring = false;
|
|
59
|
-
if (this.timer) {
|
|
60
|
-
clearInterval(this.timer);
|
|
61
|
-
this.timer = null;
|
|
62
|
-
}
|
|
63
|
-
|
|
64
|
-
if (this.samples.length === 0) {
|
|
65
|
-
return { avg: null, peak: null, samplesCount: 0 };
|
|
66
|
-
}
|
|
67
|
-
|
|
68
|
-
const max = Math.max(...this.samples);
|
|
69
|
-
const sum = this.samples.reduce((a, b) => a + b, 0);
|
|
70
|
-
const avg = Math.round(sum / this.samples.length);
|
|
71
|
-
|
|
72
|
-
return {
|
|
73
|
-
avg,
|
|
74
|
-
peak: max,
|
|
75
|
-
samplesCount: this.samples.length,
|
|
76
|
-
samples: this.samples,
|
|
77
|
-
};
|
|
78
|
-
}
|
|
79
|
-
}
|
|
80
|
-
|
|
81
|
-
export class ExecutionTracer {
|
|
82
|
-
constructor(name = "Batch Inference") {
|
|
83
|
-
this.name = name;
|
|
84
|
-
this.stages = [];
|
|
85
|
-
this.currentStage = null;
|
|
86
|
-
this.startTime = Date.now();
|
|
87
|
-
}
|
|
88
|
-
|
|
89
|
-
startStage(stageName, category = "CPU") {
|
|
90
|
-
const now = performance.now();
|
|
91
|
-
if (this.currentStage) {
|
|
92
|
-
this.currentStage.duration = now - this.currentStage.startTime;
|
|
93
|
-
this.stages.push(this.currentStage);
|
|
94
|
-
}
|
|
95
|
-
this.currentStage = {
|
|
96
|
-
name: stageName,
|
|
97
|
-
category, // "CPU" or "GPU"
|
|
98
|
-
startTime: now,
|
|
99
|
-
duration: 0,
|
|
100
|
-
};
|
|
101
|
-
}
|
|
102
|
-
|
|
103
|
-
endStage() {
|
|
104
|
-
const now = performance.now();
|
|
105
|
-
if (this.currentStage) {
|
|
106
|
-
this.currentStage.duration = now - this.currentStage.startTime;
|
|
107
|
-
this.stages.push(this.currentStage);
|
|
108
|
-
this.currentStage = null;
|
|
109
|
-
}
|
|
110
|
-
}
|
|
111
|
-
|
|
112
|
-
getSummary(gpuStats = null) {
|
|
113
|
-
this.endStage();
|
|
114
|
-
const totalMs = this.stages.reduce((sum, s) => sum + s.duration, 0);
|
|
115
|
-
|
|
116
|
-
const breakdown = this.stages.map((s) => ({
|
|
117
|
-
name: s.name,
|
|
118
|
-
category: s.category,
|
|
119
|
-
durationMs: parseFloat(s.duration.toFixed(2)),
|
|
120
|
-
pct: totalMs > 0 ? parseFloat(((s.duration / totalMs) * 100).toFixed(1)) : 0,
|
|
121
|
-
}));
|
|
122
|
-
|
|
123
|
-
const cpuTime = this.stages.filter((s) => s.category === "CPU").reduce((sum, s) => sum + s.duration, 0);
|
|
124
|
-
const gpuTime = this.stages.filter((s) => s.category === "GPU").reduce((sum, s) => sum + s.duration, 0);
|
|
125
|
-
|
|
126
|
-
return {
|
|
127
|
-
name: this.name,
|
|
128
|
-
totalMs: parseFloat(totalMs.toFixed(2)),
|
|
129
|
-
cpuMs: parseFloat(cpuTime.toFixed(2)),
|
|
130
|
-
gpuMs: parseFloat(gpuTime.toFixed(2)),
|
|
131
|
-
cpuPct: totalMs > 0 ? parseFloat(((cpuTime / totalMs) * 100).toFixed(1)) : 0,
|
|
132
|
-
gpuPct: totalMs > 0 ? parseFloat(((gpuTime / totalMs) * 100).toFixed(1)) : 0,
|
|
133
|
-
breakdown,
|
|
134
|
-
gpuStats,
|
|
135
|
-
};
|
|
136
|
-
}
|
|
137
|
-
|
|
138
|
-
printTraceReport(gpuStats = null, minGpuThreshold = 0) {
|
|
139
|
-
const summary = this.getSummary(gpuStats);
|
|
140
|
-
const line = "─".repeat(65);
|
|
141
|
-
|
|
142
|
-
|
|
143
|
-
|
|
144
|
-
|
|
145
|
-
|
|
146
|
-
|
|
147
|
-
|
|
148
|
-
|
|
149
|
-
|
|
150
|
-
|
|
151
|
-
|
|
152
|
-
|
|
153
|
-
|
|
154
|
-
|
|
155
|
-
|
|
156
|
-
|
|
157
|
-
|
|
158
|
-
|
|
159
|
-
|
|
160
|
-
|
|
161
|
-
|
|
162
|
-
|
|
163
|
-
|
|
164
|
-
|
|
165
|
-
|
|
166
|
-
|
|
1
|
+
import { execFile } from "node:child_process";
|
|
2
|
+
import { promisify } from "node:util";
|
|
3
|
+
|
|
4
|
+
const execFileAsync = promisify(execFile);
|
|
5
|
+
|
|
6
|
+
export async function getGpuUtilizationAsync() {
|
|
7
|
+
try {
|
|
8
|
+
const { stdout } = await execFileAsync("nvidia-smi",
|
|
9
|
+
["--query-gpu=utilization.gpu", "--format=csv,noheader,nounits"],
|
|
10
|
+
{ timeout: 1000 }
|
|
11
|
+
);
|
|
12
|
+
const val = parseInt(stdout.trim(), 10);
|
|
13
|
+
if (!isNaN(val)) return val;
|
|
14
|
+
} catch {}
|
|
15
|
+
|
|
16
|
+
if (process.platform === "win32") {
|
|
17
|
+
try {
|
|
18
|
+
const { stdout } = await execFileAsync("powershell",
|
|
19
|
+
["-NoProfile", "-Command", "(Get-CimInstance Win32_PerfFormattedData_GPUPerformanceCounters_GPUEngine | Measure-Object -Property UtilizationPercentage -Sum).Sum"],
|
|
20
|
+
{ timeout: 1500 }
|
|
21
|
+
);
|
|
22
|
+
const val = parseInt(stdout.trim(), 10);
|
|
23
|
+
if (!isNaN(val)) return Math.min(100, val);
|
|
24
|
+
} catch {}
|
|
25
|
+
}
|
|
26
|
+
|
|
27
|
+
return null;
|
|
28
|
+
}
|
|
29
|
+
|
|
30
|
+
export class GpuMonitor {
|
|
31
|
+
constructor(sampleIntervalMs = 100) {
|
|
32
|
+
this.intervalMs = sampleIntervalMs;
|
|
33
|
+
this.samples = [];
|
|
34
|
+
this.timer = null;
|
|
35
|
+
this.isMonitoring = false;
|
|
36
|
+
}
|
|
37
|
+
|
|
38
|
+
start() {
|
|
39
|
+
this.samples = [];
|
|
40
|
+
this.isMonitoring = true;
|
|
41
|
+
this.sample();
|
|
42
|
+
|
|
43
|
+
this.timer = setInterval(() => {
|
|
44
|
+
if (this.isMonitoring) {
|
|
45
|
+
this.sample();
|
|
46
|
+
}
|
|
47
|
+
}, this.intervalMs);
|
|
48
|
+
}
|
|
49
|
+
|
|
50
|
+
async sample() {
|
|
51
|
+
const util = await getGpuUtilizationAsync();
|
|
52
|
+
if (util !== null) {
|
|
53
|
+
this.samples.push(util);
|
|
54
|
+
}
|
|
55
|
+
}
|
|
56
|
+
|
|
57
|
+
stop() {
|
|
58
|
+
this.isMonitoring = false;
|
|
59
|
+
if (this.timer) {
|
|
60
|
+
clearInterval(this.timer);
|
|
61
|
+
this.timer = null;
|
|
62
|
+
}
|
|
63
|
+
|
|
64
|
+
if (this.samples.length === 0) {
|
|
65
|
+
return { avg: null, peak: null, samplesCount: 0 };
|
|
66
|
+
}
|
|
67
|
+
|
|
68
|
+
const max = Math.max(...this.samples);
|
|
69
|
+
const sum = this.samples.reduce((a, b) => a + b, 0);
|
|
70
|
+
const avg = Math.round(sum / this.samples.length);
|
|
71
|
+
|
|
72
|
+
return {
|
|
73
|
+
avg,
|
|
74
|
+
peak: max,
|
|
75
|
+
samplesCount: this.samples.length,
|
|
76
|
+
samples: this.samples,
|
|
77
|
+
};
|
|
78
|
+
}
|
|
79
|
+
}
|
|
80
|
+
|
|
81
|
+
export class ExecutionTracer {
|
|
82
|
+
constructor(name = "Batch Inference") {
|
|
83
|
+
this.name = name;
|
|
84
|
+
this.stages = [];
|
|
85
|
+
this.currentStage = null;
|
|
86
|
+
this.startTime = Date.now();
|
|
87
|
+
}
|
|
88
|
+
|
|
89
|
+
startStage(stageName, category = "CPU") {
|
|
90
|
+
const now = performance.now();
|
|
91
|
+
if (this.currentStage) {
|
|
92
|
+
this.currentStage.duration = now - this.currentStage.startTime;
|
|
93
|
+
this.stages.push(this.currentStage);
|
|
94
|
+
}
|
|
95
|
+
this.currentStage = {
|
|
96
|
+
name: stageName,
|
|
97
|
+
category, // "CPU" or "GPU"
|
|
98
|
+
startTime: now,
|
|
99
|
+
duration: 0,
|
|
100
|
+
};
|
|
101
|
+
}
|
|
102
|
+
|
|
103
|
+
endStage() {
|
|
104
|
+
const now = performance.now();
|
|
105
|
+
if (this.currentStage) {
|
|
106
|
+
this.currentStage.duration = now - this.currentStage.startTime;
|
|
107
|
+
this.stages.push(this.currentStage);
|
|
108
|
+
this.currentStage = null;
|
|
109
|
+
}
|
|
110
|
+
}
|
|
111
|
+
|
|
112
|
+
getSummary(gpuStats = null) {
|
|
113
|
+
this.endStage();
|
|
114
|
+
const totalMs = this.stages.reduce((sum, s) => sum + s.duration, 0);
|
|
115
|
+
|
|
116
|
+
const breakdown = this.stages.map((s) => ({
|
|
117
|
+
name: s.name,
|
|
118
|
+
category: s.category,
|
|
119
|
+
durationMs: parseFloat(s.duration.toFixed(2)),
|
|
120
|
+
pct: totalMs > 0 ? parseFloat(((s.duration / totalMs) * 100).toFixed(1)) : 0,
|
|
121
|
+
}));
|
|
122
|
+
|
|
123
|
+
const cpuTime = this.stages.filter((s) => s.category === "CPU").reduce((sum, s) => sum + s.duration, 0);
|
|
124
|
+
const gpuTime = this.stages.filter((s) => s.category === "GPU").reduce((sum, s) => sum + s.duration, 0);
|
|
125
|
+
|
|
126
|
+
return {
|
|
127
|
+
name: this.name,
|
|
128
|
+
totalMs: parseFloat(totalMs.toFixed(2)),
|
|
129
|
+
cpuMs: parseFloat(cpuTime.toFixed(2)),
|
|
130
|
+
gpuMs: parseFloat(gpuTime.toFixed(2)),
|
|
131
|
+
cpuPct: totalMs > 0 ? parseFloat(((cpuTime / totalMs) * 100).toFixed(1)) : 0,
|
|
132
|
+
gpuPct: totalMs > 0 ? parseFloat(((gpuTime / totalMs) * 100).toFixed(1)) : 0,
|
|
133
|
+
breakdown,
|
|
134
|
+
gpuStats,
|
|
135
|
+
};
|
|
136
|
+
}
|
|
137
|
+
|
|
138
|
+
printTraceReport(gpuStats = null, minGpuThreshold = 0) {
|
|
139
|
+
const summary = this.getSummary(gpuStats);
|
|
140
|
+
const line = "─".repeat(65);
|
|
141
|
+
// stderr, never stdout: stdout is the MCP JSON-RPC channel and any stray
|
|
142
|
+
// write there corrupts the protocol stream.
|
|
143
|
+
const out = (msg) => console.error(msg);
|
|
144
|
+
|
|
145
|
+
out(`\n┌${line}┐`);
|
|
146
|
+
out(`│ CPU/GPU OPERATION TRACE & BOTTLENECK PROFILE: ${summary.name.padEnd(16)} │`);
|
|
147
|
+
out(`├${line}┤`);
|
|
148
|
+
|
|
149
|
+
for (const b of summary.breakdown) {
|
|
150
|
+
const icon = b.category === "GPU" ? "⚡ [GPU]" : "💻 [CPU]";
|
|
151
|
+
const label = `${icon} ${b.name}`.padEnd(42);
|
|
152
|
+
const timeStr = `${b.durationMs.toFixed(1)}ms (${b.pct.toFixed(1)}%)`.padStart(18);
|
|
153
|
+
out(`│ ${label}${timeStr} │`);
|
|
154
|
+
}
|
|
155
|
+
|
|
156
|
+
out(`├${line}┤`);
|
|
157
|
+
out(`│ Total Batch Time: ${summary.totalMs.toFixed(1)}ms | CPU Time: ${summary.cpuMs.toFixed(1)}ms (${summary.cpuPct}%) | GPU Engine Time: ${summary.gpuMs.toFixed(1)}ms (${summary.gpuPct}%) │`);
|
|
158
|
+
|
|
159
|
+
if (gpuStats && gpuStats.peak !== null) {
|
|
160
|
+
const statsBadge = `Peak: ${gpuStats.peak}%, Avg: ${gpuStats.avg}%`;
|
|
161
|
+
out(`│ GPU Hardware Load: ${statsBadge.padEnd(44)} │`);
|
|
162
|
+
out(`└${line}┘\n`);
|
|
163
|
+
} else {
|
|
164
|
+
out(`└${line}┘\n`);
|
|
165
|
+
}
|
|
166
|
+
|
|
167
|
+
return summary;
|
|
168
|
+
}
|
|
169
|
+
}
|
|
@@ -255,7 +255,7 @@ export function formatInputText(text, isQuery = false, modelName = null, instruc
|
|
|
255
255
|
return cleanText;
|
|
256
256
|
}
|
|
257
257
|
|
|
258
|
-
export async function embedText(text, isQuery = false, modelName = null, progressCallback = null, instruction = null) {
|
|
258
|
+
export async function embedText(text, isQuery = false, modelName = null, progressCallback = null, instruction = null, vectorDimension = null) {
|
|
259
259
|
const targetModel = modelName || getConfig().embeddingModel || "Xenova/multilingual-e5-small";
|
|
260
260
|
const extractor = await getExtractor(targetModel, progressCallback);
|
|
261
261
|
const formattedText = formatInputText(text, isQuery, targetModel, instruction);
|
|
@@ -270,14 +270,14 @@ export async function embedText(text, isQuery = false, modelName = null, progres
|
|
|
270
270
|
max_length: maxLen,
|
|
271
271
|
});
|
|
272
272
|
|
|
273
|
-
const result = output.data.slice();
|
|
273
|
+
const result = applyFixedDimension(output.data.slice(), vectorDimension);
|
|
274
274
|
if (typeof output.dispose === 'function') {
|
|
275
275
|
output.dispose();
|
|
276
276
|
}
|
|
277
277
|
return result;
|
|
278
278
|
}
|
|
279
279
|
|
|
280
|
-
export async function embedBatch(texts, isQuery = false, modelName = null, progressCallback = null, instruction = null, traceOptions = {}) {
|
|
280
|
+
export async function embedBatch(texts, isQuery = false, modelName = null, progressCallback = null, instruction = null, traceOptions = {}, vectorDimension = null) {
|
|
281
281
|
if (!texts || texts.length === 0) return [];
|
|
282
282
|
const targetModel = modelName || getConfig().embeddingModel || "Xenova/multilingual-e5-small";
|
|
283
283
|
|
|
@@ -356,7 +356,7 @@ export async function embedBatch(texts, isQuery = false, modelName = null, progr
|
|
|
356
356
|
|
|
357
357
|
for (let i = 0; i < batchSize; i++) {
|
|
358
358
|
const byteOffset = i * vectorDim;
|
|
359
|
-
allResults.push(rawData.slice(byteOffset, byteOffset + vectorDim));
|
|
359
|
+
allResults.push(applyFixedDimension(rawData.slice(byteOffset, byteOffset + vectorDim), vectorDimension));
|
|
360
360
|
}
|
|
361
361
|
|
|
362
362
|
if (typeof output.dispose === 'function') {
|
|
@@ -474,6 +474,19 @@ export function vectorToBuffer(float32Array) {
|
|
|
474
474
|
return Buffer.from(float32Array.buffer, float32Array.byteOffset, float32Array.byteLength);
|
|
475
475
|
}
|
|
476
476
|
|
|
477
|
+
export function resizeVector(float32Array, targetDim) {
|
|
478
|
+
const dim = Number(targetDim);
|
|
479
|
+
if (!dim || dim <= 0 || !float32Array || float32Array.length === dim) return float32Array;
|
|
480
|
+
const out = new Float32Array(dim);
|
|
481
|
+
out.set(float32Array.subarray(0, Math.min(float32Array.length, dim)));
|
|
482
|
+
return out;
|
|
483
|
+
}
|
|
484
|
+
|
|
485
|
+
function applyFixedDimension(float32Array, overrideDim = null) {
|
|
486
|
+
const dim = overrideDim !== null && overrideDim !== undefined ? overrideDim : (getConfig().vectorDimension || 0);
|
|
487
|
+
return resizeVector(float32Array, dim);
|
|
488
|
+
}
|
|
489
|
+
|
|
477
490
|
export function bufferToVector(buffer) {
|
|
478
491
|
return new Float32Array(
|
|
479
492
|
buffer.buffer,
|