diff --git a/apps/sim/app/api/knowledge/search/utils.test.ts b/apps/sim/app/api/knowledge/search/utils.test.ts index 7106308b9d5..83217cbd6b2 100644 --- a/apps/sim/app/api/knowledge/search/utils.test.ts +++ b/apps/sim/app/api/knowledge/search/utils.test.ts @@ -212,11 +212,14 @@ describe('Knowledge Search Utils', () => { describe('handleTagAndVectorSearch', () => { it('returns only bounded ranked rows without first materializing every matching tag ID', async () => { resetDbChainMock() - queueTableRows( - schemaMock.embedding, - Array.from({ length: 201 }, (_, index) => ({ id: `candidate-${index}` })) - ) - queueTableRows(schemaMock.embedding, [makeResult('second', 0.2), makeResult('first', 0.1)]) + dbChainMockFns.execute.mockImplementation(async (query) => { + 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')) + return [makeResult('second', 0.2), makeResult('first', 0.1)] + return [{ id: 'doc-first' }, { id: 'doc-second' }] + }) queueTableRows(schemaMock.embedding, [makeResult('second', 0.2), makeResult('first', 0.1)]) const results = await handleTagAndVectorSearch({ @@ -231,11 +234,13 @@ describe('Knowledge Search Utils', () => { }) expect(results.map((row) => row.id)).toEqual(['first', 'second']) - expect(dbChainMockFns.select).toHaveBeenCalledTimes(3) - expect(Object.keys(dbChainMockFns.select.mock.calls[0][0])).toEqual(['id']) - expect(dbChainMockFns.limit).toHaveBeenNthCalledWith(1, 400) - expect(dbChainMockFns.select.mock.calls[1][0]).toHaveProperty('distance') - expect(dbChainMockFns.limit).toHaveBeenCalledWith(20) + /** Only hydration reads through the query builder; ranking never materializes tag IDs. */ + expect(dbChainMockFns.select).toHaveBeenCalledTimes(1) + expect(dbChainMockFns.select.mock.calls[0][0]).toHaveProperty('distance') + 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) }) it('should throw error when no filters provided', async () => { @@ -552,13 +557,13 @@ describe('Knowledge Search Utils', () => { }) expect(results.map((r) => r.id)).toEqual(['vector-hit']) - expect(dbChainMockFns.select).toHaveBeenCalledTimes(2) + expect(dbChainMockFns.select).toHaveBeenCalledTimes(1) }) it('runs both legs and fuses them in hybrid mode', async () => { /** - * The raw vector probe does not consume a table chain. Keyword ranking and - * hydration complete before vector exact ranking and content hydration. + * Vector ranking is raw SQL throughout and consumes no table chain. Keyword ranking and + * hydration complete before vector content hydration. */ dbChainMockFns.execute.mockResolvedValue([{ id: 'vector-hit' }]) queueTableRows(schemaMock.embedding, [{ id: 'keyword-hit', keywordRank: 0.9 }]) @@ -576,7 +581,7 @@ describe('Knowledge Search Utils', () => { }) expect(results.map((r) => r.id).sort()).toEqual(['keyword-hit', 'vector-hit']) - expect(dbChainMockFns.select).toHaveBeenCalledTimes(4) + expect(dbChainMockFns.select).toHaveBeenCalledTimes(3) }) it('propagates unexpected keyword errors after the vector leg finishes', async () => { @@ -601,7 +606,7 @@ describe('Knowledge Search Utils', () => { queryVector: { vector: JSON.stringify(TEST_EMBEDDING), dimensions: 1536 }, }) ).rejects.toBe(failure) - expect(dbChainMockFns.select).toHaveBeenCalledTimes(3) + expect(dbChainMockFns.select).toHaveBeenCalledTimes(2) }) it('skips both query legs when only tag filters are provided', async () => { diff --git a/apps/sim/lib/knowledge/__integration__/kb-block-search.integration.ts b/apps/sim/lib/knowledge/__integration__/kb-block-search.integration.ts index b54266037d7..5abf70817cd 100644 --- a/apps/sim/lib/knowledge/__integration__/kb-block-search.integration.ts +++ b/apps/sim/lib/knowledge/__integration__/kb-block-search.integration.ts @@ -91,7 +91,7 @@ describe('API-key KB block fan-out', () => { const previousDebug = db.$client.options.debug const statements: string[] = [] db.$client.options.debug = (_connection, query) => { - if (statements.length < 250) statements.push(query) + if (statements.length < 1000) statements.push(query) } try { const results = await Promise.all( @@ -124,9 +124,26 @@ describe('API-key KB block fan-out', () => { expect(result.rows[0].knowledgeBaseId).toBe(bases[index].id) expect(result.rows[0].distance).toBeCloseTo(0) } - expect(statements.filter((query) => query.includes('statement_timeout'))).toHaveLength(54) - expect(statements.filter((query) => query.includes('+ 0'))).toHaveLength(18) - expect(statements.some((query) => query.includes('hnsw.iterative_scan'))).toBe(false) + const matching = (fragment: string) => + statements.filter((query) => query.includes(fragment)) + /** + * Every statement runs under the leg's deadline: the candidate search reinstates it after + * tuning the scan, and the probe, the exact ranking, the rerank and hydration each open + * with one of their own. + */ + expect(matching('statement_timeout')).toHaveLength(bases.length * 6) + /** + * A scope this small leaves the bounded traversal short of its candidate limit, so every + * search probes once and rescues once — never a widening retry loop. + */ + 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) + /** The probe enumerates visible documents; it never ranks them. */ + expect( + statements.filter((query) => query.includes('AS id FROM') && !query.includes('ORDER BY')) + ).toHaveLength(bases.length) } finally { db.$client.options.debug = previousDebug } diff --git a/apps/sim/lib/knowledge/search/budget.test.ts b/apps/sim/lib/knowledge/search/budget.test.ts index 941fbdbd98e..40f6d7a791e 100644 --- a/apps/sim/lib/knowledge/search/budget.test.ts +++ b/apps/sim/lib/knowledge/search/budget.test.ts @@ -44,6 +44,23 @@ describe('search SQL deadline', () => { } }) + it('caps a step without letting it shorten or outlive the leg it came from', () => { + vi.spyOn(performance, 'now').mockReturnValue(0) + const controller = new AbortController() + const leg = new SearchBudget('vector', 1000, controller.signal) + const step = leg.capped(600) + expect(step.deadline).toBe(600) + expect(step.leg).toBe('vector') + expect(step.signal).toBe(controller.signal) + /** Spending the step is the step's own business; the leg keeps its deadline and its verdict. */ + step.isTimeout(new SearchDeadlineError()) + expect(step.timedOut).toBe(true) + expect(leg.timedOut).toBe(false) + expect(leg.remaining()).toBe(1000) + /** A step longer than what the leg has left cannot extend it. */ + expect(new SearchBudget('vector', 200).capped(600).deadline).toBe(200) + }) + it('preserves cancellation and unexpected errors instead of labeling them incomplete evidence', async () => { const controller = new AbortController() const budget = new SearchBudget('keyword', performance.now() - 1, controller.signal) diff --git a/apps/sim/lib/knowledge/search/budget.ts b/apps/sim/lib/knowledge/search/budget.ts index a80279ec3e1..fcb1ac9e5cd 100644 --- a/apps/sim/lib/knowledge/search/budget.ts +++ b/apps/sim/lib/knowledge/search/budget.ts @@ -39,6 +39,18 @@ export class SearchBudget { return remaining } + /** + * A shorter deadline for one step of the leg, carrying the same leg and cancellation signal. + * Spending it is the step's own business: the leg's deadline is untouched and still governs. + */ + capped(milliseconds: number): SearchBudget { + return new SearchBudget( + this.leg, + Math.min(this.deadline, performance.now() + milliseconds), + this.signal + ) + } + isTimeout(error: unknown): boolean { this.signal?.throwIfAborted() if (error instanceof SearchDeadlineError || getPostgresErrorCode(error) === '57014') { diff --git a/apps/sim/lib/knowledge/search/diagnostics.ts b/apps/sim/lib/knowledge/search/diagnostics.ts index 293d0acb0e7..8406c131027 100644 --- a/apps/sim/lib/knowledge/search/diagnostics.ts +++ b/apps/sim/lib/knowledge/search/diagnostics.ts @@ -47,6 +47,7 @@ export type SearchStage = | 'vector.settings' | 'vector.probe' | 'vector.rerank' + | 'vector.exact_candidates' | 'vector.exact' | 'vector.candidate_search' | 'source_overview' @@ -76,7 +77,7 @@ export interface SearchDiagnosticMetadata { searchMode?: 'hybrid' | 'vector' boostRecency?: boolean embeddingDimensions?: number - vectorRanking?: 'exact' | 'candidate-rerank' + vectorRanking?: 'exact' | 'exact-candidates' | 'candidate-rerank' vectorCandidateStorage?: 'stored-halfvec' /** * Whether the bounded traversal filled its candidate limit. `underfilled` means visibility @@ -86,6 +87,8 @@ export interface SearchDiagnosticMetadata { vectorCandidateScan?: 'planned' | 'underfilled' vectorBudgetMs?: number vectorCandidateLimit?: number + /** Visible documents the tractability probe enumerated, capped at its own document limit. */ + vectorProbeDocumentCount?: number vectorCandidateCount?: number vectorCandidateDimensions?: number resultCount?: number diff --git a/apps/sim/lib/knowledge/search/queries.test.ts b/apps/sim/lib/knowledge/search/queries.test.ts index a3a3360b0e4..ce1f0031c63 100644 --- a/apps/sim/lib/knowledge/search/queries.test.ts +++ b/apps/sim/lib/knowledge/search/queries.test.ts @@ -30,6 +30,7 @@ import { handleVectorOnlySearch, retrieveKnowledgeSearch, type SearchParams, + VECTOR_PROBE_DOCUMENT_LIMIT, } from '@/lib/knowledge/search/queries' import type { StructuredFilter } from '@/lib/knowledge/types' @@ -57,7 +58,7 @@ describe('retrieval leg budgets', () => { vi.spyOn(SearchBudget.prototype, 'remaining').mockImplementation(function ( this: SearchBudget ) { - deadlines.set(this.leg, this.deadline) + if (!deadlines.has(this.leg)) deadlines.set(this.leg, this.deadline) return remaining.call(this) }) const access: UserAccessScope = { @@ -106,6 +107,18 @@ function render(condition: unknown) { return (condition as { toSQL: () => { sql: string; params: unknown[] } }).toSQL() } +/** The document probe is the only vector statement that selects ids without ordering them. */ +function isProbeStatement(sql: string) { + return sql.includes('AS id FROM') && !sql.includes('ORDER BY') +} + +/** `+ 0` is what keeps the exact ranking off the ANN index, so it also identifies the statement. */ +function isExactRanking(sql: string) { + return sql.includes(') + 0 LIMIT') +} + +const statements = () => dbChainMockFns.execute.mock.calls.map(([query]) => render(query)) + function renderOne(filters: StructuredFilter[]) { const conditions = getStructuredTagFilters(filters, embeddingTable) expect(conditions).toHaveLength(1) @@ -354,6 +367,8 @@ describe('workspace-scoped vector retrieval', () => { }, ] let probeRows: Array<{ id: string }> + let traversedRows: Array<{ id: string; initial_count?: number }> + let exactRows: Array<{ id: string }> let failSettings: unknown let failCandidates: unknown @@ -361,20 +376,23 @@ describe('workspace-scoped vector retrieval', () => { resetDbChainMock() getForConnectors.mockReset() probeRows = probe + traversedRows = candidates + exactRows = probe failSettings = undefined failCandidates = undefined dbChainMockFns.execute.mockImplementation(async (query) => { const statement = render(query).sql - if (statement.includes('SELECT scoped_chunk.id')) return probeRows if (statement.includes('hnsw.iterative_scan')) { if (failSettings) throw failSettings return [] } if (statement.includes('AS visible')) { if (failCandidates) throw failCandidates - return candidates + return traversedRows } if (statement.includes('WITH scored_search_candidates')) return ranked + if (isExactRanking(statement)) return exactRows + if (statement.includes('AS id FROM')) return probeRows return [] }) }) @@ -383,8 +401,6 @@ describe('workspace-scoped vector retrieval', () => { vi.useRealTimers() }) - const statements = () => dbChainMockFns.execute.mock.calls.map(([query]) => render(query)) - it.each([handleVectorOnlySearch, handleTagAndVectorSearch])( 'does not acquire a connection or start SQL after the KB retrieval deadline', async (search) => { @@ -405,32 +421,102 @@ describe('workspace-scoped vector retrieval', () => { } ) - it('ranks an exhausted visible scope exactly and rechecks access before returning content', async () => { - probeRows = ranked - queueTableRows(schemaMock.embedding, ranked) + it('rescues an underfilled traversal by ranking the permitted set exactly', async () => { + traversedRows = ranked + probeRows = [{ id: 'near-doc' }, { id: 'far-doc' }] + exactRows = ranked queueTableRows(schemaMock.embedding, [...ranked].reverse()) expect((await handleVectorOnlySearch(params)).map((row) => row.id)).toEqual(['near', 'far']) - expect(statements()).toHaveLength(1) - expect(statements()[0].sql).not.toContain('<=>') - expect(render(dbChainMockFns.orderBy.mock.calls[0][0]).sql).toContain('+ 0') - for (const [condition] of dbChainMockFns.where.mock.calls) { - expect( - hasMockCondition( - condition, - (node) => - node.type === 'inArray' && - node.column === schemaMock.embedding.id && - Array.isArray(node.values) && - node.values.length === 2 && - node.values.includes('near') - ) - ).toBe(true) - expect(JSON.stringify(condition)).toContain('required_clause') - expect(JSON.stringify(condition)).toContain('aclVerifiedAt') - } + const exact = statements().find((query) => isExactRanking(query.sql))! + expect(exact.sql).not.toContain('CROSS JOIN LATERAL') + expect(JSON.stringify(exact)).toContain('near-doc') + const probeStatement = statements().find((query) => isProbeStatement(query.sql))! + expect(probeStatement.sql).not.toContain('<=>') + expect(probeStatement.params).toContain(VECTOR_PROBE_DOCUMENT_LIMIT + 1) + expect(JSON.stringify(probeStatement)).toContain('required_clause') + expect(getForConnectors).not.toHaveBeenCalled() + }) + + it('counts only chunks the search can return when a tag filter decides a document', async () => { + traversedRows = ranked + probeRows = [{ id: 'near-doc' }] + exactRows = ranked + queueTableRows(schemaMock.embedding, [...ranked].reverse()) + await handleTagAndVectorSearch({ + ...params, + structuredFilters: [{ tagSlot: 'tag1', fieldType: 'text', operator: 'eq', value: 'common' }], + }) + /** + * A document whose only tagged chunk is disabled contributes no candidate, so admitting it + * would spend the probe's document bound on a document the ranking then discards. + */ + const probe = JSON.stringify(statements().find((query) => isProbeStatement(query.sql))!) + expect(probe).toContain(String(schemaMock.embedding.tag1)) + expect(probe).toContain(`"left":"${schemaMock.embedding.enabled}","right":true`) + }) + + it('keeps the tuple budget an order of magnitude above the beam so the scan can iterate', async () => { + queueTableRows(schemaMock.embedding, [...ranked].reverse()) + await handleVectorOnlySearch(params) + const settings = statements().find((query) => query.sql.includes('hnsw.iterative_scan'))! + const [maxScanTuples, efSearch] = settings.params as string[] + /** + * `max_scan_tuples` excludes the first beam, so a budget at or below `ef_search` is spent + * before `ResumeScanItems` can widen anything and the scan never iterates (pgvector#912). + */ + expect(Number(maxScanTuples)).toBeGreaterThanOrEqual(Number(efSearch) * 10) + /** A beam is uninterruptible, so its width is also this leg's cancellation floor. */ + expect(Number(efSearch)).toBeLessThanOrEqual(400) + }) + + it('leaves a filled traversal alone instead of probing for an exact ranking', async () => { + queueTableRows(schemaMock.embedding, [...ranked].reverse()) + expect((await handleVectorOnlySearch(params)).map((row) => row.id)).toEqual(['near', 'far']) + expect(statements().filter((query) => isProbeStatement(query.sql))).toHaveLength(0) + expect(statements().filter((query) => isExactRanking(query.sql))).toHaveLength(0) + }) + + it('keeps an underfilled traversal when the permitted set is too large to rank exactly', async () => { + traversedRows = ranked + probeRows = new Array(VECTOR_PROBE_DOCUMENT_LIMIT + 1).fill({ id: 'doc' }) + queueTableRows(schemaMock.embedding, [...ranked].reverse()) + expect((await handleVectorOnlySearch(params)).map((row) => row.id)).toEqual(['near', 'far']) + expect(statements().filter((query) => isExactRanking(query.sql))).toHaveLength(0) + }) + + it('finishes a scope the probe finds nothing in without ranking anything', async () => { + traversedRows = [] + probeRows = [] + expect(await handleVectorOnlySearch(params)).toEqual([]) + expect(statements().filter((query) => isExactRanking(query.sql))).toHaveLength(0) + expect( + statements().filter((query) => query.sql.includes('WITH scored_search_candidates')) + ).toHaveLength(0) expect(getForConnectors).not.toHaveBeenCalled() }) + it('spends only its own share of the leg on a probe that runs long', async () => { + const budget = new SearchBudget('vector', performance.now() + 8000) + traversedRows = ranked + exactRows = ranked + const execute = dbChainMockFns.execute.getMockImplementation()! + dbChainMockFns.execute.mockImplementation(async (query) => { + const statement = render(query).sql + if (isProbeStatement(statement)) + throw new Error('canceling statement due to statement timeout', { + cause: { code: '57014' }, + }) + return execute(query) + }) + queueTableRows(schemaMock.embedding, [...ranked].reverse()) + expect((await handleVectorOnlySearch({ ...params, budget })).map((row) => row.id)).toEqual([ + 'near', + 'far', + ]) + expect(budget.timedOut).toBe(false) + expect(statements().filter((query) => isExactRanking(query.sql))).toHaveLength(0) + }) + 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']) @@ -504,14 +590,13 @@ describe('workspace-scoped vector retrieval', () => { }) it('does not turn a broad tag filter into exhaustive full-vector ranking', async () => { - queueTableRows(schemaMock.embedding, probe) queueTableRows(schemaMock.embedding, ranked) const rows = await handleTagAndVectorSearch({ ...params, structuredFilters: [{ tagSlot: 'tag1', fieldType: 'text', operator: 'eq', value: 'common' }], }) expect(rows.map((row) => row.id)).toEqual(['near', 'far']) - expect(Object.keys(dbChainMockFns.select.mock.calls[0][0])).toEqual(['id']) + expect(statements().filter((query) => isExactRanking(query.sql))).toHaveLength(0) 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)) @@ -534,7 +619,7 @@ describe('workspace-scoped vector retrieval', () => { expect(dbChainMockFns.transaction).toHaveBeenCalledOnce() }) - it.each(['vector.probe', 'vector.candidate_search', 'vector.rerank', 'vector.sql'] as const)( + it.each(['vector.candidate_search', 'vector.rerank', 'vector.sql'] as const)( 'reports a %s timeout as partial, not a complete empty search', async (failedStage) => { const query = SearchBudget.prototype.query @@ -565,7 +650,6 @@ describe('workspace-scoped vector retrieval', () => { run: (executor: SearchExecutor) => PromiseLike ) { const result = await (query.bind(this) as SearchBudget['query'])(stage, run) - if (stage === 'vector.probe') vi.spyOn(performance, 'now').mockReturnValue(30) if (stage === 'vector.candidate_search') vi.spyOn(performance, 'now').mockReturnValue(60) if (stage === 'vector.rerank') vi.spyOn(performance, 'now').mockReturnValue(80) return result @@ -576,7 +660,7 @@ describe('workspace-scoped vector retrieval', () => { statements() .filter((query) => query.sql.includes('statement_timeout')) .map((query) => query.params[0]) - ).toEqual(['100', '70', '70', '40', '20']) + ).toEqual(['100', '100', '40', '20']) }) it('does not convert an unexpected candidate failure into partial retrieval', async () => { @@ -778,6 +862,7 @@ describe('live repository authorization follows ranked candidates', () => { } const probePages: Array> = [] + const exactPages: Array> = [] const candidatePages: Array> = [] const rerankPages: Array>> = [] const keywordPages: Array>> = [] @@ -791,20 +876,19 @@ describe('live repository authorization follows ranked candidates', () => { beforeEach(() => { resetDbChainMock() probePages.length = 0 + exactPages.length = 0 candidatePages.length = 0 rerankPages.length = 0 keywordPages.length = 0 - dbChainMockFns.execute.mockImplementation(async (query) => - render(query).sql.includes('SELECT scoped_chunk.id') - ? (probePages.shift() ?? []) - : render(query).sql.includes('AS visible') - ? (candidatePages.shift() ?? []) - : render(query).sql.includes('WITH scored_search_candidates') - ? (rerankPages.shift() ?? []) - : render(query).sql.includes('WITH visible_keyword_documents') - ? (keywordPages.shift() ?? []) - : [] - ) + dbChainMockFns.execute.mockImplementation(async (query) => { + const statement = render(query).sql + if (statement.includes('AS visible')) return candidatePages.shift() ?? [] + if (statement.includes('WITH scored_search_candidates')) return rerankPages.shift() ?? [] + if (statement.includes('WITH visible_keyword_documents')) return keywordPages.shift() ?? [] + if (isExactRanking(statement)) return exactPages.shift() ?? [] + if (statement.includes('AS id FROM')) return probePages.shift() ?? [] + return [] + }) getForConnectors.mockReset().mockResolvedValue(allowed) }) @@ -847,73 +931,58 @@ describe('live repository authorization follows ranked candidates', () => { ) }) - it('finishes empty scopes after the bounded probe without scanning HNSW or calling providers', async () => { + it('finishes a scope the probe finds nothing in without ranking or calling providers', async () => { probePages.push([]) expect(await handleVectorOnlySearch({ ...params, structuredFilters: undefined })).toEqual([]) - expect(dbChainMockFns.select).not.toHaveBeenCalled() - const probe = render(dbChainMockFns.execute.mock.calls[0][0]) - expect(probe.sql).toContain('CROSS JOIN LATERAL') - expect(probe.params.filter((value) => value === 400)).toHaveLength(2) - expect(dbChainMockFns.orderBy).not.toHaveBeenCalled() + const probe = statements().find((query) => isProbeStatement(query.sql))! + expect(probe.params).toContain(VECTOR_PROBE_DOCUMENT_LIMIT + 1) + expect(probe.sql).not.toContain('<=>') + expect(JSON.stringify(probe)).toContain('required_clause') + expect( + statements().filter((query) => query.sql.includes('WITH scored_search_candidates')) + ).toHaveLength(0) expect(getForConnectors).not.toHaveBeenCalled() }) - it('reads vectors only for the bounded IDs when a broad scope has few candidates', async () => { - probePages.push([candidate('selected', 'allowed-source')]) - queueTableRows(schemaMock.embedding, [candidate('selected', 'allowed-source')]) + it('ranks the permitted set exactly when the traversal comes back underfilled', async () => { + probePages.push([{ id: 'doc-selected' }]) + exactPages.push([{ id: 'selected' }]) + queueRerank([candidate('selected', 'allowed-source')]) queueTableRows(schemaMock.embedding, [ { id: 'selected', content: 'Verified small scope', distance: 0.1 }, ]) expect(await handleVectorOnlySearch({ ...params, structuredFilters: undefined })).toEqual([ { id: 'selected', content: 'Verified small scope', distance: 0.1 }, ]) - const probe = dbChainMockFns.execute.mock.calls[0][0] - expect(render(probe).sql).toContain('SELECT scoped_chunk.id') - expect(render(probe).sql).not.toContain('<=>') - expect(JSON.stringify(probe)).toContain('required_clause') - expect( - hasMockCondition( - dbChainMockFns.where.mock.calls[0][0], - (node) => - node.type === 'inArray' && - node.column === schemaMock.embedding.id && - Array.isArray(node.values) && - node.values.length === 1 && - node.values[0] === 'selected' - ) - ).toBe(true) + const exact = statements().find((query) => isExactRanking(query.sql))! + expect(exact.sql).not.toContain('CROSS JOIN LATERAL') + expect(JSON.stringify(exact)).toContain('doc-selected') expect(getForConnectors).toHaveBeenCalledExactlyOnceWith(['allowed-source'], undefined) }) - it.each([199, 200, 399])( - 'ranks an exhausted scope of %s chunks once without repeating candidate search', + it.each([1, 200, 399])( + 'ranks a permitted set of %s documents exactly without repeating the traversal', async (count) => { - const probe = Array.from({ length: count }, (_, index) => ({ id: `chunk-${index}` })) - probePages.push(probe) - queueTableRows(schemaMock.embedding, [candidate('chunk-0', 'allowed-source')]) + probePages.push(Array.from({ length: count }, (_, index) => ({ id: `doc-${index}` }))) + exactPages.push([{ id: 'chunk-0' }]) + queueRerank([candidate('chunk-0', 'allowed-source')]) queueTableRows(schemaMock.embedding, [ { id: 'chunk-0', content: 'Authorized passage', distance: 0.1 }, ]) const rows = await handleVectorOnlySearch({ ...params, structuredFilters: undefined }) expect(rows.map((row) => row.id)).toEqual(['chunk-0']) - expect(dbChainMockFns.execute).toHaveBeenCalledOnce() - expect( - hasMockCondition( - dbChainMockFns.where.mock.calls[0][0], - (node) => - node.type === 'inArray' && - node.column === schemaMock.embedding.id && - Array.isArray(node.values) && - node.values.length === count - ) - ).toBe(true) - expect(dbChainMockFns.orderBy).toHaveBeenCalledOnce() + expect(statements().filter((query) => query.sql.includes('AS visible'))).toHaveLength(1) + expect(statements().filter((query) => isExactRanking(query.sql))).toHaveLength(1) } ) - it('keeps an underfilled ANN result instead of rescoring the whole projection', async () => { - probePages.push(Array.from({ length: 400 }, (_, index) => ({ id: `probe-${index}` }))) + it('keeps an underfilled traversal when the permitted set is past the probe bound', async () => { queueCandidates([{ id: 'selected' }], 1) + probePages.push( + Array.from({ length: VECTOR_PROBE_DOCUMENT_LIMIT + 1 }, (_, index) => ({ + id: `doc-${index}`, + })) + ) queueRerank([candidate('selected', 'allowed-source')]) queueTableRows(schemaMock.embedding, [ { id: 'selected', content: 'Verified fallback', distance: 0.1 }, @@ -921,13 +990,12 @@ describe('live repository authorization follows ranked candidates', () => { expect(await handleVectorOnlySearch({ ...params, structuredFilters: undefined })).toEqual([ { id: 'selected', content: 'Verified fallback', distance: 0.1 }, ]) - const candidateQuery = dbChainMockFns.execute.mock.calls.find(([query]) => - render(query).sql.includes('AS visible') - )![0] + const candidateQuery = statements().find((query) => query.sql.includes('AS visible'))! /** 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(candidateQuery.sql).not.toContain('UNION ALL') + expect(candidateQuery.sql).not.toContain('filtered_scores') + expect(candidateQuery.sql).toContain('CROSS JOIN LATERAL') + expect(statements().filter((query) => isExactRanking(query.sql))).toHaveLength(0) expect(JSON.stringify(dbChainMockFns.where.mock.calls.at(-1)![0])).toContain( 'github_read_grant' ) @@ -1010,10 +1078,13 @@ describe('live repository authorization follows ranked candidates', () => { '%s ranks identifiers before verification and loads content under the full predicate', async (mode) => { const candidates = [candidate('selected', 'allowed-source')] - if (mode === 'vector' || mode === 'tag-vector') - queueTableRows(schemaMock.embedding, candidates) + if (mode === 'vector' || mode === 'tag-vector') { + probePages.push([{ id: 'doc-selected' }]) + exactPages.push([{ id: 'selected' }]) + queueRerank(candidates) + } if (mode === 'keyword') keywordPages.push(candidates) - else queueTableRows(schemaMock.embedding, candidates) + if (mode === 'tags') queueTableRows(schemaMock.embedding, candidates) queueTableRows(schemaMock.embedding, [{ id: 'selected', content: 'verified result' }]) const rows = mode === 'vector' @@ -1035,23 +1106,19 @@ describe('live repository authorization follows ranked candidates', () => { expect(ranking).toContain('ORDER BY keyword_rank DESC, id LIMIT') expect(ranking).not.toContain('<=>') expect(ranking).not.toContain('"content"') - } else { - expect( - Object.keys(dbChainMockFns.select.mock.calls[mode === 'tags' ? 0 : 1][0]).sort() - ).toEqual( - [ - 'id', - 'documentId', - 'connectorId', - 'liveAuthorizationSource', - ...(mode === 'tags' ? [] : ['distance']), - ].sort() + } else if (mode === 'tags') { + expect(Object.keys(dbChainMockFns.select.mock.calls[0][0]).sort()).toEqual( + ['id', 'documentId', 'connectorId', 'liveAuthorizationSource'].sort() ) + } else { + const ranking = statements().find((query) => isExactRanking(query.sql))! + expect(ranking.sql).toContain('AS id FROM') + expect(ranking.sql).not.toContain('"content"') } const rankingOrder = - mode === 'keyword' - ? dbChainMockFns.execute.mock.invocationCallOrder[0] - : dbChainMockFns.select.mock.invocationCallOrder[0] + mode === 'tags' + ? dbChainMockFns.select.mock.invocationCallOrder[0] + : dbChainMockFns.execute.mock.invocationCallOrder[0] expect(rankingOrder).toBeLessThan(getForConnectors.mock.invocationCallOrder[0]) expect(getForConnectors.mock.invocationCallOrder[0]).toBeLessThan( dbChainMockFns.select.mock.invocationCallOrder.at(-1)! @@ -1080,10 +1147,13 @@ describe('live repository authorization follows ranked candidates', () => { async (mode) => { getForConnectors.mockResolvedValue(identity) const candidates = [{ ...candidate('gmail', 'gmail-source'), installationSource: false }] - if (mode === 'vector' || mode === 'tag-vector') - queueTableRows(schemaMock.embedding, candidates) + if (mode === 'vector' || mode === 'tag-vector') { + probePages.push([{ id: 'doc-gmail' }]) + exactPages.push([{ id: 'gmail' }]) + queueRerank(candidates) + } if (mode === 'keyword') keywordPages.push(candidates) - else queueTableRows(schemaMock.embedding, candidates) + if (mode === 'tags') queueTableRows(schemaMock.embedding, candidates) const hydrated = [{ id: 'gmail', content: 'current permitted content' }] queueTableRows(schemaMock.embedding, hydrated) const searchParams = { ...params, filters: { source: 'gmail' } } @@ -1111,9 +1181,8 @@ describe('live repository authorization follows ranked candidates', () => { const hydration = JSON.stringify(dbChainMockFns.where.mock.calls.at(-1)![0]) expect(hydration).toContain('acl') expect(hydration).toContain('knowledgeConnectorMember') - expect(dbChainMockFns.select).toHaveBeenCalledTimes( - mode === 'keyword' ? 1 : mode === 'tags' ? 2 : 3 - ) + /** Vector ranking is raw SQL throughout; only hydration reads through the query builder. */ + expect(dbChainMockFns.select).toHaveBeenCalledTimes(mode === 'tags' ? 2 : 1) } ) diff --git a/apps/sim/lib/knowledge/search/queries.ts b/apps/sim/lib/knowledge/search/queries.ts index 341c62d91d8..bb5ae17f674 100644 --- a/apps/sim/lib/knowledge/search/queries.ts +++ b/apps/sim/lib/knowledge/search/queries.ts @@ -12,6 +12,7 @@ import { and, eq, inArray, isNull, type SQL, sql } from 'drizzle-orm' import { knowledgeAccessCondition, knowledgeMetadataCandidateAccessCondition, + textArrayLiteral, } from '@/lib/knowledge/access/predicate' import type { KnowledgeAccessProvider, KnowledgeAccessScope } from '@/lib/knowledge/access/types' import type { KbEmbeddingDimensions } from '@/lib/knowledge/embedding-models' @@ -50,20 +51,53 @@ const UNDEFINED_OBJECT_SQLSTATE = '42704' /** Bound candidate pages retained while live permissions are checked. */ const MAX_AUTHORIZED_SEARCH_CANDIDATES = 20_000 /** - * 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. + * Approximate iterative-visit threshold for a permission-starved graph walk. It excludes + * pgvector's initial beam, so it only takes effect once `ResumeScanItems` starts widening. * - * Raising it trades recall for latency far more steeply than its size suggests. Measured on a - * search index where visibility admitted a tenth of the corpus, visiting four times as many tuples - * took 5.1s and 9.7s on consecutive identical runs, against ~115ms for the bounded walk — the walk - * degrades superlinearly with depth, and unpredictably. Re-measure before changing it. + * It must therefore stay roughly an order of magnitude above `ef_search`, or the first beam + * already exhausts the tuple budget and the scan stops before it can iterate at all — pgvector's + * maintainer says as much in pgvector#912. Measured on a production-shaped corpus, the previous + * pairing of a 1,000-wide beam against a 1,000-tuple budget returned fewer candidates than a + * narrower beam allowed to iterate, and spent longer inside the one uninterruptible beam. */ -const CANDIDATE_HNSW_MAX_SCAN_TUPLES = '1000' -const CANDIDATE_HNSW_EF_SEARCH = '1000' +const CANDIDATE_HNSW_MAX_SCAN_TUPLES = '20000' +/** + * Beam width per iteration. A beam is the granularity of cancellation: pgvector calls + * `CHECK_FOR_INTERRUPTS` only while building an index, never inside `hnswgettuple`, so neither + * `statement_timeout` nor a cancellation request can interrupt one. A narrower beam that iterates + * therefore bounds the leg's uninterruptible floor as well as widening its reach. + */ +const CANDIDATE_HNSW_EF_SEARCH = '200' const CANDIDATE_HNSW_SCAN_MEM_MULTIPLIER = '2' const MIN_VECTOR_RERANK_CANDIDATES = 400 const MAX_VECTOR_RERANK_CANDIDATES = 1600 const VECTOR_RERANK_OVERSAMPLING = 32 +/** + * The probe's share of the leg. It ranks nothing, so it must never be why the leg misses its + * own deadline. + * + * Its share comes out of what the rescue can claim from the narrowest live budget, + * `DIRECT_SEARCH_VECTOR_BUDGET_MS`, before live authorization, hydration and the exact rerank + * need the rest. + */ +const VECTOR_PROBE_BUDGET_MS = 600 +/** + * What one document costs the probe, measured on a corpus shaped like a search index under + * comparable cache pressure: the access predicate, evaluated once per document. + */ +const VECTOR_PROBE_MICROSECONDS_PER_DOCUMENT = 6 +/** + * Documents the probe enumerates before it concludes the permitted set is too large to rank + * exactly. Derived so that reaching it is what spends the probe's budget, rather than a separate + * number that a change to that budget could silently invalidate. + * + * It bounds the rescue's second step too: exact ranking of the `halfvec` projection measures at + * around half the probe's per-document cost, so a permitted set within this bound is affordable + * by construction. + */ +export const VECTOR_PROBE_DOCUMENT_LIMIT = Math.round( + (VECTOR_PROBE_BUDGET_MS * 1000) / VECTOR_PROBE_MICROSECONDS_PER_DOCUMENT +) /** How long to stop trying the iterative-scan settings after the server rejected them. */ const HNSW_SETTINGS_UNSUPPORTED_RETRY_MS = 10 * 60 * 1000 @@ -781,10 +815,58 @@ export async function handleVectorOnlySearch(params: SearchParams): Promise { + const probeBudget = budget?.capped(VECTOR_PROBE_BUDGET_MS) + try { + const probed = await runSearchQuery(probeBudget, 'vector.probe', (executor) => + executor.execute<{ id: string }>(sql` + SELECT ${document.id} AS id FROM ${document} + WHERE ${and(...conditions)} + LIMIT ${VECTOR_PROBE_DOCUMENT_LIMIT + 1} + `) + ) + if (probed.length > VECTOR_PROBE_DOCUMENT_LIMIT) return null + return probed.map(({ id }) => id) + } catch (error) { + if (!budget || !probeBudget?.isTimeout(error)) throw error + /** Only the probe's share was spent; the leg's own deadline still governs. */ + budget.remaining() + return null + } +} + +/** + * Tags live on chunks, so a row qualifies when a chunk it joins to carries them — and only a + * chunk the search can actually return counts, or a document whose sole match is disabled would + * be admitted by a check that ranking then discards. + */ +function chunkTagCondition(join: SQL, tagConditions: SQL[]): SQL | undefined { + if (!tagConditions.length) return undefined + return sql`EXISTS ( + SELECT 1 FROM ${embedding} + WHERE ${and(join, eq(embedding.enabled, true), ...tagConditions)} + )` +} + +/** + * Select a bounded candidate pool and rerank it against the original vectors. + * + * A bounded ANN traversal fills that pool. When visibility leaves the traversal short of its + * limit, a bounded probe decides whether the permitted set is small enough to rank exactly + * instead, which recovers the candidates the traversal's post-filter discarded. + * * Live source authorization and content hydration still run after candidate ranking. */ async function selectVectorResults(params: SearchParams): Promise { @@ -796,13 +878,11 @@ async function selectVectorResults(params: SearchParams): Promise } | undefined return selectAuthorizedSearchResults({ leg: 'vector', access: params.access, @@ -836,21 +926,15 @@ async function selectVectorResults(params: SearchParams): Promise { + /** Explicit document IDs are already a bounded scope, and retain exhaustive ordering. */ + const exactPage = async () => { annotateSearchDiagnostics({ vectorRanking: 'exact' }) const candidates = await runSearchQuery(params.budget, 'vector.exact', (executor) => executor .select({ ...SEARCH_READ_CANDIDATE_FIELDS, distance: distance.as('distance') }) .from(embedding) .innerJoin(document, eq(embedding.documentId, document.id)) - .where( - and( - ...conditions, - ...visibility, - candidateIds ? inArray(embedding.id, candidateIds) : undefined - ) - ) + .where(and(...conditions, ...visibility)) .orderBy(sql`(${distance}) + 0`, embedding.id) .limit(limit) .offset(offset) @@ -858,69 +942,25 @@ async function selectVectorResults(params: SearchParams): Promise - tagConditions.length - ? executor - .select({ id: embedding.id }) - .from(embedding) - .innerJoin(document, eq(document.id, embedding.documentId)) - .where( - and( - inArray(embedding.knowledgeBaseId, params.knowledgeBaseIds), - ...visibility, - ...tagConditions - ) - ) - .limit(candidateLimit) - : executor.execute<{ id: string }>(sql` - SELECT scoped_chunk.id FROM ${document} - CROSS JOIN LATERAL ( - SELECT ${embeddingSearch.id} AS id FROM ${embeddingSearch} - WHERE ${and( - eq(embeddingSearch.documentId, document.id), - inArray(embeddingSearch.knowledgeBaseId, params.knowledgeBaseIds), - eq(embeddingSearch.enabled, true) - )} - LIMIT ${candidateLimit} - ) AS scoped_chunk - WHERE ${and(...candidateDocumentVisibility)} - LIMIT ${candidateLimit} - `) - ) - if (probe.length === 0) return { candidates: [], nextOffset: offset } - if (probe.length < candidateLimit) { - return exactPage(probe.map((candidate) => candidate.id)) - } - annotateSearchDiagnostics({ - vectorRanking: 'candidate-rerank', - vectorCandidateStorage: 'stored-halfvec', - vectorCandidateLimit: candidateLimit, - vectorCandidateScan: 'planned', - vectorCandidateDimensions: embeddingCandidateDimensions( - queryVector.dimensions, - queryVector.model - ), - }) - /** - * The bounded ANN traversal is the whole candidate set. LIMIT keeps document authorization - * downstream of the traversal, with a primary-key lookup per candidate. - * - * An underfilled traversal yields fewer candidates rather than widening the search. Widening - * it has no affordable form here: rescoring the projection exhaustively is O(corpus) and a - * deeper `hnsw.max_scan_tuples` is worse still. Measured where visibility admitted a tenth of - * the corpus, the exhaustive rescan took 1.9s while visiting four times as many tuples took - * 5.1s and 9.7s on consecutive identical runs. Both exceed the retrieval budget once the - * corpus grows, and a leg that exceeds its budget returns nothing at all, so fewer candidates - * strictly beats every widening strategy available. - */ - const identities = await withVectorScanSettings( - (executor) => - executor.execute<{ id: string }>(sql` + const excludedKey = excludedSources.join('\u0000') + if (candidatePool?.excludedKey !== excludedKey) { + annotateSearchDiagnostics({ + vectorRanking: 'candidate-rerank', + vectorCandidateStorage: 'stored-halfvec', + vectorCandidateLimit: candidateLimit, + vectorCandidateScan: 'planned', + vectorCandidateDimensions: embeddingCandidateDimensions( + queryVector.dimensions, + queryVector.model + ), + }) + /** + * The bounded ANN traversal is the whole candidate set. LIMIT keeps document + * authorization downstream of the traversal, with a primary-key lookup per candidate. + */ + const traversed = await withVectorScanSettings( + (executor) => + executor.execute<{ id: string }>(sql` SELECT ${embeddingSearch.id} AS id FROM ${embeddingSearch} CROSS JOIN LATERAL ( SELECT 1 FROM ${document} @@ -933,12 +973,59 @@ async function selectVectorResults(params: SearchParams): Promise = traversed + if (traversed.length < candidateLimit) { + const visibleDocumentIds = await probeVisibleDocuments( + [...candidateDocumentVisibility, documentTagCondition], + params.budget + ) + if (visibleDocumentIds) { + annotateSearchDiagnostics({ + vectorRanking: 'exact-candidates', + vectorProbeDocumentCount: visibleDocumentIds.length, + }) + /** + * `+ 0` keeps the planner off the ANN index, and the probed identities keep the scan + * on `embedding_search_document_lookup_idx`, so this reads what the permitted set + * costs rather than re-deriving permission across the whole index. Exact ranking also + * honours `statement_timeout`, which a traversal cannot. + */ + selected = visibleDocumentIds.length + ? await runSearchQuery(params.budget, 'vector.exact_candidates', (executor) => + executor.execute<{ id: string }>(sql` + SELECT ${embeddingSearch.id} AS id FROM ${embeddingSearch} + WHERE ${and( + inArray(embeddingSearch.knowledgeBaseId, params.knowledgeBaseIds), + eq(embeddingSearch.enabled, true), + sql`${embeddingSearch.documentId} = ANY(${textArrayLiteral(visibleDocumentIds)})`, + candidateTagCondition + )} + ORDER BY (${candidateDistance}) + 0 LIMIT ${candidateLimit} + `) + ) + : [] + } + } + candidatePool = { excludedKey, identities: selected } + annotateSearchDiagnostics({ + vectorCandidateCount: selected.length, + vectorCandidateScan: selected.length < candidateLimit ? 'underfilled' : 'planned', + }) + } + const { identities } = candidatePool if (!identities.length) return { candidates: [], nextOffset: offset } /** Score each bounded candidate once; sorting the materialized scalar cannot invoke HNSW again. */ const page = await runSearchQuery(params.budget, 'vector.rerank', (executor) =>