improvement(knowledge): admit before embedding, budget the fill check, pass over emptied slices

Author: waleedlatif1

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

@@ -134,10 +134,12 @@ async function exactProjectionNeighbors(vector: number[], limit: number, readerC
     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
+  /** Unaliased: the distance expression qualifies its column with the table's own name. */
+  return db.execute<{ id: string }>(sql`SELECT ${embeddingSearch.id} AS id FROM ${embeddingSearch}
+    INNER JOIN document d ON d.id = ${embeddingSearch.documentId}
+    WHERE ${embeddingSearch.knowledgeBaseId} = ${ids.knowledgeBaseId}
+      AND ${embeddingSearch.enabled} ${readerClause ?? sql``}
+    ORDER BY (${distance}) + 0, ${embeddingSearch.id}
     LIMIT ${limit}`)
 }
 const captured: CapturedQuery[] = []
@@ -1197,7 +1199,7 @@ describe.skipIf(!enabled)('Knowledge search latency on a realistic indexed corpu
             const expected = await exactProjectionNeighbors(
               queryVector,
               actual.length,
-              sql`AND s.document_id IN (${sql.join(
+              sql`AND ${embeddingSearch.documentId} IN (${sql.join(
                 documentIds.map((id) => sql`${id}`),
                 sql`, `
               )})`

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()
   })
 

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

@@ -285,8 +285,9 @@ const searchKnowledgeUseCase = defineAuthorizedKnowledgeUseCase({
     )
     /**
      * Whether the organization may search at all, and whether this payer still may: neither
-     * depends on the query, so both run beside the embedding call below instead of ahead of it.
-     * A refusal still ends the search before any retrieval.
+     * depends on the query, so both run beside the scope and defaults reads below instead of
+     * ahead of them. Admission stays ahead of the embedding call, which a refused search must
+     * never make.
      */
     const admit = async (): Promise<BillingAttributionSnapshot | undefined> => {
       if (context.organizationId)
@@ -368,41 +369,41 @@ const searchKnowledgeUseCase = defineAuthorizedKnowledgeUseCase({
       : undefined
     const resultSecretRegistry = preparedRegistry ?? input.resultSecretRegistry
     input.signal?.throwIfAborted()
-    const [queryEmbedding, access, searchDefaults, billingAttribution, tagDefinitions] =
-      await Promise.all([
-        hasQuery
-          ? measureSearchStage('embedding', () =>
-              runWithKnowledgeModelInputProvenance(resultSecretRegistry, () =>
-                generateSearchEmbedding(
-                  input.query!,
-                  embeddingTarget!,
-                  context.workspaceId,
-                  input.signal
-                )
-              )
-            )
-          : Promise.resolve(null),
-        measureSearchStage('access_scope', () => context.access.get()),
-        measureSearchStage('defaults', () =>
-          resolveKnowledgeSearchDefaults({
-            workspaceId: context.workspaceId,
-            organizationId: context.organizationId,
+    const [access, searchDefaults, billingAttribution, tagDefinitions] = await Promise.all([
+      measureSearchStage('access_scope', () => context.access.get()),
+      measureSearchStage('defaults', () =>
+        resolveKnowledgeSearchDefaults({
+          workspaceId: context.workspaceId,
+          organizationId: context.organizationId,
 
-            /** The signed-in person, if any; never the billing owner or a key's creator. */
-            userId: resolvePrincipalSubjectUserId(principal) ?? undefined,
-            requestedMode: input.searchMode,
-          })
-        ),
-        admit(),
-        /** The tag names the results are labelled with depend on the bases alone. */
-        filters.length === 0
-          ? measureSearchStage('tag_definitions', () =>
-              getDocumentTagDefinitionsByKnowledgeBaseIds(knowledgeBaseIds)
-            )
-          : Promise.resolve(definitionsByKnowledgeBase),
-      ])
+          /** The signed-in person, if any; never the billing owner or a key's creator. */
+          userId: resolvePrincipalSubjectUserId(principal) ?? undefined,
+          requestedMode: input.searchMode,
+        })
+      ),
+      admit(),
+      /** The tag names the results are labelled with depend on the bases alone. */
+      filters.length === 0
+        ? measureSearchStage('tag_definitions', () =>
+            getDocumentTagDefinitionsByKnowledgeBaseIds(knowledgeBaseIds)
+          )
+        : Promise.resolve(definitionsByKnowledgeBase),
+    ])
     definitionsByKnowledgeBase = tagDefinitions
     input.signal?.throwIfAborted()
+    const queryEmbedding = hasQuery
+      ? await measureSearchStage('embedding', () =>
+          runWithKnowledgeModelInputProvenance(resultSecretRegistry, () =>
+            generateSearchEmbedding(
+              input.query!,
+              embeddingTarget!,
+              context.workspaceId,
+              input.signal
+            )
+          )
+        )
+      : null
+    input.signal?.throwIfAborted()
     annotateSearchDiagnostics({
       accessScopeKind: access.kind,
       searchMode: searchDefaults.searchMode,

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

@@ -51,6 +51,8 @@ export type SearchStage =
   | 'vector.settings'
   | 'vector.probe'
   | 'vector.page'
+  | 'vector.projection_filled'
+  | 'keyword.projection_filled'
   | 'vector.exact_candidates'
   | 'vector.exact'
   | 'vector.candidate_search'

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

@@ -559,6 +559,21 @@ describe('workspace-scoped vector retrieval', () => {
     expect(JSON.stringify(dbChainMockFns.leftJoin.mock.calls)).toContain('embeddingSearch')
   })
 
+  it('passes over a slice whose documents went away instead of ending the pool there', async () => {
+    const execute = dbChainMockFns.execute.getMockImplementation()!
+    let pages = 0
+    dbChainMockFns.execute.mockImplementation(async (query) => {
+      const statement = render(query)
+      /** The first slice's documents are gone; the next slice still has the readable rows. */
+      if (isPageStatement(statement.sql)) return pages++ === 0 ? [] : ranked
+      return execute(query)
+    })
+    queueTableRows(schemaMock.embedding, [...ranked].reverse())
+    expect((await handleVectorOnlySearch(params)).map((row) => row.id)).toEqual(['near', 'far'])
+    expect(pages).toBe(2)
+    expect(statements().filter((query) => isWalk(query.sql))).toHaveLength(1)
+  })
+
   it('sizes the pool to the pages asked for, doubling a pool the pages outran', () => {
     expect(vectorCandidatePoolLimit(20, undefined)).toBe(200)
     expect(vectorCandidatePoolLimit(150, undefined)).toBe(300)

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

@@ -117,13 +117,20 @@ const PROJECTION_FILLED_TTL_MS = 60_000
  * Whether the ranking projection still holds rows the backfill has not filled. Read off the
  * unfilled-rows index in microseconds and remembered briefly: the answer only ever changes once.
  */
-const projectionFilled = new LRUCache<ProjectionSourceAclTable, boolean>({
+const projectionFilled = new LRUCache<
+  ProjectionSourceAclTable,
+  boolean,
+  { budget: SearchBudget | undefined; stage: SearchStage }
+>({
   max: PROJECTION_SOURCE_ACL_TABLES.length,
   ttl: PROJECTION_FILLED_TTL_MS,
-  fetchMethod: async (projection) => {
+  /** The read that misses the cache spends the leg's own budget, like every other read of the leg. */
+  fetchMethod: async (projection, _stale, { context }) => {
     const table = projection === 'embedding_search' ? embeddingSearch : embeddingKeywordTin
-    const [row] = await db.execute<{ unfilled: boolean }>(sql`
+    const [row] = await runSearchQuery(context.budget, context.stage, (executor) =>
+      executor.execute<{ unfilled: boolean }>(sql`
       SELECT EXISTS (SELECT 1 FROM ${table} WHERE ${table.acl} IS NULL) AS unfilled`)
+    )
     return !row?.unfilled
   },
 })
