feat(knowledge): add embedding model selection and Cohere reranker

Author: waleedlatif1

apps/docs/content/docs/en/tools/knowledge.mdx

@@ -47,6 +47,8 @@ Search for similar content in a knowledge base using vector similarity
 | `properties` | string | No | No description |
 | `tagName` | string | No | No description |
 | `tagValue` | string | No | No description |
+| `rerankerEnabled` | boolean | No | Whether to apply Cohere reranking to vector search results |
+| `rerankerModel` | string | No | Cohere rerank model to use \(one of: $\{SUPPORTED_RERANKER_MODELS.join\(', '\)\}\) |
 | `tagFilters` | string | No | No description |
 
 #### Output

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

@@ -213,9 +213,11 @@ export const POST = withRouteHandler(
           accessCheck.knowledgeBase?.workspaceId
         )
 
+        const chunkEmbeddingModel =
+          accessCheck.knowledgeBase?.embeddingModel ?? 'text-embedding-3-small'
         let cost = null
         try {
-          cost = calculateCost('text-embedding-3-small', newChunk.tokenCount, 0, false)
+          cost = calculateCost(chunkEmbeddingModel, newChunk.tokenCount, 0, false)
         } catch (error) {
           logger.warn(`[${requestId}] Failed to calculate cost for chunk upload`, {
             error: error instanceof Error ? error.message : 'Unknown error',
@@ -240,7 +242,7 @@ export const POST = withRouteHandler(
                       completion: 0,
                       total: newChunk.tokenCount,
                     },
-                    model: 'text-embedding-3-small',
+                    model: chunkEmbeddingModel,
                     pricing: cost.pricing,
                   },
                 }

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

@@ -6,6 +6,11 @@ import { getSession } from '@/lib/auth'
 import { PlatformEvents } from '@/lib/core/telemetry'
 import { generateRequestId } from '@/lib/core/utils/request'
 import { withRouteHandler } from '@/lib/core/utils/with-route-handler'
+import {
+  DEFAULT_EMBEDDING_MODEL,
+  EMBEDDING_DIMENSIONS,
+  SUPPORTED_EMBEDDING_MODEL_IDS,
+} from '@/lib/knowledge/embeddings'
 import {
   createKnowledgeBase,
   getKnowledgeBases,
@@ -20,8 +25,10 @@ const CreateKnowledgeBaseSchema = z.object({
   name: z.string().min(1, 'Name is required'),
   description: z.string().optional(),
   workspaceId: z.string().min(1, 'Workspace ID is required'),
-  embeddingModel: z.literal('text-embedding-3-small').default('text-embedding-3-small'),
-  embeddingDimension: z.literal(1536).default(1536),
+  embeddingModel: z
+    .enum(SUPPORTED_EMBEDDING_MODEL_IDS as [string, ...string[]])
+    .default(DEFAULT_EMBEDDING_MODEL),
+  embeddingDimension: z.literal(EMBEDDING_DIMENSIONS).default(EMBEDDING_DIMENSIONS),
   chunkingConfig: z
     .object({
       maxSize: z.number().min(100).max(4000).default(1024),

apps/sim/app/api/knowledge/search/route.ts

@@ -7,6 +7,7 @@ import { PlatformEvents } from '@/lib/core/telemetry'
 import { generateRequestId } from '@/lib/core/utils/request'
 import { withRouteHandler } from '@/lib/core/utils/with-route-handler'
 import { ALL_TAG_SLOTS } from '@/lib/knowledge/constants'
+import { DEFAULT_RERANKER_MODEL, rerank, SUPPORTED_RERANKER_MODELS } from '@/lib/knowledge/reranker'
 import { getDocumentTagDefinitions } from '@/lib/knowledge/tags/service'
 import { buildUndefinedTagsError, validateTagValue } from '@/lib/knowledge/tags/utils'
 import type { StructuredFilter } from '@/lib/knowledge/types'
@@ -21,6 +22,7 @@ import {
   type SearchResult,
 } from '@/app/api/knowledge/search/utils'
 import { checkKnowledgeBaseAccess } from '@/app/api/knowledge/utils'
+import { getRerankModelPricing } from '@/providers/models'
 import { calculateCost } from '@/providers/utils'
 
 const logger = createLogger('VectorSearchAPI')
@@ -59,6 +61,11 @@ const VectorSearchSchema = z
       .optional()
       .nullable()
       .transform((val) => val || undefined),
+    rerankerEnabled: z.boolean().optional().default(false),
+    rerankerModel: z
+      .enum(SUPPORTED_RERANKER_MODELS as unknown as [string, ...string[]])
+      .optional()
+      .default(DEFAULT_RERANKER_MODEL),
   })
   .refine(
     (data) => {
@@ -235,11 +242,40 @@ export const POST = withRouteHandler(async (request: NextRequest) => {
         )
       }
 
-      const workspaceId = accessChecks.find((ac) => ac?.hasAccess)?.knowledgeBase?.workspaceId
+      const accessibleKbs = accessChecks
+        .filter(
+          (
+            ac
+          ): ac is {
+            hasAccess: true
+            knowledgeBase: {
+              id: string
+              embeddingModel: string
+              workspaceId?: string | null
+            }
+          } => Boolean(ac?.hasAccess)
+        )
+        .map((ac) => ac.knowledgeBase)
+      const workspaceId = accessibleKbs[0]?.workspaceId
+
+      const useReranker = validatedData.rerankerEnabled && Boolean(validatedData.query?.trim())
+      const rerankerModel = useReranker ? validatedData.rerankerModel : null
+
+      const embeddingModels = Array.from(new Set(accessibleKbs.map((kb) => kb.embeddingModel)))
+      if (embeddingModels.length > 1) {
+        return NextResponse.json(
+          {
+            error:
+              'Selected knowledge bases use different embedding models and cannot be searched together. Search them separately.',
+          },
+          { status: 400 }
+        )
+      }
+      const queryEmbeddingModel = embeddingModels[0]
 
       const hasQuery = validatedData.query && validatedData.query.trim().length > 0
       const queryEmbeddingPromise = hasQuery
-        ? generateSearchEmbedding(validatedData.query!, undefined, workspaceId)
+        ? generateSearchEmbedding(validatedData.query!, queryEmbeddingModel, workspaceId)
         : Promise.resolve(null)
 
       // Check if any requested knowledge bases were not accessible
@@ -278,6 +314,10 @@ export const POST = withRouteHandler(async (request: NextRequest) => {
 
       const hasFilters = structuredFilters && structuredFilters.length > 0
 
+      // Oversample candidates when reranking so the reranker has more to choose from.
+      // Cap at 100 to bound Cohere request cost (1 search unit = ≤100 docs).
+      const candidateTopK = useReranker ? Math.min(100, validatedData.topK * 4) : validatedData.topK
+
       if (!hasQuery && hasFilters) {
         // Tag-only search without vector similarity
         results = await handleTagOnlySearch({
@@ -291,24 +331,24 @@ export const POST = withRouteHandler(async (request: NextRequest) => {
           `[${requestId}] Executing tag + vector search with filters:`,
           structuredFilters
         )
-        const strategy = getQueryStrategy(accessibleKbIds.length, validatedData.topK)
+        const strategy = getQueryStrategy(accessibleKbIds.length, candidateTopK)
         const queryVector = JSON.stringify(await queryEmbeddingPromise)
 
         results = await handleTagAndVectorSearch({
           knowledgeBaseIds: accessibleKbIds,
-          topK: validatedData.topK,
+          topK: candidateTopK,
           structuredFilters,
           queryVector,
           distanceThreshold: strategy.distanceThreshold,
         })
       } else if (hasQuery && !hasFilters) {
         // Vector-only search
-        const strategy = getQueryStrategy(accessibleKbIds.length, validatedData.topK)
+        const strategy = getQueryStrategy(accessibleKbIds.length, candidateTopK)
         const queryVector = JSON.stringify(await queryEmbeddingPromise)
 
         results = await handleVectorOnlySearch({
           knowledgeBaseIds: accessibleKbIds,
-          topK: validatedData.topK,
+          topK: candidateTopK,
           queryVector,
           distanceThreshold: strategy.distanceThreshold,
         })
@@ -323,13 +363,54 @@ export const POST = withRouteHandler(async (request: NextRequest) => {
         )
       }
 
+      // Optional Cohere rerank pass on top of vector results.
+      const rerankedScores = new Map<string, number>()
+      let rerankApplied = false
+      if (useReranker && rerankerModel && results.length > 0) {
+        const candidateCount = results.length
+        try {
+          const ranked = await rerank(
+            validatedData.query!,
+            results.map((r) => ({ id: r.id, text: r.content })),
+            { model: rerankerModel, topN: validatedData.topK, workspaceId }
+          )
+          if (ranked.length === 0) {
+            logger.warn(
+              `[${requestId}] Reranker returned 0 results; falling back to vector ordering`,
+              { model: rerankerModel, candidateCount }
+            )
+            results = results.slice(0, validatedData.topK)
+          } else {
+            const idToResult = new Map(results.map((r) => [r.id, r]))
+            results = ranked
+              .map((r) => idToResult.get(r.item.id))
+              .filter((r): r is SearchResult => Boolean(r))
+            for (const r of ranked) rerankedScores.set(r.item.id, r.relevanceScore)
+            rerankApplied = true
+            logger.info(`[${requestId}] Reranked ${candidateCount} → ${results.length} results`, {
+              model: rerankerModel,
+            })
+          }
+        } catch (error) {
+          logger.warn(`[${requestId}] Reranker failed; falling back to vector ordering`, {
+            error: error instanceof Error ? error.message : 'Unknown error',
+            model: rerankerModel,
+            candidateCount,
+            workspaceId,
+          })
+          results = results.slice(0, validatedData.topK)
+        }
+      } else if (useReranker) {
+        results = results.slice(0, validatedData.topK)
+      }
+
       // Calculate cost for the embedding (with fallback if calculation fails)
       let cost = null
       let tokenCount = null
       if (hasQuery) {
         try {
           tokenCount = estimateTokenCount(validatedData.query!, 'openai')
-          cost = calculateCost('text-embedding-3-small', tokenCount.count, 0, false)
+          cost = calculateCost(queryEmbeddingModel, tokenCount.count, 0, false)
         } catch (error) {
           logger.warn(`[${requestId}] Failed to calculate cost for search query`, {
             error: error instanceof Error ? error.message : 'Unknown error',
@@ -338,6 +419,31 @@ export const POST = withRouteHandler(async (request: NextRequest) => {
         }
       }
 
+      // Add Cohere rerank cost (1 search unit per call, since we cap candidates ≤100).
+      let rerankerCost = 0
+      if (rerankApplied && rerankerModel) {
+        const pricing = getRerankModelPricing(rerankerModel)
+        if (pricing) {
+          rerankerCost = pricing.perSearchUnit
+          if (cost) {
+            cost = {
+              ...cost,
+              input: cost.input + rerankerCost,
+              total: cost.total + rerankerCost,
+            }
+          } else {
+            cost = {
+              input: rerankerCost,
+              output: 0,
+              total: rerankerCost,
+              pricing: { input: 0, output: 0, updatedAt: pricing.updatedAt },
+            }
+          }
+        } else {
+          logger.warn(`[${requestId}] No pricing entry for rerank model ${rerankerModel}`)
+        }
+      }
+
       // Fetch tag definitions for display name mapping (reuse the same fetch from filtering)
       const tagDefsResults = await Promise.all(
         accessibleKbIds.map(async (kbId) => {
@@ -400,33 +506,36 @@ export const POST = withRouteHandler(async (request: NextRequest) => {
               }
             })
 
+            const rerankerScore = rerankedScores.get(result.id)
             return {
               documentId: result.documentId,
               documentName: documentNameMap[result.documentId] || undefined,
               content: result.content,
               chunkIndex: result.chunkIndex,
               metadata: tags, // Clean display name mapped tags
               similarity: hasQuery ? 1 - result.distance : 1, // Perfect similarity for tag-only searches
+              ...(rerankerScore !== undefined && { rerankerScore }),
             }
           }),
           query: validatedData.query || '',
           knowledgeBaseIds: accessibleKbIds,
           knowledgeBaseId: accessibleKbIds[0],
           topK: validatedData.topK,
           totalResults: results.length,
-          ...(cost && tokenCount
+          ...(cost
             ? {
                 cost: {
                   input: cost.input,
                   output: cost.output,
                   total: cost.total,
                   tokens: {
-                    prompt: tokenCount.count,
+                    prompt: tokenCount?.count ?? 0,
                     completion: 0,
-                    total: tokenCount.count,
+                    total: tokenCount?.count ?? 0,
                   },
-                  model: 'text-embedding-3-small',
+                  model: queryEmbeddingModel,
                   pricing: cost.pricing,
+                  ...(rerankApplied ? { rerankerCost, rerankerModel, rerankerSearchUnits: 1 } : {}),
                 },
               }
             : {}),

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

@@ -103,7 +103,10 @@ export interface EmbeddingData {
 
 export interface KnowledgeBaseAccessResult {
   hasAccess: true
-  knowledgeBase: Pick<KnowledgeBaseData, 'id' | 'userId' | 'workspaceId' | 'name'>
+  knowledgeBase: Pick<
+    KnowledgeBaseData,
+    'id' | 'userId' | 'workspaceId' | 'name' | 'embeddingModel'
+  >
 }
 
 export interface KnowledgeBaseAccessDenied {
@@ -117,7 +120,10 @@ export type KnowledgeBaseAccessCheck = KnowledgeBaseAccessResult | KnowledgeBase
 export interface DocumentAccessResult {
   hasAccess: true
   document: DocumentData
-  knowledgeBase: Pick<KnowledgeBaseData, 'id' | 'userId' | 'workspaceId' | 'name'>
+  knowledgeBase: Pick<
+    KnowledgeBaseData,
+    'id' | 'userId' | 'workspaceId' | 'name' | 'embeddingModel'
+  >
 }
 
 export interface DocumentAccessDenied {
@@ -132,7 +138,10 @@ export interface ChunkAccessResult {
   hasAccess: true
   chunk: EmbeddingData
   document: DocumentData
-  knowledgeBase: Pick<KnowledgeBaseData, 'id' | 'userId' | 'workspaceId' | 'name'>
+  knowledgeBase: Pick<
+    KnowledgeBaseData,
+    'id' | 'userId' | 'workspaceId' | 'name' | 'embeddingModel'
+  >
 }
 
 export interface ChunkAccessDenied {
@@ -156,6 +165,7 @@ export async function checkKnowledgeBaseAccess(
       userId: knowledgeBase.userId,
       workspaceId: knowledgeBase.workspaceId,
       name: knowledgeBase.name,
+      embeddingModel: knowledgeBase.embeddingModel,
     })
     .from(knowledgeBase)
     .where(and(eq(knowledgeBase.id, knowledgeBaseId), isNull(knowledgeBase.deletedAt)))
@@ -200,6 +210,7 @@ export async function checkKnowledgeBaseWriteAccess(
       userId: knowledgeBase.userId,
       workspaceId: knowledgeBase.workspaceId,
       name: knowledgeBase.name,
+      embeddingModel: knowledgeBase.embeddingModel,
     })
     .from(knowledgeBase)
     .where(and(eq(knowledgeBase.id, knowledgeBaseId), isNull(knowledgeBase.deletedAt)))

apps/sim/app/api/v1/knowledge/route.ts

@@ -2,6 +2,11 @@ import { AuditAction, AuditResourceType, recordAudit } from '@sim/audit'
 import { type NextRequest, NextResponse } from 'next/server'
 import { z } from 'zod'
 import { withRouteHandler } from '@/lib/core/utils/with-route-handler'
+import {
+  DEFAULT_EMBEDDING_MODEL,
+  EMBEDDING_DIMENSIONS,
+  SUPPORTED_EMBEDDING_MODEL_IDS,
+} from '@/lib/knowledge/embeddings'
 import { createKnowledgeBase, getKnowledgeBases } from '@/lib/knowledge/service'
 import {
   authenticateRequest,
@@ -29,6 +34,9 @@ const CreateKBSchema = z.object({
   workspaceId: z.string().min(1, 'Workspace ID is required'),
   name: z.string().min(1, 'Name is required').max(255, 'Name must be 255 characters or less'),
   description: z.string().max(1000, 'Description must be 1000 characters or less').optional(),
+  embeddingModel: z
+    .enum(SUPPORTED_EMBEDDING_MODEL_IDS as [string, ...string[]])
+    .default(DEFAULT_EMBEDDING_MODEL),
   chunkingConfig: ChunkingConfigSchema.optional().default({
     maxSize: 1024,
     minSize: 100,
@@ -81,7 +89,7 @@ export const POST = withRouteHandler(async (request: NextRequest) => {
     const validation = validateSchema(CreateKBSchema, body.data)
     if (!validation.success) return validation.response
 
-    const { workspaceId, name, description, chunkingConfig } = validation.data
+    const { workspaceId, name, description, embeddingModel, chunkingConfig } = validation.data
 
     const accessError = await validateWorkspaceAccess(rateLimit, userId, workspaceId, 'write')
     if (accessError) return accessError
@@ -92,8 +100,8 @@ export const POST = withRouteHandler(async (request: NextRequest) => {
         description,
         workspaceId,
         userId,
-        embeddingModel: 'text-embedding-3-small',
-        embeddingDimension: 1536,
+        embeddingModel,
+        embeddingDimension: EMBEDDING_DIMENSIONS,
         chunkingConfig: chunkingConfig ?? { maxSize: 1024, minSize: 100, overlap: 200 },
       },
       requestId