diff --git a/apps/sim/lib/knowledge/__integration__/search-latency.integration.ts b/apps/sim/lib/knowledge/__integration__/search-latency.integration.ts index a5edda72119..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. */ @@ -294,7 +297,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(), @@ -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/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..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' @@ -902,57 +906,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. */