fix(knowledge): fan the keyword leg out per knowledge base

The vector leg caps candidates per base once getQueryStrategy sets useParallel,
but the keyword leg always ran one global query with a single LIMIT. Searching
several bases at once let whichever one ranks strongest lexically consume every
slot, so an exact-token hit in a smaller base never reached fusion — the case
hybrid exists to serve.

The keyword leg now uses the same strategy: per-base queries under the same
parallel limit, re-ranked globally on a selected ts_rank_cd. Both legs draw
candidates the same way, so fusion combines rankings over the same pool.

Author: waleedlatif1

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

@@ -46,9 +46,11 @@ afterEach(() => {
 })
 
 import {
+  executeKeywordSearch,
   executeKnowledgeSearch,
   fuseByReciprocalRank,
   generateSearchEmbedding,
+  getQueryStrategy,
   handleTagAndVectorSearch,
   handleTagOnlySearch,
   handleVectorOnlySearch,
@@ -356,6 +358,57 @@ describe('Knowledge Search Utils', () => {
     })
   })
 
+  describe('executeKeywordSearch', () => {
+    beforeEach(() => {
+      resetDbChainMock()
+    })
+
+    it('returns nothing for a whitespace-only query without touching the database', async () => {
+      const results = await executeKeywordSearch({
+        knowledgeBaseIds: ['kb-123'],
+        topK: 10,
+        query: '   ',
+        queryVector: JSON.stringify([0.1, 0.2, 0.3]),
+      })
+
+      expect(results).toEqual([])
+      expect(dbChainMockFns.select).not.toHaveBeenCalled()
+    })
+
+    it('issues one query per knowledge base once the parallel threshold is crossed', async () => {
+      const knowledgeBaseIds = ['kb-1', 'kb-2', 'kb-3', 'kb-4', 'kb-5']
+      expect(getQueryStrategy(knowledgeBaseIds.length, 10).useParallel).toBe(true)
+
+      await executeKeywordSearch({
+        knowledgeBaseIds,
+        topK: 10,
+        query: 'PROJ-1234',
+        queryVector: JSON.stringify([0.1, 0.2, 0.3]),
+      })
+
+      /**
+       * A single global LIMIT would let the lexically strongest base consume
+       * every slot, so an exact-token hit in a smaller base never reaches
+       * fusion. The vector leg already fans out here; both legs must match.
+       */
+      expect(dbChainMockFns.select).toHaveBeenCalledTimes(knowledgeBaseIds.length)
+    })
+
+    it('uses a single query when the parallel threshold is not crossed', async () => {
+      const knowledgeBaseIds = ['kb-1', 'kb-2']
+      expect(getQueryStrategy(knowledgeBaseIds.length, 10).useParallel).toBe(false)
+
+      await executeKeywordSearch({
+        knowledgeBaseIds,
+        topK: 10,
+        query: 'PROJ-1234',
+        queryVector: JSON.stringify([0.1, 0.2, 0.3]),
+      })
+
+      expect(dbChainMockFns.select).toHaveBeenCalledTimes(1)
+    })
+  })
+
   describe('executeKnowledgeSearch', () => {
     beforeEach(() => {
       resetDbChainMock()

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

@@ -564,6 +564,12 @@ export interface KeywordSearchParams {
  * can report `similarity` for rows only the lexical leg found. Unlike the vector
  * leg there is no distance threshold — surfacing exact-token matches that are
  * semantically distant is the entire point of this leg.
+ *
+ * Candidate gathering mirrors the vector leg's `getQueryStrategy`: across many
+ * knowledge bases a single global `LIMIT` lets whichever base ranks strongest
+ * lexically consume every slot, so an exact-token hit in a smaller base would
+ * never reach fusion. Both legs must draw candidates the same way, or rank
+ * fusion is combining rankings taken over differently-shaped pools.
  */
 export async function executeKeywordSearch(params: KeywordSearchParams): Promise<SearchResult[]> {
   const { knowledgeBaseIds, topK, query, queryVector, structuredFilters } = params
@@ -573,16 +579,51 @@ export async function executeKeywordSearch(params: KeywordSearchParams): Promise
   }
 
   const tsQuery = sql`websearch_to_tsquery(${FTS_CONFIG}, ${query})`
+  const rankExpr = sql<number>`ts_rank_cd(${embedding.contentTsv}, ${tsQuery})`
   const tagFilterConditions = structuredFilters?.length
     ? getStructuredTagFilters(structuredFilters, embedding)
     : []
 
-  return await db
-    .select(
-      getSearchResultFields(
-        sql<number>`${embedding.embedding} <=> ${queryVector}::vector`.as('distance')
+  /** Selected alongside the row so per-base batches can be re-ranked globally. */
+  const selectFields = {
+    ...getSearchResultFields(
+      sql<number>`${embedding.embedding} <=> ${queryVector}::vector`.as('distance')
+    ),
+    keywordRank: rankExpr.as('keyword_rank'),
+  }
+
+  const strategy = getQueryStrategy(knowledgeBaseIds.length, topK)
+
+  if (strategy.useParallel) {
+    const parallelLimit = Math.ceil(topK / knowledgeBaseIds.length) + 5
+
+    const perBase = await Promise.all(
+      knowledgeBaseIds.map((kbId) =>
+        db
+          .select(selectFields)
+          .from(embedding)
+          .innerJoin(document, eq(embedding.documentId, document.id))
+          .where(
+            and(
+              eq(embedding.knowledgeBaseId, kbId),
+              ...getVisibilityConditions(),
+              sql`${embedding.contentTsv} @@ ${tsQuery}`,
+              ...tagFilterConditions
+            )
+          )
+          .orderBy(sql`${rankExpr} DESC`)
+          .limit(parallelLimit)
       )
     )
+
+    return perBase
+      .flat()
+      .sort((a, b) => b.keywordRank - a.keywordRank)
+      .slice(0, topK)
+  }
+
+  return await db
+    .select(selectFields)
     .from(embedding)
     .innerJoin(document, eq(embedding.documentId, document.id))
     .where(
@@ -593,7 +634,7 @@ export async function executeKeywordSearch(params: KeywordSearchParams): Promise
         ...tagFilterConditions
       )
     )
-    .orderBy(sql`ts_rank_cd(${embedding.contentTsv}, ${tsQuery}) DESC`)
+    .orderBy(sql`${rankExpr} DESC`)
     .limit(topK)
 }