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( ipcMain.handle(
ipcChannels.knowledgeDeleteChunk, ipcChannels.knowledgeDeleteChunk,
(event, input: unknown) => { async (event, input: unknown) => {
assertTrustedSender(event, window) assertTrustedSender(event, window)
const deleted = knowledgeService.deleteChunk( const deleted = await knowledgeService.deleteChunk(
knowledgeChunkDeleteInputSchema.parse(input) knowledgeChunkDeleteInputSchema.parse(input)
) )
if (!deleted) { if (!deleted) {
@@ -558,6 +558,65 @@ describe('extraction strategies', () => {
).rejects.toThrow('Model extraction is unavailable') ).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 () => { it('honors cancellation before and after the injected model callback', async () => {
const preCancelled = new AbortController() const preCancelled = new AbortController()
preCancelled.abort() preCancelled.abort()
+42 -15
View File
@@ -21,6 +21,7 @@ export const GRAPH_LIMITS = {
maximumWarnings: 20, maximumWarnings: 20,
maximumWarningLength: 240 maximumWarningLength: 240
} as const } as const
const maximumEvidencePerRecord = 20
export type ExtractionStrategy = 'rules' | 'model' | 'hybrid' | 'ask' export type ExtractionStrategy = 'rules' | 'model' | 'hybrid' | 'ask'
@@ -192,7 +193,7 @@ function throwIfAborted(signal?: AbortSignal): void {
function prepareChunks(chunks: readonly GraphChunk[]): GraphChunk[] { function prepareChunks(chunks: readonly GraphChunk[]): GraphChunk[] {
const ids = new Set<string>() const ids = new Set<string>()
const prepared: GraphChunk[] = [] 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) const id = truncate(chunk.id.trim(), GRAPH_LIMITS.maximumFieldLength)
if (!id || ids.has(id)) { if (!id || ids.has(id)) {
continue continue
@@ -219,6 +220,9 @@ function mergeEvidence(
const key = evidenceKey(evidence) const key = evidenceKey(evidence)
if (!merged.has(key)) { if (!merged.has(key)) {
merged.set(key, evidence) merged.set(key, evidence)
if (merged.size >= maximumEvidencePerRecord) {
break
}
} }
} }
return [...merged.values()] return [...merged.values()]
@@ -805,24 +809,39 @@ export async function extractKnowledgeGraph(
const context = createOntologyContext(options.ontology) const context = createOntologyContext(options.ontology)
throwIfAborted(options.signal) throwIfAborted(options.signal)
const prepared = prepareChunks(chunks) const prepared = prepareChunks(chunks)
const rules = if (prepared.length === 0) {
strategy === 'rules' || strategy === 'hybrid' || strategy === 'ask'
? extractGraphWithRulesInternal(prepared, options.signal, context)
: emptyGraph()
if (strategy === 'rules' || strategy === 'ask') {
return { return {
...rules, ...emptyGraph(),
strategy, strategy,
requiresModelApproval: strategy === 'ask', requiresModelApproval: strategy === 'ask',
warnings: [...context.warnings] warnings: [...context.warnings]
} }
} }
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) { if (!options.extractStructured) {
throw new Error('Model extraction is unavailable') throw new Error('Model extraction is unavailable')
} }
const output = await options.extractStructured( const output = await options.extractStructured(
createModelPrompt(prepared, context.settings), createModelPrompt(batch, context.settings),
options.signal options.signal
) )
throwIfAborted(options.signal) throwIfAborted(options.signal)
@@ -830,16 +849,24 @@ export async function extractKnowledgeGraph(
if (!modelEnvelopeSchema.safeParse(parsedOutput).success) { if (!modelEnvelopeSchema.safeParse(parsedOutput).success) {
throw new Error('模型返回的图谱结构无效') throw new Error('模型返回的图谱结构无效')
} }
const model = validateModelGraphInternal(parsedOutput, prepared, context) const model = validateModelGraphInternal(
const graph = parsedOutput,
batch,
batchContext
)
graph = mergeKnowledgeGraphsInternal(
graph,
strategy === 'hybrid' strategy === 'hybrid'
? mergeKnowledgeGraphsInternal(rules, model, context) ? mergeKnowledgeGraphsInternal(rules, model, batchContext)
: model : model,
context
)
}
return { return {
...graph, ...graph,
strategy, strategy,
requiresModelApproval: false, requiresModelApproval: strategy === 'ask',
warnings: [...context.warnings] warnings: [...context.warnings].slice(0, GRAPH_LIMITS.maximumWarnings)
} }
} }
+199 -6
View File
@@ -86,7 +86,7 @@ describe('KnowledgeDatabase', () => {
const inspection = new DatabaseSync(path) const inspection = new DatabaseSync(path)
expect( expect(
inspection.prepare('PRAGMA user_version').get() inspection.prepare('PRAGMA user_version').get()
).toEqual({ user_version: 10 }) ).toEqual({ user_version: 11 })
expect( expect(
inspection inspection
.prepare('SELECT version FROM schema_migrations ORDER BY version') .prepare('SELECT version FROM schema_migrations ORDER BY version')
@@ -101,7 +101,8 @@ describe('KnowledgeDatabase', () => {
{ version: 7 }, { version: 7 },
{ version: 8 }, { version: 8 },
{ version: 9 }, { version: 9 },
{ version: 10 } { version: 10 },
{ version: 11 }
]) ])
inspection.close() 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 { database, path } = await createDatabase()
const knowledgeBase = database.createKnowledgeBase({ const knowledgeBase = database.createKnowledgeBase({
name: 'Version one data', name: 'Version one data',
@@ -210,7 +211,7 @@ describe('KnowledgeDatabase', () => {
DROP TABLE chunk_embeddings; DROP TABLE chunk_embeddings;
DROP TABLE knowledge_tasks; DROP TABLE knowledge_tasks;
DELETE FROM schema_migrations 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; PRAGMA user_version = 1;
`) `)
downgrade.close() downgrade.close()
@@ -220,7 +221,7 @@ describe('KnowledgeDatabase', () => {
upgraded.initialize() upgraded.initialize()
const inspection = new DatabaseSync(path) const inspection = new DatabaseSync(path)
expect(inspection.prepare('PRAGMA user_version').get()).toEqual({ expect(inspection.prepare('PRAGMA user_version').get()).toEqual({
user_version: 10 user_version: 11
}) })
expect( expect(
inspection inspection
@@ -327,7 +328,7 @@ describe('KnowledgeDatabase', () => {
WHERE new.enabled = 1 AND new.role <> 'parent'; WHERE new.enabled = 1 AND new.role <> 'parent';
END; END;
INSERT INTO chunks_fts(chunks_fts) VALUES ('rebuild'); 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; PRAGMA user_version = 8;
`) `)
downgrade downgrade
@@ -700,6 +701,22 @@ describe('KnowledgeDatabase', () => {
).toEqual([]) ).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 () => { it('edits graph records and merges entities while retaining evidence and locks', async () => {
const { database } = await createDatabase() const { database } = await createDatabase()
const knowledgeBase = database.createKnowledgeBase({ const knowledgeBase = database.createKnowledgeBase({
@@ -1403,6 +1420,182 @@ describe('KnowledgeDatabase', () => {
database.discardDocumentEmbeddingReplacement(replacementId) 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 () => { it('persists embedding index jobs independently by knowledge base', async () => {
const created = await createDatabase() const created = await createDatabase()
let database = created.database let database = created.database
+452 -19
View File
@@ -80,7 +80,7 @@ import type {
VectorSearchOptions VectorSearchOptions
} from './types' } from './types'
const DATABASE_VERSION = 10 const DATABASE_VERSION = 11
const MAX_ID_LENGTH = 128 const MAX_ID_LENGTH = 128
const MAX_NAME_LENGTH = 512 const MAX_NAME_LENGTH = 512
const MAX_LOCATION_LENGTH = 8192 const MAX_LOCATION_LENGTH = 8192
@@ -118,6 +118,22 @@ export type HybridSearchResultPage = {
vectorScannedCount: number 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> type Row = Record<string, null | number | bigint | string | Uint8Array>
const taskScopes = knowledgeTaskScopeSchema.options const taskScopes = knowledgeTaskScopeSchema.options
@@ -1635,6 +1651,22 @@ export class KnowledgeDatabase {
.map(mapSource) .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 { getSource(id: string): KnowledgeSource | undefined {
const row = this.requireDatabase() const row = this.requireDatabase()
.prepare('SELECT * FROM knowledge_sources WHERE id = ?') .prepare('SELECT * FROM knowledge_sources WHERE id = ?')
@@ -1654,6 +1686,14 @@ export class KnowledgeDatabase {
upsertDocument( upsertDocument(
input: UpsertDocumentInput, input: UpsertDocumentInput,
chunks: ReplaceChunkInput[] chunks: ReplaceChunkInput[]
): Document {
return this.publishDocument(input, chunks)
}
publishDocument(
input: UpsertDocumentInput,
chunks: ReplaceChunkInput[],
options: DocumentPublicationOptions = {}
): Document { ): Document {
if (!Array.isArray(chunks) || chunks.length > MAX_CHUNKS) { if (!Array.isArray(chunks) || chunks.length > MAX_CHUNKS) {
throw new RangeError(`chunks must contain at most ${MAX_CHUNKS} items`) 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 checksum = optionalString(input.checksum, 'checksum', 512)
const metadata = jsonObject(input.metadata, 'metadata') 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() const now = new Date().toISOString()
this.transaction(database, () => { this.transaction(database, () => {
@@ -1749,7 +1841,7 @@ export class KnowledgeDatabase {
now now
) )
database database
.prepare('DELETE FROM embedding_index_state WHERE document_id = ?') .prepare('DELETE FROM graph_evidence WHERE document_id = ?')
.run(id) .run(id)
database.prepare('DELETE FROM chunks WHERE document_id = ?').run(id) database.prepare('DELETE FROM chunks WHERE document_id = ?').run(id)
const insertChunk = database.prepare( const insertChunk = database.prepare(
@@ -1794,6 +1886,29 @@ export class KnowledgeDatabase {
indexContent 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) return this.requiredDocument(id)
} }
@@ -1851,6 +1966,10 @@ export class KnowledgeDatabase {
.map(mapDocument) .map(mapDocument)
} }
listDocumentsForSnapshot(knowledgeBaseId: string): Document[] {
return this.listDocumentsForLibraryRebuild(knowledgeBaseId)
}
getDocumentChunkCounts(knowledgeBaseId: string): Map<string, number> { getDocumentChunkCounts(knowledgeBaseId: string): Map<string, number> {
const normalizedId = requiredString( const normalizedId = requiredString(
knowledgeBaseId, knowledgeBaseId,
@@ -2365,6 +2484,17 @@ export class KnowledgeDatabase {
return randomUUID() 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( appendDocumentEmbeddingBatch(
replacementId: string, replacementId: string,
documentId: string, documentId: string,
@@ -2465,26 +2595,95 @@ export class KnowledgeDatabase {
} }
return { chunkId, checksum, ...vector } return { chunkId, checksum, ...vector }
}) })
const insert = database.prepare( this.insertEmbeddingReplacementBatch(database, {
`INSERT INTO embedding_rebuild_staging replacementId: normalizedReplacementId,
(replacement_id, document_id, provider, model, chunk_id, documentId: normalizedDocumentId,
dimensions, content_checksum, vector, magnitude) provider: normalizedProvider,
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)` 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, () => { const normalizedDocumentId = requiredString(
for (const item of normalized) { documentId,
insert.run( 'documentId',
normalizedReplacementId, MAX_ID_LENGTH
normalizedDocumentId, )
normalizedProvider, const normalizedProvider = requiredString(
normalizedModel, provider,
item.chunkId, 'provider',
item.dimensions, MAX_EMBEDDING_PROVIDER_LENGTH
item.checksum, )
item.bytes, const normalizedModel = requiredString(
item.magnitude 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()) .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(`PRAGMA user_version = ${DATABASE_VERSION}`)
database.exec('COMMIT') database.exec('COMMIT')
} catch (error) { } 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<{ private normalizeChunks(chunks: ReplaceChunkInput[]): Array<{
id: string id: string
ordinal: number ordinal: number
@@ -5175,6 +5604,10 @@ export class KnowledgeDatabase {
} }
private transaction(database: DatabaseSync, operation: () => void): void { private transaction(database: DatabaseSync, operation: () => void): void {
if (database.isTransaction) {
operation()
return
}
database.exec('BEGIN IMMEDIATE') database.exec('BEGIN IMMEDIATE')
try { try {
operation() operation()
+226 -1
View File
@@ -464,6 +464,104 @@ describe('KnowledgeService', () => {
expect(service.snapshot(library.id).evidence).toEqual([]) 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 () => { it('fails hybrid reextraction when model extraction fails', async () => {
const extractStructured = vi.fn(async () => { const extractStructured = vi.fn(async () => {
throw new Error('模型未返回图谱内容') 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 () => { it('aborts standalone document rebuild work and keeps child capabilities honest', async () => {
let parserSignal: AbortSignal | undefined let parserSignal: AbortSignal | undefined
const directory = await mkdtemp( const directory = await mkdtemp(
@@ -1334,7 +1559,7 @@ describe('KnowledgeService', () => {
displayName: 'Original title', displayName: 'Original title',
status: 'ready' status: 'ready'
}) })
vi.spyOn(service.database, 'upsertDocument').mockImplementationOnce( vi.spyOn(service.database, 'publishDocument').mockImplementationOnce(
() => { () => {
throw new Error('synthetic indexing failure') throw new Error('synthetic indexing failure')
} }
+527 -100
View File
@@ -66,9 +66,13 @@ import {
extractKnowledgeGraph, extractKnowledgeGraph,
normalizeEntityAlias, normalizeEntityAlias,
type ExtractStructured, type ExtractStructured,
type GraphChunk,
type GraphExtractionResult type GraphExtractionResult
} from './graph-extractor' } from './graph-extractor'
import { KnowledgeDatabase } from './knowledge-database' import {
KnowledgeDatabase,
type PreparedEmbeddingReplacement
} from './knowledge-database'
import { import {
containsHanText, containsHanText,
contextualIndexText, contextualIndexText,
@@ -154,6 +158,20 @@ const maximumSourceBytes = 500 * 1024 * 1024
const maximumFilesPerSource = 2_000 const maximumFilesPerSource = 2_000
const maximumEmbeddingChunksPerBatch = 32 const maximumEmbeddingChunksPerBatch = 32
type PreparedDocumentPublication = {
input: Parameters<KnowledgeDatabase['publishDocument']>[0]
chunks: ReplaceChunkInput[]
graph?: GraphExtractionResult
embeddingReplacement?: PreparedEmbeddingReplacement
embeddingFailure?: {
provider: string
model: string
message: string
}
embeddingTaskId: string
graphTaskId: string
}
interface PreparedQueryEmbedding { interface PreparedQueryEmbedding {
provider?: EmbeddingProvider provider?: EmbeddingProvider
providerStorageKey?: string providerStorageKey?: string
@@ -184,11 +202,17 @@ export class KnowledgeService {
string, string,
Promise<Document> Promise<Document>
>() >()
private readonly documentMutationTails = new Map<string, Promise<void>>()
private readonly backgroundEmbeddingReindexes = new Map< private readonly backgroundEmbeddingReindexes = new Map<
string, string,
Promise<void> Promise<void>
>() >()
private readonly pendingEmbeddingReindexes = new Set<string>() private readonly pendingEmbeddingReindexes = new Set<string>()
private readonly backgroundGraphReindexes = new Map<
string,
Promise<void>
>()
private readonly pendingGraphReindexes = new Set<string>()
private readonly taskControllers = new Map<string, AbortController>() private readonly taskControllers = new Map<string, AbortController>()
private readonly taskOperations = new Map<string, Promise<unknown>>() private readonly taskOperations = new Map<string, Promise<unknown>>()
private readonly sourceSyncTaskIds = new Map<string, string>() private readonly sourceSyncTaskIds = new Map<string, string>()
@@ -229,7 +253,7 @@ export class KnowledgeService {
await mkdir(this.managedRoot, { recursive: true }) await mkdir(this.managedRoot, { recursive: true })
this.database.initialize() this.database.initialize()
for (const library of this.database.listKnowledgeBases()) { for (const library of this.database.listKnowledgeBases()) {
for (const source of this.database.listSources(library.id)) { for (const source of this.database.listSourcesForSnapshot(library.id)) {
if ( if (
library.storageMode === 'reference' && library.storageMode === 'reference' &&
source.type !== 'url' && source.type !== 'url' &&
@@ -274,7 +298,9 @@ export class KnowledgeService {
await Promise.allSettled([ await Promise.allSettled([
...this.activeSyncs.values(), ...this.activeSyncs.values(),
...this.activeDocumentRebuilds.values(), ...this.activeDocumentRebuilds.values(),
...this.documentMutationTails.values(),
...this.backgroundEmbeddingReindexes.values(), ...this.backgroundEmbeddingReindexes.values(),
...this.backgroundGraphReindexes.values(),
...this.taskOperations.values(), ...this.taskOperations.values(),
...embeddingCompletions ...embeddingCompletions
]) ])
@@ -282,6 +308,7 @@ export class KnowledgeService {
this.taskOperations.clear() this.taskOperations.clear()
this.sourceSyncTaskIds.clear() this.sourceSyncTaskIds.clear()
this.libraryRebuildControllers.clear() this.libraryRebuildControllers.clear()
this.documentMutationTails.clear()
this.database.close() this.database.close()
} }
@@ -543,6 +570,343 @@ export class KnowledgeService {
await Promise.allSettled(operations) await Promise.allSettled(operations)
} }
private async withDocumentMutation<T>(
documentId: string,
operation: () => Promise<T> | T,
signal?: AbortSignal
): Promise<T> {
const previous = this.documentMutationTails.get(documentId)
let release: (() => void) | undefined
const tail = new Promise<void>((resolve) => {
release = resolve
})
this.documentMutationTails.set(documentId, tail)
try {
await previous
signal?.throwIfAborted()
this.lifecycleController.signal.throwIfAborted()
return await operation()
} finally {
release?.()
if (this.documentMutationTails.get(documentId) === tail) {
this.documentMutationTails.delete(documentId)
}
}
}
private async prepareDocumentEmbeddings(
documentId: string,
chunks: readonly ReplaceChunkInput[],
signal?: AbortSignal
): Promise<PreparedEmbeddingReplacement | undefined> {
const provider = this.embeddingProvider
if (!provider) {
return undefined
}
const effectiveSignal = signal
? AbortSignal.any([signal, this.lifecycleController.signal])
: this.lifecycleController.signal
const providerStorageKey = embeddingStorageProvider(provider)
const indexableChunks = chunks.filter(
(chunk) =>
(chunk.enabled ?? true) &&
(chunk.role ?? 'standalone') !== 'parent'
)
const replacementId =
this.database.beginPreparedDocumentEmbeddingReplacement(
documentId,
providerStorageKey,
provider.model
)
try {
let expectedDimensions: number | undefined
for (
let offset = 0;
offset < indexableChunks.length;
offset += this.embeddingBatchSize
) {
effectiveSignal.throwIfAborted()
const batch = indexableChunks.slice(
offset,
offset + this.embeddingBatchSize
)
const contents = batch.map((chunk) =>
contextualIndexText(
chunk.content,
chunk.metadata?.contextPrefix
)
)
const vectors = await provider.embed(contents, effectiveSignal)
effectiveSignal.throwIfAborted()
if (vectors.length !== batch.length) {
throw new Error('Embedding provider returned an invalid result count')
}
const embeddings = batch.map((chunk, index) => {
const vector = vectors[index]
if (!vector) {
throw new Error('Embedding provider returned an incomplete batch')
}
if (expectedDimensions === undefined) {
expectedDimensions = vector.length
} else if (vector.length !== expectedDimensions) {
throw new Error(
'Embedding provider returned inconsistent dimensions'
)
}
return {
chunkId: chunk.id!,
contentChecksum: createHash('sha256')
.update(contents[index]!)
.digest('hex'),
vector
}
})
this.database.appendPreparedDocumentEmbeddingBatch(
replacementId,
documentId,
providerStorageKey,
provider.model,
embeddings
)
}
effectiveSignal.throwIfAborted()
if (this.embeddingProvider !== provider) {
throw new Error('向量模型配置已变化')
}
return {
replacementId,
provider: providerStorageKey,
model: provider.model
}
} catch (error) {
this.database.discardDocumentEmbeddingReplacement(replacementId)
throw error
}
}
private async prepareDocumentPublication(
input: PreparedDocumentPublication['input'],
parsed: ParsedDocument,
library: KnowledgeBase,
signal?: AbortSignal,
parentTaskId?: string
): Promise<PreparedDocumentPublication> {
signal?.throwIfAborted()
const documentId = input.id ?? randomUUID()
const chunks = this.createDocumentChunks(parsed, library)
let embeddingReplacement: PreparedEmbeddingReplacement | undefined
let embeddingFailure: PreparedDocumentPublication['embeddingFailure']
const embeddingProvider = this.embeddingProvider
const embeddingTask = this.createKnowledgeTask({
libraryId: library.id,
parentTaskId,
sourceId: input.sourceId,
documentName: input.title,
scope: 'document',
kind: 'embedding'
})
if (!embeddingProvider) {
this.updateKnowledgeTask(embeddingTask.id, {
status: 'skipped',
message: '未启用向量化'
})
} else {
this.updateKnowledgeTask(embeddingTask.id, {
status: 'running',
stage: 'embedding',
progress: 5,
message: '正在生成候选向量索引'
})
}
try {
embeddingReplacement = await this.prepareDocumentEmbeddings(
documentId,
chunks,
signal
)
if (embeddingReplacement) {
this.updateKnowledgeTask(embeddingTask.id, {
progress: 90,
message: `已生成 ${chunks.length} 个候选分块向量`
})
}
} catch (error) {
if (signal?.aborted) {
this.database.cancelKnowledgeTask(embeddingTask.id, '文档处理已取消')
throw signal.reason
}
const safeError = classifyEmbeddingError(error)
if (embeddingProvider) {
this.failKnowledgeTask(embeddingTask.id, safeError)
embeddingFailure = {
provider: embeddingStorageProvider(embeddingProvider),
model: embeddingProvider.model,
message: safeError.message
}
} else {
this.updateKnowledgeTask(embeddingTask.id, {
status: 'skipped',
message: '未启用向量化'
})
}
}
let graph: GraphExtractionResult | undefined
const graphTask = this.createKnowledgeTask({
libraryId: library.id,
parentTaskId,
sourceId: input.sourceId,
documentName: input.title,
scope: 'document',
kind: 'graph'
})
try {
signal?.throwIfAborted()
if (library.graphEnabled && library.graphStrategy !== 'ask') {
this.updateKnowledgeTask(graphTask.id, {
status: 'running',
stage: 'graph',
progress: 10,
message: '正在抽取候选知识图谱'
})
graph = await this.extractGraphChunks(library, chunks, signal)
this.updateKnowledgeTask(graphTask.id, {
progress: 90,
message: `已生成 ${graph.entities.length} 个候选实体、${graph.relations.length} 条候选关系`
})
} else {
this.updateKnowledgeTask(graphTask.id, {
status: 'skipped',
message: library.graphEnabled
? '按需询问策略不自动抽取'
: '知识图谱未启用'
})
}
signal?.throwIfAborted()
} catch (error) {
if (embeddingReplacement) {
this.database.discardDocumentEmbeddingReplacement(
embeddingReplacement.replacementId
)
if (signal?.aborted) {
this.database.cancelKnowledgeTask(
embeddingTask.id,
'文档处理已取消'
)
} else {
this.updateKnowledgeTask(embeddingTask.id, {
status: 'skipped',
message: '因文档发布前处理失败而未保存向量'
})
}
} else if (
!embeddingFailure &&
!['skipped', 'failed'].includes(
this.database.getKnowledgeTask(embeddingTask.id)?.status ?? ''
)
) {
if (signal?.aborted) {
this.database.cancelKnowledgeTask(
embeddingTask.id,
'文档处理已取消'
)
} else {
this.updateKnowledgeTask(embeddingTask.id, {
status: 'skipped',
message: '因文档发布前处理失败而未保存向量'
})
}
}
if (signal?.aborted) {
this.database.cancelKnowledgeTask(graphTask.id, '文档处理已取消')
} else {
this.failKnowledgeTask(graphTask.id, error)
}
throw error
}
return {
input: { ...input, id: documentId },
chunks,
graph,
embeddingReplacement,
embeddingFailure,
embeddingTaskId: embeddingTask.id,
graphTaskId: graphTask.id
}
}
private publishPreparedDocument(
library: KnowledgeBase,
prepared: PreparedDocumentPublication
): Document {
try {
const document = this.database.publishDocument(
prepared.input,
prepared.chunks,
{
embeddingReplacement: prepared.embeddingReplacement,
embeddingError: prepared.embeddingFailure,
afterChunksInserted: prepared.graph
? (document) => {
this.storeExtractedGraph(library, document, prepared.graph!)
}
: undefined
}
)
if (prepared.embeddingReplacement) {
this.updateKnowledgeTask(prepared.embeddingTaskId, {
documentId: document.id,
documentName: document.title,
status: 'succeeded',
message: `已发布 ${prepared.chunks.length} 个分块向量`
})
} else {
this.updateKnowledgeTask(prepared.embeddingTaskId, {
documentId: document.id,
documentName: document.title
})
}
if (prepared.graph) {
this.updateKnowledgeTask(prepared.graphTaskId, {
documentId: document.id,
documentName: document.title,
status: 'succeeded',
message: `已发布 ${prepared.graph.entities.length} 个实体、${prepared.graph.relations.length} 条关系`
})
} else {
this.updateKnowledgeTask(prepared.graphTaskId, {
documentId: document.id,
documentName: document.title
})
}
return document
} catch (error) {
const currentEmbeddingTask =
this.database.getKnowledgeTask(prepared.embeddingTaskId)
if (prepared.embeddingReplacement) {
try {
this.database.discardDocumentEmbeddingReplacement(
prepared.embeddingReplacement.replacementId
)
} catch {
// The enclosing database transaction may already have rolled back.
}
this.failKnowledgeTask(prepared.embeddingTaskId, error)
} else if (
currentEmbeddingTask?.status === 'queued' ||
currentEmbeddingTask?.status === 'running'
) {
this.updateKnowledgeTask(prepared.embeddingTaskId, {
status: 'skipped',
message: '因文档发布失败而未保存向量'
})
}
if (prepared.graph) {
this.failKnowledgeTask(prepared.graphTaskId, error)
}
throw error
}
}
private reconcileEmbeddingTask( private reconcileEmbeddingTask(
libraryId: string, libraryId: string,
job: EmbeddingIndexJob | null job: EmbeddingIndexJob | null
@@ -627,7 +991,7 @@ export class KnowledgeService {
if (!library) { if (!library) {
return false return false
} }
for (const source of this.database.listSources(id)) { for (const source of this.database.listSourcesForSnapshot(id)) {
this.stopWatcher(source.id) this.stopWatcher(source.id)
} }
await this.cancelTasks( await this.cancelTasks(
@@ -681,7 +1045,8 @@ export class KnowledgeService {
embeddingCoordinator.status().job embeddingCoordinator.status().job
) )
} }
const libraryDocuments = this.database.listDocuments(libraryId) const libraryDocuments =
this.database.listDocumentsForSnapshot(libraryId)
const documentCountsBySource = new Map<string, number>() const documentCountsBySource = new Map<string, number>()
for (const document of libraryDocuments) { for (const document of libraryDocuments) {
documentCountsBySource.set( documentCountsBySource.set(
@@ -689,7 +1054,9 @@ export class KnowledgeService {
(documentCountsBySource.get(document.sourceId) ?? 0) + 1 (documentCountsBySource.get(document.sourceId) ?? 0) + 1
) )
} }
const sources = this.database.listSources(libraryId).map((source) => ({ const sources = this.database
.listSourcesForSnapshot(libraryId)
.map((source) => ({
...source, ...source,
documentCount: documentCountsBySource.get(source.id) ?? 0, documentCount: documentCountsBySource.get(source.id) ?? 0,
progress: progress:
@@ -1219,8 +1586,9 @@ export class KnowledgeService {
async updateChunk( async updateChunk(
rawInput: KnowledgeChunkUpdateInput rawInput: KnowledgeChunkUpdateInput
) { ): Promise<ReturnType<KnowledgeDatabase['updateChunk']>> {
const input = knowledgeChunkUpdateInputSchema.parse(rawInput) const input = knowledgeChunkUpdateInputSchema.parse(rawInput)
return this.withDocumentMutation(input.documentId, async () => {
const current = this.database.getChunkForReference( const current = this.database.getChunkForReference(
input.knowledgeBaseId, input.knowledgeBaseId,
input.documentId, input.documentId,
@@ -1237,29 +1605,32 @@ export class KnowledgeService {
) )
) { ) {
this.scheduleEmbeddingReindex(input.documentId) this.scheduleEmbeddingReindex(input.documentId)
const library = this.requireLibrary(input.knowledgeBaseId)
if (library.graphEnabled && library.graphStrategy !== 'ask') {
this.scheduleGraphReindex(input.documentId)
} }
if (
current &&
input.content !== undefined &&
input.content !== current.content
) {
this.database.pruneUnreferencedGeneratedGraph(
input.knowledgeBaseId
)
} }
this.database.pruneUnreferencedGeneratedGraph(input.knowledgeBaseId)
return chunk return chunk
})
} }
deleteChunk(rawInput: KnowledgeChunkDeleteInput): boolean { async deleteChunk(rawInput: KnowledgeChunkDeleteInput): Promise<boolean> {
const input = knowledgeChunkDeleteInputSchema.parse(rawInput) const input = knowledgeChunkDeleteInputSchema.parse(rawInput)
return this.withDocumentMutation(input.documentId, async () => {
const deleted = this.database.deleteChunk(input) const deleted = this.database.deleteChunk(input)
if (deleted) { if (deleted) {
this.scheduleEmbeddingReindex(input.documentId) this.scheduleEmbeddingReindex(input.documentId)
const library = this.requireLibrary(input.knowledgeBaseId)
if (library.graphEnabled && library.graphStrategy !== 'ask') {
this.scheduleGraphReindex(input.documentId)
}
this.database.pruneUnreferencedGeneratedGraph( this.database.pruneUnreferencedGeneratedGraph(
input.knowledgeBaseId input.knowledgeBaseId
) )
} }
return deleted return deleted
})
} }
getReferenceContext(rawInput: KnowledgeReferenceContextInput) { getReferenceContext(rawInput: KnowledgeReferenceContextInput) {
@@ -1552,7 +1923,9 @@ export class KnowledgeService {
signal: AbortSignal, signal: AbortSignal,
sourceId?: string, sourceId?: string,
graphStrategy?: Exclude<GraphStrategy, 'ask'>, graphStrategy?: Exclude<GraphStrategy, 'ask'>,
parentTaskId?: string parentTaskId?: string,
mutationDocumentId?: string,
documentMutationHeld = false
): Promise<void> { ): Promise<void> {
const library = this.requireLibrary(knowledgeBaseId) const library = this.requireLibrary(knowledgeBaseId)
const parsingTask = this.createKnowledgeTask({ const parsingTask = this.createKnowledgeTask({
@@ -1618,14 +1991,17 @@ export class KnowledgeService {
knowledgeBaseId, knowledgeBaseId,
graphStrategy graphStrategy
) )
const previousDocument = this.database const previousDocument = mutationDocumentId
? this.database.getDocument(mutationDocumentId)
: this.database
.listDocumentsForSource(source.id) .listDocumentsForSource(source.id)
.find((document) => document.externalId === result.url) .find((document) => document.externalId === result.url)
if (previousDocument) { const documentId =
this.database.removeEvidenceForDocument(previousDocument.id) mutationDocumentId ?? previousDocument?.id ?? randomUUID()
} const publish = async (): Promise<Document> => {
const document = this.database.upsertDocument( const prepared = await this.prepareDocumentPublication(
{ {
id: documentId,
knowledgeBaseId, knowledgeBaseId,
sourceId: source.id, sourceId: source.id,
externalId: result.url, externalId: result.url,
@@ -1640,37 +2016,27 @@ export class KnowledgeService {
size: Buffer.byteLength(result.document.content) size: Buffer.byteLength(result.document.content)
} }
}, },
this.createDocumentChunks(
result.document, result.document,
effectiveLibrary effectiveLibrary,
effectiveSignal,
parsingTask.id
) )
return this.publishPreparedDocument(effectiveLibrary, prepared)
}
const document = documentMutationHeld
? await publish()
: await this.withDocumentMutation(
documentId,
publish,
effectiveSignal
) )
this.database.pruneUnreferencedGeneratedGraph(library.id)
this.updateKnowledgeTask(parsingTask.id, { this.updateKnowledgeTask(parsingTask.id, {
documentId: document.id, documentId: document.id,
documentName: document.title, documentName: document.title,
stage: 'embedding', stage: 'finalizing',
progress: 80,
message: '网页解析完成,正在生成索引'
})
await this.indexDocumentEmbeddings(
document,
undefined,
parsingTask.id,
effectiveSignal
)
effectiveSignal.throwIfAborted()
this.updateKnowledgeTask(parsingTask.id, {
stage: 'graph',
progress: 92, progress: 92,
message: '正在处理知识图谱' message: '网页解析与索引完成'
}) })
await this.extractGraph(
effectiveLibrary,
document,
parsingTask.id,
effectiveSignal
)
effectiveSignal.throwIfAborted() effectiveSignal.throwIfAborted()
source = this.database.upsertSource({ source = this.database.upsertSource({
...source, ...source,
@@ -1764,7 +2130,6 @@ export class KnowledgeService {
retryOfTaskId?: string retryOfTaskId?: string
): Promise<Document> { ): Promise<Document> {
const input = knowledgeDocumentRebuildInputSchema.parse(rawInput) const input = knowledgeDocumentRebuildInputSchema.parse(rawInput)
if (!parentTaskId) {
const existing = this.activeDocumentRebuilds.get(input.documentId) const existing = this.activeDocumentRebuilds.get(input.documentId)
if (existing) { if (existing) {
return existing return existing
@@ -1772,7 +2137,7 @@ export class KnowledgeService {
const operation = this.performRebuildDocument( const operation = this.performRebuildDocument(
input, input,
signal, signal,
undefined, parentTaskId,
retryOfTaskId retryOfTaskId
).finally(() => { ).finally(() => {
if (this.activeDocumentRebuilds.get(input.documentId) === operation) { if (this.activeDocumentRebuilds.get(input.documentId) === operation) {
@@ -1782,15 +2147,27 @@ export class KnowledgeService {
this.activeDocumentRebuilds.set(input.documentId, operation) this.activeDocumentRebuilds.set(input.documentId, operation)
return operation return operation
} }
return this.performRebuildDocument(
private async performRebuildDocument(
input: KnowledgeDocumentRebuildInput,
signal?: AbortSignal,
parentTaskId?: string,
retryOfTaskId?: string
): Promise<Document> {
return this.withDocumentMutation(
input.documentId,
() =>
this.performRebuildDocumentMutation(
input, input,
signal, signal,
parentTaskId, parentTaskId,
retryOfTaskId retryOfTaskId
),
signal
) )
} }
private async performRebuildDocument( private async performRebuildDocumentMutation(
input: KnowledgeDocumentRebuildInput, input: KnowledgeDocumentRebuildInput,
signal?: AbortSignal, signal?: AbortSignal,
parentTaskId?: string, parentTaskId?: string,
@@ -1849,7 +2226,9 @@ export class KnowledgeService {
effectiveSignal, effectiveSignal,
source.id, source.id,
undefined, undefined,
task.id task.id,
document.id,
true
) )
effectiveSignal.throwIfAborted() effectiveSignal.throwIfAborted()
const rebuilt = this.database.getDocument(document.id) ?? document const rebuilt = this.database.getDocument(document.id) ?? document
@@ -1898,8 +2277,7 @@ export class KnowledgeService {
progress: 45, progress: 45,
message: '正在切分文档' message: '正在切分文档'
}) })
this.database.removeEvidenceForDocument(document.id) const prepared = await this.prepareDocumentPublication(
const rebuilt = this.database.upsertDocument(
{ {
...document, ...document,
checksum: createHash('sha256').update(buffer).digest('hex'), checksum: createHash('sha256').update(buffer).digest('hex'),
@@ -1909,28 +2287,18 @@ export class KnowledgeService {
size: fileStat.size size: fileStat.size
} }
}, },
this.createDocumentChunks(parsed, library) parsed,
library,
effectiveSignal,
task.id
) )
this.database.pruneUnreferencedGeneratedGraph(library.id)
this.updateKnowledgeTask(task.id, { this.updateKnowledgeTask(task.id, {
stage: 'embedding', stage: 'finalizing',
progress: 65, progress: 92,
message: '正在生成向量索引' message: '正在发布重建结果'
}) })
await this.indexDocumentEmbeddings(
rebuilt,
undefined,
task.id,
effectiveSignal
)
effectiveSignal.throwIfAborted()
this.updateKnowledgeTask(task.id, {
stage: 'graph',
progress: 88,
message: '正在处理知识图谱'
})
await this.extractGraph(library, rebuilt, task.id, effectiveSignal)
effectiveSignal.throwIfAborted() effectiveSignal.throwIfAborted()
const rebuilt = this.publishPreparedDocument(library, prepared)
this.updateKnowledgeTask(task.id, { this.updateKnowledgeTask(task.id, {
status: 'succeeded', status: 'succeeded',
stage: 'finalizing', stage: 'finalizing',
@@ -2238,6 +2606,9 @@ export class KnowledgeService {
}) })
tasks.push(task) tasks.push(task)
try { try {
await this.withDocumentMutation(
document.id,
async () => {
this.updateKnowledgeTask(task.id, { this.updateKnowledgeTask(task.id, {
status: 'running', status: 'running',
progress: 10, progress: 10,
@@ -2261,6 +2632,9 @@ export class KnowledgeService {
status: 'succeeded', status: 'succeeded',
message: `已抽取 ${result.entities.length} 个实体、${result.relations.length} 条关系` message: `已抽取 ${result.entities.length} 个实体、${result.relations.length} 条关系`
}) })
},
effectiveSignal
)
this.updateKnowledgeTask(parentTask.id, { this.updateKnowledgeTask(parentTask.id, {
progress: ((index + 1) / Math.max(documents.length, 1)) * 100, progress: ((index + 1) / Math.max(documents.length, 1)) * 100,
completedItems: index + 1, completedItems: index + 1,
@@ -2481,7 +2855,14 @@ export class KnowledgeService {
for (const document of existing) { for (const document of existing) {
signal.throwIfAborted() signal.throwIfAborted()
if (!currentExternalIds.has(document.externalId)) { if (!currentExternalIds.has(document.externalId)) {
this.database.removeDocument(document.id) await this.withDocumentMutation(
document.id,
() => {
this.cancelScheduledEmbeddingReindex(document.id)
return this.database.removeDocument(document.id)
},
signal
)
} }
} }
if (existing.some( if (existing.some(
@@ -2551,8 +2932,13 @@ export class KnowledgeService {
library.id, library.id,
graphStrategy graphStrategy
) )
const document = this.database.upsertDocument( const documentId = previous?.id ?? randomUUID()
const document = await this.withDocumentMutation(
documentId,
async () => {
const prepared = await this.prepareDocumentPublication(
{ {
id: documentId,
knowledgeBaseId: library.id, knowledgeBaseId: library.id,
sourceId: source.id, sourceId: source.id,
externalId: file.relativePath, externalId: file.relativePath,
@@ -2568,35 +2954,22 @@ export class KnowledgeService {
size: file.size size: file.size
} }
}, },
this.createDocumentChunks(parsed, effectiveLibrary) parsed,
effectiveLibrary,
signal,
parsingTask.id
)
return this.publishPreparedDocument(effectiveLibrary, prepared)
},
signal
) )
this.database.pruneUnreferencedGeneratedGraph(library.id)
this.updateKnowledgeTask(parsingTask.id, { this.updateKnowledgeTask(parsingTask.id, {
documentId: document.id, documentId: document.id,
documentName: document.title, documentName: document.title,
stage: 'embedding', stage: 'finalizing',
progress: 80,
message: '文档解析完成,正在生成索引'
})
this.database.removeEvidenceForDocument(document.id)
await this.indexDocumentEmbeddings(
document,
undefined,
parsingTask.id,
signal
)
signal.throwIfAborted()
this.updateKnowledgeTask(parsingTask.id, {
stage: 'graph',
progress: 92, progress: 92,
message: '正在处理知识图谱' message: '文档解析与索引完成'
}) })
await this.extractGraph(
effectiveLibrary,
document,
parsingTask.id,
signal
)
signal.throwIfAborted() signal.throwIfAborted()
this.updateKnowledgeTask(parsingTask.id, { this.updateKnowledgeTask(parsingTask.id, {
status: 'succeeded', status: 'succeeded',
@@ -2693,6 +3066,45 @@ export class KnowledgeService {
} }
} }
private scheduleGraphReindex(documentId: string): void {
if (this.backgroundGraphReindexes.has(documentId)) {
this.pendingGraphReindexes.add(documentId)
return
}
const operation = (async () => {
do {
this.pendingGraphReindexes.delete(documentId)
const document = this.database.getDocument(documentId)
if (!document || this.lifecycleController.signal.aborted) {
return
}
const library = this.database.getKnowledgeBase(
document.knowledgeBaseId
)
if (
!library?.graphEnabled ||
library.graphStrategy === 'ask'
) {
return
}
await this.extractGraphMutation(
library,
document,
undefined,
this.lifecycleController.signal
)
} while (this.pendingGraphReindexes.delete(documentId))
})()
.catch(() => undefined)
.finally(() => {
this.pendingGraphReindexes.delete(documentId)
if (this.backgroundGraphReindexes.get(documentId) === operation) {
this.backgroundGraphReindexes.delete(documentId)
}
})
this.backgroundGraphReindexes.set(documentId, operation)
}
private async indexDocumentEmbeddings( private async indexDocumentEmbeddings(
document: Document, document: Document,
requestedProvider?: EmbeddingProvider, requestedProvider?: EmbeddingProvider,
@@ -2849,7 +3261,7 @@ export class KnowledgeService {
} }
} }
private async extractGraph( private async extractGraphMutation(
library: KnowledgeBase, library: KnowledgeBase,
document: Document, document: Document,
operationTaskId?: string, operationTaskId?: string,
@@ -2927,12 +3339,27 @@ export class KnowledgeService {
document: Document, document: Document,
signal?: AbortSignal signal?: AbortSignal
): Promise<GraphExtractionResult> { ): Promise<GraphExtractionResult> {
const chunks = this.database return this.extractGraphChunks(
.listChunks(document.id, 10_000) library,
.filter((chunk) => chunk.enabled && chunk.role !== 'parent') this.database.listChunks(document.id, 10_000),
signal
)
}
private extractGraphChunks(
library: KnowledgeBase,
chunks: readonly Pick<ReplaceChunkInput, 'id' | 'content' | 'enabled' | 'role'>[],
signal?: AbortSignal
): Promise<GraphExtractionResult> {
return extractKnowledgeGraph( return extractKnowledgeGraph(
chunks.map((chunk) => ({ chunks
id: chunk.id, .filter(
(chunk) =>
(chunk.enabled ?? true) &&
(chunk.role ?? 'standalone') !== 'parent'
)
.map((chunk): GraphChunk => ({
id: chunk.id!,
content: chunk.content content: chunk.content
})), })),
{ {