fix: publish knowledge rebuilds atomically
This commit is contained in:
lofyer
+2
-2
@@ -4454,9 +4454,9 @@ export function registerIpcHandlers(
|
||||
|
||||
ipcMain.handle(
|
||||
ipcChannels.knowledgeDeleteChunk,
|
||||
(event, input: unknown) => {
|
||||
async (event, input: unknown) => {
|
||||
assertTrustedSender(event, window)
|
||||
const deleted = knowledgeService.deleteChunk(
|
||||
const deleted = await knowledgeService.deleteChunk(
|
||||
knowledgeChunkDeleteInputSchema.parse(input)
|
||||
)
|
||||
if (!deleted) {
|
||||
|
||||
@@ -558,6 +558,65 @@ describe('extraction strategies', () => {
|
||||
).rejects.toThrow('Model extraction is unavailable')
|
||||
})
|
||||
|
||||
it('extracts bounded batches beyond the first chunk window', async () => {
|
||||
const chunks = Array.from(
|
||||
{ length: GRAPH_LIMITS.maximumChunks + 1 },
|
||||
(_, index) => ({
|
||||
id: `batch-${index}`,
|
||||
content:
|
||||
index === GRAPH_LIMITS.maximumChunks
|
||||
? '# Late Batch Entity'
|
||||
: 'ordinary text'
|
||||
})
|
||||
)
|
||||
const modelCalls = vi.fn(async (prompt: string) => {
|
||||
const parsed = JSON.parse(
|
||||
prompt
|
||||
.split('<UNTRUSTED_DOCUMENT_JSON>')[1]!
|
||||
.split('</UNTRUSTED_DOCUMENT_JSON>')[0]!
|
||||
) as Array<{ chunkId: string; content: string }>
|
||||
const chunk = parsed[0]!
|
||||
return {
|
||||
entities: chunk.content.includes('Late Batch')
|
||||
? [{
|
||||
id: 'late',
|
||||
name: 'Late Batch Entity',
|
||||
evidence: [{
|
||||
chunkId: chunk.chunkId,
|
||||
start: 2,
|
||||
end: chunk.content.length
|
||||
}]
|
||||
}]
|
||||
: [],
|
||||
relations: []
|
||||
}
|
||||
})
|
||||
|
||||
const rules = await extractKnowledgeGraph(chunks, { strategy: 'rules' })
|
||||
const model = await extractKnowledgeGraph(chunks, {
|
||||
strategy: 'model',
|
||||
extractStructured: modelCalls
|
||||
})
|
||||
|
||||
expect(rules.entities.some((entity) =>
|
||||
entity.name === 'Late Batch Entity'
|
||||
)).toBe(true)
|
||||
expect(model.entities.some((entity) =>
|
||||
entity.name === 'Late Batch Entity'
|
||||
)).toBe(true)
|
||||
expect(modelCalls).toHaveBeenCalledTimes(2)
|
||||
expect(
|
||||
modelCalls.mock.calls.every(([prompt]) => {
|
||||
const parsed = JSON.parse(
|
||||
prompt
|
||||
.split('<UNTRUSTED_DOCUMENT_JSON>')[1]!
|
||||
.split('</UNTRUSTED_DOCUMENT_JSON>')[0]!
|
||||
) as unknown[]
|
||||
return parsed.length <= GRAPH_LIMITS.maximumChunks
|
||||
})
|
||||
).toBe(true)
|
||||
})
|
||||
|
||||
it('honors cancellation before and after the injected model callback', async () => {
|
||||
const preCancelled = new AbortController()
|
||||
preCancelled.abort()
|
||||
|
||||
@@ -21,6 +21,7 @@ export const GRAPH_LIMITS = {
|
||||
maximumWarnings: 20,
|
||||
maximumWarningLength: 240
|
||||
} as const
|
||||
const maximumEvidencePerRecord = 20
|
||||
|
||||
export type ExtractionStrategy = 'rules' | 'model' | 'hybrid' | 'ask'
|
||||
|
||||
@@ -192,7 +193,7 @@ function throwIfAborted(signal?: AbortSignal): void {
|
||||
function prepareChunks(chunks: readonly GraphChunk[]): GraphChunk[] {
|
||||
const ids = new Set<string>()
|
||||
const prepared: GraphChunk[] = []
|
||||
for (const chunk of chunks.slice(0, GRAPH_LIMITS.maximumChunks)) {
|
||||
for (const chunk of chunks) {
|
||||
const id = truncate(chunk.id.trim(), GRAPH_LIMITS.maximumFieldLength)
|
||||
if (!id || ids.has(id)) {
|
||||
continue
|
||||
@@ -219,6 +220,9 @@ function mergeEvidence(
|
||||
const key = evidenceKey(evidence)
|
||||
if (!merged.has(key)) {
|
||||
merged.set(key, evidence)
|
||||
if (merged.size >= maximumEvidencePerRecord) {
|
||||
break
|
||||
}
|
||||
}
|
||||
}
|
||||
return [...merged.values()]
|
||||
@@ -805,24 +809,39 @@ 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]
|
||||
}
|
||||
}
|
||||
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(prepared, context.settings),
|
||||
createModelPrompt(batch, context.settings),
|
||||
options.signal
|
||||
)
|
||||
throwIfAborted(options.signal)
|
||||
@@ -830,16 +849,24 @@ export async function extractKnowledgeGraph(
|
||||
if (!modelEnvelopeSchema.safeParse(parsedOutput).success) {
|
||||
throw new Error('模型返回的图谱结构无效')
|
||||
}
|
||||
const model = validateModelGraphInternal(parsedOutput, prepared, context)
|
||||
const graph =
|
||||
const model = validateModelGraphInternal(
|
||||
parsedOutput,
|
||||
batch,
|
||||
batchContext
|
||||
)
|
||||
graph = mergeKnowledgeGraphsInternal(
|
||||
graph,
|
||||
strategy === 'hybrid'
|
||||
? mergeKnowledgeGraphsInternal(rules, model, context)
|
||||
: model
|
||||
? mergeKnowledgeGraphsInternal(rules, model, batchContext)
|
||||
: model,
|
||||
context
|
||||
)
|
||||
}
|
||||
return {
|
||||
...graph,
|
||||
strategy,
|
||||
requiresModelApproval: false,
|
||||
warnings: [...context.warnings]
|
||||
requiresModelApproval: strategy === 'ask',
|
||||
warnings: [...context.warnings].slice(0, GRAPH_LIMITS.maximumWarnings)
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -80,7 +80,7 @@ import type {
|
||||
VectorSearchOptions
|
||||
} from './types'
|
||||
|
||||
const DATABASE_VERSION = 10
|
||||
const DATABASE_VERSION = 11
|
||||
const MAX_ID_LENGTH = 128
|
||||
const MAX_NAME_LENGTH = 512
|
||||
const MAX_LOCATION_LENGTH = 8192
|
||||
@@ -118,6 +118,22 @@ export type HybridSearchResultPage = {
|
||||
vectorScannedCount: number
|
||||
}
|
||||
|
||||
export type PreparedEmbeddingReplacement = {
|
||||
replacementId: string
|
||||
provider: string
|
||||
model: string
|
||||
}
|
||||
|
||||
export type DocumentPublicationOptions = {
|
||||
embeddingReplacement?: PreparedEmbeddingReplacement
|
||||
embeddingError?: {
|
||||
provider: string
|
||||
model: string
|
||||
message: string
|
||||
}
|
||||
afterChunksInserted?: (document: Document) => void
|
||||
}
|
||||
|
||||
type Row = Record<string, null | number | bigint | string | Uint8Array>
|
||||
|
||||
const taskScopes = knowledgeTaskScopeSchema.options
|
||||
@@ -1635,6 +1651,22 @@ export class KnowledgeDatabase {
|
||||
.map(mapSource)
|
||||
}
|
||||
|
||||
listSourcesForSnapshot(knowledgeBaseId: string): KnowledgeSource[] {
|
||||
const normalizedId = requiredString(
|
||||
knowledgeBaseId,
|
||||
'knowledgeBaseId',
|
||||
MAX_ID_LENGTH
|
||||
)
|
||||
return this.requireDatabase()
|
||||
.prepare(
|
||||
`SELECT * FROM knowledge_sources
|
||||
WHERE knowledge_base_id = ?
|
||||
ORDER BY created_at ASC, id ASC`
|
||||
)
|
||||
.all(normalizedId)
|
||||
.map(mapSource)
|
||||
}
|
||||
|
||||
getSource(id: string): KnowledgeSource | undefined {
|
||||
const row = this.requireDatabase()
|
||||
.prepare('SELECT * FROM knowledge_sources WHERE id = ?')
|
||||
@@ -1654,6 +1686,14 @@ export class KnowledgeDatabase {
|
||||
upsertDocument(
|
||||
input: UpsertDocumentInput,
|
||||
chunks: ReplaceChunkInput[]
|
||||
): Document {
|
||||
return this.publishDocument(input, chunks)
|
||||
}
|
||||
|
||||
publishDocument(
|
||||
input: UpsertDocumentInput,
|
||||
chunks: ReplaceChunkInput[],
|
||||
options: DocumentPublicationOptions = {}
|
||||
): Document {
|
||||
if (!Array.isArray(chunks) || chunks.length > MAX_CHUNKS) {
|
||||
throw new RangeError(`chunks must contain at most ${MAX_CHUNKS} items`)
|
||||
@@ -1716,6 +1756,58 @@ export class KnowledgeDatabase {
|
||||
)
|
||||
const checksum = optionalString(input.checksum, 'checksum', 512)
|
||||
const metadata = jsonObject(input.metadata, 'metadata')
|
||||
const embeddingReplacement = options.embeddingReplacement
|
||||
? {
|
||||
replacementId: requiredString(
|
||||
options.embeddingReplacement.replacementId,
|
||||
'embeddingReplacement.replacementId',
|
||||
MAX_ID_LENGTH
|
||||
),
|
||||
provider: requiredString(
|
||||
options.embeddingReplacement.provider,
|
||||
'embeddingReplacement.provider',
|
||||
MAX_EMBEDDING_PROVIDER_LENGTH
|
||||
),
|
||||
model: requiredString(
|
||||
options.embeddingReplacement.model,
|
||||
'embeddingReplacement.model',
|
||||
MAX_EMBEDDING_MODEL_LENGTH
|
||||
)
|
||||
}
|
||||
: undefined
|
||||
const embeddingError = options.embeddingError
|
||||
? {
|
||||
provider: requiredString(
|
||||
options.embeddingError.provider,
|
||||
'embeddingError.provider',
|
||||
MAX_EMBEDDING_PROVIDER_LENGTH
|
||||
),
|
||||
model: requiredString(
|
||||
options.embeddingError.model,
|
||||
'embeddingError.model',
|
||||
MAX_EMBEDDING_MODEL_LENGTH
|
||||
),
|
||||
message: requiredString(
|
||||
options.embeddingError.message,
|
||||
'embeddingError.message',
|
||||
MAX_EMBEDDING_ERROR_LENGTH,
|
||||
false
|
||||
)
|
||||
}
|
||||
: undefined
|
||||
if (embeddingReplacement && embeddingError) {
|
||||
throw new Error(
|
||||
'Document publication cannot include embeddings and an embedding error'
|
||||
)
|
||||
}
|
||||
if (embeddingReplacement) {
|
||||
this.validatePreparedDocumentEmbeddings(
|
||||
database,
|
||||
embeddingReplacement,
|
||||
id,
|
||||
normalizedChunks
|
||||
)
|
||||
}
|
||||
const now = new Date().toISOString()
|
||||
|
||||
this.transaction(database, () => {
|
||||
@@ -1749,7 +1841,7 @@ export class KnowledgeDatabase {
|
||||
now
|
||||
)
|
||||
database
|
||||
.prepare('DELETE FROM embedding_index_state WHERE document_id = ?')
|
||||
.prepare('DELETE FROM graph_evidence WHERE document_id = ?')
|
||||
.run(id)
|
||||
database.prepare('DELETE FROM chunks WHERE document_id = ?').run(id)
|
||||
const insertChunk = database.prepare(
|
||||
@@ -1794,6 +1886,29 @@ export class KnowledgeDatabase {
|
||||
indexContent
|
||||
)
|
||||
}
|
||||
options.afterChunksInserted?.(this.requiredDocument(id))
|
||||
database
|
||||
.prepare(
|
||||
'DELETE FROM embedding_index_state WHERE document_id = ?'
|
||||
)
|
||||
.run(id)
|
||||
if (embeddingReplacement) {
|
||||
this.publishPreparedDocumentEmbeddings(
|
||||
database,
|
||||
embeddingReplacement,
|
||||
id,
|
||||
knowledgeBaseId,
|
||||
now
|
||||
)
|
||||
} else if (embeddingError) {
|
||||
this.recordEmbeddingIndexError(
|
||||
id,
|
||||
embeddingError.provider,
|
||||
embeddingError.model,
|
||||
embeddingError.message
|
||||
)
|
||||
}
|
||||
this.pruneUnreferencedGeneratedGraph(knowledgeBaseId)
|
||||
})
|
||||
return this.requiredDocument(id)
|
||||
}
|
||||
@@ -1851,6 +1966,10 @@ export class KnowledgeDatabase {
|
||||
.map(mapDocument)
|
||||
}
|
||||
|
||||
listDocumentsForSnapshot(knowledgeBaseId: string): Document[] {
|
||||
return this.listDocumentsForLibraryRebuild(knowledgeBaseId)
|
||||
}
|
||||
|
||||
getDocumentChunkCounts(knowledgeBaseId: string): Map<string, number> {
|
||||
const normalizedId = requiredString(
|
||||
knowledgeBaseId,
|
||||
@@ -2365,6 +2484,17 @@ export class KnowledgeDatabase {
|
||||
return randomUUID()
|
||||
}
|
||||
|
||||
beginPreparedDocumentEmbeddingReplacement(
|
||||
documentId: string,
|
||||
provider: string,
|
||||
model: string
|
||||
): string {
|
||||
requiredString(documentId, 'documentId', MAX_ID_LENGTH)
|
||||
requiredString(provider, 'provider', MAX_EMBEDDING_PROVIDER_LENGTH)
|
||||
requiredString(model, 'model', MAX_EMBEDDING_MODEL_LENGTH)
|
||||
return randomUUID()
|
||||
}
|
||||
|
||||
appendDocumentEmbeddingBatch(
|
||||
replacementId: string,
|
||||
documentId: string,
|
||||
@@ -2465,26 +2595,95 @@ export class KnowledgeDatabase {
|
||||
}
|
||||
return { chunkId, checksum, ...vector }
|
||||
})
|
||||
const insert = database.prepare(
|
||||
`INSERT INTO embedding_rebuild_staging
|
||||
(replacement_id, document_id, provider, model, chunk_id,
|
||||
dimensions, content_checksum, vector, magnitude)
|
||||
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)`
|
||||
this.insertEmbeddingReplacementBatch(database, {
|
||||
replacementId: normalizedReplacementId,
|
||||
documentId: normalizedDocumentId,
|
||||
provider: normalizedProvider,
|
||||
model: normalizedModel,
|
||||
embeddings: normalized
|
||||
})
|
||||
}
|
||||
|
||||
appendPreparedDocumentEmbeddingBatch(
|
||||
replacementId: string,
|
||||
documentId: string,
|
||||
provider: string,
|
||||
model: string,
|
||||
embeddings: readonly ChunkEmbeddingInput[]
|
||||
): void {
|
||||
const normalizedReplacementId = requiredString(
|
||||
replacementId,
|
||||
'replacementId',
|
||||
MAX_ID_LENGTH
|
||||
)
|
||||
this.transaction(database, () => {
|
||||
for (const item of normalized) {
|
||||
insert.run(
|
||||
normalizedReplacementId,
|
||||
normalizedDocumentId,
|
||||
normalizedProvider,
|
||||
normalizedModel,
|
||||
item.chunkId,
|
||||
item.dimensions,
|
||||
item.checksum,
|
||||
item.bytes,
|
||||
item.magnitude
|
||||
const normalizedDocumentId = requiredString(
|
||||
documentId,
|
||||
'documentId',
|
||||
MAX_ID_LENGTH
|
||||
)
|
||||
const normalizedProvider = requiredString(
|
||||
provider,
|
||||
'provider',
|
||||
MAX_EMBEDDING_PROVIDER_LENGTH
|
||||
)
|
||||
const normalizedModel = requiredString(
|
||||
model,
|
||||
'model',
|
||||
MAX_EMBEDDING_MODEL_LENGTH
|
||||
)
|
||||
if (
|
||||
!Array.isArray(embeddings) ||
|
||||
embeddings.length < 1 ||
|
||||
embeddings.length > MAX_EMBEDDING_BATCH
|
||||
) {
|
||||
throw new RangeError(
|
||||
`embeddings must contain between 1 and ${MAX_EMBEDDING_BATCH} items`
|
||||
)
|
||||
}
|
||||
const database = this.requireDatabase()
|
||||
const existingDimensions = database
|
||||
.prepare(
|
||||
`SELECT dimensions FROM embedding_rebuild_staging
|
||||
WHERE replacement_id = ? LIMIT 1`
|
||||
)
|
||||
.get(normalizedReplacementId)
|
||||
let dimensions = existingDimensions
|
||||
? asNumber(existingDimensions, 'dimensions')
|
||||
: undefined
|
||||
const seen = new Set<string>()
|
||||
const normalized = embeddings.map((embedding, index) => {
|
||||
const chunkId = requiredString(
|
||||
embedding.chunkId,
|
||||
`embeddings[${index}].chunkId`,
|
||||
MAX_ID_LENGTH
|
||||
)
|
||||
if (seen.has(chunkId)) {
|
||||
throw new Error('Embedding chunk IDs must be unique within a batch')
|
||||
}
|
||||
seen.add(chunkId)
|
||||
const checksum = normalizedChecksum(
|
||||
embedding.contentChecksum,
|
||||
`embeddings[${index}].contentChecksum`
|
||||
)
|
||||
const vector = normalizeVector(
|
||||
embedding.vector,
|
||||
`embeddings[${index}].vector`
|
||||
)
|
||||
if (dimensions === undefined) {
|
||||
dimensions = vector.dimensions
|
||||
} else if (dimensions !== vector.dimensions) {
|
||||
throw new Error(
|
||||
'Document embeddings must have consistent dimensions'
|
||||
)
|
||||
}
|
||||
return { chunkId, checksum, ...vector }
|
||||
})
|
||||
this.insertEmbeddingReplacementBatch(database, {
|
||||
replacementId: normalizedReplacementId,
|
||||
documentId: normalizedDocumentId,
|
||||
provider: normalizedProvider,
|
||||
model: normalizedModel,
|
||||
embeddings: normalized
|
||||
})
|
||||
}
|
||||
|
||||
@@ -4324,6 +4523,14 @@ export class KnowledgeDatabase {
|
||||
)
|
||||
.run(10, new Date().toISOString())
|
||||
}
|
||||
if (currentVersion < 11) {
|
||||
this.migrateToVersion11(database)
|
||||
database
|
||||
.prepare(
|
||||
'INSERT INTO schema_migrations (version, applied_at) VALUES (?, ?)'
|
||||
)
|
||||
.run(11, new Date().toISOString())
|
||||
}
|
||||
database.exec(`PRAGMA user_version = ${DATABASE_VERSION}`)
|
||||
database.exec('COMMIT')
|
||||
} catch (error) {
|
||||
@@ -4947,6 +5154,228 @@ export class KnowledgeDatabase {
|
||||
`)
|
||||
}
|
||||
|
||||
private migrateToVersion11(database: DatabaseSync): void {
|
||||
database.exec('DROP TABLE embedding_rebuild_staging;')
|
||||
database.exec(`
|
||||
CREATE TABLE embedding_rebuild_staging (
|
||||
replacement_id TEXT NOT NULL,
|
||||
document_id TEXT NOT NULL,
|
||||
provider TEXT NOT NULL,
|
||||
model TEXT NOT NULL,
|
||||
chunk_id TEXT NOT NULL,
|
||||
dimensions INTEGER NOT NULL
|
||||
CHECK (dimensions >= 1 AND dimensions <= 8192),
|
||||
content_checksum TEXT NOT NULL
|
||||
CHECK (length(content_checksum) = 64),
|
||||
vector BLOB NOT NULL,
|
||||
magnitude REAL NOT NULL CHECK (magnitude > 0),
|
||||
PRIMARY KEY (replacement_id, chunk_id)
|
||||
);
|
||||
CREATE INDEX embedding_rebuild_staging_document_idx
|
||||
ON embedding_rebuild_staging(
|
||||
document_id, provider, model, replacement_id
|
||||
);
|
||||
`)
|
||||
}
|
||||
|
||||
private validatePreparedDocumentEmbeddings(
|
||||
database: DatabaseSync,
|
||||
replacement: PreparedEmbeddingReplacement,
|
||||
documentId: string,
|
||||
chunks: ReturnType<KnowledgeDatabase['normalizeChunks']>
|
||||
): void {
|
||||
const indexableChunks = chunks.filter(
|
||||
(chunk) => chunk.enabled && chunk.role !== 'parent'
|
||||
)
|
||||
if (indexableChunks.length === 0) {
|
||||
throw new Error(
|
||||
'Prepared embedding replacement requires indexable chunks'
|
||||
)
|
||||
}
|
||||
const stagedRows = database
|
||||
.prepare(
|
||||
`SELECT chunk_id, content_checksum, dimensions
|
||||
FROM embedding_rebuild_staging
|
||||
WHERE replacement_id = ? AND document_id = ?
|
||||
AND provider = ? AND model = ?
|
||||
ORDER BY chunk_id`
|
||||
)
|
||||
.all(
|
||||
replacement.replacementId,
|
||||
documentId,
|
||||
replacement.provider,
|
||||
replacement.model
|
||||
)
|
||||
if (stagedRows.length !== indexableChunks.length) {
|
||||
throw new Error('Embeddings must cover every candidate document chunk')
|
||||
}
|
||||
const chunksById = new Map(
|
||||
indexableChunks.map((chunk) => [
|
||||
chunk.id,
|
||||
contentChecksum(
|
||||
chunkIndexContent(chunk.content, parseObject(chunk.metadata))
|
||||
)
|
||||
])
|
||||
)
|
||||
let dimensions: number | undefined
|
||||
for (const row of stagedRows) {
|
||||
const chunkId = asString(row, 'chunk_id')
|
||||
const checksum = chunksById.get(chunkId)
|
||||
if (
|
||||
checksum === undefined ||
|
||||
checksum !== asString(row, 'content_checksum')
|
||||
) {
|
||||
throw new Error('Prepared embedding content does not match the chunk')
|
||||
}
|
||||
chunksById.delete(chunkId)
|
||||
const rowDimensions = asNumber(row, 'dimensions')
|
||||
if (dimensions === undefined) {
|
||||
dimensions = rowDimensions
|
||||
} else if (dimensions !== rowDimensions) {
|
||||
throw new Error(
|
||||
'Document embeddings must have consistent dimensions'
|
||||
)
|
||||
}
|
||||
}
|
||||
if (chunksById.size > 0) {
|
||||
throw new Error('Embeddings must cover every candidate document chunk')
|
||||
}
|
||||
}
|
||||
|
||||
private publishPreparedDocumentEmbeddings(
|
||||
database: DatabaseSync,
|
||||
replacement: PreparedEmbeddingReplacement,
|
||||
documentId: string,
|
||||
knowledgeBaseId: string,
|
||||
now: string
|
||||
): void {
|
||||
let dimensions: number | undefined
|
||||
const indexHash = createHash('sha256')
|
||||
let firstChecksum = true
|
||||
for (const row of database
|
||||
.prepare(
|
||||
`SELECT chunk_id, content_checksum, dimensions
|
||||
FROM embedding_rebuild_staging
|
||||
WHERE replacement_id = ? AND document_id = ?
|
||||
AND provider = ? AND model = ?
|
||||
ORDER BY chunk_id`
|
||||
)
|
||||
.iterate(
|
||||
replacement.replacementId,
|
||||
documentId,
|
||||
replacement.provider,
|
||||
replacement.model
|
||||
)) {
|
||||
const chunkId = asString(row, 'chunk_id')
|
||||
const checksum = asString(row, 'content_checksum')
|
||||
if (!firstChecksum) {
|
||||
indexHash.update('\n')
|
||||
}
|
||||
indexHash.update(`${chunkId}\0${checksum}`)
|
||||
firstChecksum = false
|
||||
dimensions ??= asNumber(row, 'dimensions')
|
||||
}
|
||||
database
|
||||
.prepare(
|
||||
`INSERT INTO chunk_embeddings
|
||||
(chunk_id, knowledge_base_id, provider, model, dimensions,
|
||||
content_checksum, vector, magnitude, created_at, updated_at)
|
||||
SELECT chunk_id, ?, provider, model, dimensions,
|
||||
content_checksum, vector, magnitude, ?, ?
|
||||
FROM embedding_rebuild_staging
|
||||
WHERE replacement_id = ? AND document_id = ?
|
||||
AND provider = ? AND model = ?`
|
||||
)
|
||||
.run(
|
||||
knowledgeBaseId,
|
||||
now,
|
||||
now,
|
||||
replacement.replacementId,
|
||||
documentId,
|
||||
replacement.provider,
|
||||
replacement.model
|
||||
)
|
||||
database
|
||||
.prepare(
|
||||
`INSERT INTO embedding_index_state
|
||||
(document_id, knowledge_base_id, provider, model, dimensions,
|
||||
content_checksum, status, last_error, updated_at)
|
||||
VALUES (?, ?, ?, ?, ?, ?, 'ready', NULL, ?)`
|
||||
)
|
||||
.run(
|
||||
documentId,
|
||||
knowledgeBaseId,
|
||||
replacement.provider,
|
||||
replacement.model,
|
||||
dimensions ?? null,
|
||||
indexHash.digest('hex'),
|
||||
now
|
||||
)
|
||||
database
|
||||
.prepare(
|
||||
'DELETE FROM embedding_rebuild_staging WHERE replacement_id = ?'
|
||||
)
|
||||
.run(replacement.replacementId)
|
||||
}
|
||||
|
||||
private insertEmbeddingReplacementBatch(
|
||||
database: DatabaseSync,
|
||||
input: {
|
||||
replacementId: string
|
||||
documentId: string
|
||||
provider: string
|
||||
model: string
|
||||
embeddings: ReadonlyArray<{
|
||||
chunkId: string
|
||||
checksum: string
|
||||
bytes: Buffer
|
||||
dimensions: number
|
||||
magnitude: number
|
||||
}>
|
||||
}
|
||||
): void {
|
||||
const insert = database.prepare(
|
||||
`INSERT INTO embedding_rebuild_staging
|
||||
(replacement_id, document_id, provider, model, chunk_id,
|
||||
dimensions, content_checksum, vector, magnitude)
|
||||
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)`
|
||||
)
|
||||
this.transaction(database, () => {
|
||||
const existingTarget = database
|
||||
.prepare(
|
||||
`SELECT document_id, provider, model
|
||||
FROM embedding_rebuild_staging
|
||||
WHERE replacement_id = ? LIMIT 1`
|
||||
)
|
||||
.get(input.replacementId)
|
||||
if (
|
||||
existingTarget &&
|
||||
(
|
||||
asString(existingTarget, 'document_id') !== input.documentId ||
|
||||
asString(existingTarget, 'provider') !== input.provider ||
|
||||
asString(existingTarget, 'model') !== input.model
|
||||
)
|
||||
) {
|
||||
throw new Error(
|
||||
'Embedding replacement target cannot change between batches'
|
||||
)
|
||||
}
|
||||
for (const item of input.embeddings) {
|
||||
insert.run(
|
||||
input.replacementId,
|
||||
input.documentId,
|
||||
input.provider,
|
||||
input.model,
|
||||
item.chunkId,
|
||||
item.dimensions,
|
||||
item.checksum,
|
||||
item.bytes,
|
||||
item.magnitude
|
||||
)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
private normalizeChunks(chunks: ReplaceChunkInput[]): Array<{
|
||||
id: string
|
||||
ordinal: number
|
||||
@@ -5175,6 +5604,10 @@ export class KnowledgeDatabase {
|
||||
}
|
||||
|
||||
private transaction(database: DatabaseSync, operation: () => void): void {
|
||||
if (database.isTransaction) {
|
||||
operation()
|
||||
return
|
||||
}
|
||||
database.exec('BEGIN IMMEDIATE')
|
||||
try {
|
||||
operation()
|
||||
|
||||
@@ -464,6 +464,104 @@ describe('KnowledgeService', () => {
|
||||
expect(service.snapshot(library.id).evidence).toEqual([])
|
||||
})
|
||||
|
||||
it('keeps a committed chunk edit successful when graph refresh fails', async () => {
|
||||
const extractStructured = vi.fn(async () => {
|
||||
throw new Error('synthetic edit graph failure')
|
||||
})
|
||||
const { directory, service } = await createService(
|
||||
undefined,
|
||||
undefined,
|
||||
extractStructured
|
||||
)
|
||||
const sourcePath = join(directory, 'edit-graph-failure.md')
|
||||
await writeFile(sourcePath, 'initial edit content', 'utf8')
|
||||
const library = service.createLibrary({
|
||||
name: 'Edit graph failure',
|
||||
storageMode: 'reference',
|
||||
graphEnabled: false,
|
||||
graphStrategy: 'model'
|
||||
})
|
||||
await service.importPaths(library.id, [sourcePath])
|
||||
service.database.updateKnowledgeBase(library.id, { graphEnabled: true })
|
||||
const document = service.snapshot(library.id).documents[0]!
|
||||
const chunk = service.database
|
||||
.listChunks(document.id)
|
||||
.find((candidate) => candidate.role !== 'parent')!
|
||||
|
||||
await expect(service.updateChunk({
|
||||
knowledgeBaseId: library.id,
|
||||
documentId: document.id,
|
||||
chunkId: chunk.id,
|
||||
content: 'committed edit survives graph failure'
|
||||
})).resolves.toMatchObject({
|
||||
content: 'committed edit survives graph failure'
|
||||
})
|
||||
await vi.waitFor(() => expect(extractStructured).toHaveBeenCalled())
|
||||
expect(service.search(library.id, 'committed')).toHaveLength(1)
|
||||
})
|
||||
|
||||
it('coalesces edit graph refreshes so stale extraction cannot publish last', async () => {
|
||||
const resolvers: Array<(value: {
|
||||
entities: unknown[]
|
||||
relations: unknown[]
|
||||
}) => void> = []
|
||||
const extractStructured: ExtractStructured = () =>
|
||||
new Promise((resolve) => {
|
||||
resolvers.push(resolve)
|
||||
})
|
||||
const { directory, service } = await createService(
|
||||
undefined,
|
||||
undefined,
|
||||
extractStructured
|
||||
)
|
||||
const sourcePath = join(directory, 'coalesced-edit-graph.md')
|
||||
await writeFile(sourcePath, 'initial graph content', 'utf8')
|
||||
const library = service.createLibrary({
|
||||
name: 'Coalesced edit graph',
|
||||
storageMode: 'reference',
|
||||
graphEnabled: false,
|
||||
graphStrategy: 'model'
|
||||
})
|
||||
await service.importPaths(library.id, [sourcePath])
|
||||
service.database.updateKnowledgeBase(library.id, { graphEnabled: true })
|
||||
const document = service.snapshot(library.id).documents[0]!
|
||||
const chunk = service.database
|
||||
.listChunks(document.id)
|
||||
.find((candidate) => candidate.role !== 'parent')!
|
||||
|
||||
await service.updateChunk({
|
||||
knowledgeBaseId: library.id,
|
||||
documentId: document.id,
|
||||
chunkId: chunk.id,
|
||||
content: 'first graph edit'
|
||||
})
|
||||
await vi.waitFor(() => expect(resolvers).toHaveLength(1))
|
||||
await service.updateChunk({
|
||||
knowledgeBaseId: library.id,
|
||||
documentId: document.id,
|
||||
chunkId: chunk.id,
|
||||
content: 'second graph edit'
|
||||
})
|
||||
await new Promise((resolve) => setTimeout(resolve, 0))
|
||||
expect(resolvers).toHaveLength(1)
|
||||
|
||||
resolvers.shift()?.({ entities: [], relations: [] })
|
||||
await vi.waitFor(() => expect(resolvers).toHaveLength(1))
|
||||
resolvers.shift()?.({ entities: [], relations: [] })
|
||||
await vi.waitFor(() => {
|
||||
expect(
|
||||
service.snapshot(library.id).tasks.filter(
|
||||
(task) => task.kind === 'graph' && task.status === 'succeeded'
|
||||
)
|
||||
).toHaveLength(1)
|
||||
})
|
||||
expect(service.database.listChunks(document.id)).toEqual(
|
||||
expect.arrayContaining([
|
||||
expect.objectContaining({ content: 'second graph edit' })
|
||||
])
|
||||
)
|
||||
})
|
||||
|
||||
it('fails hybrid reextraction when model extraction fails', async () => {
|
||||
const extractStructured = vi.fn(async () => {
|
||||
throw new Error('模型未返回图谱内容')
|
||||
@@ -1062,6 +1160,133 @@ describe('KnowledgeService', () => {
|
||||
)
|
||||
})
|
||||
|
||||
it('preserves old chunks, vectors, and graph evidence when rebuild extraction fails', async () => {
|
||||
let failExtraction = false
|
||||
const extractStructured: ExtractStructured = async () => {
|
||||
if (failExtraction) {
|
||||
throw new Error('synthetic graph failure')
|
||||
}
|
||||
return {
|
||||
entities: [],
|
||||
relations: []
|
||||
}
|
||||
}
|
||||
const provider: EmbeddingProvider = {
|
||||
provider: 'atomic-provider',
|
||||
model: 'atomic-model',
|
||||
embed: async (input) => input.map(() => [1, 0])
|
||||
}
|
||||
const { directory, service } = await createService(
|
||||
undefined,
|
||||
provider,
|
||||
extractStructured
|
||||
)
|
||||
const sourcePath = join(directory, 'atomic-rebuild.md')
|
||||
await writeFile(sourcePath, '# Original Entity\nold searchable text', 'utf8')
|
||||
const library = service.createLibrary({
|
||||
name: 'Atomic rebuild',
|
||||
storageMode: 'reference',
|
||||
graphEnabled: true,
|
||||
graphStrategy: 'hybrid'
|
||||
})
|
||||
await service.importPaths(library.id, [sourcePath])
|
||||
const before = service.snapshot(library.id)
|
||||
const document = before.documents[0]!
|
||||
const oldChunk = service.database.listChunks(document.id)[0]!
|
||||
const oldEvidence = before.evidence
|
||||
failExtraction = true
|
||||
await writeFile(sourcePath, '# Replacement Entity\nnew searchable text', 'utf8')
|
||||
|
||||
await expect(service.rebuildDocument({
|
||||
knowledgeBaseId: library.id,
|
||||
documentId: document.id
|
||||
})).rejects.toThrow('synthetic graph failure')
|
||||
|
||||
expect(service.search(library.id, 'old')[0]?.chunk.id).toBe(oldChunk.id)
|
||||
expect(service.search(library.id, 'new')).toEqual([])
|
||||
expect(service.database.vectorSearch({
|
||||
knowledgeBaseId: library.id,
|
||||
provider: embeddingStorageProvider(provider),
|
||||
model: provider.model,
|
||||
vector: [1, 0],
|
||||
limit: 1
|
||||
})[0]?.chunk.id).toBe(oldChunk.id)
|
||||
expect(service.snapshot(library.id).evidence).toEqual(oldEvidence)
|
||||
})
|
||||
|
||||
it('serializes a rebuild and chunk edit so the latest mutation wins', async () => {
|
||||
let releaseParser: (() => void) | undefined
|
||||
let parserStarted: (() => void) | undefined
|
||||
const started = new Promise<void>((resolve) => {
|
||||
parserStarted = resolve
|
||||
})
|
||||
let blockParser = false
|
||||
const directory = await mkdtemp(
|
||||
join(tmpdir(), 'goodbuddy-knowledge-service-')
|
||||
)
|
||||
temporaryDirectories.push(directory)
|
||||
const service = new KnowledgeService({
|
||||
databasePath: join(directory, 'knowledge.sqlite'),
|
||||
managedRoot: join(directory, 'managed'),
|
||||
parseDocument: async (name, buffer) => {
|
||||
if (blockParser) {
|
||||
parserStarted?.()
|
||||
await new Promise<void>((resolve) => {
|
||||
releaseParser = resolve
|
||||
})
|
||||
}
|
||||
const content = blockParser
|
||||
? 'rebuilt content'
|
||||
: buffer.toString('utf8')
|
||||
return {
|
||||
title: name,
|
||||
content,
|
||||
sourceFormat: 'text',
|
||||
sections: [{
|
||||
locator: '全文',
|
||||
content
|
||||
}],
|
||||
warnings: []
|
||||
}
|
||||
}
|
||||
})
|
||||
await service.initialize()
|
||||
services.push(service)
|
||||
const sourcePath = join(directory, 'mutation-gate.txt')
|
||||
await writeFile(sourcePath, 'initial content', 'utf8')
|
||||
const library = service.createLibrary({
|
||||
name: 'Mutation gate',
|
||||
storageMode: 'reference',
|
||||
graphEnabled: false
|
||||
})
|
||||
await service.importPaths(library.id, [sourcePath])
|
||||
const document = service.snapshot(library.id).documents[0]!
|
||||
const originalChunk = service.database
|
||||
.listChunks(document.id)
|
||||
.find((chunk) => chunk.role !== 'parent')!
|
||||
blockParser = true
|
||||
await writeFile(sourcePath, 'rebuilt source bytes', 'utf8')
|
||||
|
||||
const rebuilding = service.rebuildDocument({
|
||||
knowledgeBaseId: library.id,
|
||||
documentId: document.id
|
||||
})
|
||||
await started
|
||||
const editing = service.updateChunk({
|
||||
knowledgeBaseId: library.id,
|
||||
documentId: document.id,
|
||||
chunkId: originalChunk.id,
|
||||
content: 'manual latest content'
|
||||
})
|
||||
releaseParser?.()
|
||||
await rebuilding
|
||||
await expect(editing).rejects.toThrow(
|
||||
'Chunk must belong to the requested document'
|
||||
)
|
||||
expect(service.search(library.id, 'rebuilt')).toHaveLength(1)
|
||||
expect(service.search(library.id, 'manual')).toEqual([])
|
||||
})
|
||||
|
||||
it('aborts standalone document rebuild work and keeps child capabilities honest', async () => {
|
||||
let parserSignal: AbortSignal | undefined
|
||||
const directory = await mkdtemp(
|
||||
@@ -1334,7 +1559,7 @@ describe('KnowledgeService', () => {
|
||||
displayName: 'Original title',
|
||||
status: 'ready'
|
||||
})
|
||||
vi.spyOn(service.database, 'upsertDocument').mockImplementationOnce(
|
||||
vi.spyOn(service.database, 'publishDocument').mockImplementationOnce(
|
||||
() => {
|
||||
throw new Error('synthetic indexing failure')
|
||||
}
|
||||
|
||||
@@ -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<KnowledgeDatabase['publishDocument']>[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<Document>
|
||||
>()
|
||||
private readonly documentMutationTails = new Map<string, Promise<void>>()
|
||||
private readonly backgroundEmbeddingReindexes = new Map<
|
||||
string,
|
||||
Promise<void>
|
||||
>()
|
||||
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 taskOperations = new Map<string, Promise<unknown>>()
|
||||
private readonly sourceSyncTaskIds = new Map<string, string>()
|
||||
@@ -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<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(
|
||||
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<string, number>()
|
||||
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:
|
||||
@@ -1219,8 +1586,9 @@ export class KnowledgeService {
|
||||
|
||||
async updateChunk(
|
||||
rawInput: KnowledgeChunkUpdateInput
|
||||
) {
|
||||
): Promise<ReturnType<KnowledgeDatabase['updateChunk']>> {
|
||||
const input = knowledgeChunkUpdateInputSchema.parse(rawInput)
|
||||
return this.withDocumentMutation(input.documentId, async () => {
|
||||
const current = this.database.getChunkForReference(
|
||||
input.knowledgeBaseId,
|
||||
input.documentId,
|
||||
@@ -1237,29 +1605,32 @@ export class KnowledgeService {
|
||||
)
|
||||
) {
|
||||
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
|
||||
})
|
||||
}
|
||||
|
||||
deleteChunk(rawInput: KnowledgeChunkDeleteInput): boolean {
|
||||
async deleteChunk(rawInput: KnowledgeChunkDeleteInput): Promise<boolean> {
|
||||
const input = knowledgeChunkDeleteInputSchema.parse(rawInput)
|
||||
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<GraphStrategy, 'ask'>,
|
||||
parentTaskId?: string
|
||||
parentTaskId?: string,
|
||||
mutationDocumentId?: string,
|
||||
documentMutationHeld = false
|
||||
): Promise<void> {
|
||||
const library = this.requireLibrary(knowledgeBaseId)
|
||||
const parsingTask = this.createKnowledgeTask({
|
||||
@@ -1618,14 +1991,17 @@ export class KnowledgeService {
|
||||
knowledgeBaseId,
|
||||
graphStrategy
|
||||
)
|
||||
const previousDocument = this.database
|
||||
const previousDocument = mutationDocumentId
|
||||
? this.database.getDocument(mutationDocumentId)
|
||||
: this.database
|
||||
.listDocumentsForSource(source.id)
|
||||
.find((document) => document.externalId === result.url)
|
||||
if (previousDocument) {
|
||||
this.database.removeEvidenceForDocument(previousDocument.id)
|
||||
}
|
||||
const document = this.database.upsertDocument(
|
||||
const documentId =
|
||||
mutationDocumentId ?? previousDocument?.id ?? randomUUID()
|
||||
const publish = async (): Promise<Document> => {
|
||||
const prepared = await this.prepareDocumentPublication(
|
||||
{
|
||||
id: documentId,
|
||||
knowledgeBaseId,
|
||||
sourceId: source.id,
|
||||
externalId: result.url,
|
||||
@@ -1640,37 +2016,27 @@ export class KnowledgeService {
|
||||
size: Buffer.byteLength(result.document.content)
|
||||
}
|
||||
},
|
||||
this.createDocumentChunks(
|
||||
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, {
|
||||
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,7 +2130,6 @@ export class KnowledgeService {
|
||||
retryOfTaskId?: string
|
||||
): Promise<Document> {
|
||||
const input = knowledgeDocumentRebuildInputSchema.parse(rawInput)
|
||||
if (!parentTaskId) {
|
||||
const existing = this.activeDocumentRebuilds.get(input.documentId)
|
||||
if (existing) {
|
||||
return existing
|
||||
@@ -1772,7 +2137,7 @@ export class KnowledgeService {
|
||||
const operation = this.performRebuildDocument(
|
||||
input,
|
||||
signal,
|
||||
undefined,
|
||||
parentTaskId,
|
||||
retryOfTaskId
|
||||
).finally(() => {
|
||||
if (this.activeDocumentRebuilds.get(input.documentId) === operation) {
|
||||
@@ -1782,15 +2147,27 @@ export class KnowledgeService {
|
||||
this.activeDocumentRebuilds.set(input.documentId, 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,
|
||||
signal,
|
||||
parentTaskId,
|
||||
retryOfTaskId
|
||||
),
|
||||
signal
|
||||
)
|
||||
}
|
||||
|
||||
private async performRebuildDocument(
|
||||
private async performRebuildDocumentMutation(
|
||||
input: KnowledgeDocumentRebuildInput,
|
||||
signal?: AbortSignal,
|
||||
parentTaskId?: string,
|
||||
@@ -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,6 +2606,9 @@ export class KnowledgeService {
|
||||
})
|
||||
tasks.push(task)
|
||||
try {
|
||||
await this.withDocumentMutation(
|
||||
document.id,
|
||||
async () => {
|
||||
this.updateKnowledgeTask(task.id, {
|
||||
status: 'running',
|
||||
progress: 10,
|
||||
@@ -2261,6 +2632,9 @@ export class KnowledgeService {
|
||||
status: 'succeeded',
|
||||
message: `已抽取 ${result.entities.length} 个实体、${result.relations.length} 条关系`
|
||||
})
|
||||
},
|
||||
effectiveSignal
|
||||
)
|
||||
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,8 +2932,13 @@ export class KnowledgeService {
|
||||
library.id,
|
||||
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,
|
||||
sourceId: source.id,
|
||||
externalId: file.relativePath,
|
||||
@@ -2568,35 +2954,22 @@ export class KnowledgeService {
|
||||
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, {
|
||||
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,12 +3339,27 @@ export class KnowledgeService {
|
||||
document: Document,
|
||||
signal?: AbortSignal
|
||||
): Promise<GraphExtractionResult> {
|
||||
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<ReplaceChunkInput, 'id' | 'content' | 'enabled' | 'role'>[],
|
||||
signal?: AbortSignal
|
||||
): Promise<GraphExtractionResult> {
|
||||
return extractKnowledgeGraph(
|
||||
chunks.map((chunk) => ({
|
||||
id: chunk.id,
|
||||
chunks
|
||||
.filter(
|
||||
(chunk) =>
|
||||
(chunk.enabled ?? true) &&
|
||||
(chunk.role ?? 'standalone') !== 'parent'
|
||||
)
|
||||
.map((chunk): GraphChunk => ({
|
||||
id: chunk.id!,
|
||||
content: chunk.content
|
||||
})),
|
||||
{
|
||||
|
||||
Reference in New Issue
Block a user