fix(knowledge): keep a bundle's document token totals on import

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([
       {

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,

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) => {

apps/sim/lib/knowledge/transfer/import-archive.test.ts

@@ -51,7 +51,12 @@ function manifest(overrides: Record<string, unknown> = {}, document: Record<stri
   }
 }
 
-const CHUNKS = { chunks: `chunks/${DOCUMENT_ID}.ndjson`, chunkCount: 2, characterCount: 10 }
+const CHUNKS = {
+  chunks: `chunks/${DOCUMENT_ID}.ndjson`,
+  chunkCount: 2,
+  characterCount: 10,
+  tokenCount: 4,
+}
 
 function chunkLine(index: number, extra: Record<string, unknown> = {}): 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',
       () => {

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<KnowledgeBundleDocument, 'chunkCount' | 'characterCount'> & { chunks: string },
+    document: Pick<KnowledgeBundleDocument, 'chunkCount' | 'characterCount' | 'tokenCount'> & {
+      chunks: string
+    },
     vectors: boolean
   ): Promise<PreparedDocumentIndex>
   close(): void
@@ -231,8 +233,10 @@ export async function openKnowledgeBundle(path: string): Promise<KnowledgeBundle
           chunks.push(chunk)
           if (vector) decoded.push(vector)
         })
-        const { characterCount } = document
-        return vectors ? { chunks, vectors: decoded, characterCount } : { chunks, characterCount }
+        const { characterCount, tokenCount } = document
+        return vectors
+          ? { chunks, vectors: decoded, characterCount, tokenCount }
+          : { chunks, characterCount, tokenCount }
       },
       close: () => zip.close(),
     }

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
         )