From 96c67853014bd6ab13d22bc63dba7e1a5f2a4d63 Mon Sep 17 00:00:00 2001 From: Waleed Latif Date: Thu, 17 Sep 2026 21:15:16 -0700 Subject: [PATCH 1/2] fix(knowledge): keep vector candidate selection bounded An underfilled ANN traversal rescored the whole compact projection and joined the visible document set to it. That fallback is O(corpus): on a 132k-chunk index it took 1.9s, and on a corpus an order of magnitude larger it exceeds the retrieval budget, so the vector leg returned nothing at all rather than fewer rows. Widening the search has no affordable form here. Measured at 10% visibility on the same index, scanning 6.5k tuples instead of 1.5k took 5.1s and 9.7s on consecutive identical runs, and joining visibility before scoring took 8.4s because it turns a sequential scan into a random lookup per document. The bounded traversal itself costs ~115ms whether or not it fills. The traversal is now the whole candidate set. An underfilled one yields fewer candidates and is reported as such, which strictly beats a leg that times out. --- .../search-latency.integration.ts | 2 +- apps/sim/lib/knowledge/search/diagnostics.ts | 9 ++- apps/sim/lib/knowledge/search/queries.test.ts | 39 ++++------- apps/sim/lib/knowledge/search/queries.ts | 66 +++++++------------ 4 files changed, 44 insertions(+), 72 deletions(-) diff --git a/apps/sim/lib/knowledge/__integration__/search-latency.integration.ts b/apps/sim/lib/knowledge/__integration__/search-latency.integration.ts index a5edda72119..10ec7476bc8 100644 --- a/apps/sim/lib/knowledge/__integration__/search-latency.integration.ts +++ b/apps/sim/lib/knowledge/__integration__/search-latency.integration.ts @@ -294,7 +294,7 @@ const diagnosticSchema = z vectorBudgetMs: z.number().positive(), vectorCandidateDimensions: z.number().optional(), vectorCandidateLimit: z.number().optional(), - vectorCandidateScan: z.enum(['planned', 'filtered']).optional(), + vectorCandidateScan: z.enum(['planned', 'underfilled']).optional(), retrievalStatus: z.enum(['complete', 'partial']), timedOutLegs: z.array(z.enum(['vector', 'keyword', 'tags'])), toolResultBytes: z.number().int().nonnegative().optional(), diff --git a/apps/sim/lib/knowledge/search/diagnostics.ts b/apps/sim/lib/knowledge/search/diagnostics.ts index 0a6868f0288..293d0acb0e7 100644 --- a/apps/sim/lib/knowledge/search/diagnostics.ts +++ b/apps/sim/lib/knowledge/search/diagnostics.ts @@ -78,13 +78,16 @@ export interface SearchDiagnosticMetadata { embeddingDimensions?: number vectorRanking?: 'exact' | 'candidate-rerank' vectorCandidateStorage?: 'stored-halfvec' - /** Requested strategy, not an assertion about the physical index selected by PostgreSQL. */ - vectorCandidateScan?: 'planned' | 'filtered' + /** + * Whether the bounded traversal filled its candidate limit. `underfilled` means visibility + * removed enough neighbours that the rerank pool is smaller than requested, which lowers recall + * without widening the scan. Not an assertion about the physical index PostgreSQL selected. + */ + vectorCandidateScan?: 'planned' | 'underfilled' vectorBudgetMs?: number vectorCandidateLimit?: number vectorCandidateCount?: number vectorCandidateDimensions?: number - vectorInitialCandidateCount?: number resultCount?: number /** Tool output before the executor's final egress projection; counts only, never content. */ toolResultBytes?: number diff --git a/apps/sim/lib/knowledge/search/queries.test.ts b/apps/sim/lib/knowledge/search/queries.test.ts index dc84ba8ffed..a3a3360b0e4 100644 --- a/apps/sim/lib/knowledge/search/queries.test.ts +++ b/apps/sim/lib/knowledge/search/queries.test.ts @@ -370,7 +370,7 @@ describe('workspace-scoped vector retrieval', () => { if (failSettings) throw failSettings return [] } - if (statement.includes('WITH visible_search_documents')) { + if (statement.includes('AS visible')) { if (failCandidates) throw failCandidates return candidates } @@ -434,9 +434,7 @@ describe('workspace-scoped vector retrieval', () => { it('uses compact candidates for a large KB and applies full workspace access before its limit', async () => { queueTableRows(schemaMock.embedding, [...ranked].reverse()) expect((await handleVectorOnlySearch(params)).map((row) => row.id)).toEqual(['near', 'far']) - const candidate = statements().find((query) => - query.sql.includes('WITH visible_search_documents') - )! + const candidate = statements().find((query) => query.sql.includes('AS visible'))! expect(candidate.sql).toContain('CROSS JOIN LATERAL') expect(candidate.sql).toContain('LIMIT 1') const serialized = JSON.stringify(candidate) @@ -501,9 +499,7 @@ describe('workspace-scoped vector retrieval', () => { const rows = await handleVectorOnlySearch({ ...params, topK: 1 }) expect(rows.map((row) => row.id)).toEqual(['far']) - expect( - statements().filter((query) => query.sql.includes('WITH visible_search_documents')) - ).toHaveLength(1) + expect(statements().filter((query) => query.sql.includes('AS visible'))).toHaveLength(1) expect(getForConnectors).not.toHaveBeenCalled() }) @@ -516,9 +512,7 @@ describe('workspace-scoped vector retrieval', () => { }) expect(rows.map((row) => row.id)).toEqual(['near', 'far']) expect(Object.keys(dbChainMockFns.select.mock.calls[0][0])).toEqual(['id']) - const candidate = statements().find((query) => - query.sql.includes('WITH visible_search_documents') - )! + const candidate = statements().find((query) => query.sql.includes('AS visible'))! expect(JSON.stringify(candidate)).toContain('common') expect(JSON.stringify(candidate)).toContain(String(schemaMock.embedding.tag1)) expect(JSON.stringify(dbChainMockFns.where.mock.calls.at(-1)![0])).toContain('common') @@ -534,9 +528,7 @@ describe('workspace-scoped vector retrieval', () => { const rows = await handleVectorOnlySearch({ ...params, knowledgeBaseIds }) expect(rows.map((row) => row.id)).toEqual(['near', 'far']) expect(rows.every((row) => row.knowledgeBaseId === 'kb-1')).toBe(true) - const candidateQueries = statements().filter((query) => - query.sql.includes('WITH visible_search_documents') - ) + const candidateQueries = statements().filter((query) => query.sql.includes('AS visible')) expect(candidateQueries).toHaveLength(1) for (const id of knowledgeBaseIds) expect(JSON.stringify(candidateQueries[0])).toContain(id) expect(dbChainMockFns.transaction).toHaveBeenCalledOnce() @@ -616,9 +608,7 @@ describe('workspace-scoped vector retrieval', () => { expect(statements().filter((query) => query.sql.includes('hnsw.iterative_scan'))).toHaveLength( 1 ) - const queries = statements().filter((query) => - query.sql.includes('WITH visible_search_documents') - ) + const queries = statements().filter((query) => query.sql.includes('AS visible')) expect(queries).toHaveLength(2) expect(JSON.stringify(queries[0])).toBe(JSON.stringify(queries[1])) await vi.advanceTimersByTimeAsync(10 * 60 * 1000 + 1) @@ -807,7 +797,7 @@ describe('live repository authorization follows ranked candidates', () => { dbChainMockFns.execute.mockImplementation(async (query) => render(query).sql.includes('SELECT scoped_chunk.id') ? (probePages.shift() ?? []) - : render(query).sql.includes('WITH visible_search_documents') + : render(query).sql.includes('AS visible') ? (candidatePages.shift() ?? []) : render(query).sql.includes('WITH scored_search_candidates') ? (rerankPages.shift() ?? []) @@ -841,9 +831,8 @@ describe('live repository authorization follows ranked candidates', () => { }) expect(rows.map((row) => row.id)).toEqual(['near']) const candidateQuery = dbChainMockFns.execute.mock.calls.find(([query]) => - render(query).sql.includes('WITH visible_search_documents') + render(query).sql.includes('AS visible') )![0] - expect(render(candidateQuery).sql).toContain('MATERIALIZED') expect(render(candidateQuery).sql).toContain('CROSS JOIN LATERAL') expect(render(candidateQuery).sql).toContain('LIMIT 1') expect(JSON.stringify(candidateQuery)).toContain('required_clause') @@ -922,7 +911,7 @@ describe('live repository authorization follows ranked candidates', () => { } ) - it('scans the filtered projection when ANN cannot fill its limit', async () => { + it('keeps an underfilled ANN result instead of rescoring the whole projection', async () => { probePages.push(Array.from({ length: 400 }, (_, index) => ({ id: `probe-${index}` }))) queueCandidates([{ id: 'selected' }], 1) queueRerank([candidate('selected', 'allowed-source')]) @@ -933,12 +922,12 @@ describe('live repository authorization follows ranked candidates', () => { { id: 'selected', content: 'Verified fallback', distance: 0.1 }, ]) const candidateQuery = dbChainMockFns.execute.mock.calls.find(([query]) => - render(query).sql.includes('WITH visible_search_documents') + render(query).sql.includes('AS visible') )![0] - expect(render(candidateQuery).sql).toContain('UNION ALL') - expect(render(candidateQuery).sql).toContain('+ 0') - expect(render(candidateQuery).sql).toContain('filtered_scores AS MATERIALIZED') - expect(render(candidateQuery).sql).toContain('ORDER BY filtered_scores.distance + 0') + /** Widening the scan on underfill is what made this leg exceed its budget on a large corpus. */ + expect(render(candidateQuery).sql).not.toContain('UNION ALL') + expect(render(candidateQuery).sql).not.toContain('filtered_scores') + expect(render(candidateQuery).sql).toContain('CROSS JOIN LATERAL') expect(JSON.stringify(dbChainMockFns.where.mock.calls.at(-1)![0])).toContain( 'github_read_grant' ) diff --git a/apps/sim/lib/knowledge/search/queries.ts b/apps/sim/lib/knowledge/search/queries.ts index e2b7dc91c17..f24cc7a282d 100644 --- a/apps/sim/lib/knowledge/search/queries.ts +++ b/apps/sim/lib/knowledge/search/queries.ts @@ -902,57 +902,37 @@ async function selectVectorResults(params: SearchParams): Promise - executor.execute<{ id: string; initial_count: number }>(sql` - WITH visible_search_documents AS MATERIALIZED ( - SELECT ${document.id} AS id FROM ${document} - WHERE ${and(...candidateDocumentVisibility)} - ), initial_candidates AS MATERIALIZED ( - SELECT ${embeddingSearch.id} AS id FROM ${embeddingSearch} - CROSS JOIN LATERAL ( - SELECT 1 FROM ${document} - WHERE ${and(eq(document.id, embeddingSearch.documentId), ...candidateDocumentVisibility, candidateTagCondition)} - LIMIT 1 - ) AS visible - WHERE ${and( - inArray(embeddingSearch.knowledgeBaseId, params.knowledgeBaseIds), - eq(embeddingSearch.enabled, true) - )} - ORDER BY ${candidateDistance} LIMIT ${candidateLimit} - ), filtered_scores AS MATERIALIZED ( - SELECT ${embeddingSearch.id} AS id, ${embeddingSearch.documentId} AS document_id, - ${candidateDistance} AS distance FROM ${embeddingSearch} - WHERE ${and( - inArray(embeddingSearch.knowledgeBaseId, params.knowledgeBaseIds), - eq(embeddingSearch.enabled, true), - candidateTagCondition - )} - AND (SELECT count(*) FROM initial_candidates) < ${candidateLimit} - ), candidates AS ( - SELECT id FROM initial_candidates - WHERE (SELECT count(*) FROM initial_candidates) >= ${candidateLimit} - UNION ALL ( - SELECT filtered_scores.id FROM filtered_scores - INNER JOIN visible_search_documents ON visible_search_documents.id = filtered_scores.document_id - WHERE (SELECT count(*) FROM initial_candidates) < ${candidateLimit} - ORDER BY filtered_scores.distance + 0, filtered_scores.id - LIMIT ${candidateLimit} - ) - ) SELECT id, (SELECT count(*)::int FROM initial_candidates) AS initial_count FROM candidates + executor.execute<{ id: string }>(sql` + SELECT ${embeddingSearch.id} AS id FROM ${embeddingSearch} + CROSS JOIN LATERAL ( + SELECT 1 FROM ${document} + WHERE ${and(eq(document.id, embeddingSearch.documentId), ...candidateDocumentVisibility, candidateTagCondition)} + LIMIT 1 + ) AS visible + WHERE ${and( + inArray(embeddingSearch.knowledgeBaseId, params.knowledgeBaseIds), + eq(embeddingSearch.enabled, true) + )} + ORDER BY ${candidateDistance} LIMIT ${candidateLimit} `), params.budget ) - const initialCount = identities[0]?.initial_count ?? 0 annotateSearchDiagnostics({ vectorCandidateCount: identities.length, - vectorInitialCandidateCount: initialCount, - vectorCandidateScan: initialCount < candidateLimit ? 'filtered' : 'planned', + vectorCandidateScan: identities.length < candidateLimit ? 'underfilled' : 'planned', }) if (!identities.length) return { candidates: [], nextOffset: offset } /** Score each bounded candidate once; sorting the materialized scalar cannot invoke HNSW again. */ From 1af545d763baa446206c7a29ed191dc936ecf05f Mon Sep 17 00:00:00 2001 From: Waleed Latif Date: Thu, 17 Sep 2026 21:24:28 -0700 Subject: [PATCH 2/2] fix(knowledge): match the latency plan assertions to the bounded traversal The plan assertion still required the removed CTEs, and the predicates selecting which captured query to assert against matched the old CTE name with a lowercase fallback the rendered statement never produces. The candidate assertions would have stopped running rather than failing, so they now key on the visibility lateral through one shared predicate. The plan check drops the fallback-specific expectations and gains the one that guards this change: the traversal must never reach the projection by document lookup or sequential scan, at any candidate count. --- .../search-latency.integration.ts | 61 +++++++++---------- apps/sim/lib/knowledge/search/queries.ts | 6 +- 2 files changed, 35 insertions(+), 32 deletions(-) diff --git a/apps/sim/lib/knowledge/__integration__/search-latency.integration.ts b/apps/sim/lib/knowledge/__integration__/search-latency.integration.ts index 10ec7476bc8..e555d9af3e7 100644 --- a/apps/sim/lib/knowledge/__integration__/search-latency.integration.ts +++ b/apps/sim/lib/knowledge/__integration__/search-latency.integration.ts @@ -208,6 +208,15 @@ function explainNodes(node: ExplainNode): ExplainNode[] { } /** Broad ranking must stop the ordered ANN scan instead of sorting every accessible chunk. */ +/** + * The bounded ANN traversal, identified by the visibility lateral it alone carries. The probe + * 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. + */ +function isVectorCandidateQuery(statement: string) { + return statement.toLowerCase().includes(') as visible') +} + function assertIndexedCandidates( plan: ExplainNode, candidateLimit: number, @@ -218,31 +227,25 @@ function assertIndexedCandidates( ? 'embedding_search_cosine_hnsw_idx' : `embedding_search_${width}_cosine_hnsw_idx` const nodes = explainNodes(plan) - const initial = nodes.find((node) => node['Subplan Name'] === 'CTE initial_candidates') - expect(initial).toBeDefined() - const candidateNodes = explainNodes(initial!) - expect( - candidateNodes.some((node) => node['Index Name'] === indexName && node['Actual Loops'] > 0) - ).toBe(true) - expect(candidateNodes.some((node) => node['Node Type'] === 'Sort')).toBe(false) + expect(nodes.some((node) => node['Index Name'] === indexName && node['Actual Loops'] > 0)).toBe( + true + ) + /** The graph walk supplies the order, so a Sort here means the index ordering was discarded. */ + expect(nodes.some((node) => node['Node Type'] === 'Sort')).toBe(false) + /** + * The traversal is the whole candidate set. Reaching the projection by document lookup or by + * sequential scan is the corpus-wide rescan this query exists to avoid, at any candidate count. + */ + expect(nodes.some((node) => node['Index Name'] === 'embedding_search_document_lookup_idx')).toBe( + false + ) expect( - candidateNodes.some((node) => node['Index Name'] === 'embedding_search_document_lookup_idx') + nodes.some( + (node) => node['Relation Name'] === 'embedding_search' && node['Node Type'] === 'Seq Scan' + ) ).toBe(false) - const filtered = nodes.find((node) => node['Subplan Name'] === 'CTE filtered_scores') - expect(filtered).toBeDefined() - expect(filtered!.Output).toHaveLength(3) - expect(filtered!.Output![2]).toContain('<=>') - if (initial!['Actual Rows'] >= candidateLimit) { - for (const node of nodes.filter( - (item) => - item['Subplan Name'] === 'CTE visible_search_documents' || - item['CTE Name'] === 'visible_search_documents' || - item['Subplan Name'] === 'CTE filtered_scores' || - item['CTE Name'] === 'filtered_scores' - )) { - expect(node['Actual Loops']).toBe(0) - } - } + const traversed = nodes.find((node) => node['Index Name'] === indexName)! + expect(traversed['Actual Rows']).toBeLessThanOrEqual(candidateLimit) } /** Small scopes must seek chunk metadata by document without reading the full vector projection. */ @@ -472,7 +475,7 @@ async function sample( (item.query.includes('order by') || item.query.includes('limit') || item.query.includes('CROSS JOIN LATERAL') || - item.query.includes('WITH visible_search_documents') || + isVectorCandidateQuery(item.query) || item.query.includes('WITH scored_search_candidates') || item.query.includes('WITH visible_keyword_documents')) ) @@ -497,10 +500,7 @@ async function sample( await tx.unsafe('SET LOCAL jit = off') await tx.unsafe("SET LOCAL hnsw.iterative_scan = 'relaxed_order'") await tx.unsafe('SET LOCAL hnsw.max_scan_tuples = 20000') - if ( - query.query.includes('WITH visible_search_documents') || - (query.query.includes('from "embedding_search"') && query.query.includes('order by')) - ) { + if (isVectorCandidateQuery(query.query)) { await tx.unsafe('SET LOCAL hnsw.max_scan_tuples = 1000') await tx.unsafe('SET LOCAL hnsw.ef_search = 1000') await tx.unsafe('SET LOCAL hnsw.scan_mem_multiplier = 2') @@ -514,8 +514,7 @@ async function sample( plans.push({ kind: query.query.includes('keyword_rank') ? 'keyword' - : query.query.includes('WITH visible_search_documents') || - (query.query.includes('from "embedding_search"') && query.query.includes('order by')) + : isVectorCandidateQuery(query.query) ? 'vector' : query.query.includes('order by') || query.query.includes('WITH scored_search_candidates') @@ -526,7 +525,7 @@ async function sample( plan: parsedPlan, }) saveReport() - if (query.query.includes('WITH visible_search_documents')) { + if (isVectorCandidateQuery(query.query)) { const width = diagnostics.vectorCandidateDimensions! expect(query.query).toContain( `"embedding_search"."${width === 1536 ? 'vector' : `vector_${width}`}"` diff --git a/apps/sim/lib/knowledge/search/queries.ts b/apps/sim/lib/knowledge/search/queries.ts index f24cc7a282d..447240ee874 100644 --- a/apps/sim/lib/knowledge/search/queries.ts +++ b/apps/sim/lib/knowledge/search/queries.ts @@ -50,8 +50,12 @@ const UNDEFINED_OBJECT_SQLSTATE = '42704' /** Bound candidate pages retained while live permissions are checked. */ const MAX_AUTHORIZED_SEARCH_CANDIDATES = 20_000 /** - * Stop a permission-starved graph walk early enough to scan the filtered projection instead. + * Bounds a permission-starved graph walk, which returns fewer candidates rather than widening. * This approximate iterative-visit threshold excludes pgvector's initial scan; it is not a row limit. + * + * Raising it trades recall for latency far more steeply than its size suggests: on a 132k-chunk + * index at 10% visibility, visiting 6.5k tuples instead of 1.5k took 5.1s and 9.7s on consecutive + * identical runs, against ~115ms for the bounded walk. Re-measure before changing it. */ const CANDIDATE_HNSW_MAX_SCAN_TUPLES = '1000' const CANDIDATE_HNSW_EF_SEARCH = '1000'