diff --git a/apps/sim/lib/knowledge/__integration__/bundle-import.integration.ts b/apps/sim/lib/knowledge/__integration__/bundle-import.integration.ts index 20d62af4932..259f92e7fdb 100644 --- a/apps/sim/lib/knowledge/__integration__/bundle-import.integration.ts +++ b/apps/sim/lib/knowledge/__integration__/bundle-import.integration.ts @@ -202,6 +202,7 @@ async function importedDocuments(knowledgeBaseId: string) { enabled: document.enabled, processingStatus: document.processingStatus, chunkCount: document.chunkCount, + tokenCount: document.tokenCount, tag1: document.tag1, }) .from(document) @@ -308,6 +309,7 @@ describe('knowledge base bundle import', () => { enabled: true, processingStatus: 'completed', chunkCount: 2, + tokenCount: 8, tag1: 'public', }) expect(handbook.fileUrl).toMatch(/^\/api\/files\/serve\/kb\/.+\?context=knowledge-base$/) @@ -328,6 +330,7 @@ describe('knowledge base bundle import', () => { enabled: false, processingStatus: 'completed', chunkCount: 1, + tokenCount: 4, }) expect( ( @@ -380,7 +383,7 @@ describe('knowledge base bundle import', () => { ) ) expect(queued).toHaveLength(1) - expect(notes.processingStatus).toBe('completed') + expect(notes).toMatchObject({ processingStatus: 'completed', tokenCount: 4 }) expect(fixture.embedded).toBe(1) expect(await importedChunks(notes.id)).toEqual([ { diff --git a/apps/sim/lib/knowledge/documents/service.ts b/apps/sim/lib/knowledge/documents/service.ts index e67cef63585..799de16ee8f 100644 --- a/apps/sim/lib/knowledge/documents/service.ts +++ b/apps/sim/lib/knowledge/documents/service.ts @@ -1459,12 +1459,17 @@ export interface PreparedDocumentIndex { readonly vectors?: readonly number[][] /** Characters in the document's source text; chunks overlap, so their lengths do not sum to it. */ readonly characterCount: number + /** + * The document total its chunker estimated for these chunks. Chunkers estimate + * differently from the embedding tokenizer the pass applies to each chunk, so + * summing per-chunk counts would not match the source document. + */ + readonly tokenCount: number } /** - * The processor's result shape for chunks that were never parsed here. Token - * counts are left to the pass, which estimates every chunk against this - * base's model whichever way the chunks arrived. + * The processor's result shape for chunks that were never parsed here. The + * pass still estimates each chunk against this base's model. */ function preparedProcessingResult(prepared: PreparedDocumentIndex) { return { @@ -1474,7 +1479,7 @@ function preparedProcessingResult(prepared: PreparedDocumentIndex) { })), metadata: { chunkCount: prepared.chunks.length, - tokenCount: null, + tokenCount: prepared.tokenCount, characterCount: prepared.characterCount, }, } @@ -2037,7 +2042,9 @@ export async function processDocumentAsync( const now = new Date() logger.info( - `[${documentId}] Document parsed successfully, generating embeddings for ${processed.chunks.length} chunks` + prepared?.vectors + ? `[${documentId}] Indexing ${processed.chunks.length} chunks with their bundled vectors` + : `[${documentId}] Document parsed successfully, generating embeddings for ${processed.chunks.length} chunks` ) const chunkTexts = processed.chunks.map((chunk) => chunk.text) @@ -2297,9 +2304,7 @@ export async function processDocumentAsync( .update(document) .set({ chunkCount: processed.metadata.chunkCount, - tokenCount: - processed.metadata.tokenCount ?? - chunkTokenCounts.reduce((sum, count) => sum + count, 0), + tokenCount: processed.metadata.tokenCount, characterCount: processed.metadata.characterCount, processingStatus: 'completed', processingCompletedAt: now, diff --git a/apps/sim/lib/knowledge/transfer/bundle.ts b/apps/sim/lib/knowledge/transfer/bundle.ts index 33a8faff334..abf62ab160e 100644 --- a/apps/sim/lib/knowledge/transfer/bundle.ts +++ b/apps/sim/lib/knowledge/transfer/bundle.ts @@ -57,6 +57,9 @@ const MAX_BUNDLE_TEXT_LENGTH = 255 */ const MAX_BUNDLE_TAG_VALUE_LENGTH = KNOWLEDGE_BASE_DESCRIPTION_MAX_LENGTH +/** Largest value the `integer` count columns an import writes can hold. */ +const MAX_INTEGER_COLUMN = 2_147_483_647 + const tagSlotSchema = z.enum(ALL_TAG_SLOTS) export const knowledgeBundleTagSchema = z @@ -81,8 +84,8 @@ export const knowledgeBundleDocumentSchema = z file: z.string().nullable(), chunks: z.string().nullable(), chunkCount: z.number().int().min(0).max(MAX_DOCUMENT_CHUNKS), - tokenCount: z.number().int().min(0), - characterCount: z.number().int().min(0), + tokenCount: z.number().int().min(0).max(MAX_INTEGER_COLUMN), + characterCount: z.number().int().min(0).max(MAX_INTEGER_COLUMN), }) .strict() .superRefine((document, ctx) => { diff --git a/apps/sim/lib/knowledge/transfer/import-archive.test.ts b/apps/sim/lib/knowledge/transfer/import-archive.test.ts index 373a391f79c..9fd7e29ef40 100644 --- a/apps/sim/lib/knowledge/transfer/import-archive.test.ts +++ b/apps/sim/lib/knowledge/transfer/import-archive.test.ts @@ -51,7 +51,12 @@ function manifest(overrides: Record = {}, document: Record = {}): string { return JSON.stringify({ @@ -139,6 +144,11 @@ describe('openKnowledgeBundle', () => { }, 'duplicate document id doc-1', ], + [ + 'a document count no integer column can hold', + () => zipOf(validEntries(undefined, manifest({}, { tokenCount: 2_147_483_648 }))), + 'manifest documents.0.tokenCount', + ], [ 'a manifest naming an entry the archive lacks', () => { diff --git a/apps/sim/lib/knowledge/transfer/import-archive.ts b/apps/sim/lib/knowledge/transfer/import-archive.ts index 2a3900edbf7..f3fe065005c 100644 --- a/apps/sim/lib/knowledge/transfer/import-archive.ts +++ b/apps/sim/lib/knowledge/transfer/import-archive.ts @@ -44,7 +44,9 @@ export interface KnowledgeBundleReader { * line must carry a vector of the manifest's dimension; without, none is decoded. */ readChunks( - document: Pick & { chunks: string }, + document: Pick & { + chunks: string + }, vectors: boolean ): Promise close(): void @@ -231,8 +233,10 @@ export async function openKnowledgeBundle(path: string): Promise zip.close(), } diff --git a/apps/sim/lib/knowledge/transfer/import-runner.ts b/apps/sim/lib/knowledge/transfer/import-runner.ts index 6c8e707a108..023f2d8cc6f 100644 --- a/apps/sim/lib/knowledge/transfer/import-runner.ts +++ b/apps/sim/lib/knowledge/transfer/import-runner.ts @@ -153,6 +153,7 @@ async function importDocument( chunks: document.chunks, chunkCount: document.chunkCount, characterCount: document.characterCount, + tokenCount: document.tokenCount, }, reuseVectors )