fix(knowledge): store compact vector search candidates

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

@@ -87,11 +87,14 @@ const reused = reuseFile ? readFixtureReport(reuseFile) : undefined
 const ids = reused?.fixture ?? createKnowledgeAclFixtureIds()
 const unrelated = reused?.unrelatedFixture ?? createKnowledgeAclFixtureIds()
 const organizationChatId = generateId()
-const queryVector = Array.from({ length: dimensions }, (_, index) =>
-  Math.sin((index + 1) * 12.9898)
-)
-const queryMagnitude = Math.hypot(...queryVector)
-for (let index = 0; index < queryVector.length; index++) queryVector[index] /= queryMagnitude
+function topicVector(topic = 0) {
+  const vector = Array.from({ length: dimensions }, (_, index) =>
+    Math.sin((index + 1) * (topic + 1) * 12.9898)
+  )
+  const magnitude = Math.hypot(...vector)
+  return vector.map((value) => value / magnitude)
+}
+const queryVector = topicVector()
 const captured: CapturedQuery[] = []
 const report: Record<string, unknown> = {
   fixture: ids,
@@ -124,6 +127,8 @@ interface ExplainNode {
   'Node Type': string
   'Actual Rows': number
   'Index Name'?: string
+  'Relation Name'?: string
+  Output?: string[]
   Plans?: ExplainNode[]
 }
 
@@ -133,6 +138,8 @@ const explainNodeSchema: z.ZodType<ExplainNode> = z.lazy(() =>
       'Node Type': z.string(),
       'Actual Rows': z.number(),
       'Index Name': z.string().optional(),
+      'Relation Name': z.string().optional(),
+      Output: z.array(z.string()).optional(),
       Plans: z.array(explainNodeSchema).optional(),
     })
     .passthrough()
