dsh-memu-memory 1.1.2 → 1.1.3
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/AGENTS.md +8 -6
- package/README.md +21 -5
- package/README.zh.md +21 -5
- package/docs/contributing.md +5 -3
- package/docs/env-and-commands.md +13 -0
- package/docs/host-api.md +35 -3
- package/docs/memory-semantics.md +32 -5
- package/docs/troubleshooting.md +36 -7
- package/lib/activity-log.js +190 -0
- package/lib/backup.js +1103 -0
- package/lib/binding-diagnostics.js +630 -0
- package/lib/client.js +697 -49
- package/lib/curator.js +81 -37
- package/lib/cycle.js +3 -0
- package/lib/embedding-catalog.js +763 -0
- package/lib/embedding-download.js +407 -0
- package/lib/embedding-runtime.js +710 -0
- package/lib/engine.js +2 -2
- package/lib/index.js +36 -0
- package/lib/inject.js +131 -30
- package/lib/local-embedding-entry.js +145 -0
- package/lib/local-embedding-manager.js +249 -0
- package/lib/local-embedding-protocol.js +173 -0
- package/lib/local-embedding-worker.js +652 -0
- package/lib/local-vector-service.js +283 -0
- package/lib/memory.js +385 -45
- package/lib/memu-config.js +10 -22
- package/lib/memu.js +53 -19
- package/lib/paths.js +24 -0
- package/lib/private.js +139 -13
- package/lib/rpc.js +45 -4
- package/lib/state.js +109 -161
- package/lib/tools.js +9 -5
- package/lib/vector-migration.js +878 -0
- package/lib/vector-profile.js +941 -0
- package/package.json +3 -2
- package/scripts/acceptance-local.mjs +5 -4
- package/scripts/chain-matrix.mjs +527 -0
- package/scripts/check-all.mjs +19 -0
- package/scripts/check-doc-drift.mjs +353 -0
- package/scripts/diagnose-bindings.mjs +176 -0
- package/scripts/local-runtime-install.mjs +142 -0
- package/scripts/local-runtime-verify.mjs +56 -0
- package/scripts/local-vector-verify.mjs +93 -0
- package/scripts/recall-bench.mjs +813 -0
- package/scripts/repetition-ab-scenarios.mjs +93 -0
- package/scripts/repetition-detector-snapshot.mjs +314 -0
- package/scripts/repetition-eval-live.mjs +360 -0
- package/scripts/repetition-eval.mjs +702 -0
|
@@ -0,0 +1,652 @@
|
|
|
1
|
+
import http from 'node:http'
|
|
2
|
+
import fs from 'node:fs'
|
|
3
|
+
import path from 'node:path'
|
|
4
|
+
import crypto from 'node:crypto'
|
|
5
|
+
import { createRequire } from 'node:module'
|
|
6
|
+
import { pathToFileURL } from 'node:url'
|
|
7
|
+
import {
|
|
8
|
+
EMBEDDINGS_PATH,
|
|
9
|
+
HEALTH_PATH,
|
|
10
|
+
DEFAULT_MAX_BODY_BYTES,
|
|
11
|
+
DEFAULT_MAX_BATCH,
|
|
12
|
+
ENCODINGS,
|
|
13
|
+
MAX_VECTOR_DIMENSION,
|
|
14
|
+
LocalEmbeddingError,
|
|
15
|
+
assertVector,
|
|
16
|
+
buildEmbeddingResponse,
|
|
17
|
+
errorBody,
|
|
18
|
+
errorStatus,
|
|
19
|
+
parseEmbeddingRequest,
|
|
20
|
+
} from './local-embedding-protocol.js'
|
|
21
|
+
|
|
22
|
+
// §14.5: loopback only, port 0, random high-entropy token from the parent. Nothing
|
|
23
|
+
// here ever reaches the network, downloads a model, or logs a token or input text.
|
|
24
|
+
export const LOOPBACK_HOSTS = Object.freeze(['127.0.0.1', '::1', 'localhost'])
|
|
25
|
+
export const MIN_TOKEN_LENGTH = 16
|
|
26
|
+
const SMOKE_TEXT = 'embedding readiness probe'
|
|
27
|
+
|
|
28
|
+
function asError(message, code, status, detail = null) {
|
|
29
|
+
return new LocalEmbeddingError(message, code, status, detail)
|
|
30
|
+
}
|
|
31
|
+
|
|
32
|
+
function isLoopbackHostname(hostname) {
|
|
33
|
+
const text = String(hostname || '').toLowerCase()
|
|
34
|
+
return LOOPBACK_HOSTS.includes(text)
|
|
35
|
+
}
|
|
36
|
+
|
|
37
|
+
// The model id is what MemU passes as MEMU_EMBED_MODEL: the stable catalog id plus
|
|
38
|
+
// the pinned revision and processing protocol, so a different weight/quota/pooling
|
|
39
|
+
// build can never silently answer for the old one (§14.8: same dimension but a
|
|
40
|
+
// different space is still a different space).
|
|
41
|
+
export function manifestModelId(manifest) {
|
|
42
|
+
if (!manifest || typeof manifest !== 'object') return ''
|
|
43
|
+
if (typeof manifest.model === 'string' && manifest.model.trim()) return manifest.model.trim()
|
|
44
|
+
const id = typeof manifest.id === 'string' ? manifest.id.trim() : ''
|
|
45
|
+
if (!id) return ''
|
|
46
|
+
const revision = typeof manifest.revision === 'string' ? manifest.revision.trim() : ''
|
|
47
|
+
const protocol = typeof manifest.protocol === 'string' ? manifest.protocol.trim() : ''
|
|
48
|
+
const suffix = [revision, protocol].filter(Boolean).join('+')
|
|
49
|
+
return suffix ? id + '@' + suffix : id
|
|
50
|
+
}
|
|
51
|
+
|
|
52
|
+
export function manifestDimension(manifest) {
|
|
53
|
+
const candidates = [
|
|
54
|
+
manifest && manifest.dimension,
|
|
55
|
+
manifest && manifest.expectedDimension,
|
|
56
|
+
manifest && manifest.output && manifest.output.dimension,
|
|
57
|
+
manifest && manifest.embedding && manifest.embedding.dimension,
|
|
58
|
+
]
|
|
59
|
+
for (const value of candidates) {
|
|
60
|
+
const number = Number(value)
|
|
61
|
+
if (Number.isInteger(number) && number > 0) {
|
|
62
|
+
if (number > MAX_VECTOR_DIMENSION) throw asError('manifest dimension ' + number + ' exceeds the protocol limit', 'LOCAL_VECTOR_MANIFEST_INCOMPLETE', 503)
|
|
63
|
+
return number
|
|
64
|
+
}
|
|
65
|
+
}
|
|
66
|
+
return 0
|
|
67
|
+
}
|
|
68
|
+
|
|
69
|
+
// A model is servable only when the installer marked it verified. pending /
|
|
70
|
+
// candidate / downloading entries exist in the catalog, but a half-installed
|
|
71
|
+
// artifact must never answer a request, and nothing here starts a download to
|
|
72
|
+
// fix that (§14.4).
|
|
73
|
+
export function assertServableManifest({ manifest, modelDir } = {}) {
|
|
74
|
+
if (!manifest || typeof manifest !== 'object' || Array.isArray(manifest)) {
|
|
75
|
+
throw asError('a model manifest is required to start the embedding worker', 'LOCAL_VECTOR_MANIFEST_INVALID', 503)
|
|
76
|
+
}
|
|
77
|
+
const state = String(manifest.status || manifest.state || manifest.verification?.status || '').toLowerCase()
|
|
78
|
+
if (manifest.pending === true || state === 'pending' || state === 'candidate' || state === 'downloading') {
|
|
79
|
+
throw asError('model ' + (manifestModelId(manifest) || '(unnamed)') + ' is ' + (state || 'pending') + ', not verified; finish the managed download first', 'LOCAL_VECTOR_MANIFEST_UNVERIFIED', 503)
|
|
80
|
+
}
|
|
81
|
+
if (manifest.verified !== true && state !== 'verified') {
|
|
82
|
+
throw asError('model ' + (manifestModelId(manifest) || '(unnamed)') + ' is not verified; refusing to serve it', 'LOCAL_VECTOR_MANIFEST_UNVERIFIED', 503)
|
|
83
|
+
}
|
|
84
|
+
const dir = typeof modelDir === 'string' ? modelDir.trim() : ''
|
|
85
|
+
if (!dir) throw asError('modelDir is required', 'LOCAL_VECTOR_MODEL_DIR_MISSING', 503)
|
|
86
|
+
let stat = null
|
|
87
|
+
try { stat = fs.statSync(dir) } catch { stat = null }
|
|
88
|
+
if (!stat || !stat.isDirectory()) throw asError('modelDir ' + dir + ' is missing; installed model files are required', 'LOCAL_VECTOR_MODEL_DIR_MISSING', 503)
|
|
89
|
+
const model = manifestModelId(manifest)
|
|
90
|
+
if (!model) throw asError('manifest is missing a stable id', 'LOCAL_VECTOR_MANIFEST_INCOMPLETE', 503)
|
|
91
|
+
const dimension = manifestDimension(manifest)
|
|
92
|
+
if (!dimension) throw asError('manifest is missing expected dimension', 'LOCAL_VECTOR_MANIFEST_INCOMPLETE', 503)
|
|
93
|
+
return { model, dimension, modelDir: path.resolve(dir) }
|
|
94
|
+
}
|
|
95
|
+
|
|
96
|
+
// External-data graphs, tokenizer files and the model directory itself are all
|
|
97
|
+
// attacker-adjacent input: absolute paths, .. segments, symlinks or junctions that
|
|
98
|
+
// resolve outside the model directory are refused before anything is opened.
|
|
99
|
+
export function resolveInsideModelDir(modelDir, relPath) {
|
|
100
|
+
const dir = typeof modelDir === 'string' ? modelDir.trim() : ''
|
|
101
|
+
if (!dir) throw asError('modelDir is required', 'LOCAL_VECTOR_MODEL_DIR_MISSING', 503)
|
|
102
|
+
let root = null
|
|
103
|
+
try { root = fs.realpathSync(path.resolve(dir)) } catch { root = null }
|
|
104
|
+
if (!root) throw asError('modelDir ' + dir + ' is missing', 'LOCAL_VECTOR_MODEL_DIR_MISSING', 503)
|
|
105
|
+
const rel = typeof relPath === 'string' ? relPath.trim() : ''
|
|
106
|
+
if (!rel) throw asError('manifest names no model file', 'LOCAL_VECTOR_MANIFEST_INCOMPLETE', 503)
|
|
107
|
+
if (path.isAbsolute(rel) || /^[A-Za-z]:/.test(rel)) {
|
|
108
|
+
throw asError('model file ' + rel + ' must be relative to the model directory', 'LOCAL_VECTOR_MANIFEST_INCOMPLETE', 503)
|
|
109
|
+
}
|
|
110
|
+
const segments = rel.split(/[\\/]+/)
|
|
111
|
+
if (segments.some((segment) => segment === '..')) {
|
|
112
|
+
throw asError('model file ' + rel + ' escapes the model directory', 'LOCAL_VECTOR_MANIFEST_INCOMPLETE', 503)
|
|
113
|
+
}
|
|
114
|
+
const target = path.resolve(root, rel)
|
|
115
|
+
if (target !== root && !target.startsWith(root + path.sep)) {
|
|
116
|
+
throw asError('model file ' + rel + ' escapes the model directory', 'LOCAL_VECTOR_MANIFEST_INCOMPLETE', 503)
|
|
117
|
+
}
|
|
118
|
+
let real = null
|
|
119
|
+
try { real = fs.realpathSync(target) } catch { real = null }
|
|
120
|
+
if (!real) throw asError('model file ' + rel + ' is missing', 'LOCAL_VECTOR_MODEL_FILE_MISSING', 503)
|
|
121
|
+
if (real !== root && !real.startsWith(root + path.sep)) {
|
|
122
|
+
throw asError('model file ' + rel + ' resolves outside the model directory', 'LOCAL_VECTOR_MANIFEST_INCOMPLETE', 503)
|
|
123
|
+
}
|
|
124
|
+
return real
|
|
125
|
+
}
|
|
126
|
+
|
|
127
|
+
function runtimeUnavailable(kind, specifier, runtimeDir, error) {
|
|
128
|
+
const where = runtimeDir ? ' in ' + runtimeDir : ''
|
|
129
|
+
const reason = String((error && (error.code || error.message)) || error || 'not found')
|
|
130
|
+
return asError(
|
|
131
|
+
kind + ' module "' + specifier + '" is not installed' + where + ' (' + reason + '); install the pinned artifact instead of downloading at runtime',
|
|
132
|
+
'LOCAL_VECTOR_RUNTIME_UNAVAILABLE',
|
|
133
|
+
503,
|
|
134
|
+
{ module: specifier, kind },
|
|
135
|
+
)
|
|
136
|
+
}
|
|
137
|
+
|
|
138
|
+
// Resolve only from the explicitly installed managed runtime. Runtime imports never
|
|
139
|
+
// consult the host application's node_modules and never trigger an install.
|
|
140
|
+
async function importFrom(specifier, runtimeDir) {
|
|
141
|
+
const value = String(specifier || '').trim()
|
|
142
|
+
if (!value) throw new Error('module specifier is empty')
|
|
143
|
+
if (path.isAbsolute(value) || value.startsWith('file:')) return import(value)
|
|
144
|
+
if (!runtimeDir || !String(runtimeDir).trim()) {
|
|
145
|
+
const error = new Error('managed runtimeDir is required for module ' + value)
|
|
146
|
+
error.code = 'ERR_MODULE_NOT_FOUND'
|
|
147
|
+
throw error
|
|
148
|
+
}
|
|
149
|
+
const requireFrom = createRequire(path.join(path.resolve(String(runtimeDir).trim()), 'entry.cjs'))
|
|
150
|
+
const resolved = requireFrom.resolve(value)
|
|
151
|
+
return import(pathToFileURL(resolved).href)
|
|
152
|
+
}
|
|
153
|
+
|
|
154
|
+
function manifestFilePath(manifest, role, fallback) {
|
|
155
|
+
const files = Array.isArray(manifest?.files) ? manifest.files : []
|
|
156
|
+
const match = files.find(file => file && file.role === role)
|
|
157
|
+
const declared = role === 'graph' ? manifest?.graph?.path : (manifest?.tokenizer?.file || manifest?.tokenizer?.path)
|
|
158
|
+
return declared || match?.path || fallback
|
|
159
|
+
}
|
|
160
|
+
|
|
161
|
+
async function openTokenizer(mod, modelDir, manifest = {}) {
|
|
162
|
+
const tokenizerPath = resolveInsideModelDir(modelDir, manifestFilePath(manifest, 'tokenizer', 'tokenizer.json'))
|
|
163
|
+
const configPath = resolveInsideModelDir(modelDir, manifestFilePath(manifest, 'tokenizer-config', 'tokenizer_config.json'))
|
|
164
|
+
const Tokenizer = mod.Tokenizer || (mod.default && mod.default.Tokenizer)
|
|
165
|
+
if (Tokenizer && typeof Tokenizer.fromFile === 'function') return Tokenizer.fromFile(tokenizerPath)
|
|
166
|
+
if (Tokenizer && typeof Tokenizer === 'function') {
|
|
167
|
+
const tokenizerJson = JSON.parse(fs.readFileSync(tokenizerPath, 'utf8'))
|
|
168
|
+
let config = {}
|
|
169
|
+
try { config = JSON.parse(fs.readFileSync(configPath, 'utf8')) } catch {}
|
|
170
|
+
return new Tokenizer(tokenizerJson, config)
|
|
171
|
+
}
|
|
172
|
+
const fromFile = mod.fromFile || (mod.default && mod.default.fromFile)
|
|
173
|
+
if (typeof fromFile === 'function') return fromFile(tokenizerPath)
|
|
174
|
+
throw new Error('tokenizer module exposes no supported local tokenizer constructor')
|
|
175
|
+
}
|
|
176
|
+
|
|
177
|
+
function padRows(rows, pad, requestedWidth = null) {
|
|
178
|
+
const width = Number.isInteger(requestedWidth) && requestedWidth > 0 ? requestedWidth : rows.reduce((max, row) => Math.max(max, row.length), 0)
|
|
179
|
+
return rows.map((row) => {
|
|
180
|
+
if (row.length === width) return row
|
|
181
|
+
if (row.length > width) return row.slice(0, width)
|
|
182
|
+
const next = row.slice()
|
|
183
|
+
while (next.length < width) next.push(pad)
|
|
184
|
+
return next
|
|
185
|
+
})
|
|
186
|
+
}
|
|
187
|
+
|
|
188
|
+
// Two real tokenizer shapes: the Rust binding (encode -> Encoding with getIds) and
|
|
189
|
+
// transformers.js (await tokenizer(texts, opts) -> {input_ids, attention_mask}).
|
|
190
|
+
async function tokenizeBatch(tokenizer, texts, maxTokens) {
|
|
191
|
+
const limit = Number.isInteger(maxTokens) && maxTokens > 0 ? maxTokens : 512
|
|
192
|
+
if (tokenizer && typeof tokenizer.encode === 'function') {
|
|
193
|
+
const ids = []
|
|
194
|
+
const mask = []
|
|
195
|
+
const types = []
|
|
196
|
+
for (const text of texts) {
|
|
197
|
+
let encoding = null
|
|
198
|
+
try { encoding = await tokenizer.encode(String(text), { add_special_tokens: true, return_token_type_ids: true }) } catch { encoding = await tokenizer.encode(String(text)) }
|
|
199
|
+
const rawIds = typeof encoding?.getIds === 'function' ? Array.from(encoding.getIds(), Number) : Array.isArray(encoding?.ids) ? Array.from(encoding.ids, Number) : null
|
|
200
|
+
if (!rawIds) throw new Error('tokenizer.encode did not return token ids')
|
|
201
|
+
const rawMask = typeof encoding?.getAttentionMask === 'function' ? Array.from(encoding.getAttentionMask(), Number) : Array.isArray(encoding?.attention_mask) ? Array.from(encoding.attention_mask, Number) : rawIds.map(() => 1)
|
|
202
|
+
const rawTypes = typeof encoding?.getTypeIds === 'function' ? Array.from(encoding.getTypeIds(), Number) : Array.isArray(encoding?.token_type_ids) ? Array.from(encoding.token_type_ids, Number) : rawIds.map(() => 0)
|
|
203
|
+
const truncate = (row) => row.length <= limit ? row : row.slice(0, Math.max(0, limit - 1)).concat(row.slice(-1))
|
|
204
|
+
const rowIds = truncate(rawIds)
|
|
205
|
+
const rowMask = truncate(rawMask)
|
|
206
|
+
const rowTypes = truncate(rawTypes)
|
|
207
|
+
ids.push(rowIds)
|
|
208
|
+
mask.push(rowMask)
|
|
209
|
+
types.push(rowTypes)
|
|
210
|
+
}
|
|
211
|
+
return { ids: padRows(ids, 0), mask: padRows(mask, 0), types: padRows(types, 0) }
|
|
212
|
+
}
|
|
213
|
+
if (typeof tokenizer === 'function') {
|
|
214
|
+
const out = await tokenizer(texts, { padding: true, truncation: true, max_length: limit, return_tensor: false })
|
|
215
|
+
const rows = (value) => (Array.isArray(value) && Array.isArray(value[0]) ? value : [value])
|
|
216
|
+
const idsRows = rows(out && out.input_ids).map((row) => Array.from(row, Number).slice(0, limit))
|
|
217
|
+
const maskRows = rows(out && (out.attention_mask || out.input_ids)).map((row) => Array.from(row, Number).slice(0, limit))
|
|
218
|
+
return { ids: padRows(idsRows, 0), mask: padRows(maskRows, 0), types: idsRows.map((row) => row.map(() => 0)) }
|
|
219
|
+
}
|
|
220
|
+
throw new Error('unsupported tokenizer interface')
|
|
221
|
+
}
|
|
222
|
+
|
|
223
|
+
function poolTensor(tensor, { pooling, mask, normalize }) {
|
|
224
|
+
const dims = Array.from((tensor && tensor.dims) || [])
|
|
225
|
+
const data = tensor && tensor.data
|
|
226
|
+
if (!data || dims.length < 2) throw new Error('graph output is not a hidden-state tensor')
|
|
227
|
+
const batch = dims[0]
|
|
228
|
+
const hidden = dims[dims.length - 1]
|
|
229
|
+
const tokens = dims.length >= 3 ? dims[1] : 1
|
|
230
|
+
const out = []
|
|
231
|
+
for (let row = 0; row < batch; row++) {
|
|
232
|
+
const vector = new Array(hidden).fill(0)
|
|
233
|
+
const base = row * tokens * hidden
|
|
234
|
+
if (pooling === 'mean' && tokens > 1) {
|
|
235
|
+
const weights = mask && mask[row] ? mask[row] : null
|
|
236
|
+
let total = 0
|
|
237
|
+
for (let token = 0; token < tokens; token++) {
|
|
238
|
+
const weight = weights ? Number(weights[token]) || 0 : 1
|
|
239
|
+
if (!weight) continue
|
|
240
|
+
total += weight
|
|
241
|
+
for (let unit = 0; unit < hidden; unit++) vector[unit] += Number(data[base + token * hidden + unit]) * weight
|
|
242
|
+
}
|
|
243
|
+
for (let unit = 0; unit < hidden; unit++) vector[unit] = total > 0 ? vector[unit] / total : 0
|
|
244
|
+
} else {
|
|
245
|
+
for (let unit = 0; unit < hidden; unit++) vector[unit] = Number(data[base + unit])
|
|
246
|
+
}
|
|
247
|
+
if (normalize !== false) {
|
|
248
|
+
let norm = 0
|
|
249
|
+
for (const value of vector) norm += value * value
|
|
250
|
+
norm = Math.sqrt(norm)
|
|
251
|
+
if (norm > 0) for (let unit = 0; unit < hidden; unit++) vector[unit] /= norm
|
|
252
|
+
}
|
|
253
|
+
out.push(vector)
|
|
254
|
+
}
|
|
255
|
+
return out
|
|
256
|
+
}
|
|
257
|
+
|
|
258
|
+
function toBigInt64(rows) {
|
|
259
|
+
const flat = new BigInt64Array(rows.length * rows[0].length)
|
|
260
|
+
for (let row = 0; row < rows.length; row++) {
|
|
261
|
+
for (let column = 0; column < rows[row].length; column++) flat[row * rows[row].length + column] = BigInt(Math.trunc(Number(rows[row][column]) || 0))
|
|
262
|
+
}
|
|
263
|
+
return flat
|
|
264
|
+
}
|
|
265
|
+
|
|
266
|
+
// Real inference path: managed CPU ONNX Runtime + a pinned tokenizer. It is only
|
|
267
|
+
// reached when no adapter is injected, and every failure is reported as
|
|
268
|
+
// LOCAL_VECTOR_RUNTIME_UNAVAILABLE rather than papered over by a fallback model.
|
|
269
|
+
export async function createOnnxAdapter({ modelDir, manifest = {}, runtimeDir, tokenizerModule, runtimeModule } = {}) {
|
|
270
|
+
const graphPath = resolveInsideModelDir(modelDir, manifestFilePath(manifest, 'graph', manifest.onnx?.graph || 'model.onnx'))
|
|
271
|
+
const runtimeSpec = runtimeModule || (manifest.runtime && manifest.runtime.module) || 'onnxruntime-node'
|
|
272
|
+
let ort = null
|
|
273
|
+
try { ort = await importFrom(runtimeSpec, runtimeDir) } catch (error) { throw runtimeUnavailable('runtime', runtimeSpec, runtimeDir, error) }
|
|
274
|
+
const InferenceSession = ort.InferenceSession || (ort.default && ort.default.InferenceSession)
|
|
275
|
+
if (!InferenceSession || typeof InferenceSession.create !== 'function') {
|
|
276
|
+
throw runtimeUnavailable('runtime', runtimeSpec, runtimeDir, new Error('InferenceSession.create is missing'))
|
|
277
|
+
}
|
|
278
|
+
const tokenizerSpec = tokenizerModule || (manifest.tokenizer && manifest.tokenizer.module) || '@huggingface/tokenizers'
|
|
279
|
+
let tokenizer = null
|
|
280
|
+
try {
|
|
281
|
+
const mod = await importFrom(tokenizerSpec, runtimeDir)
|
|
282
|
+
tokenizer = await openTokenizer(mod, modelDir, manifest)
|
|
283
|
+
} catch (error) {
|
|
284
|
+
if (error instanceof LocalEmbeddingError) throw error
|
|
285
|
+
throw runtimeUnavailable('tokenizer', tokenizerSpec, runtimeDir, error)
|
|
286
|
+
}
|
|
287
|
+
|
|
288
|
+
const threads = Number(manifest.runtime && manifest.runtime.threads)
|
|
289
|
+
let session = null
|
|
290
|
+
try {
|
|
291
|
+
session = await InferenceSession.create(graphPath, {
|
|
292
|
+
executionProviders: ['cpu'],
|
|
293
|
+
graphOptimizationLevel: 'all',
|
|
294
|
+
...(Number.isInteger(threads) && threads > 0 ? { intraOpNumThreads: threads } : {}),
|
|
295
|
+
})
|
|
296
|
+
} catch (error) {
|
|
297
|
+
throw asError('embedding graph failed to load: ' + String((error && error.message) || error), 'LOCAL_VECTOR_RUNTIME_UNAVAILABLE', 503, { graph: path.basename(graphPath) })
|
|
298
|
+
}
|
|
299
|
+
|
|
300
|
+
const maxTokens = Number(manifest.maxTokens || (manifest.tokenizer && manifest.tokenizer.maxTokens) || 512)
|
|
301
|
+
const pooling = String(manifest.pooling || 'cls').toLowerCase()
|
|
302
|
+
const normalize = manifest.normalization ? String(manifest.normalization).toLowerCase() === 'l2' : manifest.normalize !== false
|
|
303
|
+
const inputNames = Array.from(session.inputNames || Object.keys(session.inputMetadata || {}))
|
|
304
|
+
|
|
305
|
+
const encode = async (texts) => {
|
|
306
|
+
const batch = await tokenizeBatch(tokenizer, texts, maxTokens)
|
|
307
|
+
const dims = [batch.ids.length, batch.ids[0].length]
|
|
308
|
+
const feeds = {}
|
|
309
|
+
for (const name of inputNames) {
|
|
310
|
+
const key = name.toLowerCase()
|
|
311
|
+
if (/token_type|segment/.test(key)) feeds[name] = new ort.Tensor('int64', toBigInt64(batch.types), dims)
|
|
312
|
+
else if (/attention|mask/.test(key)) feeds[name] = new ort.Tensor('int64', toBigInt64(batch.mask), dims)
|
|
313
|
+
else if (/input_ids|(^|_)ids$/.test(key)) feeds[name] = new ort.Tensor('int64', toBigInt64(batch.ids), dims)
|
|
314
|
+
}
|
|
315
|
+
const outputs = await session.run(feeds)
|
|
316
|
+
const names = Array.from(session.outputNames || Object.keys(outputs))
|
|
317
|
+
const tensor = outputs[names[0]] || Object.values(outputs)[0]
|
|
318
|
+
return poolTensor(tensor, { pooling, mask: batch.mask, normalize })
|
|
319
|
+
}
|
|
320
|
+
|
|
321
|
+
return {
|
|
322
|
+
// Quantized activations depend on other batch rows. Keep native inference
|
|
323
|
+
// single-row so a document and the same query occupy one stable vector space.
|
|
324
|
+
encode: async (texts) => {
|
|
325
|
+
const vectors = []
|
|
326
|
+
for (const text of texts) vectors.push(...await encode([text]))
|
|
327
|
+
return vectors
|
|
328
|
+
},
|
|
329
|
+
close: async () => {
|
|
330
|
+
try { await session.release() } catch {}
|
|
331
|
+
},
|
|
332
|
+
}
|
|
333
|
+
}
|
|
334
|
+
|
|
335
|
+
function bindAdapter(adapter) {
|
|
336
|
+
if (!adapter || typeof adapter !== 'object') {
|
|
337
|
+
throw asError('an embedding adapter with encode(inputs) is required', 'LOCAL_VECTOR_RUNTIME_UNAVAILABLE', 503)
|
|
338
|
+
}
|
|
339
|
+
const fn = [adapter.encode, adapter.embed, adapter.embedBatch].find((candidate) => typeof candidate === 'function')
|
|
340
|
+
if (!fn) throw asError('embedding adapter must expose encode(inputs)', 'LOCAL_VECTOR_RUNTIME_UNAVAILABLE', 503)
|
|
341
|
+
return {
|
|
342
|
+
encode: (inputs) => fn.call(adapter, inputs),
|
|
343
|
+
close: typeof adapter.close === 'function' ? () => adapter.close() : null,
|
|
344
|
+
}
|
|
345
|
+
}
|
|
346
|
+
|
|
347
|
+
// Accepts the shapes a runtime or a test adapter plausibly returns, but never
|
|
348
|
+
// guesses the count: a short/long result is an error, because a shifted batch
|
|
349
|
+
// would silently attach the wrong vector to the wrong memory segment.
|
|
350
|
+
function adapterVectors(output, expected, errorFactory) {
|
|
351
|
+
let list = null
|
|
352
|
+
if (Array.isArray(output)) list = output
|
|
353
|
+
else if (output && typeof output === 'object') {
|
|
354
|
+
if (Array.isArray(output.vectors)) list = output.vectors
|
|
355
|
+
else if (Array.isArray(output.embeddings)) list = output.embeddings
|
|
356
|
+
else if (Array.isArray(output.data)) list = output.data.map((entry) => (entry && typeof entry === 'object' && 'embedding' in entry ? entry.embedding : entry))
|
|
357
|
+
}
|
|
358
|
+
if (!list) throw errorFactory('embedding adapter returned no vectors', 'LOCAL_VECTOR_INVALID_VECTOR')
|
|
359
|
+
if (list.length !== expected) {
|
|
360
|
+
throw errorFactory('embedding adapter returned ' + list.length + ' vectors for ' + expected + ' inputs', 'LOCAL_VECTOR_VECTOR_COUNT_MISMATCH')
|
|
361
|
+
}
|
|
362
|
+
const usage = output && typeof output === 'object' && output.usage && typeof output.usage === 'object' ? output.usage : null
|
|
363
|
+
return { vectors: list, usage }
|
|
364
|
+
}
|
|
365
|
+
|
|
366
|
+
function tokenMatches(header, token) {
|
|
367
|
+
if (typeof header !== 'string' || typeof token !== 'string' || !token) return false
|
|
368
|
+
const match = /^Bearer\s+(.+)$/i.exec(header.trim())
|
|
369
|
+
if (!match) return false
|
|
370
|
+
const provided = Buffer.from(match[1].trim(), 'utf8')
|
|
371
|
+
const expected = Buffer.from(token, 'utf8')
|
|
372
|
+
if (provided.length !== expected.length) return false
|
|
373
|
+
return crypto.timingSafeEqual(provided, expected)
|
|
374
|
+
}
|
|
375
|
+
|
|
376
|
+
function sendJson(res, status, value, extraHeaders = null) {
|
|
377
|
+
if (res.writableEnded || res.headersSent) return
|
|
378
|
+
const payload = Buffer.from(JSON.stringify(value), 'utf8')
|
|
379
|
+
res.writeHead(status, {
|
|
380
|
+
'content-type': 'application/json; charset=utf-8',
|
|
381
|
+
'content-length': String(payload.length),
|
|
382
|
+
'cache-control': 'no-store',
|
|
383
|
+
'x-content-type-options': 'nosniff',
|
|
384
|
+
...(extraHeaders || {}),
|
|
385
|
+
})
|
|
386
|
+
res.end(payload)
|
|
387
|
+
}
|
|
388
|
+
|
|
389
|
+
function readBody(req, limit) {
|
|
390
|
+
return new Promise((resolve, reject) => {
|
|
391
|
+
let size = 0
|
|
392
|
+
let settled = false
|
|
393
|
+
const chunks = []
|
|
394
|
+
const finish = (error, value) => {
|
|
395
|
+
if (settled) return
|
|
396
|
+
settled = true
|
|
397
|
+
if (error) reject(error)
|
|
398
|
+
else resolve(value)
|
|
399
|
+
}
|
|
400
|
+
req.on('data', (chunk) => {
|
|
401
|
+
size += chunk.length
|
|
402
|
+
if (size > limit) {
|
|
403
|
+
finish(asError('request body exceeds ' + limit + ' bytes', 'LOCAL_VECTOR_BODY_TOO_LARGE', 413))
|
|
404
|
+
// Stop reading but keep the socket alive: the caller needs the 413 body,
|
|
405
|
+
// so the connection is only torn down after the response is flushed.
|
|
406
|
+
try { req.pause() } catch {}
|
|
407
|
+
return
|
|
408
|
+
}
|
|
409
|
+
chunks.push(chunk)
|
|
410
|
+
})
|
|
411
|
+
req.on('end', () => finish(null, Buffer.concat(chunks).toString('utf8')))
|
|
412
|
+
req.on('error', (error) => finish(asError('request body could not be read', 'LOCAL_VECTOR_BODY_UNREADABLE', 400, String((error && error.code) || error))))
|
|
413
|
+
})
|
|
414
|
+
}
|
|
415
|
+
|
|
416
|
+
/**
|
|
417
|
+
* §14.5 managed embedding worker: one loopback HTTP server speaking the OpenAI
|
|
418
|
+
* embeddings subset MemU already uses, backed by either an injected adapter (tests
|
|
419
|
+
* only) or the pinned managed ONNX runtime.
|
|
420
|
+
*
|
|
421
|
+
* @returns {Promise<{baseUrl: string, url: string, healthUrl: string, model: string, dimension: number, close: () => Promise<void>, lastError: object|null}>}
|
|
422
|
+
* baseUrl is what MEMU_BASE_URL wants (http://host:port/v1); url is the full
|
|
423
|
+
* embeddings endpoint; healthUrl is the readiness probe outside /v1. lastError is
|
|
424
|
+
* live: set on any 5xx/inference failure and cleared by the next successful
|
|
425
|
+
* request so a healthy worker stops reporting.
|
|
426
|
+
*/
|
|
427
|
+
export async function startEmbeddingWorker(options = {}) {
|
|
428
|
+
const {
|
|
429
|
+
modelDir,
|
|
430
|
+
manifest,
|
|
431
|
+
runtimeDir,
|
|
432
|
+
token,
|
|
433
|
+
adapter = null,
|
|
434
|
+
tokenizerModule,
|
|
435
|
+
runtimeModule,
|
|
436
|
+
host = '127.0.0.1',
|
|
437
|
+
port = 0,
|
|
438
|
+
logger = console,
|
|
439
|
+
signal,
|
|
440
|
+
maxBodyBytes = DEFAULT_MAX_BODY_BYTES,
|
|
441
|
+
maxBatch = DEFAULT_MAX_BATCH,
|
|
442
|
+
} = options
|
|
443
|
+
|
|
444
|
+
const { model, dimension } = assertServableManifest({ manifest, modelDir })
|
|
445
|
+
|
|
446
|
+
if (typeof token !== 'string' || token.length < MIN_TOKEN_LENGTH) {
|
|
447
|
+
throw asError('a random token of at least ' + MIN_TOKEN_LENGTH + ' characters is required', 'LOCAL_VECTOR_TOKEN_INVALID', 400)
|
|
448
|
+
}
|
|
449
|
+
if (!isLoopbackHostname(host)) {
|
|
450
|
+
throw asError('embedding worker only binds loopback, refused ' + String(host), 'LOCAL_VECTOR_BAD_HOST', 400)
|
|
451
|
+
}
|
|
452
|
+
const port_ = Number(port)
|
|
453
|
+
if (!Number.isInteger(port_) || port_ < 0 || port_ > 65535) {
|
|
454
|
+
throw asError('port must be an integer between 0 and 65535', 'LOCAL_VECTOR_BAD_PORT', 400)
|
|
455
|
+
}
|
|
456
|
+
const bodyLimit = Number.isInteger(maxBodyBytes) && maxBodyBytes > 0 ? maxBodyBytes : DEFAULT_MAX_BODY_BYTES
|
|
457
|
+
const batchLimit = Number.isInteger(maxBatch) && maxBatch > 0 ? maxBatch : DEFAULT_MAX_BATCH
|
|
458
|
+
|
|
459
|
+
const encoder = adapter
|
|
460
|
+
? bindAdapter(adapter)
|
|
461
|
+
: await createOnnxAdapter({ modelDir, manifest, runtimeDir, tokenizerModule, runtimeModule })
|
|
462
|
+
|
|
463
|
+
const fail = (message, code, status, detail) => asError(message, code, status, detail)
|
|
464
|
+
const encodeVectors = async (inputs) => {
|
|
465
|
+
let output = null
|
|
466
|
+
try {
|
|
467
|
+
output = await encoder.encode(inputs)
|
|
468
|
+
} catch (error) {
|
|
469
|
+
if (error instanceof LocalEmbeddingError) throw error
|
|
470
|
+
throw fail('embedding inference failed: ' + String((error && error.message) || error), 'LOCAL_VECTOR_INFERENCE_FAILED', 500)
|
|
471
|
+
}
|
|
472
|
+
return adapterVectors(output, inputs.length, fail)
|
|
473
|
+
}
|
|
474
|
+
|
|
475
|
+
let lastError = null
|
|
476
|
+
const noteFailure = (error) => {
|
|
477
|
+
if (errorStatus(error) < 500) return
|
|
478
|
+
lastError = {
|
|
479
|
+
at: new Date().toISOString(),
|
|
480
|
+
code: String((error && error.code) || 'LOCAL_VECTOR_ERROR'),
|
|
481
|
+
message: String((error && error.message) || 'embedding request failed'),
|
|
482
|
+
}
|
|
483
|
+
}
|
|
484
|
+
|
|
485
|
+
// Ready means the graph/tokenizer answered once with a finite, correctly sized,
|
|
486
|
+
// non-zero vector. Without this smoke test the first recall pays the failure.
|
|
487
|
+
let encoderClosed = false
|
|
488
|
+
try {
|
|
489
|
+
const smoke = await encodeVectors([SMOKE_TEXT])
|
|
490
|
+
assertVector(smoke.vectors[0], dimension, 0)
|
|
491
|
+
} catch (error) {
|
|
492
|
+
try { if (encoder.close && !encoderClosed) { encoderClosed = true; await encoder.close() } } catch {}
|
|
493
|
+
if (error instanceof LocalEmbeddingError) throw error
|
|
494
|
+
throw fail('embedding smoke test failed: ' + String((error && error.message) || error), 'LOCAL_VECTOR_SMOKE_FAILED', 503)
|
|
495
|
+
}
|
|
496
|
+
|
|
497
|
+
const startedAt = Date.now()
|
|
498
|
+
const sockets = new Set()
|
|
499
|
+
let closed = false
|
|
500
|
+
|
|
501
|
+
const assertLocalOrigin = (req) => {
|
|
502
|
+
// Any Origin header means a browser page is talking to us: refuse, and never
|
|
503
|
+
// send CORS headers, so a web app cannot use the local index as an oracle.
|
|
504
|
+
if (req.headers.origin !== undefined) throw fail('cross-origin requests are not accepted', 'LOCAL_VECTOR_CORS_REJECTED', 403)
|
|
505
|
+
}
|
|
506
|
+
const assertLocalHost = (req) => {
|
|
507
|
+
const raw = String(req.headers.host || '')
|
|
508
|
+
const hostname = raw.startsWith('[') ? raw.slice(1, raw.indexOf(']')) : raw.split(':')[0]
|
|
509
|
+
if (!hostname) throw fail('host header is required', 'LOCAL_VECTOR_BAD_HOST', 403)
|
|
510
|
+
if (!isLoopbackHostname(hostname)) throw fail('host header must be loopback', 'LOCAL_VECTOR_BAD_HOST', 403)
|
|
511
|
+
}
|
|
512
|
+
const assertAuthorized = (req) => {
|
|
513
|
+
if (tokenMatches(req.headers.authorization, token)) return
|
|
514
|
+
throw fail('missing or invalid bearer token', 'LOCAL_VECTOR_UNAUTHORIZED', 401)
|
|
515
|
+
}
|
|
516
|
+
const methodNotAllowed = (res, allowed) => {
|
|
517
|
+
sendJson(res, 405, errorBody(fail('method not allowed', 'LOCAL_VECTOR_METHOD_NOT_ALLOWED', 405)), { allow: allowed })
|
|
518
|
+
}
|
|
519
|
+
|
|
520
|
+
const respondHealth = (req, res) => {
|
|
521
|
+
if (req.method !== 'GET' && req.method !== 'HEAD') return methodNotAllowed(res, 'GET')
|
|
522
|
+
lastError = null
|
|
523
|
+
if (req.method === 'HEAD') {
|
|
524
|
+
res.writeHead(200, { 'content-type': 'application/json; charset=utf-8', 'cache-control': 'no-store' })
|
|
525
|
+
res.end()
|
|
526
|
+
return
|
|
527
|
+
}
|
|
528
|
+
sendJson(res, 200, {
|
|
529
|
+
object: 'local_embedding_health',
|
|
530
|
+
status: 'ok',
|
|
531
|
+
ready: !closed,
|
|
532
|
+
model,
|
|
533
|
+
dimension,
|
|
534
|
+
maxBatch: batchLimit,
|
|
535
|
+
encodings: ENCODINGS,
|
|
536
|
+
uptimeMs: Date.now() - startedAt,
|
|
537
|
+
})
|
|
538
|
+
}
|
|
539
|
+
|
|
540
|
+
const respondEmbeddings = async (req, res) => {
|
|
541
|
+
if (req.method !== 'POST') return methodNotAllowed(res, 'POST')
|
|
542
|
+
const contentType = String(req.headers['content-type'] || '')
|
|
543
|
+
if (!/^application\/json\b/i.test(contentType)) {
|
|
544
|
+
throw fail('content-type must be application/json', 'LOCAL_VECTOR_UNSUPPORTED_MEDIA_TYPE', 415)
|
|
545
|
+
}
|
|
546
|
+
const declared = Number(req.headers['content-length'])
|
|
547
|
+
if (Number.isFinite(declared) && declared > bodyLimit) {
|
|
548
|
+
throw fail('request body exceeds ' + bodyLimit + ' bytes', 'LOCAL_VECTOR_BODY_TOO_LARGE', 413)
|
|
549
|
+
}
|
|
550
|
+
const raw = await readBody(req, bodyLimit)
|
|
551
|
+
let body = null
|
|
552
|
+
try { body = JSON.parse(raw) } catch { throw fail('request body is not valid JSON', 'LOCAL_VECTOR_BAD_REQUEST', 400) }
|
|
553
|
+
const { inputs, encoding } = parseEmbeddingRequest(body, { model, maxBatch: batchLimit })
|
|
554
|
+
const encoded = await encodeVectors(inputs)
|
|
555
|
+
const vectors = encoded.vectors.map((vector, index) => assertVector(vector, dimension, index))
|
|
556
|
+
sendJson(res, 200, buildEmbeddingResponse({ model, vectors, encoding, usage: encoded.usage }))
|
|
557
|
+
lastError = null
|
|
558
|
+
}
|
|
559
|
+
|
|
560
|
+
const server = http.createServer((req, res) => {
|
|
561
|
+
const route = String(req.url || '/')
|
|
562
|
+
const serve = async () => {
|
|
563
|
+
assertLocalOrigin(req)
|
|
564
|
+
assertLocalHost(req)
|
|
565
|
+
assertAuthorized(req)
|
|
566
|
+
const pathname = new URL(route, 'http://127.0.0.1').pathname
|
|
567
|
+
if (pathname === HEALTH_PATH) return respondHealth(req, res)
|
|
568
|
+
if (pathname === EMBEDDINGS_PATH) return respondEmbeddings(req, res)
|
|
569
|
+
throw fail('unknown endpoint', 'LOCAL_VECTOR_NOT_FOUND', 404)
|
|
570
|
+
}
|
|
571
|
+
serve().catch((error) => {
|
|
572
|
+
noteFailure(error)
|
|
573
|
+
const status = errorStatus(error)
|
|
574
|
+
try {
|
|
575
|
+
logger.warn?.('[dsh-memu-memory] local embedding worker ' + status + ' ' + String((error && error.code) || 'LOCAL_VECTOR_ERROR') + ' ' + String((error && error.message) || error) + ' (' + pathnameOf(route) + ')')
|
|
576
|
+
} catch {}
|
|
577
|
+
sendJson(res, status, errorBody(error), status === 401 ? { 'www-authenticate': 'Bearer' } : null)
|
|
578
|
+
if (status === 413) res.on('finish', () => { try { req.destroy() } catch {} })
|
|
579
|
+
})
|
|
580
|
+
})
|
|
581
|
+
const pathnameOf = (route) => {
|
|
582
|
+
try { return new URL(route, 'http://127.0.0.1').pathname } catch { return '/' }
|
|
583
|
+
}
|
|
584
|
+
server.maxHeadersCount = 64
|
|
585
|
+
server.headersTimeout = 20_000
|
|
586
|
+
server.requestTimeout = 120_000
|
|
587
|
+
server.on('connection', (socket) => {
|
|
588
|
+
sockets.add(socket)
|
|
589
|
+
socket.on('close', () => sockets.delete(socket))
|
|
590
|
+
})
|
|
591
|
+
|
|
592
|
+
const close = async () => {
|
|
593
|
+
if (closed) return
|
|
594
|
+
closed = true
|
|
595
|
+
if (signal) signal.removeEventListener?.('abort', onAbort)
|
|
596
|
+
await new Promise((resolve) => {
|
|
597
|
+
try {
|
|
598
|
+
server.close(() => resolve())
|
|
599
|
+
} catch {
|
|
600
|
+
resolve()
|
|
601
|
+
}
|
|
602
|
+
for (const socket of sockets) {
|
|
603
|
+
try { socket.destroy() } catch {}
|
|
604
|
+
}
|
|
605
|
+
sockets.clear()
|
|
606
|
+
setTimeout(resolve, 2_000).unref?.()
|
|
607
|
+
})
|
|
608
|
+
if (encoder.close && !encoderClosed) {
|
|
609
|
+
encoderClosed = true
|
|
610
|
+
try { await encoder.close() } catch {}
|
|
611
|
+
}
|
|
612
|
+
}
|
|
613
|
+
const onAbort = () => { void close() }
|
|
614
|
+
|
|
615
|
+
await new Promise((resolve, reject) => {
|
|
616
|
+
const onError = (error) => {
|
|
617
|
+
server.removeListener('listening', onListening)
|
|
618
|
+
reject(fail('embedding worker could not listen on ' + host + ':' + port_ + ' (' + String((error && error.code) || error) + ')', 'LOCAL_VECTOR_LISTEN_FAILED', 503))
|
|
619
|
+
}
|
|
620
|
+
const onListening = () => {
|
|
621
|
+
server.removeListener('error', onError)
|
|
622
|
+
resolve()
|
|
623
|
+
}
|
|
624
|
+
server.once('error', onError)
|
|
625
|
+
server.once('listening', onListening)
|
|
626
|
+
server.listen(port_, host)
|
|
627
|
+
})
|
|
628
|
+
|
|
629
|
+
if (signal) {
|
|
630
|
+
if (signal.aborted) {
|
|
631
|
+
await close()
|
|
632
|
+
throw fail('embedding worker start aborted', 'LOCAL_VECTOR_ABORTED', 499)
|
|
633
|
+
}
|
|
634
|
+
signal.addEventListener?.('abort', onAbort, { once: true })
|
|
635
|
+
}
|
|
636
|
+
|
|
637
|
+
const address = server.address()
|
|
638
|
+
const actualPort = address && typeof address === 'object' ? address.port : port_
|
|
639
|
+
const hostForUrl = host === '::1' ? '[::1]' : host
|
|
640
|
+
const origin = 'http://' + hostForUrl + ':' + actualPort
|
|
641
|
+
const handle = {
|
|
642
|
+
baseUrl: origin + '/v1',
|
|
643
|
+
url: origin + EMBEDDINGS_PATH,
|
|
644
|
+
healthUrl: origin + HEALTH_PATH,
|
|
645
|
+
model,
|
|
646
|
+
dimension,
|
|
647
|
+
close,
|
|
648
|
+
}
|
|
649
|
+
Object.defineProperty(handle, 'lastError', { enumerable: true, get: () => lastError })
|
|
650
|
+
try { logger.info?.('[dsh-memu-memory] local embedding worker ready on ' + origin + ' model=' + model + ' dimension=' + dimension) } catch {}
|
|
651
|
+
return handle
|
|
652
|
+
}
|