@aion0/bastion 0.1.12 → 0.1.14
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/README.md +27 -3
- package/README.zh.md +26 -3
- package/dist/classifier/model-manager.d.ts +43 -0
- package/dist/classifier/model-manager.d.ts.map +1 -0
- package/dist/classifier/model-manager.js +161 -0
- package/dist/classifier/model-manager.js.map +1 -0
- package/dist/classifier/onnx-provider.d.ts +22 -0
- package/dist/classifier/onnx-provider.d.ts.map +1 -0
- package/dist/classifier/onnx-provider.js +131 -0
- package/dist/classifier/onnx-provider.js.map +1 -0
- package/dist/cli/commands/plugins.d.ts +3 -0
- package/dist/cli/commands/plugins.d.ts.map +1 -0
- package/dist/cli/commands/plugins.js +201 -0
- package/dist/cli/commands/plugins.js.map +1 -0
- package/dist/cli/commands/start.d.ts.map +1 -1
- package/dist/cli/commands/start.js +0 -18
- package/dist/cli/commands/start.js.map +1 -1
- package/dist/cli/index.js +2 -2
- package/dist/cli/index.js.map +1 -1
- package/dist/config/paths.d.ts +0 -1
- package/dist/config/paths.d.ts.map +1 -1
- package/dist/config/paths.js +0 -1
- package/dist/config/paths.js.map +1 -1
- package/dist/config/schema.d.ts +28 -0
- package/dist/config/schema.d.ts.map +1 -1
- package/dist/core/bootstrap.d.ts.map +1 -1
- package/dist/core/bootstrap.js +5 -0
- package/dist/core/bootstrap.js.map +1 -1
- package/dist/dashboard/api-routes.d.ts.map +1 -1
- package/dist/dashboard/api-routes.js +91 -50
- package/dist/dashboard/api-routes.js.map +1 -1
- package/dist/dashboard/page.d.ts.map +1 -1
- package/dist/dashboard/page.js +122 -67
- package/dist/dashboard/page.js.map +1 -1
- package/dist/dlp/ai-validator.d.ts.map +1 -1
- package/dist/dlp/ai-validator.js +11 -13
- package/dist/dlp/ai-validator.js.map +1 -1
- package/dist/dlp/heuristic-validator.d.ts +24 -0
- package/dist/dlp/heuristic-validator.d.ts.map +1 -0
- package/dist/dlp/heuristic-validator.js +97 -0
- package/dist/dlp/heuristic-validator.js.map +1 -0
- package/dist/plugins/builtin/dlp-scanner.d.ts.map +1 -1
- package/dist/plugins/builtin/dlp-scanner.js +3 -0
- package/dist/plugins/builtin/dlp-scanner.js.map +1 -1
- package/dist/plugins/builtin/pi-classifier.d.ts +22 -0
- package/dist/plugins/builtin/pi-classifier.d.ts.map +1 -0
- package/dist/plugins/builtin/pi-classifier.js +177 -0
- package/dist/plugins/builtin/pi-classifier.js.map +1 -0
- package/dist/plugins/builtin/threat-scorer.d.ts +6 -0
- package/dist/plugins/builtin/threat-scorer.d.ts.map +1 -0
- package/dist/plugins/builtin/threat-scorer.js +266 -0
- package/dist/plugins/builtin/threat-scorer.js.map +1 -0
- package/dist/plugins/builtin/tool-guard.d.ts.map +1 -1
- package/dist/plugins/builtin/tool-guard.js +14 -43
- package/dist/plugins/builtin/tool-guard.js.map +1 -1
- package/dist/plugins/types.d.ts +3 -0
- package/dist/plugins/types.d.ts.map +1 -1
- package/dist/storage/migrations.d.ts.map +1 -1
- package/dist/storage/migrations.js +43 -0
- package/dist/storage/migrations.js.map +1 -1
- package/dist/storage/repositories/taint-marks.d.ts +26 -0
- package/dist/storage/repositories/taint-marks.d.ts.map +1 -0
- package/dist/storage/repositories/taint-marks.js +27 -0
- package/dist/storage/repositories/taint-marks.js.map +1 -0
- package/dist/storage/repositories/threat-score-events.d.ts +27 -0
- package/dist/storage/repositories/threat-score-events.d.ts.map +1 -0
- package/dist/storage/repositories/threat-score-events.js +24 -0
- package/dist/storage/repositories/threat-score-events.js.map +1 -0
- package/dist/storage/repositories/threat-scores.d.ts +26 -0
- package/dist/storage/repositories/threat-scores.d.ts.map +1 -0
- package/dist/storage/repositories/threat-scores.js +42 -0
- package/dist/storage/repositories/threat-scores.js.map +1 -0
- package/dist/storage/repositories/tool-chain-detections.d.ts +24 -0
- package/dist/storage/repositories/tool-chain-detections.d.ts.map +1 -0
- package/dist/storage/repositories/tool-chain-detections.js +27 -0
- package/dist/storage/repositories/tool-chain-detections.js.map +1 -0
- package/dist/tool-guard/chain-detector.d.ts +18 -0
- package/dist/tool-guard/chain-detector.d.ts.map +1 -0
- package/dist/tool-guard/chain-detector.js +55 -0
- package/dist/tool-guard/chain-detector.js.map +1 -0
- package/dist/tool-guard/chain-rules.d.ts +10 -0
- package/dist/tool-guard/chain-rules.d.ts.map +1 -0
- package/dist/tool-guard/chain-rules.js +30 -0
- package/dist/tool-guard/chain-rules.js.map +1 -0
- package/dist/tool-guard/taint-tracker.d.ts +24 -0
- package/dist/tool-guard/taint-tracker.d.ts.map +1 -0
- package/dist/tool-guard/taint-tracker.js +70 -0
- package/dist/tool-guard/taint-tracker.js.map +1 -0
- package/package.json +1 -1
- package/dist/cli/commands/pro.d.ts +0 -3
- package/dist/cli/commands/pro.d.ts.map +0 -1
- package/dist/cli/commands/pro.js +0 -261
- package/dist/cli/commands/pro.js.map +0 -1
- package/dist/license/pro-license.d.ts +0 -11
- package/dist/license/pro-license.d.ts.map +0 -1
- package/dist/license/pro-license.js +0 -88
- package/dist/license/pro-license.js.map +0 -1
- package/dist/license/verify.d.ts +0 -18
- package/dist/license/verify.d.ts.map +0 -1
- package/dist/license/verify.js +0 -71
- package/dist/license/verify.js.map +0 -1
package/README.md
CHANGED
|
@@ -104,6 +104,15 @@ Monitors and blocks dangerous tool calls made by AI agents in real-time. Interce
|
|
|
104
104
|
|
|
105
105
|
Action modes: `audit` (log and alert) or `block` (intercept in real-time, including streaming responses). Desktop notifications and webhook alerts (Slack, Discord) for high-severity matches.
|
|
106
106
|
|
|
107
|
+
### 🔍 Threat Intelligence
|
|
108
|
+
|
|
109
|
+
Session-level threat scoring with multi-signal correlation. Aggregates events from DLP, Tool Guard, and Prompt Injection detection into a unified threat score per session.
|
|
110
|
+
|
|
111
|
+
- **Chain detection** — detects multi-step attack patterns (e.g., credential access followed by network exfiltration)
|
|
112
|
+
- **Taint tracking** — tracks sensitive data fingerprints across tool calls to detect data flow violations
|
|
113
|
+
- **Threat levels** — normal / elevated / high / critical with exponential decay scoring
|
|
114
|
+
- **3 built-in chain rules** — credential→exfil, execution→destruction, credential→publish
|
|
115
|
+
|
|
107
116
|
### 📝 Audit Logger
|
|
108
117
|
|
|
109
118
|
Full request/response history for every AI interaction, encrypted at rest. Session-based timeline with DLP and Tool Guard tags. Any security event automatically creates an audit entry — even if the audit plugin is disabled.
|
|
@@ -116,11 +125,11 @@ Configurable retention with automatic purge. Formatted viewer in the dashboard f
|
|
|
116
125
|
|
|
117
126
|
Real-time security dashboard at `http://127.0.0.1:8420/dashboard`:
|
|
118
127
|
|
|
119
|
-
- **Overview** — Request metrics, cost, tokens, per-provider/model/session breakdown
|
|
128
|
+
- **Overview** — Request metrics, cost, tokens, per-provider/model/session breakdown, threat gauge
|
|
120
129
|
- **DLP** — Findings, config, signature management, standalone test scanner with trace log
|
|
121
|
-
- **Tool Guard** — Tool call history, severity, rule management (built-in + custom)
|
|
130
|
+
- **Tool Guard** — Tool call history, severity, rule management (built-in + custom), threat sessions, chain detections
|
|
122
131
|
- **Audit** — Session timeline, security-tagged entries, formatted request/response viewer
|
|
123
|
-
- **Settings** — Toggle plugins,
|
|
132
|
+
- **Settings** — Toggle plugins, optional features management — all changes apply without restart
|
|
124
133
|
|
|
125
134
|
## How It Works
|
|
126
135
|
|
|
@@ -168,6 +177,20 @@ bash ~/.openclaw/skills/bastion/scripts/docker-setup.sh
|
|
|
168
177
|
|
|
169
178
|
See [@aion0/bastion-skills on GitHub](https://github.com/aiwatching/bastion-skills)
|
|
170
179
|
|
|
180
|
+
## Optional Plugins
|
|
181
|
+
|
|
182
|
+
Bastion's core detection is fully free. For ML-based detection with heavier dependencies (ONNX Runtime, ~436 MB models), install the optional plugin pack:
|
|
183
|
+
|
|
184
|
+
```bash
|
|
185
|
+
bastion plugins install # auto-detects sibling bastion-plugin-api
|
|
186
|
+
bastion plugins list # show installed plugins
|
|
187
|
+
bastion plugins uninstall # remove (prompts to also delete models)
|
|
188
|
+
```
|
|
189
|
+
|
|
190
|
+
Or via the Dashboard → Settings → Optional Features.
|
|
191
|
+
|
|
192
|
+
See [@aiwatching/bastion-pro](https://github.com/aiwatching/bastion-plugin-api) for details.
|
|
193
|
+
|
|
171
194
|
## Documentation
|
|
172
195
|
|
|
173
196
|
| Doc | Description |
|
|
@@ -191,6 +214,7 @@ Everything stays on your machine in `~/.bastion/`:
|
|
|
191
214
|
config.yaml # Your config overrides
|
|
192
215
|
ca.key / ca.crt / certs/ # Local CA & certificates
|
|
193
216
|
.key # AES encryption key for audit data
|
|
217
|
+
models/ # ML models (optional, ~436 MB if plugins installed)
|
|
194
218
|
```
|
|
195
219
|
|
|
196
220
|
## Contributing
|
package/README.zh.md
CHANGED
|
@@ -207,12 +207,12 @@ bastion trust-ca
|
|
|
207
207
|
网关运行时,在浏览器中打开 `http://127.0.0.1:8420/dashboard`。
|
|
208
208
|
|
|
209
209
|
6 个标签页:
|
|
210
|
-
- **Overview** — 请求指标、费用、token
|
|
210
|
+
- **Overview** — 请求指标、费用、token 数、按提供商/模型/会话分类的统计、威胁仪表盘
|
|
211
211
|
- **DLP** — 子标签:Findings(方向、片段、钻取至审计记录)、Config(引擎开关、动作模式、AI 验证、语义规则)、Signatures(远程同步状态、版本追踪、变更日志、模式管理)、Test(独立扫描器,含预设、trace 日志)
|
|
212
|
-
- **Tool Guard** — 子标签:Calls(最近工具调用历史,含严重级别、规则匹配、执行动作)、Rules(26 条内置规则 +
|
|
212
|
+
- **Tool Guard** — 子标签:Calls(最近工具调用历史,含严重级别、规则匹配、执行动作)、Rules(26 条内置规则 + 自定义规则管理,可逐条启用/禁用)、威胁会话、链式检测
|
|
213
213
|
- **Optimizer** — 缓存命中率、节省的 token 数
|
|
214
214
|
- **Audit** — 基于会话的时间线、DLP/Tool Guard 标记条目、摘要预览、格式化请求/响应查看器
|
|
215
|
-
- **Settings** — 切换 plugin
|
|
215
|
+
- **Settings** — 切换 plugin、可选功能管理,运行时修改无需重启
|
|
216
216
|
|
|
217
217
|
## 工作原理
|
|
218
218
|
|
|
@@ -329,6 +329,14 @@ plugins:
|
|
|
329
329
|
- **响应缓存** — 对相同请求进行精确匹配缓存(AES-256-GCM 加密)
|
|
330
330
|
- **空白压缩** — 折叠多余空白以节省 token
|
|
331
331
|
|
|
332
|
+
### 威胁情报
|
|
333
|
+
会话级威胁评分,多信号关联分析。聚合 DLP、Tool Guard、Prompt Injection 检测事件,为每个会话生成统一威胁分数。
|
|
334
|
+
|
|
335
|
+
- **链式检测** — 检测多步攻击模式(如凭证访问后跟网络外泄)
|
|
336
|
+
- **污点追踪** — 跟踪敏感数据指纹在工具调用间的流动,检测数据流违规
|
|
337
|
+
- **威胁等级** — normal / elevated / high / critical,指数衰减评分
|
|
338
|
+
- **3 条内置链式规则** — 凭证→外泄、执行→破坏、凭证→发布
|
|
339
|
+
|
|
332
340
|
### Audit Logger
|
|
333
341
|
存储请求/响应内容(静态加密)以供在 Dashboard 中查看。可配置保留期限,自动清理。即使此 plugin 未启用,DLP 命中也会自动记录审计日志。
|
|
334
342
|
|
|
@@ -443,6 +451,20 @@ BASTION_LOG_LEVEL=debug bastion start
|
|
|
443
451
|
| `GET` | `/api/config` | Current configuration + plugin status |
|
|
444
452
|
| `PUT` | `/api/config` | Update configuration at runtime |
|
|
445
453
|
|
|
454
|
+
## 可选插件
|
|
455
|
+
|
|
456
|
+
Bastion 核心检测能力完全免费。如需基于 ML 的检测(ONNX Runtime,~436 MB 模型),可安装可选插件包:
|
|
457
|
+
|
|
458
|
+
```bash
|
|
459
|
+
bastion plugins install # 自动检测同级 bastion-plugin-api 目录
|
|
460
|
+
bastion plugins list # 查看已安装插件
|
|
461
|
+
bastion plugins uninstall # 卸载(会提示是否同时删除模型文件)
|
|
462
|
+
```
|
|
463
|
+
|
|
464
|
+
也可通过 Dashboard → Settings → Optional Features 管理。
|
|
465
|
+
|
|
466
|
+
详见 [@aiwatching/bastion-pro](https://github.com/aiwatching/bastion-plugin-api)。
|
|
467
|
+
|
|
446
468
|
## 文档
|
|
447
469
|
|
|
448
470
|
- [DLP 引擎架构](docs/dlp.zh.md) — 5 层检测管线详解
|
|
@@ -468,4 +490,5 @@ BASTION_LOG_LEVEL=debug bastion start
|
|
|
468
490
|
.key # AES encryption key for cache & audit
|
|
469
491
|
bastion.pid # Daemon PID file
|
|
470
492
|
bastion.log # Daemon log file
|
|
493
|
+
models/ # ML 模型(可选,安装插件后 ~436 MB)
|
|
471
494
|
```
|
|
@@ -0,0 +1,43 @@
|
|
|
1
|
+
export interface ModelInfo {
|
|
2
|
+
modelId: string;
|
|
3
|
+
dir: string;
|
|
4
|
+
ready: boolean;
|
|
5
|
+
files: string[];
|
|
6
|
+
}
|
|
7
|
+
export interface DownloadProgress {
|
|
8
|
+
file: string;
|
|
9
|
+
received: number;
|
|
10
|
+
total: number;
|
|
11
|
+
percent: number;
|
|
12
|
+
}
|
|
13
|
+
export type ProgressCallback = (progress: DownloadProgress) => void;
|
|
14
|
+
/**
|
|
15
|
+
* Resolve the local directory for a given model ID.
|
|
16
|
+
* Model IDs use '/' which gets replaced with '--' for filesystem safety.
|
|
17
|
+
*/
|
|
18
|
+
export declare function getModelDir(modelId: string): string;
|
|
19
|
+
/**
|
|
20
|
+
* Check if a model is already downloaded and complete.
|
|
21
|
+
*/
|
|
22
|
+
export declare function isModelReady(modelId: string): Promise<boolean>;
|
|
23
|
+
/**
|
|
24
|
+
* Get info about a local model.
|
|
25
|
+
*/
|
|
26
|
+
export declare function getModelInfo(modelId: string): Promise<ModelInfo>;
|
|
27
|
+
/**
|
|
28
|
+
* Download model files from Hugging Face Hub.
|
|
29
|
+
* Files: model.onnx (from onnx/ subdirectory), tokenizer.json, config.json
|
|
30
|
+
*/
|
|
31
|
+
export declare function downloadModel(modelId: string, onProgress?: ProgressCallback, options?: {
|
|
32
|
+
onnxSubdir?: string;
|
|
33
|
+
}): Promise<string>;
|
|
34
|
+
/**
|
|
35
|
+
* Ensure a model is available locally. Downloads if missing.
|
|
36
|
+
*/
|
|
37
|
+
export declare function ensureModel(modelId: string, onProgress?: ProgressCallback, options?: {
|
|
38
|
+
onnxSubdir?: string;
|
|
39
|
+
logger?: {
|
|
40
|
+
info: (msg: string, meta?: Record<string, unknown>) => void;
|
|
41
|
+
};
|
|
42
|
+
}): Promise<string>;
|
|
43
|
+
//# sourceMappingURL=model-manager.d.ts.map
|
|
@@ -0,0 +1 @@
|
|
|
1
|
+
{"version":3,"file":"model-manager.d.ts","sourceRoot":"","sources":["../../src/classifier/model-manager.ts"],"names":[],"mappings":"AASA,MAAM,WAAW,SAAS;IACxB,OAAO,EAAE,MAAM,CAAC;IAChB,GAAG,EAAE,MAAM,CAAC;IACZ,KAAK,EAAE,OAAO,CAAC;IACf,KAAK,EAAE,MAAM,EAAE,CAAC;CACjB;AAED,MAAM,WAAW,gBAAgB;IAC/B,IAAI,EAAE,MAAM,CAAC;IACb,QAAQ,EAAE,MAAM,CAAC;IACjB,KAAK,EAAE,MAAM,CAAC;IACd,OAAO,EAAE,MAAM,CAAC;CACjB;AAED,MAAM,MAAM,gBAAgB,GAAG,CAAC,QAAQ,EAAE,gBAAgB,KAAK,IAAI,CAAC;AAEpE;;;GAGG;AACH,wBAAgB,WAAW,CAAC,OAAO,EAAE,MAAM,GAAG,MAAM,CAGnD;AAED;;GAEG;AACH,wBAAsB,YAAY,CAAC,OAAO,EAAE,MAAM,GAAG,OAAO,CAAC,OAAO,CAAC,CAWpE;AAED;;GAEG;AACH,wBAAsB,YAAY,CAAC,OAAO,EAAE,MAAM,GAAG,OAAO,CAAC,SAAS,CAAC,CAiBtE;AAoFD;;;GAGG;AACH,wBAAsB,aAAa,CACjC,OAAO,EAAE,MAAM,EACf,UAAU,CAAC,EAAE,gBAAgB,EAC7B,OAAO,CAAC,EAAE;IAAE,UAAU,CAAC,EAAE,MAAM,CAAA;CAAE,GAChC,OAAO,CAAC,MAAM,CAAC,CA8BjB;AAED;;GAEG;AACH,wBAAsB,WAAW,CAC/B,OAAO,EAAE,MAAM,EACf,UAAU,CAAC,EAAE,gBAAgB,EAC7B,OAAO,CAAC,EAAE;IAAE,UAAU,CAAC,EAAE,MAAM,CAAC;IAAC,MAAM,CAAC,EAAE;QAAE,IAAI,EAAE,CAAC,GAAG,EAAE,MAAM,EAAE,IAAI,CAAC,EAAE,MAAM,CAAC,MAAM,EAAE,OAAO,CAAC,KAAK,IAAI,CAAA;KAAE,CAAA;CAAE,GAC1G,OAAO,CAAC,MAAM,CAAC,CAQjB"}
|
|
@@ -0,0 +1,161 @@
|
|
|
1
|
+
"use strict";
|
|
2
|
+
var __importDefault = (this && this.__importDefault) || function (mod) {
|
|
3
|
+
return (mod && mod.__esModule) ? mod : { "default": mod };
|
|
4
|
+
};
|
|
5
|
+
Object.defineProperty(exports, "__esModule", { value: true });
|
|
6
|
+
exports.getModelDir = getModelDir;
|
|
7
|
+
exports.isModelReady = isModelReady;
|
|
8
|
+
exports.getModelInfo = getModelInfo;
|
|
9
|
+
exports.downloadModel = downloadModel;
|
|
10
|
+
exports.ensureModel = ensureModel;
|
|
11
|
+
const node_path_1 = require("node:path");
|
|
12
|
+
const node_os_1 = require("node:os");
|
|
13
|
+
const promises_1 = require("node:fs/promises");
|
|
14
|
+
const node_https_1 = __importDefault(require("node:https"));
|
|
15
|
+
const MODELS_DIR = (0, node_path_1.join)((0, node_os_1.homedir)(), '.bastion', 'models');
|
|
16
|
+
const REQUIRED_FILES = ['model.onnx', 'tokenizer.json', 'config.json'];
|
|
17
|
+
/**
|
|
18
|
+
* Resolve the local directory for a given model ID.
|
|
19
|
+
* Model IDs use '/' which gets replaced with '--' for filesystem safety.
|
|
20
|
+
*/
|
|
21
|
+
function getModelDir(modelId) {
|
|
22
|
+
const safeName = modelId.replace(/\//g, '--');
|
|
23
|
+
return (0, node_path_1.join)(MODELS_DIR, safeName);
|
|
24
|
+
}
|
|
25
|
+
/**
|
|
26
|
+
* Check if a model is already downloaded and complete.
|
|
27
|
+
*/
|
|
28
|
+
async function isModelReady(modelId) {
|
|
29
|
+
const dir = getModelDir(modelId);
|
|
30
|
+
for (const file of REQUIRED_FILES) {
|
|
31
|
+
try {
|
|
32
|
+
const s = await (0, promises_1.stat)((0, node_path_1.join)(dir, file));
|
|
33
|
+
if (s.size === 0)
|
|
34
|
+
return false;
|
|
35
|
+
}
|
|
36
|
+
catch {
|
|
37
|
+
return false;
|
|
38
|
+
}
|
|
39
|
+
}
|
|
40
|
+
return true;
|
|
41
|
+
}
|
|
42
|
+
/**
|
|
43
|
+
* Get info about a local model.
|
|
44
|
+
*/
|
|
45
|
+
async function getModelInfo(modelId) {
|
|
46
|
+
const dir = getModelDir(modelId);
|
|
47
|
+
const files = [];
|
|
48
|
+
for (const file of REQUIRED_FILES) {
|
|
49
|
+
try {
|
|
50
|
+
await (0, promises_1.stat)((0, node_path_1.join)(dir, file));
|
|
51
|
+
files.push(file);
|
|
52
|
+
}
|
|
53
|
+
catch {
|
|
54
|
+
// file missing
|
|
55
|
+
}
|
|
56
|
+
}
|
|
57
|
+
return {
|
|
58
|
+
modelId,
|
|
59
|
+
dir,
|
|
60
|
+
ready: files.length === REQUIRED_FILES.length,
|
|
61
|
+
files,
|
|
62
|
+
};
|
|
63
|
+
}
|
|
64
|
+
/**
|
|
65
|
+
* Download a single file from Hugging Face Hub.
|
|
66
|
+
*/
|
|
67
|
+
function getHfToken() {
|
|
68
|
+
return process.env['HF_TOKEN'] ?? process.env['HUGGING_FACE_HUB_TOKEN'];
|
|
69
|
+
}
|
|
70
|
+
function downloadFile(url, destPath, onProgress, fileName) {
|
|
71
|
+
return new Promise((resolve, reject) => {
|
|
72
|
+
const request = (targetUrl) => {
|
|
73
|
+
const headers = { 'User-Agent': 'bastion-pro/0.1.0' };
|
|
74
|
+
const hfToken = getHfToken();
|
|
75
|
+
if (hfToken && targetUrl.includes('huggingface.co')) {
|
|
76
|
+
headers['Authorization'] = `Bearer ${hfToken}`;
|
|
77
|
+
}
|
|
78
|
+
node_https_1.default.get(targetUrl, { headers }, (res) => {
|
|
79
|
+
// Follow redirects (HF Hub returns 302 to CDN)
|
|
80
|
+
if (res.statusCode && res.statusCode >= 301 && res.statusCode <= 308 && res.statusCode !== 304) {
|
|
81
|
+
let location = res.headers.location;
|
|
82
|
+
if (!location)
|
|
83
|
+
return reject(new Error(`Redirect without location for ${url}`));
|
|
84
|
+
// Handle relative redirects — resolve against current URL
|
|
85
|
+
if (location.startsWith('/')) {
|
|
86
|
+
const base = new URL(targetUrl);
|
|
87
|
+
location = `${base.protocol}//${base.host}${location}`;
|
|
88
|
+
}
|
|
89
|
+
res.resume();
|
|
90
|
+
request(location);
|
|
91
|
+
return;
|
|
92
|
+
}
|
|
93
|
+
if (res.statusCode !== 200) {
|
|
94
|
+
res.resume();
|
|
95
|
+
return reject(new Error(`HTTP ${res.statusCode} downloading ${url}`));
|
|
96
|
+
}
|
|
97
|
+
const total = parseInt(res.headers['content-length'] ?? '0', 10);
|
|
98
|
+
let received = 0;
|
|
99
|
+
const chunks = [];
|
|
100
|
+
res.on('data', (chunk) => {
|
|
101
|
+
chunks.push(chunk);
|
|
102
|
+
received += chunk.length;
|
|
103
|
+
if (onProgress) {
|
|
104
|
+
onProgress({
|
|
105
|
+
file: fileName ?? destPath,
|
|
106
|
+
received,
|
|
107
|
+
total,
|
|
108
|
+
percent: total > 0 ? Math.round((received / total) * 100) : 0,
|
|
109
|
+
});
|
|
110
|
+
}
|
|
111
|
+
});
|
|
112
|
+
res.on('end', async () => {
|
|
113
|
+
try {
|
|
114
|
+
await (0, promises_1.writeFile)(destPath, Buffer.concat(chunks));
|
|
115
|
+
resolve();
|
|
116
|
+
}
|
|
117
|
+
catch (err) {
|
|
118
|
+
reject(err);
|
|
119
|
+
}
|
|
120
|
+
});
|
|
121
|
+
res.on('error', reject);
|
|
122
|
+
}).on('error', reject);
|
|
123
|
+
};
|
|
124
|
+
request(url);
|
|
125
|
+
});
|
|
126
|
+
}
|
|
127
|
+
/**
|
|
128
|
+
* Build a Hugging Face Hub download URL for a file in a model repo.
|
|
129
|
+
* Format: https://huggingface.co/{model_id}/resolve/main/{filename}
|
|
130
|
+
*/
|
|
131
|
+
function hfFileUrl(modelId, filename) {
|
|
132
|
+
return `https://huggingface.co/${modelId}/resolve/main/${filename}`;
|
|
133
|
+
}
|
|
134
|
+
/**
|
|
135
|
+
* Download model files from Hugging Face Hub.
|
|
136
|
+
* Files: model.onnx (from onnx/ subdirectory), tokenizer.json, config.json
|
|
137
|
+
*/
|
|
138
|
+
async function downloadModel(modelId, onProgress, options) {
|
|
139
|
+
const dir = getModelDir(modelId);
|
|
140
|
+
await (0, promises_1.mkdir)(dir, { recursive: true });
|
|
141
|
+
const onnxSubdir = options?.onnxSubdir ?? 'onnx';
|
|
142
|
+
// Download config.json and tokenizer.json from repo root
|
|
143
|
+
await downloadFile(hfFileUrl(modelId, 'config.json'), (0, node_path_1.join)(dir, 'config.json'), onProgress, 'config.json');
|
|
144
|
+
await downloadFile(hfFileUrl(modelId, 'tokenizer.json'), (0, node_path_1.join)(dir, 'tokenizer.json'), onProgress, 'tokenizer.json');
|
|
145
|
+
// Download model.onnx from onnx subdirectory
|
|
146
|
+
await downloadFile(hfFileUrl(modelId, `${onnxSubdir}/model.onnx`), (0, node_path_1.join)(dir, 'model.onnx'), onProgress, 'model.onnx');
|
|
147
|
+
return dir;
|
|
148
|
+
}
|
|
149
|
+
/**
|
|
150
|
+
* Ensure a model is available locally. Downloads if missing.
|
|
151
|
+
*/
|
|
152
|
+
async function ensureModel(modelId, onProgress, options) {
|
|
153
|
+
if (await isModelReady(modelId)) {
|
|
154
|
+
const dir = getModelDir(modelId);
|
|
155
|
+
options?.logger?.info('Model already downloaded, skipping', { modelId, dir });
|
|
156
|
+
return dir;
|
|
157
|
+
}
|
|
158
|
+
options?.logger?.info('Model not found locally, downloading...', { modelId });
|
|
159
|
+
return downloadModel(modelId, onProgress, { onnxSubdir: options?.onnxSubdir });
|
|
160
|
+
}
|
|
161
|
+
//# sourceMappingURL=model-manager.js.map
|
|
@@ -0,0 +1 @@
|
|
|
1
|
+
{"version":3,"file":"model-manager.js","sourceRoot":"","sources":["../../src/classifier/model-manager.ts"],"names":[],"mappings":";;;;;AA6BA,kCAGC;AAKD,oCAWC;AAKD,oCAiBC;AAwFD,sCAkCC;AAKD,kCAYC;AAjND,yCAAiC;AACjC,qCAAkC;AAClC,+CAA0D;AAC1D,4DAA+B;AAE/B,MAAM,UAAU,GAAG,IAAA,gBAAI,EAAC,IAAA,iBAAO,GAAE,EAAE,UAAU,EAAE,QAAQ,CAAC,CAAC;AAEzD,MAAM,cAAc,GAAG,CAAC,YAAY,EAAE,gBAAgB,EAAE,aAAa,CAAC,CAAC;AAkBvE;;;GAGG;AACH,SAAgB,WAAW,CAAC,OAAe;IACzC,MAAM,QAAQ,GAAG,OAAO,CAAC,OAAO,CAAC,KAAK,EAAE,IAAI,CAAC,CAAC;IAC9C,OAAO,IAAA,gBAAI,EAAC,UAAU,EAAE,QAAQ,CAAC,CAAC;AACpC,CAAC;AAED;;GAEG;AACI,KAAK,UAAU,YAAY,CAAC,OAAe;IAChD,MAAM,GAAG,GAAG,WAAW,CAAC,OAAO,CAAC,CAAC;IACjC,KAAK,MAAM,IAAI,IAAI,cAAc,EAAE,CAAC;QAClC,IAAI,CAAC;YACH,MAAM,CAAC,GAAG,MAAM,IAAA,eAAI,EAAC,IAAA,gBAAI,EAAC,GAAG,EAAE,IAAI,CAAC,CAAC,CAAC;YACtC,IAAI,CAAC,CAAC,IAAI,KAAK,CAAC;gBAAE,OAAO,KAAK,CAAC;QACjC,CAAC;QAAC,MAAM,CAAC;YACP,OAAO,KAAK,CAAC;QACf,CAAC;IACH,CAAC;IACD,OAAO,IAAI,CAAC;AACd,CAAC;AAED;;GAEG;AACI,KAAK,UAAU,YAAY,CAAC,OAAe;IAChD,MAAM,GAAG,GAAG,WAAW,CAAC,OAAO,CAAC,CAAC;IACjC,MAAM,KAAK,GAAa,EAAE,CAAC;IAC3B,KAAK,MAAM,IAAI,IAAI,cAAc,EAAE,CAAC;QAClC,IAAI,CAAC;YACH,MAAM,IAAA,eAAI,EAAC,IAAA,gBAAI,EAAC,GAAG,EAAE,IAAI,CAAC,CAAC,CAAC;YAC5B,KAAK,CAAC,IAAI,CAAC,IAAI,CAAC,CAAC;QACnB,CAAC;QAAC,MAAM,CAAC;YACP,eAAe;QACjB,CAAC;IACH,CAAC;IACD,OAAO;QACL,OAAO;QACP,GAAG;QACH,KAAK,EAAE,KAAK,CAAC,MAAM,KAAK,cAAc,CAAC,MAAM;QAC7C,KAAK;KACN,CAAC;AACJ,CAAC;AAED;;GAEG;AACH,SAAS,UAAU;IACjB,OAAO,OAAO,CAAC,GAAG,CAAC,UAAU,CAAC,IAAI,OAAO,CAAC,GAAG,CAAC,wBAAwB,CAAC,CAAC;AAC1E,CAAC;AAED,SAAS,YAAY,CACnB,GAAW,EACX,QAAgB,EAChB,UAA6B,EAC7B,QAAiB;IAEjB,OAAO,IAAI,OAAO,CAAC,CAAC,OAAO,EAAE,MAAM,EAAE,EAAE;QACrC,MAAM,OAAO,GAAG,CAAC,SAAiB,EAAE,EAAE;YACpC,MAAM,OAAO,GAA2B,EAAE,YAAY,EAAE,mBAAmB,EAAE,CAAC;YAC9E,MAAM,OAAO,GAAG,UAAU,EAAE,CAAC;YAC7B,IAAI,OAAO,IAAI,SAAS,CAAC,QAAQ,CAAC,gBAAgB,CAAC,EAAE,CAAC;gBACpD,OAAO,CAAC,eAAe,CAAC,GAAG,UAAU,OAAO,EAAE,CAAC;YACjD,CAAC;YACD,oBAAK,CAAC,GAAG,CAAC,SAAS,EAAE,EAAE,OAAO,EAAE,EAAE,CAAC,GAAG,EAAE,EAAE;gBACxC,+CAA+C;gBAC/C,IAAI,GAAG,CAAC,UAAU,IAAI,GAAG,CAAC,UAAU,IAAI,GAAG,IAAI,GAAG,CAAC,UAAU,IAAI,GAAG,IAAI,GAAG,CAAC,UAAU,KAAK,GAAG,EAAE,CAAC;oBAC/F,IAAI,QAAQ,GAAG,GAAG,CAAC,OAAO,CAAC,QAAQ,CAAC;oBACpC,IAAI,CAAC,QAAQ;wBAAE,OAAO,MAAM,CAAC,IAAI,KAAK,CAAC,iCAAiC,GAAG,EAAE,CAAC,CAAC,CAAC;oBAChF,0DAA0D;oBAC1D,IAAI,QAAQ,CAAC,UAAU,CAAC,GAAG,CAAC,EAAE,CAAC;wBAC7B,MAAM,IAAI,GAAG,IAAI,GAAG,CAAC,SAAS,CAAC,CAAC;wBAChC,QAAQ,GAAG,GAAG,IAAI,CAAC,QAAQ,KAAK,IAAI,CAAC,IAAI,GAAG,QAAQ,EAAE,CAAC;oBACzD,CAAC;oBACD,GAAG,CAAC,MAAM,EAAE,CAAC;oBACb,OAAO,CAAC,QAAQ,CAAC,CAAC;oBAClB,OAAO;gBACT,CAAC;gBAED,IAAI,GAAG,CAAC,UAAU,KAAK,GAAG,EAAE,CAAC;oBAC3B,GAAG,CAAC,MAAM,EAAE,CAAC;oBACb,OAAO,MAAM,CAAC,IAAI,KAAK,CAAC,QAAQ,GAAG,CAAC,UAAU,gBAAgB,GAAG,EAAE,CAAC,CAAC,CAAC;gBACxE,CAAC;gBAED,MAAM,KAAK,GAAG,QAAQ,CAAC,GAAG,CAAC,OAAO,CAAC,gBAAgB,CAAC,IAAI,GAAG,EAAE,EAAE,CAAC,CAAC;gBACjE,IAAI,QAAQ,GAAG,CAAC,CAAC;gBACjB,MAAM,MAAM,GAAa,EAAE,CAAC;gBAE5B,GAAG,CAAC,EAAE,CAAC,MAAM,EAAE,CAAC,KAAa,EAAE,EAAE;oBAC/B,MAAM,CAAC,IAAI,CAAC,KAAK,CAAC,CAAC;oBACnB,QAAQ,IAAI,KAAK,CAAC,MAAM,CAAC;oBACzB,IAAI,UAAU,EAAE,CAAC;wBACf,UAAU,CAAC;4BACT,IAAI,EAAE,QAAQ,IAAI,QAAQ;4BAC1B,QAAQ;4BACR,KAAK;4BACL,OAAO,EAAE,KAAK,GAAG,CAAC,CAAC,CAAC,CAAC,IAAI,CAAC,KAAK,CAAC,CAAC,QAAQ,GAAG,KAAK,CAAC,GAAG,GAAG,CAAC,CAAC,CAAC,CAAC,CAAC;yBAC9D,CAAC,CAAC;oBACL,CAAC;gBACH,CAAC,CAAC,CAAC;gBAEH,GAAG,CAAC,EAAE,CAAC,KAAK,EAAE,KAAK,IAAI,EAAE;oBACvB,IAAI,CAAC;wBACH,MAAM,IAAA,oBAAS,EAAC,QAAQ,EAAE,MAAM,CAAC,MAAM,CAAC,MAAM,CAAC,CAAC,CAAC;wBACjD,OAAO,EAAE,CAAC;oBACZ,CAAC;oBAAC,OAAO,GAAG,EAAE,CAAC;wBACb,MAAM,CAAC,GAAG,CAAC,CAAC;oBACd,CAAC;gBACH,CAAC,CAAC,CAAC;gBAEH,GAAG,CAAC,EAAE,CAAC,OAAO,EAAE,MAAM,CAAC,CAAC;YAC1B,CAAC,CAAC,CAAC,EAAE,CAAC,OAAO,EAAE,MAAM,CAAC,CAAC;QACzB,CAAC,CAAC;QAEF,OAAO,CAAC,GAAG,CAAC,CAAC;IACf,CAAC,CAAC,CAAC;AACL,CAAC;AAED;;;GAGG;AACH,SAAS,SAAS,CAAC,OAAe,EAAE,QAAgB;IAClD,OAAO,0BAA0B,OAAO,iBAAiB,QAAQ,EAAE,CAAC;AACtE,CAAC;AAED;;;GAGG;AACI,KAAK,UAAU,aAAa,CACjC,OAAe,EACf,UAA6B,EAC7B,OAAiC;IAEjC,MAAM,GAAG,GAAG,WAAW,CAAC,OAAO,CAAC,CAAC;IACjC,MAAM,IAAA,gBAAK,EAAC,GAAG,EAAE,EAAE,SAAS,EAAE,IAAI,EAAE,CAAC,CAAC;IAEtC,MAAM,UAAU,GAAG,OAAO,EAAE,UAAU,IAAI,MAAM,CAAC;IAEjD,yDAAyD;IACzD,MAAM,YAAY,CAChB,SAAS,CAAC,OAAO,EAAE,aAAa,CAAC,EACjC,IAAA,gBAAI,EAAC,GAAG,EAAE,aAAa,CAAC,EACxB,UAAU,EACV,aAAa,CACd,CAAC;IAEF,MAAM,YAAY,CAChB,SAAS,CAAC,OAAO,EAAE,gBAAgB,CAAC,EACpC,IAAA,gBAAI,EAAC,GAAG,EAAE,gBAAgB,CAAC,EAC3B,UAAU,EACV,gBAAgB,CACjB,CAAC;IAEF,6CAA6C;IAC7C,MAAM,YAAY,CAChB,SAAS,CAAC,OAAO,EAAE,GAAG,UAAU,aAAa,CAAC,EAC9C,IAAA,gBAAI,EAAC,GAAG,EAAE,YAAY,CAAC,EACvB,UAAU,EACV,YAAY,CACb,CAAC;IAEF,OAAO,GAAG,CAAC;AACb,CAAC;AAED;;GAEG;AACI,KAAK,UAAU,WAAW,CAC/B,OAAe,EACf,UAA6B,EAC7B,OAA2G;IAE3G,IAAI,MAAM,YAAY,CAAC,OAAO,CAAC,EAAE,CAAC;QAChC,MAAM,GAAG,GAAG,WAAW,CAAC,OAAO,CAAC,CAAC;QACjC,OAAO,EAAE,MAAM,EAAE,IAAI,CAAC,oCAAoC,EAAE,EAAE,OAAO,EAAE,GAAG,EAAE,CAAC,CAAC;QAC9E,OAAO,GAAG,CAAC;IACb,CAAC;IACD,OAAO,EAAE,MAAM,EAAE,IAAI,CAAC,yCAAyC,EAAE,EAAE,OAAO,EAAE,CAAC,CAAC;IAC9E,OAAO,aAAa,CAAC,OAAO,EAAE,UAAU,EAAE,EAAE,UAAU,EAAE,OAAO,EAAE,UAAU,EAAE,CAAC,CAAC;AACjF,CAAC"}
|
|
@@ -0,0 +1,22 @@
|
|
|
1
|
+
import type { ClassifierProvider, ClassificationResult } from '../plugin-api/types.js';
|
|
2
|
+
export interface OnnxProviderConfig {
|
|
3
|
+
modelDir: string;
|
|
4
|
+
maxLength?: number;
|
|
5
|
+
}
|
|
6
|
+
export declare class OnnxClassifierProvider implements ClassifierProvider {
|
|
7
|
+
readonly name = "onnx";
|
|
8
|
+
private config;
|
|
9
|
+
private maxLength;
|
|
10
|
+
private session;
|
|
11
|
+
private tokenizer;
|
|
12
|
+
private labels;
|
|
13
|
+
private _modelName;
|
|
14
|
+
constructor(config: OnnxProviderConfig);
|
|
15
|
+
get modelName(): string;
|
|
16
|
+
get ready(): boolean;
|
|
17
|
+
initialize(): Promise<void>;
|
|
18
|
+
classify(text: string): Promise<ClassificationResult>;
|
|
19
|
+
classifyBatch(texts: string[]): Promise<ClassificationResult[]>;
|
|
20
|
+
destroy(): Promise<void>;
|
|
21
|
+
}
|
|
22
|
+
//# sourceMappingURL=onnx-provider.d.ts.map
|
|
@@ -0,0 +1 @@
|
|
|
1
|
+
{"version":3,"file":"onnx-provider.d.ts","sourceRoot":"","sources":["../../src/classifier/onnx-provider.ts"],"names":[],"mappings":"AAAA,OAAO,KAAK,EAAE,kBAAkB,EAAE,oBAAoB,EAAE,MAAM,wBAAwB,CAAC;AAEvF,MAAM,WAAW,kBAAkB;IACjC,QAAQ,EAAE,MAAM,CAAC;IACjB,SAAS,CAAC,EAAE,MAAM,CAAC;CACpB;AAaD,qBAAa,sBAAuB,YAAW,kBAAkB;IAC/D,QAAQ,CAAC,IAAI,UAAU;IACvB,OAAO,CAAC,MAAM,CAAqB;IACnC,OAAO,CAAC,SAAS,CAAS;IAC1B,OAAO,CAAC,OAAO,CAA4D;IAC3E,OAAO,CAAC,SAAS,CAA4D;IAC7E,OAAO,CAAC,MAAM,CAAgB;IAC9B,OAAO,CAAC,UAAU,CAAM;gBAEZ,MAAM,EAAE,kBAAkB;IAKtC,IAAI,SAAS,IAAI,MAAM,CAEtB;IAED,IAAI,KAAK,IAAI,OAAO,CAEnB;IAEK,UAAU,IAAI,OAAO,CAAC,IAAI,CAAC;IA0C3B,QAAQ,CAAC,IAAI,EAAE,MAAM,GAAG,OAAO,CAAC,oBAAoB,CAAC;IA6DrD,aAAa,CAAC,KAAK,EAAE,MAAM,EAAE,GAAG,OAAO,CAAC,oBAAoB,EAAE,CAAC;IAK/D,OAAO,IAAI,OAAO,CAAC,IAAI,CAAC;CAO/B"}
|
|
@@ -0,0 +1,131 @@
|
|
|
1
|
+
"use strict";
|
|
2
|
+
Object.defineProperty(exports, "__esModule", { value: true });
|
|
3
|
+
exports.OnnxClassifierProvider = void 0;
|
|
4
|
+
// Lazy-loaded native dependencies (only imported when initialize() is called)
|
|
5
|
+
let ort = null;
|
|
6
|
+
let TokenizerModule = null;
|
|
7
|
+
function softmax(logits) {
|
|
8
|
+
const max = logits.reduce((a, b) => Math.max(a, b), -Infinity);
|
|
9
|
+
const exps = logits.map(v => Math.exp(v - max));
|
|
10
|
+
const sum = exps.reduce((a, b) => a + b, 0);
|
|
11
|
+
return exps.map(v => v / sum);
|
|
12
|
+
}
|
|
13
|
+
class OnnxClassifierProvider {
|
|
14
|
+
name = 'onnx';
|
|
15
|
+
config;
|
|
16
|
+
maxLength;
|
|
17
|
+
session = null;
|
|
18
|
+
tokenizer = null;
|
|
19
|
+
labels = [];
|
|
20
|
+
_modelName = '';
|
|
21
|
+
constructor(config) {
|
|
22
|
+
this.config = config;
|
|
23
|
+
this.maxLength = config.maxLength ?? 512;
|
|
24
|
+
}
|
|
25
|
+
get modelName() {
|
|
26
|
+
return this._modelName;
|
|
27
|
+
}
|
|
28
|
+
get ready() {
|
|
29
|
+
return this.session !== null && this.tokenizer !== null;
|
|
30
|
+
}
|
|
31
|
+
async initialize() {
|
|
32
|
+
const { join } = await import('node:path');
|
|
33
|
+
const { readFile } = await import('node:fs/promises');
|
|
34
|
+
// Lazy-load native modules
|
|
35
|
+
if (!ort)
|
|
36
|
+
ort = await import('onnxruntime-node');
|
|
37
|
+
if (!TokenizerModule)
|
|
38
|
+
TokenizerModule = await import('@huggingface/tokenizers');
|
|
39
|
+
const modelPath = join(this.config.modelDir, 'model.onnx');
|
|
40
|
+
const tokenizerPath = join(this.config.modelDir, 'tokenizer.json');
|
|
41
|
+
const configPath = join(this.config.modelDir, 'config.json');
|
|
42
|
+
// Load model config for label mapping
|
|
43
|
+
let modelConfig = {};
|
|
44
|
+
try {
|
|
45
|
+
const configRaw = await readFile(configPath, 'utf-8');
|
|
46
|
+
modelConfig = JSON.parse(configRaw);
|
|
47
|
+
if (modelConfig.id2label) {
|
|
48
|
+
const id2label = modelConfig.id2label;
|
|
49
|
+
const maxId = Math.max(...Object.keys(id2label).map(Number));
|
|
50
|
+
this.labels = new Array(maxId + 1);
|
|
51
|
+
for (const [id, label] of Object.entries(id2label)) {
|
|
52
|
+
this.labels[Number(id)] = label;
|
|
53
|
+
}
|
|
54
|
+
}
|
|
55
|
+
this._modelName = (modelConfig._name_or_path ?? modelConfig.model_type ?? 'unknown');
|
|
56
|
+
}
|
|
57
|
+
catch {
|
|
58
|
+
this.labels = ['SAFE', 'INJECTION'];
|
|
59
|
+
this._modelName = 'unknown';
|
|
60
|
+
}
|
|
61
|
+
// Load tokenizer — @huggingface/tokenizers v0.1.x uses new Tokenizer(tokenizerJson, configJson)
|
|
62
|
+
const tokenizerRaw = await readFile(tokenizerPath, 'utf-8');
|
|
63
|
+
const tokenizerJson = JSON.parse(tokenizerRaw);
|
|
64
|
+
this.tokenizer = new TokenizerModule.Tokenizer(tokenizerJson, modelConfig);
|
|
65
|
+
// Create ONNX inference session
|
|
66
|
+
this.session = await ort.InferenceSession.create(modelPath, {
|
|
67
|
+
executionProviders: ['cpu'],
|
|
68
|
+
});
|
|
69
|
+
}
|
|
70
|
+
async classify(text) {
|
|
71
|
+
if (!this.session || !this.tokenizer) {
|
|
72
|
+
throw new Error('OnnxClassifierProvider not initialized — call initialize() first');
|
|
73
|
+
}
|
|
74
|
+
const start = performance.now();
|
|
75
|
+
// Tokenize — @huggingface/tokenizers v0.1.x returns plain object { ids, attention_mask }
|
|
76
|
+
const encoded = this.tokenizer.encode(text);
|
|
77
|
+
const inputIds = encoded.ids;
|
|
78
|
+
const attentionMask = encoded.attention_mask;
|
|
79
|
+
// Truncate to maxLength
|
|
80
|
+
const length = Math.min(inputIds.length, this.maxLength);
|
|
81
|
+
const ids = BigInt64Array.from(inputIds.slice(0, length).map((v) => BigInt(v)));
|
|
82
|
+
const mask = BigInt64Array.from(attentionMask.slice(0, length).map((v) => BigInt(v)));
|
|
83
|
+
// Create tensors
|
|
84
|
+
const inputIdsTensor = new ort.Tensor('int64', ids, [1, length]);
|
|
85
|
+
const attentionMaskTensor = new ort.Tensor('int64', mask, [1, length]);
|
|
86
|
+
// Run inference
|
|
87
|
+
const feeds = {
|
|
88
|
+
input_ids: inputIdsTensor,
|
|
89
|
+
attention_mask: attentionMaskTensor,
|
|
90
|
+
};
|
|
91
|
+
const output = await this.session.run(feeds);
|
|
92
|
+
// Get logits from first output
|
|
93
|
+
const outputName = this.session.outputNames[0];
|
|
94
|
+
const logits = output[outputName].data;
|
|
95
|
+
// Apply softmax
|
|
96
|
+
const probs = softmax(logits);
|
|
97
|
+
// Build label scores
|
|
98
|
+
const labels = [];
|
|
99
|
+
let bestIdx = 0;
|
|
100
|
+
let bestScore = 0;
|
|
101
|
+
for (let i = 0; i < probs.length; i++) {
|
|
102
|
+
const label = this.labels[i] ?? `CLASS_${i}`;
|
|
103
|
+
const score = probs[i];
|
|
104
|
+
labels.push({ label, score });
|
|
105
|
+
if (score > bestScore) {
|
|
106
|
+
bestScore = score;
|
|
107
|
+
bestIdx = i;
|
|
108
|
+
}
|
|
109
|
+
}
|
|
110
|
+
const latencyMs = Math.round(performance.now() - start);
|
|
111
|
+
return {
|
|
112
|
+
label: this.labels[bestIdx] ?? `CLASS_${bestIdx}`,
|
|
113
|
+
score: bestScore,
|
|
114
|
+
labels,
|
|
115
|
+
latencyMs,
|
|
116
|
+
};
|
|
117
|
+
}
|
|
118
|
+
async classifyBatch(texts) {
|
|
119
|
+
// Simple sequential batch — ONNX batch inference requires padding and is more complex
|
|
120
|
+
return Promise.all(texts.map(text => this.classify(text)));
|
|
121
|
+
}
|
|
122
|
+
async destroy() {
|
|
123
|
+
if (this.session) {
|
|
124
|
+
await this.session.release();
|
|
125
|
+
this.session = null;
|
|
126
|
+
}
|
|
127
|
+
this.tokenizer = null;
|
|
128
|
+
}
|
|
129
|
+
}
|
|
130
|
+
exports.OnnxClassifierProvider = OnnxClassifierProvider;
|
|
131
|
+
//# sourceMappingURL=onnx-provider.js.map
|
|
@@ -0,0 +1 @@
|
|
|
1
|
+
{"version":3,"file":"onnx-provider.js","sourceRoot":"","sources":["../../src/classifier/onnx-provider.ts"],"names":[],"mappings":";;;AAOA,8EAA8E;AAC9E,IAAI,GAAG,GAA6C,IAAI,CAAC;AACzD,IAAI,eAAe,GAAoD,IAAI,CAAC;AAE5E,SAAS,OAAO,CAAC,MAAoB;IACnC,MAAM,GAAG,GAAG,MAAM,CAAC,MAAM,CAAC,CAAC,CAAC,EAAE,CAAC,EAAE,EAAE,CAAC,IAAI,CAAC,GAAG,CAAC,CAAC,EAAE,CAAC,CAAC,EAAE,CAAC,QAAQ,CAAC,CAAC;IAC/D,MAAM,IAAI,GAAG,MAAM,CAAC,GAAG,CAAC,CAAC,CAAC,EAAE,CAAC,IAAI,CAAC,GAAG,CAAC,CAAC,GAAG,GAAG,CAAC,CAAC,CAAC;IAChD,MAAM,GAAG,GAAG,IAAI,CAAC,MAAM,CAAC,CAAC,CAAC,EAAE,CAAC,EAAE,EAAE,CAAC,CAAC,GAAG,CAAC,EAAE,CAAC,CAAC,CAAC;IAC5C,OAAO,IAAI,CAAC,GAAG,CAAC,CAAC,CAAC,EAAE,CAAC,CAAC,GAAG,GAAG,CAA4B,CAAC;AAC3D,CAAC;AAED,MAAa,sBAAsB;IACxB,IAAI,GAAG,MAAM,CAAC;IACf,MAAM,CAAqB;IAC3B,SAAS,CAAS;IAClB,OAAO,GAAuD,IAAI,CAAC;IACnE,SAAS,GAAuD,IAAI,CAAC;IACrE,MAAM,GAAa,EAAE,CAAC;IACtB,UAAU,GAAG,EAAE,CAAC;IAExB,YAAY,MAA0B;QACpC,IAAI,CAAC,MAAM,GAAG,MAAM,CAAC;QACrB,IAAI,CAAC,SAAS,GAAG,MAAM,CAAC,SAAS,IAAI,GAAG,CAAC;IAC3C,CAAC;IAED,IAAI,SAAS;QACX,OAAO,IAAI,CAAC,UAAU,CAAC;IACzB,CAAC;IAED,IAAI,KAAK;QACP,OAAO,IAAI,CAAC,OAAO,KAAK,IAAI,IAAI,IAAI,CAAC,SAAS,KAAK,IAAI,CAAC;IAC1D,CAAC;IAED,KAAK,CAAC,UAAU;QACd,MAAM,EAAE,IAAI,EAAE,GAAG,MAAM,MAAM,CAAC,WAAW,CAAC,CAAC;QAC3C,MAAM,EAAE,QAAQ,EAAE,GAAG,MAAM,MAAM,CAAC,kBAAkB,CAAC,CAAC;QAEtD,2BAA2B;QAC3B,IAAI,CAAC,GAAG;YAAE,GAAG,GAAG,MAAM,MAAM,CAAC,kBAAkB,CAAC,CAAC;QACjD,IAAI,CAAC,eAAe;YAAE,eAAe,GAAG,MAAM,MAAM,CAAC,yBAAyB,CAAC,CAAC;QAEhF,MAAM,SAAS,GAAG,IAAI,CAAC,IAAI,CAAC,MAAM,CAAC,QAAQ,EAAE,YAAY,CAAC,CAAC;QAC3D,MAAM,aAAa,GAAG,IAAI,CAAC,IAAI,CAAC,MAAM,CAAC,QAAQ,EAAE,gBAAgB,CAAC,CAAC;QACnE,MAAM,UAAU,GAAG,IAAI,CAAC,IAAI,CAAC,MAAM,CAAC,QAAQ,EAAE,aAAa,CAAC,CAAC;QAE7D,sCAAsC;QACtC,IAAI,WAAW,GAA4B,EAAE,CAAC;QAC9C,IAAI,CAAC;YACH,MAAM,SAAS,GAAG,MAAM,QAAQ,CAAC,UAAU,EAAE,OAAO,CAAC,CAAC;YACtD,WAAW,GAAG,IAAI,CAAC,KAAK,CAAC,SAAS,CAAC,CAAC;YACpC,IAAI,WAAW,CAAC,QAAQ,EAAE,CAAC;gBACzB,MAAM,QAAQ,GAAG,WAAW,CAAC,QAAkC,CAAC;gBAChE,MAAM,KAAK,GAAG,IAAI,CAAC,GAAG,CAAC,GAAG,MAAM,CAAC,IAAI,CAAC,QAAQ,CAAC,CAAC,GAAG,CAAC,MAAM,CAAC,CAAC,CAAC;gBAC7D,IAAI,CAAC,MAAM,GAAG,IAAI,KAAK,CAAC,KAAK,GAAG,CAAC,CAAC,CAAC;gBACnC,KAAK,MAAM,CAAC,EAAE,EAAE,KAAK,CAAC,IAAI,MAAM,CAAC,OAAO,CAAC,QAAQ,CAAC,EAAE,CAAC;oBACnD,IAAI,CAAC,MAAM,CAAC,MAAM,CAAC,EAAE,CAAC,CAAC,GAAG,KAAK,CAAC;gBAClC,CAAC;YACH,CAAC;YACD,IAAI,CAAC,UAAU,GAAG,CAAC,WAAW,CAAC,aAAa,IAAI,WAAW,CAAC,UAAU,IAAI,SAAS,CAAW,CAAC;QACjG,CAAC;QAAC,MAAM,CAAC;YACP,IAAI,CAAC,MAAM,GAAG,CAAC,MAAM,EAAE,WAAW,CAAC,CAAC;YACpC,IAAI,CAAC,UAAU,GAAG,SAAS,CAAC;QAC9B,CAAC;QAED,gGAAgG;QAChG,MAAM,YAAY,GAAG,MAAM,QAAQ,CAAC,aAAa,EAAE,OAAO,CAAC,CAAC;QAC5D,MAAM,aAAa,GAAG,IAAI,CAAC,KAAK,CAAC,YAAY,CAAC,CAAC;QAC/C,IAAI,CAAC,SAAS,GAAG,IAAI,eAAe,CAAC,SAAS,CAAC,aAAa,EAAE,WAAW,CAAC,CAAC;QAE3E,gCAAgC;QAChC,IAAI,CAAC,OAAO,GAAG,MAAM,GAAG,CAAC,gBAAgB,CAAC,MAAM,CAAC,SAAS,EAAE;YAC1D,kBAAkB,EAAE,CAAC,KAAK,CAAC;SAC5B,CAAC,CAAC;IACL,CAAC;IAED,KAAK,CAAC,QAAQ,CAAC,IAAY;QACzB,IAAI,CAAC,IAAI,CAAC,OAAO,IAAI,CAAC,IAAI,CAAC,SAAS,EAAE,CAAC;YACrC,MAAM,IAAI,KAAK,CAAC,kEAAkE,CAAC,CAAC;QACtF,CAAC;QAED,MAAM,KAAK,GAAG,WAAW,CAAC,GAAG,EAAE,CAAC;QAEhC,yFAAyF;QACzF,MAAM,OAAO,GAAG,IAAI,CAAC,SAAS,CAAC,MAAM,CAAC,IAAI,CAAgD,CAAC;QAC3F,MAAM,QAAQ,GAAG,OAAO,CAAC,GAAG,CAAC;QAC7B,MAAM,aAAa,GAAG,OAAO,CAAC,cAAc,CAAC;QAE7C,wBAAwB;QACxB,MAAM,MAAM,GAAG,IAAI,CAAC,GAAG,CAAC,QAAQ,CAAC,MAAM,EAAE,IAAI,CAAC,SAAS,CAAC,CAAC;QACzD,MAAM,GAAG,GAAG,aAAa,CAAC,IAAI,CAAC,QAAQ,CAAC,KAAK,CAAC,CAAC,EAAE,MAAM,CAAC,CAAC,GAAG,CAAC,CAAC,CAAS,EAAE,EAAE,CAAC,MAAM,CAAC,CAAC,CAAC,CAAC,CAAC,CAAC;QACxF,MAAM,IAAI,GAAG,aAAa,CAAC,IAAI,CAAC,aAAa,CAAC,KAAK,CAAC,CAAC,EAAE,MAAM,CAAC,CAAC,GAAG,CAAC,CAAC,CAAS,EAAE,EAAE,CAAC,MAAM,CAAC,CAAC,CAAC,CAAC,CAAC,CAAC;QAE9F,iBAAiB;QACjB,MAAM,cAAc,GAAG,IAAI,GAAI,CAAC,MAAM,CAAC,OAAO,EAAE,GAAG,EAAE,CAAC,CAAC,EAAE,MAAM,CAAC,CAAC,CAAC;QAClE,MAAM,mBAAmB,GAAG,IAAI,GAAI,CAAC,MAAM,CAAC,OAAO,EAAE,IAAI,EAAE,CAAC,CAAC,EAAE,MAAM,CAAC,CAAC,CAAC;QAExE,gBAAgB;QAChB,MAAM,KAAK,GAAsD;YAC/D,SAAS,EAAE,cAAc;YACzB,cAAc,EAAE,mBAAmB;SACpC,CAAC;QAEF,MAAM,MAAM,GAAG,MAAM,IAAI,CAAC,OAAO,CAAC,GAAG,CAAC,KAAK,CAAC,CAAC;QAE7C,+BAA+B;QAC/B,MAAM,UAAU,GAAG,IAAI,CAAC,OAAO,CAAC,WAAW,CAAC,CAAC,CAAC,CAAC;QAC/C,MAAM,MAAM,GAAG,MAAM,CAAC,UAAU,CAAC,CAAC,IAAoB,CAAC;QAEvD,gBAAgB;QAChB,MAAM,KAAK,GAAG,OAAO,CAAC,MAAM,CAAC,CAAC;QAE9B,qBAAqB;QACrB,MAAM,MAAM,GAA4C,EAAE,CAAC;QAC3D,IAAI,OAAO,GAAG,CAAC,CAAC;QAChB,IAAI,SAAS,GAAG,CAAC,CAAC;QAElB,KAAK,IAAI,CAAC,GAAG,CAAC,EAAE,CAAC,GAAG,KAAK,CAAC,MAAM,EAAE,CAAC,EAAE,EAAE,CAAC;YACtC,MAAM,KAAK,GAAG,IAAI,CAAC,MAAM,CAAC,CAAC,CAAC,IAAI,SAAS,CAAC,EAAE,CAAC;YAC7C,MAAM,KAAK,GAAG,KAAK,CAAC,CAAC,CAAC,CAAC;YACvB,MAAM,CAAC,IAAI,CAAC,EAAE,KAAK,EAAE,KAAK,EAAE,CAAC,CAAC;YAC9B,IAAI,KAAK,GAAG,SAAS,EAAE,CAAC;gBACtB,SAAS,GAAG,KAAK,CAAC;gBAClB,OAAO,GAAG,CAAC,CAAC;YACd,CAAC;QACH,CAAC;QAED,MAAM,SAAS,GAAG,IAAI,CAAC,KAAK,CAAC,WAAW,CAAC,GAAG,EAAE,GAAG,KAAK,CAAC,CAAC;QAExD,OAAO;YACL,KAAK,EAAE,IAAI,CAAC,MAAM,CAAC,OAAO,CAAC,IAAI,SAAS,OAAO,EAAE;YACjD,KAAK,EAAE,SAAS;YAChB,MAAM;YACN,SAAS;SACV,CAAC;IACJ,CAAC;IAED,KAAK,CAAC,aAAa,CAAC,KAAe;QACjC,sFAAsF;QACtF,OAAO,OAAO,CAAC,GAAG,CAAC,KAAK,CAAC,GAAG,CAAC,IAAI,CAAC,EAAE,CAAC,IAAI,CAAC,QAAQ,CAAC,IAAI,CAAC,CAAC,CAAC,CAAC;IAC7D,CAAC;IAED,KAAK,CAAC,OAAO;QACX,IAAI,IAAI,CAAC,OAAO,EAAE,CAAC;YACjB,MAAM,IAAI,CAAC,OAAO,CAAC,OAAO,EAAE,CAAC;YAC7B,IAAI,CAAC,OAAO,GAAG,IAAI,CAAC;QACtB,CAAC;QACD,IAAI,CAAC,SAAS,GAAG,IAAI,CAAC;IACxB,CAAC;CACF;AAzID,wDAyIC"}
|
|
@@ -0,0 +1 @@
|
|
|
1
|
+
{"version":3,"file":"plugins.d.ts","sourceRoot":"","sources":["../../../src/cli/commands/plugins.ts"],"names":[],"mappings":"AAAA,OAAO,EAAE,KAAK,OAAO,EAAE,MAAM,WAAW,CAAC;AA0DzC,wBAAgB,sBAAsB,CAAC,OAAO,EAAE,OAAO,GAAG,IAAI,CA0J7D"}
|