From 076b3aaa7bfb1310e5bc9be359e72422d2026375 Mon Sep 17 00:00:00 2001 From: mesalogo Date: Sun, 23 Aug 2026 11:26:44 +0800 Subject: [PATCH] feat: add managed Granite embedding models MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit GoodBuddy previously required an external OpenAI-compatible endpoint for vector retrieval. It can now install, verify, import, remove, and run the built-in Granite multilingual embedding model in an isolated local inference process, while retaining custom embedding connections and migrating existing settings. The model connection UI now manages built-in and custom vector connections consistently, reports the actual retrieval path, preserves credentials across failed saves, and renders the vector tab synchronously. Related settings copy, switch styling, metadata typography, and renderer bundle guardrails are aligned with the resulting workflow. Release note: 新增可管理的内置 Granite 多语言向量模型，支持本地下载、离线导入、连接切换与诊断；同时优化向量设置及相关界面的文案、开关和状态显示。 --- AGENTS.md | 5 + BUILD.md | 12 + build/check-renderer-bundle.cjs | 19 +- .../local-text-embedding-model-design.md | 636 +++++++------- .../model-download-source-design.md | 24 +- .../knowledge-rag-enhancement-prd.md | 43 +- electron.vite.config.ts | 18 +- package-lock.json | 7 + package.json | 9 +- src/main/embedding-inference-bootstrap.ts | 33 + src/main/index.ts | 307 ++++++- src/main/ipc.test.ts | 32 +- src/main/ipc.ts | 263 +++++- .../embedding-index-coordinator.test.ts | 53 ++ .../knowledge/embedding-index-coordinator.ts | 50 +- .../embedding-inference-broker.test.ts | 243 ++++++ .../knowledge/embedding-inference-broker.ts | 455 ++++++++++ .../embedding-inference-contracts.ts | 262 ++++++ .../embedding-inference-worker.test.ts | 251 ++++++ .../knowledge/embedding-inference-worker.ts | 309 +++++++ src/main/knowledge/embedding-model-catalog.ts | 136 +++ .../embedding-model-contracts.test.ts | 68 ++ .../knowledge/embedding-model-contracts.ts | 220 +++++ .../knowledge/embedding-model-manager.test.ts | 492 +++++++++++ src/main/knowledge/embedding-model-manager.ts | 824 ++++++++++++++++++ .../knowledge/embedding-provider-key.test.ts | 53 +- src/main/knowledge/embedding-provider-key.ts | 32 + .../granite-embedding-engine.test.ts | 150 ++++ .../knowledge/granite-embedding-engine.ts | 400 +++++++++ src/main/knowledge/knowledge-service.test.ts | 50 ++ src/main/knowledge/knowledge-service.ts | 32 +- .../knowledge/openai-embedding-client.test.ts | 92 ++ src/main/knowledge/openai-embedding-client.ts | 161 +++- src/main/knowledge/types.ts | 8 + src/main/model-archive.ts | 2 +- src/main/runtime-settings-store.test.ts | 126 ++- src/main/runtime-settings-store.ts | 485 +++++++++-- src/preload/index.ts | 36 +- src/preload/preload-sandbox.test.ts | 23 + src/renderer/src/ActivityPanel.test.tsx | 19 +- src/renderer/src/ActivityPanel.tsx | 60 +- src/renderer/src/App.test.tsx | 64 +- src/renderer/src/App.tsx | 15 +- .../src/ChannelProjectSettingsFields.tsx | 4 - .../src/ChannelSettingsSection.test.tsx | 7 +- src/renderer/src/ChannelSettingsSection.tsx | 1 - .../DocumentParsingSettingsSection.test.tsx | 6 + .../src/DocumentParsingSettingsSection.tsx | 54 +- .../src/EmbeddingSettingsSection.test.tsx | 262 ++++-- src/renderer/src/EmbeddingSettingsSection.tsx | 459 ++++++++-- src/renderer/src/HeartbeatCenter.test.tsx | 18 + src/renderer/src/HeartbeatCenter.tsx | 22 - .../src/KnowledgeRetrievalWorkbench.test.tsx | 16 +- .../src/KnowledgeRetrievalWorkbench.tsx | 68 -- src/renderer/src/KnowledgeWorkspace.test.tsx | 8 + src/renderer/src/KnowledgeWorkspace.tsx | 24 +- src/renderer/src/MagicNotesWorkspace.test.tsx | 5 +- src/renderer/src/MagicNotesWorkspace.tsx | 19 +- src/renderer/src/McpSettingsSection.tsx | 41 +- .../src/PlatformFeaturesSettingsSection.tsx | 19 +- src/renderer/src/RightAssistantSidebar.tsx | 58 +- src/renderer/src/SettingsPanel.test.tsx | 258 +++++- src/renderer/src/SettingsPanel.tsx | 460 +++++++--- .../src/SpeechModelSettingsSection.test.tsx | 3 + .../src/SpeechModelSettingsSection.tsx | 12 - src/renderer/src/WorkspacePrimitives.test.tsx | 20 + .../src/i18n/locales/en-US/activity.ts | 19 +- src/renderer/src/i18n/locales/en-US/app.ts | 19 +- .../src/i18n/locales/en-US/heartbeat.ts | 12 +- .../src/i18n/locales/en-US/integrations.ts | 24 +- .../src/i18n/locales/en-US/knowledge.ts | 52 +- .../src/i18n/locales/en-US/magicNotes.ts | 9 +- .../src/i18n/locales/en-US/settings.ts | 23 +- .../i18n/locales/en-US/settingsSections.ts | 75 +- .../src/i18n/locales/en-US/workspace.ts | 23 +- .../src/i18n/locales/zh-CN/activity.ts | 19 +- src/renderer/src/i18n/locales/zh-CN/app.ts | 19 +- .../src/i18n/locales/zh-CN/heartbeat.ts | 14 +- .../src/i18n/locales/zh-CN/integrations.ts | 21 +- .../src/i18n/locales/zh-CN/knowledge.ts | 45 +- .../src/i18n/locales/zh-CN/magicNotes.ts | 9 +- .../src/i18n/locales/zh-CN/settings.ts | 23 +- .../i18n/locales/zh-CN/settingsSections.ts | 70 +- .../src/i18n/locales/zh-CN/workspace.ts | 22 +- src/renderer/src/styles.css | 6 +- src/shared/contracts.ts | 155 +++- src/shared/embedding-contracts.test.ts | 43 + src/shared/embedding-contracts.ts | 124 ++- src/shared/ipc-channels.ts | 6 + tests/electron-vite-config.test.ts | 11 + 90 files changed, 7820 insertions(+), 1446 deletions(-) create mode 100644 src/main/embedding-inference-bootstrap.ts create mode 100644 src/main/knowledge/embedding-inference-broker.test.ts create mode 100644 src/main/knowledge/embedding-inference-broker.ts create mode 100644 src/main/knowledge/embedding-inference-contracts.ts create mode 100644 src/main/knowledge/embedding-inference-worker.test.ts create mode 100644 src/main/knowledge/embedding-inference-worker.ts create mode 100644 src/main/knowledge/embedding-model-catalog.ts create mode 100644 src/main/knowledge/embedding-model-contracts.test.ts create mode 100644 src/main/knowledge/embedding-model-contracts.ts create mode 100644 src/main/knowledge/embedding-model-manager.test.ts create mode 100644 src/main/knowledge/embedding-model-manager.ts create mode 100644 src/main/knowledge/granite-embedding-engine.test.ts create mode 100644 src/main/knowledge/granite-embedding-engine.ts diff --git a/AGENTS.md b/AGENTS.md index ac51d443..1c86bacc 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -89,6 +89,11 @@ Keep Electron security boundaries intact: blocking page state, or an error with an immediate local recovery action. - Do not show the same event both inline and as an application notification. Preserve user input and actionable error context when an operation fails. +- Treat renderer bundle budgets as regression guardrails, not UX targets. Do + not lazy-load one lightweight peer tab solely to satisfy a byte ceiling. + When a genuinely heavy view needs deferred loading, preserve peer + interaction semantics, render a non-empty layout-stable local fallback, + and test its first interaction rather than only its eventual appearance. ## Commit Messages and Release Notes diff --git a/BUILD.md b/BUILD.md index 82a0ff89..cd95dea2 100644 --- a/BUILD.md +++ b/BUILD.md @@ -316,3 +316,15 @@ Windows 与 POSIX 构建输出。构建专用 module manifest 只记录项目相 绝对路径；检查成功后会删除该诊断文件，检查失败时保留以便排查。标准 Vite manifest 会保留在输出中。预算以干净生产构建为基线并保留有限余量；若业务确需 提高预算，必须先检查 manifest 闭包和产物差异，不能只为通过 CI 调大数值。 + +Bundle 大小是发现意外依赖和持续膨胀的粗粒度工程护栏，不是用户体验指标，也不能替代 +Electron 中的实际交互验证。Settings 的硬上限以 2026-08-23 验证构建 +`648.78 kB raw / 97.79 kB gzip` 为基线，增加约 15% 评审余量后取整为 +`750 kB raw / 115 kB gzip`。该上限用于拦截显著增长，不要求实现长期贴线运行。 + +调整 Renderer 拆包时遵守以下优先级： + +1. 同级页签的交互一致性、首个绘制帧内反馈和输入保留优先于包体数字。 +2. 不得仅为通过预算，把一个原本轻量的同级页签单独改成首次点击才加载。 +3. 确需动态加载重型功能时，必须提供非空、无布局跳变的局部状态，并增加首次交互测试。 +4. 修改预算必须记录干净生产构建基线、增量闭包差异和 Electron 实测结果。 diff --git a/build/check-renderer-bundle.cjs b/build/check-renderer-bundle.cjs index f2f528bc..8676b8d3 100644 --- a/build/check-renderer-bundle.cjs +++ b/build/check-renderer-bundle.cjs @@ -2,13 +2,17 @@ const { existsSync, readFileSync, rmSync } = require('node:fs') const { gzipSync } = require('node:zlib') const { resolve } = require('node:path') +// These are hard regression ceilings, not user-experience targets. In +// particular, do not lazy-load one peer tab solely to stay below a ceiling. +// The Settings ceiling is the validated 2026-08-23 production baseline +// (648.78 kB raw / 97.79 kB gzip), rounded up with about 15% review headroom. const rendererBundleBudgets = Object.freeze({ initial: Object.freeze({ raw: 3_500_000, gzip: 720_000 }), knowledge: Object.freeze({ raw: 330_000, gzip: 50_000 }), graph: Object.freeze({ raw: 3_500_000, gzip: 720_000 }), activity: Object.freeze({ raw: 55_000, gzip: 9_000 }), magicNotes: Object.freeze({ raw: 630_000, gzip: 130_000 }), - settings: Object.freeze({ raw: 650_000, gzip: 105_000 }) + settings: Object.freeze({ raw: 750_000, gzip: 115_000 }) }) function normalizePath(value) { @@ -29,10 +33,21 @@ function findEntryKey(manifest) { function findSourceKey(manifest, sourceSuffix) { const normalizedSuffix = normalizePath(sourceSuffix) - const matches = Object.entries(manifest).filter(([key, item]) => { + let matches = Object.entries(manifest).filter(([key, item]) => { const source = normalizePath(item.src || key) return source.endsWith(normalizedSuffix) }) + if (matches.length === 0) { + const sourceName = normalizedSuffix + .split('/') + .at(-1) + ?.replace(/\.[^.]+$/u, '') + matches = Object.entries(manifest).filter( + ([, item]) => + item?.isDynamicEntry && + item.name === sourceName + ) + } if (matches.length !== 1) { throw new Error( `Expected one manifest entry for ${sourceSuffix}, found ${matches.length}` diff --git a/docs/architecture/local-text-embedding-model-design.md b/docs/architecture/local-text-embedding-model-design.md index a5074ca1..e2552a6e 100644 --- a/docs/architecture/local-text-embedding-model-design.md +++ b/docs/architecture/local-text-embedding-model-design.md @@ -5,82 +5,83 @@ | 项目 | 内容 | | --- | --- | | 文档类型 | 跨功能技术与产品架构 | -| 状态 | 设计中 | -| 版本 | 0.1 | -| 日期 | 2026-08-19 | +| 状态 | 已实施 | +| 版本 | 0.4 | +| 日期 | 2026-08-22 | | 适用产品 | GoodBuddy 桌面端 | | 目标平台 | Windows、macOS、Linux，x64 与 arm64 | -| 相关基线 | [知识库检索与分块增强 PRD](../prd/knowledge/knowledge-rag-enhancement-prd.md)、[模型下载源设计](./model-download-source-design.md)、[知识检索评估](../quality/knowledge-retrieval-evaluation.md)、[跨平台助手产品设计](../design/cross-platform-assistant-product-design.md)、[统一界面设计系统](../../UI-DESIGN.md) | +| 相关基线 | [知识库检索与分块增强 PRD](../prd/knowledge/knowledge-rag-enhancement-prd.md)、[模型下载源设计](./model-download-source-design.md)、[知识检索评估](../quality/knowledge-retrieval-evaluation.md)、[统一界面设计系统](../../UI-DESIGN.md) | -本文定义 GoodBuddy 的文本向量模型产品形态，包括应用托管的轻量本地模型、用户自行安装的 -Ollama 或其他自托管服务、云端 OpenAI 兼容服务、模型管理界面、进程边界、编码契约、 -索引兼容性、失败状态和验收方法。 +本文定义 GoodBuddy 的文本向量模型产品形态，包括 GoodBuddy 提供的内置本地连接、 +用户添加的 OpenAI 兼容连接、模型管理界面、进程边界、编码契约、索引兼容性、失败状态 +和验收方法。 -本文所称“支持多种连接”是指用户可以显式选择不同向量执行路径，不代表应用可以在它们 -之间自动切换。**GoodBuddy 不设计静默替换向量 Provider、模型、精度、数据位置或编码 -方式。** +本文中的向量模型是 **GoodBuddy Main 进程所管理的知识检索基础设施**，不是 LLM，也不是 +Agent Runtime。它不会成为直连模型、OpenCode、Continue、DeepSeek Harness 或后续 Runtime +的模型来源。Runtime 只能获得知识服务返回的有界文本证据、引用和检索诊断，不能获得 +向量连接、模型文件、Tokenizer、原始向量或凭据。 + +内置连接与用户连接使用同一连接列表、选择、诊断和索引兼容语义。本文所称“支持多种连接” +不代表应用可以在它们之间自动切换。**GoodBuddy 不设计静默替换向量连接、模型、精度、 +数据位置或编码方式。** --- ## 1. 摘要与核心决策 -1. GoodBuddy 提供三种明确的向量执行路径： - - **应用托管本地模型**：由 GoodBuddy 下载、校验、安装并在受控进程中运行。 - - **Ollama / 自托管服务**：由用户自行安装和管理服务及模型，GoodBuddy 只连接用户 - 明确填写的 Endpoint 和模型。 - - **云端兼容服务**：使用用户明确配置的 OpenAI 兼容 Embeddings Endpoint。 -2. 首个应用托管默认模型选择 - `ibm-granite/granite-embedding-97m-multilingual-r2` 的固定修订和官方 INT8 ONNX - 工件。它是目录中的“推荐”模型，不在未获用户操作时自动下载或启用。 -3. 高性能设备或更高准确性需求不通过应用内自动路由解决。用户可以自行安装 Ollama， - 显式选择 Qwen3-Embedding、BGE-M3、其他 BGE 或其认可的向量模型。 -4. Ollama 模型、云端模型和应用托管模型没有回退顺序。当前路径不可用时，向量通道明确 - 标记不可用；GoodBuddy 不自动换到另一路径或另一模型。 -5. 向量通道失败时，全文、中文 CJK 和图谱通道可以按已保存检索设置继续工作，但界面必须 +1. 向量连接只服务于 GoodBuddy 的知识索引与检索链路，不进入 Agent Runtime 的模型选择、 + 配置、进程环境或凭据边界。 +2. GoodBuddy 提供两种明确的向量连接类型： + - **内置连接**：GoodBuddy 提供的只读系统连接，负责下载、校验、安装并在受控进程中 + 运行固定本地模型。 + - **OpenAI 兼容连接**：用户填写 Endpoint、模型名称和可选 API Key，可以指向 + loopback、本地网络、自托管或外部服务。UI 不为具体服务品牌建立专属类型。 +3. 内置连接与 LLM 模型连接使用相同的连接列表交互。这里复用的是设置界面的管理模式， + 不表示向量连接可分配给 LLM 或 Agent Runtime。向量页面不增加二级 + `SegmentedControl`。 +4. 首个内置连接固定使用 + `ibm-granite/granite-embedding-97m-multilingual-r2`。模型权重不随安装包交付， + 用户需要先下载或导入，再完成真实诊断和保存。 +5. 内置连接和用户连接没有回退顺序。当前连接不可用时，向量通道明确标记不可用； + GoodBuddy 不自动换到另一连接或另一模型。 +6. 向量通道失败时，全文、中文 CJK 和图谱通道可以按已保存检索设置继续工作，但界面必须 显示“未使用向量检索”、失败原因和实际使用通道。这是可见的检索通道降级，不是向量 Provider 回退。 -6. 模型权重不随安装包交付。下载使用固定修订、已知字节数和 SHA-256；内网环境支持 - ZIP 导入、导出。 -7. 应用托管推理不在 UI Renderer 或 Main 事件循环内执行。Main 通过受控推理进程调用 +7. 模型权重不随安装包交付。下载使用固定修订、已知字节数和 SHA-256；内网环境支持 + ZIP 导入。 +8. 内置连接推理不在 UI Renderer 或 Main 事件循环内执行。Main 通过受控推理进程调用 ONNX CPU Runtime，Renderer 只展示状态和发起经过校验的操作。 -8. 文档和查询编码必须具有显式角色。Tokenizer、特殊 Token、Pooling、归一化、前缀、 +9. 文档和查询编码必须具有显式角色。Tokenizer、特殊 Token、Pooling、归一化、前缀、 最大序列长度、精度和输出维度共同构成可复现的 `EmbeddingEncodingRecipe`。 -9. Provider Fingerprint 必须覆盖完整编码配方和数据路径。任一兼容性字段变化后，旧向量 +10. Provider Fingerprint 必须覆盖完整编码配方和数据路径。任一兼容性字段变化后，旧向量 不参与新请求，用户需要显式重建。 -10. 保存模型切换不自动重建所有知识库。界面先说明受影响的知识库数量，再由用户逐库或 +11. 保存模型切换不自动重建所有知识库。界面先说明受影响的知识库数量，再由用户逐库或 显式批量发起重建。 --- -## 2. 背景与当前基础 +## 2. 实施基线 -### 2.1 当前已经具备的能力 +### 2.1 当前实现 GoodBuddy 当前已经具备： -- `EmbeddingProvider`、诊断、取消、超时和批量 Embeddings 调用。 -- OpenAI 兼容 Embeddings Endpoint、模型名称和可选 API Key 设置。 -- Main-only 凭据读取和系统安全存储。 -- Provider、Model、维度和 Fingerprint 参与的索引隔离。 +- 内置 Granite 系统连接与多个 OpenAI 兼容用户连接。 +- 与 LLM 模型连接一致的左侧连接列表、右侧详情管理模式。 +- 内置模型固定下载、进度、取消、SHA-256 校验、原子安装、ZIP 导入和删除。 +- `@huggingface/tokenizers`、`onnxruntime-web/wasm`、CLS Pooling、L2 归一化和 + 384 维输出组成的真实本地推理链路。 +- 独立 Utility Process、私有协议、超时、取消、崩溃恢复预算和有界结果校验。 +- `embedQuery()` 与 `embedDocuments()` 角色感知调用，以及旧 `embed()` 文档角色兼容。 +- Main-only OpenAI 兼容 Endpoint、模型名称和每连接可选 API Key。 +- Provider、Model、维度、数据路径和完整编码配方参与的 Fingerprint 与索引隔离。 - Float32 SQLite BLOB 存储、向量模长和有界分页余弦扫描。 - 文档级暂存与原子替换、失败和取消后的旧索引保护。 - 每个知识库的向量覆盖率、重建任务和检索通道诊断。 -- 本地 OCR 与本地语音模型已经使用的受管模型目录、固定下载、SHA-256、ZIP 迁移和 - 操作进度模式。 -- `onnxruntime-web` 及 Renderer OCR Worker 证明的 ONNX/WASM 打包基础。 -当前实现仍有以下限制： - -- 只有一个全局 OpenAI 兼容向量连接，没有应用托管的本地向量模型。 -- 设置界面把“本地 loopback 服务”和“云端服务”表示为同一 Endpoint 表单，数据路径 - 不够直观。 -- `EmbeddingProvider.embed()` 没有区分 Query 和 Document 角色。 -- Fingerprint 主要由 Endpoint 和模型名称组成，尚未覆盖 Tokenizer、Pooling、归一化、 - 精度和角色格式。 -- 当前分块按字符计数，Provider 可能按 Token 截断，应用无法证明实际编码内容完整。 -- 模型变更后的受影响知识库数量和重建影响未在全局设置中集中显示。 -- “失败时自动回退”的设置文案没有准确表达实际使用通道，容易让用户误解为应用会自动 - 切换向量模型。 +当前仍待后续知识增强阶段完成的能力包括全局展示受影响知识库数量、模型切换确认与逐库 +重建入口，以及 Token 感知分块。内置 Granite 客户端不会静默截断超过 32,768 Token 的 +输入，而是明确拒绝并返回诊断。 ### 2.2 当前存储和搜索约束 @@ -109,8 +110,8 @@ GoodBuddy 当前已经具备： - 不安装其他服务即可按需获得可用的本地语义检索。 - 清楚知道向量计算发生在本机、局域网自托管服务还是云端。 -- 在高性能硬件上自行选择 Ollama 中的 Qwen3-Embedding、BGE 等模型。 -- 在断网设备上通过 ZIP 导入应用托管模型。 +- 在高性能硬件或自托管环境中通过 OpenAI 兼容连接使用自行选择的向量模型。 +- 在断网设备上通过 ZIP 导入内置模型。 - 切换模型前知道哪些知识库会失去兼容向量，以及接下来需要做什么。 - 查看模型下载、校验、加载、真实推理测试和索引重建的独立状态。 - 失败时得到准确原因，不被应用暗中换模型、换精度或上传数据。 @@ -120,18 +121,50 @@ GoodBuddy 当前已经具备： - 提供一个六平台共同支持的轻量本地默认向量能力。 - 复用现有知识索引、模型管理、安全存储、任务中心和通知体系。 - 保持应用安装包不包含模型权重，避免无条件增加下载和安装体积。 -- 对本地、Ollama、自托管和云端连接使用同一上层 Provider 和索引兼容契约。 +- 对内置连接和任意 OpenAI 兼容连接使用同一上层 Provider 和索引兼容契约。 - 让模型升级、量化变化和角色格式变化都可见、可诊断、可重建。 - 保持 Renderer 无 Node 集成，不向 Renderer 暴露数据库、模型目录或长期凭据。 ### 3.3 质量目标 -- 应用托管模型对当前中文、中英文混合和跨语言检索样例具有稳定收益。 +- 内置模型对当前中文、中英文混合和跨语言检索样例具有稳定收益。 - 同一模型包、编码配方和输入在同一 CPU Runtime 下产生可复现的向量维度和近似值。 - UI 线程不执行 Tokenization、ONNX 推理或大向量序列化。 - 取消、应用退出和模型切换不会留下半安装目录、孤立进程或半替换索引。 - 任何向量不可用、索引不兼容或输入过长都进入有界诊断，不能表现为正常的空向量结果。 +### 3.4 适用范围与 Agent Runtime 边界 + +首期允许调用向量连接的生产服务只有 GoodBuddy Main 中的知识能力： + +- 知识库文档分块向量化。 +- 知识查询向量化。 +- 向量召回、混合检索以及 GraphRAG 中明确设计的语义召回。 +- 不读取用户知识库的固定连接诊断和检索质量评测。 + +对话需要知识时，数据链路固定为： + +```text +用户问题 +→ Main / KnowledgeService 生成查询向量并检索 +→ Main 选择有界文本证据、引用和诊断 +→ Runtime Orchestrator 把这些不可信证据作为上下文提供给当前 Agent Runtime +``` + +边界规则： + +- 向量连接不进入 `AgentRuntimeSelection`、Runtime Model Profile 或项目 Runtime 设置。 +- OpenCode、Continue、DeepSeek Harness、直连 LLM 和远程 Runtime 都不能直接创建或调用 + `EmbeddingProvider`。 +- 向量 Endpoint、API Key、模型目录、Tokenizer、向量和索引 Fingerprint 不进入 Runtime + 进程环境、配置文件、模型网关或远程主机。 +- `knowledge_search` 等 Runtime 工具只能调用 Main 暴露的有界知识服务，并获得文本结果、 + 引用和脱敏诊断。 +- Agent Runtime 的选择或切换不改变当前向量连接；向量连接的选择或切换也不改变当前 + Agent Runtime。 +- 未来记忆、笔记或其他本地语义检索若要复用向量连接，必须经过独立产品设计、权限与 + 数据路径评审后显式接入，不能因存在全局向量连接而自动获得使用权。 + --- ## 4. 非目标 @@ -139,15 +172,22 @@ GoodBuddy 当前已经具备： 首期不包含： - 把 Granite、Qwen、BGE 或其他向量权重直接打入 GoodBuddy 安装包。 -- 由 GoodBuddy 自动安装 Ollama、执行 `ollama pull` 或管理 Ollama 服务生命周期。 -- 根据硬件、网络、延迟或准确率自动选择应用托管模型、Ollama 或云端模型。 -- 在 Ollama 不可用时自动连接云端，或在云端失败时自动加载应用托管模型。 +- 由 GoodBuddy 安装、启动、停止或升级用户填写的 OpenAI 兼容服务。 +- 根据硬件、网络、延迟或准确率自动选择内置连接或用户连接。 +- 在用户连接不可用时自动启用内置连接，或在内置连接失败时自动连接外部服务。 - 在模型失败时自动改用另一个模型名称、量化档位、输出维度或 Query 前缀。 -- 为 Ollama 模型质量作统一保证。用户选择的具体模型需要通过真实诊断和知识评测验证。 +- 为用户连接背后的具体服务或模型质量作统一保证。用户选择的模型需要通过真实诊断和 + 知识评测验证。 - 首期引入专用向量数据库、GPU 索引或平台原生 SQLite 向量扩展。 - 在后台自动更新模型版本或覆盖当前正在使用的模型包。 - 静默截断超过模型上下文的分块。 - 把知识库文档、查询、向量或模型凭据用于训练、遥测或外部质量分析。 +- 把向量连接作为聊天生成模型、Agent Runtime 模型、Runtime 模型网关或远程 Runtime + 依赖。 +- 向 Agent Runtime 暴露 Embedding Provider、模型文件、Tokenizer、原始向量、Endpoint + 或凭据。 +- 在没有独立设计和显式接入的情况下，把当前向量连接自动扩展到记忆、自动化、项目文件 + 搜索或其他应用服务。 --- @@ -155,28 +195,27 @@ GoodBuddy 当前已经具备： | 术语 | 定义 | | --- | --- | -| `managed-local` | GoodBuddy 管理模型包和推理进程，推理数据不离开设备 | -| `ollama` | 用户管理的 Ollama 服务；loopback Endpoint 才标记为本机 | +| `builtin` | GoodBuddy 提供的只读系统连接，管理模型包和推理进程，推理数据不离开设备 | | `openai-compatible` | 用户配置的通用 OpenAI 兼容 Embeddings 服务 | | 模型包 | 固定修订的 ONNX、Tokenizer、配置、许可和校验清单 | | 编码配方 | Query/Document 格式、Tokenizer、Pooling、归一化、精度、维度和 Token 限制 | -| Provider Profile | 用户保存的执行路径和模型配置，不含凭据正文 | -| Runtime Snapshot | 一次诊断、索引任务或查询实际使用的冻结配置 | -| Provider Fingerprint | 对执行路径和编码配方的规范化摘要，用于索引隔离 | +| Connection Profile | 一个可选择的内置或用户连接配置，不含凭据正文 | +| Embedding Operation Snapshot | 一次向量诊断、索引任务或查询实际使用的冻结配置；与 Agent Runtime 无关 | +| Provider Fingerprint | 对连接、数据路径和编码配方的规范化摘要，用于索引隔离 | | 兼容索引 | Provider Fingerprint、模型、维度和内容校验均匹配的向量集合 | | 可见通道降级 | 向量不可用时继续使用已配置的其他检索通道，并明确显示差异 | -Ollama Endpoint 为 `127.0.0.1`、`localhost` 或 `[::1]` 时，可以标记“本机 Ollama”。 -局域网地址、VPN 地址、域名或公网地址统一标记“自托管网络服务”，不得仅因 Provider 类型 -是 Ollama 就声称数据不离开本机。 +OpenAI 兼容 Endpoint 为 `127.0.0.1`、`localhost` 或 `[::1]` 时，可以标记“本机服务”。 +局域网地址、VPN 地址、域名或公网地址统一标记“网络服务：<主机名>”。GoodBuddy 不根据 +URL 或模型名称猜测具体服务品牌，也不把非 loopback Endpoint 描述为数据不离开本机。 --- ## 6. 模型选择 -### 6.1 应用托管默认模型 +### 6.1 内置默认模型 -首个应用托管目录只包含一个推荐模型： +首个内置连接只包含一个固定模型： | 字段 | 值 | | --- | --- | @@ -185,21 +224,38 @@ Ollama Endpoint 为 `127.0.0.1`、`localhost` 或 `[::1]` 时，可以标记“ | 输出维度 | 384 | | 上下文 | 32K Token | | 语言 | 200+，重点增强 52 种语言的检索 | -| 权重 | 固定修订的官方 INT8 ONNX | -| 量化权重体积 | 约 91 MB，以目录清单中的实际固定字节数为准 | +| 权重 | 固定修订的通用 INT8 ONNX 工件 | +| 量化权重体积 | 97,858,099 字节，不含 Tokenizer 和配置 | | License | Apache 2.0 | -| UI 标签 | 推荐、轻量、多语言、本地、INT8 | +| 编码配方 | Query/Document 空前缀、CLS Pooling、L2 归一化 | +| UI 标签 | 内置、轻量、多语言、本地、INT8 | 选择理由： 1. 384 维降低当前 SQLite 线性扫描和长期索引成本。 2. 32K 上下文覆盖 GoodBuddy 当前常见分块，不要求为了默认模型全面缩短分块。 3. 中文、英文和混合内容可以使用同一索引，不需要按语言拆分。 -4. 官方提供量化 ONNX，适合 CPU 路径和按需下载。 +4. ONNX Community 提供面向 Transformers.js 的通用 INT8 转换，并在 Hugging Face 与 + ModelScope 上提供 SHA-256 相同的固定工件。 5. Apache 2.0 适合应用内提供目录元数据和用户自行下载使用。 -“推荐”表示 GoodBuddy 的默认集成目标，不表示模型已安装、已启用或适合所有资料。用户仍需 -执行下载或导入、保存选择和真实测试。 +“内置”表示 GoodBuddy 提供并维护该系统连接，不表示权重已经包含在安装包中。用户仍需 +执行下载或导入、真实测试和保存选择。 + +IBM 官方仓库当前提供的轻量工件名为 `model_quint8_avx2.onnx`，明确面向 AVX2，不能未经 +验证作为六平台统一工件。当前内置连接使用 +`onnx-community/granite-embedding-97m-multilingual-r2-ONNX` 中固定修订的 +`onnx/model_quantized.onnx`： + +- Hugging Face Revision：`536a9f241cb3f02a9c5995a1e708c784bd274859`。 +- ModelScope Revision：`2741cd30a7448219ec2699afdf373a44df5aaa33`。 +- 文件大小：`97,858,099` 字节。 +- SHA-256：`704c1ebca5fbb7cd83ced41827658ac4c9990c64f7f2874d22b78044e5022e22`。 + +该工件是社区自动转换。Windows x64 已使用固定 Tokenizer 和真实 ONNX 文件完成 +Query/Document 推理，确认输出为 384 维、CLS Pooling 和 L2 归一化链路可运行。其余五个 +发布目标仍由对应原生 Runner 验证安装、加载、推理、取消和退出；任一目标验证失败时必须 +修复同一路径，不能静默改用平台特定工件或用户连接。 ### 6.2 为什么首期不托管多个默认模型 @@ -207,7 +263,7 @@ Ollama Endpoint 为 `127.0.0.1`、`localhost` 或 `[::1]` 时，可以标记“ 当前默认中文分块，并增加每条索引的维度成本。首期不把它设为第二个自动候选，也不建立 Granite → BGE 的回退关系。 -未来只有在以下条件全部满足后，才把其他模型加入应用托管目录： +未来只有在以下条件全部满足后，才增加其他内置连接： - 固定修订和全部文件 SHA-256 已独立验证。 - 六个平台均能加载和运行。 @@ -218,14 +274,12 @@ Granite → BGE 的回退关系。 ### 6.3 高硬件或准确性需求 -高性能用户可以自行安装 Ollama，并在 GoodBuddy 中显式选择： +高性能用户可以自行准备提供 OpenAI 兼容 Embeddings API 的本机、网络或外部服务，并在 +GoodBuddy 中添加连接，显式填写 Endpoint、模型名称、可选凭据和编码预设。 -- Qwen3-Embedding 的合适参数规模。 -- BGE-M3 或其他 BGE Embedding 模型。 -- Ollama 目录中其他声明支持 Embedding 的模型。 - -GoodBuddy 不代替用户判断模型的显存、内存、速度、语言和 License。设置界面必须显示 -Ollama 返回的实际模型标识、输出维度和真实诊断结果，不能仅根据名称宣称质量更高。 +GoodBuddy 不代替用户判断服务背后模型的显存、内存、速度、语言和 License。设置界面必须 +显示用户填写的模型标识、实际 Endpoint 主机、输出维度和真实诊断结果，不能仅根据名称 +宣称质量更高。 --- @@ -237,9 +291,9 @@ Ollama 返回的实际模型标识、输出维度和真实诊断结果，不能 设置 └─ 模型连接 └─ 向量模型 - ├─ 应用托管 - ├─ Ollama / 自托管 - └─ 云端服务 + ├─ GoodBuddy 内置连接 + ├─ 用户添加的 OpenAI 兼容连接 + └─ 当前连接详情与诊断 ``` “模型连接”中的模型类型继续使用现有 `SegmentedControl`： @@ -248,8 +302,10 @@ Ollama 返回的实际模型标识、输出维度和真实诊断结果，不能 [LLM 模型] [向量模型] [重排模型] [语音输入] ``` -向量模型内部的“应用托管 / Ollama / 云端服务”也是当前面板内的互斥模式，使用共享 -`SegmentedControl`，不使用 `PageTabs`，不新增页面专属切换样式。 +向量模型内部不增加第二组模式切换。它复用 LLM 模型连接的列表、当前连接、添加、编辑、 +测试和删除交互。复用只限界面模式和通用表单行为，不复用 LLM/Agent Runtime 的连接契约、 +选择状态、凭据注入或路由。内置连接只是列表中的一种只读系统连接，不是独立页面或独立 +模式。 知识库设置继续显示当前知识库的向量覆盖率和重建入口，不复制全局模型下载、Endpoint 或 凭据编辑表单。 @@ -262,56 +318,44 @@ Ollama 返回的实际模型标识、输出维度和真实诊断结果，不能 向量模型面板按以下顺序组织： -1. 启用状态和当前生效摘要。 -2. 执行路径选择。 -3. 持续可见的数据路径说明。 -4. 当前路径的配置或模型目录。 -5. 待保存变化和索引影响。 -6. 真实诊断。 -7. 模型目录和低频维护操作。 +1. 标题、说明、启用开关和“添加自定义”。 +2. 左侧为内置连接与用户连接组成的单一连接列表。 +3. 右侧为所选连接详情；列表选择只切换详情，不改变当前连接。 +4. 当前连接单选、连接编辑和添加/删除均保存在页面草稿中，只随“保存设置”提交。 +5. 内置详情提供下载、取消、ZIP 导入、删除和真实诊断。 +6. 用户连接详情提供名称、Endpoint、模型、认证、API Key、删除和真实诊断。 + +停用向量检索只阻止把连接设为当前和执行诊断，不阻止选择详情、编辑用户连接或管理内置 +模型文件。 示意： ```text -向量模型 -用于知识库语义检索与 GraphRAG [保存设置] +向量模型 [启用向量检索] [添加自定义] -[启用向量模型 开关] -当前生效：应用托管 · Granite Embedding 97M R2 · INT8 · 384 维 +┌ 可用连接 ───────────────┬ 连接详情 ─────────────────────────────┐ +│ GoodBuddy 内置向量模型 │ GoodBuddy 内置向量模型 (当前连接) │ +│ Granite 97M R2 │ 模型：Granite 97M R2 │ +│ [当前] │ 状态：已安装 │ +│ 自定义向量模型 │ [测试] [从 ZIP 导入] [移除本地模型] │ +│ nomic-embed-text │ │ +└──────────────────────────┴──────────────────────────────────────┘ -执行路径 -[ 应用托管 ] [ Ollama / 自托管 ] [ 云端服务 ] +用户连接详情： +名称 / Embeddings Endpoint / 模型名称 / 认证方式 / API Key +[删除] [测试] -┌ 数据路径 ───────────────────────────────────────────────┐ -│ 本机知识分块 → GoodBuddy 本地推理进程 → 本地 SQLite │ -│ 模型安装完成后，推理不需要联网，文档和查询不离开本机。 │ -└────────────────────────────────────────────────────────┘ - -当前模型 -┌────────────────────────────────────────────────────────┐ -│ Granite Embedding 97M Multilingual R2 │ -│ 推荐 · 多语言 · INT8 · 384 维 · 32K · Apache 2.0 │ -│ 状态：已安装 │ -│ [设为当前模型] [导出 ZIP] [打开模型仓库] [删除] │ -└────────────────────────────────────────────────────────┘ - -配置影响 -3 个知识库没有与待保存配置兼容的向量，需要显式重建。 -[查看受影响的知识库] - -连接诊断 -最近测试：可用 · 384 维 · 178 毫秒 · 2026-08-19 10:30 -[测试待保存配置] +页面级 [保存设置] 提交启用状态、当前连接和用户连接草稿。 ``` ### 8.2 当前生效与待保存 界面必须同时区分： -- **当前生效**：Main 正在用于新查询和新索引任务的已保存 Profile。 -- **正在编辑**：表单中的草稿。 -- **已验证草稿**：最近一次诊断与草稿 Fingerprint 完全匹配。 -- **待保存**：草稿与当前生效 Profile 不同。 +- **当前生效**：Main 正在用于新查询和新索引任务的已保存 Connection Profile。 +- **正在编辑**：所选用户连接的表单草稿，或准备设为当前的内置连接。 +- **已验证草稿**：最近一次诊断与连接草稿 Fingerprint 完全匹配。 +- **待保存**：当前连接选择或用户连接草稿与已保存状态不同。 用户修改路径、Endpoint、模型、精度、维度或编码预设后，旧诊断立即显示： @@ -319,82 +363,71 @@ Ollama 返回的实际模型标识、输出维度和真实诊断结果，不能 不能把旧模型的“测试成功”继续显示在新模型旁边。 -### 8.3 应用托管模式 +### 8.3 内置连接 -应用托管模式显示： +内置连接是 GoodBuddy 提供的只读系统 Profile，与用户连接出现在同一列表。它显示： -- 模型目录根路径。 -- 当前目录中的模型选择器。 -- 单个模型详情卡，不同时铺开大量卡片。 -- 下载、取消、导入 ZIP、导出 ZIP、打开目录和删除。 -- 固定版本、量化、维度、上下文、语言、下载体积、预计内存、License 和仓库链接。 +- 固定连接名称“GoodBuddy 内置”和固定模型名称。 +- 下载、取消、导入 ZIP 和删除。 +- 固定模型名称和安装状态。 - “已安装”“待保存”“正在使用”“下载中”“校验中”“损坏”等状态。 -首期只有一个目录条目时仍保留选择器和目录契约，为后续经评测的模型预留兼容结构。 +内置连接不能编辑 Endpoint、模型名称、凭据、精度、维度或编码配方。未来增加其他内置 +模型时，为每个经过验证的模型提供独立系统连接，不在同一个连接中建立自动模型路由。 未安装时主操作是“下载”，次操作是“导入 ZIP”。下载完成后可以自动把该模型选入草稿， 但不能自动保存、启用或重建知识库。 -### 8.4 Ollama / 自托管模式 +### 8.4 OpenAI 兼容连接 -Ollama 模式显示： +用户可以添加多个 OpenAI 兼容连接。添加和编辑表单复用 LLM 模型连接的结构，显示： -- Endpoint，loopback 默认建议为 `http://127.0.0.1:11434/v1/embeddings`。 -- 模型名称，例如 `qwen3-embedding:0.6b` 或 `bge-m3`。 -- 可选 API Key，只在用户的 Ollama 网关需要认证时使用。 +- 用户可编辑的连接名称。 +- 完整 Embeddings Endpoint。 +- 模型名称。 +- 可选 API Key 和凭据来源。 - 编码预设和模型最大输入 Token。 -- 数据位置：“本机 Ollama”或“自托管网络服务：主机名”。 -- “检测服务”“测试向量生成”“打开 Ollama Embeddings 文档”。 +- 根据 Endpoint 显示“本机服务”或“网络服务：<主机名>”。 +- “测试向量生成”和连接删除操作。 -GoodBuddy 不显示“安装 Ollama”主操作，也不自动执行 CLI。可以提供只读说明或复制命令， -但执行安装和拉取模型必须由用户在其环境中完成。 +UI 不显示具体服务的专属类型、安装入口、品牌文案或检测按钮。用户可以在任何兼容服务中 +自行准备模型，再将其 Endpoint 作为普通 OpenAI 兼容连接添加。 -当服务支持模型摘要或 Digest 时，诊断结果显示并纳入 Fingerprint。无法读取稳定 Digest -时，界面显示： +通用接口无法提供稳定模型摘要时，界面显示： > 服务未提供可验证的模型修订。GoodBuddy 将按 Endpoint、模型名称和编码配置隔离索引； > 服务端原地替换同名模型后，需要重新测试并重建。 -### 8.5 云端服务模式 - -云端模式沿用 OpenAI 兼容字段： - -- 完整 Embeddings Endpoint。 -- 模型名称。 -- API Key 和凭据来源。 -- 编码预设。 -- Provider 或账号备注，不保存凭据正文到 Profile。 - -Endpoint 下方持续显示： +非 loopback Endpoint 下方持续显示： > 建立索引时会向 `api.example.com` 发送已启用知识库的分块文本；检索时会发送用户查询。 -> GoodBuddy 不会自动切换到本地模型或其他云端模型。 +> GoodBuddy 不会自动切换到内置连接或其他用户连接。 设置页不把“已配置 API Key”表示为“服务可用”。只有真实 Embeddings 请求可以产生 “测试成功”状态。 -### 8.6 数据路径卡 +### 8.5 数据去向 -数据路径卡不能折叠，也不能只用“本地”“云端”徽标表达： +内置连接名称和模型管理操作已经明确表达本地执行，不再常驻展示固定的进程与 SQLite +路径。下载模型时访问仓库属于模型安装操作，不等同于推理时发送知识内容。 + +用户连接必须在 Endpoint 下方显示实际数据去向： | 路径 | 文案 | | --- | --- | -| 应用托管 | 本机知识分块 → GoodBuddy 本地推理进程 → 本地 SQLite | -| loopback Ollama | 本机知识分块 → 本机 Ollama → 本地 SQLite | -| 网络自托管 | 本机知识分块 → `host:port` → 本地 SQLite | -| 云端 | 本机知识分块 → `provider host` → 本地 SQLite | +| Loopback 兼容连接 | 知识分块和查询将发送到此设备上的 `host:port` | +| 非 Loopback 兼容连接 | 建立索引时会向 `host:port` 发送知识分块，检索时会发送查询 | -下载模型时访问模型仓库与推理时发送文档是不同数据行为。应用托管模式应说明“下载阶段需要 -访问固定模型地址；安装后推理不联网”，不能简单写“完全不联网”。 +非 Loopback 提示持续可见，不能只用“网络服务”徽标表达。 -### 8.7 状态 +### 8.6 状态 | 状态 | 界面表达 | 可用操作 | | --- | --- | --- | | 未安装 | 需要下载或导入后使用 | 下载、导入 ZIP | | 下载中 | 文件、已下载字节、总字节、百分比 | 取消 | | 校验中 | 正在校验大小和 SHA-256 | 取消 | -| 已安装 | 已安装，尚未使用 | 设为当前模型、导出、删除 | +| 已安装 | 已安装，尚未使用 | 设为当前模型、导入、删除 | | 待保存 | 选择尚未生效 | 保存设置 | | 可用 | 路径、模型、维度、最近测试时间 | 重新测试 | | 加载失败 | 脱敏原因和处理建议 | 重试当前配置 | @@ -405,7 +438,7 @@ Endpoint 下方持续显示： 状态必须同时使用图标和文字，不只依赖颜色。 -### 8.8 真实诊断 +### 8.7 真实诊断 诊断使用固定、不含用户数据的测试文本，至少验证： @@ -430,7 +463,7 @@ Endpoint 下方持续显示： 诊断不读取用户知识库、不改变索引，也不证明真实资料的召回质量。 -### 8.9 保存和切换确认 +### 8.8 保存和切换确认 保存会改变 Fingerprint 的配置前显示确认对话框： @@ -450,9 +483,9 @@ Endpoint 下方持续显示： 保存前存在运行中的向量重建任务时，界面必须说明任务将被取消并丢弃未发布的暂存结果。 用户确认后先完成取消，再替换活动 Profile。 -### 8.10 删除模型 +### 8.9 删除模型 -- 当前生效的应用托管模型不能直接删除。 +- 当前生效的内置模型不能直接删除。 - 用户必须先停用向量模型或显式选择并保存其他 Profile。 - 删除只移除模型文件，不自动删除旧向量索引。 - 删除确认显示模型名称、版本和磁盘空间。 @@ -481,7 +514,7 @@ Endpoint 下方持续显示： 规则： - 当前 Profile 和索引 Fingerprint 完全匹配时才显示“兼容”。 -- 重建始终使用启动任务时冻结的 Runtime Snapshot。 +- 重建始终使用启动任务时冻结的 Embedding Operation Snapshot。 - 重建期间全局 Profile 变化时取消旧任务，不让任务后半段切到新模型。 - 同一 Fingerprint 的重建失败时，可以继续保留上一次完整就绪索引。 - Fingerprint 已变化时，旧索引可以保留在数据库中等待清理，但不参与当前向量召回。 @@ -498,9 +531,9 @@ Endpoint 下方持续显示： 以下变化不得在诊断、索引任务或查询中静默发生： -- 应用托管、Ollama、自托管网络服务和云端之间切换。 +- 内置连接与任何 OpenAI 兼容连接之间切换。 - Endpoint、Provider、账号或数据位置切换。 -- 模型名称、模型修订、Ollama Digest 或模型文件变化。 +- 模型名称、模型修订、服务端模型摘要或模型文件变化。 - INT8、FP16、FP32 或其他精度变化。 - 输出维度变化或 Matryoshka 截断维度变化。 - Query/Document 前缀、指令或角色格式变化。 @@ -510,9 +543,9 @@ Endpoint 下方持续显示： ### 10.2 允许的内部恢复 -以下操作可以自动执行，但必须保持同一 Runtime Snapshot 和语义： +以下操作可以自动执行，但必须保持同一 Embedding Operation Snapshot 和语义： -- 同一应用托管模型进程的有限重启。 +- 同一内置模型进程的有限重启。 - 同一 Endpoint、模型、Digest 和编码配方的有限网络重试。 - 动态批大小、线程调度、缓存和空闲卸载等不改变向量语义的资源优化。 - 同一索引任务内对可重试批次进行有界重试。 @@ -555,17 +588,22 @@ FTS、CJK、向量和图谱是用户已配置的并列检索通道。向量失 └───────────────┬──────────────────────┬────────────────────┘ │ │ Managed utility process HTTP(S) provider - Tokenizer + ONNX CPU Ollama / self-hosted / cloud + Tokenizer + ONNX CPU OpenAI-compatible endpoint │ Managed model directory manifest + fixed artifacts ``` +Agent Runtime 不在向量 Provider 架构内。对话编排只能通过 `KnowledgeService` 获取有界 +检索结果；图中的 Provider、推理进程、HTTP Endpoint、凭据和向量存储都保持在 Main 知识 +服务边界内。 + ### 11.1 Renderer Renderer 负责： -- 显示目录、下载进度、当前/草稿 Profile 和数据路径。 +- 显示内置模型状态与下载进度、当前/草稿 Profile，并仅为用户连接显示根据 Endpoint + 判断的实际数据去向。 - 发起下载、导入、导出、选择、删除、诊断和重建。 - 展示知识库影响、真实维度和脱敏错误。 - 使用应用通知显示短期成功和非字段异步错误。 @@ -577,7 +615,7 @@ Renderer 不负责： - 读取 API Key。 - 为知识库分块执行 Tokenization 或推理。 - 决定 Provider 回退。 -- 直接访问 Ollama、云端 Endpoint 或 SQLite。 +- 直接访问用户 Endpoint 或 SQLite。 ### 11.2 Main @@ -587,17 +625,17 @@ Main 负责： - 管理模型目录、固定下载、ZIP 和文件摘要。 - 保存不含凭据正文的 Profile，并在系统安全存储中绑定凭据。 - 规范化 Endpoint 和计算 Fingerprint。 -- 创建冻结的 Provider Runtime Snapshot。 +- 创建冻结的 Embedding Operation Snapshot。 - 选择唯一的 Adapter，不执行自动 Adapter 路由。 - 将查询、索引、取消、超时和应用关闭传播到推理进程或 HTTP Provider。 - 对模型切换、任务取消和知识索引状态执行一致的生命周期管理。 -### 11.3 应用托管推理进程 +### 11.3 内置连接推理进程 -应用托管模型运行在 GoodBuddy 管理的独立 Utility Process： +内置模型运行在 GoodBuddy 管理的独立 Utility Process： - 不监听端口。 -- 不继承云端 API Key、Runtime 凭据或无关环境变量。 +- 不继承云端 API Key、Agent Runtime 凭据或无关环境变量。 - 只接受 Main 通过受控消息通道发送的有界请求。 - 模型路径由 Main 从已校验 Manifest 解析，不接受 Renderer 或请求中的任意路径。 - Tokenization、ONNX 推理、Pooling 和归一化都在该进程完成。 @@ -608,14 +646,14 @@ Main 负责： 初始推理后端固定为 ONNX CPU。实现阶段必须验证所选 ONNX JavaScript/Native Runtime 在 六个平台的打包、许可、算子、线程和退出行为；某个平台未通过时，该平台显示“不支持此 -本地模型”，不能改连 Ollama 或云端继续。 +本地模型”，不能改用用户连接继续。 现有 Renderer OCR Worker 不直接复用，因为知识分块不应为了推理被批量发送到 UI Renderer。可以复用其 ONNX 资源固定、Worker 消息边界和有界错误经验。 -### 11.4 HTTP Provider +### 11.4 OpenAI 兼容 Provider -Ollama、自托管和云端首期复用 Main 中的 OpenAI 兼容 Embeddings Client，并保持： +所有用户连接复用 Main 中的 OpenAI 兼容 Embeddings Client，并保持： - HTTP(S) 协议限制。 - Endpoint 长度和规范化。 @@ -624,8 +662,20 @@ Ollama、自托管和云端首期复用 Main 中的 OpenAI 兼容 Embeddings Cli - 返回数量、索引顺序、维度、有限值和非零范数校验。 - 取消传播和脱敏错误。 -Ollama 的服务检测、模型 Digest 和上下文元数据可以使用独立的 Ollama Adapter，但最终 -Embedding 调用仍遵守统一 `EmbeddingProvider` 契约。 +首期不增加品牌专属 Adapter，不探测服务安装状态、已安装模型列表或管理服务生命周期。 + +### 11.5 Agent Runtime 集成边界 + +知识检索发生在 Agent Runtime 启动前的 Main 预检索，或 Runtime 通过受控 +`knowledge_search` 工具请求 Main 检索时。两种路径都由 `KnowledgeService` 调用当前 +Embedding Operation Snapshot，Runtime 只接收： + +- 有界的证据文本和来源引用。 +- 请求通道、实际通道和是否降级。 +- 脱敏且有长度上限的检索诊断。 + +Runtime 不接收连接 Profile、Endpoint、凭据引用、模型路径、Tokenizer、Float32 向量或 +索引存储键。远程 Runtime 不通过模型网关转发 Embeddings 请求，也不在远端复制向量配置。 --- @@ -655,8 +705,8 @@ interface EmbeddingProvider { 索引分块始终使用 `document`，检索问题始终使用 `query`。诊断分别发送固定的 Query 和 Document 测试输入，不能偷用文档角色代替查询测试，也不创建第三种模型角色。 -应用托管 Granite 使用模型包内固定的官方编码配置。Ollama 或通用服务通过用户显式选择的 -编码预设决定是否添加前缀或指令。不能只根据模型名称在后台猜测并改变输入。 +内置 Granite 使用模型包内固定的编码配置。OpenAI 兼容连接通过用户显式选择的编码预设 +决定是否添加前缀或指令。不能只根据模型名称在后台猜测并改变输入。 ### 12.2 编码预设 @@ -665,14 +715,14 @@ Profile 至少支持： - `symmetric`：Query 和 Document 不添加不同前缀。 - `query-passage`：固定 Query 与 Passage 前缀。 - `instruction-query`：Query 使用固定检索指令，Document 使用固定文档格式。 -- `managed-manifest`：完全由受信任模型包 Manifest 决定，用户不可局部修改。 +- `builtin-manifest`：完全由受信任模型包 Manifest 决定，用户不可局部修改。 -已知 Ollama 模型可以提供推荐预设，但 UI 必须在保存前显示实际选择。自定义预设中的前缀 -有长度上限，不允许插入凭据或读取知识内容。 +UI 必须在保存前显示用户连接的实际编码预设。自定义预设中的前缀有长度上限，不允许插入 +凭据或读取知识内容。 ### 12.3 Tokenizer、Pooling 和归一化 -应用托管模型必须从固定模型包读取： +内置模型必须从固定模型包读取： - `tokenizer.json` 及其摘要。 - 特殊 Token ID。 @@ -691,7 +741,7 @@ SQLite 当前会保存模长并计算余弦相似度。即使模型输出已经 ### 12.4 输入长度 - 每个 Profile 声明 `maximumSequenceTokens`。 -- 应用托管路径在发送 ONNX 前完成 Token 计数。 +- 内置路径在发送 ONNX 前完成 Token 计数。 - 批处理同时受最大项目数和最大总 Token 数限制。 - 单条输入超过上限时返回 `input_too_long`，不截断后继续表示成功。 - 已知模型上下文小于当前知识库分块时，设置和知识库页面持续显示不兼容提示。 @@ -708,15 +758,15 @@ GoodBuddy 继续执行客户端字符上限，但不声称已经证明服务端 建议扩展 `src/shared/embedding-contracts.ts`： ```ts -type EmbeddingProviderKind = - | 'managed-local' - | 'ollama' - | 'openai-compatible' +type EmbeddingConnectionKind = 'builtin' | 'openai-compatible' type EmbeddingDataLocation = | { kind: 'device' } - | { kind: 'self-hosted'; endpointHost: string } - | { kind: 'cloud'; endpointHost: string; region?: string } + | { + kind: 'endpoint' + endpointHost: string + scope: 'loopback' | 'network' + } type EmbeddingEncodingRecipe = { recipeId: string @@ -737,10 +787,11 @@ type EmbeddingModelReference = { precision?: 'int8' | 'fp16' | 'fp32' | 'provider-managed' } -type EmbeddingProviderProfile = { +type EmbeddingConnectionProfile = { id: string name: string - kind: EmbeddingProviderKind + kind: EmbeddingConnectionKind + system: boolean endpoint?: string accountRef?: string credentialRef?: string @@ -749,17 +800,19 @@ type EmbeddingProviderProfile = { dataLocation: EmbeddingDataLocation } -type EmbeddingRuntimeSnapshot = { - profile: EmbeddingProviderProfile +type EmbeddingOperationSnapshot = { + profile: EmbeddingConnectionProfile fingerprint: string actualDimensions: number createdAt: string } ``` -`accountRef` 和 `credentialRef` 只引用 Main 加密设置，不包含密钥正文。Endpoint 进入 -Profile 前删除用户名、密码和 Fragment；查询参数只有在 Provider API 明确需要且经过 -脱敏策略验证时才允许保留。 +内置连接固定为 `kind: 'builtin'` 和 `system: true`，不能携带 Endpoint、账号或凭据。 +用户连接固定为 `kind: 'openai-compatible'` 和 `system: false`。`accountRef` 和 +`credentialRef` 只引用 Main 加密设置，不包含密钥正文。Endpoint 进入 Profile 前删除 +用户名、密码和 Fragment；查询参数只有在 Provider API 明确需要且经过脱敏策略验证时 +才允许保留。 模型目录契约复用语音模型管理模式，并增加适合文本向量的文件角色： @@ -782,14 +835,14 @@ Fingerprint 使用规范化 JSON 计算 SHA-256，至少覆盖： - Provider Kind。 - 规范化 Endpoint 和数据位置。 - 模型 ID。 -- 固定 Revision、模型文件摘要或 Ollama Digest。 +- 固定 Revision、模型文件摘要或服务端稳定模型摘要。 - 精度。 - 输出维度。 - Tokenizer 摘要。 - Pooling 和归一化。 - Query/Document 模板。 - 最大序列长度。 -- 影响向量数值的 Runtime 配置版本。 +- 影响向量数值的推理配置版本。 Fingerprint 不包含： @@ -813,8 +866,8 @@ Fingerprint 本身必须来自上述完整规范化结构。 5. 分块内容校验和匹配。 6. 文档级索引状态为 Ready。 -不能因为两个模型都返回 384 维就视为兼容，也不能因为 Ollama 中模型名称相同就忽略 -Digest 变化。 +不能因为两个模型都返回 384 维就视为兼容，也不能因为用户连接中的模型名称相同就忽略 +Endpoint、编码配置或服务端稳定模型摘要变化。 ### 14.3 模型切换 @@ -837,19 +890,17 @@ Digest 变化。 ### 15.1 目录 -建议受管目录： +受管目录： ```text /models/embedding/ ├─ granite-embedding-97m-multilingual-r2-int8/ │ ├─ manifest.json -│ ├─ model.int8.onnx +│ ├─ model_quantized.onnx │ ├─ tokenizer.json │ ├─ tokenizer_config.json │ ├─ config.json -│ ├─ modules.json -│ ├─ LICENSE -│ └─ NOTICE +│ └─ special_tokens_map.json └─ .staging-/ ``` @@ -871,57 +922,52 @@ Digest 变化。 - 两个来源可以维护不同的固定大小和 SHA-256，但文件名和运行角色必须兼容；一次模型包 只能完整使用一个来源，不混合来源文件。 -### 15.3 ZIP 导入与导出 +### 15.3 ZIP 导入 - 导入先在暂存目录展开并验证 Manifest 和全部文件摘要。 - 不接受 Manifest 之外的可执行文件。 - 不接受路径穿越、符号链接、超大文件和压缩炸弹。 -- 导出只包含已验证模型包和许可文件，不包含设置、API Key、知识向量或用户数据。 - 导入的模型 ID、Revision 和摘要必须与目录条目一致。 --- -## 16. Ollama 与自托管服务 +## 16. OpenAI 兼容连接 ### 16.1 产品边界 -GoodBuddy 对 Ollama 负责： +GoodBuddy 对用户连接负责： -- 保存 Endpoint、模型名称和编码预设。 -- 可选检测 Ollama 服务和已安装模型元数据。 +- 保存连接名称、Endpoint、模型名称、可选凭据和编码预设。 - 发送有界 Embeddings 请求。 - 校验向量并展示真实诊断。 -- 根据 Endpoint、模型、Digest、维度和编码配置隔离索引。 +- 根据 Endpoint、模型、稳定模型摘要、维度和编码配置隔离索引。 GoodBuddy 不负责： -- 安装、启动、停止或升级 Ollama。 -- 下载、删除或更新 Ollama 模型。 -- 修改 Ollama 的监听地址、认证、代理或 GPU 设置。 +- 安装、启动、停止或升级 Endpoint 背后的服务。 +- 下载、删除或更新用户服务中的模型。 +- 修改用户服务的监听地址、认证、代理或硬件设置。 - 判断第三方模型 License 是否适合用户用途。 -### 16.2 显式高性能路径 +### 16.2 统一连接表单 -设置页可以提供以下说明： +添加连接只显示通用字段： -> 如果设备具有更多内存或显存，或你希望自行评估更高准确性的模型，可以安装 Ollama， -> 拉取 Qwen3-Embedding、BGE-M3 等 Embedding 模型，然后在这里填写模型名称。 -> GoodBuddy 不会自动安装、更新或切换这些模型。 +- 连接名称。 +- 完整 OpenAI 兼容 Embeddings Endpoint。 +- 模型名称。 +- 可选 API Key。 +- 编码预设和可选最大输入 Token。 -这段说明不能把任意 Qwen 或 BGE 模型标记为必然优于应用托管默认模型。真实质量以用户 -知识库评测为准。 +GoodBuddy 不根据 Endpoint、端口或模型名称猜测具体服务品牌。连接名称由用户决定， +GoodBuddy 只持续显示规范化 Endpoint 主机、实际维度和诊断结果。 -### 16.3 服务检测 +### 16.3 连接诊断 -用户点击“检测 Ollama”后，Main 可以读取： - -- 服务版本。 -- 已安装模型列表。 -- 所选模型的 Digest。 -- 模型能力和可用时的上下文元数据。 - -检测结果只用于当前草稿，不自动保存或改变模型。服务返回的自由文本经过长度限制和 -脱敏，不直接写入日志或 DOM。 +用户点击“测试”后，Main 发送固定、不含用户数据的真实 Embeddings 请求。诊断验证返回 +数量、维度、有限值、非零范数、Query/Document 编码和配置确定性。诊断结果只绑定当前 +连接草稿，不自动保存、切换或修改模型。服务返回的自由文本经过长度限制和脱敏，不直接 +写入日志或 DOM。 ### 16.4 Endpoint 安全 @@ -930,7 +976,7 @@ GoodBuddy 不负责： - HTTPS 自托管地址按 TLS 校验，不忽略证书错误。 - HTTP 非 loopback 地址显示明文传输警告，但不自动改写为其他地址。 - 不跟随重定向。 -- 不把 Ollama API 响应中的模型路径或环境信息暴露给 Renderer。 +- 不把服务响应中的内部路径、环境信息或未受限原文暴露给 Renderer。 --- @@ -940,10 +986,9 @@ GoodBuddy 不负责： | 路径 | 文档分块 | 查询 | 向量 | 凭据 | | --- | --- | --- | --- | --- | -| 应用托管 | Main → 本地 Utility Process | Main → 本地 Utility Process | 返回 Main，保存 SQLite | 不需要 | -| loopback Ollama | Main → 本机 Ollama | Main → 本机 Ollama | 返回 Main，保存 SQLite | 可选，Main-only | -| 网络自托管 | Main → 指定主机 | Main → 指定主机 | 返回 Main，保存 SQLite | Main-only | -| 云端 | Main → 指定云端 | Main → 指定云端 | 返回 Main，保存 SQLite | Main-only | +| 内置连接 | Main → 本地 Utility Process | Main → 本地 Utility Process | 返回 Main，保存 SQLite | 不需要 | +| Loopback 兼容连接 | Main → 本机服务 | Main → 本机服务 | 返回 Main，保存 SQLite | 可选，Main-only | +| 非 Loopback 兼容连接 | Main → 指定主机 | Main → 指定主机 | 返回 Main，保存 SQLite | 可选，Main-only | 默认不记录原始分块、查询正文、返回向量或 Authorization Header。诊断只保存有界统计和 脱敏错误。 @@ -962,7 +1007,7 @@ GoodBuddy 不负责： - 模型目录元数据进入源码审查。 - Revision、大小和 SHA-256 必须固定。 -- License、NOTICE 和上游来源随模型包导出。 +- 模型许可和上游来源随应用资源与目录元数据交付。 - 下载工件不执行脚本。 - ONNX 外部数据文件必须显式列入 Manifest。 - 模型加载前再次验证已安装 Manifest，不能只在首次下载时验证。 @@ -973,12 +1018,12 @@ GoodBuddy 不负责： | 场景 | 行为 | | --- | --- | -| 应用托管模型未安装 | 阻止启用，提供下载或 ZIP 导入 | -| 模型文件损坏 | 阻止加载，显示重新导入或删除，不连接 Ollama/云端 | +| 内置模型未安装 | 阻止启用，提供下载或 ZIP 导入 | +| 模型文件损坏 | 阻止加载，显示重新导入或删除，不启用用户连接 | | Utility Process 崩溃 | 同一 Snapshot 有限重启，耗尽后当前操作失败 | -| Ollama 未运行 | 显示服务不可用，保留配置，不启动或改连其他服务 | -| Ollama 模型不存在 | 显示准确模型名和用户处理建议，不自动选择已安装模型 | -| 云端认证失败 | 显示认证错误，保留配置，不尝试本地模型 | +| 用户 Endpoint 不可用 | 显示连接不可用，保留配置，不启动或改连其他服务 | +| 用户连接模型不存在 | 显示准确模型名和处理建议，不自动选择其他模型 | +| 用户连接认证失败 | 显示认证错误，保留配置，不尝试内置连接 | | Provider 限流或超时 | 同一 Provider 有界重试，之后明确失败 | | 返回维度变化 | 当前请求失败，索引标记不兼容，不截断或填充向量 | | 输入超过上下文 | 标记文档或查询失败，提示调整分块，不静默截断 | @@ -997,7 +1042,7 @@ GoodBuddy 不负责： ### 19.1 推理生命周期 - 本地推理进程按首次诊断、查询或索引任务惰性启动。 -- 同一时间只加载当前应用托管模型。 +- 同一时间只加载当前内置模型。 - 空闲达到有界时间后可以卸载同一模型，下一次使用重新加载。 - 卸载和重载不改变 Fingerprint。 - 查询优先于后台索引批次，避免知识重建长期阻塞交互。 @@ -1040,7 +1085,7 @@ GoodBuddy 不负责： 5. **平台评测**：六个平台的安装、加载、取消、退出和资源。 现有[知识检索评估](../quality/knowledge-retrieval-evaluation.md)中的确定性内存 Provider -用于验证生产检索管道，不代表真实模型质量。应用托管模型需要单独的可选本地评测命令， +用于验证生产检索管道，不代表真实模型质量。内置模型需要单独的可选本地评测命令， 该命令不读取用户数据库、不联网、不进入默认 CI。 ### 20.2 固定质量样例 @@ -1068,52 +1113,57 @@ GoodBuddy 不负责： ### 20.3 对比组 - FTS + CJK。 -- 应用托管 Granite 97M R2 INT8。 -- 用户显式配置时的 Ollama Qwen3-Embedding。 -- 用户显式配置时的 Ollama BGE。 -- 当前 OpenAI 兼容 Provider。 +- 内置 Granite 97M R2 INT8。 +- 用户显式配置的本机 OpenAI 兼容连接。 +- 用户显式配置的网络 OpenAI 兼容连接。 -Ollama 和云端组只在用户明确安装、配置并授权评测时运行。评测报告必须记录完整 -Fingerprint，不能只写模型显示名。 +用户连接组只在用户明确配置并授权评测时运行。评测报告必须记录完整 Fingerprint，不能 +只写连接名称或模型显示名。 --- -## 21. 实施阶段 +## 21. 实施状态 -### 阶段 1：契约和模型目录 +### 已完成：契约和模型目录 - 扩展角色感知 `EmbeddingProvider`。 -- 定义 Profile、Encoding Recipe、Runtime Snapshot 和完整 Fingerprint。 -- 新增应用托管模型目录、Manifest、下载、取消、ZIP 和删除。 +- 定义 Profile、Encoding Recipe、Embedding Operation Snapshot 和完整 Fingerprint。 +- 新增内置模型目录、Manifest、下载、取消、ZIP 和删除。 - 固定 Granite 97M R2 INT8 的 Revision、字节数、SHA-256、Tokenizer 和许可清单。 -### 阶段 2：本地推理 +### 已完成：本地推理 - 新增 Utility Process 和私有消息协议。 - 实现 Tokenizer、ONNX CPU、Pooling、归一化和动态 Token 批处理。 - 实现加载、诊断、取消、崩溃恢复和退出。 -- 完成六平台推理与打包门。 +- Windows x64 已完成真实推理；其余目标由六平台发布 Runner 继续验证。 -### 阶段 3：设置界面 +### 已完成：设置界面 -- 把现有向量表单重组为应用托管、Ollama/自托管、云端三种显式路径。 -- 新增模型详情、下载进度、数据路径、当前/草稿状态和真实诊断。 -- 新增模型切换确认和受影响知识库入口。 +- 把现有单一向量表单重组为与 LLM 一致的连接列表，加入只读内置连接和用户可管理的 + OpenAI 兼容连接。 +- 新增内置模型详情、下载进度、当前/草稿状态和真实诊断；内置连接不重复展示固定本地 + 路径。 - 更正“自动回退”文案。 -### 阶段 4：索引兼容与知识库 +### 已完成：Provider 与索引兼容 - 将完整 Fingerprint 接入索引存储和查询。 -- 冻结每个诊断、查询和重建任务的 Runtime Snapshot。 -- 显示当前模型、索引模型、兼容性和重建影响。 +- 冻结每个诊断、查询和重建任务的 Embedding Operation Snapshot。 - 验证模型切换、任务取消和文档级原子替换。 -### 阶段 5：Ollama 增强 +### 已完成：多连接管理 -- 为 Ollama 增加显式服务检测、模型列表和 Digest。 -- 提供 Qwen3-Embedding、BGE 等已知模型的可见编码预设建议。 -- 保持安装、拉取、更新和硬件管理由用户负责。 -- 增加用户授权的 Ollama 质量评测入口。 +- 支持添加、编辑、测试和删除多个 OpenAI 兼容连接。 +- 将现有单一 Endpoint、模型和凭据迁移为一个用户 Connection Profile。 +- 用户连接持续显示 Loopback 与非 Loopback 数据去向，不引入品牌专属检测。 + +### 后续知识增强阶段 + +- 在全局设置中显示模型切换的受影响知识库数量和确认。 +- 提供逐库或显式批量重建入口。 +- 在知识库界面集中显示当前模型、索引模型和完整兼容性诊断。 +- 增加用户授权的兼容连接质量评测入口。 --- @@ -1121,10 +1171,12 @@ Fingerprint，不能只写模型显示名。 ### 22.1 产品 -- 用户可以不安装外部服务，下载或导入应用托管 Granite 模型并完成真实诊断。 +- 用户可以不安装外部服务，下载或导入内置 Granite 模型并完成真实诊断。 - 未经用户操作不会下载、启用、切换或更新模型。 -- 用户可以自行安装 Ollama，并明确配置 Qwen3-Embedding、BGE 或其他 Embedding 模型。 -- 三种路径持续显示数据去向，非 loopback Ollama 不标记为“本机”。 +- 内置连接与用户添加的 OpenAI 兼容连接出现在同一连接列表，不存在第二组路径切换。 +- 向量连接不出现在项目 Runtime、Agent Runtime 或远程 Runtime 的模型选择中。 +- 切换 Agent Runtime 不改变向量连接，切换向量连接也不改变 Agent Runtime。 +- 用户连接持续显示数据去向，非 Loopback Endpoint 不标记为“本机”。 - 保存模型切换前显示当前模型、新模型和受影响知识库。 - 模型切换不会自动重建知识库。 - 向量失败时显示请求通道、实际通道和原因，不自动换模型。 @@ -1133,7 +1185,7 @@ Fingerprint，不能只写模型显示名。 ### 22.2 推理与索引 - Query 和 Document 使用显式角色。 -- 应用托管模型的 Tokenizer、Pooling、归一化、精度和维度全部来自固定 Manifest。 +- 内置模型的 Tokenizer、Pooling、归一化、精度和维度全部来自固定 Manifest。 - 超长输入不会被 GoodBuddy 静默截断。 - Fingerprint 任一兼容字段变化后，旧向量不参与当前召回。 - 384 维向量继续以 Float32 和有效模长保存。 @@ -1143,6 +1195,8 @@ Fingerprint，不能只写模型显示名。 ### 22.3 安全与交付 - Renderer 不读取模型目录、知识数据库或长期凭据。 +- Agent Runtime 只能获得有界检索文本、引用和脱敏诊断，不能获得 Provider、模型文件、 + Tokenizer、原始向量、Endpoint 或凭据。 - 模型下载、导入和启动都验证 Manifest、大小和 SHA-256。 - ZIP 导入拒绝路径穿越、符号链接、可执行文件和超限内容。 - Utility Process 不监听端口、不继承无关密钥，应用退出时正常终止。 @@ -1160,8 +1214,8 @@ npm run lint npm run build ``` -真实本地模型评测、Ollama 和云端调用使用独立门控命令，只有在模型已安装且用户明确授权时 -运行。 +真实本地模型评测和用户连接调用使用独立门控命令，只有在模型已安装或连接已配置且用户 +明确授权时运行。 --- @@ -1177,13 +1231,5 @@ npm run build - BGE Small Chinese v1.5 模型卡： -- Ollama Embeddings： - -- Ollama Embed API： - -- Ollama OpenAI 兼容 API： - -- Ollama Qwen3-Embedding： - -- Ollama BGE-M3： - +- ONNX Community Granite 97M Multilingual R2： + diff --git a/docs/architecture/model-download-source-design.md b/docs/architecture/model-download-source-design.md index 9b3897b9..e43c8e51 100644 --- a/docs/architecture/model-download-source-design.md +++ b/docs/architecture/model-download-source-design.md @@ -232,8 +232,7 @@ [ 通用设置 ] [ 魔法笔记 ] -┌ 本地模型 ──────────────────────────────────────────────┐ -│ 模型下载源 │ +┌ 模型下载源 ────────────────────────────────────────────┐ │ 选择 GoodBuddy 托管本地模型后续下载使用的平台。 │ │ 已安装模型和 ZIP 导入不受影响。 │ │ │ @@ -243,7 +242,6 @@ │ ○ Hugging Face │ │ 适合可以稳定访问 Hugging Face 的网络环境 │ │ │ -│ 当前选择：ModelScope │ └────────────────────────────────────────────────────────┘ ``` @@ -299,14 +297,10 @@ - 当前模型切换。 - 活动下载取消。 -### 7.5 活动下载说明 +### 7.5 活动下载 -来源选项下方持续说明： - -> 正在进行的模型下载会继续使用启动时的来源；新的下载使用当前选择。 - -首期不为这段说明新增跨语音、OCR 和向量管理器的活动任务聚合 IPC。具体活动操作继续在 -对应模型卡显示冻结来源、进度和取消入口。 +不常驻展示活动下载说明。具体下载操作在对应模型卡显示进度和取消入口；需要显示来源时 +必须使用任务启动时冻结的来源。切换全局来源不改变已经开始的操作。 --- @@ -651,12 +645,12 @@ resolveModelDownloadPackage(files, source) ### 12.1 来源状态 -语音、OCR、向量和后续模型卡统一显示： +语音和 OCR 的当前模型详情各保留一个全局下载源标签。向量内置连接不重复显示平台功能中 +已经明确的全局来源，也不显示固定本地安装路径；下载仍使用已保存的全局来源。 -- 当前全局下载源。 -- 当前模型在该来源是否可下载。 -- 活动操作的冻结来源。 -- 已安装模型的安装和校验状态。 +所有模型管理界面都显示当前模型是否可下载、活动操作的进度与取消入口，以及已安装模型 +的安装和校验状态。活动操作需要显示来源时，使用任务启动时冻结的来源，不能读取后来 +切换的全局值。 未安装且来源可用： diff --git a/docs/prd/knowledge/knowledge-rag-enhancement-prd.md b/docs/prd/knowledge/knowledge-rag-enhancement-prd.md index 5fa7b26e..f8ec65ac 100644 --- a/docs/prd/knowledge/knowledge-rag-enhancement-prd.md +++ b/docs/prd/knowledge/knowledge-rag-enhancement-prd.md @@ -5,8 +5,8 @@ | 项目 | 内容 | | --- | --- | | 状态 | 实施中 | -| 版本 | 0.1 | -| 日期 | 2026-08-11 | +| 版本 | 0.2 | +| 日期 | 2026-08-22 | | 适用产品 | GoodBuddy 桌面端 | | 实施范围 | 第一阶段：可用、可见、可诊断；第二阶段：可调、可优化、可维护 | | 相关设计 | [本地文本向量模型与连接设计](../../architecture/local-text-embedding-model-design.md) | @@ -40,7 +40,7 @@ OpenAI 兼容向量模型、RRF 混合检索、知识图谱、任务状态和来 预检索，不能只依赖提示词要求模型调用工具。 4. 知识库新建后不默认启用全部已有知识库；对话中的范围继续由用户显式选择。 5. 向量服务不可用时保留全文与中文检索，但必须返回明确降级状态。这是可见的检索通道 - 降级，不得自动切换应用托管模型、Ollama、云端 Provider 或其他向量模型。 + 降级，不得自动切换内置连接、用户添加的 OpenAI 兼容连接或其他向量模型。 6. 中文召回使用应用内可控的 CJK n-gram 索引，不新增远程服务依赖。 7. 混合检索保留 RRF 候选融合，并增加本地确定性重排、可选的 Cohere/Jina 兼容学习型重排、最低相关度和上下文预算。学习型重排失败时 @@ -60,6 +60,12 @@ OpenAI 兼容向量模型、RRF 混合检索、知识图谱、任务状态和来 用户配置决定。 14. Ask 的运行时边界保持只读。知识库内容始终被标记为不可信证据， 不能成为系统指令。 +15. 向量连接是 GoodBuddy Main 管理的知识服务依赖，不是 Agent Runtime 模型。它不进入 + 直连模型、OpenCode、Continue、DeepSeek Harness、远程 Runtime 或项目 Runtime 的 + 模型选择、配置和凭据边界。 +16. Runtime 只能获得 Main 检索后选出的有界文本证据、来源引用和脱敏诊断，不能获得 + Embedding Provider、连接 Profile、Endpoint、API Key、模型目录、Tokenizer、原始向量 + 或索引 Fingerprint。 ## 3. 目标 @@ -102,6 +108,9 @@ OpenAI 兼容向量模型、RRF 混合检索、知识图谱、任务状态和来 - 自动问题生成、FAQ 生成和训练数据标注平台。 - 完整 RAG 离线评测平台。第二阶段只提供手动检索测试与可导出的诊断信息。 - 在应用内高保真渲染所有原始 Office 和 PDF 文档。 +- 把向量连接用作 LLM、Agent Runtime 模型或 Runtime 模型网关。 +- 让 Runtime 直接调用 Embeddings Endpoint、内置推理进程或向量数据库。 +- 未经独立设计就把知识库向量连接扩展到记忆、自动化、项目文件搜索或其他应用服务。 ## 5. 竞品基线与 GoodBuddy 定位 @@ -160,10 +169,25 @@ OpenAI 兼容向量模型、RRF 混合检索、知识图谱、任务状态和来 “检索测试”是当前知识库的高频诊断操作，通过知识库标题区次操作打开独立 工作台，不新增第五个一级页签。 -全局向量模型仍在“设置 → 模型连接 → 向量模型”中配置。应用托管本地模型、 -用户自行安装的 Ollama/自托管服务和云端兼容服务的界面、数据路径及切换语义以 +全局向量模型仍在“设置 → 模型连接 → 向量模型”中配置。GoodBuddy 内置连接和用户添加的 +OpenAI 兼容连接使用同一连接列表；其界面、数据路径及切换语义以 [本地文本向量模型与连接设计](../../architecture/local-text-embedding-model-design.md) -为准，知识库页面只显示当前模型、索引兼容性、覆盖率和重建操作。 +为准，知识库页面只显示当前模型、索引兼容性、覆盖率和重建操作。“模型连接”是设置入口 +名称，不表示向量连接可以分配给 Agent Runtime；项目 Runtime 和 Agent Runtime 的模型 +选择中不显示向量连接。 + +知识检索与 Runtime 的关系固定为： + +```text +用户问题 +→ Main / KnowledgeService 使用当前向量连接检索 +→ Main 生成有界文本证据、引用和检索诊断 +→ 当前 Runtime 仅消费这些不可信检索结果 +``` + +`always` 模式由 Main 在启动 Runtime 前完成检索；`auto` 模式中的 `knowledge_search` +也只允许 Runtime 调用 Main 的有界知识工具。两种模式都不把向量连接、凭据、模型文件、 +Tokenizer、原始向量或索引访问能力传入 Runtime。切换 Runtime 不改变当前向量连接。 对话输入区的知识范围弹层包含： @@ -567,7 +591,12 @@ GoodBuddy 不上传私人检索查询或文档内容。本地诊断至少记录 - “每次先检索”在 Runtime 启动前产生检索诊断和引用，即使模型未调用工具。 - 未配置向量模型时，中文改写问题仍能通过 CJK 索引召回相关分块。 - 向量查询失败时回答可继续，界面明确显示已降级。 -- 应用托管模型、Ollama 和云端向量连接之间不会自动切换；实际数据路径持续可见。 +- 内置连接和用户添加的 OpenAI 兼容连接之间不会自动切换；用户连接的实际数据去向持续 + 可见，内置连接不重复展示固定本地路径。 +- 项目 Runtime 和 Agent Runtime 的模型选择不显示向量连接。 +- Runtime 只获得有界文本证据、引用和脱敏诊断，不获得向量 Provider、凭据或原始向量。 +- 切换直连模型、OpenCode、Continue、DeepSeek Harness 或远程 Runtime 不改变当前向量 + 连接。 - 10,000 个分块的向量测试能够返回正确 Top K，不出现固定上限空结果。 - 同名模型切换端点后，不会读取 Fingerprint 不匹配的旧向量。 - 重建失败时，上一版已就绪向量仍能继续召回。 diff --git a/electron.vite.config.ts b/electron.vite.config.ts index f5b36491..7034f70a 100644 --- a/electron.vite.config.ts +++ b/electron.vite.config.ts @@ -133,6 +133,17 @@ function rendererBundleModuleManifestPlugin(): Plugin { } } +export function stableMainEntryFileName(chunk: { + readonly name: string +}): string { + return [ + 'deepseek-harness-host-bootstrap', + 'embedding-inference-bootstrap' + ].includes(chunk.name) + ? `${chunk.name}.js` + : '[name].js' +} + export default defineConfig({ main: { plugins: [ @@ -177,6 +188,9 @@ export default defineConfig({ ), 'deepseek-harness-host-bootstrap': resolve( 'src/main/deepseek-harness-host-bootstrap.ts' + ), + 'embedding-inference-bootstrap': resolve( + 'src/main/embedding-inference-bootstrap.ts' ) }, external: [ @@ -186,9 +200,7 @@ export default defineConfig({ ], output: { entryFileNames(chunk) { - return chunk.name === 'deepseek-harness-host-bootstrap' - ? 'deepseek-harness-host-bootstrap.js' - : '[name].js' + return stableMainEntryFileName(chunk) }, chunkFileNames(chunk) { const moduleIds = chunk.moduleIds.join('\n') diff --git a/package-lock.json b/package-lock.json index 8976b0d7..7c44c319 100644 --- a/package-lock.json +++ b/package-lock.json @@ -34,6 +34,7 @@ "@deepseek-ai/dsh-tool-skill": "0.1.0-rc.6", "@deepseek-ai/dsh-tools": "0.1.0-rc.6", "@deepseek-ai/dsh-user-approval": "0.1.0-rc.6", + "@huggingface/tokenizers": "^0.1.3", "@modelcontextprotocol/sdk": "^1.30.0", "@napi-rs/canvas": "1.0.3", "@opencode-ai/sdk": "^1.18.9", @@ -3087,6 +3088,12 @@ "hono": "^4" } }, + "node_modules/@huggingface/tokenizers": { + "version": "0.1.3", + "resolved": "https://registry.npmjs.org/@huggingface/tokenizers/-/tokenizers-0.1.3.tgz", + "integrity": "sha512-8rF/RRT10u+kn7YuUbUg0OF30K8rjTc78aHpxT+qJ1uWSqxT1MHi8+9ltwYfkFYJzT/oS+qw3JVfHtNMGAdqyA==", + "license": "Apache-2.0" + }, "node_modules/@humanfs/core": { "version": "0.19.2", "resolved": "https://registry.npmjs.org/@humanfs/core/-/core-0.19.2.tgz", diff --git a/package.json b/package.json index d04fed77..a3953259 100644 --- a/package.json +++ b/package.json @@ -54,6 +54,7 @@ "asarUnpack": [ "out/main/package.json", "out/main/deepseek-harness-*", + "out/main/embedding-inference-bootstrap.js", "out/main/chunks/**/*", "node_modules/node-pty/lib/**/*", "node_modules/node-pty/package.json", @@ -62,7 +63,8 @@ "node_modules/koffi/**/*", "node_modules/@koromix/koffi-*/**/*", "node_modules/@napi-rs/canvas{,/**/*}", - "node_modules/@napi-rs/canvas-*/**/*" + "node_modules/@napi-rs/canvas-*/**/*", + "node_modules/onnxruntime-web/dist/**/*.wasm" ], "npmRebuild": false, "compression": "maximum", @@ -172,6 +174,10 @@ "from": "node_modules/onnxruntime-web/LICENSE", "to": "licenses/onnxruntime-web-MIT.txt" }, + { + "from": "node_modules/@huggingface/tokenizers/LICENSE", + "to": "licenses/huggingface-tokenizers-Apache-2.0.txt" + }, { "from": "node_modules/katex/LICENSE", "to": "licenses/katex-MIT.txt" @@ -254,6 +260,7 @@ "@deepseek-ai/dsh-tool-skill": "0.1.0-rc.6", "@deepseek-ai/dsh-tools": "0.1.0-rc.6", "@deepseek-ai/dsh-user-approval": "0.1.0-rc.6", + "@huggingface/tokenizers": "^0.1.3", "@modelcontextprotocol/sdk": "^1.30.0", "@napi-rs/canvas": "1.0.3", "@opencode-ai/sdk": "^1.18.9", diff --git a/src/main/embedding-inference-bootstrap.ts b/src/main/embedding-inference-bootstrap.ts new file mode 100644 index 00000000..2254c8d3 --- /dev/null +++ b/src/main/embedding-inference-bootstrap.ts @@ -0,0 +1,33 @@ +import path from 'node:path' +import { GraniteEmbeddingEngine } from './knowledge/granite-embedding-engine' +import { + createEmbeddingInferenceWorkerPort, + EmbeddingInferenceWorker +} from './knowledge/embedding-inference-worker' + +const modelDirectory = process.argv[2] +if ( + typeof modelDirectory !== 'string' || + modelDirectory.length < 1 || + modelDirectory.length > 4_096 || + !path.isAbsolute(modelDirectory) +) { + throw new TypeError('Expected a bounded absolute model directory in argv') +} + +const worker = new EmbeddingInferenceWorker({ + modelDirectory, + engine: new GraniteEmbeddingEngine(), + port: createEmbeddingInferenceWorkerPort(process.parentPort) +}) + +async function shutdown(): Promise { + await worker.shutdown() +} + +process.once('disconnect', () => { + void shutdown() +}) +process.once('SIGTERM', () => { + void shutdown() +}) diff --git a/src/main/index.ts b/src/main/index.ts index 46e7758e..1fd21a99 100644 --- a/src/main/index.ts +++ b/src/main/index.ts @@ -37,6 +37,13 @@ import { KnowledgeService } from './knowledge/knowledge-service' import { AssistantDatabase } from './assistant/assistant-database' import { createModelGraphExtractor } from './knowledge/model-extractor' import { OpenAIEmbeddingClient } from './knowledge/openai-embedding-client' +import { embeddingProviderFingerprint } from './knowledge/embedding-provider-key' +import type { EmbeddingProvider } from './knowledge/types' +import { EmbeddingModelManager } from './knowledge/embedding-model-manager' +import { + EmbeddingInferenceBroker, + type EmbeddingInferenceTransport +} from './knowledge/embedding-inference-broker' import { CohereRerankClient } from './knowledge/cohere-rerank-client' import { RuntimeSettingsStore } from './runtime-settings-store' import type { ResolvedRuntimeSettings } from './runtime-settings-store' @@ -101,6 +108,7 @@ import { ShortcutSettingsService } from './shortcut-settings-service' import { defaultGlobalShortcutSettings } from '../shared/shortcut' import { requestProcessTreeTermination } from './agent/child-process-termination' import { createContinueUtilityProcessChild } from './agent/continue-utility-process-adapter' +import type { RuntimeSettings, RuntimeSettingsInput } from '../shared/contracts' const legacyDefaultShortcut = defaultGlobalShortcutSettings.accelerator @@ -145,19 +153,220 @@ let browserService: BrowserService | undefined let globalTlsPolicy: GlobalTlsPolicy | undefined let documentOcrBroker: DocumentOcrBroker | undefined let documentOcrModelManager: DocumentOcrModelManager | undefined +let embeddingModelManager: EmbeddingModelManager | undefined +const embeddingBrokers = new Set() let stopRuntimeReconfiguration: (() => Promise) | undefined let dshExtensionInstaller: DshNpmExtensionInstaller | undefined -function createEmbeddingProvider( - settings: ResolvedRuntimeSettings -): OpenAIEmbeddingClient | undefined { - return settings.knowledgeEmbeddingEnabled - ? new OpenAIEmbeddingClient({ - endpoint: settings.knowledgeEmbeddingBaseUrl, - model: settings.knowledgeEmbeddingModel, - apiKey: settings.knowledgeEmbeddingApiKey - }) - : undefined +type ManagedEmbeddingProvider = EmbeddingProvider & { + dispose?: () => void | Promise +} + +function createEmbeddingUtilityTransport( + modulePath: string, + modelDirectory: string +): EmbeddingInferenceTransport { + const child = utilityProcess.fork( + modulePath, + [modelDirectory], + { + serviceName: 'GoodBuddy Embedding Inference', + stdio: 'pipe', + allowLoadingUnsignedLibraries: false + } + ) + return { + postMessage: (message) => child.postMessage(message), + onMessage: (listener) => { + child.on('message', listener) + return () => child.removeListener('message', listener) + }, + onClose: (listener) => { + let closed = false + const notify = (): void => { + if (closed) { + return + } + closed = true + listener() + } + child.on('exit', notify) + child.on('error', notify) + return () => { + child.removeListener('exit', notify) + child.removeListener('error', notify) + } + }, + close: () => { + child.kill() + } + } +} + +async function createEmbeddingProvider( + settings: ResolvedRuntimeSettings, + manager: EmbeddingModelManager, + connectionId = settings.activeEmbeddingConnectionId +): Promise { + if (!settings.knowledgeEmbeddingEnabled) { + return undefined + } + const connection = settings.embeddingConnections?.find( + (candidate) => candidate.id === connectionId + ) + if (!connection) { + throw new Error('当前向量连接不存在') + } + if (connection.kind !== 'builtin') { + return new OpenAIEmbeddingClient({ + endpoint: connection.baseUrl, + model: connection.modelName, + apiKey: connection.apiKey + }) + } + const snapshot = await manager.getSnapshot() + const model = + snapshot.catalog.find((candidate) => candidate.recommended) ?? + snapshot.catalog[0] + if (!model) { + throw new Error('内置向量模型目录为空') + } + const modelDirectory = + await manager.getVerifiedModelDirectory(model.id) + const installed = snapshot.installed.find( + (candidate) => candidate.id === model.id + ) + if (!installed) { + throw new Error('内置向量模型尚未安装') + } + const tokenizerDigest = installed.files.find( + (file) => file.role === 'tokenizer' + )?.sha256 + const modelDigest = installed.files.find( + (file) => file.role === 'model' + )?.sha256 + if (!tokenizerDigest || !modelDigest) { + throw new Error('内置向量模型安装清单不完整') + } + const broker = new EmbeddingInferenceBroker({ + createTransport: () => + createEmbeddingUtilityTransport( + join(mainModuleDirectory, 'embedding-inference-bootstrap.js'), + modelDirectory + ) + }) + embeddingBrokers.add(broker) + return { + provider: 'builtin', + model: model.id, + fingerprint: embeddingProviderFingerprint({ + provider: 'builtin', + dataPath: { kind: 'device' }, + model: model.id, + dimensions: model.dimensions, + encodingRecipe: { + recipeId: 'granite-embedding-97m-r2', + artifactDigest: modelDigest, + tokenizerDigest, + pooling: 'cls', + normalization: 'l2', + queryTemplate: '{text}', + documentTemplate: '{text}', + maximumSequenceTokens: model.contextTokens + } + }), + embed: (texts, signal) => + broker.embed(texts, 'document', signal), + embedQuery: (texts, signal) => + broker.embed(texts, 'query', signal), + embedDocuments: (texts, signal) => + broker.embed(texts, 'document', signal), + dispose: async () => { + try { + await broker.shutdown() + } finally { + embeddingBrokers.delete(broker) + } + } + } +} + +function runtimeSettingsInputWithEmbeddingSelection( + settings: RuntimeSettings, + activeEmbeddingConnectionId: string +): RuntimeSettingsInput { + const configured = settings.configured ?? settings + const defaultProfile = + settings.modelProfiles.find( + (profile) => profile.id === settings.defaultModelProfileId + ) ?? settings.modelProfiles[0] + if (!defaultProfile || !settings.embeddingConnections) { + throw new Error('模型连接设置不完整') + } + return { + provider: settings.provider, + modelBaseUrl: defaultProfile.baseUrl, + modelName: defaultProfile.modelName, + modelProtocol: defaultProfile.protocol, + modelAuthentication: defaultProfile.authentication, + imageGenerationQuality: defaultProfile.imageGenerationQuality, + opencodeBaseUrl: configured.opencodeBaseUrl, + opencodeEmbedded: settings.opencodeEmbedded, + opencodeBinaryPath: configured.opencodeBinaryPath, + opencodeConfigPath: configured.opencodeConfigPath, + continueBinaryPath: configured.continueBinaryPath, + continueConfigPath: configured.continueConfigPath, + continueMode: settings.continueMode, + subagentSmartRoutingEnabled: settings.subagentSmartRoutingEnabled, + knowledgeEmbeddingEnabled: settings.knowledgeEmbeddingEnabled, + knowledgeEmbeddingBaseUrl: settings.knowledgeEmbeddingBaseUrl, + knowledgeEmbeddingModel: settings.knowledgeEmbeddingModel, + knowledgeEmbeddingApiKey: { action: 'keep' }, + embeddingConnections: settings.embeddingConnections.map( + (connection) => + connection.kind === 'builtin' + ? { + id: connection.id, + name: connection.name, + kind: connection.kind + } + : { + id: connection.id, + name: connection.name, + kind: connection.kind, + baseUrl: connection.baseUrl, + modelName: connection.modelName, + authentication: connection.authentication, + apiKey: { action: 'keep' as const } + } + ), + activeEmbeddingConnectionId, + knowledgeRerankEnabled: settings.knowledgeRerankEnabled ?? false, + knowledgeRerankEndpoint: settings.knowledgeRerankEndpoint ?? '', + knowledgeRerankModel: settings.knowledgeRerankModel ?? '', + knowledgeRerankApiKey: { action: 'keep' }, + contextCompression: settings.contextCompression, + workspacePath: configured.workspacePath, + apiKey: { action: 'keep' }, + modelProfiles: settings.modelProfiles.map((profile) => ({ + id: profile.id, + name: profile.name, + baseUrl: profile.baseUrl, + modelName: profile.modelName, + protocol: profile.protocol, + authentication: profile.authentication, + supportsImageInput: profile.supportsImageInput, + contextWindowTokens: profile.contextWindowTokens, + imageGenerationQuality: profile.imageGenerationQuality, + apiKey: { action: 'keep' as const } + })), + defaultModelProfileId: settings.defaultModelProfileId, + opencodeModelSource: configured.opencodeModelSource, + continueModelSource: configured.continueModelSource, + deepseekHarnessModelSource: + configured.deepseekHarnessModelSource ?? { kind: 'platform' }, + toolApproval: settings.toolApproval + } } function createRerankProvider( @@ -410,6 +619,13 @@ if (hasSingleInstanceLock) { const applicationSettingsStore = new ApplicationSettingsStore( join(app.getPath('userData'), 'application-settings.json') ) + const startupEmbeddingModelManager = new EmbeddingModelManager({ + userDataDirectory: app.getPath('userData'), + fetch: globalThis.fetch, + getDownloadSource: async () => + (await applicationSettingsStore.get()).modelDownloadSource + }) + embeddingModelManager = startupEmbeddingModelManager const releaseNotesService = new ReleaseNotesService({ currentVersion: app.getVersion(), filePath: app.isPackaged @@ -505,7 +721,7 @@ if (hasSingleInstanceLock) { }) knowledgeService = startupKnowledgeService let activeEmbeddingProvider: - | ReturnType + | Awaited> | undefined let activeRerankProvider: | ReturnType @@ -621,8 +837,9 @@ if (hasSingleInstanceLock) { }, initializeKnowledgeAndGateway: async () => { await startupKnowledgeService.initialize() - const embeddingProvider = createEmbeddingProvider( - initialResolvedSettings + const embeddingProvider = await createEmbeddingProvider( + initialResolvedSettings, + startupEmbeddingModelManager ) const rerankProvider = createRerankProvider( initialResolvedSettings @@ -706,7 +923,10 @@ if (hasSingleInstanceLock) { } const settings = await settingsStore.getResolvedSettings() const nextEmbeddingProvider = - createEmbeddingProvider(settings) + await createEmbeddingProvider( + settings, + startupEmbeddingModelManager + ) const nextRerankProvider = createRerankProvider(settings) let nextRuntime: AgentRuntime | undefined let nextSubagentRuntime: AgentRuntime | undefined @@ -760,8 +980,10 @@ if (hasSingleInstanceLock) { nextSubagentProfileRuntimes ) await selectedRuntimeManager?.reset() + const previousEmbeddingProvider = activeEmbeddingProvider activeEmbeddingProvider = nextEmbeddingProvider activeRerankProvider = nextRerankProvider + await previousEmbeddingProvider?.dispose?.() } catch (activationError) { const rollbackResults = knowledgeService ? await Promise.allSettled([ @@ -774,6 +996,7 @@ if (hasSingleInstanceLock) { ]) : [] await Promise.allSettled([ + nextEmbeddingProvider?.dispose?.(), runtimeConsumed ? undefined : nextRuntime?.dispose(), subagentRuntimesConsumed ? undefined @@ -804,6 +1027,41 @@ if (hasSingleInstanceLock) { runtimeReconfigurationClosing = true await runtimeReconfigurationQueue } + const setCurrentEmbeddingConnection = async ( + connectionId: string + ): Promise => { + const previous = await settingsStore.getPublicSettings() + if (previous.activeEmbeddingConnectionId === connectionId) { + return + } + await settingsStore.update( + runtimeSettingsInputWithEmbeddingSelection( + previous, + connectionId + ) + ) + try { + await reconfigureRuntimes() + } catch (activationError) { + try { + await settingsStore.update( + runtimeSettingsInputWithEmbeddingSelection( + previous, + previous.activeEmbeddingConnectionId ?? + connectionId + ) + ) + await reconfigureRuntimes() + } catch (rollbackError) { + throw new AggregateError( + [activationError, rollbackError], + '向量模型连接激活失败，且回滚未能完成', + { cause: rollbackError } + ) + } + throw activationError + } + } removeIpcHandlers = registerIpcHandlers( mainWindow, @@ -837,7 +1095,21 @@ if (hasSingleInstanceLock) { releaseNotesService, goodbuddyConfigService, runtimeExtensionStore, - shortcutSettingsService + shortcutSettingsService, + startupEmbeddingModelManager, + async (connectionId) => { + const settings = await settingsStore.getResolvedSettings() + const provider = await createEmbeddingProvider( + settings, + startupEmbeddingModelManager, + connectionId + ) + if (!provider) { + throw new Error('请先启用并保存向量模型设置') + } + return provider + }, + setCurrentEmbeddingConnection ) loadMainWindow(mainWindow) setImmediate(() => { @@ -927,6 +1199,11 @@ app.on('before-quit', (event) => { () => selectedRuntimeManager?.dispose(), () => browserService?.dispose(), () => globalTlsPolicy?.dispose(), + () => + Promise.allSettled( + [...embeddingBrokers].map((broker) => broker.shutdown()) + ).then(() => undefined), + () => embeddingModelManager?.dispose(), () => documentOcrModelManager?.dispose(), () => documentOcrBroker?.dispose() ], diff --git a/src/main/ipc.test.ts b/src/main/ipc.test.ts index 0cd10fd2..67d71acb 100644 --- a/src/main/ipc.test.ts +++ b/src/main/ipc.test.ts @@ -1,5 +1,11 @@ import { afterEach, describe, expect, it, vi } from 'vitest' -import { mkdtemp, realpath, rm, writeFile } from 'node:fs/promises' +import { + mkdtemp, + readFile, + realpath, + rm, + writeFile +} from 'node:fs/promises' import { tmpdir } from 'node:os' import { join } from 'node:path' import { ipcChannels } from '../shared/ipc-channels' @@ -12,6 +18,30 @@ import { defaultKnowledgeOntologySettings } from '../shared/knowledge-ontology' import { AssistantDatabase } from './assistant/assistant-database' import { registerIpcHandlers } from './ipc' +describe('embedding IPC boundary', () => { + it('validates explicit IDs and keeps ZIP paths in the native dialog', async () => { + const source = await readFile( + join(process.cwd(), 'src', 'main', 'ipc.ts'), + 'utf8' + ) + const region = + source.match( + /const requireEmbeddingProvider[\s\S]*?registerHandler\(ipcChannels\.speechModelsGet/u + )?.[0] ?? '' + expect(region).toContain('embeddingConnectionIdRequestSchema.parse(input)') + expect(region).toContain('embeddingModelInstallInputSchema.parse(input)') + expect(region).toContain('embeddingModelActionInputSchema.parse(input)') + expect(region).toContain("properties: ['openFile']") + expect(region).toContain('filters: modelArchiveDialogFilters') + expect(region).not.toMatch( + /embeddingModelsImportArchive[\s\S]{0,300}input\.(?:path|archivePath|directory)/u + ) + expect(region).not.toMatch( + /(?:agent|selectedRuntimes|remoteModelGateway).{0,80}(?:embedding|vector)/iu + ) + }) +}) + type InvokeHandler = (event: unknown, input?: unknown) => unknown type KnowledgeGrantMock = ( requestId: string, diff --git a/src/main/ipc.ts b/src/main/ipc.ts index ed508d64..d40e3d14 100644 --- a/src/main/ipc.ts +++ b/src/main/ipc.ts @@ -119,6 +119,10 @@ import { speechModelSelectionInputSchema } from '../shared/speech-model-contracts' import { + embeddingConnectionIdRequestSchema, + embeddingModelActionInputSchema, + embeddingModelInstallInputSchema, + embeddingModelSnapshotSchema, embeddingSettingsSnapshotSchema, knowledgeEmbeddingIndexCancelRequestSchema, knowledgeEmbeddingIndexRequestSchema, @@ -256,6 +260,8 @@ import type { SpeechModelManager } from './speech/speech-model-manager' import type { SpeechTranscriptionService } from './speech/speech-transcription-service' import { diagnoseEmbeddingProvider } from './knowledge/embedding-index-coordinator' import type { EmbeddingIndexCoordinator } from './knowledge/embedding-index-coordinator' +import type { EmbeddingModelManager } from './knowledge/embedding-model-manager' +import type { EmbeddingProvider } from './knowledge/types' import type { DocumentParsingService } from './document-parsing-service' import type { DocumentOcrModelManager } from './document-ocr-model-manager' import type { DocumentOcrBroker } from './document-ocr-broker' @@ -264,7 +270,6 @@ import type { GoodBuddyConfigApplyEvent, GoodBuddyConfigService } from './goodbuddy-config-service' -import { OpenAIEmbeddingClient } from './knowledge/openai-embedding-client' import { magicNotePlainText, validateMagicNoteRichContent @@ -915,7 +920,19 @@ export function registerIpcHandlers( releaseNotesService?: ReleaseNotesService, goodbuddyConfigService?: GoodBuddyConfigService, runtimeExtensionStore?: RuntimeExtensionStore, - shortcutSettingsService?: ShortcutSettingsService + shortcutSettingsService?: ShortcutSettingsService, + embeddingModelManager?: EmbeddingModelManager & { + importArchive?( + modelId: string, + archivePath: string + ): Promise + }, + resolveEmbeddingProvider?: ( + connectionId: string + ) => Promise, + setCurrentEmbeddingConnection?: ( + connectionId: string + ) => Promise ): () => Promise { type ActiveRequestLease = { controller: AbortController @@ -4339,38 +4356,209 @@ export function registerIpcHandlers( } ) - const requireEmbeddingProvider = async (): Promise => { + const requireEmbeddingProvider = async ( + connectionId: string + ): Promise => { const settings = await settingsStore.getResolvedSettings() if (!settings.knowledgeEmbeddingEnabled) { throw new Error('请先启用并保存向量模型设置') } - return new OpenAIEmbeddingClient({ - endpoint: settings.knowledgeEmbeddingBaseUrl, - model: settings.knowledgeEmbeddingModel, - apiKey: settings.knowledgeEmbeddingApiKey + if ( + !settings.embeddingConnections?.some( + (connection) => connection.id === connectionId + ) + ) { + throw new Error('当前向量连接不存在') + } + if (!resolveEmbeddingProvider) { + throw new Error('向量模型服务不可用') + } + return resolveEmbeddingProvider(connectionId) + } + + const getEmbeddingSettingsSnapshot = async () => { + if (!embeddingModelManager) { + throw new Error('内置向量模型服务不可用') + } + const settings = await settingsStore.getPublicSettings() + const models = await embeddingModelManager.getSnapshot() + const connection = settings.embeddingConnections?.find( + (candidate) => + candidate.id === settings.activeEmbeddingConnectionId + ) + if (!connection || !settings.activeEmbeddingConnectionId) { + throw new Error('当前向量连接不存在') + } + const builtinModel = + models.catalog.find((model) => model.recommended) ?? + models.catalog[0] + if (!builtinModel) { + throw new Error('内置向量模型目录为空') + } + return embeddingSettingsSnapshotSchema.parse({ + configuration: + connection?.kind === 'builtin' + ? { + provider: 'builtin', + model: builtinModel.id, + credentialConfigured: false + } + : { + provider: 'openai-compatible', + model: + connection?.modelName ?? + settings.knowledgeEmbeddingModel, + endpoint: + connection?.baseUrl ?? + settings.knowledgeEmbeddingBaseUrl, + credentialConfigured: + connection?.apiKeyConfigured ?? + settings.knowledgeEmbeddingApiKeyConfigured + }, + connections: (settings.embeddingConnections ?? []).map( + (candidate) => + candidate.kind === 'builtin' + ? { + id: candidate.id, + name: candidate.name, + kind: candidate.kind, + model: builtinModel.id, + credentialConfigured: false + } + : { + id: candidate.id, + name: candidate.name, + kind: candidate.kind, + model: candidate.modelName, + endpoint: candidate.baseUrl, + authentication: candidate.authentication, + credentialConfigured: + candidate.apiKeyConfigured ?? false + } + ), + currentConnectionId: settings.activeEmbeddingConnectionId, + models }) } registerHandler(ipcChannels.embeddingSettingsGet, async (event) => { assertTrustedSender(event, window) - const settings = await settingsStore.getPublicSettings() - return embeddingSettingsSnapshotSchema.parse({ - configuration: { - provider: 'openai-compatible', - model: settings.knowledgeEmbeddingModel, - endpoint: settings.knowledgeEmbeddingBaseUrl, - credentialConfigured: - settings.knowledgeEmbeddingApiKeyConfigured - } - }) + return getEmbeddingSettingsSnapshot() }) - registerHandler(ipcChannels.embeddingDiagnose, async (event) => { - assertTrustedSender(event, window) - return diagnoseEmbeddingProvider( - await requireEmbeddingProvider() - ) - }) + registerHandler( + ipcChannels.embeddingDiagnose, + async (event, input: unknown) => { + assertTrustedSender(event, window) + const { connectionId } = + embeddingConnectionIdRequestSchema.parse(input) + const provider = await requireEmbeddingProvider(connectionId) + try { + return await diagnoseEmbeddingProvider(provider) + } finally { + await ( + provider as EmbeddingProvider & { + dispose?: () => void | Promise + } + ).dispose?.() + } + } + ) + + registerHandler( + ipcChannels.embeddingSetCurrent, + async (event, input: unknown) => { + assertTrustedSender(event, window) + if (!setCurrentEmbeddingConnection) { + throw new Error('向量模型设置服务不可用') + } + const { connectionId } = + embeddingConnectionIdRequestSchema.parse(input) + await setCurrentEmbeddingConnection(connectionId) + return getEmbeddingSettingsSnapshot() + } + ) + + registerHandler( + ipcChannels.embeddingModelsInstall, + async (event, input: unknown) => { + assertTrustedSender(event, window) + if (!embeddingModelManager || !applicationSettingsStore) { + throw new Error('内置向量模型服务不可用') + } + const { modelId, expectedDownloadSource } = + embeddingModelInstallInputSchema.parse(input) + const { modelDownloadSource } = + await applicationSettingsStore.get() + if (modelDownloadSource !== expectedDownloadSource) { + throw new Error('模型下载源已变化，请刷新后重试') + } + await embeddingModelManager.install( + modelId, + expectedDownloadSource + ) + return embeddingModelSnapshotSchema.parse( + await embeddingModelManager.getSnapshot() + ) + } + ) + + registerHandler( + ipcChannels.embeddingModelsCancel, + (event, input: unknown) => { + assertTrustedSender(event, window) + if (!embeddingModelManager) { + throw new Error('内置向量模型服务不可用') + } + const { modelId } = + embeddingModelActionInputSchema.parse(input) + return embeddingModelManager.cancel(modelId) + } + ) + + registerHandler( + ipcChannels.embeddingModelsImportArchive, + async (event, input: unknown) => { + assertTrustedSender(event, window) + if (!embeddingModelManager?.importArchive) { + throw new Error('当前版本不支持导入向量模型 ZIP') + } + const { modelId } = + embeddingModelActionInputSchema.parse(input) + const result = await dialog.showOpenDialog(window, { + title: '导入向量模型 ZIP', + properties: ['openFile'], + filters: modelArchiveDialogFilters + }) + const archivePath = result.filePaths[0] + if (result.canceled || !archivePath) { + return undefined + } + await embeddingModelManager.importArchive( + modelId, + archivePath + ) + return embeddingModelSnapshotSchema.parse( + await embeddingModelManager.getSnapshot() + ) + } + ) + + registerHandler( + ipcChannels.embeddingModelsRemove, + async (event, input: unknown) => { + assertTrustedSender(event, window) + if (!embeddingModelManager) { + throw new Error('内置向量模型服务不可用') + } + const { modelId } = + embeddingModelActionInputSchema.parse(input) + await embeddingModelManager.remove(modelId) + return embeddingModelSnapshotSchema.parse( + await embeddingModelManager.getSnapshot() + ) + } + ) registerHandler(ipcChannels.speechModelsGet, (event) => { assertTrustedSender(event, window) @@ -6091,13 +6279,28 @@ export function registerIpcHandlers( const embeddingConfiguration = async () => { const settings = await settingsStore.getPublicSettings() - return { - provider: 'openai-compatible', - model: settings.knowledgeEmbeddingModel, - endpoint: settings.knowledgeEmbeddingBaseUrl, - credentialConfigured: - settings.knowledgeEmbeddingApiKeyConfigured - } + const connection = settings.embeddingConnections?.find( + (candidate) => + candidate.id === settings.activeEmbeddingConnectionId + ) + return connection?.kind === 'builtin' + ? { + provider: 'builtin', + model: 'granite-embedding-97m-multilingual-r2', + credentialConfigured: false + } + : { + provider: 'openai-compatible', + model: + connection?.modelName ?? + settings.knowledgeEmbeddingModel, + endpoint: + connection?.baseUrl ?? + settings.knowledgeEmbeddingBaseUrl, + credentialConfigured: + connection?.apiKeyConfigured ?? + settings.knowledgeEmbeddingApiKeyConfigured + } } registerHandler( diff --git a/src/main/knowledge/embedding-index-coordinator.test.ts b/src/main/knowledge/embedding-index-coordinator.test.ts index ed032469..90fd2704 100644 --- a/src/main/knowledge/embedding-index-coordinator.test.ts +++ b/src/main/knowledge/embedding-index-coordinator.test.ts @@ -156,6 +156,59 @@ describe('EmbeddingIndexCoordinator', () => { ) }) + it('uses explicit query and document provider entry points when available', async () => { + const repository = new MemoryRepository() + const embed = vi.fn(async () => [[9, 9]]) + const embedQuery = vi.fn(async () => [[1, 2]]) + const embedDocuments = vi.fn(async (input: readonly string[]) => + input.map(() => [3, 4]) + ) + const roleAwareProvider = { + ...provider(embed), + embedQuery, + embedDocuments + } + const coordinator = new EmbeddingIndexCoordinator(repository, { + createId: () => 'job-role-aware' + }) + + await coordinator.diagnose(roleAwareProvider) + coordinator.startRebuild(roleAwareProvider) + await coordinator.waitForCompletion() + + expect(embedQuery).toHaveBeenCalledWith( + ['GoodBuddy 向量模型连接测试'], + undefined + ) + expect(embedDocuments).toHaveBeenCalledWith( + ['GoodBuddy 知识文档向量连接测试'], + undefined + ) + expect(embedDocuments).toHaveBeenCalledWith(['alpha'], expect.any(AbortSignal)) + expect(embedDocuments).toHaveBeenCalledWith( + ['beta', 'gamma'], + expect.any(AbortSignal) + ) + expect(embed).not.toHaveBeenCalled() + }) + + it('rejects diagnostics whose query and document dimensions differ', async () => { + const repository = new MemoryRepository() + const coordinator = new EmbeddingIndexCoordinator(repository) + const roleAwareProvider = { + ...provider(async () => [[9, 9]]), + embedQuery: vi.fn(async () => [[1, 2]]), + embedDocuments: vi.fn(async () => [[1, 2, 3]]) + } + + await expect(coordinator.diagnose(roleAwareProvider)).resolves.toMatchObject({ + status: 'unavailable', + error: { + code: 'invalid_response' + } + }) + }) + it('reports a safe diagnostic failure instead of treating config as success', async () => { const repository = new MemoryRepository() const coordinator = new EmbeddingIndexCoordinator(repository, { diff --git a/src/main/knowledge/embedding-index-coordinator.ts b/src/main/knowledge/embedding-index-coordinator.ts index 4674d8bc..36547623 100644 --- a/src/main/knowledge/embedding-index-coordinator.ts +++ b/src/main/knowledge/embedding-index-coordinator.ts @@ -19,6 +19,14 @@ export interface EmbeddingIndexProvider { readonly model: string readonly fingerprint?: string embed(input: readonly string[], signal?: AbortSignal): Promise + embedQuery?( + input: readonly string[], + signal?: AbortSignal + ): Promise + embedDocuments?( + input: readonly string[], + signal?: AbortSignal + ): Promise } export interface EmbeddingIndexItem { @@ -161,11 +169,13 @@ export async function diagnoseEmbeddingProvider( const now = options.now ?? Date.now const startedAt = now() try { - const vectors = await provider.embed( - [options.probeText ?? 'GoodBuddy 向量模型连接测试'], - options.signal - ) - if (vectors.length !== 1 || !vectors[0]) { + const queryDiagnosticInput = [ + options.probeText ?? 'GoodBuddy 向量模型连接测试' + ] + const queryVectors = provider.embedQuery + ? await provider.embedQuery(queryDiagnosticInput, options.signal) + : await provider.embed(queryDiagnosticInput, options.signal) + if (queryVectors.length !== 1 || !queryVectors[0]) { throw new EmbeddingOperationError({ code: 'invalid_response', message: '向量服务返回了无效结果。', @@ -173,7 +183,27 @@ export async function diagnoseEmbeddingProvider( remedy: '请确认服务为每个输入返回一个有效向量。' }) } - const dimensions = validateVector(vectors[0]) + const dimensions = validateVector(queryVectors[0]) + if (provider.embedQuery || provider.embedDocuments) { + const documentVectors = provider.embedDocuments + ? await provider.embedDocuments( + ['GoodBuddy 知识文档向量连接测试'], + options.signal + ) + : await provider.embed( + ['GoodBuddy 知识文档向量连接测试'], + options.signal + ) + if (documentVectors.length !== 1 || !documentVectors[0]) { + throw new EmbeddingOperationError({ + code: 'invalid_response', + message: '向量服务返回了无效结果。', + retryable: false, + remedy: '请确认服务为每个输入返回一个有效向量。' + }) + } + validateVector(documentVectors[0], dimensions) + } const checkedAt = now() return { status: 'available', @@ -394,10 +424,10 @@ export class EmbeddingIndexCoordinator { offset, offset + this.batchSize ) - const vectors = await provider.embed( - batch.map((item) => item.content), - signal - ) + const embeddingInput = batch.map((item) => item.content) + const vectors = provider.embedDocuments + ? await provider.embedDocuments(embeddingInput, signal) + : await provider.embed(embeddingInput, signal) if (vectors.length !== batch.length) { throw new EmbeddingOperationError({ code: 'invalid_response', diff --git a/src/main/knowledge/embedding-inference-broker.test.ts b/src/main/knowledge/embedding-inference-broker.test.ts new file mode 100644 index 00000000..6093cb6d --- /dev/null +++ b/src/main/knowledge/embedding-inference-broker.test.ts @@ -0,0 +1,243 @@ +import { describe, expect, it, vi } from 'vitest' +import { + EMBEDDING_INFERENCE_PROTOCOL, + EMBEDDING_INFERENCE_PROTOCOL_VERSION, + type EmbeddingInferenceRequest +} from './embedding-inference-contracts' +import { + EmbeddingInferenceBroker, + type EmbeddingInferenceTransport +} from './embedding-inference-broker' + +class TestTransport implements EmbeddingInferenceTransport { + readonly sent: EmbeddingInferenceRequest[] = [] + readonly close = vi.fn() + private readonly messageListeners = new Set<(message: unknown) => void>() + private readonly closeListeners = new Set<() => void>() + + postMessage(message: EmbeddingInferenceRequest): void { + this.sent.push(message) + } + + onMessage(listener: (message: unknown) => void): () => void { + this.messageListeners.add(listener) + return () => this.messageListeners.delete(listener) + } + + onClose(listener: () => void): () => void { + this.closeListeners.add(listener) + return () => this.closeListeners.delete(listener) + } + + receive(message: unknown): void { + for (const listener of this.messageListeners) { + listener(message) + } + } + + crash(): void { + for (const listener of this.closeListeners) { + listener() + } + } +} + +function response(message: Record): Record { + return { + protocol: EMBEDDING_INFERENCE_PROTOCOL, + version: EMBEDDING_INFERENCE_PROTOCOL_VERSION, + ...message + } +} + +function embedRequest(transport: TestTransport): Extract< + EmbeddingInferenceRequest, + { type: 'embed' } +> { + const request = transport.sent.find((item) => item.type === 'embed') + if (!request || request.type !== 'embed') { + throw new Error('No embed request was sent') + } + return request +} + +describe('EmbeddingInferenceBroker', () => { + it('uses UUID requests with roles and validates bounded results', async () => { + const transport = new TestTransport() + const broker = new EmbeddingInferenceBroker({ + createTransport: () => transport + }) + + const pending = broker.embed(['first', 'second'], 'query') + await vi.waitFor(() => expect(transport.sent).toHaveLength(1)) + const request = embedRequest(transport) + + expect(request).toMatchObject({ + type: 'embed', + role: 'query', + texts: ['first', 'second'] + }) + expect(request.requestId).toMatch( + /^[0-9a-f]{8}-[0-9a-f]{4}-[1-8][0-9a-f]{3}-[89ab][0-9a-f]{3}-[0-9a-f]{12}$/i + ) + expect(request).not.toHaveProperty('modelDirectory') + expect(request).not.toHaveProperty('env') + + transport.receive( + response({ + type: 'result', + requestId: request.requestId, + vectors: [ + [1, 2], + [3, 4] + ] + }) + ) + await expect(pending).resolves.toEqual([ + [1, 2], + [3, 4] + ]) + const shuttingDown = broker.shutdown() + transport.receive(response({ type: 'shutdown-complete' })) + await shuttingDown + }) + + it('rejects malformed input and untrusted vector responses', async () => { + const transport = new TestTransport() + const broker = new EmbeddingInferenceBroker({ + createTransport: () => transport + }) + + await expect(broker.embed([], 'document')).rejects.toThrow( + 'texts must contain' + ) + const pending = broker.embed(['safe input'], 'document') + await vi.waitFor(() => expect(transport.sent).toHaveLength(1)) + const request = embedRequest(transport) + transport.receive( + response({ + type: 'result', + requestId: request.requestId, + vectors: [[0, 0]] + }) + ) + + await expect(pending).rejects.toMatchObject({ + code: 'PROTOCOL_VIOLATION' + }) + expect(transport.close).toHaveBeenCalledOnce() + }) + + it('propagates cancellation and timeout without accepting late results', async () => { + vi.useFakeTimers() + try { + const transport = new TestTransport() + const broker = new EmbeddingInferenceBroker({ + createTransport: () => transport, + timeoutMs: 100 + }) + const controller = new AbortController() + const cancelled = broker.embed( + ['cancel this'], + 'query', + controller.signal + ) + await vi.advanceTimersByTimeAsync(0) + const cancelledRequest = embedRequest(transport) + controller.abort() + await expect(cancelled).rejects.toMatchObject({ name: 'AbortError' }) + expect(transport.sent.at(-1)).toMatchObject({ + type: 'cancel', + requestId: cancelledRequest.requestId + }) + + transport.receive( + response({ + type: 'error', + requestId: cancelledRequest.requestId, + code: 'CANCELLED', + message: 'Embedding inference was cancelled' + }) + ) + expect(transport.close).not.toHaveBeenCalled() + + const timedOut = broker.embed(['time this'], 'document') + await vi.advanceTimersByTimeAsync(0) + const timedRequest = transport.sent.find( + (item) => + item.type === 'embed' && + item.requestId !== cancelledRequest.requestId + ) + if (!timedRequest || timedRequest.type !== 'embed') { + throw new Error('No second embed request was sent') + } + const timedOutExpectation = expect(timedOut).rejects.toMatchObject({ + name: 'TimeoutError' + }) + await vi.advanceTimersByTimeAsync(100) + await timedOutExpectation + expect(transport.sent.at(-1)).toMatchObject({ + type: 'cancel', + requestId: timedRequest.requestId + }) + const shuttingDown = broker.shutdown() + transport.receive(response({ type: 'shutdown-complete' })) + await shuttingDown + } finally { + vi.useRealTimers() + } + }) + + it('rejects on process close and enforces its lifetime restart bound', async () => { + const transports = [ + new TestTransport(), + new TestTransport(), + new TestTransport() + ] + const factory = vi.fn(() => { + const transport = transports[factory.mock.calls.length - 1] + if (!transport) { + throw new Error('unexpected restart') + } + return transport + }) + const broker = new EmbeddingInferenceBroker({ + createTransport: factory, + maxRestarts: 1 + }) + + const first = broker.embed(['first'], 'query') + await vi.waitFor(() => expect(transports[0]?.sent).toHaveLength(1)) + transports[0]?.crash() + await expect(first).rejects.toMatchObject({ code: 'PROCESS_CLOSED' }) + + const second = broker.embed(['second'], 'document') + await vi.waitFor(() => expect(transports[1]?.sent).toHaveLength(1)) + transports[1]?.crash() + await expect(second).rejects.toMatchObject({ code: 'PROCESS_CLOSED' }) + + await expect(broker.embed(['third'], 'query')).rejects.toMatchObject({ + code: 'RESTART_LIMIT' + }) + expect(factory).toHaveBeenCalledTimes(2) + }) + + it('performs graceful shutdown and rejects future work', async () => { + const transport = new TestTransport() + const broker = new EmbeddingInferenceBroker({ + createTransport: () => transport + }) + const pending = broker.embed(['active'], 'query') + await vi.waitFor(() => expect(transport.sent).toHaveLength(1)) + + const shuttingDown = broker.shutdown() + await expect(pending).rejects.toMatchObject({ code: 'CLOSED' }) + expect(transport.sent.at(-1)).toMatchObject({ type: 'shutdown' }) + transport.receive(response({ type: 'shutdown-complete' })) + await expect(shuttingDown).resolves.toBeUndefined() + expect(transport.close).toHaveBeenCalledOnce() + await expect(broker.embed(['late'], 'query')).rejects.toMatchObject({ + code: 'CLOSED' + }) + }) +}) diff --git a/src/main/knowledge/embedding-inference-broker.ts b/src/main/knowledge/embedding-inference-broker.ts new file mode 100644 index 00000000..528379f9 --- /dev/null +++ b/src/main/knowledge/embedding-inference-broker.ts @@ -0,0 +1,455 @@ +import { randomUUID } from 'node:crypto' +import { + EMBEDDING_INFERENCE_MAX_IN_FLIGHT, + embeddingInferenceProtocolBase, + parseEmbeddingInferenceResponse, + validateEmbeddingInferenceTexts, + validateEmbeddingInferenceVectors, + type EmbeddingInferenceRequest, + type EmbeddingInferenceRole +} from './embedding-inference-contracts' + +const DEFAULT_TIMEOUT_MS = 120_000 +const MIN_TIMEOUT_MS = 100 +const MAX_TIMEOUT_MS = 300_000 +const DEFAULT_MAX_RESTARTS = 2 +const MAX_RESTARTS = 10 +const SHUTDOWN_GRACE_MS = 1_000 +const MAX_RETIRED_REQUESTS = 256 + +export interface EmbeddingInferenceTransport { + postMessage(message: EmbeddingInferenceRequest): void + onMessage(listener: (message: unknown) => void): () => void + onClose(listener: () => void): () => void + close(): void +} + +export type EmbeddingInferenceTransportFactory = () => + | EmbeddingInferenceTransport + | Promise + +export interface EmbeddingInferenceBrokerOptions { + readonly createTransport: EmbeddingInferenceTransportFactory + readonly timeoutMs?: number + /** + * Maximum replacement processes over this broker's lifetime. The initial + * process is not a restart and failed requests are never replayed. + */ + readonly maxRestarts?: number +} + +type PendingRequest = { + readonly expectedCount: number + readonly resolve: (vectors: number[][]) => void + readonly reject: (error: Error) => void + readonly timeout: ReturnType + removeAbortListener?: () => void +} + +type ActiveTransport = { + readonly transport: EmbeddingInferenceTransport + readonly removeMessageListener: () => void + readonly removeCloseListener: () => void +} + +export class EmbeddingInferenceBrokerError extends Error { + constructor( + readonly code: + | 'CLOSED' + | 'PROCESS_CLOSED' + | 'PROTOCOL_VIOLATION' + | 'RESTART_LIMIT' + | 'START_FAILURE' + | 'ENGINE_FAILURE' + | 'INVALID_RESULT' + | 'OVERLOADED' + ) { + super(`Embedding inference failed (${code})`) + this.name = 'EmbeddingInferenceBrokerError' + } +} + +function boundedInteger( + value: number, + field: string, + minimum: number, + maximum: number +): number { + if (!Number.isSafeInteger(value) || value < minimum || value > maximum) { + throw new RangeError( + `${field} must be an integer between ${minimum} and ${maximum}` + ) + } + return value +} + +function abortError(): Error { + const error = new Error('Embedding inference was cancelled') + error.name = 'AbortError' + return error +} + +function timeoutError(): Error { + const error = new Error('Embedding inference timed out') + error.name = 'TimeoutError' + return error +} + +export class EmbeddingInferenceBroker { + private readonly createTransport: EmbeddingInferenceTransportFactory + private readonly timeoutMs: number + private readonly maxRestarts: number + private readonly pending = new Map() + private readonly retiredRequestIds = new Set() + private active?: ActiveTransport + private starting?: Promise + private restarts = 0 + private everStarted = false + private closed = false + private shutdownComplete?: () => void + + constructor(options: EmbeddingInferenceBrokerOptions) { + if (typeof options.createTransport !== 'function') { + throw new TypeError('createTransport must be a function') + } + this.createTransport = options.createTransport + this.timeoutMs = boundedInteger( + options.timeoutMs ?? DEFAULT_TIMEOUT_MS, + 'timeoutMs', + MIN_TIMEOUT_MS, + MAX_TIMEOUT_MS + ) + this.maxRestarts = boundedInteger( + options.maxRestarts ?? DEFAULT_MAX_RESTARTS, + 'maxRestarts', + 0, + MAX_RESTARTS + ) + } + + async embed( + texts: readonly string[], + role: EmbeddingInferenceRole, + signal?: AbortSignal + ): Promise { + const normalizedTexts = validateEmbeddingInferenceTexts(texts) + if (role !== 'query' && role !== 'document') { + throw new TypeError('role must be query or document') + } + if (this.closed) { + throw new EmbeddingInferenceBrokerError('CLOSED') + } + if (signal?.aborted) { + throw abortError() + } + if (this.pending.size >= EMBEDDING_INFERENCE_MAX_IN_FLIGHT) { + throw new EmbeddingInferenceBrokerError('OVERLOADED') + } + + const active = await this.ensureTransport() + if (this.closed) { + throw new EmbeddingInferenceBrokerError('CLOSED') + } + if (signal?.aborted) { + throw abortError() + } + if (this.pending.size >= EMBEDDING_INFERENCE_MAX_IN_FLIGHT) { + throw new EmbeddingInferenceBrokerError('OVERLOADED') + } + + const requestId = randomUUID() + return new Promise((resolve, reject) => { + const timeout = setTimeout(() => { + this.cancelPending(requestId, timeoutError()) + }, this.timeoutMs) + const pending: PendingRequest = { + expectedCount: normalizedTexts.length, + resolve, + reject, + timeout + } + if (signal) { + const onAbort = (): void => { + this.cancelPending(requestId, abortError()) + } + signal.addEventListener('abort', onAbort, { once: true }) + pending.removeAbortListener = () => + signal.removeEventListener('abort', onAbort) + } + this.pending.set(requestId, pending) + + try { + active.transport.postMessage({ + ...embeddingInferenceProtocolBase(), + type: 'embed', + requestId, + role, + texts: normalizedTexts + }) + } catch { + this.failTransport( + active.transport, + new EmbeddingInferenceBrokerError('PROCESS_CLOSED') + ) + } + }) + } + + async shutdown(): Promise { + if (this.closed) { + return + } + this.closed = true + for (const requestId of this.pending.keys()) { + this.retireRequestId(requestId) + } + this.rejectAll(new EmbeddingInferenceBrokerError('CLOSED')) + + let active = this.active + if (!active && this.starting) { + try { + active = await this.starting + } catch { + return + } + } + if (!active || this.active !== active) { + return + } + + await new Promise((resolve) => { + let settled = false + const finish = (): void => { + if (settled) { + return + } + settled = true + clearTimeout(timeout) + this.shutdownComplete = undefined + this.disposeTransport(active.transport, true) + resolve() + } + const timeout = setTimeout(finish, SHUTDOWN_GRACE_MS) + this.shutdownComplete = finish + try { + active.transport.postMessage({ + ...embeddingInferenceProtocolBase(), + type: 'shutdown' + }) + } catch { + finish() + } + }) + } + + private async ensureTransport(): Promise { + if (this.active) { + return this.active + } + if (this.starting) { + return this.starting + } + if (this.everStarted && this.restarts >= this.maxRestarts) { + throw new EmbeddingInferenceBrokerError('RESTART_LIMIT') + } + if (this.everStarted) { + this.restarts += 1 + } + this.everStarted = true + + const starting = Promise.resolve() + .then(() => this.createTransport()) + .then((transport) => { + if (this.closed) { + transport.close() + throw new EmbeddingInferenceBrokerError('CLOSED') + } + const active: ActiveTransport = { + transport, + removeMessageListener: transport.onMessage((message) => + this.handleMessage(transport, message) + ), + removeCloseListener: transport.onClose(() => + this.failTransport( + transport, + new EmbeddingInferenceBrokerError('PROCESS_CLOSED') + ) + ) + } + this.active = active + return active + }) + .catch((error: unknown) => { + if (error instanceof EmbeddingInferenceBrokerError) { + throw error + } + throw new EmbeddingInferenceBrokerError('START_FAILURE') + }) + .finally(() => { + if (this.starting === starting) { + this.starting = undefined + } + }) + this.starting = starting + return starting + } + + private handleMessage( + transport: EmbeddingInferenceTransport, + rawMessage: unknown + ): void { + if (this.active?.transport !== transport) { + return + } + const message = parseEmbeddingInferenceResponse(rawMessage) + if (!message) { + this.failTransport( + transport, + new EmbeddingInferenceBrokerError('PROTOCOL_VIOLATION') + ) + return + } + if (message.type === 'shutdown-complete') { + if (!this.closed || !this.shutdownComplete) { + this.failTransport( + transport, + new EmbeddingInferenceBrokerError('PROTOCOL_VIOLATION') + ) + return + } + this.shutdownComplete() + return + } + if (message.type === 'fatal') { + this.failTransport( + transport, + new EmbeddingInferenceBrokerError('PROTOCOL_VIOLATION') + ) + return + } + + const pending = this.pending.get(message.requestId) + if (!pending) { + if (this.retiredRequestIds.delete(message.requestId)) { + return + } + this.failTransport( + transport, + new EmbeddingInferenceBrokerError('PROTOCOL_VIOLATION') + ) + return + } + + if (message.type === 'error') { + this.finishPending( + message.requestId, + message.code === 'CANCELLED' + ? abortError() + : new EmbeddingInferenceBrokerError(message.code) + ) + return + } + try { + const vectors = validateEmbeddingInferenceVectors( + message.vectors, + pending.expectedCount + ) + this.finishPending(message.requestId, undefined, vectors) + } catch { + this.failTransport( + transport, + new EmbeddingInferenceBrokerError('PROTOCOL_VIOLATION') + ) + } + } + + private cancelPending(requestId: string, error: Error): void { + const pending = this.pending.get(requestId) + if (!pending) { + return + } + this.retireRequestId(requestId) + this.finishPending(requestId, error) + try { + this.active?.transport.postMessage({ + ...embeddingInferenceProtocolBase(), + type: 'cancel', + requestId + }) + } catch { + if (this.active) { + this.failTransport( + this.active.transport, + new EmbeddingInferenceBrokerError('PROCESS_CLOSED') + ) + } + } + } + + private finishPending( + requestId: string, + error?: Error, + vectors?: number[][] + ): void { + const pending = this.pending.get(requestId) + if (!pending) { + return + } + this.pending.delete(requestId) + clearTimeout(pending.timeout) + pending.removeAbortListener?.() + if (error) { + pending.reject(error) + } else { + pending.resolve(vectors as number[][]) + } + } + + private retireRequestId(requestId: string): void { + this.retiredRequestIds.add(requestId) + if (this.retiredRequestIds.size > MAX_RETIRED_REQUESTS) { + const oldest = this.retiredRequestIds.values().next().value + if (typeof oldest === 'string') { + this.retiredRequestIds.delete(oldest) + } + } + } + + private failTransport( + transport: EmbeddingInferenceTransport, + error: Error + ): void { + if (this.active?.transport !== transport) { + return + } + this.disposeTransport(transport, true) + this.rejectAll(error) + this.shutdownComplete?.() + } + + private disposeTransport( + transport: EmbeddingInferenceTransport, + close: boolean + ): void { + const active = this.active + if (!active || active.transport !== transport) { + return + } + this.active = undefined + active.removeMessageListener() + active.removeCloseListener() + if (close) { + try { + transport.close() + } catch { + // The process is already unavailable. + } + } + } + + private rejectAll(error: Error): void { + for (const requestId of [...this.pending.keys()]) { + this.finishPending(requestId, error) + } + } +} diff --git a/src/main/knowledge/embedding-inference-contracts.ts b/src/main/knowledge/embedding-inference-contracts.ts new file mode 100644 index 00000000..10da1bc9 --- /dev/null +++ b/src/main/knowledge/embedding-inference-contracts.ts @@ -0,0 +1,262 @@ +export const EMBEDDING_INFERENCE_PROTOCOL = + 'goodbuddy.embedding-inference' +export const EMBEDDING_INFERENCE_PROTOCOL_VERSION = 1 + +export const EMBEDDING_INFERENCE_MAX_TEXTS = 64 +export const EMBEDDING_INFERENCE_MAX_TEXT_LENGTH = 16_000 +export const EMBEDDING_INFERENCE_MAX_TOTAL_TEXT_LENGTH = 128_000 +export const EMBEDDING_INFERENCE_MAX_DIMENSIONS = 8_192 +export const EMBEDDING_INFERENCE_MAX_IN_FLIGHT = 8 +export const EMBEDDING_INFERENCE_MAX_ERROR_LENGTH = 256 + +export type EmbeddingInferenceRole = 'query' | 'document' + +type ProtocolBase = { + readonly protocol: typeof EMBEDDING_INFERENCE_PROTOCOL + readonly version: typeof EMBEDDING_INFERENCE_PROTOCOL_VERSION +} + +export type EmbeddingInferenceRequest = + | (ProtocolBase & { + readonly type: 'embed' + readonly requestId: string + readonly role: EmbeddingInferenceRole + readonly texts: readonly string[] + }) + | (ProtocolBase & { + readonly type: 'cancel' + readonly requestId: string + }) + | (ProtocolBase & { + readonly type: 'shutdown' + }) + +export type EmbeddingInferenceErrorCode = + | 'CANCELLED' + | 'ENGINE_FAILURE' + | 'INVALID_RESULT' + | 'OVERLOADED' + +export type EmbeddingInferenceResponse = + | (ProtocolBase & { + readonly type: 'result' + readonly requestId: string + readonly vectors: readonly (readonly number[])[] + }) + | (ProtocolBase & { + readonly type: 'error' + readonly requestId: string + readonly code: EmbeddingInferenceErrorCode + readonly message: string + }) + | (ProtocolBase & { + readonly type: 'fatal' + readonly code: 'PROTOCOL_VIOLATION' + }) + | (ProtocolBase & { + readonly type: 'shutdown-complete' + }) + +const BASE_KEYS = ['protocol', 'version', 'type'] as const +const REQUEST_ID_KEYS = [...BASE_KEYS, 'requestId'] as const +const EMBED_KEYS = [...REQUEST_ID_KEYS, 'role', 'texts'] as const +const RESULT_KEYS = [...REQUEST_ID_KEYS, 'vectors'] as const +const ERROR_KEYS = [...REQUEST_ID_KEYS, 'code', 'message'] as const +const FATAL_KEYS = [...BASE_KEYS, 'code'] as const + +const UUID_PATTERN = + /^[0-9a-f]{8}-[0-9a-f]{4}-[1-8][0-9a-f]{3}-[89ab][0-9a-f]{3}-[0-9a-f]{12}$/i + +function isRecord(value: unknown): value is Record { + if (typeof value !== 'object' || value === null || Array.isArray(value)) { + return false + } + const prototype = Object.getPrototypeOf(value) + return prototype === Object.prototype || prototype === null +} + +function hasExactKeys( + value: Record, + expected: readonly string[] +): boolean { + const keys = Object.keys(value) + return ( + keys.length === expected.length && + expected.every((key) => Object.prototype.hasOwnProperty.call(value, key)) + ) +} + +function hasProtocolBase(value: Record): boolean { + return ( + value.protocol === EMBEDDING_INFERENCE_PROTOCOL && + value.version === EMBEDDING_INFERENCE_PROTOCOL_VERSION && + typeof value.type === 'string' + ) +} + +export function isEmbeddingInferenceRequestId( + value: unknown +): value is string { + return typeof value === 'string' && UUID_PATTERN.test(value) +} + +export function validateEmbeddingInferenceTexts( + value: unknown +): readonly string[] { + if ( + !Array.isArray(value) || + value.length < 1 || + value.length > EMBEDDING_INFERENCE_MAX_TEXTS + ) { + throw new RangeError( + `texts must contain between 1 and ${EMBEDDING_INFERENCE_MAX_TEXTS} items` + ) + } + + let totalLength = 0 + for (const [index, text] of value.entries()) { + if (typeof text !== 'string' || text.length < 1) { + throw new TypeError(`texts[${index}] must be a non-empty string`) + } + if (text.length > EMBEDDING_INFERENCE_MAX_TEXT_LENGTH) { + throw new RangeError( + `texts[${index}] must be at most ${EMBEDDING_INFERENCE_MAX_TEXT_LENGTH} characters` + ) + } + totalLength += text.length + if (totalLength > EMBEDDING_INFERENCE_MAX_TOTAL_TEXT_LENGTH) { + throw new RangeError( + `texts must total at most ${EMBEDDING_INFERENCE_MAX_TOTAL_TEXT_LENGTH} characters` + ) + } + } + return value as readonly string[] +} + +export function validateEmbeddingInferenceVectors( + value: unknown, + expectedCount: number +): number[][] { + if (!Array.isArray(value) || value.length !== expectedCount) { + throw new RangeError('Embedding result count does not match the request') + } + + let dimensions: number | undefined + return value.map((rawVector, vectorIndex) => { + if ( + !Array.isArray(rawVector) || + rawVector.length < 1 || + rawVector.length > EMBEDDING_INFERENCE_MAX_DIMENSIONS + ) { + throw new RangeError( + `Embedding ${vectorIndex} has invalid dimensions` + ) + } + if (dimensions === undefined) { + dimensions = rawVector.length + } else if (rawVector.length !== dimensions) { + throw new RangeError('Embeddings have inconsistent dimensions') + } + + let magnitudeSquared = 0 + const vector = rawVector.map((component) => { + if (typeof component !== 'number' || !Number.isFinite(component)) { + throw new TypeError('Embeddings must contain finite numbers') + } + magnitudeSquared += component * component + return component + }) + if (!Number.isFinite(magnitudeSquared) || magnitudeSquared <= 0) { + throw new RangeError('Embeddings must have a finite non-zero norm') + } + return vector + }) +} + +export function parseEmbeddingInferenceRequest( + value: unknown +): EmbeddingInferenceRequest | undefined { + if (!isRecord(value) || !hasProtocolBase(value)) { + return undefined + } + if (value.type === 'shutdown') { + return hasExactKeys(value, BASE_KEYS) + ? (value as EmbeddingInferenceRequest) + : undefined + } + if ( + value.type === 'cancel' && + hasExactKeys(value, REQUEST_ID_KEYS) && + isEmbeddingInferenceRequestId(value.requestId) + ) { + return value as EmbeddingInferenceRequest + } + if ( + value.type !== 'embed' || + !hasExactKeys(value, EMBED_KEYS) || + !isEmbeddingInferenceRequestId(value.requestId) || + (value.role !== 'query' && value.role !== 'document') + ) { + return undefined + } + try { + validateEmbeddingInferenceTexts(value.texts) + return value as EmbeddingInferenceRequest + } catch { + return undefined + } +} + +export function parseEmbeddingInferenceResponse( + value: unknown +): EmbeddingInferenceResponse | undefined { + if (!isRecord(value) || !hasProtocolBase(value)) { + return undefined + } + if (value.type === 'shutdown-complete') { + return hasExactKeys(value, BASE_KEYS) + ? (value as EmbeddingInferenceResponse) + : undefined + } + if ( + value.type === 'fatal' && + hasExactKeys(value, FATAL_KEYS) && + value.code === 'PROTOCOL_VIOLATION' + ) { + return value as EmbeddingInferenceResponse + } + if ( + value.type === 'result' && + hasExactKeys(value, RESULT_KEYS) && + isEmbeddingInferenceRequestId(value.requestId) && + Array.isArray(value.vectors) + ) { + return value as EmbeddingInferenceResponse + } + const errorCodes: readonly EmbeddingInferenceErrorCode[] = [ + 'CANCELLED', + 'ENGINE_FAILURE', + 'INVALID_RESULT', + 'OVERLOADED' + ] + if ( + value.type === 'error' && + hasExactKeys(value, ERROR_KEYS) && + isEmbeddingInferenceRequestId(value.requestId) && + typeof value.code === 'string' && + errorCodes.includes(value.code as EmbeddingInferenceErrorCode) && + typeof value.message === 'string' && + value.message.length > 0 && + value.message.length <= EMBEDDING_INFERENCE_MAX_ERROR_LENGTH + ) { + return value as EmbeddingInferenceResponse + } + return undefined +} + +export function embeddingInferenceProtocolBase(): ProtocolBase { + return { + protocol: EMBEDDING_INFERENCE_PROTOCOL, + version: EMBEDDING_INFERENCE_PROTOCOL_VERSION + } +} diff --git a/src/main/knowledge/embedding-inference-worker.test.ts b/src/main/knowledge/embedding-inference-worker.test.ts new file mode 100644 index 00000000..6bef9085 --- /dev/null +++ b/src/main/knowledge/embedding-inference-worker.test.ts @@ -0,0 +1,251 @@ +import path from 'node:path' +import { describe, expect, it, vi } from 'vitest' +import { + EMBEDDING_INFERENCE_PROTOCOL, + EMBEDDING_INFERENCE_PROTOCOL_VERSION, + type EmbeddingInferenceRequest, + type EmbeddingInferenceResponse +} from './embedding-inference-contracts' +import { + EmbeddingInferenceWorker, + type EmbeddingInferenceEngine, + type EmbeddingInferenceWorkerPort +} from './embedding-inference-worker' + +const FIRST_ID = '123e4567-e89b-42d3-a456-426614174000' +const SECOND_ID = '123e4567-e89b-42d3-a456-426614174001' + +type RequestWithoutBase = T extends unknown + ? Omit + : never + +class TestPort implements EmbeddingInferenceWorkerPort { + readonly sent: EmbeddingInferenceResponse[] = [] + private readonly listeners = new Set<(message: unknown) => void>() + + postMessage(message: EmbeddingInferenceResponse): void { + this.sent.push(message) + } + + onMessage(listener: (message: unknown) => void): () => void { + this.listeners.add(listener) + return () => this.listeners.delete(listener) + } + + receive(message: unknown): void { + for (const listener of this.listeners) { + listener(message) + } + } +} + +function request( + message: RequestWithoutBase +): EmbeddingInferenceRequest { + return { + protocol: EMBEDDING_INFERENCE_PROTOCOL, + version: EMBEDDING_INFERENCE_PROTOCOL_VERSION, + ...message + } as EmbeddingInferenceRequest +} + +function engine( + embed: EmbeddingInferenceEngine['embed'] +): EmbeddingInferenceEngine { + return { + tokenizer: { tokenize: async () => ({}) }, + onnx: { run: async () => ({}) }, + embed + } +} + +describe('EmbeddingInferenceWorker', () => { + it('uses only its trusted model directory and forwards role to the engine', async () => { + const port = new TestPort() + const modelDirectory = path.resolve('trusted-model') + const embed = vi.fn( + async () => [[1, 2, 3]] + ) + const worker = new EmbeddingInferenceWorker({ + modelDirectory, + engine: engine(embed), + port + }) + + port.receive( + request({ + type: 'embed', + requestId: FIRST_ID, + role: 'document', + texts: ['bounded text'] + }) + ) + await vi.waitFor(() => expect(port.sent).toHaveLength(1)) + + expect(embed).toHaveBeenCalledWith({ + modelDirectory, + texts: ['bounded text'], + role: 'document', + signal: expect.any(AbortSignal) + }) + expect(port.sent[0]).toMatchObject({ + type: 'result', + requestId: FIRST_ID, + vectors: [[1, 2, 3]] + }) + await worker.shutdown() + }) + + it('rejects invalid engine output instead of fabricating vectors', async () => { + const port = new TestPort() + const worker = new EmbeddingInferenceWorker({ + modelDirectory: path.resolve('trusted-model'), + engine: engine(async () => [[Number.NaN]]), + port + }) + + port.receive( + request({ + type: 'embed', + requestId: FIRST_ID, + role: 'query', + texts: ['input'] + }) + ) + await vi.waitFor(() => + expect(port.sent[0]).toMatchObject({ + type: 'error', + requestId: FIRST_ID, + code: 'INVALID_RESULT' + }) + ) + await worker.shutdown() + }) + + it('aborts an active engine request when cancelled', async () => { + const port = new TestPort() + let observedSignal: AbortSignal | undefined + const worker = new EmbeddingInferenceWorker({ + modelDirectory: path.resolve('trusted-model'), + engine: engine( + async ({ signal }) => + new Promise((_resolve, reject) => { + observedSignal = signal + signal.addEventListener('abort', () => reject(signal.reason), { + once: true + }) + }) + ), + port + }) + + port.receive( + request({ + type: 'embed', + requestId: FIRST_ID, + role: 'query', + texts: ['input'] + }) + ) + port.receive(request({ type: 'cancel', requestId: FIRST_ID })) + await vi.waitFor(() => expect(observedSignal?.aborted).toBe(true)) + await vi.waitFor(() => + expect(port.sent[0]).toMatchObject({ + type: 'error', + requestId: FIRST_ID, + code: 'CANCELLED' + }) + ) + await worker.shutdown() + }) + + it('treats malformed, credential-bearing, and duplicate requests as violations', async () => { + const malformedPort = new TestPort() + const malformedWorker = new EmbeddingInferenceWorker({ + modelDirectory: path.resolve('trusted-model'), + engine: engine(async () => [[1]]), + port: malformedPort + }) + malformedPort.receive({ + ...request({ + type: 'embed', + requestId: FIRST_ID, + role: 'query', + texts: ['input'] + }), + apiKey: 'must-not-cross-boundary' + }) + await vi.waitFor(() => + expect(malformedPort.sent[0]).toMatchObject({ + type: 'fatal', + code: 'PROTOCOL_VIOLATION' + }) + ) + await malformedWorker.shutdown() + + let finish: + | ((vectors: readonly (readonly number[])[]) => void) + | undefined + const duplicatePort = new TestPort() + const duplicateWorker = new EmbeddingInferenceWorker({ + modelDirectory: path.resolve('trusted-model'), + engine: engine( + async () => + new Promise((resolve) => { + finish = resolve + }) + ), + port: duplicatePort + }) + const duplicate = request({ + type: 'embed', + requestId: SECOND_ID, + role: 'query', + texts: ['input'] + }) + duplicatePort.receive(duplicate) + duplicatePort.receive(duplicate) + await vi.waitFor(() => + expect(duplicatePort.sent[0]).toMatchObject({ + type: 'fatal', + code: 'PROTOCOL_VIOLATION' + }) + ) + finish?.([[1]]) + await duplicateWorker.shutdown() + }) + + it('acknowledges shutdown only after disposing its engine', async () => { + const port = new TestPort() + const dispose = vi.fn(async () => undefined) + new EmbeddingInferenceWorker({ + modelDirectory: path.resolve('trusted-model'), + engine: { + ...engine(async () => [[1]]), + dispose + }, + port + }) + + port.receive(request({ type: 'shutdown' })) + await vi.waitFor(() => + expect(port.sent[0]).toMatchObject({ type: 'shutdown-complete' }) + ) + expect(dispose).toHaveBeenCalledOnce() + }) + + it('requires explicit tokenizer and ONNX engine implementations', () => { + expect( + () => + new EmbeddingInferenceWorker({ + modelDirectory: path.resolve('trusted-model'), + engine: { + tokenizer: undefined, + onnx: undefined, + embed: async () => [[1]] + } as unknown as EmbeddingInferenceEngine, + port: new TestPort() + }) + ).toThrow('tokenizer and ONNX') + }) +}) diff --git a/src/main/knowledge/embedding-inference-worker.ts b/src/main/knowledge/embedding-inference-worker.ts new file mode 100644 index 00000000..b8906f4e --- /dev/null +++ b/src/main/knowledge/embedding-inference-worker.ts @@ -0,0 +1,309 @@ +import path from 'node:path' +import { + EMBEDDING_INFERENCE_MAX_ERROR_LENGTH, + EMBEDDING_INFERENCE_MAX_IN_FLIGHT, + embeddingInferenceProtocolBase, + parseEmbeddingInferenceRequest, + validateEmbeddingInferenceVectors, + type EmbeddingInferenceRequest, + type EmbeddingInferenceResponse, + type EmbeddingInferenceRole +} from './embedding-inference-contracts' + +export interface EmbeddingTokenizer { + tokenize(options: { + readonly texts: readonly string[] + readonly role: EmbeddingInferenceRole + readonly signal: AbortSignal + }): Promise +} + +export interface EmbeddingOnnxRuntime { + run(options: { + readonly modelDirectory: string + readonly inputs: unknown + readonly signal: AbortSignal + }): Promise +} + +/** + * Production engines must own both a tokenizer and an ONNX runtime. This + * boundary deliberately supplies no fallback or synthetic vector engine. + */ +export interface EmbeddingInferenceEngine { + readonly tokenizer: EmbeddingTokenizer + readonly onnx: EmbeddingOnnxRuntime + embed(options: { + readonly modelDirectory: string + readonly texts: readonly string[] + readonly role: EmbeddingInferenceRole + readonly signal: AbortSignal + }): Promise + dispose?(): void | Promise +} + +export interface EmbeddingInferenceWorkerPort { + postMessage(message: EmbeddingInferenceResponse): void + onMessage(listener: (message: unknown) => void): () => void +} + +export interface EmbeddingInferenceWorkerOptions { + /** + * Trusted Main configuration, never read from a protocol message. + */ + readonly modelDirectory: string + readonly engine: EmbeddingInferenceEngine + readonly port: EmbeddingInferenceWorkerPort +} + +type RunningRequest = { + readonly controller: AbortController + readonly done: Promise +} + +function trustedModelDirectory(value: string): string { + if ( + typeof value !== 'string' || + value.length < 1 || + value.length > 4_096 || + !path.isAbsolute(value) + ) { + throw new TypeError('modelDirectory must be a bounded absolute path') + } + return path.normalize(value) +} + +function boundedMessage(value: string): string { + return value.slice(0, EMBEDDING_INFERENCE_MAX_ERROR_LENGTH) +} + +export class EmbeddingInferenceWorker { + private readonly modelDirectory: string + private readonly engine: EmbeddingInferenceEngine + private readonly port: EmbeddingInferenceWorkerPort + private readonly running = new Map() + private readonly removeMessageListener: () => void + private shuttingDown = false + private disposed = false + private shutdownTask?: Promise + + constructor(options: EmbeddingInferenceWorkerOptions) { + this.modelDirectory = trustedModelDirectory(options.modelDirectory) + this.engine = options.engine + this.port = options.port + if ( + typeof options.engine.tokenizer?.tokenize !== 'function' || + typeof options.engine.onnx?.run !== 'function' + ) { + throw new TypeError('engine must provide tokenizer and ONNX implementations') + } + this.removeMessageListener = this.port.onMessage((message) => { + this.handleMessage(message) + }) + } + + async shutdown(): Promise { + await this.stop(false) + } + + private handleMessage(rawMessage: unknown): void { + if (this.disposed) { + return + } + const message = parseEmbeddingInferenceRequest(rawMessage) + if (!message) { + this.protocolViolation() + return + } + if (message.type === 'shutdown') { + if (this.shuttingDown) { + this.protocolViolation() + return + } + this.shuttingDown = true + void this.stop(true).then( + () => undefined, + () => { + this.post({ + ...embeddingInferenceProtocolBase(), + type: 'fatal', + code: 'PROTOCOL_VIOLATION' + }) + } + ) + return + } + if (message.type === 'cancel') { + this.running.get(message.requestId)?.controller.abort() + return + } + this.startEmbedding(message) + } + + private startEmbedding( + message: Extract + ): void { + if ( + this.shuttingDown || + this.running.has(message.requestId) + ) { + this.protocolViolation() + return + } + if (this.running.size >= EMBEDDING_INFERENCE_MAX_IN_FLIGHT) { + this.postError( + message.requestId, + 'OVERLOADED', + 'Embedding inference worker is busy' + ) + return + } + + const controller = new AbortController() + const done = this.runEmbedding(message, controller) + this.running.set(message.requestId, { controller, done }) + void done.finally(() => { + this.running.delete(message.requestId) + }) + } + + private async runEmbedding( + message: Extract, + controller: AbortController + ): Promise { + try { + const rawVectors = await this.engine.embed({ + modelDirectory: this.modelDirectory, + texts: message.texts, + role: message.role, + signal: controller.signal + }) + if (controller.signal.aborted) { + this.postError( + message.requestId, + 'CANCELLED', + 'Embedding inference was cancelled' + ) + return + } + let vectors: number[][] + try { + vectors = validateEmbeddingInferenceVectors( + rawVectors, + message.texts.length + ) + } catch { + this.postError( + message.requestId, + 'INVALID_RESULT', + 'Embedding engine returned an invalid result' + ) + return + } + this.post({ + ...embeddingInferenceProtocolBase(), + type: 'result', + requestId: message.requestId, + vectors + }) + } catch { + this.postError( + message.requestId, + controller.signal.aborted ? 'CANCELLED' : 'ENGINE_FAILURE', + controller.signal.aborted + ? 'Embedding inference was cancelled' + : 'Embedding engine failed' + ) + } + } + + private postError( + requestId: string, + code: 'CANCELLED' | 'ENGINE_FAILURE' | 'INVALID_RESULT' | 'OVERLOADED', + message: string + ): void { + this.post({ + ...embeddingInferenceProtocolBase(), + type: 'error', + requestId, + code, + message: boundedMessage(message) + }) + } + + private protocolViolation(): void { + if (this.disposed) { + return + } + this.shuttingDown = true + this.post({ + ...embeddingInferenceProtocolBase(), + type: 'fatal', + code: 'PROTOCOL_VIOLATION' + }) + void this.shutdown() + } + + private stop(sendCompletion: boolean): Promise { + if (this.shutdownTask) { + return this.shutdownTask + } + this.shuttingDown = true + const task = (async (): Promise => { + try { + for (const request of this.running.values()) { + request.controller.abort() + } + await Promise.allSettled( + [...this.running.values()].map((request) => request.done) + ) + await this.engine.dispose?.() + if (sendCompletion) { + this.post({ + ...embeddingInferenceProtocolBase(), + type: 'shutdown-complete' + }) + } + } finally { + this.disposed = true + this.removeMessageListener() + } + })() + this.shutdownTask = task + return task + } + + private post(message: EmbeddingInferenceResponse): void { + try { + this.port.postMessage(message) + } catch { + void this.shutdown() + } + } +} + +export interface ElectronParentPortLike { + postMessage(message: unknown): void + on( + event: 'message', + listener: (event: { readonly data: unknown }) => void + ): unknown + removeListener( + event: 'message', + listener: (event: { readonly data: unknown }) => void + ): unknown +} + +export function createEmbeddingInferenceWorkerPort( + parentPort: ElectronParentPortLike +): EmbeddingInferenceWorkerPort { + return { + postMessage: (message) => parentPort.postMessage(message), + onMessage: (listener) => { + const onMessage = (event: { readonly data: unknown }): void => + listener(event.data) + parentPort.on('message', onMessage) + return () => parentPort.removeListener('message', onMessage) + } + } +} diff --git a/src/main/knowledge/embedding-model-catalog.ts b/src/main/knowledge/embedding-model-catalog.ts new file mode 100644 index 00000000..d2e2b48d --- /dev/null +++ b/src/main/knowledge/embedding-model-catalog.ts @@ -0,0 +1,136 @@ +import { + embeddingModelCatalogEntrySchema, + type EmbeddingModelCatalogEntry +} from './embedding-model-contracts' +import { + huggingFaceTarget, + modelScopeTarget +} from '../model-download-targets' + +const repository = 'onnx-community/granite-embedding-97m-multilingual-r2-ONNX' +const huggingFaceRevision = '536a9f241cb3f02a9c5995a1e708c784bd274859' +const modelScopeRevision = '2741cd30a7448219ec2699afdf373a44df5aaa33' + +export const EMBEDDING_MODEL_CATALOG: readonly EmbeddingModelCatalogEntry[] = + embeddingModelCatalogEntrySchema.array().parse([ + { + id: 'granite-embedding-97m-multilingual-r2-int8', + displayName: 'Granite Embedding 97M Multilingual R2', + description: + '面向中文、英文及跨语言检索的 384 维本地 INT8 向量模型。', + languages: ['200+ languages'], + runtime: 'onnxruntime-web/wasm', + dimensions: 384, + contextTokens: 32_768, + quantization: 'INT8', + recommended: true, + available: true, + repositoryUrls: { + modelscope: `https://modelscope.cn/models/${repository}`, + 'hugging-face': `https://huggingface.co/${repository}` + }, + license: { + name: 'Apache License 2.0', + notice: + '模型由 IBM Granite 发布，ONNX 工件由 ONNX Community 转换，依据 Apache License 2.0 使用。', + url: 'https://huggingface.co/ibm-granite/granite-embedding-97m-multilingual-r2/blob/main/LICENSE' + }, + files: [ + { + name: 'model_quantized.onnx', + role: 'model', + size: 97_858_099, + sha256: + '704c1ebca5fbb7cd83ced41827658ac4c9990c64f7f2874d22b78044e5022e22', + targets: { + modelscope: modelScopeTarget( + repository, + modelScopeRevision, + 'onnx/model_quantized.onnx' + ), + 'hugging-face': huggingFaceTarget( + repository, + huggingFaceRevision, + 'onnx/model_quantized.onnx' + ) + } + }, + { + name: 'tokenizer.json', + role: 'tokenizer', + size: 25_301_671, + sha256: + '51947676cae1f991fa51c6b9a24e14ee5460e5f0b9f692f13bb3159829d1592a', + targets: { + modelscope: modelScopeTarget( + repository, + modelScopeRevision, + 'tokenizer.json' + ), + 'hugging-face': huggingFaceTarget( + repository, + huggingFaceRevision, + 'tokenizer.json' + ) + } + }, + { + name: 'tokenizer_config.json', + role: 'tokenizer-configuration', + size: 12_860, + sha256: + '6ed69389e30a8ecabfce2f9ebcdf0c908b34056f24d994340f2f216521c057d5', + targets: { + modelscope: modelScopeTarget( + repository, + modelScopeRevision, + 'tokenizer_config.json' + ), + 'hugging-face': huggingFaceTarget( + repository, + huggingFaceRevision, + 'tokenizer_config.json' + ) + } + }, + { + name: 'config.json', + role: 'configuration', + size: 1_215, + sha256: + 'ae74d55a56f779774cb9a8e63d3c2da9ae1af83c00229ffdff43d0b38407a0ee', + targets: { + modelscope: modelScopeTarget( + repository, + modelScopeRevision, + 'config.json' + ), + 'hugging-face': huggingFaceTarget( + repository, + huggingFaceRevision, + 'config.json' + ) + } + }, + { + name: 'special_tokens_map.json', + role: 'tokenizer-configuration', + size: 871, + sha256: + '013787ee251ff611722479197c00853b62113ad303cb0a36524231783c676c69', + targets: { + modelscope: modelScopeTarget( + repository, + modelScopeRevision, + 'special_tokens_map.json' + ), + 'hugging-face': huggingFaceTarget( + repository, + huggingFaceRevision, + 'special_tokens_map.json' + ) + } + } + ] + } + ]) diff --git a/src/main/knowledge/embedding-model-contracts.test.ts b/src/main/knowledge/embedding-model-contracts.test.ts new file mode 100644 index 00000000..37c1ee8e --- /dev/null +++ b/src/main/knowledge/embedding-model-contracts.test.ts @@ -0,0 +1,68 @@ +import { describe, expect, it } from 'vitest' +import { embeddingModelCatalogEntrySchema } from './embedding-model-contracts' + +describe('embedding model contracts', () => { + it('requires a reason when a catalog package is unavailable', () => { + const result = embeddingModelCatalogEntrySchema.safeParse({ + id: 'test-model', + displayName: 'Test', + description: 'Test model.', + languages: ['English'], + runtime: 'onnxruntime-web/wasm', + dimensions: 4, + contextTokens: 128, + quantization: 'int8', + recommended: false, + available: false, + repositoryUrls: {}, + license: { + name: 'Test', + notice: 'Test only.', + url: 'https://example.com/license' + }, + files: [ + { + name: 'model.onnx', + role: 'model', + size: 1, + sha256: 'a'.repeat(64), + targets: {} + } + ] + }) + + expect(result.success).toBe(false) + }) + + it('requires all runtime artifacts for an available package', () => { + const result = embeddingModelCatalogEntrySchema.safeParse({ + id: 'test-model', + displayName: 'Test', + description: 'Test model.', + languages: ['English'], + runtime: 'onnxruntime-web/wasm', + dimensions: 4, + contextTokens: 128, + quantization: 'int8', + recommended: false, + available: true, + repositoryUrls: {}, + license: { + name: 'Test', + notice: 'Test only.', + url: 'https://example.com/license' + }, + files: [ + { + name: 'model.onnx', + role: 'model', + size: 1, + sha256: 'a'.repeat(64), + targets: {} + } + ] + }) + + expect(result.success).toBe(false) + }) +}) diff --git a/src/main/knowledge/embedding-model-contracts.ts b/src/main/knowledge/embedding-model-contracts.ts new file mode 100644 index 00000000..f0f57f66 --- /dev/null +++ b/src/main/knowledge/embedding-model-contracts.ts @@ -0,0 +1,220 @@ +import { z } from 'zod' +import { + MODEL_DOWNLOAD_SOURCES, + modelArtifactTargetsSchema, + modelDownloadAvailabilitySchema, + modelDownloadSourceSchema +} from '../../shared/model-download-contracts' + +const sha256Schema = z.string().regex(/^[a-f0-9]{64}$/u) +const modelFileNameSchema = z + .string() + .min(1) + .max(255) + .regex(/^[^/\\:\0]+$/u) + +export const embeddingModelIdSchema = z + .string() + .min(1) + .max(96) + .regex(/^[a-z0-9]+(?:-[a-z0-9]+)*$/u) + +export const embeddingModelFileRoleSchema = z.enum([ + 'model', + 'tokenizer', + 'tokenizer-configuration', + 'configuration', + 'license' +]) + +export const embeddingModelFileSchema = z + .object({ + name: modelFileNameSchema, + role: embeddingModelFileRoleSchema, + size: z.number().int().positive().safe(), + sha256: sha256Schema, + targets: modelArtifactTargetsSchema + }) + .strict() + +const repositoryUrlsSchema = z + .object({ + modelscope: z.url().optional(), + 'hugging-face': z.url().optional() + }) + .strict() + +const embeddingModelLicenseSchema = z + .object({ + name: z.string().trim().min(1).max(120), + notice: z.string().trim().min(1).max(1_000), + url: z.url() + }) + .strict() + +export const embeddingModelCatalogEntrySchema = z + .object({ + id: embeddingModelIdSchema, + displayName: z.string().trim().min(1).max(120), + description: z.string().trim().min(1).max(1_000), + languages: z.array(z.string().trim().min(1).max(80)).min(1).max(32), + runtime: z.literal('onnxruntime-web/wasm'), + dimensions: z.number().int().positive().max(65_536), + contextTokens: z.number().int().positive().max(10_000_000), + quantization: z.string().trim().min(1).max(40), + recommended: z.boolean(), + available: z.boolean(), + unavailableReason: z.string().trim().min(1).max(1_000).optional(), + repositoryUrls: repositoryUrlsSchema, + license: embeddingModelLicenseSchema, + files: z.array(embeddingModelFileSchema).min(1).max(32) + }) + .strict() + .superRefine((entry, context) => { + if (!entry.available && !entry.unavailableReason) { + context.addIssue({ + code: 'custom', + path: ['unavailableReason'], + message: '不可用的内置向量模型必须说明原因' + }) + } + if (entry.available && entry.unavailableReason) { + context.addIssue({ + code: 'custom', + path: ['unavailableReason'], + message: '可用的内置向量模型不能包含不可用原因' + }) + } + if (entry.available) { + for (const role of [ + 'model', + 'tokenizer', + 'tokenizer-configuration', + 'configuration' + ] as const) { + if (!entry.files.some((file) => file.role === role)) { + context.addIssue({ + code: 'custom', + path: ['files'], + message: `可安装的向量模型缺少 ${role} 工件` + }) + } + } + } + if ( + new Set(entry.files.map((file) => file.name.toLowerCase())).size !== + entry.files.length + ) { + context.addIssue({ + code: 'custom', + path: ['files'], + message: '向量模型目录包含重复文件名' + }) + } + }) + +const embeddingModelCatalogViewFileSchema = embeddingModelFileSchema.omit({ + targets: true +}) + +export const embeddingModelCatalogViewEntrySchema = z + .object({ + id: embeddingModelCatalogEntrySchema.shape.id, + displayName: embeddingModelCatalogEntrySchema.shape.displayName, + description: embeddingModelCatalogEntrySchema.shape.description, + languages: embeddingModelCatalogEntrySchema.shape.languages, + runtime: embeddingModelCatalogEntrySchema.shape.runtime, + dimensions: embeddingModelCatalogEntrySchema.shape.dimensions, + contextTokens: embeddingModelCatalogEntrySchema.shape.contextTokens, + quantization: embeddingModelCatalogEntrySchema.shape.quantization, + recommended: embeddingModelCatalogEntrySchema.shape.recommended, + available: embeddingModelCatalogEntrySchema.shape.available, + unavailableReason: + embeddingModelCatalogEntrySchema.shape.unavailableReason, + license: embeddingModelCatalogEntrySchema.shape.license, + files: z.array(embeddingModelCatalogViewFileSchema).min(1).max(32), + downloadAvailability: z + .array(modelDownloadAvailabilitySchema) + .length(MODEL_DOWNLOAD_SOURCES.length) + }) + .strict() + +export const installedEmbeddingModelFileSchema = z + .object({ + name: modelFileNameSchema, + role: embeddingModelFileRoleSchema, + size: z.number().int().positive().safe(), + sha256: sha256Schema + }) + .strict() + +export const installedEmbeddingModelSchema = z + .object({ + id: embeddingModelIdSchema, + displayName: z.string().trim().min(1).max(120), + source: z.enum(['download', 'local']), + installedAt: z.string().datetime(), + files: z.array(installedEmbeddingModelFileSchema).min(1).max(32) + }) + .strict() + +export const embeddingModelOperationSchema = z + .object({ + modelId: embeddingModelIdSchema, + kind: z.enum(['download', 'import']), + phase: z.enum(['preparing', 'transferring', 'installing']), + currentFile: modelFileNameSchema.nullable(), + completedBytes: z.number().int().nonnegative().safe(), + totalBytes: z.number().int().positive().safe().nullable(), + downloadSource: modelDownloadSourceSchema.optional() + }) + .strict() + +export const embeddingModelSnapshotSchema = z + .object({ + selectedDownloadSource: modelDownloadSourceSchema, + catalog: z.array(embeddingModelCatalogViewEntrySchema), + installed: z.array(installedEmbeddingModelSchema), + operations: z.array(embeddingModelOperationSchema) + }) + .strict() + +export const embeddingModelProgressSnapshotSchema = z + .object({ + operations: z.array(embeddingModelOperationSchema) + }) + .strict() + +export const embeddingModelStatusSchema = z + .object({ + id: embeddingModelIdSchema, + displayName: z.string().trim().min(1).max(120), + catalogAvailable: z.boolean(), + installed: z.boolean(), + verified: z.boolean(), + detail: z.string().trim().min(1).max(1_000) + }) + .strict() + +export type EmbeddingModelCatalogEntry = z.infer< + typeof embeddingModelCatalogEntrySchema +> +export type EmbeddingModelCatalogViewEntry = z.infer< + typeof embeddingModelCatalogViewEntrySchema +> +export type EmbeddingModelFile = z.infer +export type InstalledEmbeddingModel = z.infer< + typeof installedEmbeddingModelSchema +> +export type EmbeddingModelOperation = z.infer< + typeof embeddingModelOperationSchema +> +export type EmbeddingModelSnapshot = z.infer< + typeof embeddingModelSnapshotSchema +> +export type EmbeddingModelProgressSnapshot = z.infer< + typeof embeddingModelProgressSnapshotSchema +> +export type EmbeddingModelStatus = z.infer< + typeof embeddingModelStatusSchema +> diff --git a/src/main/knowledge/embedding-model-manager.test.ts b/src/main/knowledge/embedding-model-manager.test.ts new file mode 100644 index 00000000..fd747be5 --- /dev/null +++ b/src/main/knowledge/embedding-model-manager.test.ts @@ -0,0 +1,492 @@ +import { createHash } from 'node:crypto' +import { + mkdtemp, + mkdir, + readFile, + readdir, + rm, + symlink, + writeFile +} from 'node:fs/promises' +import { tmpdir } from 'node:os' +import { join } from 'node:path' +import { zipSync } from 'fflate' +import { afterEach, describe, expect, it } from 'vitest' +import { + MODEL_DOWNLOAD_SOURCES, + resolveModelDownloadPackage +} from '../../shared/model-download-contracts' +import type { EmbeddingModelCatalogEntry } from './embedding-model-contracts' +import { EMBEDDING_MODEL_CATALOG } from './embedding-model-catalog' +import { EmbeddingModelManager } from './embedding-model-manager' + +const roots: string[] = [] + +function sha256(contents: string): string { + return createHash('sha256').update(contents).digest('hex') +} + +function testCatalog( + contents: Record +): EmbeddingModelCatalogEntry[] { + const roles = [ + 'model', + 'tokenizer', + 'tokenizer-configuration', + 'configuration', + 'license' + ] as const + return [ + { + id: 'complete-test-model', + displayName: 'Complete Test Model', + description: 'A complete injected test model.', + languages: ['English'], + runtime: 'onnxruntime-web/wasm', + dimensions: 4, + contextTokens: 128, + quantization: 'test', + recommended: false, + available: true, + repositoryUrls: {}, + license: { + name: 'Test license', + notice: 'Tests only.', + url: 'https://example.com/license' + }, + files: roles.map((role, index) => { + const name = Object.keys(contents)[index]! + const value = contents[name]! + return { + name, + role, + size: Buffer.byteLength(value), + sha256: sha256(value), + targets: { + modelscope: { + url: `https://modelscope.cn/models/test/model/resolve/${'a'.repeat(40)}/${name}`, + repositoryUrl: + 'https://modelscope.cn/models/test/model', + revision: 'a'.repeat(40), + redirectHosts: [] + }, + 'hugging-face': { + url: `https://huggingface.co/test/model/resolve/${'b'.repeat(40)}/${name}`, + repositoryUrl: 'https://huggingface.co/test/model', + revision: 'b'.repeat(40), + redirectHosts: [] + } + } + } + }) + } + ] +} + +function embeddingArchive( + contents: Record, + overrides: { + kind?: string + modelId?: string + extraEntries?: Record + hashOverrides?: Record + } = {} +): Uint8Array { + const entry = testCatalog(contents)[0]! + const files = entry.files.map(({ name, role, size, sha256 }) => ({ + name, + role, + size, + sha256: overrides.hashOverrides?.[name] ?? sha256 + })) + return zipSync({ + 'goodbuddy-model.json': Buffer.from( + JSON.stringify({ + format: 'goodbuddy-model-archive', + version: 1, + kind: overrides.kind ?? 'embedding', + modelId: overrides.modelId ?? entry.id, + displayName: entry.displayName, + exportedAt: '2026-08-20T00:00:00.000Z', + files + }) + ), + ...Object.fromEntries( + Object.entries(contents).map(([name, value]) => [ + name, + Buffer.from(value) + ]) + ), + ...overrides.extraEntries + }) +} + +async function fixture(): Promise<{ + root: string + source: string + manager: EmbeddingModelManager + contents: Record +}> { + const root = await mkdtemp(join(tmpdir(), 'embedding-model-')) + roots.push(root) + const source = join(root, 'source') + await mkdir(source) + const contents = { + 'model.onnx': 'model', + 'tokenizer.json': 'tokenizer', + 'tokenizer_config.json': 'tokenizer config', + 'config.json': 'model config', + LICENSE: 'license' + } + for (const [name, value] of Object.entries(contents)) { + await writeFile(join(source, name), value) + } + return { + root, + source, + contents, + manager: new EmbeddingModelManager({ + userDataDirectory: root, + catalog: testCatalog(contents) + }) + } +} + +afterEach(async () => { + await Promise.all( + roots.splice(0).map((root) => + rm(root, { recursive: true, force: true }) + ) + ) +}) + +describe('EmbeddingModelManager', () => { + it('publishes a complete downloadable production Granite package', async () => { + expect(EMBEDDING_MODEL_CATALOG).toHaveLength(1) + expect(EMBEDDING_MODEL_CATALOG[0]).toMatchObject({ + available: true, + runtime: 'onnxruntime-web/wasm' + }) + expect( + EMBEDDING_MODEL_CATALOG[0]!.files.map( + ({ name, size, sha256 }) => ({ name, size, sha256 }) + ) + ).toEqual([ + { + name: 'model_quantized.onnx', + size: 97_858_099, + sha256: + '704c1ebca5fbb7cd83ced41827658ac4c9990c64f7f2874d22b78044e5022e22' + }, + { + name: 'tokenizer.json', + size: 25_301_671, + sha256: + '51947676cae1f991fa51c6b9a24e14ee5460e5f0b9f692f13bb3159829d1592a' + }, + { + name: 'tokenizer_config.json', + size: 12_860, + sha256: + '6ed69389e30a8ecabfce2f9ebcdf0c908b34056f24d994340f2f216521c057d5' + }, + { + name: 'config.json', + size: 1_215, + sha256: + 'ae74d55a56f779774cb9a8e63d3c2da9ae1af83c00229ffdff43d0b38407a0ee' + }, + { + name: 'special_tokens_map.json', + size: 871, + sha256: + '013787ee251ff611722479197c00853b62113ad303cb0a36524231783c676c69' + } + ]) + for (const source of MODEL_DOWNLOAD_SOURCES) { + const resolved = resolveModelDownloadPackage( + EMBEDDING_MODEL_CATALOG[0]!.files, + source + ) + expect(resolved.files.map((file) => file.name)).toEqual([ + 'model_quantized.onnx', + 'tokenizer.json', + 'tokenizer_config.json', + 'config.json', + 'special_tokens_map.json' + ]) + expect(resolved.files.every((file) => file.target.revision.length === 40)) + .toBe(true) + } + + const root = await mkdtemp(join(tmpdir(), 'embedding-catalog-')) + roots.push(root) + const snapshot = await new EmbeddingModelManager({ + userDataDirectory: root + }).getSnapshot() + expect(snapshot.catalog[0]!.downloadAvailability).toEqual([ + expect.objectContaining({ available: true }), + expect.objectContaining({ available: true }) + ]) + }) + + it('downloads, verifies, and atomically publishes a package', async () => { + const root = await mkdtemp(join(tmpdir(), 'embedding-download-')) + roots.push(root) + const contents: Record = { + 'model.onnx': 'model', + 'tokenizer.json': 'tokenizer', + 'tokenizer_config.json': 'tokenizer config', + 'config.json': 'model config', + LICENSE: 'license' + } + const requests: string[] = [] + const manager = new EmbeddingModelManager({ + userDataDirectory: root, + catalog: testCatalog(contents), + fetch: async (input) => { + const url = input.toString() + requests.push(url) + const name = decodeURIComponent(new URL(url).pathname.split('/').at(-1)!) + const body = contents[name] + if (body === undefined) { + return new Response(null, { status: 404 }) + } + return new Response(body, { + headers: { 'content-length': String(Buffer.byteLength(body)) } + }) + } + }) + + const installed = await manager.install( + 'complete-test-model', + 'modelscope' + ) + expect(installed).toMatchObject({ + source: 'download', + files: expect.arrayContaining([ + expect.objectContaining({ name: 'model.onnx' }), + expect.objectContaining({ name: 'tokenizer.json' }) + ]) + }) + expect(requests).toHaveLength(5) + expect(await manager.getVerifiedModelDirectory('complete-test-model')) + .toBe(join(manager.rootDirectory, 'complete-test-model')) + expect( + await readFile( + join(manager.rootDirectory, 'complete-test-model', 'config.json'), + 'utf8' + ) + ).toBe(contents['config.json']) + expect( + ( + await readdir(manager.rootDirectory) + ).some((name) => name.startsWith('.install-')) + ).toBe(false) + }) + + it('installs a complete local package and exposes only a verified Main path', async () => { + const { manager, source } = await fixture() + const installed = await manager.registerLocalDirectory( + 'complete-test-model', + source + ) + expect(installed.files).toHaveLength(5) + + const snapshot = await manager.getSnapshot() + expect(snapshot).not.toHaveProperty('rootDirectory') + expect(JSON.stringify(snapshot)).not.toContain(manager.rootDirectory) + expect(snapshot.catalog[0]).not.toHaveProperty('repositoryUrls') + expect(snapshot.catalog[0]!.files[0]).not.toHaveProperty('targets') + expect(snapshot.installed).toHaveLength(1) + + const trusted = await manager.getVerifiedModelDirectory( + 'complete-test-model' + ) + expect(trusted).toBe( + join(manager.rootDirectory, 'complete-test-model') + ) + expect(await manager.getStatus('complete-test-model')).toMatchObject({ + catalogAvailable: true, + installed: true, + verified: true + }) + }) + + it('imports and atomically publishes a verified embedding ZIP', async () => { + const { manager, root, contents } = await fixture() + const archive = join(root, 'embedding.zip') + await writeFile(archive, embeddingArchive(contents)) + + await expect( + manager.importArchive('complete-test-model', archive) + ).resolves.toMatchObject({ + id: 'complete-test-model', + source: 'local', + files: expect.arrayContaining([ + expect.objectContaining({ name: 'model.onnx' }), + expect.objectContaining({ name: 'tokenizer.json' }) + ]) + }) + + const directory = await manager.getVerifiedModelDirectory( + 'complete-test-model' + ) + expect(await readdir(directory)).toHaveLength(6) + await expect( + readFile(join(directory, 'goodbuddy-model.json')) + ).rejects.toThrow() + expect( + (await readdir(manager.rootDirectory)).some((name) => + name.startsWith('.install-') + ) + ).toBe(false) + }) + + it('rejects mismatched ZIP descriptors and hashes without publishing', async () => { + const { manager, root, contents } = await fixture() + const wrongKind = join(root, 'wrong-kind.zip') + const wrongHash = join(root, 'wrong-hash.zip') + await Promise.all([ + writeFile( + wrongKind, + embeddingArchive(contents, { kind: 'speech' }) + ), + writeFile( + wrongHash, + embeddingArchive(contents, { + hashOverrides: { 'model.onnx': 'a'.repeat(64) } + }) + ) + ]) + + await expect( + manager.importArchive('complete-test-model', wrongKind) + ).rejects.toThrow(/类型或模型 ID 不匹配/u) + await expect( + manager.importArchive('complete-test-model', wrongHash) + ).rejects.toThrow(/校验失败/u) + expect((await manager.getSnapshot()).installed).toEqual([]) + expect( + (await readdir(manager.rootDirectory)).some((name) => + name.startsWith('.install-') + ) + ).toBe(false) + }) + + it('rejects undeclared ZIP entries and preserves an installed model', async () => { + const { manager, source, root, contents } = await fixture() + await manager.registerLocalDirectory('complete-test-model', source) + const directory = await manager.getVerifiedModelDirectory( + 'complete-test-model' + ) + const archive = join(root, 'unsafe.zip') + await writeFile( + archive, + embeddingArchive(contents, { + extraEntries: { + '../outside.exe': Buffer.from('unsafe') + } + }) + ) + + await expect( + manager.importArchive('complete-test-model', archive) + ).rejects.toThrow(/已安装/u) + await expect( + manager.getVerifiedModelDirectory('complete-test-model') + ).resolves.toBe(directory) + await expect( + readFile(join(directory, 'model.onnx'), 'utf8') + ).resolves.toBe(contents['model.onnx']) + }) + + it('honors an already-aborted ZIP import and cleans staging', async () => { + const { manager, root, contents } = await fixture() + const archive = join(root, 'embedding.zip') + await writeFile(archive, embeddingArchive(contents)) + const controller = new AbortController() + controller.abort() + + await expect( + manager.importArchive( + 'complete-test-model', + archive, + controller.signal + ) + ).rejects.toMatchObject({ name: 'AbortError' }) + expect(manager.getProgressSnapshot().operations).toEqual([]) + expect( + (await readdir(manager.rootDirectory)).some((name) => + name.startsWith('.install-') + ) + ).toBe(false) + }) + + it('rejects wrong hashes and leaves no installed model', async () => { + const { manager, source } = await fixture() + await writeFile(join(source, 'model.onnx'), 'tampered') + await expect( + manager.registerLocalDirectory('complete-test-model', source) + ).rejects.toThrow(/校验失败/u) + expect((await manager.getSnapshot()).installed).toEqual([]) + }) + + it('re-verifies installed bytes before returning the trusted directory', async () => { + const { manager, source } = await fixture() + await manager.registerLocalDirectory('complete-test-model', source) + const directory = join( + manager.rootDirectory, + 'complete-test-model' + ) + await writeFile(join(directory, 'tokenizer.json'), 'tampered') + await expect( + manager.getVerifiedModelDirectory('complete-test-model') + ).rejects.toThrow(/校验失败/u) + expect(await manager.getStatus('complete-test-model')).toMatchObject({ + installed: true, + verified: false + }) + }) + + it('rejects source symlinks', async () => { + const { manager, source, root } = await fixture() + const linked = join(root, 'linked') + await symlink(source, linked, 'junction') + await expect( + manager.registerLocalDirectory('complete-test-model', linked) + ).rejects.toThrow(/普通目录/u) + }) + + it('honors an already-aborted local import signal', async () => { + const { manager, source } = await fixture() + const controller = new AbortController() + controller.abort() + await expect( + manager.registerLocalDirectory( + 'complete-test-model', + source, + controller.signal + ) + ).rejects.toMatchObject({ name: 'AbortError' }) + expect(manager.getProgressSnapshot().operations).toEqual([]) + }) + + it('does not trust a modified manifest', async () => { + const { manager, source } = await fixture() + await manager.registerLocalDirectory('complete-test-model', source) + const manifestPath = join( + manager.rootDirectory, + 'complete-test-model', + 'manifest.json' + ) + const manifest = JSON.parse(await readFile(manifestPath, 'utf8')) as { + id: string + } + manifest.id = 'other-model' + await writeFile(manifestPath, JSON.stringify(manifest)) + await expect( + manager.getVerifiedModelDirectory('complete-test-model') + ).rejects.toThrow() + }) +}) diff --git a/src/main/knowledge/embedding-model-manager.ts b/src/main/knowledge/embedding-model-manager.ts new file mode 100644 index 00000000..f3a94b05 --- /dev/null +++ b/src/main/knowledge/embedding-model-manager.ts @@ -0,0 +1,824 @@ +import { createHash } from 'node:crypto' +import { + copyFile, + lstat, + mkdir, + open, + readFile, + readdir, + rename, + rm, + stat, + writeFile +} from 'node:fs/promises' +import { resolve } from 'node:path' +import { + MODEL_DOWNLOAD_SOURCES, + getMaximumModelPackageBytes, + getModelDownloadAvailability, + resolveModelDownloadPackage, + type ModelDownloadAvailability, + type ModelDownloadSource, + type ResolvedModelArtifactFile +} from '../../shared/model-download-contracts' +import { extractModelArchive } from '../model-archive' +import { fetchModelDownloadResponse } from '../model-download-transport' +import { + MODEL_PARTIAL_SUFFIX, + attachModelAbortSignal, + cleanupStaleModelInstallArtifacts, + createModelStagingDirectory, + ensureModelOperationNotAborted, + hashModelFile, + managedModelChild, + writeModelBuffer +} from '../model-package-utils' +import { isMissingFileError } from '../settings-file-utils' +import { EMBEDDING_MODEL_CATALOG } from './embedding-model-catalog' +import { + embeddingModelCatalogEntrySchema, + embeddingModelCatalogViewEntrySchema, + embeddingModelIdSchema, + embeddingModelProgressSnapshotSchema, + embeddingModelSnapshotSchema, + embeddingModelStatusSchema, + installedEmbeddingModelSchema, + type EmbeddingModelCatalogEntry, + type EmbeddingModelCatalogViewEntry, + type EmbeddingModelOperation, + type EmbeddingModelProgressSnapshot, + type EmbeddingModelSnapshot, + type EmbeddingModelStatus, + type InstalledEmbeddingModel +} from './embedding-model-contracts' + +const MANIFEST_FILE_NAME = 'manifest.json' +const ARCHIVE_MANIFEST_FILE_NAME = 'goodbuddy-model.json' +const DEFAULT_MAX_FILE_BYTES = 2 * 1024 * 1024 * 1024 +const MAXIMUM_ARCHIVE_BYTES = 4 * 1024 * 1024 * 1024 - 1 +const ARCHIVE_OVERHEAD_BYTES = 1024 * 1024 +const executableExtensionPattern = + /\.(?:app|bat|bin|cmd|com|cpl|dll|dmg|exe|hta|iso|jar|js|jse|lnk|msi|ps1|reg|scr|sh|sys|vb|vbe|vbs|wsf)$/iu + +type ActiveOperation = { + controller: AbortController + progress: EmbeddingModelOperation +} + +export type EmbeddingModelManagerOptions = { + userDataDirectory: string + catalog?: readonly EmbeddingModelCatalogEntry[] + fetch?: typeof fetch + getDownloadSource?: () => + | ModelDownloadSource + | Promise + maxFileBytes?: number +} + +function safeChild(parent: string, name: string): string { + return managedModelChild( + parent, + name, + '向量模型路径超出受管目录' + ) +} + +function unavailableDownloadAvailability( + source: ModelDownloadSource, + reason: string +): ModelDownloadAvailability { + return { source, available: false, unavailableReason: reason } +} + +function toCatalogView( + entry: EmbeddingModelCatalogEntry +): EmbeddingModelCatalogViewEntry { + const { repositoryUrls, files, ...metadata } = entry + void repositoryUrls + return embeddingModelCatalogViewEntrySchema.parse({ + ...metadata, + files: files.map(({ targets, ...file }) => { + void targets + return file + }), + downloadAvailability: MODEL_DOWNLOAD_SOURCES.map((source) => + entry.available + ? getModelDownloadAvailability(files, source) + : unavailableDownloadAvailability( + source, + entry.unavailableReason ?? + '当前内置向量模型尚不可安装' + ) + ) + }) +} + +function packageMatches( + entry: EmbeddingModelCatalogEntry, + files: readonly { + name: string + role: string + size: number + sha256: string + }[] +): boolean { + return ( + files.length === entry.files.length && + entry.files.every((expected) => + files.some( + (actual) => + actual.name === expected.name && + actual.role === expected.role && + actual.size === expected.size && + actual.sha256 === expected.sha256 + ) + ) + ) +} + +export class EmbeddingModelManager { + readonly rootDirectory: string + + private readonly catalog: EmbeddingModelCatalogEntry[] + private readonly catalogViews: EmbeddingModelCatalogViewEntry[] + private readonly transport: typeof fetch + private readonly getDownloadSource: () => + | ModelDownloadSource + | Promise + private readonly maxFileBytes: number + private readonly operations = new Map() + + constructor(options: EmbeddingModelManagerOptions) { + if (!options.userDataDirectory.trim()) { + throw new Error('userDataDirectory is required') + } + this.rootDirectory = resolve( + options.userDataDirectory, + 'models', + 'embedding' + ) + this.catalog = (options.catalog ?? EMBEDDING_MODEL_CATALOG).map( + (entry) => embeddingModelCatalogEntrySchema.parse(entry) + ) + if ( + new Set(this.catalog.map((entry) => entry.id)).size !== + this.catalog.length + ) { + throw new Error('向量模型目录包含重复 ID') + } + this.catalogViews = this.catalog.map(toCatalogView) + this.transport = options.fetch ?? globalThis.fetch + this.getDownloadSource = + options.getDownloadSource ?? (() => 'modelscope') + this.maxFileBytes = options.maxFileBytes ?? DEFAULT_MAX_FILE_BYTES + if ( + !Number.isSafeInteger(this.maxFileBytes) || + this.maxFileBytes <= 0 || + this.maxFileBytes > 8 * 1024 * 1024 * 1024 + ) { + throw new RangeError('maxFileBytes must be a positive safe integer') + } + } + + async getSnapshot(): Promise { + await this.ensureRoot() + await this.cleanupStaleArtifacts() + const [selectedDownloadSource, installed] = await Promise.all([ + this.getDownloadSource(), + this.readInstalled() + ]) + return embeddingModelSnapshotSchema.parse({ + selectedDownloadSource, + catalog: this.catalogViews, + installed, + operations: [...this.operations.values()].map(({ progress }) => ({ + ...progress + })) + }) + } + + async snapshot(): Promise { + return this.getSnapshot() + } + + getProgressSnapshot(): EmbeddingModelProgressSnapshot { + return embeddingModelProgressSnapshotSchema.parse({ + operations: [...this.operations.values()].map(({ progress }) => ({ + ...progress + })) + }) + } + + async getStatus(modelId: string): Promise { + const id = embeddingModelIdSchema.parse(modelId) + const entry = this.catalog.find((candidate) => candidate.id === id) + if (!entry) { + return embeddingModelStatusSchema.parse({ + id, + displayName: id, + catalogAvailable: false, + installed: false, + verified: false, + detail: '当前版本不提供此内置向量模型' + }) + } + try { + await this.verifyInstalledModel(entry) + return embeddingModelStatusSchema.parse({ + id, + displayName: entry.displayName, + catalogAvailable: entry.available, + installed: true, + verified: true, + detail: '模型已安装并通过 SHA-256 校验，可供本地推理使用' + }) + } catch { + const installed = await this.hasInstalledDirectory(id) + return embeddingModelStatusSchema.parse({ + id, + displayName: entry.displayName, + catalogAvailable: entry.available, + installed, + verified: false, + detail: installed + ? '模型文件校验失败' + : entry.available + ? '模型尚未安装' + : entry.unavailableReason! + }) + } + } + + /** + * Returns a Main-owned absolute path only after re-verifying every package + * file. This path must not be sent to the renderer. + */ + async getVerifiedModelDirectory(modelId: string): Promise { + const entry = this.requireAvailableEntry(modelId) + await this.verifyInstalledModel(entry) + return this.modelDirectory(entry.id) + } + + async registerLocalDirectory( + modelId: string, + sourceDirectory: string, + externalSignal?: AbortSignal + ): Promise { + const entry = this.requireAvailableEntry(modelId) + const source = resolve(sourceDirectory) + const operation = this.beginOperation(entry.id, 'import', null) + const detachAbort = attachModelAbortSignal( + externalSignal, + operation.controller + ) + let stagingDirectory: string | undefined + try { + await this.ensureRoot() + await this.assertNotInstalled(entry.id) + await this.validateLocalDirectory( + source, + entry, + operation.controller.signal + ) + stagingDirectory = await createModelStagingDirectory( + this.rootDirectory, + entry.id, + '向量模型路径超出受管目录' + ) + operation.progress.phase = 'transferring' + for (const file of entry.files) { + ensureModelOperationNotAborted(operation.controller.signal) + operation.progress.currentFile = file.name + const destination = safeChild(stagingDirectory, file.name) + await copyFile(safeChild(source, file.name), destination) + ensureModelOperationNotAborted(operation.controller.signal) + operation.progress.completedBytes += (await stat(destination)).size + } + operation.progress.totalBytes = operation.progress.completedBytes + operation.progress.phase = 'installing' + operation.progress.currentFile = null + const installed = await this.writeVerifiedManifest( + entry, + 'local', + stagingDirectory, + operation.controller.signal + ) + ensureModelOperationNotAborted(operation.controller.signal) + await rename(stagingDirectory, this.modelDirectory(entry.id)) + stagingDirectory = undefined + return installed + } finally { + detachAbort() + this.operations.delete(entry.id) + if (stagingDirectory) { + await rm(stagingDirectory, { recursive: true, force: true }) + } + } + } + + async install( + modelId: string, + downloadSource?: ModelDownloadSource, + externalSignal?: AbortSignal + ): Promise { + const entry = this.requireAvailableEntry(modelId) + const selectedDownloadSource = + downloadSource ?? (await this.getDownloadSource()) + const resolvedPackage = resolveModelDownloadPackage( + entry.files, + selectedDownloadSource + ) + const operation = this.beginOperation( + entry.id, + 'download', + resolvedPackage.totalBytes, + resolvedPackage.source + ) + const detachAbort = attachModelAbortSignal( + externalSignal, + operation.controller + ) + let stagingDirectory: string | undefined + try { + await this.ensureRoot() + await this.assertNotInstalled(entry.id) + stagingDirectory = await createModelStagingDirectory( + this.rootDirectory, + entry.id, + '向量模型路径超出受管目录' + ) + for (const file of resolvedPackage.files) { + ensureModelOperationNotAborted(operation.controller.signal) + operation.progress.phase = 'transferring' + operation.progress.currentFile = file.name + await this.downloadFile( + file, + safeChild(stagingDirectory, file.name), + operation, + operation.controller.signal + ) + } + operation.progress.phase = 'installing' + operation.progress.currentFile = null + const installed = await this.writeVerifiedManifest( + entry, + 'download', + stagingDirectory, + operation.controller.signal + ) + ensureModelOperationNotAborted(operation.controller.signal) + await rename(stagingDirectory, this.modelDirectory(entry.id)) + stagingDirectory = undefined + return installed + } finally { + detachAbort() + this.operations.delete(entry.id) + if (stagingDirectory) { + await rm(stagingDirectory, { recursive: true, force: true }) + } + } + } + + async importArchive( + modelId: string, + archivePath: string, + externalSignal?: AbortSignal + ): Promise { + const entry = this.requireAvailableEntry(modelId) + if ( + entry.files.some((file) => + executableExtensionPattern.test(file.name) + ) + ) { + throw new Error('向量模型目录包含不安全文件') + } + const expectedTotal = getMaximumModelPackageBytes(entry.files) + const maximumTotalBytes = Math.min( + MAXIMUM_ARCHIVE_BYTES, + expectedTotal + ARCHIVE_OVERHEAD_BYTES + ) + const operation = this.beginOperation( + entry.id, + 'import', + expectedTotal + ) + const detachAbort = attachModelAbortSignal( + externalSignal, + operation.controller + ) + let stagingDirectory: string | undefined + try { + await this.ensureRoot() + await this.assertNotInstalled(entry.id) + stagingDirectory = await createModelStagingDirectory( + this.rootDirectory, + entry.id, + '向量模型路径超出受管目录' + ) + operation.progress.phase = 'transferring' + const descriptor = await extractModelArchive({ + archivePath, + destinationDirectory: stagingDirectory, + expectedKind: 'embedding', + expectedModelId: entry.id, + expectedFiles: entry.files.map((file) => ({ + name: file.name, + role: file.role + })), + maximumArchiveBytes: Math.min( + MAXIMUM_ARCHIVE_BYTES, + maximumTotalBytes + ARCHIVE_OVERHEAD_BYTES + ), + maximumFileBytes: this.maxFileBytes, + maximumTotalBytes, + signal: operation.controller.signal, + onProgress: (completedBytes) => { + operation.progress.completedBytes = completedBytes + } + }) + if ( + descriptor.files.some( + (file) => file.size > this.maxFileBytes + ) || + !packageMatches(entry, descriptor.files) + ) { + throw new Error('向量模型 ZIP 与当前模型目录不匹配') + } + operation.progress.totalBytes = descriptor.files.reduce( + (total, file) => total + file.size, + 0 + ) + operation.progress.phase = 'installing' + operation.progress.currentFile = null + await rm( + safeChild(stagingDirectory, ARCHIVE_MANIFEST_FILE_NAME), + { force: true } + ) + const installed = await this.writeVerifiedManifest( + entry, + 'local', + stagingDirectory, + operation.controller.signal + ) + ensureModelOperationNotAborted(operation.controller.signal) + await rename(stagingDirectory, this.modelDirectory(entry.id)) + stagingDirectory = undefined + return installed + } finally { + detachAbort() + this.operations.delete(entry.id) + if (stagingDirectory) { + await rm(stagingDirectory, { recursive: true, force: true }) + } + } + } + + cancel(modelId: string): boolean { + const id = embeddingModelIdSchema.parse(modelId) + const operation = this.operations.get(id) + if (!operation) { + return false + } + operation.controller.abort() + return true + } + + async remove(modelId: string): Promise { + const id = embeddingModelIdSchema.parse(modelId) + this.cancel(id) + await rm(this.modelDirectory(id), { recursive: true, force: true }) + } + + dispose(): void { + for (const operation of this.operations.values()) { + operation.controller.abort() + } + this.operations.clear() + } + + private async ensureRoot(): Promise { + await mkdir(this.rootDirectory, { recursive: true }) + } + + private modelDirectory(modelId: string): string { + return safeChild( + this.rootDirectory, + embeddingModelIdSchema.parse(modelId) + ) + } + + private requireAvailableEntry( + modelId: string + ): EmbeddingModelCatalogEntry { + const id = embeddingModelIdSchema.parse(modelId) + const entry = this.catalog.find((candidate) => candidate.id === id) + if (!entry) { + throw new Error('未知的向量模型') + } + if (!entry.available) { + throw new Error( + entry.unavailableReason ?? '当前向量模型尚不可安装' + ) + } + return entry + } + + private beginOperation( + modelId: string, + kind: EmbeddingModelOperation['kind'], + totalBytes: number | null, + downloadSource?: ModelDownloadSource + ): ActiveOperation { + if (this.operations.has(modelId)) { + throw new Error('该向量模型已有进行中的操作') + } + const operation: ActiveOperation = { + controller: new AbortController(), + progress: { + modelId, + kind, + phase: 'preparing', + currentFile: null, + completedBytes: 0, + totalBytes, + ...(downloadSource ? { downloadSource } : {}) + } + } + this.operations.set(modelId, operation) + return operation + } + + private async downloadFile( + file: ResolvedModelArtifactFile< + EmbeddingModelCatalogEntry['files'][number]['role'] + >, + destination: string, + operation: ActiveOperation, + signal: AbortSignal + ): Promise { + if (file.size <= 0 || file.size > this.maxFileBytes) { + throw new RangeError(`向量模型文件过大：${file.name}`) + } + const response = await fetchModelDownloadResponse({ + transport: this.transport, + initialUrl: file.target.url, + redirectHosts: file.target.redirectHosts, + signal, + modelLabel: '向量模型' + }) + if (!response.ok) { + await response.body?.cancel().catch(() => undefined) + throw new Error(`向量模型下载失败：HTTP ${response.status}`) + } + if (!response.body) { + throw new Error('向量模型下载响应没有内容') + } + const declaredLength = response.headers.get('content-length') + if (declaredLength !== null) { + const parsedLength = Number(declaredLength) + if ( + !Number.isSafeInteger(parsedLength) || + parsedLength !== file.size + ) { + await response.body.cancel().catch(() => undefined) + throw new Error(`向量模型文件大小不匹配：${file.name}`) + } + } + + const partialPath = `${destination}${MODEL_PARTIAL_SUFFIX}` + const handle = await open(partialPath, 'wx') + const reader = response.body.getReader() + const hash = createHash('sha256') + let written = 0 + try { + while (true) { + ensureModelOperationNotAborted(signal) + const result = await reader.read() + if (result.done) { + break + } + if ( + written + result.value.byteLength > file.size || + written + result.value.byteLength > this.maxFileBytes + ) { + await reader.cancel() + throw new RangeError(`向量模型文件过大：${file.name}`) + } + written += await writeModelBuffer( + handle, + result.value, + (persisted) => { + hash.update(persisted) + operation.progress.completedBytes += persisted.byteLength + } + ) + } + } catch (error) { + await reader.cancel().catch(() => undefined) + throw error + } finally { + await handle.close() + } + if (written !== file.size) { + throw new Error(`向量模型文件大小不匹配：${file.name}`) + } + if (hash.digest('hex') !== file.sha256) { + throw new Error(`向量模型文件校验失败：${file.name}`) + } + await rename(partialPath, destination) + } + + private async assertNotInstalled(modelId: string): Promise { + try { + await lstat(this.modelDirectory(modelId)) + throw new Error('向量模型已安装') + } catch (error) { + if (isMissingFileError(error)) { + return + } + throw error + } + } + + private async hasInstalledDirectory(modelId: string): Promise { + try { + const info = await lstat(this.modelDirectory(modelId)) + return info.isDirectory() && !info.isSymbolicLink() + } catch (error) { + if (isMissingFileError(error)) { + return false + } + throw error + } + } + + private async validateLocalDirectory( + source: string, + entry: EmbeddingModelCatalogEntry, + signal: AbortSignal + ): Promise { + const sourceInfo = await lstat(source) + if (!sourceInfo.isDirectory() || sourceInfo.isSymbolicLink()) { + throw new Error('本地向量模型来源必须是普通目录') + } + const directoryEntries = await readdir(source, { + withFileTypes: true + }) + if (directoryEntries.length !== entry.files.length) { + throw new Error('本地向量模型目录包含未声明或缺失的文件') + } + for (const localEntry of directoryEntries) { + ensureModelOperationNotAborted(signal) + if ( + localEntry.isSymbolicLink() || + !localEntry.isFile() || + executableExtensionPattern.test(localEntry.name) + ) { + throw new Error('本地向量模型目录包含不安全文件') + } + } + for (const file of entry.files) { + ensureModelOperationNotAborted(signal) + const path = safeChild(source, file.name) + const info = await lstat(path) + if ( + !info.isFile() || + info.isSymbolicLink() || + info.size <= 0 || + info.size > this.maxFileBytes + ) { + throw new Error(`向量模型文件无效：${file.name}`) + } + const actual = await hashModelFile(path, signal) + if ( + actual.size !== file.size || + actual.sha256 !== file.sha256 + ) { + throw new Error(`向量模型文件校验失败：${file.name}`) + } + } + } + + private async writeVerifiedManifest( + entry: EmbeddingModelCatalogEntry, + source: InstalledEmbeddingModel['source'], + directory: string, + signal: AbortSignal + ): Promise { + const files: InstalledEmbeddingModel['files'] = [] + for (const file of entry.files) { + ensureModelOperationNotAborted(signal) + const actual = await hashModelFile( + safeChild(directory, file.name), + signal + ) + files.push({ name: file.name, role: file.role, ...actual }) + } + if (!packageMatches(entry, files)) { + throw new Error('向量模型文件与当前目录不匹配') + } + const installed = installedEmbeddingModelSchema.parse({ + id: entry.id, + displayName: entry.displayName, + source, + installedAt: new Date().toISOString(), + files + }) + await writeFile( + safeChild(directory, MANIFEST_FILE_NAME), + `${JSON.stringify(installed, null, 2)}\n`, + { encoding: 'utf8', flag: 'wx' } + ) + return installed + } + + private async readInstalled(): Promise { + const entries = await readdir(this.rootDirectory, { + withFileTypes: true + }) + const installed: InstalledEmbeddingModel[] = [] + for (const entry of entries) { + if ( + !entry.isDirectory() || + entry.isSymbolicLink() || + entry.name.startsWith('.install-') || + !embeddingModelIdSchema.safeParse(entry.name).success + ) { + continue + } + try { + const manifest = installedEmbeddingModelSchema.parse( + JSON.parse( + await readFile( + safeChild( + this.modelDirectory(entry.name), + MANIFEST_FILE_NAME + ), + 'utf8' + ) + ) as unknown + ) + if (manifest.id === entry.name) { + installed.push(manifest) + } + } catch { + // Ignore incomplete or externally modified directories. + } + } + return installed + } + + private async verifyInstalledModel( + entry: EmbeddingModelCatalogEntry + ): Promise { + const directory = this.modelDirectory(entry.id) + const directoryInfo = await lstat(directory) + if (!directoryInfo.isDirectory() || directoryInfo.isSymbolicLink()) { + throw new Error('向量模型目录无效') + } + const manifestPath = safeChild(directory, MANIFEST_FILE_NAME) + const manifestInfo = await lstat(manifestPath) + if (!manifestInfo.isFile() || manifestInfo.isSymbolicLink()) { + throw new Error('向量模型清单无效') + } + const manifest = installedEmbeddingModelSchema.parse( + JSON.parse(await readFile(manifestPath, 'utf8')) as unknown + ) + if (manifest.id !== entry.id || !packageMatches(entry, manifest.files)) { + throw new Error('向量模型清单与目录不匹配') + } + const actualEntries = await readdir(directory, { withFileTypes: true }) + if (actualEntries.length !== entry.files.length + 1) { + throw new Error('向量模型目录包含未声明文件') + } + for (const file of entry.files) { + const path = safeChild(directory, file.name) + const info = await lstat(path) + if (!info.isFile() || info.isSymbolicLink()) { + throw new Error(`向量模型文件无效：${file.name}`) + } + const actual = await hashModelFile(path) + if ( + actual.size !== file.size || + actual.sha256 !== file.sha256 + ) { + throw new Error(`向量模型文件校验失败：${file.name}`) + } + } + } + + private cleanupStaleArtifacts(): Promise { + return cleanupStaleModelInstallArtifacts({ + rootDirectory: this.rootDirectory, + isModelId: (value) => + embeddingModelIdSchema.safeParse(value).success, + activeModelIds: new Set(this.operations.keys()), + partialFileNames: new Set( + this.catalog.flatMap((entry) => + entry.files.map((file) => file.name) + ) + ), + escapeMessage: '向量模型路径超出受管目录' + }) + } +} diff --git a/src/main/knowledge/embedding-provider-key.test.ts b/src/main/knowledge/embedding-provider-key.test.ts index 998eefe4..66ad8419 100644 --- a/src/main/knowledge/embedding-provider-key.test.ts +++ b/src/main/knowledge/embedding-provider-key.test.ts @@ -1,5 +1,8 @@ import { describe, expect, it } from 'vitest' -import { embeddingStorageProvider } from './embedding-provider-key' +import { + embeddingProviderFingerprint, + embeddingStorageProvider +} from './embedding-provider-key' describe('embeddingStorageProvider', () => { it('preserves legacy provider keys without a fingerprint', () => { @@ -23,4 +26,52 @@ describe('embeddingStorageProvider', () => { expect(second).not.toContain('two.invalid') expect(first.length).toBeLessThanOrEqual(128) }) + + it('canonically fingerprints every vector compatibility field', () => { + const fields = { + provider: 'openai-compatible', + endpoint: 'https://vectors.invalid/v1/embeddings', + dataPath: { scope: 'network', kind: 'endpoint' }, + model: 'embed-v2', + dimensions: 768, + encodingRecipe: { + recipeId: 'query-passage', + queryTemplate: 'query: {text}', + documentTemplate: 'passage: {text}' + } + } as const + const reordered = { + dimensions: 768, + model: 'embed-v2', + provider: 'openai-compatible', + encodingRecipe: { + documentTemplate: 'passage: {text}', + queryTemplate: 'query: {text}', + recipeId: 'query-passage' + }, + dataPath: { kind: 'endpoint', scope: 'network' }, + endpoint: 'https://vectors.invalid/v1/embeddings' + } as const + + expect(embeddingProviderFingerprint(fields)).toBe( + embeddingProviderFingerprint(reordered) + ) + for (const changed of [ + { ...fields, endpoint: 'https://other.invalid/v1/embeddings' }, + { ...fields, dataPath: { kind: 'endpoint', scope: 'loopback' } }, + { ...fields, model: 'embed-v3' }, + { ...fields, dimensions: 384 }, + { + ...fields, + encodingRecipe: { + ...fields.encodingRecipe, + queryTemplate: 'search: {text}' + } + } + ]) { + expect(embeddingProviderFingerprint(changed)).not.toBe( + embeddingProviderFingerprint(fields) + ) + } + }) }) diff --git a/src/main/knowledge/embedding-provider-key.ts b/src/main/knowledge/embedding-provider-key.ts index f52c73d8..63e26c32 100644 --- a/src/main/knowledge/embedding-provider-key.ts +++ b/src/main/knowledge/embedding-provider-key.ts @@ -5,6 +5,38 @@ export interface EmbeddingProviderIdentity { readonly fingerprint?: string } +export interface EmbeddingProviderFingerprintFields { + readonly provider: string + readonly endpoint?: string + readonly dataPath?: unknown + readonly model: string + readonly dimensions?: number + readonly encodingRecipe?: unknown +} + +function canonicalJson(value: unknown): string { + if (value === null || typeof value !== 'object') { + return JSON.stringify(value) + } + if (Array.isArray(value)) { + return `[${value.map(canonicalJson).join(',')}]` + } + const record = value as Record + return `{${Object.keys(record) + .filter((key) => record[key] !== undefined) + .sort() + .map((key) => `${JSON.stringify(key)}:${canonicalJson(record[key])}`) + .join(',')}}` +} + +export function embeddingProviderFingerprint( + fields: EmbeddingProviderFingerprintFields +): string { + return `embedding-v1:${createHash('sha256') + .update(canonicalJson(fields)) + .digest('hex')}` +} + export function embeddingStorageProvider( provider: EmbeddingProviderIdentity ): string { diff --git a/src/main/knowledge/granite-embedding-engine.test.ts b/src/main/knowledge/granite-embedding-engine.test.ts new file mode 100644 index 00000000..1549781c --- /dev/null +++ b/src/main/knowledge/granite-embedding-engine.test.ts @@ -0,0 +1,150 @@ +import path from 'node:path' +import { describe, expect, it, vi } from 'vitest' +import { + buildGraniteInputBatch, + GraniteEmbeddingEngine, + poolGraniteClsEmbeddings, + type GraniteSessionLike, + type GraniteTokenizerLike +} from './granite-embedding-engine' + +function tokenizer(encodings: Record): GraniteTokenizerLike { + return { + encode(text) { + const ids = encodings[text] + if (!ids) { + throw new Error('Unexpected test text') + } + return { + ids, + attention_mask: ids.map(() => 1), + token_type_ids: ids.map(() => 0) + } + }, + token_to_id: () => 7 + } +} + +describe('GraniteEmbeddingEngine', () => { + it('pads int64 model inputs and includes token types only when requested', () => { + const batch = buildGraniteInputBatch({ + tokenizer: tokenizer({ short: [101, 9, 102], long: [101, 8, 6, 102] }), + texts: ['short', 'long'], + padTokenId: 7, + includeTokenTypeIds: true + }) + + expect(batch.sequenceLength).toBe(4) + expect([...batch.inputIds]).toEqual([ + 101n, 9n, 102n, 7n, 101n, 8n, 6n, 102n + ]) + expect([...batch.attentionMask]).toEqual([ + 1n, 1n, 1n, 0n, 1n, 1n, 1n, 1n + ]) + expect(batch.tokenTypeIds).toBeInstanceOf(BigInt64Array) + }) + + it('rejects tokenized input beyond the Granite context limit', () => { + const ids = Array.from({ length: 32_769 }, () => 1) + expect(() => + buildGraniteInputBatch({ + tokenizer: tokenizer({ oversized: ids }), + texts: ['oversized'], + padTokenId: 0, + includeTokenTypeIds: false + }) + ).toThrow(/32768 token limit/u) + }) + + it('selects last_hidden_state, CLS pools, and normalizes 384 dimensions', () => { + const data = new Float32Array(2 * 384) + data[0] = 3 + data[1] = 4 + data[384] = 100 + + const vectors = poolGraniteClsEmbeddings({ + outputs: { + other: { + type: 'float32', + dims: [1, 2, 384], + data: new Float32Array(768).fill(1) + }, + last_hidden_state: { + type: 'float32', + dims: [1, 2, 384], + data + } + }, + outputNames: ['other', 'last_hidden_state'], + batchSize: 1, + sequenceLength: 2 + }) + + expect(vectors[0]?.slice(0, 3)).toEqual([0.6, 0.8, 0]) + expect(vectors[0]).toHaveLength(384) + }) + + it('loads once and runs an injected ONNX session with int64 inputs', async () => { + const run = vi.fn(async (feeds) => { + const input = feeds.input_ids as { + readonly dimensions: readonly [number, number] + } + const [batchSize, sequenceLength] = input.dimensions + const output = new Float32Array( + batchSize * sequenceLength * 384 + ) + output[0] = 1 + if (batchSize > 1) { + output[sequenceLength * 384] = 1 + } + return { + last_hidden_state: { + type: 'float32', + dims: [batchSize, sequenceLength, 384], + data: output + } + } + }) + const session: GraniteSessionLike = { + inputNames: ['input_ids', 'attention_mask'], + outputNames: ['last_hidden_state'], + run + } + const readFile = vi.fn(async (filePath: string) => + new TextEncoder().encode( + filePath.endsWith('.onnx') + ? 'model' + : filePath.endsWith('tokenizer_config.json') + ? '{"pad_token":"[PAD]"}' + : '{}' + ) + ) + const createTensor = vi.fn( + (data: BigInt64Array, dimensions: readonly [number, number]) => ({ + data, + dimensions, + type: 'int64' + }) + ) + const engine = new GraniteEmbeddingEngine({ + readFile, + loadTokenizer: () => tokenizer({ a: [1, 2], b: [3] }), + createSession: async () => session, + createInt64Tensor: createTensor + }) + + const options = { + modelDirectory: path.resolve('granite-model'), + role: 'query' as const, + signal: new AbortController().signal + } + const first = await engine.embed({ ...options, texts: ['a', 'b'] }) + await engine.embed({ ...options, texts: ['a'] }) + + expect(first).toHaveLength(2) + expect(readFile).toHaveBeenCalledTimes(3) + expect(createTensor).toHaveBeenCalledTimes(4) + expect(run).toHaveBeenCalledTimes(2) + expect(run.mock.calls[0]?.[0]).not.toHaveProperty('token_type_ids') + }) +}) diff --git a/src/main/knowledge/granite-embedding-engine.ts b/src/main/knowledge/granite-embedding-engine.ts new file mode 100644 index 00000000..12e7d97e --- /dev/null +++ b/src/main/knowledge/granite-embedding-engine.ts @@ -0,0 +1,400 @@ +import { readFile } from 'node:fs/promises' +import path from 'node:path' +import { Tokenizer } from '@huggingface/tokenizers' +import * as ort from 'onnxruntime-web/wasm' +import type { + EmbeddingInferenceEngine, + EmbeddingOnnxRuntime, + EmbeddingTokenizer +} from './embedding-inference-worker' + +const MODEL_FILE = 'model_quantized.onnx' +const TOKENIZER_FILE = 'tokenizer.json' +const TOKENIZER_CONFIG_FILE = 'tokenizer_config.json' +const MAX_TOKENS = 32_768 +const EMBEDDING_DIMENSIONS = 384 + +type Encoding = { + readonly ids: readonly number[] + readonly attention_mask: readonly number[] + readonly token_type_ids?: readonly number[] +} + +export interface GraniteTokenizerLike { + encode(text: string, options?: { + readonly add_special_tokens?: boolean + readonly return_token_type_ids?: boolean | null + }): Encoding + token_to_id(token: string): number | undefined +} + +export interface GraniteTensorLike { + readonly type: string + readonly dims: readonly number[] + readonly data: ArrayLike +} + +export interface GraniteSessionLike { + readonly inputNames: readonly string[] + readonly outputNames: readonly string[] + run( + feeds: Readonly> + ): Promise>> + release?(): void +} + +type LoadedEngine = { + readonly tokenizer: GraniteTokenizerLike + readonly padTokenId: number + readonly session: GraniteSessionLike +} + +export interface GraniteEmbeddingEngineDependencies { + readonly loadTokenizer?: ( + tokenizerJson: unknown, + tokenizerConfig: unknown + ) => GraniteTokenizerLike + readonly createSession?: (model: Uint8Array) => Promise + readonly createInt64Tensor?: ( + data: BigInt64Array, + dimensions: readonly [number, number] + ) => unknown + readonly readFile?: (filePath: string) => Promise +} + +type TokenizedBatch = { + readonly batchSize: number + readonly sequenceLength: number + readonly inputIds: BigInt64Array + readonly attentionMask: BigInt64Array + readonly tokenTypeIds?: BigInt64Array +} + +function abortIfRequested(signal: AbortSignal): void { + if (signal.aborted) { + throw signal.reason ?? new Error('Embedding inference was cancelled') + } +} + +function asJsonObject(value: unknown, name: string): Record { + if (typeof value !== 'object' || value === null || Array.isArray(value)) { + throw new TypeError(`${name} must contain a JSON object`) + } + return value as Record +} + +function resolvePadTokenId( + tokenizer: GraniteTokenizerLike, + rawConfig: unknown +): number { + const config = asJsonObject(rawConfig, TOKENIZER_CONFIG_FILE) + const configuredId = config.pad_token_id + if ( + typeof configuredId === 'number' && + Number.isSafeInteger(configuredId) && + configuredId >= 0 + ) { + return configuredId + } + if (typeof config.pad_token === 'string') { + const tokenId = tokenizer.token_to_id(config.pad_token) + if (tokenId !== undefined) { + return tokenId + } + } + throw new Error('Tokenizer configuration does not define a valid pad token') +} + +function parseJson(bytes: Uint8Array, name: string): unknown { + try { + return JSON.parse(new TextDecoder().decode(bytes)) as unknown + } catch { + throw new Error(`${name} is not valid JSON`) + } +} + +export function buildGraniteInputBatch(options: { + readonly tokenizer: GraniteTokenizerLike + readonly texts: readonly string[] + readonly padTokenId: number + readonly includeTokenTypeIds: boolean +}): TokenizedBatch { + if (options.texts.length === 0) { + throw new RangeError('At least one text is required') + } + const encodings = options.texts.map((text) => + options.tokenizer.encode(text, { + add_special_tokens: true, + return_token_type_ids: options.includeTokenTypeIds + }) + ) + const sequenceLength = Math.max(...encodings.map(({ ids }) => ids.length)) + if (sequenceLength > MAX_TOKENS) { + throw new RangeError( + `Embedding input exceeds the ${MAX_TOKENS} token limit` + ) + } + if (sequenceLength < 1) { + throw new RangeError('Tokenizer produced an empty input') + } + + const itemCount = options.texts.length * sequenceLength + const inputIds = new BigInt64Array(itemCount) + inputIds.fill(BigInt(options.padTokenId)) + const attentionMask = new BigInt64Array(itemCount) + const tokenTypeIds = options.includeTokenTypeIds + ? new BigInt64Array(itemCount) + : undefined + + for (const [batchIndex, encoding] of encodings.entries()) { + if ( + encoding.attention_mask.length !== encoding.ids.length || + (options.includeTokenTypeIds && + encoding.token_type_ids !== undefined && + encoding.token_type_ids.length !== encoding.ids.length) + ) { + throw new Error('Tokenizer produced inconsistent input lengths') + } + const offset = batchIndex * sequenceLength + for (let index = 0; index < encoding.ids.length; index += 1) { + const tokenId = encoding.ids[index] + const mask = encoding.attention_mask[index] + if ( + tokenId === undefined || + mask === undefined || + !Number.isSafeInteger(tokenId) || + !Number.isSafeInteger(mask) + ) { + throw new Error('Tokenizer produced an invalid token') + } + inputIds[offset + index] = BigInt(tokenId) + attentionMask[offset + index] = BigInt(mask) + if (tokenTypeIds) { + tokenTypeIds[offset + index] = BigInt( + encoding.token_type_ids?.[index] ?? 0 + ) + } + } + } + + return { + batchSize: options.texts.length, + sequenceLength, + inputIds, + attentionMask, + tokenTypeIds + } +} + +export function poolGraniteClsEmbeddings(options: { + readonly outputs: Readonly> + readonly outputNames: readonly string[] + readonly batchSize: number + readonly sequenceLength: number +}): number[][] { + const preferred = options.outputs.last_hidden_state + const output = + preferred?.dims.length === 3 && + (preferred.type === 'float32' || preferred.type === 'float64') + ? preferred + : options.outputNames + .map((name) => options.outputs[name]) + .find( + (candidate) => + candidate?.dims.length === 3 && + (candidate.type === 'float32' || + candidate.type === 'float64') + ) + if (!output) { + throw new Error('ONNX model did not return a rank-3 float output') + } + const [batchSize, sequenceLength, dimensions] = output.dims + if ( + batchSize !== options.batchSize || + sequenceLength !== options.sequenceLength || + dimensions !== EMBEDDING_DIMENSIONS || + output.data.length !== batchSize * sequenceLength * dimensions + ) { + throw new Error('ONNX embedding output has invalid dimensions') + } + + const vectors: number[][] = [] + for (let batchIndex = 0; batchIndex < batchSize; batchIndex += 1) { + const offset = batchIndex * sequenceLength * dimensions + const vector = Array.from( + { length: dimensions }, + (_, index) => output.data[offset + index] ?? Number.NaN + ) + const magnitude = Math.hypot(...vector) + if (!Number.isFinite(magnitude) || magnitude <= 0) { + throw new Error('ONNX embedding output has an invalid CLS vector') + } + vectors.push(vector.map((component) => component / magnitude)) + } + return vectors +} + +export class GraniteEmbeddingEngine implements EmbeddingInferenceEngine { + readonly tokenizer: EmbeddingTokenizer + readonly onnx: EmbeddingOnnxRuntime + + private readonly dependencies: Required + private loadTask?: Promise + private modelDirectory?: string + private disposed = false + + constructor(dependencies: GraniteEmbeddingEngineDependencies = {}) { + this.dependencies = { + loadTokenizer: + dependencies.loadTokenizer ?? + ((tokenizerJson, tokenizerConfig) => + new Tokenizer( + asJsonObject(tokenizerJson, TOKENIZER_FILE), + asJsonObject(tokenizerConfig, TOKENIZER_CONFIG_FILE) + )), + createSession: + dependencies.createSession ?? + (async (model) => + (await ort.InferenceSession.create(model, { + executionProviders: ['wasm'] + })) as GraniteSessionLike), + createInt64Tensor: + dependencies.createInt64Tensor ?? + ((data, dimensions) => new ort.Tensor('int64', data, dimensions)), + readFile: + dependencies.readFile ?? + (async (filePath) => new Uint8Array(await readFile(filePath))) + } + this.tokenizer = { + tokenize: async ({ texts, signal }) => { + const loaded = await this.load(this.requiredModelDirectory(), signal) + return buildGraniteInputBatch({ + tokenizer: loaded.tokenizer, + texts, + padTokenId: loaded.padTokenId, + includeTokenTypeIds: + loaded.session.inputNames.includes('token_type_ids') + }) + } + } + this.onnx = { + run: async ({ modelDirectory, inputs, signal }) => { + const loaded = await this.load(modelDirectory, signal) + abortIfRequested(signal) + const output = await loaded.session.run( + inputs as Readonly> + ) + abortIfRequested(signal) + return output + } + } + } + + async embed(options: Parameters[0]) { + void options.role + const loaded = await this.load(options.modelDirectory, options.signal) + abortIfRequested(options.signal) + const batch = buildGraniteInputBatch({ + tokenizer: loaded.tokenizer, + texts: options.texts, + padTokenId: loaded.padTokenId, + includeTokenTypeIds: + loaded.session.inputNames.includes('token_type_ids') + }) + abortIfRequested(options.signal) + const dimensions: [number, number] = [ + batch.batchSize, + batch.sequenceLength + ] + const feeds: Record = { + input_ids: this.dependencies.createInt64Tensor( + batch.inputIds, + dimensions + ), + attention_mask: this.dependencies.createInt64Tensor( + batch.attentionMask, + dimensions + ) + } + if (batch.tokenTypeIds) { + feeds.token_type_ids = this.dependencies.createInt64Tensor( + batch.tokenTypeIds, + dimensions + ) + } + const outputs = await loaded.session.run(feeds) + abortIfRequested(options.signal) + return poolGraniteClsEmbeddings({ + outputs, + outputNames: loaded.session.outputNames, + batchSize: batch.batchSize, + sequenceLength: batch.sequenceLength + }) + } + + async dispose(): Promise { + if (this.disposed) { + return + } + this.disposed = true + const loaded = await this.loadTask?.catch(() => undefined) + loaded?.session.release?.() + } + + private requiredModelDirectory(): string { + if (!this.modelDirectory) { + throw new Error('The Granite model has not been selected') + } + return this.modelDirectory + } + + private load( + modelDirectory: string, + signal: AbortSignal + ): Promise { + abortIfRequested(signal) + if (this.disposed) { + throw new Error('Granite embedding engine is disposed') + } + const normalizedDirectory = path.normalize(modelDirectory) + if ( + this.modelDirectory && + this.modelDirectory !== normalizedDirectory + ) { + throw new Error('Granite embedding engine cannot change model directory') + } + this.modelDirectory = normalizedDirectory + this.loadTask ??= this.loadFiles(normalizedDirectory) + return this.loadTask.then((loaded) => { + abortIfRequested(signal) + return loaded + }) + } + + private async loadFiles(modelDirectory: string): Promise { + const [tokenizerBytes, tokenizerConfigBytes, modelBytes] = + await Promise.all([ + this.dependencies.readFile( + path.join(modelDirectory, TOKENIZER_FILE) + ), + this.dependencies.readFile( + path.join(modelDirectory, TOKENIZER_CONFIG_FILE) + ), + this.dependencies.readFile(path.join(modelDirectory, MODEL_FILE)) + ]) + const tokenizerJson = parseJson(tokenizerBytes, TOKENIZER_FILE) + const tokenizerConfig = parseJson( + tokenizerConfigBytes, + TOKENIZER_CONFIG_FILE + ) + const tokenizer = this.dependencies.loadTokenizer( + tokenizerJson, + tokenizerConfig + ) + const [padTokenId, session] = await Promise.all([ + Promise.resolve(resolvePadTokenId(tokenizer, tokenizerConfig)), + this.dependencies.createSession(modelBytes) + ]) + return { tokenizer, padTokenId, session } + } +} diff --git a/src/main/knowledge/knowledge-service.test.ts b/src/main/knowledge/knowledge-service.test.ts index 84bf411d..4e059cb9 100644 --- a/src/main/knowledge/knowledge-service.test.ts +++ b/src/main/knowledge/knowledge-service.test.ts @@ -48,6 +48,56 @@ afterEach(async () => { }) describe('KnowledgeService', () => { + it('uses document and query embedding roles without exposing providers to runtimes', async () => { + const embed = vi.fn(async (input: readonly string[]) => + input.map(() => [9, 9]) + ) + const embedDocuments = vi.fn(async (input: readonly string[]) => + input.map(() => [1, 0]) + ) + const embedQuery = vi.fn(async (input: readonly string[]) => + input.map(() => [1, 0]) + ) + const provider: EmbeddingProvider = { + provider: 'openai-compatible', + model: 'role-aware', + embed, + embedDocuments, + embedQuery + } + const { directory, service } = await createService( + undefined, + provider + ) + const sourcePath = join(directory, 'role-aware.md') + await writeFile(sourcePath, 'role aware knowledge content', 'utf8') + const library = service.createLibrary({ + name: 'Role-aware library', + storageMode: 'reference', + graphEnabled: false + }) + + await service.importPaths(library.id, [sourcePath]) + expect(embedDocuments).toHaveBeenCalled() + expect(embed).not.toHaveBeenCalled() + + await service.retrieve({ + knowledgeBaseId: library.id, + query: 'role aware', + settings: { + ...library.retrievalSettings, + ftsWeight: 0, + vectorWeight: 1, + graphWeight: 0 + } + }) + expect(embedQuery).toHaveBeenCalledWith( + ['role aware'], + expect.any(AbortSignal) + ) + expect(embed).not.toHaveBeenCalled() + }) + it('rebuilds vectors only for the selected knowledge base', async () => { const embed = vi.fn(async (input: readonly string[]) => input.map(() => [1, 0]) diff --git a/src/main/knowledge/knowledge-service.ts b/src/main/knowledge/knowledge-service.ts index 0c3f9dc2..9c0eb0fe 100644 --- a/src/main/knowledge/knowledge-service.ts +++ b/src/main/knowledge/knowledge-service.ts @@ -180,6 +180,26 @@ interface PreparedQueryEmbedding { degradationReason?: string } +function embedKnowledgeDocuments( + provider: EmbeddingProvider, + input: readonly string[], + signal?: AbortSignal +): Promise { + return provider.embedDocuments + ? provider.embedDocuments(input, signal) + : provider.embed(input, signal) +} + +function embedKnowledgeQuery( + provider: EmbeddingProvider, + input: readonly string[], + signal?: AbortSignal +): Promise { + return provider.embedQuery + ? provider.embedQuery(input, signal) + : provider.embed(input, signal) +} + export class KnowledgeService { readonly database: KnowledgeDatabase private readonly managedRoot: string @@ -636,7 +656,11 @@ export class KnowledgeService { chunk.metadata?.contextPrefix ) ) - const vectors = await provider.embed(contents, effectiveSignal) + const vectors = await embedKnowledgeDocuments( + provider, + contents, + effectiveSignal + ) effectiveSignal.throwIfAborted() if (vectors.length !== batch.length) { throw new Error('Embedding provider returned an invalid result count') @@ -1740,7 +1764,8 @@ export class KnowledgeService { ? AbortSignal.any([signal, this.lifecycleController.signal]) : this.lifecycleController.signal try { - const result = await provider.embed( + const result = await embedKnowledgeQuery( + provider, [query], effectiveSignal ) @@ -3164,7 +3189,8 @@ export class KnowledgeService { ) { effectiveSignal.throwIfAborted() const batch = chunks.slice(offset, offset + this.embeddingBatchSize) - const vectors = await provider.embed( + const vectors = await embedKnowledgeDocuments( + provider, batch.map((chunk) => chunk.content), effectiveSignal ) diff --git a/src/main/knowledge/openai-embedding-client.test.ts b/src/main/knowledge/openai-embedding-client.test.ts index c09e319e..13b8b70b 100644 --- a/src/main/knowledge/openai-embedding-client.test.ts +++ b/src/main/knowledge/openai-embedding-client.test.ts @@ -68,6 +68,98 @@ describe('OpenAIEmbeddingClient', () => { ) }) + it('encodes query and document roles explicitly and fingerprints the recipe', async () => { + const bodies: unknown[] = [] + const transport = vi.fn(async (_input, init) => { + const body = JSON.parse(String(init?.body)) as { + input: string[] + } + bodies.push(body) + return new Response( + JSON.stringify({ + data: body.input.map((_, index) => ({ + index, + embedding: [1, 2, 3] + })) + }) + ) + }) + const options = { + endpoint: 'https://vectors.example/v1/embeddings', + model: 'embed-v2', + dimensions: 3, + encodingRecipe: { + recipeId: 'query-passage', + queryTemplate: 'query: {text}', + documentTemplate: 'passage: {text}' + }, + fetch: transport + } as const + const client = new OpenAIEmbeddingClient(options) + const changedRecipe = new OpenAIEmbeddingClient({ + ...options, + encodingRecipe: { + ...options.encodingRecipe, + queryTemplate: 'search: {text}' + } + }) + + await client.embedQuery(['question']) + await client.embedDocuments(['answer']) + await client.embed(['legacy document']) + + expect(bodies).toEqual([ + { + model: 'embed-v2', + input: ['query: question'], + dimensions: 3 + }, + { + model: 'embed-v2', + input: ['passage: answer'], + dimensions: 3 + }, + { + model: 'embed-v2', + input: ['passage: legacy document'], + dimensions: 3 + } + ]) + expect(client.fingerprint).not.toBe(changedRecipe.fingerprint) + expect(client.fingerprint).not.toContain('vectors.example') + }) + + it('rejects vectors that do not match configured dimensions', async () => { + const client = new OpenAIEmbeddingClient({ + endpoint: 'https://vectors.example/v1/embeddings', + model: 'embed-v2', + dimensions: 2, + fetch: async () => + new Response( + JSON.stringify({ + data: [{ index: 0, embedding: [1, 2, 3] }] + }) + ) + }) + + await expect(client.embedQuery(['question'])).rejects.toThrow( + 'configured dimensions' + ) + }) + + it('rejects unsupported encoding behavior instead of silently changing it', () => { + expect( + () => + new OpenAIEmbeddingClient({ + endpoint: 'https://vectors.example/v1/embeddings', + model: 'embed-v2', + encodingRecipe: { + pooling: 'mean' + } as never + }) + ).toThrow('pooling must be provider-managed') + }) + it('distinguishes its request timeout from caller cancellation', async () => { const waitForAbort = vi.fn( async (_input, init) => diff --git a/src/main/knowledge/openai-embedding-client.ts b/src/main/knowledge/openai-embedding-client.ts index 40516387..c7f40e95 100644 --- a/src/main/knowledge/openai-embedding-client.ts +++ b/src/main/knowledge/openai-embedding-client.ts @@ -1,4 +1,5 @@ import type { EmbeddingProvider } from './types' +import { embeddingProviderFingerprint } from './embedding-provider-key' const MAX_INPUTS = 256 const MAX_BATCH_SIZE = 32 @@ -16,11 +17,33 @@ export interface OpenAIEmbeddingClientOptions { endpoint: string model: string apiKey?: string + dimensions?: number + encodingRecipe?: Partial batchSize?: number timeoutMs?: number fetch?: typeof fetch } +export type EmbeddingInputRole = 'query' | 'document' + +export interface OpenAIEmbeddingEncodingRecipe { + recipeId: string + tokenizerDigest?: string + pooling: 'provider-managed' + normalization: 'provider-managed' + queryTemplate: string + documentTemplate: string + maximumSequenceTokens?: number +} + +const DEFAULT_ENCODING_RECIPE: OpenAIEmbeddingEncodingRecipe = { + recipeId: 'symmetric', + pooling: 'provider-managed', + normalization: 'provider-managed', + queryTemplate: '{text}', + documentTemplate: '{text}' +} + function boundedInteger( value: number, field: string, @@ -56,6 +79,71 @@ function normalizedEndpoint(input: string): string { return url.toString() } +function endpointScope(endpoint: string): 'loopback' | 'network' { + const hostname = new URL(endpoint).hostname.toLowerCase() + return hostname === 'localhost' || + hostname === '127.0.0.1' || + hostname === '[::1]' + ? 'loopback' + : 'network' +} + +function normalizedTemplate(value: string, field: string): string { + const template = requiredString(value, field, MAX_INPUT_LENGTH) + if (template.split('{text}').length !== 2) { + throw new RangeError(`${field} must contain exactly one {text} placeholder`) + } + return template +} + +function normalizedEncodingRecipe( + input: Partial | undefined +): OpenAIEmbeddingEncodingRecipe { + const recipe = { ...DEFAULT_ENCODING_RECIPE, ...input } + if (recipe.pooling !== 'provider-managed') { + throw new RangeError( + 'encodingRecipe.pooling must be provider-managed' + ) + } + if (recipe.normalization !== 'provider-managed') { + throw new RangeError( + 'encodingRecipe.normalization must be provider-managed' + ) + } + return { + recipeId: requiredString(recipe.recipeId, 'encodingRecipe.recipeId', 256), + ...(recipe.tokenizerDigest + ? { + tokenizerDigest: requiredString( + recipe.tokenizerDigest, + 'encodingRecipe.tokenizerDigest', + 256 + ) + } + : {}), + pooling: 'provider-managed', + normalization: 'provider-managed', + queryTemplate: normalizedTemplate( + recipe.queryTemplate, + 'encodingRecipe.queryTemplate' + ), + documentTemplate: normalizedTemplate( + recipe.documentTemplate, + 'encodingRecipe.documentTemplate' + ), + ...(recipe.maximumSequenceTokens !== undefined + ? { + maximumSequenceTokens: boundedInteger( + recipe.maximumSequenceTokens, + 'encodingRecipe.maximumSequenceTokens', + 1, + 1_000_000 + ) + } + : {}) + } +} + function embeddingAbortError( requestSignal: AbortSignal, timeoutError: Error @@ -174,6 +262,8 @@ export class OpenAIEmbeddingClient implements EmbeddingProvider { readonly model: string readonly fingerprint: string private readonly endpoint: string + private readonly dimensions?: number + private readonly encodingRecipe: OpenAIEmbeddingEncodingRecipe private readonly apiKey?: string private readonly batchSize: number private readonly timeoutMs: number @@ -183,7 +273,22 @@ export class OpenAIEmbeddingClient implements EmbeddingProvider { this.endpoint = normalizedEndpoint(options.endpoint) this.model = requiredString(options.model, 'model', MAX_MODEL_LENGTH) this.apiKey = options.apiKey?.trim() || undefined - this.fingerprint = `${this.provider}:${this.endpoint}:${this.model}` + this.dimensions = + options.dimensions === undefined + ? undefined + : boundedInteger(options.dimensions, 'dimensions', 1, MAX_DIMENSIONS) + this.encodingRecipe = normalizedEncodingRecipe(options.encodingRecipe) + this.fingerprint = embeddingProviderFingerprint({ + provider: this.provider, + endpoint: this.endpoint, + dataPath: { + kind: 'endpoint', + scope: endpointScope(this.endpoint) + }, + model: this.model, + dimensions: this.dimensions, + encodingRecipe: this.encodingRecipe + }) this.batchSize = boundedInteger( options.batchSize ?? 16, 'batchSize', @@ -206,6 +311,31 @@ export class OpenAIEmbeddingClient implements EmbeddingProvider { input: readonly string[], signal?: AbortSignal ): Promise { + return this.embedRole(input, 'document', signal) + } + + async embedQuery( + input: readonly string[], + signal?: AbortSignal + ): Promise { + return this.embedRole(input, 'query', signal) + } + + async embedDocuments( + input: readonly string[], + signal?: AbortSignal + ): Promise { + return this.embedRole(input, 'document', signal) + } + + async embedRole( + input: readonly string[], + role: EmbeddingInputRole, + signal?: AbortSignal + ): Promise { + if (role !== 'query' && role !== 'document') { + throw new RangeError('role must be query or document') + } if (!Array.isArray(input) || input.length < 1 || input.length > MAX_INPUTS) { throw new RangeError(`input must contain between 1 and ${MAX_INPUTS} items`) } @@ -218,7 +348,17 @@ export class OpenAIEmbeddingClient implements EmbeddingProvider { `input[${index}] must be at most ${MAX_INPUT_LENGTH} characters` ) } - return item + const template = + role === 'query' + ? this.encodingRecipe.queryTemplate + : this.encodingRecipe.documentTemplate + const encoded = template.replace('{text}', item) + if (encoded.length > MAX_INPUT_LENGTH) { + throw new RangeError( + `encoded input[${index}] must be at most ${MAX_INPUT_LENGTH} characters` + ) + } + return encoded }) const embeddings: number[][] = [] @@ -284,7 +424,13 @@ export class OpenAIEmbeddingClient implements EmbeddingProvider { response = await this.transport(this.endpoint, { method: 'POST', headers, - body: JSON.stringify({ model: this.model, input }), + body: JSON.stringify({ + model: this.model, + input, + ...(this.dimensions === undefined + ? {} + : { dimensions: this.dimensions }) + }), redirect: 'error', signal: requestSignal }) @@ -293,10 +439,17 @@ export class OpenAIEmbeddingClient implements EmbeddingProvider { `Embedding request failed with HTTP ${response.status}` ) } - return validateEmbeddings( + const vectors = validateEmbeddings( await readBoundedJson(response), input.length ) + if ( + this.dimensions !== undefined && + vectors.some((vector) => vector.length !== this.dimensions) + ) { + throw new Error('Embeddings do not match configured dimensions') + } + return vectors } catch (error) { if (requestSignal.aborted) { throw embeddingAbortError(requestSignal, timeoutError) diff --git a/src/main/knowledge/types.ts b/src/main/knowledge/types.ts index 7df8879b..476f96ad 100644 --- a/src/main/knowledge/types.ts +++ b/src/main/knowledge/types.ts @@ -158,6 +158,14 @@ export interface EmbeddingProvider { readonly model: string readonly fingerprint?: string embed(input: readonly string[], signal?: AbortSignal): Promise + embedQuery?( + input: readonly string[], + signal?: AbortSignal + ): Promise + embedDocuments?( + input: readonly string[], + signal?: AbortSignal + ): Promise } export interface RerankProviderResult { diff --git a/src/main/model-archive.ts b/src/main/model-archive.ts index 29ac74ae..2633c731 100644 --- a/src/main/model-archive.ts +++ b/src/main/model-archive.ts @@ -47,7 +47,7 @@ const modelArchiveFileSchema = z const modelArchiveDescriptorSchema = z .object({ - kind: z.enum(['speech', 'document-ocr']), + kind: z.enum(['speech', 'document-ocr', 'embedding']), modelId: z .string() .min(1) diff --git a/src/main/runtime-settings-store.test.ts b/src/main/runtime-settings-store.test.ts index 4f6faef2..4b931253 100644 --- a/src/main/runtime-settings-store.test.ts +++ b/src/main/runtime-settings-store.test.ts @@ -9,6 +9,8 @@ import { tmpdir } from 'node:os' import { join } from 'node:path' import { afterEach, describe, expect, it, vi } from 'vitest' import { + builtinEmbeddingConnectionId, + legacyEmbeddingConnectionId, runtimeSettingsInputSchema, type RuntimeSettingsInput } from '../shared/contracts' @@ -479,7 +481,7 @@ describe('RuntimeSettingsStore', () => { const persisted = JSON.parse( await readFile(filePath, 'utf8') ) as Record - expect(persisted.version).toBe(18) + expect(persisted.version).toBe(19) expect(persisted).not.toHaveProperty( 'deepseekHarnessBinaryPath' ) @@ -758,7 +760,7 @@ describe('RuntimeSettingsStore', () => { const persisted = JSON.parse(await readFile(filePath, 'utf8')) as { version: number } - expect(persisted.version).toBe(18) + expect(persisted.version).toBe(19) }) it('migrates version 11 and removes the obsolete intranet toggle', async () => { @@ -778,7 +780,7 @@ describe('RuntimeSettingsStore', () => { version: number intranetCompatibilityEnabled?: boolean } - expect(persisted.version).toBe(18) + expect(persisted.version).toBe(19) expect(persisted).not.toHaveProperty('intranetCompatibilityEnabled') }) @@ -947,6 +949,118 @@ describe('RuntimeSettingsStore', () => { }) }) + it('migrates version 18 embedding settings into a separate user connection', async () => { + const { filePath, store } = await createStore() + await store.update( + settings({ + knowledgeEmbeddingEnabled: true, + knowledgeEmbeddingBaseUrl: + 'https://legacy-vectors.example/v1/embeddings', + knowledgeEmbeddingModel: 'legacy/embed-large', + knowledgeEmbeddingApiKey: { + action: 'replace', + value: 'legacy-vector-secret' + } + }) + ) + const previous = JSON.parse( + await readFile(filePath, 'utf8') + ) as Record + previous.version = 18 + delete previous.embeddingConnections + delete previous.activeEmbeddingConnectionId + await writeFile(filePath, JSON.stringify(previous), 'utf8') + + const migrated = new RuntimeSettingsStore(filePath, cipher, {}) + await expect(migrated.getPublicSettings()).resolves.toMatchObject({ + activeEmbeddingConnectionId: legacyEmbeddingConnectionId, + embeddingConnections: [ + { + id: builtinEmbeddingConnectionId, + kind: 'builtin', + apiKeyConfigured: false, + credentialSource: 'none' + }, + { + id: legacyEmbeddingConnectionId, + kind: 'openai-compatible', + baseUrl: 'https://legacy-vectors.example/v1/embeddings', + modelName: 'legacy/embed-large', + apiKeyConfigured: true, + credentialSource: 'encrypted' + } + ], + knowledgeEmbeddingBaseUrl: + 'https://legacy-vectors.example/v1/embeddings', + knowledgeEmbeddingModel: 'legacy/embed-large' + }) + await expect(migrated.getResolvedSettings()).resolves.toMatchObject({ + activeEmbeddingConnection: { + id: legacyEmbeddingConnectionId, + apiKey: 'legacy-vector-secret' + }, + knowledgeEmbeddingApiKey: 'legacy-vector-secret' + }) + }) + + it('round-trips embedding connections without assigning them to Runtime profiles', async () => { + const { filePath, store } = await createStore() + const connectionId = '00000000-0000-4000-8000-000000000081' + await store.update( + settings({ + embeddingConnections: [ + { + id: builtinEmbeddingConnectionId, + name: 'GoodBuddy 内置向量模型', + kind: 'builtin' + }, + { + id: connectionId, + name: '团队向量模型', + kind: 'openai-compatible', + baseUrl: 'https://team-vectors.example/v1/embeddings', + modelName: 'team/embed-v2', + authentication: 'api-key', + apiKey: { + action: 'replace', + value: 'team-vector-secret' + } + } + ], + activeEmbeddingConnectionId: connectionId + }) + ) + + const reloaded = new RuntimeSettingsStore(filePath, cipher, {}) + await expect(reloaded.getResolvedSettings()).resolves.toMatchObject({ + activeEmbeddingConnectionId: connectionId, + activeEmbeddingConnection: { + id: connectionId, + kind: 'openai-compatible', + apiKey: 'team-vector-secret' + }, + modelProfiles: [ + expect.not.objectContaining({ id: connectionId }) + ] + }) + const publicSettings = await reloaded.getPublicSettings() + expect(JSON.stringify(publicSettings)).not.toContain( + 'team-vector-secret' + ) + expect(publicSettings.embeddingConnections).toContainEqual( + expect.objectContaining({ + id: connectionId, + apiKeyConfigured: true, + credentialSource: 'encrypted' + }) + ) + expect( + (JSON.parse(await readFile(filePath, 'utf8')) as { + version: number + }).version + ).toBe(19) + }) + it('migrates version 13 with reranking disabled by default', async () => { const { filePath, store } = await createStore() await store.update(settings()) @@ -1451,7 +1565,7 @@ describe('RuntimeSettingsStore', () => { version: number modelProfiles: Array> } - expect(persisted.version).toBe(18) + expect(persisted.version).toBe(19) expect(persisted.modelProfiles).toContainEqual( expect.objectContaining({ id: imageId, @@ -1750,7 +1864,7 @@ describe('RuntimeSettingsStore', () => { unknown > expect(saved).toMatchObject({ - version: 18, + version: 19, provider: 'model', continueBinaryPath: '', continueMode: 'chat', @@ -2029,7 +2143,7 @@ describe('RuntimeSettingsStore', () => { version: number modelProfiles: Array> } - expect(persisted.version).toBe(18) + expect(persisted.version).toBe(19) expect(persisted.modelProfiles[0]).not.toHaveProperty('credential') }) diff --git a/src/main/runtime-settings-store.ts b/src/main/runtime-settings-store.ts index 1a830d8e..c75b03fd 100644 --- a/src/main/runtime-settings-store.ts +++ b/src/main/runtime-settings-store.ts @@ -7,12 +7,14 @@ import { homedir } from 'node:os' import { z } from 'zod' import { continueModeSchema, + builtinEmbeddingConnectionId, contextCompressionSettingsSchema, defaultContextCompressionSettings, defaultModelProfileId, defaultRuntimeCustomizationSettings, defaultRuntimeSettings, imageGenerationQualitySchema, + legacyEmbeddingConnectionId, isAgentRuntimeModelProtocol, isDeepSeekHarnessModelProfile, minimumModelContextWindowTokens, @@ -227,13 +229,47 @@ const version17StoredSettingsSchema = version16StoredSettingsSchema contextCompression: contextCompressionSettingsSchema }) -const storedSettingsSchema = version17StoredSettingsSchema +const version18StoredSettingsSchema = version17StoredSettingsSchema .omit({ version: true }) .extend({ version: z.literal(18), runtimeCustomization: runtimeCustomizationSettingsSchema }) +const storedEmbeddingConnectionSchema = z.discriminatedUnion('kind', [ + z + .object({ + id: z.literal(builtinEmbeddingConnectionId), + name: z.string().trim().min(1).max(64), + kind: z.literal('builtin') + }) + .strict(), + z + .object({ + id: z.string().uuid().refine( + (id) => id !== builtinEmbeddingConnectionId + ), + name: z.string().trim().min(1).max(64), + kind: z.literal('openai-compatible'), + baseUrl: z.string().url().max(2_048), + modelName: z.string().trim().min(1).max(256), + authentication: modelAuthenticationSchema, + credential: credentialSchema + }) + .strict() +]) + +const storedSettingsSchema = version18StoredSettingsSchema + .omit({ version: true }) + .extend({ + version: z.literal(19), + embeddingConnections: z + .array(storedEmbeddingConnectionSchema) + .min(1) + .max(20), + activeEmbeddingConnectionId: z.string().uuid() + }) + type StoredSettings = z.infer export type RuntimeSettingsRollback = { publicSettings: RuntimeSettings @@ -242,6 +278,9 @@ export type RuntimeSettingsRollback = { type Version17StoredSettings = z.infer< typeof version17StoredSettingsSchema > +type Version18StoredSettings = z.infer< + typeof version18StoredSettingsSchema +> type Version16StoredSettings = z.infer< typeof version16StoredSettingsSchema > @@ -356,6 +395,9 @@ export type ResolvedRuntimeSettings = { knowledgeEmbeddingBaseUrl: string knowledgeEmbeddingModel: string knowledgeEmbeddingApiKey?: string + embeddingConnections?: ResolvedEmbeddingConnection[] + activeEmbeddingConnectionId?: string + activeEmbeddingConnection?: ResolvedEmbeddingConnection knowledgeRerankEnabled: boolean knowledgeRerankEndpoint: string knowledgeRerankModel: string @@ -384,6 +426,22 @@ export type ResolvedModelProfile = { apiKey?: string } +export type ResolvedEmbeddingConnection = + | { + id: typeof builtinEmbeddingConnectionId + name: string + kind: 'builtin' + } + | { + id: string + name: string + kind: 'openai-compatible' + baseUrl: string + modelName: string + authentication: RuntimeSettings['modelAuthentication'] + apiKey?: string + } + type ResolvedModelCredential = { activeApiKey?: string configured: boolean @@ -391,7 +449,7 @@ type ResolvedModelCredential = { } const defaultSettings: StoredSettings = { - version: 18, + version: 19, provider: defaultRuntimeSettings.provider, modelProfiles: [ { @@ -431,6 +489,22 @@ const defaultSettings: StoredSettings = { defaultRuntimeSettings.knowledgeEmbeddingBaseUrl, knowledgeEmbeddingModel: defaultRuntimeSettings.knowledgeEmbeddingModel, + embeddingConnections: [ + { + id: builtinEmbeddingConnectionId, + name: 'GoodBuddy 内置向量模型', + kind: 'builtin' + }, + { + id: legacyEmbeddingConnectionId, + name: '自定义向量模型', + kind: 'openai-compatible', + baseUrl: defaultRuntimeSettings.knowledgeEmbeddingBaseUrl, + modelName: defaultRuntimeSettings.knowledgeEmbeddingModel, + authentication: 'api-key' + } + ], + activeEmbeddingConnectionId: builtinEmbeddingConnectionId, knowledgeRerankEnabled: defaultRuntimeSettings.knowledgeRerankEnabled, knowledgeRerankEndpoint: @@ -528,13 +602,13 @@ function migrateVersion14( ...current } = settings void _obsolete - return { + return migrateVersion18({ ...current, version: 18, deepseekHarnessModelSource: { kind: 'platform' }, contextCompression: defaultContextCompressionSettings, runtimeCustomization: defaultRuntimeCustomizationSettings - } + }) } function migrateVersion15( @@ -547,32 +621,58 @@ function migrateVersion15( } = settings void _obsolete void _obsoleteSandbox - return { + return migrateVersion18({ ...current, version: 18, contextCompression: defaultContextCompressionSettings, runtimeCustomization: defaultRuntimeCustomizationSettings - } + }) } function migrateVersion16( settings: Version16StoredSettings ): StoredSettings { - return { + return migrateVersion18({ ...settings, version: 18, contextCompression: defaultContextCompressionSettings, runtimeCustomization: defaultRuntimeCustomizationSettings - } + }) } function migrateVersion17( settings: Version17StoredSettings ): StoredSettings { - return { + return migrateVersion18({ ...settings, version: 18, runtimeCustomization: defaultRuntimeCustomizationSettings + }) +} + +function migrateVersion18( + settings: Version18StoredSettings +): StoredSettings { + return { + ...settings, + version: 19, + embeddingConnections: [ + { + id: builtinEmbeddingConnectionId, + name: 'GoodBuddy 内置向量模型', + kind: 'builtin' + }, + { + id: legacyEmbeddingConnectionId, + name: '自定义向量模型', + kind: 'openai-compatible', + baseUrl: settings.knowledgeEmbeddingBaseUrl, + modelName: settings.knowledgeEmbeddingModel, + authentication: 'api-key', + credential: settings.knowledgeEmbeddingCredential + } + ], + activeEmbeddingConnectionId: legacyEmbeddingConnectionId } } @@ -681,6 +781,26 @@ function normalizeStoredSettings(settings: StoredSettings): StoredSettings { ? compressionSource : ({ kind: 'current' } as const) } + const userEmbeddingConnections = settings.embeddingConnections + .filter((connection) => connection.kind === 'openai-compatible') + .map((connection) => ({ + ...connection, + baseUrl: normalizeModelBaseUrl(connection.baseUrl) + })) + const embeddingConnections: StoredSettings['embeddingConnections'] = [ + { + id: builtinEmbeddingConnectionId, + name: 'GoodBuddy 内置向量模型', + kind: 'builtin' + }, + ...userEmbeddingConnections + ] + const activeEmbeddingConnectionId = embeddingConnections.some( + (connection) => + connection.id === settings.activeEmbeddingConnectionId + ) + ? settings.activeEmbeddingConnectionId + : builtinEmbeddingConnectionId return { ...settings, @@ -698,6 +818,8 @@ function normalizeStoredSettings(settings: StoredSettings): StoredSettings { settings.deepseekHarnessModelSource ), contextCompression, + embeddingConnections, + activeEmbeddingConnectionId, opencodeBaseUrl, opencodeEmbedded: !opencodeBaseUrl } @@ -873,7 +995,7 @@ export class RuntimeSettingsStore { const parsed: unknown = JSON.parse(contents) assertSupportedSettingsVersion( parsed, - 18, + 19, (version) => `当前 GoodBuddy 不支持 Runtime 设置版本 ${version}，请升级应用后重试` ) @@ -881,76 +1003,81 @@ export class RuntimeSettingsStore { if (current.success) { this.settings = current.data } else { - const version17 = - version17StoredSettingsSchema.safeParse(parsed) - if (version17.success) { - this.settings = migrateVersion17(version17.data) + const version18 = + version18StoredSettingsSchema.safeParse(parsed) + if (version18.success) { + this.settings = migrateVersion18(version18.data) } else { - const version16 = - version16StoredSettingsSchema.safeParse(parsed) - if (version16.success) { - this.settings = migrateVersion16(version16.data) + const version17 = + version17StoredSettingsSchema.safeParse(parsed) + if (version17.success) { + this.settings = migrateVersion17(version17.data) } else { - const version15 = - version15StoredSettingsSchema.safeParse(parsed) - if (version15.success) { - this.settings = migrateVersion15(version15.data) + const version16 = + version16StoredSettingsSchema.safeParse(parsed) + if (version16.success) { + this.settings = migrateVersion16(version16.data) } else { - const version14 = - version14StoredSettingsSchema.safeParse(parsed) - if (version14.success) { - this.settings = migrateVersion14(version14.data) + const version15 = + version15StoredSettingsSchema.safeParse(parsed) + if (version15.success) { + this.settings = migrateVersion15(version15.data) } else { - const version13 = - version13StoredSettingsSchema.safeParse(parsed) - if (version13.success) { - this.settings = migrateVersion13(version13.data) + const version14 = + version14StoredSettingsSchema.safeParse(parsed) + if (version14.success) { + this.settings = migrateVersion14(version14.data) } else { - const version12 = - version12StoredSettingsSchema.safeParse(parsed) - if (version12.success) { - this.settings = migrateVersion12(version12.data) + const version13 = + version13StoredSettingsSchema.safeParse(parsed) + if (version13.success) { + this.settings = migrateVersion13(version13.data) } else { - const version11 = - version11StoredSettingsSchema.safeParse(parsed) - if (version11.success) { - this.settings = migrateVersion11(version11.data) + const version12 = + version12StoredSettingsSchema.safeParse(parsed) + if (version12.success) { + this.settings = migrateVersion12(version12.data) } else { - const version10 = - version10StoredSettingsSchema.safeParse(parsed) - if (version10.success) { - this.settings = migrateVersion10(version10.data) + const version11 = + version11StoredSettingsSchema.safeParse(parsed) + if (version11.success) { + this.settings = migrateVersion11(version11.data) } else { - const version9 = - version9StoredSettingsSchema.safeParse(parsed) - if (version9.success) { - this.settings = migrateVersion9(version9.data) + const version10 = + version10StoredSettingsSchema.safeParse(parsed) + if (version10.success) { + this.settings = migrateVersion10(version10.data) } else { - const version8 = - version8StoredSettingsSchema.safeParse(parsed) - if (version8.success) { - this.settings = migrateVersion8(version8.data) + const version9 = + version9StoredSettingsSchema.safeParse(parsed) + if (version9.success) { + this.settings = migrateVersion9(version9.data) } else { - const version7 = - version7StoredSettingsSchema.safeParse(parsed) - if (version7.success) { - this.settings = migrateVersion7(version7.data) + const version8 = + version8StoredSettingsSchema.safeParse(parsed) + if (version8.success) { + this.settings = migrateVersion8(version8.data) } else { - const version6 = - version6StoredSettingsSchema.safeParse(parsed) - if (version6.success) { - this.settings = migrateVersion6(version6.data) + const version7 = + version7StoredSettingsSchema.safeParse(parsed) + if (version7.success) { + this.settings = migrateVersion7(version7.data) } else { - const version5 = - version5StoredSettingsSchema.safeParse(parsed) - if (version5.success) { - this.settings = migrateVersion5(version5.data) + const version6 = + version6StoredSettingsSchema.safeParse(parsed) + if (version6.success) { + this.settings = migrateVersion6(version6.data) } else { - const version4 = - version4StoredSettingsSchema.safeParse(parsed) - if (version4.success) { - this.settings = migrateVersion4(version4.data) + const version5 = + version5StoredSettingsSchema.safeParse(parsed) + if (version5.success) { + this.settings = migrateVersion5(version5.data) } else { + const version4 = + version4StoredSettingsSchema.safeParse(parsed) + if (version4.success) { + this.settings = migrateVersion4(version4.data) + } else { const version3 = version3StoredSettingsSchema.safeParse(parsed) if (version3.success) { @@ -1007,6 +1134,7 @@ export class RuntimeSettingsStore { }) } } + } } } } @@ -1095,6 +1223,64 @@ export class RuntimeSettingsStore { } } + private getEmbeddingConnectionApiKey( + connection: Extract< + StoredSettings['embeddingConnections'][number], + { kind: 'openai-compatible' } + > + ): string | undefined { + if (!connection.credential) { + return undefined + } + if (!this.cipher.isAvailable()) { + this.addWarning({ + code: 'runtime-embedding-credential-unreadable' + }) + return undefined + } + try { + return endpointCredentialPayloadSchema.parse( + decryptSettingsCredential(this.cipher, connection.credential) + ).apiKey + } catch { + this.addWarning({ + code: 'runtime-embedding-credential-unreadable' + }) + return undefined + } + } + + private resolveEmbeddingConnections( + settings: StoredSettings + ): ResolvedEmbeddingConnection[] { + const environmentApiKey = + this.environment.GOODBUDDY_EMBEDDING_API_KEY?.trim() + return settings.embeddingConnections.map((connection) => { + if (connection.kind === 'builtin') { + return connection + } + const storedApiKey = + connection.id === legacyEmbeddingConnectionId && + environmentApiKey + ? undefined + : this.getEmbeddingConnectionApiKey(connection) + return { + id: connection.id, + name: connection.name, + kind: connection.kind, + baseUrl: connection.baseUrl, + modelName: connection.modelName, + authentication: connection.authentication, + apiKey: + connection.authentication === 'api-key' + ? connection.id === legacyEmbeddingConnectionId + ? environmentApiKey || storedApiKey + : storedApiKey + : undefined + } + }) + } + private getStoredRerankApiKey( settings: StoredSettings ): string | undefined { @@ -1429,6 +1615,44 @@ export class RuntimeSettingsStore { const embeddingStoredApiKey = embeddingEnvironmentApiKey ? undefined : this.getStoredEmbeddingApiKey(settings) + const resolvedEmbeddingConnections = + this.resolveEmbeddingConnections(settings) + const embeddingConnections = settings.embeddingConnections.map( + (connection) => { + if (connection.kind === 'builtin') { + return { + ...connection, + apiKeyConfigured: false as const, + credentialSource: 'none' as const + } + } + const resolved = resolvedEmbeddingConnections.find( + (candidate) => candidate.id === connection.id + ) + const environmentManaged = + connection.id === legacyEmbeddingConnectionId && + Boolean(embeddingEnvironmentApiKey) + return { + id: connection.id, + name: connection.name, + kind: connection.kind, + baseUrl: connection.baseUrl, + modelName: connection.modelName, + authentication: connection.authentication, + apiKeyConfigured: + resolved?.kind === 'openai-compatible' && + Boolean(resolved.apiKey), + credentialSource: environmentManaged + ? ('environment' as const) + : resolved?.kind === 'openai-compatible' && + resolved.apiKey + ? ('encrypted' as const) + : connection.credential + ? ('unreadable' as const) + : ('none' as const) + } + } + ) const rerankEnvironmentApiKey = this.environment.GOODBUDDY_RERANK_API_KEY?.trim() const rerankStoredApiKey = rerankEnvironmentApiKey @@ -1464,6 +1688,9 @@ export class RuntimeSettingsStore { : settings.knowledgeEmbeddingCredential ? 'unreadable' : 'none', + embeddingConnections, + activeEmbeddingConnectionId: + settings.activeEmbeddingConnectionId, knowledgeRerankEnabled: settings.knowledgeRerankEnabled, knowledgeRerankEndpoint: settings.knowledgeRerankEndpoint, knowledgeRerankModel: settings.knowledgeRerankModel, @@ -1588,6 +1815,16 @@ export class RuntimeSettingsStore { settings.deepseekHarnessModelSource.profileId ) : this.resolvePlatformHarnessProfile() + const embeddingConnections = + this.resolveEmbeddingConnections(settings) + const activeEmbeddingConnection = + embeddingConnections.find( + (connection) => + connection.id === settings.activeEmbeddingConnectionId + ) + if (!activeEmbeddingConnection) { + throw new Error('当前向量连接不存在') + } return { provider: settings.provider, modelBaseUrl: effective.baseUrl, @@ -1611,6 +1848,10 @@ export class RuntimeSettingsStore { knowledgeEmbeddingApiKey: this.environment.GOODBUDDY_EMBEDDING_API_KEY?.trim() || this.getStoredEmbeddingApiKey(settings), + embeddingConnections, + activeEmbeddingConnectionId: + settings.activeEmbeddingConnectionId, + activeEmbeddingConnection, knowledgeRerankEnabled: settings.knowledgeRerankEnabled, knowledgeRerankEndpoint: settings.knowledgeRerankEndpoint, knowledgeRerankModel: settings.knowledgeRerankModel, @@ -1680,6 +1921,12 @@ export class RuntimeSettingsStore { profile.apiKey.action === 'replace' ) || input.knowledgeEmbeddingApiKey?.action === 'replace' || + input.embeddingConnections?.some( + (connection) => + connection.kind === 'openai-compatible' && + connection.authentication === 'api-key' && + connection.apiKey.action === 'replace' + ) || input.knowledgeRerankApiKey?.action === 'replace' ) && !this.cipher.isAvailable() @@ -1732,6 +1979,71 @@ export class RuntimeSettingsStore { return nextProfile }) + const embeddingConnections: StoredSettings['embeddingConnections'] = + input.embeddingConnections + ? input.embeddingConnections.map((connection) => { + if (connection.kind === 'builtin') { + return { + id: builtinEmbeddingConnectionId, + name: 'GoodBuddy 内置向量模型', + kind: 'builtin' + } + } + const existing = current.embeddingConnections.find( + (candidate) => + candidate.kind === 'openai-compatible' && + candidate.id === connection.id + ) + const nextConnection: Extract< + StoredSettings['embeddingConnections'][number], + { kind: 'openai-compatible' } + > = { + id: connection.id, + name: connection.name, + kind: connection.kind, + baseUrl: normalizeModelBaseUrl(connection.baseUrl), + modelName: connection.modelName, + authentication: connection.authentication + } + if ( + connection.apiKey.action === 'keep' && + existing?.kind === 'openai-compatible' && + existing.credential + ) { + nextConnection.credential = existing.credential + } else if ( + connection.authentication === 'api-key' && + connection.apiKey.action === 'replace' + ) { + nextConnection.credential = encryptSavedApiKey( + this.cipher, + connection.apiKey.value + ) + } + return nextConnection + }) + : current.embeddingConnections + if ( + !embeddingConnections.some( + (connection) => + connection.id === builtinEmbeddingConnectionId && + connection.kind === 'builtin' + ) + ) { + throw new Error('系统向量连接不能删除') + } + const activeEmbeddingConnectionId = + input.activeEmbeddingConnectionId ?? + current.activeEmbeddingConnectionId + if ( + !embeddingConnections.some( + (connection) => + connection.id === activeEmbeddingConnectionId + ) + ) { + throw new Error('当前向量连接不存在') + } + const embeddingEndpoint = new URL( input.knowledgeEmbeddingBaseUrl ).toString() @@ -1750,6 +2062,29 @@ export class RuntimeSettingsStore { embeddingApiKeyUpdate.value ) } + const synchronizedEmbeddingConnections = + input.embeddingConnections + ? embeddingConnections + : embeddingConnections.map((connection) => + connection.id === legacyEmbeddingConnectionId && + connection.kind === 'openai-compatible' + ? { + ...connection, + baseUrl: embeddingEndpoint, + modelName: input.knowledgeEmbeddingModel, + credential: knowledgeEmbeddingCredential + } + : connection + ) + const activeEmbeddingConnection = + synchronizedEmbeddingConnections.find( + (connection) => + connection.id === activeEmbeddingConnectionId + ) + const compatibleEmbeddingConnection = + activeEmbeddingConnection?.kind === 'openai-compatible' + ? activeEmbeddingConnection + : undefined const rerankEndpoint = new URL( input.knowledgeRerankEndpoint @@ -1893,7 +2228,7 @@ export class RuntimeSettingsStore { const next: StoredSettings = { ...current, - version: 18, + version: 19, provider: input.provider, modelProfiles, defaultModelProfileId, @@ -1912,9 +2247,17 @@ export class RuntimeSettingsStore { input.subagentSmartRoutingEnabled ?? current.subagentSmartRoutingEnabled, knowledgeEmbeddingEnabled: input.knowledgeEmbeddingEnabled, - knowledgeEmbeddingBaseUrl: embeddingEndpoint, - knowledgeEmbeddingModel: input.knowledgeEmbeddingModel, - knowledgeEmbeddingCredential, + knowledgeEmbeddingBaseUrl: + compatibleEmbeddingConnection?.baseUrl ?? embeddingEndpoint, + knowledgeEmbeddingModel: + compatibleEmbeddingConnection?.modelName ?? + input.knowledgeEmbeddingModel, + knowledgeEmbeddingCredential: + compatibleEmbeddingConnection + ? compatibleEmbeddingConnection.credential + : knowledgeEmbeddingCredential, + embeddingConnections: synchronizedEmbeddingConnections, + activeEmbeddingConnectionId, knowledgeRerankEnabled: input.knowledgeRerankEnabled, knowledgeRerankEndpoint: rerankEndpoint, knowledgeRerankModel: input.knowledgeRerankModel, @@ -1946,7 +2289,7 @@ export class RuntimeSettingsStore { const current = await this.load() const next: StoredSettings = { ...current, - version: 18, + version: 19, runtimeCustomization: parsed } await writeJsonFileAtomically(this.filePath, next) diff --git a/src/preload/index.ts b/src/preload/index.ts index 8d3b28e3..96028aff 100644 --- a/src/preload/index.ts +++ b/src/preload/index.ts @@ -84,6 +84,7 @@ import type { } from '../shared/speech-model-contracts' import type { EmbeddingDiagnosticResult, + EmbeddingModelSnapshot, EmbeddingSettingsSnapshot, KnowledgeEmbeddingIndexSnapshot } from '../shared/embedding-contracts' @@ -485,10 +486,39 @@ const desktopApi: DesktopApi = { ipcRenderer.invoke( ipcChannels.embeddingSettingsGet ) as Promise, - diagnose: () => + diagnose: (connectionId: string) => ipcRenderer.invoke( - ipcChannels.embeddingDiagnose - ) as Promise + ipcChannels.embeddingDiagnose, + { connectionId } + ) as Promise, + setCurrent: (connectionId: string) => + ipcRenderer.invoke( + ipcChannels.embeddingSetCurrent, + { connectionId } + ) as Promise, + installModel: ( + modelId: string, + expectedDownloadSource: ModelDownloadSource + ) => + ipcRenderer.invoke( + ipcChannels.embeddingModelsInstall, + { modelId, expectedDownloadSource } + ) as Promise, + cancelModelOperation: (modelId: string) => + ipcRenderer.invoke( + ipcChannels.embeddingModelsCancel, + { modelId } + ) as Promise, + importModelArchive: (modelId: string) => + ipcRenderer.invoke( + ipcChannels.embeddingModelsImportArchive, + { modelId } + ) as Promise, + removeModel: (modelId: string) => + ipcRenderer.invoke( + ipcChannels.embeddingModelsRemove, + { modelId } + ) as Promise }, documentParsing: { getSnapshot: () => diff --git a/src/preload/preload-sandbox.test.ts b/src/preload/preload-sandbox.test.ts index e921310e..dc6a11d5 100644 --- a/src/preload/preload-sandbox.test.ts +++ b/src/preload/preload-sandbox.test.ts @@ -46,8 +46,31 @@ describe('sandboxed preload', () => { expect(source).toContain('exportArchive: (modelId: string)') expect(source).toContain('importOcrModelArchive: (modelId: string)') expect(source).toContain('exportOcrModelArchive: (modelId: string)') + expect(source).toContain('importModelArchive: (modelId: string)') + expect(source).toContain('ipcChannels.embeddingModelsImportArchive') expect(source).not.toContain('importLocalDirectory:') expect(source).not.toContain('importOcrModel:') + expect(source).not.toMatch( + /importModelArchive:\s*\([^)]*(?:path|directory|archivePath)/u + ) + }) + + it('exposes explicit embedding settings operations without vectors or credentials', () => { + const source = readFileSync( + join(process.cwd(), 'src', 'preload', 'index.ts'), + 'utf8' + ) + const embeddings = + source.match( + /embeddings: \{(?[\s\S]*?)\n {2}\},\n {2}documentParsing:/u + )?.groups?.body ?? '' + expect(embeddings).toContain('getSnapshot:') + expect(embeddings).toContain('diagnose: (connectionId: string)') + expect(embeddings).toContain('setCurrent: (connectionId: string)') + expect(embeddings).toContain('installModel:') + expect(embeddings).toContain('cancelModelOperation:') + expect(embeddings).toContain('removeModel:') + expect(embeddings).not.toMatch(/\b(?:apiKey|credential|vector)\b/iu) }) it('passes an explicit parsing scenario to document diagnostics', () => { diff --git a/src/renderer/src/ActivityPanel.test.tsx b/src/renderer/src/ActivityPanel.test.tsx index a275f62d..2811ef19 100644 --- a/src/renderer/src/ActivityPanel.test.tsx +++ b/src/renderer/src/ActivityPanel.test.tsx @@ -108,6 +108,12 @@ describe('ActivityPanel', () => { expect( screen.getByRole('heading', { level: 1, name: 'Run history' }) ).toBeInTheDocument() + expect(screen.queryByText('RUN HISTORY')).not.toBeInTheDocument() + expect( + screen.queryByText( + 'Review execution details by project, task, and conversation, or browse the activity timeline and model usage.' + ) + ).not.toBeInTheDocument() expect( screen.getByRole('tab', { name: 'Tasks and conversations' }) ).toHaveAttribute('aria-selected', 'true') @@ -234,11 +240,11 @@ describe('ActivityPanel', () => { /> ) - fireEvent.click(screen.getByRole('button', { name: '进行中' })) + fireEvent.click(screen.getByRole('button', { name: '进行中 1' })) expect(screen.getByText('对话：活动 1')).toBeInTheDocument() expect(screen.queryByText('活动 2')).not.toBeInTheDocument() - fireEvent.click(screen.getByRole('button', { name: '异常' })) + fireEvent.click(screen.getByRole('button', { name: '异常 2' })) expect(screen.getByText('对话：活动 2')).toBeInTheDocument() expect(screen.getByText('对话：活动 3')).toBeInTheDocument() expect(screen.queryByText('活动 1')).not.toBeInTheDocument() @@ -299,12 +305,12 @@ describe('ActivityPanel', () => { /> ) - fireEvent.click(screen.getByRole('button', { name: '进行中' })) + fireEvent.click(screen.getByRole('button', { name: '进行中 0' })) expect(screen.getByText('没有匹配的活动')).toBeInTheDocument() fireEvent.click(screen.getByRole('button', { name: '清除筛选' })) expect( - screen.getByRole('button', { name: '全部' }) + screen.getByRole('button', { name: '全部 1' }) ).toHaveAttribute('aria-pressed', 'true') expect(screen.getByText('对话：活动 1')).toBeInTheDocument() }) @@ -521,6 +527,7 @@ describe('ActivityPanel', () => { expect(screen.getByLabelText('筛选活动')).toHaveClass( 'segmented-control' ) + expect(screen.queryByLabelText('活动统计')).not.toBeInTheDocument() expect( screen.queryByLabelText('Token 用量分组') ).not.toBeInTheDocument() @@ -609,10 +616,10 @@ describe('ActivityPanel', () => { ).toBeInTheDocument() expect(within(timeline).getByText('项目：项目甲')).toBeInTheDocument() expect( - within(timeline).getByText( + within(timeline).queryByText( '所有轨道共享同一执行顺序，节点按发生时间依次展开，点击节点查看身份和活动详情。' ) - ).toBeInTheDocument() + ).not.toBeInTheDocument() fireEvent.click( screen.getByRole('button', { diff --git a/src/renderer/src/ActivityPanel.tsx b/src/renderer/src/ActivityPanel.tsx index 7fc9bff8..b1df9f87 100644 --- a/src/renderer/src/ActivityPanel.tsx +++ b/src/renderer/src/ActivityPanel.tsx @@ -299,14 +299,6 @@ export function ActivityPanel({ subagent: t('kinds.subagent'), result: t('kinds.result') } - const filters: ReadonlyArray<{ - value: ActivityFilter - label: string - }> = [ - { value: 'all', label: t('filters.all') }, - { value: 'active', label: t('filters.active') }, - { value: 'failed', label: t('filters.failed') } - ] const views: ReadonlyArray<{ id: ActivityView label: string @@ -428,6 +420,32 @@ export function ActivityPanel({ ) const activeCount = visibleRecords.filter(isActive).length const failedCount = visibleRecords.filter(isFailed).length + const filters: ReadonlyArray<{ + value: ActivityFilter + label: string + }> = [ + { + value: 'all', + label: t('filters.all', { + count: visibleRecords.length, + formattedCount: formatCount(visibleRecords.length) + }) + }, + { + value: 'active', + label: t('filters.active', { + count: activeCount, + formattedCount: formatCount(activeCount) + }) + }, + { + value: 'failed', + label: t('filters.failed', { + count: failedCount, + formattedCount: formatCount(failedCount) + }) + } + ] const tokenTotals = useMemo( () => getTokenUsageTotals(tokenUsage), [tokenUsage] @@ -592,25 +610,6 @@ export function ActivityPanel({ ) } - const recordSummary = ( -
-
-
{t('stats.all')}
-
{formatCount(visibleRecords.length)}
-
-
-
{t('stats.active')}
-
{formatCount(activeCount)}
-
-
-
{t('stats.failed')}
-
{formatCount(failedCount)}
-
-
- ) const recordToolbar = (
} scope={{ kind: 'all-projects' }} @@ -702,7 +699,6 @@ export function ActivityPanel({ role="tabpanel" tabIndex={0} > - {recordSummary} {recordToolbar} {filteredRecords.length === 0 ? ( emptyState @@ -797,7 +793,6 @@ export function ActivityPanel({ role="tabpanel" tabIndex={0} > - {recordSummary} {recordToolbar} {filteredRecords.length === 0 ? ( emptyState @@ -806,9 +801,6 @@ export function ActivityPanel({ aria-label={t('timeline.ariaLabel')} className="activity-tracks" > -

- {t('timeline.description')} -

    { expect( screen.getByRole('button', { name: 'Chat' }) ).toBeInTheDocument() - expect(screen.getByText('Desktop workspace')).toBeInTheDocument() - expect(screen.getByText('GOODBUDDY WORKSPACE')).toBeInTheDocument() + expect(screen.queryByText('Desktop workspace')).not.toBeInTheDocument() + expect(screen.queryByText('GOODBUDDY WORKSPACE')).not.toBeInTheDocument() expect( screen.getByRole('heading', { level: 1, @@ -1074,14 +1074,18 @@ describe('App', () => { expect(screen.getByText('Project: Default project')).toHaveClass( 'scope-badge' ) - expect( - screen.getByText(/Hi, I’m GoodBuddy/u) - ).toBeInTheDocument() + expect(screen.getByText('Hi, I’m GoodBuddy.')).toBeInTheDocument() expect( screen.getByLabelText('Message GoodBuddy') ).toHaveAttribute( 'placeholder', - 'Message GoodBuddy…\nEnter to send · Shift+Enter for a new line · Ctrl+V to paste an image or text' + 'Message GoodBuddy…' + ) + expect( + screen.getByLabelText('Message GoodBuddy') + ).toHaveAttribute( + 'title', + 'Enter to send, Shift+Enter for a new line' ) } finally { cleanup() @@ -1236,8 +1240,12 @@ describe('App', () => { it('renders localized workspace branding in Chinese', async () => { render() - expect(await screen.findByText('桌面工作区')).toBeInTheDocument() - expect(screen.getByText('GOODBUDDY 工作台')).toBeInTheDocument() + await screen.findByRole('heading', { + name: '今天想一起完成什么？' + }) + expect(screen.getByText('你好，我是 GoodBuddy。')).toBeInTheDocument() + expect(screen.queryByText('桌面工作区')).not.toBeInTheDocument() + expect(screen.queryByText('GOODBUDDY 工作台')).not.toBeInTheDocument() expect( document.querySelector('.composer__runtime-toolbar') ).not.toBeInTheDocument() @@ -5316,7 +5324,7 @@ describe('App', () => { within(modeMenu) .getAllByRole('menuitemradio') .map((option) => option.querySelector('span')?.textContent) - ).toEqual(['Ask · 只读问答', 'Execute · 受控执行']) + ).toEqual(['Ask · 只读问答', 'Execute · 完全权限']) fireEvent.click(screen.getByLabelText('新建项目')) const dialog = screen.getByRole('dialog', { name: '新建项目' }) @@ -5327,7 +5335,7 @@ describe('App', () => { within(defaultMode) .getAllByRole('option') .map((option) => option.textContent) - ).toEqual(['Ask · 只读问答', 'Execute · 受控执行']) + ).toEqual(['Ask · 只读问答', 'Execute · 完全权限']) expect(screen.queryByRole('option', { name: /Plan/u })).toBeNull() }) @@ -5387,10 +5395,7 @@ describe('App', () => { ).toHaveAttribute('title', '添加附件') expect( within(contentTools).getByRole('button', { name: '语音输入' }) - ).toHaveAttribute( - 'title', - '语音转文字，转写后可编辑再发送' - ) + ).not.toHaveAttribute('title') const conversationSettings = within(composer!).getByRole( 'group', @@ -5413,7 +5418,7 @@ describe('App', () => { ).toHaveTextContent(/^Ask$/u) expect(screen.getByLabelText('向 GoodBuddy 提问')).toHaveAttribute( 'placeholder', - '给 GoodBuddy 发消息…\nEnter 发送 · Shift+Enter 换行 · Ctrl+V 粘贴图片或文本' + '给 GoodBuddy 发消息…' ) expect( within(conversationSettings).getByRole('button', { @@ -5449,7 +5454,7 @@ describe('App', () => { ).toHaveTextContent(secondProject.name) expect( screen.getByRole('button', { - name: '工作模式：Execute · 受控执行' + name: '工作模式：Execute · 完全权限' }) ).toBeEnabled() @@ -5892,9 +5897,9 @@ describe('App', () => { expect( await screen.findByText('快捷唤起：', { exact: false }) ).toBeInTheDocument() - selectComposerOption('工作模式', 'Execute · 受控执行') + selectComposerOption('工作模式', 'Execute · 完全权限') expect(mode).toHaveAccessibleName( - '工作模式：Execute · 受控执行' + '工作模式：Execute · 完全权限' ) expect(mode).toHaveTextContent(/^Execute$/u) @@ -6325,9 +6330,9 @@ describe('App', () => { name: '工作模式：Ask · 只读问答' }) expect(mode).toBeEnabled() - selectComposerOption('工作模式', 'Execute · 受控执行') + selectComposerOption('工作模式', 'Execute · 完全权限') expect(mode).toHaveAccessibleName( - '工作模式：Execute · 受控执行' + '工作模式：Execute · 完全权限' ) fireEvent.click(await screen.findByRole('button', { name: /OpenCode/u })) @@ -6359,7 +6364,7 @@ describe('App', () => { const modeMenu = openComposerMenu('工作模式') expect( within(modeMenu).getByRole('menuitemradio', { - name: /^Execute · 受控执行/u + name: /^Execute · 完全权限/u }) ).toBeDisabled() expect(mode).toHaveAccessibleName('工作模式：Ask · 只读问答') @@ -6378,7 +6383,7 @@ describe('App', () => { const mode = await screen.findByRole('button', { name: '工作模式：Ask · 只读问答' }) - selectComposerOption('工作模式', 'Execute · 受控执行') + selectComposerOption('工作模式', 'Execute · 完全权限') fireEvent.change(screen.getByLabelText('向 GoodBuddy 提问'), { target: { value: '读取项目文件' } }) @@ -6443,7 +6448,9 @@ describe('App', () => { expect(cancelledDot).not.toHaveClass('message__status-dot--active') fireEvent.click(screen.getByText('运行记录')) expect((await screen.findAllByText('已取消')).length).toBeGreaterThan(0) - fireEvent.click(screen.getByRole('button', { name: '进行中' })) + fireEvent.click( + screen.getByRole('button', { name: /^进行中 \d+$/u }) + ) expect( screen.getByText('当前没有等待中或正在运行的活动。') ).toBeInTheDocument() @@ -7626,7 +7633,7 @@ describe('App', () => { expect((await screen.findAllByText('生图')).length).toBeGreaterThan(0) expect(screen.getByLabelText('向 GoodBuddy 提问')).toHaveAttribute( 'placeholder', - '描述你想生成的图片…\nEnter 发送 · Shift+Enter 换行 · Ctrl+V 粘贴图片或文本' + '描述你想生成的图片…' ) await waitFor(() => expect(api.artifacts.list).toHaveBeenCalled() @@ -8130,17 +8137,20 @@ describe('App', () => { ).not.toBeInTheDocument() fireEvent.click(screen.getByRole('tab', { name: '工作区' })) expect( - screen.getByText(/选择文件后在当前工作区内预览/) + screen.getByRole('button', { name: '刷新工作区文件' }) ).toBeInTheDocument() + expect( + screen.queryByText(/选择文件后在当前工作区内预览/) + ).not.toBeInTheDocument() fireEvent.click(screen.getByRole('tab', { name: '浏览器' })) expect( screen.getByText(/Agent 打开网页后/) ).toBeInTheDocument() fireEvent.click(screen.getByRole('tab', { name: '成果' })) - expect(screen.getByText('生成与导入成果')).toBeInTheDocument() expect( - screen.getByText(/查看并预览由任务、自动化生成或手动导入/) + screen.getByRole('button', { name: '导入 PDF、图片或网页' }) ).toBeInTheDocument() + expect(screen.queryByText('生成与导入成果')).not.toBeInTheDocument() expect( screen.queryByRole('tab', { name: '预览' }) ).not.toBeInTheDocument() diff --git a/src/renderer/src/App.tsx b/src/renderer/src/App.tsx index 05d69f5a..6ef4b987 100644 --- a/src/renderer/src/App.tsx +++ b/src/renderer/src/App.tsx @@ -646,8 +646,7 @@ const chatBottomProximity = 96 function createConversation( projectId?: string, runtimeSelection?: AgentRuntimeSelection, - greeting = - '你好，我是 GoodBuddy。你可以直接向我提问、添加本地文件或使用知识库。启用 Agent Runtime 后，我也可以在你的授权下调用工具。' + greeting = '你好，我是 GoodBuddy。' ): Conversation { const now = Date.now() return { @@ -1025,11 +1024,7 @@ function ChatHistoryPane({
    -

    {t('chat.welcome.eyebrow')}

    {t('chat.welcome.title')}

    -

    - {t('chat.welcome.description')} -

    {quickActions.map((action) => (
    @@ -8440,13 +8434,14 @@ function App(): React.JSX.Element {