@@ -1770,7 +1777,11 @@ async function selectVectorResults(params: SearchParams): Promise<SearchResult[]
         }
         let selected: Array<{ id: string }>
         const plan = params.access.kind === 'user' ? params.accessPlan : undefined
-        const filled = plan ? ((await projectionFilled.fetch('embedding_search')) ?? false) : false
+        const filled = plan
+          ? ((await projectionFilled.fetch('embedding_search', {
+              context: { budget: params.budget, stage: 'vector.projection_filled' },
+            })) ?? false)
+          : false
         /**
          * A source the caller is a member of that has its own index is walked on its own, which
          * beats ranking it exactly once it is large enough to have earned that index.
@@ -1891,29 +1902,32 @@ async function selectVectorResults(params: SearchParams): Promise<SearchResult[]
           vectorCandidateScan: selected.length < candidateLimit ? 'underfilled' : 'planned',
         })
       }
-      const slice = candidatePool.ids.slice(offset, offset + limit)
-      if (!slice.length) return { candidates: [], nextOffset: offset }
       /**
        * The walk's order is the page's order: it ranked on the stored halfvec, and rescoring the
        * pool against the original vectors read one out-of-line vector per candidate from storage
        * no cache holds, seconds on a query nobody had run before. Only the page's identities are
-       * read here; the full read predicate follows at hydration, as before.
+       * read here; the full read predicate follows at hydration, as before. A slice whose
+       * documents all went away since the walk is passed over, not mistaken for the pool's end.
        */
