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
2026-08-23 11:26:44 +08:00
parent 58d774ae0c
commit 076b3aaa7b
90 changed files with 7820 additions and 1446 deletions
@@ -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')
})
})