fix: publish knowledge rebuilds atomically

This commit is contained in:
lofyer
2026-08-13 06:35:07 +08:00
parent 0c46afba59
commit 7a078c6ffe
7 changed files with 1618 additions and 254 deletions
+2 -2
View File
@@ -4454,9 +4454,9 @@ export function registerIpcHandlers(
ipcMain.handle(
ipcChannels.knowledgeDeleteChunk,
(event, input: unknown) => {
async (event, input: unknown) => {
assertTrustedSender(event, window)
const deleted = knowledgeService.deleteChunk(
const deleted = await knowledgeService.deleteChunk(
knowledgeChunkDeleteInputSchema.parse(input)
)
if (!deleted) {
@@ -558,6 +558,65 @@ describe('extraction strategies', () => {
).rejects.toThrow('Model extraction is unavailable')
})
it('extracts bounded batches beyond the first chunk window', async () => {
const chunks = Array.from(
{ length: GRAPH_LIMITS.maximumChunks + 1 },
(_, index) => ({
id: `batch-${index}`,
content:
index === GRAPH_LIMITS.maximumChunks
? '# Late Batch Entity'
: 'ordinary text'
})
)
const modelCalls = vi.fn(async (prompt: string) => {
const parsed = JSON.parse(
prompt
.split('<UNTRUSTED_DOCUMENT_JSON>')[1]!
.split('</UNTRUSTED_DOCUMENT_JSON>')[0]!
) as Array<{ chunkId: string; content: string }>
const chunk = parsed[0]!
return {
entities: chunk.content.includes('Late Batch')
? [{
id: 'late',
name: 'Late Batch Entity',
evidence: [{
chunkId: chunk.chunkId,
start: 2,
end: chunk.content.length
}]
}]
: [],
relations: []
}
})
const rules = await extractKnowledgeGraph(chunks, { strategy: 'rules' })
const model = await extractKnowledgeGraph(chunks, {
strategy: 'model',
extractStructured: modelCalls
})
expect(rules.entities.some((entity) =>
entity.name === 'Late Batch Entity'
)).toBe(true)
expect(model.entities.some((entity) =>
entity.name === 'Late Batch Entity'
)).toBe(true)
expect(modelCalls).toHaveBeenCalledTimes(2)
expect(
modelCalls.mock.calls.every(([prompt]) => {
const parsed = JSON.parse(
prompt
.split('<UNTRUSTED_DOCUMENT_JSON>')[1]!
.split('</UNTRUSTED_DOCUMENT_JSON>')[0]!
) as unknown[]
return parsed.length <= GRAPH_LIMITS.maximumChunks
})
).toBe(true)
})
it('honors cancellation before and after the injected model callback', async () => {
const preCancelled = new AbortController()
preCancelled.abort()
+53 -26
View File
@@ -21,6 +21,7 @@ export const GRAPH_LIMITS = {
maximumWarnings: 20,
maximumWarningLength: 240
} as const
const maximumEvidencePerRecord = 20
export type ExtractionStrategy = 'rules' | 'model' | 'hybrid' | 'ask'
@@ -192,7 +193,7 @@ function throwIfAborted(signal?: AbortSignal): void {
function prepareChunks(chunks: readonly GraphChunk[]): GraphChunk[] {
const ids = new Set<string>()
const prepared: GraphChunk[] = []
for (const chunk of chunks.slice(0, GRAPH_LIMITS.maximumChunks)) {
for (const chunk of chunks) {
const id = truncate(chunk.id.trim(), GRAPH_LIMITS.maximumFieldLength)
if (!id || ids.has(id)) {
continue
@@ -219,6 +220,9 @@ function mergeEvidence(
const key = evidenceKey(evidence)
if (!merged.has(key)) {
merged.set(key, evidence)
if (merged.size >= maximumEvidencePerRecord) {
break
}
}
}
return [...merged.values()]
@@ -805,41 +809,64 @@ export async function extractKnowledgeGraph(
const context = createOntologyContext(options.ontology)
throwIfAborted(options.signal)
const prepared = prepareChunks(chunks)
const rules =
strategy === 'rules' || strategy === 'hybrid' || strategy === 'ask'
? extractGraphWithRulesInternal(prepared, options.signal, context)
: emptyGraph()
if (strategy === 'rules' || strategy === 'ask') {
if (prepared.length === 0) {
return {
...rules,
...emptyGraph(),
strategy,
requiresModelApproval: strategy === 'ask',
warnings: [...context.warnings]
}
}
if (!options.extractStructured) {
throw new Error('Model extraction is unavailable')
let graph = emptyGraph()
for (
let offset = 0;
offset < prepared.length;
offset += GRAPH_LIMITS.maximumChunks
) {
throwIfAborted(options.signal)
const batch = prepared.slice(offset, offset + GRAPH_LIMITS.maximumChunks)
const batchContext = createOntologyContext(
context.settings,
context.warnings
)
const rules =
strategy === 'rules' || strategy === 'hybrid' || strategy === 'ask'
? extractGraphWithRulesInternal(batch, options.signal, batchContext)
: emptyGraph()
if (strategy === 'rules' || strategy === 'ask') {
graph = mergeKnowledgeGraphsInternal(graph, rules, context)
continue
}
if (!options.extractStructured) {
throw new Error('Model extraction is unavailable')
}
const output = await options.extractStructured(
createModelPrompt(batch, context.settings),
options.signal
)
throwIfAborted(options.signal)
const parsedOutput = parseModelOutput(output)
if (!modelEnvelopeSchema.safeParse(parsedOutput).success) {
throw new Error('模型返回的图谱结构无效')
}
const model = validateModelGraphInternal(
parsedOutput,
batch,
batchContext
)
graph = mergeKnowledgeGraphsInternal(
graph,
strategy === 'hybrid'
? mergeKnowledgeGraphsInternal(rules, model, batchContext)
: model,
context
)
}
const output = await options.extractStructured(
createModelPrompt(prepared, context.settings),
options.signal
)
throwIfAborted(options.signal)
const parsedOutput = parseModelOutput(output)
if (!modelEnvelopeSchema.safeParse(parsedOutput).success) {
throw new Error('模型返回的图谱结构无效')
}
const model = validateModelGraphInternal(parsedOutput, prepared, context)
const graph =
strategy === 'hybrid'
? mergeKnowledgeGraphsInternal(rules, model, context)
: model
return {
...graph,
strategy,
requiresModelApproval: false,
warnings: [...context.warnings]
requiresModelApproval: strategy === 'ask',
warnings: [...context.warnings].slice(0, GRAPH_LIMITS.maximumWarnings)
}
}
+199 -6
View File
@@ -86,7 +86,7 @@ describe('KnowledgeDatabase', () => {
const inspection = new DatabaseSync(path)
expect(
inspection.prepare('PRAGMA user_version').get()
).toEqual({ user_version: 10 })
).toEqual({ user_version: 11 })
expect(
inspection
.prepare('SELECT version FROM schema_migrations ORDER BY version')
@@ -101,7 +101,8 @@ describe('KnowledgeDatabase', () => {
{ version: 7 },
{ version: 8 },
{ version: 9 },
{ version: 10 }
{ version: 10 },
{ version: 11 }
])
inspection.close()
@@ -193,7 +194,7 @@ describe('KnowledgeDatabase', () => {
})
})
it('upgrades an existing v1 database through knowledge schema v10', async () => {
it('upgrades an existing v1 database through knowledge schema v11', async () => {
const { database, path } = await createDatabase()
const knowledgeBase = database.createKnowledgeBase({
name: 'Version one data',
@@ -210,7 +211,7 @@ describe('KnowledgeDatabase', () => {
DROP TABLE chunk_embeddings;
DROP TABLE knowledge_tasks;
DELETE FROM schema_migrations
WHERE version IN (2, 3, 4, 5, 6, 7, 8, 9, 10);
WHERE version IN (2, 3, 4, 5, 6, 7, 8, 9, 10, 11);
PRAGMA user_version = 1;
`)
downgrade.close()
@@ -220,7 +221,7 @@ describe('KnowledgeDatabase', () => {
upgraded.initialize()
const inspection = new DatabaseSync(path)
expect(inspection.prepare('PRAGMA user_version').get()).toEqual({
user_version: 10
user_version: 11
})
expect(
inspection
@@ -327,7 +328,7 @@ describe('KnowledgeDatabase', () => {
WHERE new.enabled = 1 AND new.role <> 'parent';
END;
INSERT INTO chunks_fts(chunks_fts) VALUES ('rebuild');
DELETE FROM schema_migrations WHERE version IN (9, 10);
DELETE FROM schema_migrations WHERE version IN (9, 10, 11);
PRAGMA user_version = 8;
`)
downgrade
@@ -700,6 +701,22 @@ describe('KnowledgeDatabase', () => {
).toEqual([])
})
it('lists complete source and document snapshots beyond display limits', async () => {
const { database } = await createDatabase()
const library = database.createKnowledgeBase({
name: 'Complete snapshot',
storageMode: 'reference'
})
for (let index = 0; index < 501; index += 1) {
seedDocument(database, library.id, `snapshot-${index}`)
}
expect(database.listSources(library.id)).toHaveLength(500)
expect(database.listDocuments(library.id)).toHaveLength(500)
expect(database.listSourcesForSnapshot(library.id)).toHaveLength(501)
expect(database.listDocumentsForSnapshot(library.id)).toHaveLength(501)
})
it('edits graph records and merges entities while retaining evidence and locks', async () => {
const { database } = await createDatabase()
const knowledgeBase = database.createKnowledgeBase({
@@ -1403,6 +1420,182 @@ describe('KnowledgeDatabase', () => {
database.discardDocumentEmbeddingReplacement(replacementId)
})
it('publishes candidate chunks, vectors, and graph evidence atomically', async () => {
const { database } = await createDatabase()
const library = database.createKnowledgeBase({
name: 'Atomic publication',
storageMode: 'reference',
graphEnabled: true,
graphStrategy: 'rules'
})
const seeded = seedDocument(database, library.id, 'atomic-publication')
const originalContent =
'atomic-publication contains the searchable lighthouse phrase'
database.replaceDocumentEmbeddings(
seeded.documentId,
'provider',
'model',
[{
chunkId: seeded.chunkId,
contentChecksum: createHash('sha256')
.update(originalContent)
.digest('hex'),
vector: [1, 0]
}]
)
const originalEntity = database.createEntity({
knowledgeBaseId: library.id,
name: 'Original Entity',
type: 'CONCEPT'
})
database.createEvidence({
knowledgeBaseId: library.id,
entityId: originalEntity.id,
documentId: seeded.documentId,
chunkId: seeded.chunkId,
quote: originalContent
})
const candidateChunkId = 'atomic-publication-candidate'
const candidateContent = 'candidate searchable content'
const replacementId =
database.beginPreparedDocumentEmbeddingReplacement(
seeded.documentId,
'provider',
'model'
)
database.appendPreparedDocumentEmbeddingBatch(
replacementId,
seeded.documentId,
'provider',
'model',
[{
chunkId: candidateChunkId,
contentChecksum: createHash('sha256')
.update(candidateContent)
.digest('hex'),
vector: [0, 1]
}]
)
expect(() =>
database.publishDocument(
{
id: seeded.documentId,
knowledgeBaseId: library.id,
sourceId: seeded.sourceId,
externalId: 'atomic-publication',
title: 'candidate'
},
[{
id: candidateChunkId,
ordinal: 0,
content: candidateContent
}],
{
embeddingReplacement: {
replacementId,
provider: 'provider',
model: 'model'
},
afterChunksInserted: () => {
throw new Error('synthetic graph failure')
}
}
)
).toThrow('synthetic graph failure')
database.discardDocumentEmbeddingReplacement(replacementId)
expect(database.search({
knowledgeBaseId: library.id,
query: 'lighthouse'
})[0]?.chunk.id).toBe(seeded.chunkId)
expect(database.vectorSearch({
knowledgeBaseId: library.id,
provider: 'provider',
model: 'model',
vector: [1, 0],
limit: 1
})[0]?.chunk.id).toBe(seeded.chunkId)
expect(database.listEvidence(library.id)).toEqual(
expect.arrayContaining([
expect.objectContaining({
documentId: seeded.documentId,
chunkId: seeded.chunkId
})
])
)
const successfulReplacementId =
database.beginPreparedDocumentEmbeddingReplacement(
seeded.documentId,
'provider',
'model'
)
database.appendPreparedDocumentEmbeddingBatch(
successfulReplacementId,
seeded.documentId,
'provider',
'model',
[{
chunkId: candidateChunkId,
contentChecksum: createHash('sha256')
.update(candidateContent)
.digest('hex'),
vector: [0, 1]
}]
)
const candidateEntity = database.createEntity({
knowledgeBaseId: library.id,
name: 'Candidate Entity',
type: 'CONCEPT'
})
database.publishDocument(
{
id: seeded.documentId,
knowledgeBaseId: library.id,
sourceId: seeded.sourceId,
externalId: 'atomic-publication',
title: 'candidate'
},
[{
id: candidateChunkId,
ordinal: 0,
content: candidateContent
}],
{
embeddingReplacement: {
replacementId: successfulReplacementId,
provider: 'provider',
model: 'model'
},
afterChunksInserted: (document) => {
database.createEvidence({
knowledgeBaseId: library.id,
entityId: candidateEntity.id,
documentId: document.id,
chunkId: candidateChunkId,
quote: candidateContent
})
}
}
)
expect(database.search({
knowledgeBaseId: library.id,
query: 'candidate'
})[0]?.chunk.id).toBe(candidateChunkId)
expect(database.vectorSearch({
knowledgeBaseId: library.id,
provider: 'provider',
model: 'model',
vector: [0, 1],
limit: 1
})[0]?.chunk.id).toBe(candidateChunkId)
expect(database.listEvidence(library.id)).toEqual([
expect.objectContaining({
documentId: seeded.documentId,
chunkId: candidateChunkId
})
])
})
it('persists embedding index jobs independently by knowledge base', async () => {
const created = await createDatabase()
let database = created.database
+452 -19
View File
@@ -80,7 +80,7 @@ import type {
VectorSearchOptions
} from './types'
const DATABASE_VERSION = 10
const DATABASE_VERSION = 11
const MAX_ID_LENGTH = 128
const MAX_NAME_LENGTH = 512
const MAX_LOCATION_LENGTH = 8192
@@ -118,6 +118,22 @@ export type HybridSearchResultPage = {
vectorScannedCount: number
}
export type PreparedEmbeddingReplacement = {
replacementId: string
provider: string
model: string
}
export type DocumentPublicationOptions = {
embeddingReplacement?: PreparedEmbeddingReplacement
embeddingError?: {
provider: string
model: string
message: string
}
afterChunksInserted?: (document: Document) => void
}
type Row = Record<string, null | number | bigint | string | Uint8Array>
const taskScopes = knowledgeTaskScopeSchema.options
@@ -1635,6 +1651,22 @@ export class KnowledgeDatabase {
.map(mapSource)
}
listSourcesForSnapshot(knowledgeBaseId: string): KnowledgeSource[] {
const normalizedId = requiredString(
knowledgeBaseId,
'knowledgeBaseId',
MAX_ID_LENGTH
)
return this.requireDatabase()
.prepare(
`SELECT * FROM knowledge_sources
WHERE knowledge_base_id = ?
ORDER BY created_at ASC, id ASC`
)
.all(normalizedId)
.map(mapSource)
}
getSource(id: string): KnowledgeSource | undefined {
const row = this.requireDatabase()
.prepare('SELECT * FROM knowledge_sources WHERE id = ?')
@@ -1654,6 +1686,14 @@ export class KnowledgeDatabase {
upsertDocument(
input: UpsertDocumentInput,
chunks: ReplaceChunkInput[]
): Document {
return this.publishDocument(input, chunks)
}
publishDocument(
input: UpsertDocumentInput,
chunks: ReplaceChunkInput[],
options: DocumentPublicationOptions = {}
): Document {
if (!Array.isArray(chunks) || chunks.length > MAX_CHUNKS) {
throw new RangeError(`chunks must contain at most ${MAX_CHUNKS} items`)
@@ -1716,6 +1756,58 @@ export class KnowledgeDatabase {
)
const checksum = optionalString(input.checksum, 'checksum', 512)
const metadata = jsonObject(input.metadata, 'metadata')
const embeddingReplacement = options.embeddingReplacement
? {
replacementId: requiredString(
options.embeddingReplacement.replacementId,
'embeddingReplacement.replacementId',
MAX_ID_LENGTH
),
provider: requiredString(
options.embeddingReplacement.provider,
'embeddingReplacement.provider',
MAX_EMBEDDING_PROVIDER_LENGTH
),
model: requiredString(
options.embeddingReplacement.model,
'embeddingReplacement.model',
MAX_EMBEDDING_MODEL_LENGTH
)
}
: undefined
const embeddingError = options.embeddingError
? {
provider: requiredString(
options.embeddingError.provider,
'embeddingError.provider',
MAX_EMBEDDING_PROVIDER_LENGTH
),
model: requiredString(
options.embeddingError.model,
'embeddingError.model',
MAX_EMBEDDING_MODEL_LENGTH
),
message: requiredString(
options.embeddingError.message,
'embeddingError.message',
MAX_EMBEDDING_ERROR_LENGTH,
false
)
}
: undefined
if (embeddingReplacement && embeddingError) {
throw new Error(
'Document publication cannot include embeddings and an embedding error'
)
}
if (embeddingReplacement) {
this.validatePreparedDocumentEmbeddings(
database,
embeddingReplacement,
id,
normalizedChunks
)
}
const now = new Date().toISOString()
this.transaction(database, () => {
@@ -1749,7 +1841,7 @@ export class KnowledgeDatabase {
now
)
database
.prepare('DELETE FROM embedding_index_state WHERE document_id = ?')
.prepare('DELETE FROM graph_evidence WHERE document_id = ?')
.run(id)
database.prepare('DELETE FROM chunks WHERE document_id = ?').run(id)
const insertChunk = database.prepare(
@@ -1794,6 +1886,29 @@ export class KnowledgeDatabase {
indexContent
)
}
options.afterChunksInserted?.(this.requiredDocument(id))
database
.prepare(
'DELETE FROM embedding_index_state WHERE document_id = ?'
)
.run(id)
if (embeddingReplacement) {
this.publishPreparedDocumentEmbeddings(
database,
embeddingReplacement,
id,
knowledgeBaseId,
now
)
} else if (embeddingError) {
this.recordEmbeddingIndexError(
id,
embeddingError.provider,
embeddingError.model,
embeddingError.message
)
}
this.pruneUnreferencedGeneratedGraph(knowledgeBaseId)
})
return this.requiredDocument(id)
}
@@ -1851,6 +1966,10 @@ export class KnowledgeDatabase {
.map(mapDocument)
}
listDocumentsForSnapshot(knowledgeBaseId: string): Document[] {
return this.listDocumentsForLibraryRebuild(knowledgeBaseId)
}
getDocumentChunkCounts(knowledgeBaseId: string): Map<string, number> {
const normalizedId = requiredString(
knowledgeBaseId,
@@ -2365,6 +2484,17 @@ export class KnowledgeDatabase {
return randomUUID()
}
beginPreparedDocumentEmbeddingReplacement(
documentId: string,
provider: string,
model: string
): string {
requiredString(documentId, 'documentId', MAX_ID_LENGTH)
requiredString(provider, 'provider', MAX_EMBEDDING_PROVIDER_LENGTH)
requiredString(model, 'model', MAX_EMBEDDING_MODEL_LENGTH)
return randomUUID()
}
appendDocumentEmbeddingBatch(
replacementId: string,
documentId: string,
@@ -2465,26 +2595,95 @@ export class KnowledgeDatabase {
}
return { chunkId, checksum, ...vector }
})
const insert = database.prepare(
`INSERT INTO embedding_rebuild_staging
(replacement_id, document_id, provider, model, chunk_id,
dimensions, content_checksum, vector, magnitude)
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)`
this.insertEmbeddingReplacementBatch(database, {
replacementId: normalizedReplacementId,
documentId: normalizedDocumentId,
provider: normalizedProvider,
model: normalizedModel,
embeddings: normalized
})
}
appendPreparedDocumentEmbeddingBatch(
replacementId: string,
documentId: string,
provider: string,
model: string,
embeddings: readonly ChunkEmbeddingInput[]
): void {
const normalizedReplacementId = requiredString(
replacementId,
'replacementId',
MAX_ID_LENGTH
)
this.transaction(database, () => {
for (const item of normalized) {
insert.run(
normalizedReplacementId,
normalizedDocumentId,
normalizedProvider,
normalizedModel,
item.chunkId,
item.dimensions,
item.checksum,
item.bytes,
item.magnitude
const normalizedDocumentId = requiredString(
documentId,
'documentId',
MAX_ID_LENGTH
)
const normalizedProvider = requiredString(
provider,
'provider',
MAX_EMBEDDING_PROVIDER_LENGTH
)
const normalizedModel = requiredString(
model,
'model',
MAX_EMBEDDING_MODEL_LENGTH
)
if (
!Array.isArray(embeddings) ||
embeddings.length < 1 ||
embeddings.length > MAX_EMBEDDING_BATCH
) {
throw new RangeError(
`embeddings must contain between 1 and ${MAX_EMBEDDING_BATCH} items`
)
}
const database = this.requireDatabase()
const existingDimensions = database
.prepare(
`SELECT dimensions FROM embedding_rebuild_staging
WHERE replacement_id = ? LIMIT 1`
)
.get(normalizedReplacementId)
let dimensions = existingDimensions
? asNumber(existingDimensions, 'dimensions')
: undefined
const seen = new Set<string>()
const normalized = embeddings.map((embedding, index) => {
const chunkId = requiredString(
embedding.chunkId,
`embeddings[${index}].chunkId`,
MAX_ID_LENGTH
)
if (seen.has(chunkId)) {
throw new Error('Embedding chunk IDs must be unique within a batch')
}
seen.add(chunkId)
const checksum = normalizedChecksum(
embedding.contentChecksum,
`embeddings[${index}].contentChecksum`
)
const vector = normalizeVector(
embedding.vector,
`embeddings[${index}].vector`
)
if (dimensions === undefined) {
dimensions = vector.dimensions
} else if (dimensions !== vector.dimensions) {
throw new Error(
'Document embeddings must have consistent dimensions'
)
}
return { chunkId, checksum, ...vector }
})
this.insertEmbeddingReplacementBatch(database, {
replacementId: normalizedReplacementId,
documentId: normalizedDocumentId,
provider: normalizedProvider,
model: normalizedModel,
embeddings: normalized
})
}
@@ -4324,6 +4523,14 @@ export class KnowledgeDatabase {
)
.run(10, new Date().toISOString())
}
if (currentVersion < 11) {
this.migrateToVersion11(database)
database
.prepare(
'INSERT INTO schema_migrations (version, applied_at) VALUES (?, ?)'
)
.run(11, new Date().toISOString())
}
database.exec(`PRAGMA user_version = ${DATABASE_VERSION}`)
database.exec('COMMIT')
} catch (error) {
@@ -4947,6 +5154,228 @@ export class KnowledgeDatabase {
`)
}
private migrateToVersion11(database: DatabaseSync): void {
database.exec('DROP TABLE embedding_rebuild_staging;')
database.exec(`
CREATE TABLE embedding_rebuild_staging (
replacement_id TEXT NOT NULL,
document_id TEXT NOT NULL,
provider TEXT NOT NULL,
model TEXT NOT NULL,
chunk_id TEXT NOT NULL,
dimensions INTEGER NOT NULL
CHECK (dimensions >= 1 AND dimensions <= 8192),
content_checksum TEXT NOT NULL
CHECK (length(content_checksum) = 64),
vector BLOB NOT NULL,
magnitude REAL NOT NULL CHECK (magnitude > 0),
PRIMARY KEY (replacement_id, chunk_id)
);
CREATE INDEX embedding_rebuild_staging_document_idx
ON embedding_rebuild_staging(
document_id, provider, model, replacement_id
);
`)
}
private validatePreparedDocumentEmbeddings(
database: DatabaseSync,
replacement: PreparedEmbeddingReplacement,
documentId: string,
chunks: ReturnType<KnowledgeDatabase['normalizeChunks']>
): void {
const indexableChunks = chunks.filter(
(chunk) => chunk.enabled && chunk.role !== 'parent'
)
if (indexableChunks.length === 0) {
throw new Error(
'Prepared embedding replacement requires indexable chunks'
)
}
const stagedRows = database
.prepare(
`SELECT chunk_id, content_checksum, dimensions
FROM embedding_rebuild_staging
WHERE replacement_id = ? AND document_id = ?
AND provider = ? AND model = ?
ORDER BY chunk_id`
)
.all(
replacement.replacementId,
documentId,
replacement.provider,
replacement.model
)
if (stagedRows.length !== indexableChunks.length) {
throw new Error('Embeddings must cover every candidate document chunk')
}
const chunksById = new Map(
indexableChunks.map((chunk) => [
chunk.id,
contentChecksum(
chunkIndexContent(chunk.content, parseObject(chunk.metadata))
)
])
)
let dimensions: number | undefined
for (const row of stagedRows) {
const chunkId = asString(row, 'chunk_id')
const checksum = chunksById.get(chunkId)
if (
checksum === undefined ||
checksum !== asString(row, 'content_checksum')
) {
throw new Error('Prepared embedding content does not match the chunk')
}
chunksById.delete(chunkId)
const rowDimensions = asNumber(row, 'dimensions')
if (dimensions === undefined) {
dimensions = rowDimensions
} else if (dimensions !== rowDimensions) {
throw new Error(
'Document embeddings must have consistent dimensions'
)
}
}
if (chunksById.size > 0) {
throw new Error('Embeddings must cover every candidate document chunk')
}
}
private publishPreparedDocumentEmbeddings(
database: DatabaseSync,
replacement: PreparedEmbeddingReplacement,
documentId: string,
knowledgeBaseId: string,
now: string
): void {
let dimensions: number | undefined
const indexHash = createHash('sha256')
let firstChecksum = true
for (const row of database
.prepare(
`SELECT chunk_id, content_checksum, dimensions
FROM embedding_rebuild_staging
WHERE replacement_id = ? AND document_id = ?
AND provider = ? AND model = ?
ORDER BY chunk_id`
)
.iterate(
replacement.replacementId,
documentId,
replacement.provider,
replacement.model
)) {
const chunkId = asString(row, 'chunk_id')
const checksum = asString(row, 'content_checksum')
if (!firstChecksum) {
indexHash.update('\n')
}
indexHash.update(`${chunkId}\0${checksum}`)
firstChecksum = false
dimensions ??= asNumber(row, 'dimensions')
}
database
.prepare(
`INSERT INTO chunk_embeddings
(chunk_id, knowledge_base_id, provider, model, dimensions,
content_checksum, vector, magnitude, created_at, updated_at)
SELECT chunk_id, ?, provider, model, dimensions,
content_checksum, vector, magnitude, ?, ?
FROM embedding_rebuild_staging
WHERE replacement_id = ? AND document_id = ?
AND provider = ? AND model = ?`
)
.run(
knowledgeBaseId,
now,
now,
replacement.replacementId,
documentId,
replacement.provider,
replacement.model
)
database
.prepare(
`INSERT INTO embedding_index_state
(document_id, knowledge_base_id, provider, model, dimensions,
content_checksum, status, last_error, updated_at)
VALUES (?, ?, ?, ?, ?, ?, 'ready', NULL, ?)`
)
.run(
documentId,
knowledgeBaseId,
replacement.provider,
replacement.model,
dimensions ?? null,
indexHash.digest('hex'),
now
)
database
.prepare(
'DELETE FROM embedding_rebuild_staging WHERE replacement_id = ?'
)
.run(replacement.replacementId)
}
private insertEmbeddingReplacementBatch(
database: DatabaseSync,
input: {
replacementId: string
documentId: string
provider: string
model: string
embeddings: ReadonlyArray<{
chunkId: string
checksum: string
bytes: Buffer
dimensions: number
magnitude: number
}>
}
): void {
const insert = database.prepare(
`INSERT INTO embedding_rebuild_staging
(replacement_id, document_id, provider, model, chunk_id,
dimensions, content_checksum, vector, magnitude)
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)`
)
this.transaction(database, () => {
const existingTarget = database
.prepare(
`SELECT document_id, provider, model
FROM embedding_rebuild_staging
WHERE replacement_id = ? LIMIT 1`
)
.get(input.replacementId)
if (
existingTarget &&
(
asString(existingTarget, 'document_id') !== input.documentId ||
asString(existingTarget, 'provider') !== input.provider ||
asString(existingTarget, 'model') !== input.model
)
) {
throw new Error(
'Embedding replacement target cannot change between batches'
)
}
for (const item of input.embeddings) {
insert.run(
input.replacementId,
input.documentId,
input.provider,
input.model,
item.chunkId,
item.dimensions,
item.checksum,
item.bytes,
item.magnitude
)
}
})
}
private normalizeChunks(chunks: ReplaceChunkInput[]): Array<{
id: string
ordinal: number
@@ -5175,6 +5604,10 @@ export class KnowledgeDatabase {
}
private transaction(database: DatabaseSync, operation: () => void): void {
if (database.isTransaction) {
operation()
return
}
database.exec('BEGIN IMMEDIATE')
try {
operation()
+226 -1
View File
@@ -464,6 +464,104 @@ describe('KnowledgeService', () => {
expect(service.snapshot(library.id).evidence).toEqual([])
})
it('keeps a committed chunk edit successful when graph refresh fails', async () => {
const extractStructured = vi.fn(async () => {
throw new Error('synthetic edit graph failure')
})
const { directory, service } = await createService(
undefined,
undefined,
extractStructured
)
const sourcePath = join(directory, 'edit-graph-failure.md')
await writeFile(sourcePath, 'initial edit content', 'utf8')
const library = service.createLibrary({
name: 'Edit graph failure',
storageMode: 'reference',
graphEnabled: false,
graphStrategy: 'model'
})
await service.importPaths(library.id, [sourcePath])
service.database.updateKnowledgeBase(library.id, { graphEnabled: true })
const document = service.snapshot(library.id).documents[0]!
const chunk = service.database
.listChunks(document.id)
.find((candidate) => candidate.role !== 'parent')!
await expect(service.updateChunk({
knowledgeBaseId: library.id,
documentId: document.id,
chunkId: chunk.id,
content: 'committed edit survives graph failure'
})).resolves.toMatchObject({
content: 'committed edit survives graph failure'
})
await vi.waitFor(() => expect(extractStructured).toHaveBeenCalled())
expect(service.search(library.id, 'committed')).toHaveLength(1)
})
it('coalesces edit graph refreshes so stale extraction cannot publish last', async () => {
const resolvers: Array<(value: {
entities: unknown[]
relations: unknown[]
}) => void> = []
const extractStructured: ExtractStructured = () =>
new Promise((resolve) => {
resolvers.push(resolve)
})
const { directory, service } = await createService(
undefined,
undefined,
extractStructured
)
const sourcePath = join(directory, 'coalesced-edit-graph.md')
await writeFile(sourcePath, 'initial graph content', 'utf8')
const library = service.createLibrary({
name: 'Coalesced edit graph',
storageMode: 'reference',
graphEnabled: false,
graphStrategy: 'model'
})
await service.importPaths(library.id, [sourcePath])
service.database.updateKnowledgeBase(library.id, { graphEnabled: true })
const document = service.snapshot(library.id).documents[0]!
const chunk = service.database
.listChunks(document.id)
.find((candidate) => candidate.role !== 'parent')!
await service.updateChunk({
knowledgeBaseId: library.id,
documentId: document.id,
chunkId: chunk.id,
content: 'first graph edit'
})
await vi.waitFor(() => expect(resolvers).toHaveLength(1))
await service.updateChunk({
knowledgeBaseId: library.id,
documentId: document.id,
chunkId: chunk.id,
content: 'second graph edit'
})
await new Promise((resolve) => setTimeout(resolve, 0))
expect(resolvers).toHaveLength(1)
resolvers.shift()?.({ entities: [], relations: [] })
await vi.waitFor(() => expect(resolvers).toHaveLength(1))
resolvers.shift()?.({ entities: [], relations: [] })
await vi.waitFor(() => {
expect(
service.snapshot(library.id).tasks.filter(
(task) => task.kind === 'graph' && task.status === 'succeeded'
)
).toHaveLength(1)
})
expect(service.database.listChunks(document.id)).toEqual(
expect.arrayContaining([
expect.objectContaining({ content: 'second graph edit' })
])
)
})
it('fails hybrid reextraction when model extraction fails', async () => {
const extractStructured = vi.fn(async () => {
throw new Error('模型未返回图谱内容')
@@ -1062,6 +1160,133 @@ describe('KnowledgeService', () => {
)
})
it('preserves old chunks, vectors, and graph evidence when rebuild extraction fails', async () => {
let failExtraction = false
const extractStructured: ExtractStructured = async () => {
if (failExtraction) {
throw new Error('synthetic graph failure')
}
return {
entities: [],
relations: []
}
}
const provider: EmbeddingProvider = {
provider: 'atomic-provider',
model: 'atomic-model',
embed: async (input) => input.map(() => [1, 0])
}
const { directory, service } = await createService(
undefined,
provider,
extractStructured
)
const sourcePath = join(directory, 'atomic-rebuild.md')
await writeFile(sourcePath, '# Original Entity\nold searchable text', 'utf8')
const library = service.createLibrary({
name: 'Atomic rebuild',
storageMode: 'reference',
graphEnabled: true,
graphStrategy: 'hybrid'
})
await service.importPaths(library.id, [sourcePath])
const before = service.snapshot(library.id)
const document = before.documents[0]!
const oldChunk = service.database.listChunks(document.id)[0]!
const oldEvidence = before.evidence
failExtraction = true
await writeFile(sourcePath, '# Replacement Entity\nnew searchable text', 'utf8')
await expect(service.rebuildDocument({
knowledgeBaseId: library.id,
documentId: document.id
})).rejects.toThrow('synthetic graph failure')
expect(service.search(library.id, 'old')[0]?.chunk.id).toBe(oldChunk.id)
expect(service.search(library.id, 'new')).toEqual([])
expect(service.database.vectorSearch({
knowledgeBaseId: library.id,
provider: embeddingStorageProvider(provider),
model: provider.model,
vector: [1, 0],
limit: 1
})[0]?.chunk.id).toBe(oldChunk.id)
expect(service.snapshot(library.id).evidence).toEqual(oldEvidence)
})
it('serializes a rebuild and chunk edit so the latest mutation wins', async () => {
let releaseParser: (() => void) | undefined
let parserStarted: (() => void) | undefined
const started = new Promise<void>((resolve) => {
parserStarted = resolve
})
let blockParser = false
const directory = await mkdtemp(
join(tmpdir(), 'goodbuddy-knowledge-service-')
)
temporaryDirectories.push(directory)
const service = new KnowledgeService({
databasePath: join(directory, 'knowledge.sqlite'),
managedRoot: join(directory, 'managed'),
parseDocument: async (name, buffer) => {
if (blockParser) {
parserStarted?.()
await new Promise<void>((resolve) => {
releaseParser = resolve
})
}
const content = blockParser
? 'rebuilt content'
: buffer.toString('utf8')
return {
title: name,
content,
sourceFormat: 'text',
sections: [{
locator: '全文',
content
}],
warnings: []
}
}
})
await service.initialize()
services.push(service)
const sourcePath = join(directory, 'mutation-gate.txt')
await writeFile(sourcePath, 'initial content', 'utf8')
const library = service.createLibrary({
name: 'Mutation gate',
storageMode: 'reference',
graphEnabled: false
})
await service.importPaths(library.id, [sourcePath])
const document = service.snapshot(library.id).documents[0]!
const originalChunk = service.database
.listChunks(document.id)
.find((chunk) => chunk.role !== 'parent')!
blockParser = true
await writeFile(sourcePath, 'rebuilt source bytes', 'utf8')
const rebuilding = service.rebuildDocument({
knowledgeBaseId: library.id,
documentId: document.id
})
await started
const editing = service.updateChunk({
knowledgeBaseId: library.id,
documentId: document.id,
chunkId: originalChunk.id,
content: 'manual latest content'
})
releaseParser?.()
await rebuilding
await expect(editing).rejects.toThrow(
'Chunk must belong to the requested document'
)
expect(service.search(library.id, 'rebuilt')).toHaveLength(1)
expect(service.search(library.id, 'manual')).toEqual([])
})
it('aborts standalone document rebuild work and keeps child capabilities honest', async () => {
let parserSignal: AbortSignal | undefined
const directory = await mkdtemp(
@@ -1334,7 +1559,7 @@ describe('KnowledgeService', () => {
displayName: 'Original title',
status: 'ready'
})
vi.spyOn(service.database, 'upsertDocument').mockImplementationOnce(
vi.spyOn(service.database, 'publishDocument').mockImplementationOnce(
() => {
throw new Error('synthetic indexing failure')
}
File diff suppressed because it is too large Load Diff