@@ -141,12 +148,20 @@ const explainSchema = z.array(z.object({ Plan: explainNodeSchema }).passthrough(
 
 function usesVectorIndex(node: ExplainNode): boolean {
   return (
-    node['Index Name'] === 'embedding_binary_hnsw_idx' ||
+    node['Index Name'] === 'embedding_search_binary_hnsw_idx' ||
     node['Index Name'] === 'embedding_vector_hnsw_idx' ||
     (node.Plans?.some(usesVectorIndex) ?? false)
   )
 }
 
+/** The ANN stage must not fetch full vectors, even for planner-added sort projections. */
+function assertCompactCandidates(node: ExplainNode) {
+  expect(node['Relation Name']).not.toBe('embedding')
+  for (const expression of node.Output ?? [])
+    expect(expression).not.toMatch(/binary_quantize\([^)]*embedding\.embedding/)
+  for (const child of node.Plans ?? []) assertCompactCandidates(child)
+}
+
 function saveReport() {
   const file = process.env.KNOWLEDGE_SEARCH_PERFORMANCE_REPORT_FILE
   if (file) writeFileSync(file, JSON.stringify(report, null, 2), { mode: 0o600 })
@@ -238,7 +253,7 @@ async function sample(label: string, run: () => ReturnType<typeof search>) {
   expect(captured.length).toBeLessThan(300)
   const searches = captured.filter(
     (item) =>
-      item.query.includes('from "embedding"') &&
+      (item.query.includes('from "embedding"') || item.query.includes('from "embedding_search"')) &&
       (item.query.includes('order by') || item.query.includes('limit'))
   )
   const plans = []
@@ -251,7 +266,10 @@ async function sample(label: string, run: () => ReturnType<typeof search>) {
         await tx.unsafe('SET LOCAL hnsw.ef_search = 200')
         await tx.unsafe('SET LOCAL hnsw.scan_mem_multiplier = 4')
       }
-      return tx.unsafe(`EXPLAIN (ANALYZE, BUFFERS, FORMAT JSON) ${query.query}`, query.parameters)
+      return tx.unsafe(
+        `EXPLAIN (ANALYZE, BUFFERS, VERBOSE, FORMAT JSON) ${query.query}`,
+        query.parameters
+      )
     })
     plans.push({
       kind: query.query.includes('keyword_rank')
@@ -311,7 +329,8 @@ describe.skipIf(!enabled)('Assistant search latency on a realistic indexed corpu
         .parse(JSON.parse(String(init?.body)))
       embeddingCalls += body.input.length
       const bytes = Buffer.alloc(dimensions * 4)
-      queryVector.forEach((value, index) => bytes.writeFloatLE(value, index * 4))
+      const topic = Number(/^Topic (\d+) deployment$/.exec(body.input[0])?.[1] ?? 0)
+      topicVector(topic).forEach((value, index) => bytes.writeFloatLE(value, index * 4))
       return Response.json({
         data: [{ embedding: bytes.toString('base64') }],
         usage: { total_tokens: 4 },
@@ -369,7 +388,7 @@ describe.skipIf(!enabled)('Assistant search latency on a realistic indexed corpu
         indexname: string
         indexdef: string
       }>(sql`SELECT indexname, indexdef FROM pg_indexes
-      WHERE tablename = 'embedding' AND indexdef LIKE '% USING hnsw %'`)
+      WHERE tablename IN ('embedding', 'embedding_search') AND indexdef LIKE '% USING hnsw %'`)
       for (const index of indexes)
         await db.execute(sql`DROP INDEX ${sql.identifier(index.indexname)}`)
       for (const fixture of [ids, unrelated]) {
@@ -406,6 +425,7 @@ describe.skipIf(!enabled)('Assistant search latency on a realistic indexed corpu
     }
     await db.execute(sql`ANALYZE document`)
     await db.execute(sql`ANALYZE embedding`)
+    await db.execute(sql`ANALYZE embedding_search`)
     report.server = (
       await db.execute(sql`SELECT version(), current_setting('work_mem') AS work_mem,
       (SELECT extversion FROM pg_extension WHERE extname = 'vector') AS pgvector`)
@@ -627,6 +647,7 @@ describe.skipIf(!enabled)('Assistant search latency on a realistic indexed corpu
       const vectorPlans = plans.filter((plan) => plan.kind === 'vector')
       expect(vectorPlans).toHaveLength(1)
       expect(usesVectorIndex(vectorPlans[0].plan[0].Plan)).toBe(true)
+      assertCompactCandidates(vectorPlans[0].plan[0].Plan)
       expect(plans.some((plan) => plan.kind === 'rerank')).toBe(true)
       const rerank = plans.find((plan) => plan.kind === 'rerank')!
       const actual = await db.$client.unsafe(rerank.query, rerank.parameters).values()
@@ -643,6 +664,27 @@ describe.skipIf(!enabled)('Assistant search latency on a realistic indexed corpu
     expect(embeddingCalls - before).toBe(2)
   }, 180_000)
 
+  it('preserves exact-neighbor recall across different query vectors', async () => {
+    for (const topic of [3, 11, 23]) {
+      const { plans } = await sample(`topic.${topic}`, () =>
+        search(ids.aliceId, `Topic ${topic} deployment`)
+      )
+      const candidates = plans.find((plan) => plan.kind === 'vector')!
+      assertCompactCandidates(candidates.plan[0].Plan)
+      const rerank = plans.find((plan) => plan.kind === 'rerank')!
+      const actual = await db.$client.unsafe(rerank.query, rerank.parameters).values()
+      const expected = await db.execute<{ id: string }>(sql`SELECT id FROM embedding
+        WHERE knowledge_base_id = ${ids.knowledgeBaseId} AND enabled
+        ORDER BY (embedding <=> ${JSON.stringify(topicVector(topic))}::vector) + 0, id
+        LIMIT ${actual.length}`)
+      const expectedIds = new Set(expected.map(({ id }) => id))
+      const recall = actual.filter(([id]) => expectedIds.has(id)).length / expected.length
+      expect(recall).toBeGreaterThanOrEqual(0.95)
+      report[`recall.topic.${topic}`] = { neighbors: expected.length, recall }
+      saveReport()
+    }
+  }, 180_000)
+
   it('compares the Search tab and Assistant with the same person, query and index', async () => {
     const dashboard = await sample('dashboard', async () => {
       const result = await searchScopedKnowledge.execute({

apps/sim/lib/knowledge/search/diagnostics.ts

@@ -70,6 +70,7 @@ export interface SearchDiagnosticMetadata {
   boostRecency?: boolean
   embeddingDimensions?: number
   vectorRanking?: 'exact' | 'binary-rerank'
+  vectorCandidateStorage?: 'stored-binary'
   vectorBudgetMs?: number
   vectorCandidateLimit?: number
   vectorCandidateCount?: number

apps/sim/lib/knowledge/search/queries.test.ts

@@ -580,10 +580,10 @@ describe('live repository authorization follows ranked candidates', () => {
 
   it('bounds broad vector ranking before metadata and reorders relaxed candidates before trimming', async () => {
     queueTableRows(
-      schemaMock.embedding,
+      schemaMock.embeddingSearch,
       Array.from({ length: 200 }, (_, index) => candidate(`probe-${index}`, 'allowed-source'))
     )
-    queueTableRows(schemaMock.embedding, [{ id: 'far' }, { id: 'near' }])
+    queueTableRows(schemaMock.embeddingSearch, [{ id: 'far' }, { id: 'near' }])
     queueTableRows(schemaMock.embedding, [
       { ...candidate('far', 'allowed-source'), distance: 0.3 },
       { ...candidate('near', 'allowed-source'), distance: 0.1 },
@@ -617,7 +617,7 @@ describe('live repository authorization follows ranked candidates', () => {
   })
 
   it('finishes empty scopes after the bounded probe without scanning HNSW or calling providers', async () => {
-    queueTableRows(schemaMock.embedding, [])
+    queueTableRows(schemaMock.embeddingSearch, [])
     expect(await handleVectorOnlySearch({ ...params, structuredFilters: undefined })).toEqual([])
     expect(dbChainMockFns.select).toHaveBeenCalledOnce()
     expect(dbChainMockFns.limit).toHaveBeenCalledExactlyOnceWith(200)
@@ -626,7 +626,7 @@ describe('live repository authorization follows ranked candidates', () => {
   })
 
   it('reads vectors only for the bounded IDs when a broad scope has few candidates', async () => {
-    queueTableRows(schemaMock.embedding, [candidate('selected', 'allowed-source')])
+    queueTableRows(schemaMock.embeddingSearch, [candidate('selected', 'allowed-source')])
     queueTableRows(schemaMock.embedding, [candidate('selected', 'allowed-source')])
     queueTableRows(schemaMock.embedding, [
       { id: 'selected', content: 'Verified small scope', distance: 0.1 },
@@ -652,10 +652,10 @@ describe('live repository authorization follows ranked candidates', () => {
 
   it('falls back to exact ranking when the approximate page cannot fill its limit', async () => {
     queueTableRows(
-      schemaMock.embedding,
+      schemaMock.embeddingSearch,
       Array.from({ length: 200 }, (_, index) => candidate(`probe-${index}`, 'allowed-source'))
     )
-    queueTableRows(schemaMock.embedding, [{ id: 'partial' }])
+    queueTableRows(schemaMock.embeddingSearch, [{ id: 'partial' }])
     queueTableRows(schemaMock.embedding, [candidate('partial', 'allowed-source')])
     queueTableRows(schemaMock.embedding, [candidate('selected', 'allowed-source')])
     queueTableRows(schemaMock.embedding, [
@@ -677,17 +677,17 @@ describe('live repository authorization follows ranked candidates', () => {
     const approximate = Array.from({ length: 20 }, (_, index) =>
       candidate(`approximate-${index}`, 'allowed-source')
     )
-    queueTableRows(schemaMock.embedding, probe)
+    queueTableRows(schemaMock.embeddingSearch, probe)
     queueTableRows(
-      schemaMock.embedding,
+      schemaMock.embeddingSearch,
       approximate.map(({ id }) => ({ id }))
     )
     queueTableRows(schemaMock.embedding, approximate)
     queueTableRows(schemaMock.embedding, [])
-    queueTableRows(schemaMock.embedding, probe)
-    queueTableRows(schemaMock.embedding, [])
+    queueTableRows(schemaMock.embeddingSearch, probe)
+    queueTableRows(schemaMock.embeddingSearch, [])
     queueTableRows(schemaMock.embedding, approximate)
-    queueTableRows(schemaMock.embedding, probe)
+    queueTableRows(schemaMock.embeddingSearch, probe)
     queueTableRows(schemaMock.embedding, [candidate('selected', 'allowed-source')])
     queueTableRows(schemaMock.embedding, [
       { id: 'selected', content: 'Reachable after the exact restart', distance: 0.1 },
@@ -705,15 +705,15 @@ describe('live repository authorization follows ranked candidates', () => {
     const probe = Array.from({ length: 200 }, (_, index) =>
       candidate(`probe-${index}`, 'allowed-source')
     )
-    queueTableRows(schemaMock.embedding, probe)
-    queueTableRows(schemaMock.embedding, [{ id: 'far' }])
+    queueTableRows(schemaMock.embeddingSearch, probe)
+    queueTableRows(schemaMock.embeddingSearch, [{ id: 'far' }])
     queueTableRows(schemaMock.embedding, [
       { ...candidate('far', 'allowed-source'), distance: 0.7 },
       ...Array.from({ length: 19 }, (_, index) => candidate(`hidden-${index}`, 'allowed-source')),
     ])
     queueTableRows(schemaMock.embedding, [{ id: 'far', content: 'Far result', distance: 0.7 }])
-    queueTableRows(schemaMock.embedding, probe)
-    queueTableRows(schemaMock.embedding, [])
+    queueTableRows(schemaMock.embeddingSearch, probe)
+    queueTableRows(schemaMock.embeddingSearch, [])
     queueTableRows(schemaMock.embedding, [
       candidate('near', 'allowed-source'),
       candidate('nearer', 'allowed-source'),

apps/sim/lib/knowledge/search/queries.ts

@@ -1,5 +1,5 @@
 import { db } from '@sim/db'
-import { document, embedding, knowledgeConnector } from '@sim/db/schema'
+import { document, embedding, embeddingSearch, knowledgeConnector } from '@sim/db/schema'
 import { createLogger } from '@sim/logger'
 import { getErrorMessage, getPostgresErrorCode } from '@sim/utils/errors'
 import { and, eq, inArray, isNull, type SQL, sql } from 'drizzle-orm'
@@ -448,10 +448,11 @@ const FTS_CONFIG = 'english'
 function getVisibilityConditions(
   access: KnowledgeAccessScope,
   filters?: WorkspaceSearchFilters,
-  accessCondition: SQL = knowledgeAccessCondition(access)
+  accessCondition: SQL = knowledgeAccessCondition(access),
+  enabledColumn: typeof embedding.enabled | typeof embeddingSearch.enabled = embedding.enabled
 ) {
   return [
-    eq(embedding.enabled, true),
+    eq(enabledColumn, true),
     eq(document.enabled, true),
     eq(document.processingStatus, 'completed'),
     eq(document.userExcluded, false),
@@ -833,6 +834,20 @@ async function selectLiveVectorResults(
         ),
         excludeSearchSources(excludedSources),
       ]
+      const candidateVisibility = [
+        ...getVisibilityConditions(
+          params.access,
+          params.filters,
+          knowledgeMetadataCandidateAccessCondition(params.access),
+          embeddingSearch.enabled
+        ),
+        excludeSearchSources(excludedSources),
+      ]
+      const visibleDocument = sql`LATERAL (
+        SELECT 1 FROM ${document}
+        WHERE ${and(eq(document.id, embeddingSearch.documentId), ...candidateVisibility)}
+        OFFSET 0
+      ) AS visible_document`
       /** Adding zero prevents an underfilled HNSW scan from being chosen again for fallback. */
       const exactPage = async (candidateIds?: string[]) => {
         const exactOffset = useExactRanking ? offset : 0
@@ -862,10 +877,10 @@ async function selectLiveVectorResults(
       /** Probe visibility without vector reads; revoked scopes must not detoast the corpus. */
       const probe = await runSearchQuery(params.budget, 'vector.probe', (executor) =>
         executor
-          .select({ id: embedding.id })
-          .from(embedding)
-          .innerJoin(document, eq(embedding.documentId, document.id))
-          .where(and(inArray(embedding.knowledgeBaseId, params.knowledgeBaseIds), ...visibility))
+          .select({ id: embeddingSearch.id })
+          .from(embeddingSearch)
+          .innerJoin(visibleDocument, sql`true`)
+          .where(inArray(embeddingSearch.knowledgeBaseId, params.knowledgeBaseIds))
           .limit(LIVE_SEARCH_PAGE_SIZE)
       )
       if (probe.length === 0) return { candidates: [], nextOffset: offset }
@@ -875,22 +890,16 @@ async function selectLiveVectorResults(
       if (useExactRanking) return exactPage()
       annotateSearchDiagnostics({
         vectorRanking: 'binary-rerank',
+        vectorCandidateStorage: 'stored-binary',
         vectorCandidateLimit: candidateLimit,
       })
       const identities = await withVectorScanSettings(
         (executor) =>
           executor
-            .select({ id: embedding.id })
-            .from(embedding)
-            .innerJoin(
-              sql`LATERAL (
-                SELECT 1 FROM ${document}
-                WHERE ${and(eq(document.id, embedding.documentId), ...visibility)}
-                OFFSET 0
-              ) AS visible_document`,
-              sql`true`
-            )
-            .where(inArray(embedding.knowledgeBaseId, params.knowledgeBaseIds))
+            .select({ id: embeddingSearch.id })
+            .from(embeddingSearch)
+            .innerJoin(visibleDocument, sql`true`)
+            .where(inArray(embeddingSearch.knowledgeBaseId, params.knowledgeBaseIds))
             .orderBy(candidateDistance)
             .limit(candidateLimit),
         params.budget,

apps/sim/lib/knowledge/vector-columns.test.ts

@@ -22,14 +22,13 @@ describe('embeddingVectorColumn', () => {
 })
 
 describe('embeddingCandidateDistance', () => {
-  it('matches the compact expression index for every stored width', () => {
+  it('orders stored compact vectors without recomputing full-vector quantization', () => {
     for (const width of KB_EMBEDDING_STORAGE_DIMENSIONS) {
       const rendered = embeddingCandidateDistance(width, '[1,2]').toSQL()
-      expect(rendered.sql).toContain('binary_quantize(?)::bit(?) <~>')
-      expect(rendered.params[0]).toBe(embeddingVectorColumn(width))
-      expect(rendered.params[1]).toEqual(rendered.params[3])
-      expect(JSON.stringify(rendered.params[1])).toContain(String(width))
-      expect(rendered.params[2]).toBe('[1,2]')
+      expect(rendered.sql).toContain('? <~> binary_quantize(?::vector)::bit(?)')
+      expect(rendered.params[0]).toBe(`embeddingSearch.binary${width === 1536 ? '' : width}`)
+      expect(rendered.params[1]).toBe('[1,2]')
+      expect(JSON.stringify(rendered.params[2])).toContain(String(width))
       expect(rendered.sql).not.toContain('<=>')
     }
   })

apps/sim/lib/knowledge/vector-columns.ts

@@ -5,7 +5,7 @@
  * can never drift apart.
  */
 
-import { embedding } from '@sim/db/schema'
+import { embedding, embeddingSearch } from '@sim/db/schema'
 import { type SQL, sql } from 'drizzle-orm'
 import type { KbEmbeddingDimensions } from '@/lib/knowledge/embedding-models'
 
@@ -25,6 +25,14 @@ const VECTOR_FIELD_BY_WIDTH = {
 
 const VECTOR_FIELDS = Object.values(VECTOR_FIELD_BY_WIDTH) as readonly VectorField[]
 
+const CANDIDATE_COLUMN_BY_WIDTH = {
+  384: embeddingSearch.binary384,
+  768: embeddingSearch.binary768,
+  1024: embeddingSearch.binary1024,
+  1536: embeddingSearch.binary,
+  3072: embeddingSearch.binary3072,
+} as const
+
 export function embeddingVectorColumn(dimensions: KbEmbeddingDimensions) {
   return embedding[VECTOR_FIELD_BY_WIDTH[dimensions]]
 }
@@ -76,5 +84,5 @@ export function embeddingCandidateDistance(
   queryVector: string
 ): SQL<number> {
   const width = sql.raw(String(dimensions))
-  return sql<number>`binary_quantize(${embeddingVectorColumn(dimensions)})::bit(${width}) <~> binary_quantize(${queryVector}::vector)::bit(${width})`
+  return sql<number>`${CANDIDATE_COLUMN_BY_WIDTH[dimensions]} <~> binary_quantize(${queryVector}::vector)::bit(${width})`
 }