diff --git a/apps/sim/lib/knowledge/search/queries.test.ts b/apps/sim/lib/knowledge/search/queries.test.ts index 930021250c3..3535574312e 100644 --- a/apps/sim/lib/knowledge/search/queries.test.ts +++ b/apps/sim/lib/knowledge/search/queries.test.ts @@ -41,6 +41,7 @@ import { handleTagAndVectorSearch, handleTagOnlySearch, handleVectorOnlySearch, + PERMITTED_EXACT_DOCUMENT_LIMIT, type PermittedDocuments, resolvePermittedDocuments, resolveReach, @@ -1794,6 +1795,86 @@ describe('permitted-document planner', () => { expect(ginStatements()).toHaveLength(0) }) + describe('a bounded set past the exact-ranking size', () => { + const large = Array.from({ length: PERMITTED_EXACT_DOCUMENT_LIMIT }, (_, index) => ({ + id: `doc-${index}`, + connectorId: 'src-a', + })) + const accessPlan = { + connectors: { workspace: [], admin: ['src-a'], members: [], liveProofRequired: [] }, + observers: { confirmed: [], observed: [] }, + memberSources: [], + connectorTypes: new Map(), + uploads: true, + } + + it('ranks with Tin as a narrow reader, decided on the row', async () => { + tinPages = [{ ranked: 1500, candidates: [hit('a', 'src-a')] }] + queueTableRows(schemaMock.embedding, [{ ...hit('a', 'src-a'), content: 'release notes' }]) + const results = await keyword({ + permitted: { kind: 'bounded', documents: large }, + accessPlan, + }) + expect(results.map((row) => row.id)).toEqual(['a']) + expect(mockResolveTinKeywordQuery).toHaveBeenCalledTimes(1) + expect(tinStatements()).toHaveLength(1) + expect(JSON.stringify(tinStatements()[0])).toContain('2000') + expect(JSON.stringify(tinStatements()[0])).not.toContain('doc-4999') + expect(ginStatements()).toHaveLength(0) + }) + + it('falls back to a GIN ranking that reads what the term matches, not every chunk of the set', async () => { + mockResolveTinKeywordQuery.mockResolvedValue(null) + queueTableRows(schemaMock.embedding, [{ ...hit('a', 'src-a'), content: 'release notes' }]) + await keyword({ permitted: { kind: 'bounded', documents: large }, accessPlan }) + expect(tinStatements()).toHaveLength(0) + expect(ginStatements()).toHaveLength(1) + expect(JSON.stringify(ginStatements()[0])).not.toContain('doc-4999') + }) + + it('hands the page to the GIN ranking when the widest window cannot fill it', async () => { + tinPages = [ + { ranked: 2000, candidates: [] }, + { ranked: 20_000, candidates: [] }, + ] + queueTableRows(schemaMock.embedding, [{ ...hit('a', 'src-a'), content: 'release notes' }]) + await keyword({ permitted: { kind: 'bounded', documents: large }, accessPlan }) + expect(tinStatements()).toHaveLength(2) + /** Every match is covered again, by the ranking whose cost follows the term, not the set. */ + expect(ginStatements()).toHaveLength(1) + expect(JSON.stringify(ginStatements()[0])).not.toContain('doc-4999') + }) + + it('leaves a later page short rather than resuming a different ranking at its offset', async () => { + /** The first page fills from Tin; hydration keeps half, so a second page is asked for. */ + const first = Array.from({ length: 40 }, (_, index) => hit(`t-${index}`, 'src-a')) + tinPages = [ + { ranked: 2000, candidates: first }, + { ranked: 2000, candidates: [] }, + { ranked: 20_000, candidates: [] }, + ] + queueTableRows( + schemaMock.embedding, + first.slice(0, 20).map((row) => ({ ...row, content: 'release notes' })) + ) + const results = await keyword({ + topK: 40, + permitted: { kind: 'bounded', documents: large }, + accessPlan, + }) + expect(results).toHaveLength(20) + expect(tinStatements()).toHaveLength(3) + expect(ginStatements()).toHaveLength(0) + }) + + it('keeps the bounded read for a set under the size', async () => { + mockResolveTinKeywordQuery.mockResolvedValue(null) + await keyword({ permitted: { kind: 'bounded', documents: large.slice(0, -1) }, accessPlan }) + expect(mockResolveTinKeywordQuery).not.toHaveBeenCalled() + expect(JSON.stringify(ginStatements()[0])).toContain('doc-4998') + }) + }) + it('leaves the page to the GIN ranking when the widest window cannot fill it', async () => { tinPages = [ { ranked: 2000, candidates: [] }, @@ -2400,6 +2481,74 @@ describe('filters on a resolved scope', () => { expect(JSON.stringify(walks[0])).toContain('release') }) + describe('a bounded set past the exact-ranking size', () => { + const large = Array.from({ length: PERMITTED_EXACT_DOCUMENT_LIMIT }, (_, index) => ({ + id: `doc-${index}`, + connectorId: 'src-a', + })) + const walked = Array.from({ length: 200 }, (_, index) => hit(`w-${index}`, 'src-a')) + const search = (documents: typeof large) => + handleVectorOnlySearch({ + ...params, + permitted: { kind: 'bounded', documents }, + accessPlan: plan(), + }) + beforeEach(() => { + const execute = dbChainMockFns.execute.getMockImplementation()! + dbChainMockFns.execute.mockImplementation(async (query) => { + /** The projection is filled, so a walk decides readability on the row. */ + if (render(query).sql.includes('AS unfilled')) return [{ unfilled: false }] + return execute(query) + }) + }) + + it('walks the graph on the row instead of ranking every chunk of the set', async () => { + traversedRows = walked + queueTableRows(schemaMock.embedding, [walked[0]]) + expect((await search(large)).map((row) => row.id)).toEqual(['w-0']) + const walks = statements().filter((query) => isWalk(query.sql)) + expect(walks).toHaveLength(1) + expect(statements().filter((query) => isExactRanking(query.sql))).toHaveLength(0) + /** Readability rides on the row through the plan; the set's identifiers never cross the wire. */ + expect(JSON.stringify(walks[0])).not.toContain('doc-4999') + expect(JSON.stringify(walks[0])).toContain('src-a') + }) + + it('ranks the set exactly when the walk cannot fill its pool', async () => { + traversedRows = [] + exactRows = [hit('a', 'src-a')] + queueTableRows(schemaMock.embedding, [hit('a', 'src-a')]) + expect((await search(large)).map((row) => row.id)).toEqual(['a']) + expect(statements().filter((query) => isWalk(query.sql))).toHaveLength(1) + const exact = statements().filter((query) => isExactRanking(query.sql)) + expect(exact).toHaveLength(1) + expect(JSON.stringify(exact[0])).toContain('doc-4999') + }) + + it('refills a pool the walk filled but hydration could not with the exact ranking', async () => { + traversedRows = walked + exactRows = [hit('a', 'src-a')] + /** None of the walked rows survives the document predicate; the set's own ranking then does. */ + for (let page = 0; page < 10; page++) queueTableRows(schemaMock.embedding, []) + queueTableRows(schemaMock.embedding, [hit('a', 'src-a')]) + expect((await search(large)).map((row) => row.id)).toEqual(['a']) + expect(statements().filter((query) => isWalk(query.sql))).toHaveLength(1) + const exact = statements().filter((query) => isExactRanking(query.sql)) + expect(exact).toHaveLength(1) + /** The refill ranks past the rows already read, so nothing already rejected is read twice. */ + expect(JSON.stringify(exact[0])).toContain('doc-4999') + expect(JSON.stringify(exact[0])).toContain('w-199') + }) + + it('ranks a set under the size exactly, without a walk', async () => { + exactRows = [hit('a', 'src-a')] + queueTableRows(schemaMock.embedding, [hit('a', 'src-a')]) + expect((await search(large.slice(0, -1))).map((row) => row.id)).toEqual(['a']) + expect(statements().filter((query) => isWalk(query.sql))).toHaveLength(0) + expect(statements().filter((query) => isExactRanking(query.sql))).toHaveLength(1) + }) + }) + it('ranks a date-bounded set exactly even when a member source has its own index', async () => { indexedSourceRows = [{ name: 'idx', connectorId: 'member-src' }] exactRows = [{ id: 'a' }] diff --git a/apps/sim/lib/knowledge/search/queries.ts b/apps/sim/lib/knowledge/search/queries.ts index 0798f7bacf2..a92a5eb1ec7 100644 --- a/apps/sim/lib/knowledge/search/queries.ts +++ b/apps/sim/lib/knowledge/search/queries.ts @@ -216,6 +216,16 @@ export const VECTOR_PROBE_DOCUMENT_LIMIT = Math.round( (VECTOR_PROBE_BUDGET_MS * 1000) / VECTOR_PROBE_MICROSECONDS_PER_DOCUMENT ) +/** + * Documents a bounded permitted set may hold before ranking it exactly costs more than walking + * the graph on the row. Exact ranking reads every chunk of the set, a few per document, where an + * on-row walk reads at most {@link CANDIDATE_HNSW_MAX_SCAN_TUPLES} tuples; at this size the two + * meet. A set past it is walked first and ranked exactly only if the walk cannot fill its pool, so + * its recall is never below the exact ranking's and its usual cost is the walk's. The same size + * turns the keyword leg from a read of the set's every chunk into a ranking decided on the row. + */ +export const PERMITTED_EXACT_DOCUMENT_LIMIT = 5_000 + /** How long to stop trying the iterative-scan settings after the server rejected them. */ const HNSW_SETTINGS_UNSUPPORTED_RETRY_MS = 10 * 60 * 1000 @@ -1747,7 +1757,8 @@ async function selectVectorResults(params: SearchParams): Promise { + /** Ranks the set's chunks exactly; `read` are chunks a pool already holds, ranked past. */ + const rankPermittedExactly = async (documentIds: string[], read?: readonly string[]) => { annotateSearchDiagnostics({ vectorRanking: 'exact-candidates' }) if (!documentIds.length) return [] return runSearchQuery(params.budget, 'vector.exact_candidates', (executor) => @@ -1758,6 +1769,9 @@ async function selectVectorResults(params: SearchParams): Promise + withVectorScanSettings( + (executor) => + executor.execute( + plan + ? sql` + SELECT ${PROJECTION_CANDIDATE_COLUMNS} + FROM ${embeddingSearch} /* on-row visibility */ + WHERE ${and( + scopeOfWalk, + projectionCandidateAccessCondition(embeddingSearch, params.access, plan, { filled }), + documentCondition === undefined + ? undefined + : sql`EXISTS (SELECT 1 FROM ${document} WHERE ${and(eq(document.id, embeddingSearch.documentId), documentCondition)})` + )} + ORDER BY ${candidateDistance} LIMIT ${candidateLimit} + ` + : sql` + SELECT ${embeddingSearch.id} AS id, ${embeddingSearch.documentId} AS "documentId", + visible.connector_id AS "connectorId" + FROM ${embeddingSearch} + CROSS JOIN LATERAL ( + SELECT ${document.connectorId} AS connector_id FROM ${document} + WHERE ${and(eq(document.id, embeddingSearch.documentId), ...candidateDocumentVisibility, candidateTagCondition)} + LIMIT 1 + ) AS visible + WHERE ${scopeOfWalk} + ORDER BY ${candidateDistance} LIMIT ${candidateLimit} + ` + ), + params.budget, + 'vector.candidate_search', + plan ? onRowWalkScanTuples(documentCondition, filled) : undefined + ) /** * 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. @@ -1774,6 +1833,40 @@ async function selectVectorResults(params: SearchParams): Promise plannedIndexedSources?.has(id) ?? false ) if ( + params.permitted?.kind === 'bounded' && + plan && + filled && + params.permitted.documents.length >= PERMITTED_EXACT_DOCUMENT_LIMIT + ) { + /** + * A set this large costs more to rank exactly than to walk: exact ranking reads every + * chunk of every document in it, while the walk decides readability on the rows it + * visits and stops at its tuple cap. The walk answers whenever the set is a fair share + * of the graph; where it is not, the walk underfills and the exact ranking that was + * always complete takes over, so nothing is lost but the walk's bounded cost. + * + * The walk decides readability on the projection row, which is broader than the + * document predicate hydration applies, so a pool it filled can still run short of + * readable rows. That shortfall is what refills a pool: the refill is the exact ranking, + * complete over the set, ranked past the rows already read and placed behind them, so + * the pages keep their offsets and every refill is a full window of fresh rows. + */ + const permittedIds = params.permitted.documents.map((entry) => entry.id) + const previous = + candidatePool?.excludedKey === excludedKey ? candidatePool.ids : undefined + if (previous) { + const exact = await rankPermittedExactly( + permittedIds, + previous.map((candidate) => candidate.id) + ) + selected = [...previous, ...exact] + exhausted = exact.length < candidateLimit + } else { + selected = await walkGraph() + if (selected.length < candidateLimit) + selected = await rankPermittedExactly(permittedIds) + } + } else if ( params.permitted?.kind === 'bounded' && (!walksASource || dateFilterCondition(params.filters) || params.filters?.source) ) { @@ -1806,48 +1899,7 @@ async function selectVectorResults(params: SearchParams): Promise - executor.execute( - plan - ? sql` - SELECT ${PROJECTION_CANDIDATE_COLUMNS} - FROM ${embeddingSearch} /* on-row visibility */ - WHERE ${and( - scopeOfWalk, - projectionCandidateAccessCondition(embeddingSearch, params.access, plan, { filled }), - documentCondition === undefined - ? undefined - : sql`EXISTS (SELECT 1 FROM ${document} WHERE ${and(eq(document.id, embeddingSearch.documentId), documentCondition)})` - )} - ORDER BY ${candidateDistance} LIMIT ${candidateLimit} - ` - : sql` - SELECT ${embeddingSearch.id} AS id, ${embeddingSearch.documentId} AS "documentId", - visible.connector_id AS "connectorId" - FROM ${embeddingSearch} - CROSS JOIN LATERAL ( - SELECT ${document.connectorId} AS connector_id FROM ${document} - WHERE ${and(eq(document.id, embeddingSearch.documentId), ...candidateDocumentVisibility, candidateTagCondition)} - LIMIT 1 - ) AS visible - WHERE ${scopeOfWalk} - ORDER BY ${candidateDistance} LIMIT ${candidateLimit} - ` - ), - params.budget, - 'vector.candidate_search', - plan ? onRowWalkScanTuples(documentCondition, filled) : undefined - ) + selected = await walkGraph() /** * A full traversal is already the nearest permitted chunks, so nothing else is worth * running. An underfilled one is the signal that visibility removed neighbours the graph @@ -1883,7 +1935,9 @@ async function selectVectorResults(params: SearchParams): Promise= MAX_VECTOR_CANDIDATES, + exhausted: + (exhausted ?? selected.length < candidateLimit) || + candidateLimit >= MAX_VECTOR_CANDIDATES, filled, } annotateSearchDiagnostics({ @@ -2021,10 +2075,18 @@ export async function executeKeywordSearch(params: KeywordSearchParams): Promise * A caller reaching past the permitted-set limit reads much of the index, so ranking every * match before checking access is the leg's whole cost for a common term. Where the Tin * projection is complete, BM25 ranks inside the bases first and access is checked only on the - * top of that ranking. + * top of that ranking. A bounded set past the exact-ranking size is read on the row like an + * unbounded one: the bounded read materializes every chunk of the set before it matches a + * term, where a ranking decided on the row costs what the term matches. */ + const accessPlan = access.kind === 'user' ? params.accessPlan : undefined + const largePermittedSet = + accessPlan !== undefined && + params.permitted?.kind === 'bounded' && + params.permitted.documents.length >= PERMITTED_EXACT_DOCUMENT_LIMIT + const onRowReader = params.permitted?.kind === 'unbounded' || largePermittedSet let tinQuery: Awaited> = null - if (params.permitted?.kind === 'unbounded' && tagFilterConditions.length === 0) { + if (onRowReader && tagFilterConditions.length === 0) { try { tinQuery = await resolveTinKeywordQuery( params.searchIndexOnly === true, @@ -2038,9 +2100,7 @@ export async function executeKeywordSearch(params: KeywordSearchParams): Promise return [] } } - if (params.permitted?.kind === 'unbounded') - annotateSearchDiagnostics({ keywordRanking: tinQuery ? 'tin' : 'gin' }) - const accessPlan = access.kind === 'user' ? params.accessPlan : undefined + if (onRowReader) annotateSearchDiagnostics({ keywordRanking: tinQuery ? 'tin' : 'gin' }) /** A filled projection decides readability on the ranked row alone; none of its rows needs the document. */ const tinFilled = accessPlan && tinQuery @@ -2094,12 +2154,15 @@ export async function executeKeywordSearch(params: KeywordSearchParams): Promise /** * A resolved scope decides readability on the ranked row. The windows widen while the page * is short, a narrow reader's to a wide one sooner and no further, and what the widest - * cannot fill is left short rather than handed to a ranking over every match. + * cannot fill is left short rather than handed to a ranking over every match. A large + * bounded set is the exception on its first page: its bounded read was exhaustive, so the + * widest window that still falls short hands that page to the GIN ranking, which covers + * every match. A later page stays with Tin: the two rankers order differently, so an offset + * advanced through one cannot resume the other. */ const narrow = accessPlan !== undefined && - params.permitted?.kind === 'unbounded' && - !params.permitted.broad + ((params.permitted?.kind === 'unbounded' && !params.permitted.broad) || largePermittedSet) const windows: readonly number[] = narrow ? NARROW_KEYWORD_WINDOWS : TIN_KEYWORD_WINDOWS /** * A narrow reader's page is the readable remainder of a wide ranking, and that ranking is @@ -2160,7 +2223,9 @@ export async function executeKeywordSearch(params: KeywordSearchParams): Promise if ( page.candidates.length >= limit || page.ranked < window || - (accessPlan !== undefined && window === windows[windows.length - 1]) + (accessPlan !== undefined && + !(largePermittedSet && offset === 0) && + window === windows[windows.length - 1]) ) { return { candidates: page.candidates, nextOffset: offset + page.candidates.length } } @@ -2190,7 +2255,7 @@ export async function executeKeywordSearch(params: KeywordSearchParams): Promise * still re-applies the candidate predicate, so the restriction can only narrow. */ const permittedIds = - params.permitted?.kind === 'bounded' + params.permitted?.kind === 'bounded' && !largePermittedSet ? params.permitted.documents.map((entry) => entry.id) : undefined if (permittedIds?.length === 0) return { candidates: [], nextOffset: offset }