feat: add managed Granite embedding models
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 多语言向量模型，支持本地下载、离线导入、连接切换与诊断；同时优化向量设置及相关界面的文案、开关和状态显示。
This commit is contained in:
mesalogo
@@ -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<void> {
|
||||
await worker.shutdown()
|
||||
}
|
||||
|
||||
process.once('disconnect', () => {
|
||||
void shutdown()
|
||||
})
|
||||
process.once('SIGTERM', () => {
|
||||
void shutdown()
|
||||
})
|
||||
+292
-15
@@ -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<EmbeddingInferenceBroker>()
|
||||
let stopRuntimeReconfiguration: (() => Promise<void>) | 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<void>
|
||||
}
|
||||
|
||||
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<ManagedEmbeddingProvider | undefined> {
|
||||
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<typeof createEmbeddingProvider>
|
||||
| Awaited<ReturnType<typeof createEmbeddingProvider>>
|
||||
| undefined
|
||||
let activeRerankProvider:
|
||||
| ReturnType<typeof createRerankProvider>
|
||||
@@ -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<void> => {
|
||||
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()
|
||||
],
|
||||
|
||||
+31
-1
@@ -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,
|
||||
|
||||
+233
-30
@@ -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<unknown>
|
||||
},
|
||||
resolveEmbeddingProvider?: (
|
||||
connectionId: string
|
||||
) => Promise<EmbeddingProvider>,
|
||||
setCurrentEmbeddingConnection?: (
|
||||
connectionId: string
|
||||
) => Promise<void>
|
||||
): () => Promise<void> {
|
||||
type ActiveRequestLease = {
|
||||
controller: AbortController
|
||||
@@ -4339,38 +4356,209 @@ export function registerIpcHandlers(
|
||||
}
|
||||
)
|
||||
|
||||
const requireEmbeddingProvider = async (): Promise<OpenAIEmbeddingClient> => {
|
||||
const requireEmbeddingProvider = async (
|
||||
connectionId: string
|
||||
): Promise<EmbeddingProvider> => {
|
||||
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<void>
|
||||
}
|
||||
).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(
|
||||
|
||||
@@ -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, {
|
||||
|
||||
@@ -19,6 +19,14 @@ export interface EmbeddingIndexProvider {
|
||||
readonly model: string
|
||||
readonly fingerprint?: string
|
||||
embed(input: readonly string[], signal?: AbortSignal): Promise<number[][]>
|
||||
embedQuery?(
|
||||
input: readonly string[],
|
||||
signal?: AbortSignal
|
||||
): Promise<number[][]>
|
||||
embedDocuments?(
|
||||
input: readonly string[],
|
||||
signal?: AbortSignal
|
||||
): Promise<number[][]>
|
||||
}
|
||||
|
||||
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',
|
||||
|
||||
@@ -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<string, unknown>): Record<string, unknown> {
|
||||
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'
|
||||
})
|
||||
})
|
||||
})
|
||||
@@ -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<EmbeddingInferenceTransport>
|
||||
|
||||
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<typeof setTimeout>
|
||||
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<string, PendingRequest>()
|
||||
private readonly retiredRequestIds = new Set<string>()
|
||||
private active?: ActiveTransport
|
||||
private starting?: Promise<ActiveTransport>
|
||||
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<number[][]> {
|
||||
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<number[][]>((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<void> {
|
||||
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<void>((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<ActiveTransport> {
|
||||
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)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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<string, unknown> {
|
||||
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<string, unknown>,
|
||||
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<string, unknown>): 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
|
||||
}
|
||||
}
|
||||
@@ -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> = T extends unknown
|
||||
? Omit<T, 'protocol' | 'version'>
|
||||
: 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>
|
||||
): 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<EmbeddingInferenceEngine['embed']>(
|
||||
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<readonly (readonly number[])[]>((_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<readonly (readonly number[])[]>((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')
|
||||
})
|
||||
})
|
||||
@@ -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<unknown>
|
||||
}
|
||||
|
||||
export interface EmbeddingOnnxRuntime {
|
||||
run(options: {
|
||||
readonly modelDirectory: string
|
||||
readonly inputs: unknown
|
||||
readonly signal: AbortSignal
|
||||
}): Promise<unknown>
|
||||
}
|
||||
|
||||
/**
|
||||
* 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<readonly (readonly number[])[]>
|
||||
dispose?(): void | Promise<void>
|
||||
}
|
||||
|
||||
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<void>
|
||||
}
|
||||
|
||||
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<string, RunningRequest>()
|
||||
private readonly removeMessageListener: () => void
|
||||
private shuttingDown = false
|
||||
private disposed = false
|
||||
private shutdownTask?: Promise<void>
|
||||
|
||||
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<void> {
|
||||
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<EmbeddingInferenceRequest, { type: 'embed' }>
|
||||
): 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<EmbeddingInferenceRequest, { type: 'embed' }>,
|
||||
controller: AbortController
|
||||
): Promise<void> {
|
||||
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<void> {
|
||||
if (this.shutdownTask) {
|
||||
return this.shutdownTask
|
||||
}
|
||||
this.shuttingDown = true
|
||||
const task = (async (): Promise<void> => {
|
||||
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)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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'
|
||||
)
|
||||
}
|
||||
}
|
||||
]
|
||||
}
|
||||
])
|
||||
@@ -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)
|
||||
})
|
||||
})
|
||||
@@ -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<typeof embeddingModelFileSchema>
|
||||
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
|
||||
>
|
||||
@@ -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<string, string>
|
||||
): 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<string, string>,
|
||||
overrides: {
|
||||
kind?: string
|
||||
modelId?: string
|
||||
extraEntries?: Record<string, Uint8Array>
|
||||
hashOverrides?: Record<string, string>
|
||||
} = {}
|
||||
): 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<string, string>
|
||||
}> {
|
||||
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<string, string> = {
|
||||
'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()
|
||||
})
|
||||
})
|
||||
@@ -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<ModelDownloadSource>
|
||||
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<ModelDownloadSource>
|
||||
private readonly maxFileBytes: number
|
||||
private readonly operations = new Map<string, ActiveOperation>()
|
||||
|
||||
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<EmbeddingModelSnapshot> {
|
||||
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<EmbeddingModelSnapshot> {
|
||||
return this.getSnapshot()
|
||||
}
|
||||
|
||||
getProgressSnapshot(): EmbeddingModelProgressSnapshot {
|
||||
return embeddingModelProgressSnapshotSchema.parse({
|
||||
operations: [...this.operations.values()].map(({ progress }) => ({
|
||||
...progress
|
||||
}))
|
||||
})
|
||||
}
|
||||
|
||||
async getStatus(modelId: string): Promise<EmbeddingModelStatus> {
|
||||
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<string> {
|
||||
const entry = this.requireAvailableEntry(modelId)
|
||||
await this.verifyInstalledModel(entry)
|
||||
return this.modelDirectory(entry.id)
|
||||
}
|
||||
|
||||
async registerLocalDirectory(
|
||||
modelId: string,
|
||||
sourceDirectory: string,
|
||||
externalSignal?: AbortSignal
|
||||
): Promise<InstalledEmbeddingModel> {
|
||||
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<InstalledEmbeddingModel> {
|
||||
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<InstalledEmbeddingModel> {
|
||||
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<void> {
|
||||
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<void> {
|
||||
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<void> {
|
||||
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<void> {
|
||||
try {
|
||||
await lstat(this.modelDirectory(modelId))
|
||||
throw new Error('向量模型已安装')
|
||||
} catch (error) {
|
||||
if (isMissingFileError(error)) {
|
||||
return
|
||||
}
|
||||
throw error
|
||||
}
|
||||
}
|
||||
|
||||
private async hasInstalledDirectory(modelId: string): Promise<boolean> {
|
||||
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<void> {
|
||||
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<InstalledEmbeddingModel> {
|
||||
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<InstalledEmbeddingModel[]> {
|
||||
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<void> {
|
||||
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<void> {
|
||||
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: '向量模型路径超出受管目录'
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -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)
|
||||
)
|
||||
}
|
||||
})
|
||||
})
|
||||
|
||||
@@ -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<string, unknown>
|
||||
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 {
|
||||
|
||||
@@ -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<string, readonly number[]>): 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<GraniteSessionLike['run']>(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')
|
||||
})
|
||||
})
|
||||
@@ -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<number>
|
||||
}
|
||||
|
||||
export interface GraniteSessionLike {
|
||||
readonly inputNames: readonly string[]
|
||||
readonly outputNames: readonly string[]
|
||||
run(
|
||||
feeds: Readonly<Record<string, unknown>>
|
||||
): Promise<Readonly<Record<string, GraniteTensorLike>>>
|
||||
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<GraniteSessionLike>
|
||||
readonly createInt64Tensor?: (
|
||||
data: BigInt64Array,
|
||||
dimensions: readonly [number, number]
|
||||
) => unknown
|
||||
readonly readFile?: (filePath: string) => Promise<Uint8Array>
|
||||
}
|
||||
|
||||
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<string, unknown> {
|
||||
if (typeof value !== 'object' || value === null || Array.isArray(value)) {
|
||||
throw new TypeError(`${name} must contain a JSON object`)
|
||||
}
|
||||
return value as Record<string, unknown>
|
||||
}
|
||||
|
||||
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<Record<string, GraniteTensorLike>>
|
||||
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<GraniteEmbeddingEngineDependencies>
|
||||
private loadTask?: Promise<LoadedEngine>
|
||||
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<Record<string, unknown>>
|
||||
)
|
||||
abortIfRequested(signal)
|
||||
return output
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async embed(options: Parameters<EmbeddingInferenceEngine['embed']>[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<string, unknown> = {
|
||||
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<void> {
|
||||
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<LoadedEngine> {
|
||||
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<LoadedEngine> {
|
||||
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 }
|
||||
}
|
||||
}
|
||||
@@ -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])
|
||||
|
||||
@@ -180,6 +180,26 @@ interface PreparedQueryEmbedding {
|
||||
degradationReason?: string
|
||||
}
|
||||
|
||||
function embedKnowledgeDocuments(
|
||||
provider: EmbeddingProvider,
|
||||
input: readonly string[],
|
||||
signal?: AbortSignal
|
||||
): Promise<number[][]> {
|
||||
return provider.embedDocuments
|
||||
? provider.embedDocuments(input, signal)
|
||||
: provider.embed(input, signal)
|
||||
}
|
||||
|
||||
function embedKnowledgeQuery(
|
||||
provider: EmbeddingProvider,
|
||||
input: readonly string[],
|
||||
signal?: AbortSignal
|
||||
): Promise<number[][]> {
|
||||
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
|
||||
)
|
||||
|
||||
@@ -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<typeof fetch>(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<typeof fetch>(
|
||||
async (_input, init) =>
|
||||
|
||||
@@ -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<OpenAIEmbeddingEncodingRecipe>
|
||||
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<OpenAIEmbeddingEncodingRecipe> | 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<number[][]> {
|
||||
return this.embedRole(input, 'document', signal)
|
||||
}
|
||||
|
||||
async embedQuery(
|
||||
input: readonly string[],
|
||||
signal?: AbortSignal
|
||||
): Promise<number[][]> {
|
||||
return this.embedRole(input, 'query', signal)
|
||||
}
|
||||
|
||||
async embedDocuments(
|
||||
input: readonly string[],
|
||||
signal?: AbortSignal
|
||||
): Promise<number[][]> {
|
||||
return this.embedRole(input, 'document', signal)
|
||||
}
|
||||
|
||||
async embedRole(
|
||||
input: readonly string[],
|
||||
role: EmbeddingInputRole,
|
||||
signal?: AbortSignal
|
||||
): Promise<number[][]> {
|
||||
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)
|
||||
|
||||
@@ -158,6 +158,14 @@ export interface EmbeddingProvider {
|
||||
readonly model: string
|
||||
readonly fingerprint?: string
|
||||
embed(input: readonly string[], signal?: AbortSignal): Promise<number[][]>
|
||||
embedQuery?(
|
||||
input: readonly string[],
|
||||
signal?: AbortSignal
|
||||
): Promise<number[][]>
|
||||
embedDocuments?(
|
||||
input: readonly string[],
|
||||
signal?: AbortSignal
|
||||
): Promise<number[][]>
|
||||
}
|
||||
|
||||
export interface RerankProviderResult {
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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<string, unknown>
|
||||
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<string, unknown>
|
||||
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<Record<string, unknown>>
|
||||
}
|
||||
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<Record<string, unknown>>
|
||||
}
|
||||
expect(persisted.version).toBe(18)
|
||||
expect(persisted.version).toBe(19)
|
||||
expect(persisted.modelProfiles[0]).not.toHaveProperty('credential')
|
||||
})
|
||||
|
||||
|
||||
@@ -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<typeof storedSettingsSchema>
|
||||
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)
|
||||
|
||||
Reference in New Issue
Block a user