fix(knowledge): align search fixtures with shared retrieval

Author: icecrasher321

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

@@ -217,6 +217,7 @@ describe('Knowledge Search Utils', () => {
         Array.from({ length: 201 }, (_, index) => ({ id: `candidate-${index}` }))
       )
       queueTableRows(schemaMock.embedding, [makeResult('second', 0.2), makeResult('first', 0.1)])
+      queueTableRows(schemaMock.embedding, [makeResult('second', 0.2), makeResult('first', 0.1)])
 
       const results = await handleTagAndVectorSearch({
         knowledgeBaseIds: ['kb-1', 'kb-2'],
@@ -231,9 +232,8 @@ describe('Knowledge Search Utils', () => {
 
       expect(results.map((row) => row.id)).toEqual(['first', 'second'])
       expect(dbChainMockFns.select).toHaveBeenCalledTimes(3)
-      expect(dbChainMockFns.as).toHaveBeenCalledWith('ranked_embeddings')
       expect(Object.keys(dbChainMockFns.select.mock.calls[0][0])).toEqual(['id'])
-      expect(dbChainMockFns.limit).toHaveBeenNthCalledWith(1, 201)
+      expect(dbChainMockFns.limit).toHaveBeenNthCalledWith(1, 400)
       expect(dbChainMockFns.select.mock.calls[1][0]).toHaveProperty('distance')
       expect(dbChainMockFns.limit).toHaveBeenCalledWith(2)
     })
@@ -463,11 +463,7 @@ describe('Knowledge Search Utils', () => {
         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.
-       */
+      /** Keyword retrieval preserves its existing per-base lexical candidate selection. */
       expect(dbChainMockFns.select).toHaveBeenCalledTimes(knowledgeBaseIds.length)
     })
 
@@ -542,6 +538,7 @@ describe('Knowledge Search Utils', () => {
     })
 
     it('runs a single retrieval leg in vector mode', async () => {
+      dbChainMockFns.execute.mockResolvedValue([{ id: 'vector-hit' }])
       queueTableRows(schemaMock.embedding, [{ id: 'vector-hit' }])
       queueTableRows(schemaMock.embedding, [makeResult('vector-hit')])
 
@@ -551,22 +548,22 @@ describe('Knowledge Search Utils', () => {
         topK: 10,
         searchMode: 'vector',
         query: 'PROJ-1234',
-        queryVector: JSON.stringify([0.1, 0.2, 0.3]),
+        queryVector: { vector: JSON.stringify(TEST_EMBEDDING), dimensions: 1536 },
       })
 
       expect(results.map((r) => r.id)).toEqual(['vector-hit'])
-      expect(dbChainMockFns.select).toHaveBeenCalledTimes(3)
-      expect(dbChainMockFns.as).toHaveBeenCalledWith('ranked_embeddings')
+      expect(dbChainMockFns.select).toHaveBeenCalledTimes(2)
     })
 
     it('runs both legs and fuses them in hybrid mode', async () => {
       /**
-       * Chains dequeue in creation order: keyword ranking, the budgeted vector
-       * probe, keyword hydration, then vector ranking and hydration in one query.
+       * The raw vector probe does not consume a table chain. Keyword ranking and
+       * hydration complete before vector exact ranking and content hydration.
        */
+      dbChainMockFns.execute.mockResolvedValue([{ id: 'vector-hit' }])
       queueTableRows(schemaMock.embedding, [{ id: 'keyword-hit', keywordRank: 0.9 }])
-      queueTableRows(schemaMock.embedding, [{ id: 'vector-hit' }])
       queueTableRows(schemaMock.embedding, [makeResult('keyword-hit')])
+      queueTableRows(schemaMock.embedding, [{ id: 'vector-hit' }])
       queueTableRows(schemaMock.embedding, [makeResult('vector-hit')])
 
       const results = await executeKnowledgeSearch({
@@ -575,15 +572,16 @@ describe('Knowledge Search Utils', () => {
         topK: 10,
         searchMode: 'hybrid',
         query: 'PROJ-1234',
-        queryVector: JSON.stringify([0.1, 0.2, 0.3]),
+        queryVector: { vector: JSON.stringify(TEST_EMBEDDING), dimensions: 1536 },
       })
 
       expect(results.map((r) => r.id).sort()).toEqual(['keyword-hit', 'vector-hit'])
-      expect(dbChainMockFns.select).toHaveBeenCalledTimes(5)
+      expect(dbChainMockFns.select).toHaveBeenCalledTimes(4)
     })
 
     it('propagates unexpected keyword errors after the vector leg finishes', async () => {
       /** The failing ranking chain is still built first and takes the first queued set. */
+      dbChainMockFns.execute.mockResolvedValue([{ id: 'vector-hit' }])
       queueTableRows(schemaMock.embedding, [{ id: 'never-ranked', keywordRank: 0 }])
       queueTableRows(schemaMock.embedding, [{ id: 'vector-hit' }])
       queueTableRows(schemaMock.embedding, [makeResult('vector-hit')])
@@ -600,10 +598,10 @@ describe('Knowledge Search Utils', () => {
           topK: 10,
           searchMode: 'hybrid',
           query: 'PROJ-1234',
-          queryVector: JSON.stringify([0.1, 0.2, 0.3]),
+          queryVector: { vector: JSON.stringify(TEST_EMBEDDING), dimensions: 1536 },
         })
       ).rejects.toBe(failure)
-      expect(dbChainMockFns.as).toHaveBeenCalledWith('ranked_embeddings')
+      expect(dbChainMockFns.select).toHaveBeenCalledTimes(3)
     })
 
     it('skips both query legs when only tag filters are provided', async () => {

apps/sim/lib/knowledge/__integration__/search-latency.integration.ts

@@ -1,4 +1,4 @@
-/** Real Assistant tool, application authorization, PostgreSQL/pgvector, and result processing. */
+/** Real search adapters, application authorization, PostgreSQL/pgvector, and result processing. */
 import { readFileSync, statSync, writeFileSync } from 'node:fs'
 import type { Principal } from '@sim/auth/principal'
 import { db } from '@sim/db'
@@ -9,7 +9,6 @@ import {
   document,
   embedding,
   knowledgeBase,
-  knowledgeBaseTagDefinitions,
   knowledgeConnector,
   knowledgeConnectorMember,
   knowledgeDocumentObservation,
@@ -103,6 +102,8 @@ function readFixtureReport(file: string) {
 const reused = reuseFile ? readFixtureReport(reuseFile) : undefined
 const ids = reused?.fixture ?? createKnowledgeAclFixtureIds()
 const unrelated = reused?.unrelatedFixture ?? createKnowledgeAclFixtureIds()
+const fullWidthFixture = createKnowledgeAclFixtureIds()
+const fullWidthChunkCount = 5000
 const organizationChatId = generateId()
 function topicVector(topic = 0) {
   const vector = Array.from({ length: dimensions }, (_, index) =>
@@ -124,6 +125,7 @@ const report: Record<string, unknown> = {
     fixtureVersion: 2,
     chunkCount,
     unrelatedChunkCount,
+    fullWidthChunkCount,
     dimensions,
     candidateDimensions,
     chunksPerDocument,
@@ -133,7 +135,7 @@ const report: Record<string, unknown> = {
     vectors:
       'Normalized 512-dimensional topic/noise geometry with permuted copies across 1536 dimensions; verifies prefix candidate ranking, not semantic embedding quality',
     cache: evictSharedBuffers
-      ? 'Organization samples evict PostgreSQL shared buffers before each request; operating-system cache is not cleared'
+      ? 'Selected workspace and organization samples evict PostgreSQL shared buffers; operating-system cache is not cleared'
       : 'First and repeated samples; no claim of a cold operating-system cache',
     layout: reused
       ? 'Reused fixture; physical layout is inherited from its original report'
@@ -383,17 +385,22 @@ async function searchDashboard(
 
 async function searchWorkspaceKb(
   query = 'Orion deployment',
-  options: { principal?: Principal; tagFilters?: KnowledgeSearchTagFilter[] } = {}
+  options: {
+    principal?: Principal
+    tagFilters?: KnowledgeSearchTagFilter[]
+    fixture?: typeof ids
+  } = {}
 ) {
+  const fixture = options.fixture ?? ids
   const result = await searchKnowledge.execute({
     principal: options.principal ?? {
       kind: 'workspace_api_key',
-      workspaceId: ids.workspaceId,
+      workspaceId: fixture.workspaceId,
       keyId: 'fixture-search-key',
     },
     input: {
-      workspaceId: ids.workspaceId,
-      knowledgeBaseIds: [ids.knowledgeBaseId],
+      workspaceId: fixture.workspaceId,
+      knowledgeBaseIds: [fixture.knowledgeBaseId],
       query,
       topK: 15,
       searchMode: 'vector',
@@ -683,6 +690,32 @@ describe.skipIf(!enabled)('Knowledge search latency on a realistic indexed corpu
       logger.info('Synthetic corpora loaded', { chunkCount, unrelatedChunkCount })
       for (const index of indexes) await db.execute(sql.raw(index.indexdef))
     }
+    /** A non-shortenable model must populate its own full-width projection through the write trigger. */
+    await seedKnowledgeAclFixture(fullWidthFixture, { connectorType: 'google_drive' })
+    await db
+      .update(knowledgeBase)
+      .set({ embeddingModel: 'text-embedding-ada-002' })
+      .where(eq(knowledgeBase.id, fullWidthFixture.knowledgeBaseId))
+    await db.execute(sql`
+      WITH source AS MATERIALIZED (
+        SELECT id, content, embedding FROM embedding
+        WHERE knowledge_base_id = ${ids.knowledgeBaseId} ORDER BY id LIMIT ${fullWidthChunkCount}
+      ), documents AS (
+        INSERT INTO document
+          (id, knowledge_base_id, connector_id, external_id, filename, file_url, file_size,
+            mime_type, processing_status, acl, acl_verified_at)
+        SELECT ${fullWidthFixture.workspaceId} || '-doc-' || id, ${fullWidthFixture.knowledgeBaseId},
+          ${fullWidthFixture.connectorId}, id, 'Full-width deployment guide',
+          'https://fixture.invalid/full-width', 12000, 'text/plain', 'completed',
+          ARRAY['pub']::text[], statement_timestamp() FROM source RETURNING id
+      ) INSERT INTO embedding
+        (id, knowledge_base_id, document_id, chunk_index, chunk_hash, content, content_length,
+          token_count, start_offset, end_offset, embedding)
+      SELECT ${fullWidthFixture.workspaceId} || '-chunk-' || source.id,
+        ${fullWidthFixture.knowledgeBaseId}, documents.id, 0, source.id, source.content,
+        3000, 750, 0, 3000, source.embedding
+      FROM source JOIN documents ON documents.id = ${fullWidthFixture.workspaceId} || '-doc-' || source.id
+    `)
     await db.execute(sql`ANALYZE document`)
     await db.execute(sql`ANALYZE embedding`)
     await db.execute(sql`ANALYZE embedding_search`)
@@ -710,9 +743,12 @@ describe.skipIf(!enabled)('Knowledge search latency on a realistic indexed corpu
     db.$client.options.debug = previousDebug
     vi.unstubAllGlobals()
     saveReport()
-    for (const fixture of process.env.KNOWLEDGE_SEARCH_PERFORMANCE_KEEP_DATABASE === 'true'
-      ? []
-      : [ids, unrelated]) {
+    for (const fixture of [
+      fullWidthFixture,
+      ...(process.env.KNOWLEDGE_SEARCH_PERFORMANCE_KEEP_DATABASE === 'true'
+        ? []
+        : [ids, unrelated]),
+    ]) {
       await db.delete(workspace).where(eq(workspace.id, fixture.workspaceId))
       await db.delete(organization).where(eq(organization.id, fixture.organizationId))
       await db.delete(user).where(eq(user.id, fixture.aliceId))
@@ -1281,18 +1317,12 @@ describe.skipIf(!enabled)('Knowledge search latency on a realistic indexed corpu
         report[`${label}.recall`] = { neighbors: expected.length, recall }
         saveReport()
       }
-      await db
-        .update(knowledgeBase)
-        .set({ embeddingModel: 'text-embedding-ada-002' })
-        .where(eq(knowledgeBase.id, ids.knowledgeBaseId))
-      const fullWidth = await sample('workspace-kb.full-width', () => searchWorkspaceKb())
+      const fullWidth = await sample('workspace-kb.full-width', () =>
+        searchWorkspaceKb('Orion deployment', { fixture: fullWidthFixture })
+      )
       expectCompleteVectorSearch(fullWidth.diagnostics)
       expect(fullWidth.diagnostics.vectorCandidateDimensions).toBe(dimensions)
       expect(fullWidth.result.data.results).toHaveLength(15)
-      await db
-        .update(knowledgeBase)
-        .set({ embeddingModel: 'text-embedding-3-small' })
-        .where(eq(knowledgeBase.id, ids.knowledgeBaseId))
 
       const workflowId = generateId()
       const scheduled: Principal = {
@@ -1346,17 +1376,11 @@ describe.skipIf(!enabled)('Knowledge search latency on a realistic indexed corpu
         for (const diagnostics of completed) expectCompleteVectorSearch(diagnostics)
       }
 
-      await db.insert(knowledgeBaseTagDefinitions).values({
-        id: generateId(),
-        knowledgeBaseId: ids.knowledgeBaseId,
-        tagSlot: 'tag1',
-        displayName: 'Fixture group',
-      })
       await db.execute(sql`UPDATE embedding SET tag1 = 'selected' WHERE knowledge_base_id = ${ids.knowledgeBaseId}
         AND document_id IN (SELECT id FROM document WHERE knowledge_base_id = ${ids.knowledgeBaseId} AND external_id::int < 600)`)
       const tagged = await sample('workspace-kb.tagged', () =>
         searchWorkspaceKb('Orion deployment', {
-          tagFilters: [{ tagName: 'Fixture group', operator: 'eq', value: 'selected' }],
+          tagFilters: [{ tagName: 'Fixture', operator: 'eq', value: 'selected' }],
         })
       )
       expectCompleteVectorSearch(tagged.diagnostics)
@@ -1393,16 +1417,9 @@ describe.skipIf(!enabled)('Knowledge search latency on a realistic indexed corpu
       expectCompleteVectorSearch(denied.diagnostics)
       expect(denied.result.data.results).toEqual([])
     } finally {
-      await db
-        .delete(knowledgeBaseTagDefinitions)
-        .where(eq(knowledgeBaseTagDefinitions.knowledgeBaseId, ids.knowledgeBaseId))
       await db.execute(
         sql`UPDATE embedding SET tag1 = NULL WHERE knowledge_base_id = ${ids.knowledgeBaseId} AND tag1 = 'selected'`
       )
-      await db
-        .update(knowledgeBase)
-        .set({ embeddingModel: 'text-embedding-3-small' })
-        .where(eq(knowledgeBase.id, ids.knowledgeBaseId))
       await db
         .update(knowledgeConnector)
         .set({ accessRewritePending: false })