improvement(kb): added configurability for chunks, query across multiple knowledge bases (#512)

* refactor: consolidate create modal file

* fix: identify dead processes

* fix: mark failed in DB after processing timeout

* improvement: added overlap chunks and fixed modal UI

* feat: multiselect logic

* fix: biome changes for css ordering warn instead of error

* improvement: create chunk ui

* fix: removed unused schema columns

* fix: removed references to deleted columns

* improvement: sped up vector search time

* feat: multi-kb search

* add bulk endpoint to disable/delete multiple chunks

* add bulk endpoint to disable/delete multiple chunks

* fix: removed unused schema columns

* fix: removed references to deleted columns

* made endpoints for knowledge more RESTful, added tests

* added batch operations for delete/enable/disable docs, alr have this for chunks

* added migrations

* added migrations

---------

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

Author: emir-karabeg

apps/sim/app/api/__test-utils__/utils.ts

@@ -1,7 +1,6 @@
 import { NextRequest } from 'next/server'
 import { vi } from 'vitest'
 
-// Add type definitions for better type safety
 export interface MockUser {
   id: string
   email: string
@@ -14,7 +13,6 @@ export interface MockAuthResult {
   mockUnauthenticated: () => void
 }
 
-// Database result types
 export interface DatabaseSelectResult {
   id: string
   [key: string]: any
@@ -234,7 +232,6 @@ export function createMockRequest(
 ): NextRequest {
   const url = 'http://localhost:3000/api/test'
 
-  // Use the URL constructor to create a proper URL object
   return new NextRequest(new URL(url), {
     method,
     headers: new Headers(headers),
@@ -248,7 +245,6 @@ export function mockExecutionDependencies() {
     return {
       ...(actual as any),
       decryptSecret: vi.fn().mockImplementation((encrypted: string) => {
-        // Map from encrypted to decrypted
         const entries = Object.entries(mockEnvironmentVars)
         const found = entries.find(([_, val]) => val === encrypted)
         const key = found ? found[0] : null
@@ -570,6 +566,7 @@ export function mockDrizzleOrm() {
     asc: vi.fn((field) => ({ field, type: 'asc' })),
     desc: vi.fn((field) => ({ field, type: 'desc' })),
     isNull: vi.fn((field) => ({ field, type: 'isNull' })),
+    count: vi.fn((field) => ({ field, type: 'count' })),
     sql: vi.fn((strings, ...values) => ({
       type: 'sql',
       sql: strings,
@@ -578,6 +575,57 @@ export function mockDrizzleOrm() {
   }))
 }
 
+/**
+ * Mock knowledge-related database schemas
+ */
+export function mockKnowledgeSchemas() {
+  vi.doMock('@/db/schema', () => ({
+    knowledgeBase: {
+      id: 'kb_id',
+      userId: 'user_id',
+      name: 'kb_name',
+      description: 'description',
+      tokenCount: 'token_count',
+      embeddingModel: 'embedding_model',
+      embeddingDimension: 'embedding_dimension',
+      chunkingConfig: 'chunking_config',
+      workspaceId: 'workspace_id',
+      createdAt: 'created_at',
+      updatedAt: 'updated_at',
+      deletedAt: 'deleted_at',
+    },
+    document: {
+      id: 'doc_id',
+      knowledgeBaseId: 'kb_id',
+      filename: 'filename',
+      fileUrl: 'file_url',
+      fileSize: 'file_size',
+      mimeType: 'mime_type',
+      chunkCount: 'chunk_count',
+      tokenCount: 'token_count',
+      characterCount: 'character_count',
+      processingStatus: 'processing_status',
+      processingStartedAt: 'processing_started_at',
+      processingCompletedAt: 'processing_completed_at',
+      processingError: 'processing_error',
+      enabled: 'enabled',
+      uploadedAt: 'uploaded_at',
+      deletedAt: 'deleted_at',
+    },
+    embedding: {
+      id: 'embedding_id',
+      documentId: 'doc_id',
+      knowledgeBaseId: 'kb_id',
+      chunkIndex: 'chunk_index',
+      content: 'content',
+      embedding: 'embedding',
+      tokenCount: 'token_count',
+      characterCount: 'character_count',
+      createdAt: 'created_at',
+    },
+  }))
+}
+
 /**
  * Mock console logger
  */

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

@@ -1,3 +1,4 @@
+import crypto from 'node:crypto'
 import { eq, sql } from 'drizzle-orm'
 import { type NextRequest, NextResponse } from 'next/server'
 import { z } from 'zod'
@@ -12,8 +13,6 @@ const logger = createLogger('ChunkByIdAPI')
 const UpdateChunkSchema = z.object({
   content: z.string().min(1, 'Content is required').optional(),
   enabled: z.boolean().optional(),
-  searchRank: z.number().min(0).optional(),
-  qualityScore: z.number().min(0).max(1).optional(),
 })
 
 export async function GET(
@@ -103,21 +102,27 @@ export async function PUT(
     try {
       const validatedData = UpdateChunkSchema.parse(body)
 
-      const updateData: any = {
-        updatedAt: new Date(),
-      }
+      const updateData: Partial<{
+        content: string
+        contentLength: number
+        tokenCount: number
+        chunkHash: string
+        enabled: boolean
+        updatedAt: Date
+      }> = {}
 
-      if (validatedData.content !== undefined) {
+      if (validatedData.content) {
         updateData.content = validatedData.content
         updateData.contentLength = validatedData.content.length
         // Update token count estimation (rough approximation: 4 chars per token)
         updateData.tokenCount = Math.ceil(validatedData.content.length / 4)
+        updateData.chunkHash = crypto
+          .createHash('sha256')
+          .update(validatedData.content)
+          .digest('hex')
       }
+
       if (validatedData.enabled !== undefined) updateData.enabled = validatedData.enabled
-      if (validatedData.searchRank !== undefined)
-        updateData.searchRank = validatedData.searchRank.toString()
-      if (validatedData.qualityScore !== undefined)
-        updateData.qualityScore = validatedData.qualityScore.toString()
 
       await db.update(embedding).set(updateData).where(eq(embedding.id, chunkId))
 

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

@@ -1,5 +1,5 @@
 import crypto from 'crypto'
-import { and, asc, eq, ilike, sql } from 'drizzle-orm'
+import { and, asc, eq, ilike, inArray, sql } from 'drizzle-orm'
 import { type NextRequest, NextResponse } from 'next/server'
 import { z } from 'zod'
 import { getSession } from '@/lib/auth'
@@ -11,20 +11,26 @@ import { checkDocumentAccess, generateEmbeddings } from '../../../../utils'
 
 const logger = createLogger('DocumentChunksAPI')
 
-// Schema for query parameters
 const GetChunksQuerySchema = z.object({
   search: z.string().optional(),
   enabled: z.enum(['true', 'false', 'all']).optional().default('all'),
   limit: z.coerce.number().min(1).max(100).optional().default(50),
   offset: z.coerce.number().min(0).optional().default(0),
 })
 
-// Schema for creating manual chunks
 const CreateChunkSchema = z.object({
   content: z.string().min(1, 'Content is required').max(10000, 'Content too long'),
   enabled: z.boolean().optional().default(true),
 })
 
+const BatchOperationSchema = z.object({
+  operation: z.enum(['enable', 'disable', 'delete']),
+  chunkIds: z
+    .array(z.string())
+    .min(1, 'At least one chunk ID is required')
+    .max(100, 'Cannot operate on more than 100 chunks at once'),
+})
+
 export async function GET(
   req: NextRequest,
   { params }: { params: Promise<{ id: string; documentId: string }> }
@@ -112,10 +118,7 @@ export async function GET(
         enabled: embedding.enabled,
         startOffset: embedding.startOffset,
         endOffset: embedding.endOffset,
-        overlapTokens: embedding.overlapTokens,
         metadata: embedding.metadata,
-        searchRank: embedding.searchRank,
-        qualityScore: embedding.qualityScore,
         createdAt: embedding.createdAt,
         updatedAt: embedding.updatedAt,
       })
@@ -236,12 +239,7 @@ export async function POST(
           embeddingModel: 'text-embedding-3-small',
           startOffset: 0, // Manual chunks don't have document offsets
           endOffset: validatedData.content.length,
-          overlapTokens: 0,
           metadata: { manual: true }, // Mark as manually created
-          searchRank: '1.0',
-          accessCount: 0,
-          lastAccessedAt: null,
-          qualityScore: null,
           enabled: validatedData.enabled,
           createdAt: now,
           updatedAt: now,
@@ -286,3 +284,144 @@ export async function POST(
     return NextResponse.json({ error: 'Failed to create chunk' }, { status: 500 })
   }
 }
+
+export async function PATCH(
+  req: NextRequest,
+  { params }: { params: Promise<{ id: string; documentId: string }> }
+) {
+  const requestId = crypto.randomUUID().slice(0, 8)
+  const { id: knowledgeBaseId, documentId } = await params
+
+  try {
+    const session = await getSession()
+    if (!session?.user?.id) {
+      logger.warn(`[${requestId}] Unauthorized batch chunk operation attempt`)
+      return NextResponse.json({ error: 'Unauthorized' }, { status: 401 })
+    }
+
+    const accessCheck = await checkDocumentAccess(knowledgeBaseId, documentId, session.user.id)
+
+    if (!accessCheck.hasAccess) {
+      if (accessCheck.notFound) {
+        logger.warn(
+          `[${requestId}] ${accessCheck.reason}: KB=${knowledgeBaseId}, Doc=${documentId}`
+        )
+        return NextResponse.json({ error: accessCheck.reason }, { status: 404 })
+      }
+      logger.warn(
+        `[${requestId}] User ${session.user.id} attempted unauthorized batch chunk operation: ${accessCheck.reason}`
+      )
+      return NextResponse.json({ error: 'Unauthorized' }, { status: 401 })
+    }
+
+    const body = await req.json()
+
+    try {
+      const validatedData = BatchOperationSchema.parse(body)
+      const { operation, chunkIds } = validatedData
+
+      logger.info(
+        `[${requestId}] Starting batch ${operation} operation on ${chunkIds.length} chunks for document ${documentId}`
+      )
+
+      const results = []
+      let successCount = 0
+      const errorCount = 0
+
+      if (operation === 'delete') {
+        // Handle batch delete with transaction for consistency
+        await db.transaction(async (tx) => {
+          // Get chunks to delete for statistics update
+          const chunksToDelete = await tx
+            .select({
+              id: embedding.id,
+              tokenCount: embedding.tokenCount,
+              contentLength: embedding.contentLength,
+            })
+            .from(embedding)
+            .where(and(eq(embedding.documentId, documentId), inArray(embedding.id, chunkIds)))
+
+          if (chunksToDelete.length === 0) {
+            throw new Error('No valid chunks found to delete')
+          }
+
+          // Delete chunks
+          await tx
+            .delete(embedding)
+            .where(and(eq(embedding.documentId, documentId), inArray(embedding.id, chunkIds)))
+
+          // Update document statistics
+          const totalTokens = chunksToDelete.reduce((sum, chunk) => sum + chunk.tokenCount, 0)
+          const totalCharacters = chunksToDelete.reduce(
+            (sum, chunk) => sum + chunk.contentLength,
+            0
+          )
+
+          await tx
+            .update(document)
+            .set({
+              chunkCount: sql`${document.chunkCount} - ${chunksToDelete.length}`,
+              tokenCount: sql`${document.tokenCount} - ${totalTokens}`,
+              characterCount: sql`${document.characterCount} - ${totalCharacters}`,
+            })
+            .where(eq(document.id, documentId))
+
+          successCount = chunksToDelete.length
+          results.push({
+            operation: 'delete',
+            deletedCount: chunksToDelete.length,
+            chunkIds: chunksToDelete.map((c) => c.id),
+          })
+        })
+      } else {
+        // Handle batch enable/disable
+        const enabled = operation === 'enable'
+
+        // Update chunks in a single query
+        const updateResult = await db
+          .update(embedding)
+          .set({
+            enabled,
+            updatedAt: new Date(),
+          })
+          .where(and(eq(embedding.documentId, documentId), inArray(embedding.id, chunkIds)))
+          .returning({ id: embedding.id })
+
+        successCount = updateResult.length
+        results.push({
+          operation,
+          updatedCount: updateResult.length,
+          chunkIds: updateResult.map((r) => r.id),
+        })
+      }
+
+      logger.info(
+        `[${requestId}] Batch ${operation} operation completed: ${successCount} successful, ${errorCount} errors`
+      )
+
+      return NextResponse.json({
+        success: true,
+        data: {
+          operation,
+          successCount,
+          errorCount,
+          results,
+        },
+      })
+    } catch (validationError) {
+      if (validationError instanceof z.ZodError) {
+        logger.warn(`[${requestId}] Invalid batch operation data`, {
+          errors: validationError.errors,
+        })
+        return NextResponse.json(
+          { error: 'Invalid request data', details: validationError.errors },
+          { status: 400 }
+        )
+      }
+      throw validationError
+    }
+  } catch (error) {
+    logger.error(`[${requestId}] Error in batch chunk operation`, error)
+    return NextResponse.json({ error: 'Failed to perform batch operation' }, { status: 500 })
+  }
+}