feat(knowledge): expose Cohere reranker controls on knowledge block

Add a self-hosted Cohere API key field (mirroring the agent block's hosted-key
pattern), a configurable reranker input pool size (1-100), and surface
meta.warnings from Cohere rerank responses via logger.warn. All new contract
fields are optional and nullable for full backwards compatibility.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>

Author: waleedlatif1

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

@@ -247,9 +247,15 @@ 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
+    /** Oversample vector results when reranking so the reranker has more to choose from.
+     * Cap at 100 to bound Cohere request cost (1 search unit = ≤100 docs). When the caller
+     * supplies `rerankerInputCount`, honor it but never let it drop below `topK`
+     * (which would defeat the purpose) or exceed 100 (which would split into >1 search units). */
+    const candidateTopK = useReranker
+      ? validatedData.rerankerInputCount !== undefined
+        ? Math.min(100, Math.max(validatedData.topK, validatedData.rerankerInputCount))
+        : Math.min(100, validatedData.topK * 4)
+      : validatedData.topK
 
     if (!hasQuery && hasFilters) {
       results = await handleTagOnlySearch({
@@ -300,7 +306,12 @@ export const POST = withRouteHandler(async (request: NextRequest) => {
         const { results: ranked, isBYOK } = await rerank(
           validatedData.query!,
           results.map((r) => ({ id: r.id, text: r.content })),
-          { model: rerankerModel, topN: validatedData.topK, workspaceId }
+          {
+            model: rerankerModel,
+            topN: validatedData.topK,
+            workspaceId,
+            apiKey: validatedData.rerankerApiKey,
+          }
         )
         rerankBilled = true
         rerankIsBYOK = isBYOK

apps/sim/blocks/blocks/knowledge.ts

@@ -1,6 +1,7 @@
 import { PackageSearchIcon } from '@/components/icons'
 import { DEFAULT_RERANKER_MODEL, SUPPORTED_RERANKER_MODELS } from '@/lib/knowledge/reranker-models'
 import type { BlockConfig } from '@/blocks/types'
+import { getCohereRerankerApiKeyCondition } from '@/blocks/utils'
 
 export const KnowledgeBlock: BlockConfig = {
   type: 'knowledge',
@@ -105,6 +106,28 @@ export const KnowledgeBlock: BlockConfig = {
         and: { field: 'rerankerEnabled', value: true },
       },
     },
+    {
+      id: 'rerankerInputCount',
+      title: 'Documents Sent to Reranker',
+      type: 'short-input',
+      placeholder: 'Auto (4× results, capped at 100)',
+      mode: 'advanced',
+      condition: {
+        field: 'operation',
+        value: 'search',
+        and: { field: 'rerankerEnabled', value: true },
+      },
+    },
+    {
+      id: 'apiKey',
+      title: 'Cohere API Key',
+      type: 'short-input',
+      placeholder: 'Enter your Cohere API key',
+      password: true,
+      connectionDroppable: false,
+      required: true,
+      condition: getCohereRerankerApiKeyCondition(),
+    },
 
     // --- List Documents ---
     {
@@ -419,6 +442,11 @@ export const KnowledgeBlock: BlockConfig = {
     tagFilters: { type: 'string', description: 'Tag filter criteria' },
     rerankerEnabled: { type: 'boolean', description: 'Apply Cohere reranking to search results' },
     rerankerModel: { type: 'string', description: 'Cohere rerank model identifier' },
+    rerankerInputCount: {
+      type: 'number',
+      description: 'Number of vector results sent to the Cohere reranker (1–100)',
+    },
+    apiKey: { type: 'string', description: 'Cohere API key (self-hosted only)' },
     documentTags: { type: 'string', description: 'Document tags' },
     chunkSearch: { type: 'string', description: 'Search filter for chunks' },
     chunkEnabledFilter: { type: 'string', description: 'Filter chunks by enabled status' },

apps/sim/blocks/utils.ts

@@ -184,6 +184,26 @@ export function getApiKeyCondition() {
   }
 }
 
+/**
+ * Visibility condition for the Cohere reranker API key field on the Knowledge block.
+ * On hosted Sim, the platform supplies the Cohere key (via workspace BYOK or rotating
+ * env keys), so the field is hidden. On self-hosted deployments, the field is shown
+ * whenever reranking is enabled for a search operation, mirroring the agent block's
+ * `getApiKeyCondition` pattern.
+ */
+export function getCohereRerankerApiKeyCondition() {
+  return () => {
+    if (isHosted) {
+      return { field: 'operation', value: '__never_show__' }
+    }
+    return {
+      field: 'operation',
+      value: 'search',
+      and: { field: 'rerankerEnabled', value: true },
+    }
+  }
+}
+
 /**
  * Returns the standard provider credential subblocks used by LLM-based blocks.
  * This includes: Vertex AI OAuth, API Key, Azure (OpenAI + Anthropic), Vertex AI config, and Bedrock config.

apps/sim/lib/api/contracts/knowledge/search.ts

@@ -36,6 +36,25 @@ export const knowledgeSearchBodySchema = z
       .transform((val) => val || undefined),
     rerankerEnabled: z.boolean().optional().default(false),
     rerankerModel: rerankerModelSchema.optional().default(DEFAULT_RERANKER_MODEL),
+    /**
+     * Number of vector results sent to Cohere as the documents array for reranking. Capped at 100
+     * so each rerank call stays within a single Cohere search unit (1 query × ≤100 docs); see
+     * `RERANK_MODEL_PRICING` in `providers/models.ts`.
+     */
+    rerankerInputCount: z
+      .number()
+      .int('rerankerInputCount must be an integer')
+      .min(1, 'rerankerInputCount must be at least 1')
+      .max(100, 'rerankerInputCount cannot exceed 100')
+      .optional()
+      .nullable()
+      .transform((val) => val ?? undefined),
+    rerankerApiKey: z
+      .string()
+      .min(1, 'rerankerApiKey cannot be empty')
+      .optional()
+      .nullable()
+      .transform((val) => val || undefined),
   })
   .refine(
     (data) => {

apps/sim/lib/knowledge/reranker.ts

@@ -2,6 +2,7 @@ import { createLogger } from '@sim/logger'
 import { getBYOKKey } from '@/lib/api-key/byok'
 import { getRotatingApiKey } from '@/lib/core/config/api-keys'
 import { env } from '@/lib/core/config/env'
+import { isHosted } from '@/lib/core/config/feature-flags'
 import { isRetryableError, retryWithExponentialBackoff } from '@/lib/knowledge/documents/utils'
 import {
   DEFAULT_RERANKER_MODEL,
@@ -56,8 +57,18 @@ class RerankAPIError extends Error {
 }
 
 async function resolveCohereKey(
-  workspaceId?: string | null
+  workspaceId?: string | null,
+  userApiKey?: string
 ): Promise<{ apiKey: string; isBYOK: boolean }> {
+  /**
+   * Mirrors the agent block hosted-key pattern (`injectHostedKeyIfNeeded`):
+   * on self-hosted the user-supplied key from the block field flows through
+   * unchanged; on hosted Sim we always source the key from workspace BYOK or
+   * platform env, so any user-supplied value is ignored.
+   */
+  if (!isHosted && userApiKey) {
+    return { apiKey: userApiKey, isBYOK: false }
+  }
   if (workspaceId) {
     const byokResult = await getBYOKKey(workspaceId, 'cohere')
     if (byokResult) {
@@ -77,8 +88,19 @@ async function resolveCohereKey(
   }
 }
 
+/**
+ * Subset of Cohere v2/rerank response fields we read.
+ * Reference: https://docs.cohere.com/v2/reference/rerank
+ * - `results[].index` maps back to the position in the documents we sent.
+ * - `results[].relevance_score` is normalized 0–1.
+ * - `meta.warnings` is documented as an array of strings; we surface them in logs
+ *   so issues like document truncation don't disappear silently.
+ */
 interface CohereRerankResponse {
   results: Array<{ index: number; relevance_score: number }>
+  meta?: {
+    warnings?: string[]
+  }
 }
 
 /**
@@ -92,6 +114,8 @@ export async function rerank<T extends RerankItem>(
     model: string
     topN?: number
     workspaceId?: string | null
+    /** User-supplied Cohere key from the Knowledge block field. Honored only on self-hosted. */
+    apiKey?: string
   }
 ): Promise<RerankResponse<T>> {
   if (items.length === 0) return { results: [], isBYOK: false }
@@ -100,7 +124,7 @@ export async function rerank<T extends RerankItem>(
     throw new Error(`Unsupported reranker model: ${options.model}`)
   }
 
-  const { apiKey, isBYOK } = await resolveCohereKey(options.workspaceId)
+  const { apiKey, isBYOK } = await resolveCohereKey(options.workspaceId, options.apiKey)
   const cappedItems =
     items.length > MAX_DOCUMENTS_PER_RERANK ? items.slice(0, MAX_DOCUMENTS_PER_RERANK) : items
   if (items.length > MAX_DOCUMENTS_PER_RERANK) {
@@ -151,6 +175,13 @@ export async function rerank<T extends RerankItem>(
     }
   )
 
+  if (response.meta?.warnings && response.meta.warnings.length > 0) {
+    logger.warn('Cohere rerank returned warnings', {
+      model: options.model,
+      warnings: response.meta.warnings,
+    })
+  }
+
   return {
     results: response.results
       .filter((r) => r.index >= 0 && r.index < cappedItems.length)

apps/sim/tools/knowledge/search.ts

@@ -55,6 +55,19 @@ export const knowledgeSearchTool: ToolConfig<any, KnowledgeSearchResponse> = {
       description:
         'Cohere rerank model to use (one of: rerank-v4.0-pro, rerank-v4.0-fast, rerank-v3.5)',
     },
+    rerankerInputCount: {
+      type: 'number',
+      required: false,
+      visibility: 'user-only',
+      description:
+        'Number of vector results sent to the Cohere reranker (1–100). Defaults to topK × 4 capped at 100.',
+    },
+    apiKey: {
+      type: 'string',
+      required: false,
+      visibility: 'user-only',
+      description: 'Cohere API key for reranker (self-hosted deployments only)',
+    },
   },
 
   schemaEnrichment: {
@@ -84,13 +97,27 @@ export const knowledgeSearchTool: ToolConfig<any, KnowledgeSearchResponse> = {
         typeof params.rerankerModel === 'string' && params.rerankerModel.length > 0
           ? params.rerankerModel
           : DEFAULT_RERANKER_MODEL
+      const rerankerApiKey =
+        typeof params.apiKey === 'string' && params.apiKey.length > 0 ? params.apiKey : undefined
+      const rawInputCount =
+        params.rerankerInputCount !== undefined && params.rerankerInputCount !== null
+          ? Number(params.rerankerInputCount)
+          : Number.NaN
+      const rerankerInputCount = Number.isFinite(rawInputCount)
+        ? Math.max(1, Math.min(100, Math.floor(rawInputCount)))
+        : undefined
 
       const requestBody = {
         knowledgeBaseIds,
         query: params.query,
         topK: params.topK ? Math.max(1, Math.min(100, Number(params.topK))) : 10,
         ...(structuredFilters.length > 0 && { tagFilters: structuredFilters }),
-        ...(rerankerEnabled && { rerankerEnabled: true, rerankerModel }),
+        ...(rerankerEnabled && {
+          rerankerEnabled: true,
+          rerankerModel,
+          ...(rerankerInputCount !== undefined && { rerankerInputCount }),
+          ...(rerankerApiKey && { rerankerApiKey }),
+        }),
         ...(workflowId && { workflowId }),
       }