improvement(knowledge): rank the vector page on the projection and admit searches in parallel

Author: waleedlatif1

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

@@ -216,7 +216,8 @@ describe('Knowledge Search Utils', () => {
         const statement = (query as { toSQL: () => { sql: string } }).toSQL().sql
         if (statement.includes('AS visible')) return []
         if (statement.includes(') + 0 LIMIT')) return [{ id: 'first' }, { id: 'second' }]
-        if (statement.includes('WITH scored_search_candidates'))
+        /** The page reads the pool slice's identities; the walk's order is kept client-side. */
+        if (statement.includes('AS "connectorId"') && statement.includes('= ANY('))
           return [makeResult('second', 0.2), makeResult('first', 0.1)]
         return [{ id: 'doc-first' }, { id: 'doc-second' }]
       })
@@ -240,7 +241,7 @@ describe('Knowledge Search Utils', () => {
       const exact = dbChainMockFns.execute.mock.calls
         .map(([query]) => (query as { toSQL: () => { sql: string; params: unknown[] } }).toSQL())
         .find((statement) => statement.sql.includes(') + 0 LIMIT'))!
-      expect(exact.params).toContain(400)
+      expect(exact.params).toContain(200)
     })
 
     it('should throw error when no filters provided', async () => {

apps/sim/lib/billing/core/usage-gate-cache.ts

@@ -14,9 +14,11 @@ import { coalesceLocally } from '@/lib/concurrency/singleflight'
  * every uncached call. Bulk ingestion re-checks per document and knowledge
  * search checks per query. Staleness is bounded by this TTL and fails in the
  * harmless direction: a payer who crosses their limit keeps going for at most
- * this long, which charges nobody wrongly.
+ * this long, which charges nobody wrongly. Five minutes: the sum is a few
+ * hundred milliseconds for a busy payer, and a minute made every search after
+ * a pause pay it.
  */
-export const USAGE_GATE_TTL_MS = 60 * 1000
+export const USAGE_GATE_TTL_MS = 5 * 60 * 1000
 
 /**
  * Recent gate answers, admitted and refused, with `LRUCache` supplying the TTL

apps/sim/lib/knowledge/__integration__/kb-block-search.integration.ts

@@ -143,7 +143,7 @@ describe('API-key KB block fan-out', () => {
         expect(matching('hnsw.iterative_scan')).toHaveLength(bases.length)
         expect(matching('AS visible')).toHaveLength(bases.length)
         expect(matching(') + 0 LIMIT')).toHaveLength(bases.length)
-        expect(matching('scored_search_candidates')).toHaveLength(bases.length)
+        expect(matching('"embedding_search"."id" = ANY(')).toHaveLength(bases.length)
         /** The probe enumerates visible documents and reports saturation; it never ranks them. */
         expect(
           statements.filter(

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

@@ -8,6 +8,7 @@ import {
   credentialGroup,
   document,
   embedding,
+  embeddingSearch,
   knowledgeBase,
   knowledgeConnector,
   knowledgeConnectorMember,
@@ -19,7 +20,7 @@ import {
 } from '@sim/db/schema'
 import { createLogger, Logger } from '@sim/logger'
 import { generateId } from '@sim/utils/id'
-import { and, eq, inArray, sql } from 'drizzle-orm'
+import { and, eq, inArray, type SQL, sql } from 'drizzle-orm'
 import { NextRequest } from 'next/server'
 import { afterAll, beforeAll, describe, expect, it, type MockInstance, vi } from 'vitest'
 import { z } from 'zod'
@@ -41,13 +42,15 @@ import {
   seedKnowledgeMemberFixture,
 } from '@/lib/knowledge/__integration__/seed-source-access-fixture'
 import { type KnowledgeSearchTagFilter, searchKnowledge } from '@/lib/knowledge/application/search'
+import type { KbEmbeddingDimensions } from '@/lib/knowledge/embedding-models'
 import {
   SearchBudget,
   SearchDeadlineError,
   type SearchExecutor,
 } from '@/lib/knowledge/search/budget'
 import type { SearchStage } from '@/lib/knowledge/search/diagnostics'
 import type { WorkspaceSearchFilters } from '@/lib/knowledge/search/filters'
+import { embeddingCandidateDistance } from '@/lib/knowledge/vector-columns'
 import { POST as searchRoute } from '@/app/api/knowledge/search/route'
 import { ResolvedSecretTraceRegistry } from '@/executor/utils/resolved-secret-trace-registry'
 
@@ -120,6 +123,23 @@ function topicVector(topic = 0) {
   return vector.map((value) => value / magnitude)
 }
 const queryVector = topicVector()
+
+/**
+ * The exact nearest chunks on the projection's stored halfvec, which is what the page's order
+ * is measured against: the walk ranks on that column, and nothing rescores it.
+ */
+async function exactProjectionNeighbors(vector: number[], limit: number, readerClause?: SQL) {
+  const distance = embeddingCandidateDistance(
+    dimensions as KbEmbeddingDimensions,
+    JSON.stringify(vector),
+    'text-embedding-3-small'
+  )
+  return db.execute<{ id: string }>(sql`SELECT s.id FROM ${embeddingSearch} s
+    INNER JOIN document d ON d.id = s.document_id
+    WHERE s.knowledge_base_id = ${ids.knowledgeBaseId} AND s.enabled ${readerClause ?? sql``}
+    ORDER BY (${distance}) + 0, s.id
+    LIMIT ${limit}`)
+}
 const captured: CapturedQuery[] = []
 const report: Record<string, unknown> = {
   fixture: ids,
@@ -213,6 +233,10 @@ function explainNodes(node: ExplainNode): ExplainNode[] {
  * aliases its own lateral `scoped_chunk`, so this cannot match it, and matching on the rendered
  * clause casing would silently stop these assertions from running at all.
  */
+/** The vector page reads a pool slice's identities from the projection and its documents. */
+const VECTOR_PAGE_JOIN =
+  'INNER JOIN "document" ON "document"."id" = "embedding_search"."document_id"'
+
 function isVectorCandidateQuery(statement: string) {
   return statement.toLowerCase().includes(') as visible')
 }
@@ -476,12 +500,12 @@ async function sample(
         item.query.includes('limit') ||
         item.query.includes('CROSS JOIN LATERAL') ||
         isVectorCandidateQuery(item.query) ||
-        item.query.includes('WITH scored_search_candidates') ||
+        item.query.includes(VECTOR_PAGE_JOIN) ||
         item.query.includes('WITH matched_keyword_chunks'))
   )
   const plans: Array<
     CapturedQuery & {
-      kind: 'keyword' | 'vector' | 'rerank' | 'probe'
+      kind: 'keyword' | 'vector' | 'page' | 'probe'
       plan: z.infer<typeof explainSchema>
     }
   > = []
@@ -516,9 +540,8 @@ async function sample(
         ? 'keyword'
         : isVectorCandidateQuery(query.query)
           ? 'vector'
-          : query.query.includes('order by') ||
-              query.query.includes('WITH scored_search_candidates')
-            ? 'rerank'
+          : query.query.includes('order by') || query.query.includes(VECTOR_PAGE_JOIN)
+            ? 'page'
             : 'probe',
       query: query.query,
       parameters: query.parameters,
@@ -1004,13 +1027,10 @@ describe.skipIf(!enabled)('Knowledge search latency on a realistic indexed corpu
       expect(vectorPlans).toHaveLength(1)
       expect(vectorPlans[0].plan[0].Plan['Actual Rows']).toBeGreaterThan(0)
       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()
-      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(queryVector)}::vector) + 0, id
-        LIMIT ${actual.length}`)
+      expect(plans.some((plan) => plan.kind === 'page')).toBe(true)
+      const page = plans.find((plan) => plan.kind === 'page')!
+      const actual = await db.$client.unsafe(page.query, page.parameters).values()
+      const expected = await exactProjectionNeighbors(queryVector, 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)
@@ -1028,12 +1048,9 @@ describe.skipIf(!enabled)('Knowledge search latency on a realistic indexed corpu
       expectCompleteVectorSearch(diagnostics)
       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 page = plans.find((plan) => plan.kind === 'page')!
+      const actual = await db.$client.unsafe(page.query, page.parameters).values()
+      const expected = await exactProjectionNeighbors(topicVector(topic), 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)
@@ -1082,15 +1099,14 @@ describe.skipIf(!enabled)('Knowledge search latency on a realistic indexed corpu
         )
         expectCompleteVectorSearch(diagnostics)
         expect(result.data.results).toHaveLength(15)
-        const rerank = plans.find((plan) => plan.kind === 'rerank')!
-        expect(rerank).toBeDefined()
-        const actual = await db.$client.unsafe(rerank.query, rerank.parameters).values()
-        const expected = await db.execute<{ id: string }>(sql`SELECT e.id FROM embedding e
-          INNER JOIN document d ON d.id = e.document_id
-          WHERE e.knowledge_base_id = ${ids.knowledgeBaseId} AND e.enabled
-            AND d.acl @> ARRAY[${reader}]::text[]
-          ORDER BY (e.embedding <=> ${JSON.stringify(queryVector)}::vector) + 0, e.id
-          LIMIT ${actual.length}`)
+        const page = plans.find((plan) => plan.kind === 'page')!
+        expect(page).toBeDefined()
+        const actual = await db.$client.unsafe(page.query, page.parameters).values()
+        const expected = await exactProjectionNeighbors(
+          queryVector,
+          actual.length,
+          sql`AND d.acl @> ARRAY[${reader}]::text[]`
+        )
         expect(expected.length).toBeGreaterThan(0)
         const expectedIds = new Set(expected.map(({ id }) => id))
         const recall = actual.filter(([id]) => expectedIds.has(id)).length / expected.length
@@ -1128,9 +1144,10 @@ describe.skipIf(!enabled)('Knowledge search latency on a realistic indexed corpu
         expect(probe[0].query).not.toContain('<=>')
         expect(probe[0].plan[0].Plan['Actual Rows']).toBe(12)
         expect(assertIndexedChunkProbe(probe[0].plan[0].Plan)).toBe(documentIds.length)
-        const vector = plans.filter((plan) => plan.kind === 'rerank')
-        expect(vector).toHaveLength(1)
-        expect(vector[0].query).toContain('"embedding"."id" in')
+        /** The page reads the bounded ranking's identities from the projection, never the original vectors. */
+        const page = plans.filter((plan) => plan.kind === 'page')
+        expect(page).toHaveLength(1)
+        expect(page[0].query).not.toContain('"embedding"."embedding"')
       }
     } finally {
       await db
@@ -1175,16 +1192,16 @@ describe.skipIf(!enabled)('Knowledge search latency on a realistic indexed corpu
             count < HYBRID_CANDIDATE_LIMIT ? 0 : 1
           )
           if (count > HYBRID_CANDIDATE_LIMIT) {
-            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
-                AND document_id IN (${sql.join(
-                  documentIds.map((id) => sql`${id}`),
-                  sql`, `
-                )})
-              ORDER BY (embedding <=> ${JSON.stringify(queryVector)}::vector) + 0, id
-              LIMIT ${actual.length}`)
+            const page = plans.find((plan) => plan.kind === 'page')!
+            const actual = await db.$client.unsafe(page.query, page.parameters).values()
+            const expected = await exactProjectionNeighbors(
+              queryVector,
+              actual.length,
+              sql`AND s.document_id IN (${sql.join(
+                documentIds.map((id) => sql`${id}`),
+                sql`, `
+              )})`
+            )
             const expectedIds = new Set(expected.map(({ id }) => id))
             const recall = actual.filter(([id]) => expectedIds.has(id)).length / expected.length
             expect(recall).toBeGreaterThanOrEqual(0.95)
@@ -1352,7 +1369,7 @@ describe.skipIf(!enabled)('Knowledge search latency on a realistic indexed corpu
         )
         expectCompleteVectorSearch(diagnostics)
         expect(diagnostics.accessScopeKind).toBe('workspace')
-        expect(diagnostics.vectorRanking).toBe('candidate-rerank')
+        expect(diagnostics.vectorRanking).toBe('projection-walk')
         expect(result.data.results).toHaveLength(15)
         expect(plans.some((plan) => plan.kind === 'vector')).toBe(true)
         for (const row of result.data.results) {
@@ -1441,7 +1458,7 @@ describe.skipIf(!enabled)('Knowledge search latency on a realistic indexed corpu
         })
       )
       expectCompleteVectorSearch(tagged.diagnostics)
-      expect(tagged.diagnostics.vectorRanking).toBe('candidate-rerank')
+      expect(tagged.diagnostics.vectorRanking).toBe('projection-walk')
       expect(tagged.result.data.results).toHaveLength(15)
       for (const row of tagged.result.data.results) {
         const ordinal = Number(row.documentId.split('-doc-')[1])

apps/sim/lib/knowledge/access/availability.test.ts

@@ -29,13 +29,15 @@ vi.mock('@/lib/credential-groups/scoped-availability', () => ({
 }))
 
 import {
+  forgetKnowledgeAccessAvailability,
   requireOrganizationSearchAvailable,
   resolveKnowledgeAccessAvailability,
 } from '@/lib/knowledge/access/availability'
 
 describe('knowledge access availability ownership', () => {
   beforeEach(() => {
     vi.clearAllMocks()
+    forgetKnowledgeAccessAvailability()
     mocks.featureEnabled.mockResolvedValue(true)
     mocks.enterprise.mockResolvedValue(true)
     mocks.scopedGroups.mockResolvedValue(true)
@@ -59,6 +61,17 @@ describe('knowledge access availability ownership', () => {
     expect(mocks.workspaceGroups).not.toHaveBeenCalled()
   })
 
+  it('answers the same owner from one read for a minute', async () => {
+    await resolveKnowledgeAccessAvailability({ organizationId: 'org-1' })
+    await resolveKnowledgeAccessAvailability({ organizationId: 'org-1' })
+    expect(mocks.enterprise).toHaveBeenCalledTimes(1)
+    await resolveKnowledgeAccessAvailability({ organizationId: 'org-2' })
+    expect(mocks.enterprise).toHaveBeenCalledTimes(2)
+    forgetKnowledgeAccessAvailability()
+    await resolveKnowledgeAccessAvailability({ organizationId: 'org-1' })
+    expect(mocks.enterprise).toHaveBeenCalledTimes(3)
+  })
+
   it('keeps source mirroring independent from managed identity availability', async () => {
     mocks.scopedGroups.mockResolvedValue(false)
     await expect(resolveKnowledgeAccessAvailability({ organizationId: 'org-1' })).resolves.toEqual({

apps/sim/lib/knowledge/access/availability.ts

@@ -1,3 +1,4 @@
+import { LRUCache } from 'lru-cache'
 import { isOrganizationOnEnterprisePlan } from '@/lib/billing/core/subscription'
 import {
   getWorkspaceOwnerSubscriptionAccess,
@@ -39,11 +40,44 @@ export interface KnowledgeAccessAvailability {
   memberScoped: boolean
 }
 
+/**
+ * How long a resolved availability holds. A search resolves it three times over — the gate, the
+ * defaults, the reader's scope — each a subscription and a billing read; one read per owner per
+ * minute answers all of them, and a plan or flag change lands within the minute.
+ */
+const AVAILABILITY_TTL_MS = 60 * 1000
+
+const availabilityCache = new LRUCache<
+  string,
+  KnowledgeAccessAvailability,
+  KnowledgeMemberAccessContext
+>({
+  max: 10_000,
+  ttl: AVAILABILITY_TTL_MS,
+  fetchMethod: (_key, _stale, { context }) => readKnowledgeAccessAvailability(context),
+})
+
 export async function resolveKnowledgeAccessAvailability(
   context: KnowledgeMemberAccessContext
 ): Promise<KnowledgeAccessAvailability> {
   if (context.organizationId && context.workspaceId)
     throw new Error('Knowledge access requires one resource owner')
+  /** A caller that brings its own billing snapshot is answered from that snapshot, uncached. */
+  if (context.ownerBilling) return readKnowledgeAccessAvailability(context)
+  const key = `${context.organizationId ?? ''}|${context.workspaceId ?? ''}|${context.userId ?? ''}`
+  const availability = await availabilityCache.fetch(key, { context })
+  if (!availability) throw new Error('Knowledge access availability could not be resolved')
+  return availability
+}
+
+/** Forgets every resolved availability, for tests and for a settings change that must land now. */
+export function forgetKnowledgeAccessAvailability(): void {
+  availabilityCache.clear()
+}
+
+async function readKnowledgeAccessAvailability(
+  context: KnowledgeMemberAccessContext
+): Promise<KnowledgeAccessAvailability> {
   if (
     !(await isFeatureEnabled(
       'knowledge-member-access',

apps/sim/lib/knowledge/access/predicate.ts

@@ -357,7 +357,15 @@ export function projectionCandidateAccessCondition(
     documentId: AnyPgColumn | SQL
   },
   scope: KnowledgeAccessScope | SystemAccessScope,
-  plan: SearchAccessPlan
+  plan: SearchAccessPlan,
+  options: {
+    /**
+     * Whether every row of the projection carries its mirrored source and ACL. While the fill
+     * is under way, a row it has not reached is decided on its document; once it is complete no
+     * such row exists, and the predicate is the array test alone.
+     */
+    filled?: boolean
+  } = {}
 ): SQL {
   if (scope.kind === 'system') return sql`true`
   if (scope.tokens.length === 0) return sql`false`
@@ -374,12 +382,14 @@ export function projectionCandidateAccessCondition(
   const owned = plan.uploads
     ? sql`(${projection.connectorId} IS NULL OR ${inSources(mirrored)})`
     : inSources(mirrored)
+  const onRow = sql`(${projection.acl} && ${tokens} AND ${owned})`
+  if (options.filled) return onRow
   const unfilled = sql`(${projection.acl} IS NULL AND EXISTS (
     SELECT 1 FROM ${document}
     WHERE ${document.id} = ${projection.documentId}
       AND ${knowledgeCandidateAccessConditionForConnectors(scope, plan)}
   ))`
-  return sql`(${unfilled} OR (${projection.acl} && ${tokens} AND ${owned}))`
+  return sql`(${unfilled} OR ${onRow})`
 }
 
 /**

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

@@ -356,8 +356,8 @@ describe('knowledge search application use case', () => {
     ).rejects.toThrow('Search is not enabled for this organization')
     expect(mocks.recordActivity).not.toHaveBeenCalled()
     expect(mocks.requireOrganizationSearch).toHaveBeenCalledExactlyOnceWith('org-canonical')
+    /** The gate runs beside the embedding call; billing and retrieval still never start after a refusal. */
     expect(mocks.resolveBilling).not.toHaveBeenCalled()
-    expect(mocks.generateEmbedding).not.toHaveBeenCalled()
     expect(mocks.executeSearch).not.toHaveBeenCalled()
   })