-      const ranked = new Map(slice.map((candidate, index) => [candidate.id, index]))
-      const identities = await runSearchQuery(params.budget, 'vector.page', (executor) =>
-        executor.execute<SearchReadCandidate>(sql`
+      for (let start = offset; start < candidatePool.ids.length; start += limit) {
+        const slice = candidatePool.ids.slice(start, start + limit)
+        const ranked = new Map(slice.map((candidate, index) => [candidate.id, index]))
+        const identities = await runSearchQuery(params.budget, 'vector.page', (executor) =>
+          executor.execute<SearchReadCandidate>(sql`
           SELECT ${embeddingSearch.id} AS id, ${document.id} AS "documentId",
             ${document.connectorId} AS "connectorId"
           FROM ${embeddingSearch}
           INNER JOIN ${document} ON ${document.id} = ${embeddingSearch.documentId}
           WHERE ${embeddingSearch.id} = ANY(${textArrayLiteral(slice.map((candidate) => candidate.id))})
         `)
-      )
-      const page = [...identities].sort((a, b) => (ranked.get(a.id) ?? 0) - (ranked.get(b.id) ?? 0))
-      return {
-        candidates: page,
-        nextOffset: offset + page.length,
+        )
+        if (!identities.length) continue
+        const page = [...identities].sort(
+          (a, b) => (ranked.get(a.id) ?? 0) - (ranked.get(b.id) ?? 0)
+        )
+        return { candidates: page, nextOffset: start + slice.length }
       }
+      return { candidates: [], nextOffset: candidatePool.ids.length }
     },
     hydrate: (ids, authorized) =>
       hydrateSearchCandidates(
@@ -2027,7 +2041,9 @@ export async function executeKeywordSearch(params: KeywordSearchParams): Promise
     /** A filled projection decides readability on the ranked row alone; none of its rows needs the document. */
     const tinFilled =
       accessPlan && tinQuery
-        ? ((await projectionFilled.fetch('embedding_keyword_tin')) ?? false)
+        ? ((await projectionFilled.fetch('embedding_keyword_tin', {
+            context: { budget: params.budget, stage: 'keyword.projection_filled' },
+          })) ?? false)
         : false
     /** The projection predicate over the ranked CTE's mirrored columns, plus any excluded source. */
     const onRowKeywordVisibility = (excludedSources: readonly string[]) =>

apps/sim/lib/navigation/organization-rollout.test.ts

@@ -22,13 +22,18 @@ vi.mock('@/lib/billing/core/access', () => ({
   isOrganizationBillingBlocked: vi.fn().mockResolvedValue(false),
 }))
 
-import { requireOrganizationSearchAvailable } from '@/lib/knowledge/access/availability'
+import {
+  forgetKnowledgeAccessAvailability,
+  requireOrganizationSearchAvailable,
+} from '@/lib/knowledge/access/availability'
 import { resolveAppEntryPath } from '@/lib/navigation/resolve-app-entry'
 
 afterAll(resetEnvFlagsMock)
 
 describe('organization rollout during impersonation', () => {
   beforeEach(() => {
+    /** Each case answers the same organization differently; the memo must not carry one across. */
+    forgetKnowledgeAccessAvailability()
     vi.clearAllMocks()
     setEnvFlags({ isAppConfigEnabled: true, isHosted: true })
     mocks.landing.mockImplementation(async (userId: string) =>