improvement(kb): pagination for docs + remove kb token count problematic query (#689)

* added pagination to docs in kb

* ack PR comments

* fix failing tests

* fix(kb): remove problematic query updating kb token count

---------

Co-authored-by: Waleed Latif <walif6@gmail.com>

Author: icecrasher321

apps/sim/app/api/knowledge/[id]/documents/route.test.ts

@@ -30,6 +30,8 @@ describe('Knowledge Base Documents API Route', () => {
     from: vi.fn().mockReturnThis(),
     where: vi.fn().mockReturnThis(),
     orderBy: vi.fn().mockReturnThis(),
+    limit: vi.fn().mockReturnThis(),
+    offset: vi.fn().mockReturnThis(),
     insert: vi.fn().mockReturnThis(),
     values: vi.fn().mockReturnThis(),
     update: vi.fn().mockReturnThis(),
@@ -99,7 +101,12 @@ describe('Knowledge Base Documents API Route', () => {
     it('should retrieve documents successfully for authenticated user', async () => {
       mockAuth$.mockAuthenticatedUser()
       mockCheckKnowledgeBaseAccess.mockResolvedValue({ hasAccess: true })
-      mockDbChain.orderBy.mockResolvedValue([mockDocument])
+
+      // Mock the count query (first query)
+      mockDbChain.where.mockResolvedValueOnce([{ count: 1 }])
+
+      // Mock the documents query (second query)
+      mockDbChain.offset.mockResolvedValue([mockDocument])
 
       const req = createMockRequest('GET')
       const { GET } = await import('./route')
@@ -108,16 +115,21 @@ describe('Knowledge Base Documents API Route', () => {
 
       expect(response.status).toBe(200)
       expect(data.success).toBe(true)
-      expect(data.data).toHaveLength(1)
-      expect(data.data[0].id).toBe('doc-123')
+      expect(data.data.documents).toHaveLength(1)
+      expect(data.data.documents[0].id).toBe('doc-123')
       expect(mockDbChain.select).toHaveBeenCalled()
       expect(mockCheckKnowledgeBaseAccess).toHaveBeenCalledWith('kb-123', 'user-123')
     })
 
     it('should filter disabled documents by default', async () => {
       mockAuth$.mockAuthenticatedUser()
       mockCheckKnowledgeBaseAccess.mockResolvedValue({ hasAccess: true })
-      mockDbChain.orderBy.mockResolvedValue([mockDocument])
+
+      // Mock the count query (first query)
+      mockDbChain.where.mockResolvedValueOnce([{ count: 1 }])
+
+      // Mock the documents query (second query)
+      mockDbChain.offset.mockResolvedValue([mockDocument])
 
       const req = createMockRequest('GET')
       const { GET } = await import('./route')
@@ -130,7 +142,12 @@ describe('Knowledge Base Documents API Route', () => {
     it('should include disabled documents when requested', async () => {
       mockAuth$.mockAuthenticatedUser()
       mockCheckKnowledgeBaseAccess.mockResolvedValue({ hasAccess: true })
-      mockDbChain.orderBy.mockResolvedValue([mockDocument])
+
+      // Mock the count query (first query)
+      mockDbChain.where.mockResolvedValueOnce([{ count: 1 }])
+
+      // Mock the documents query (second query)
+      mockDbChain.offset.mockResolvedValue([mockDocument])
 
       const url = 'http://localhost:3000/api/knowledge/kb-123/documents?includeDisabled=true'
       const req = new Request(url, { method: 'GET' }) as any

apps/sim/app/api/knowledge/[id]/documents/route.ts

@@ -1,5 +1,5 @@
 import crypto from 'node:crypto'
-import { and, desc, eq, inArray, isNull } from 'drizzle-orm'
+import { and, desc, eq, inArray, isNull, sql } from 'drizzle-orm'
 import { type NextRequest, NextResponse } from 'next/server'
 import { z } from 'zod'
 import { getSession } from '@/lib/auth'
@@ -210,6 +210,9 @@ export async function GET(req: NextRequest, { params }: { params: Promise<{ id:
 
     const url = new URL(req.url)
     const includeDisabled = url.searchParams.get('includeDisabled') === 'true'
+    const search = url.searchParams.get('search')
+    const limit = Number.parseInt(url.searchParams.get('limit') || '50')
+    const offset = Number.parseInt(url.searchParams.get('offset') || '0')
 
     // Build where conditions
     const whereConditions = [
@@ -222,6 +225,23 @@ export async function GET(req: NextRequest, { params }: { params: Promise<{ id:
       whereConditions.push(eq(document.enabled, true))
     }
 
+    // Add search condition if provided
+    if (search) {
+      whereConditions.push(
+        // Search in filename
+        sql`LOWER(${document.filename}) LIKE LOWER(${`%${search}%`})`
+      )
+    }
+
+    // Get total count for pagination
+    const totalResult = await db
+      .select({ count: sql<number>`COUNT(*)` })
+      .from(document)
+      .where(and(...whereConditions))
+
+    const total = totalResult[0]?.count || 0
+    const hasMore = offset + limit < total
+
     const documents = await db
       .select({
         id: document.id,
@@ -250,14 +270,24 @@ export async function GET(req: NextRequest, { params }: { params: Promise<{ id:
       .from(document)
       .where(and(...whereConditions))
       .orderBy(desc(document.uploadedAt))
+      .limit(limit)
+      .offset(offset)
 
     logger.info(
-      `[${requestId}] Retrieved ${documents.length} documents for knowledge base ${knowledgeBaseId}`
+      `[${requestId}] Retrieved ${documents.length} documents (${offset}-${offset + documents.length} of ${total}) for knowledge base ${knowledgeBaseId}`
     )
 
     return NextResponse.json({
       success: true,
-      data: documents,
+      data: {
+        documents,
+        pagination: {
+          total,
+          limit,
+          offset,
+          hasMore,
+        },
+      },
     })
   } catch (error) {
     logger.error(`[${requestId}] Error fetching documents`, error)

apps/sim/app/api/knowledge/utils.test.ts

@@ -176,7 +176,7 @@ describe('Knowledge Utils', () => {
         {}
       )
 
-      expect(dbOps.order).toEqual(['insert', 'updateDoc', 'updateKb'])
+      expect(dbOps.order).toEqual(['insert', 'updateDoc'])
 
       expect(dbOps.updatePayloads[0]).toMatchObject({
         processingStatus: 'completed',

apps/sim/app/api/knowledge/utils.ts

@@ -1,5 +1,5 @@
 import crypto from 'crypto'
-import { and, eq, isNull, sql } from 'drizzle-orm'
+import { and, eq, isNull } from 'drizzle-orm'
 import { processDocument } from '@/lib/documents/document-processor'
 import { retryWithExponentialBackoff } from '@/lib/documents/utils'
 import { env } from '@/lib/env'
@@ -522,14 +522,6 @@ export async function processDocumentAsync(
               processingError: null,
             })
             .where(eq(document.id, documentId))
-
-          await tx
-            .update(knowledgeBase)
-            .set({
-              tokenCount: sql`${knowledgeBase.tokenCount} + ${processed.metadata.tokenCount}`,
-              updatedAt: now,
-            })
-            .where(eq(knowledgeBase.id, knowledgeBaseId))
         })
       })(),
       TIMEOUTS.OVERALL_PROCESSING,

apps/sim/app/workspace/[workspaceId]/knowledge/[id]/[documentId]/document.tsx

@@ -117,7 +117,7 @@ export function Document({
         setError(null)
 
         const cachedDocuments = getCachedDocuments(knowledgeBaseId)
-        const cachedDoc = cachedDocuments?.find((d) => d.id === documentId)
+        const cachedDoc = cachedDocuments?.documents?.find((d) => d.id === documentId)
 
         if (cachedDoc) {
           setDocument(cachedDoc)