diff --git a/src/main/ipc.ts b/src/main/ipc.ts index 23bfbc6..82e60ca 100644 --- a/src/main/ipc.ts +++ b/src/main/ipc.ts @@ -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) { diff --git a/src/main/knowledge/graph-extractor.test.ts b/src/main/knowledge/graph-extractor.test.ts index 284af3d..540dae4 100644 --- a/src/main/knowledge/graph-extractor.test.ts +++ b/src/main/knowledge/graph-extractor.test.ts @@ -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('')[1]! + .split('')[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('')[1]! + .split('')[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() diff --git a/src/main/knowledge/graph-extractor.ts b/src/main/knowledge/graph-extractor.ts index eac1c83..917bd2c 100644 --- a/src/main/knowledge/graph-extractor.ts +++ b/src/main/knowledge/graph-extractor.ts @@ -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() 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) } } diff --git a/src/main/knowledge/knowledge-database.test.ts b/src/main/knowledge/knowledge-database.test.ts index 9cb927b..8213334 100644 --- a/src/main/knowledge/knowledge-database.test.ts +++ b/src/main/knowledge/knowledge-database.test.ts @@ -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 diff --git a/src/main/knowledge/knowledge-database.ts b/src/main/knowledge/knowledge-database.ts index 05f26b2..50190db 100644 --- a/src/main/knowledge/knowledge-database.ts +++ b/src/main/knowledge/knowledge-database.ts @@ -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 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 { 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() + 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 + ): 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() diff --git a/src/main/knowledge/knowledge-service.test.ts b/src/main/knowledge/knowledge-service.test.ts index 9d8d827..84bf411 100644 --- a/src/main/knowledge/knowledge-service.test.ts +++ b/src/main/knowledge/knowledge-service.test.ts @@ -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((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((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') } diff --git a/src/main/knowledge/knowledge-service.ts b/src/main/knowledge/knowledge-service.ts index d0e2c1c..0c3f9dc 100644 --- a/src/main/knowledge/knowledge-service.ts +++ b/src/main/knowledge/knowledge-service.ts @@ -66,9 +66,13 @@ import { extractKnowledgeGraph, normalizeEntityAlias, type ExtractStructured, + type GraphChunk, type GraphExtractionResult } from './graph-extractor' -import { KnowledgeDatabase } from './knowledge-database' +import { + KnowledgeDatabase, + type PreparedEmbeddingReplacement +} from './knowledge-database' import { containsHanText, contextualIndexText, @@ -154,6 +158,20 @@ const maximumSourceBytes = 500 * 1024 * 1024 const maximumFilesPerSource = 2_000 const maximumEmbeddingChunksPerBatch = 32 +type PreparedDocumentPublication = { + input: Parameters[0] + chunks: ReplaceChunkInput[] + graph?: GraphExtractionResult + embeddingReplacement?: PreparedEmbeddingReplacement + embeddingFailure?: { + provider: string + model: string + message: string + } + embeddingTaskId: string + graphTaskId: string +} + interface PreparedQueryEmbedding { provider?: EmbeddingProvider providerStorageKey?: string @@ -184,11 +202,17 @@ export class KnowledgeService { string, Promise >() + private readonly documentMutationTails = new Map>() private readonly backgroundEmbeddingReindexes = new Map< string, Promise >() private readonly pendingEmbeddingReindexes = new Set() + private readonly backgroundGraphReindexes = new Map< + string, + Promise + >() + private readonly pendingGraphReindexes = new Set() private readonly taskControllers = new Map() private readonly taskOperations = new Map>() private readonly sourceSyncTaskIds = new Map() @@ -229,7 +253,7 @@ export class KnowledgeService { await mkdir(this.managedRoot, { recursive: true }) this.database.initialize() 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 ( library.storageMode === 'reference' && source.type !== 'url' && @@ -274,7 +298,9 @@ export class KnowledgeService { await Promise.allSettled([ ...this.activeSyncs.values(), ...this.activeDocumentRebuilds.values(), + ...this.documentMutationTails.values(), ...this.backgroundEmbeddingReindexes.values(), + ...this.backgroundGraphReindexes.values(), ...this.taskOperations.values(), ...embeddingCompletions ]) @@ -282,6 +308,7 @@ export class KnowledgeService { this.taskOperations.clear() this.sourceSyncTaskIds.clear() this.libraryRebuildControllers.clear() + this.documentMutationTails.clear() this.database.close() } @@ -543,6 +570,343 @@ export class KnowledgeService { await Promise.allSettled(operations) } + private async withDocumentMutation( + documentId: string, + operation: () => Promise | T, + signal?: AbortSignal + ): Promise { + const previous = this.documentMutationTails.get(documentId) + let release: (() => void) | undefined + const tail = new Promise((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 { + 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 { + 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( libraryId: string, job: EmbeddingIndexJob | null @@ -627,7 +991,7 @@ export class KnowledgeService { if (!library) { return false } - for (const source of this.database.listSources(id)) { + for (const source of this.database.listSourcesForSnapshot(id)) { this.stopWatcher(source.id) } await this.cancelTasks( @@ -681,7 +1045,8 @@ export class KnowledgeService { embeddingCoordinator.status().job ) } - const libraryDocuments = this.database.listDocuments(libraryId) + const libraryDocuments = + this.database.listDocumentsForSnapshot(libraryId) const documentCountsBySource = new Map() for (const document of libraryDocuments) { documentCountsBySource.set( @@ -689,7 +1054,9 @@ export class KnowledgeService { (documentCountsBySource.get(document.sourceId) ?? 0) + 1 ) } - const sources = this.database.listSources(libraryId).map((source) => ({ + const sources = this.database + .listSourcesForSnapshot(libraryId) + .map((source) => ({ ...source, documentCount: documentCountsBySource.get(source.id) ?? 0, progress: @@ -702,7 +1069,7 @@ export class KnowledgeService { typeof source.metadata.lastSyncedAt === 'string' ? source.metadata.lastSyncedAt : undefined - })) + })) const chunkCounts = this.database.getDocumentChunkCounts(libraryId) const documents = libraryDocuments.map((document) => { const status = @@ -1219,47 +1586,51 @@ export class KnowledgeService { async updateChunk( rawInput: KnowledgeChunkUpdateInput - ) { + ): Promise> { const input = knowledgeChunkUpdateInputSchema.parse(rawInput) - const current = this.database.getChunkForReference( - input.knowledgeBaseId, - input.documentId, - input.chunkId - )?.chunk - const chunk = this.database.updateChunk(input) - if ( - current && - ( - (input.content !== undefined && - input.content !== current.content) || - (input.enabled !== undefined && - input.enabled !== current.enabled) - ) - ) { - this.scheduleEmbeddingReindex(input.documentId) - } - if ( - current && - input.content !== undefined && - input.content !== current.content - ) { - this.database.pruneUnreferencedGeneratedGraph( - input.knowledgeBaseId - ) - } - return chunk + return this.withDocumentMutation(input.documentId, async () => { + const current = this.database.getChunkForReference( + input.knowledgeBaseId, + input.documentId, + input.chunkId + )?.chunk + const chunk = this.database.updateChunk(input) + if ( + current && + ( + (input.content !== undefined && + input.content !== current.content) || + (input.enabled !== undefined && + input.enabled !== current.enabled) + ) + ) { + this.scheduleEmbeddingReindex(input.documentId) + const library = this.requireLibrary(input.knowledgeBaseId) + if (library.graphEnabled && library.graphStrategy !== 'ask') { + this.scheduleGraphReindex(input.documentId) + } + } + this.database.pruneUnreferencedGeneratedGraph(input.knowledgeBaseId) + return chunk + }) } - deleteChunk(rawInput: KnowledgeChunkDeleteInput): boolean { + async deleteChunk(rawInput: KnowledgeChunkDeleteInput): Promise { const input = knowledgeChunkDeleteInputSchema.parse(rawInput) - const deleted = this.database.deleteChunk(input) - if (deleted) { - this.scheduleEmbeddingReindex(input.documentId) - this.database.pruneUnreferencedGeneratedGraph( - input.knowledgeBaseId - ) - } - return deleted + return this.withDocumentMutation(input.documentId, async () => { + const deleted = this.database.deleteChunk(input) + if (deleted) { + this.scheduleEmbeddingReindex(input.documentId) + const library = this.requireLibrary(input.knowledgeBaseId) + if (library.graphEnabled && library.graphStrategy !== 'ask') { + this.scheduleGraphReindex(input.documentId) + } + this.database.pruneUnreferencedGeneratedGraph( + input.knowledgeBaseId + ) + } + return deleted + }) } getReferenceContext(rawInput: KnowledgeReferenceContextInput) { @@ -1552,7 +1923,9 @@ export class KnowledgeService { signal: AbortSignal, sourceId?: string, graphStrategy?: Exclude, - parentTaskId?: string + parentTaskId?: string, + mutationDocumentId?: string, + documentMutationHeld = false ): Promise { const library = this.requireLibrary(knowledgeBaseId) const parsingTask = this.createKnowledgeTask({ @@ -1618,59 +1991,52 @@ export class KnowledgeService { knowledgeBaseId, graphStrategy ) - const previousDocument = this.database - .listDocumentsForSource(source.id) - .find((document) => document.externalId === result.url) - if (previousDocument) { - this.database.removeEvidenceForDocument(previousDocument.id) + const previousDocument = mutationDocumentId + ? this.database.getDocument(mutationDocumentId) + : this.database + .listDocumentsForSource(source.id) + .find((document) => document.externalId === result.url) + const documentId = + mutationDocumentId ?? previousDocument?.id ?? randomUUID() + const publish = async (): Promise => { + const prepared = await this.prepareDocumentPublication( + { + id: documentId, + knowledgeBaseId, + sourceId: source.id, + externalId: result.url, + title: result.title, + mimeType: result.contentType, + sourceLocation: result.url, + checksum: createHash('sha256') + .update(result.document.content) + .digest('hex'), + metadata: { + status: 'ready', + size: Buffer.byteLength(result.document.content) + } + }, + result.document, + effectiveLibrary, + effectiveSignal, + parsingTask.id + ) + return this.publishPreparedDocument(effectiveLibrary, prepared) } - const document = this.database.upsertDocument( - { - knowledgeBaseId, - sourceId: source.id, - externalId: result.url, - title: result.title, - mimeType: result.contentType, - sourceLocation: result.url, - checksum: createHash('sha256') - .update(result.document.content) - .digest('hex'), - metadata: { - status: 'ready', - size: Buffer.byteLength(result.document.content) - } - }, - this.createDocumentChunks( - result.document, - effectiveLibrary - ) - ) - this.database.pruneUnreferencedGeneratedGraph(library.id) + const document = documentMutationHeld + ? await publish() + : await this.withDocumentMutation( + documentId, + publish, + effectiveSignal + ) this.updateKnowledgeTask(parsingTask.id, { documentId: document.id, documentName: document.title, - stage: 'embedding', - progress: 80, - message: '网页解析完成,正在生成索引' - }) - await this.indexDocumentEmbeddings( - document, - undefined, - parsingTask.id, - effectiveSignal - ) - effectiveSignal.throwIfAborted() - this.updateKnowledgeTask(parsingTask.id, { - stage: 'graph', + stage: 'finalizing', progress: 92, - message: '正在处理知识图谱' + message: '网页解析与索引完成' }) - await this.extractGraph( - effectiveLibrary, - document, - parsingTask.id, - effectiveSignal - ) effectiveSignal.throwIfAborted() source = this.database.upsertSource({ ...source, @@ -1764,30 +2130,22 @@ export class KnowledgeService { retryOfTaskId?: string ): Promise { const input = knowledgeDocumentRebuildInputSchema.parse(rawInput) - if (!parentTaskId) { - const existing = this.activeDocumentRebuilds.get(input.documentId) - if (existing) { - return existing - } - const operation = this.performRebuildDocument( - input, - signal, - undefined, - retryOfTaskId - ).finally(() => { - if (this.activeDocumentRebuilds.get(input.documentId) === operation) { - this.activeDocumentRebuilds.delete(input.documentId) - } - }) - this.activeDocumentRebuilds.set(input.documentId, operation) - return operation + const existing = this.activeDocumentRebuilds.get(input.documentId) + if (existing) { + return existing } - return this.performRebuildDocument( + const operation = this.performRebuildDocument( input, signal, parentTaskId, retryOfTaskId - ) + ).finally(() => { + if (this.activeDocumentRebuilds.get(input.documentId) === operation) { + this.activeDocumentRebuilds.delete(input.documentId) + } + }) + this.activeDocumentRebuilds.set(input.documentId, operation) + return operation } private async performRebuildDocument( @@ -1795,6 +2153,25 @@ export class KnowledgeService { signal?: AbortSignal, parentTaskId?: string, retryOfTaskId?: string + ): Promise { + return this.withDocumentMutation( + input.documentId, + () => + this.performRebuildDocumentMutation( + input, + signal, + parentTaskId, + retryOfTaskId + ), + signal + ) + } + + private async performRebuildDocumentMutation( + input: KnowledgeDocumentRebuildInput, + signal?: AbortSignal, + parentTaskId?: string, + retryOfTaskId?: string ): Promise { const document = this.database.getDocument(input.documentId) if (!document || document.knowledgeBaseId !== input.knowledgeBaseId) { @@ -1849,7 +2226,9 @@ export class KnowledgeService { effectiveSignal, source.id, undefined, - task.id + task.id, + document.id, + true ) effectiveSignal.throwIfAborted() const rebuilt = this.database.getDocument(document.id) ?? document @@ -1898,8 +2277,7 @@ export class KnowledgeService { progress: 45, message: '正在切分文档' }) - this.database.removeEvidenceForDocument(document.id) - const rebuilt = this.database.upsertDocument( + const prepared = await this.prepareDocumentPublication( { ...document, checksum: createHash('sha256').update(buffer).digest('hex'), @@ -1909,28 +2287,18 @@ export class KnowledgeService { size: fileStat.size } }, - this.createDocumentChunks(parsed, library) + parsed, + library, + effectiveSignal, + task.id ) - this.database.pruneUnreferencedGeneratedGraph(library.id) this.updateKnowledgeTask(task.id, { - stage: 'embedding', - progress: 65, - message: '正在生成向量索引' + stage: 'finalizing', + progress: 92, + 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() + const rebuilt = this.publishPreparedDocument(library, prepared) this.updateKnowledgeTask(task.id, { status: 'succeeded', stage: 'finalizing', @@ -2238,29 +2606,35 @@ export class KnowledgeService { }) tasks.push(task) try { - this.updateKnowledgeTask(task.id, { - status: 'running', - progress: 10, - message: '正在重新抽取知识图谱' - }) - const result = await this.extractGraphResult( - library, - document, + await this.withDocumentMutation( + document.id, + async () => { + this.updateKnowledgeTask(task.id, { + status: 'running', + progress: 10, + message: '正在重新抽取知识图谱' + }) + const result = await this.extractGraphResult( + library, + document, + effectiveSignal + ) + effectiveSignal.throwIfAborted() + this.updateKnowledgeTask(task.id, { + progress: 85, + message: '正在保存实体和关系' + }) + this.database.replaceEvidenceForDocument(document.id, () => { + this.storeExtractedGraph(library, document, result) + }) + effectiveSignal.throwIfAborted() + this.updateKnowledgeTask(task.id, { + status: 'succeeded', + message: `已抽取 ${result.entities.length} 个实体、${result.relations.length} 条关系` + }) + }, effectiveSignal ) - effectiveSignal.throwIfAborted() - this.updateKnowledgeTask(task.id, { - progress: 85, - message: '正在保存实体和关系' - }) - this.database.replaceEvidenceForDocument(document.id, () => { - this.storeExtractedGraph(library, document, result) - }) - effectiveSignal.throwIfAborted() - this.updateKnowledgeTask(task.id, { - status: 'succeeded', - message: `已抽取 ${result.entities.length} 个实体、${result.relations.length} 条关系` - }) this.updateKnowledgeTask(parentTask.id, { progress: ((index + 1) / Math.max(documents.length, 1)) * 100, completedItems: index + 1, @@ -2481,7 +2855,14 @@ export class KnowledgeService { for (const document of existing) { signal.throwIfAborted() 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( @@ -2551,52 +2932,44 @@ export class KnowledgeService { library.id, graphStrategy ) - const document = this.database.upsertDocument( - { - knowledgeBaseId: library.id, - sourceId: source.id, - externalId: file.relativePath, - title: parsed.title, - mimeType: mimeTypeFromFileName( - file.absolutePath, - 'text/plain' - ), - sourceLocation: file.absolutePath, - checksum, - metadata: { - status: 'ready', - size: file.size - } + const documentId = previous?.id ?? randomUUID() + const document = await this.withDocumentMutation( + documentId, + async () => { + const prepared = await this.prepareDocumentPublication( + { + id: documentId, + knowledgeBaseId: library.id, + sourceId: source.id, + externalId: file.relativePath, + title: parsed.title, + mimeType: mimeTypeFromFileName( + file.absolutePath, + 'text/plain' + ), + sourceLocation: file.absolutePath, + checksum, + metadata: { + status: 'ready', + size: file.size + } + }, + parsed, + effectiveLibrary, + signal, + parsingTask.id + ) + return this.publishPreparedDocument(effectiveLibrary, prepared) }, - this.createDocumentChunks(parsed, effectiveLibrary) + signal ) - this.database.pruneUnreferencedGeneratedGraph(library.id) this.updateKnowledgeTask(parsingTask.id, { documentId: document.id, documentName: document.title, - stage: 'embedding', - progress: 80, - message: '文档解析完成,正在生成索引' - }) - this.database.removeEvidenceForDocument(document.id) - await this.indexDocumentEmbeddings( - document, - undefined, - parsingTask.id, - signal - ) - signal.throwIfAborted() - this.updateKnowledgeTask(parsingTask.id, { - stage: 'graph', + stage: 'finalizing', progress: 92, - message: '正在处理知识图谱' + message: '文档解析与索引完成' }) - await this.extractGraph( - effectiveLibrary, - document, - parsingTask.id, - signal - ) signal.throwIfAborted() this.updateKnowledgeTask(parsingTask.id, { 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( document: Document, requestedProvider?: EmbeddingProvider, @@ -2849,7 +3261,7 @@ export class KnowledgeService { } } - private async extractGraph( + private async extractGraphMutation( library: KnowledgeBase, document: Document, operationTaskId?: string, @@ -2927,14 +3339,29 @@ export class KnowledgeService { document: Document, signal?: AbortSignal ): Promise { - const chunks = this.database - .listChunks(document.id, 10_000) - .filter((chunk) => chunk.enabled && chunk.role !== 'parent') + return this.extractGraphChunks( + library, + this.database.listChunks(document.id, 10_000), + signal + ) + } + + private extractGraphChunks( + library: KnowledgeBase, + chunks: readonly Pick[], + signal?: AbortSignal + ): Promise { return extractKnowledgeGraph( - chunks.map((chunk) => ({ - id: chunk.id, - content: chunk.content - })), + chunks + .filter( + (chunk) => + (chunk.enabled ?? true) && + (chunk.role ?? 'standalone') !== 'parent' + ) + .map((chunk): GraphChunk => ({ + id: chunk.id!, + content: chunk.content + })), { strategy: library.graphStrategy, ontology: library.ontologySettings,