From 5c2e6c72f68ee9bd47d8474ddfc7f2fa2d16f76a Mon Sep 17 00:00:00 2001 From: Vikhyath Mondreti Date: Thu, 17 Sep 2026 14:10:06 -0700 Subject: [PATCH 1/5] fix(knowledge): share compact vector retrieval across access scopes --- .../kb-block-search.integration.ts | 2 +- .../search-latency.integration.ts | 240 ++++++++- apps/sim/lib/knowledge/search/budget.test.ts | 2 +- apps/sim/lib/knowledge/search/diagnostics.ts | 2 +- apps/sim/lib/knowledge/search/queries.test.ts | 503 +++++++++--------- apps/sim/lib/knowledge/search/queries.ts | 259 +++------ 6 files changed, 554 insertions(+), 454 deletions(-) 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 dba85e74b3d..b54266037d7 100644 --- a/apps/sim/lib/knowledge/__integration__/kb-block-search.integration.ts +++ b/apps/sim/lib/knowledge/__integration__/kb-block-search.integration.ts @@ -124,7 +124,7 @@ 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(36) + 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) } finally { diff --git a/apps/sim/lib/knowledge/__integration__/search-latency.integration.ts b/apps/sim/lib/knowledge/__integration__/search-latency.integration.ts index 0c89994ee4f..b4183e0c4d3 100644 --- a/apps/sim/lib/knowledge/__integration__/search-latency.integration.ts +++ b/apps/sim/lib/knowledge/__integration__/search-latency.integration.ts @@ -1,5 +1,6 @@ /** Real Assistant tool, application authorization, PostgreSQL/pgvector, and result processing. */ import { readFileSync, statSync, writeFileSync } from 'node:fs' +import type { Principal } from '@sim/auth/principal' import { db } from '@sim/db' import { copilotChats, @@ -8,6 +9,7 @@ import { document, embedding, knowledgeBase, + knowledgeBaseTagDefinitions, knowledgeConnector, knowledgeConnectorMember, knowledgeDocumentObservation, @@ -39,6 +41,7 @@ import { seedKnowledgeAclFixture, seedKnowledgeMemberFixture, } from '@/lib/knowledge/__integration__/seed-source-access-fixture' +import { type KnowledgeSearchTagFilter, searchKnowledge } from '@/lib/knowledge/application/search' import { SearchBudget, SearchDeadlineError, @@ -124,7 +127,7 @@ const report: Record = { dimensions, candidateDimensions, chunksPerDocument, - sql: 'Captured from the real Assistant tool; no hand-written search query', + sql: 'Captured from real search application adapters; no hand-written retrieval query', providers: 'Embedding and source-permission HTTP responses are controlled; internal search and authorization code is real', vectors: @@ -197,16 +200,21 @@ function explainNodes(node: ExplainNode): ExplainNode[] { } /** Broad ranking must stop the ordered ANN scan instead of sorting every accessible chunk. */ -function assertIndexedCandidates(plan: ExplainNode, candidateLimit: number) { +function assertIndexedCandidates( + plan: ExplainNode, + candidateLimit: number, + width = candidateDimensions +) { + const indexName = + width === 1536 + ? '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'] === 'embedding_search_512_cosine_hnsw_idx' && node['Actual Loops'] > 0 - ) + candidateNodes.some((node) => node['Index Name'] === indexName && node['Actual Loops'] > 0) ).toBe(true) expect(candidateNodes.some((node) => node['Node Type'] === 'Sort')).toBe(false) expect( @@ -272,7 +280,7 @@ async function prepareOrganizationSample(label: string) { const diagnosticSchema = z .object({ - surface: z.enum(['dashboard', 'copilot']), + surface: z.enum(['dashboard', 'copilot', 'workflow', 'api']), outcome: z.enum(['success', 'partial']), elapsedMs: z.number(), vectorBudgetMs: z.number().positive(), @@ -301,7 +309,12 @@ const resultSchema = z.object({ success: z.literal(true), data: z.object({ results: z.array( - z.object({ documentId: z.string(), content: z.string(), knowledgeBaseId: z.string() }) + z.object({ + documentId: z.string(), + content: z.string(), + knowledgeBaseId: z.string(), + embeddingId: z.string().optional(), + }) ), }), }) @@ -368,6 +381,29 @@ async function searchDashboard( } } +async function searchWorkspaceKb( + query = 'Orion deployment', + options: { principal?: Principal; tagFilters?: KnowledgeSearchTagFilter[] } = {} +) { + const result = await searchKnowledge.execute({ + principal: options.principal ?? { + kind: 'workspace_api_key', + workspaceId: ids.workspaceId, + keyId: 'fixture-search-key', + }, + input: { + workspaceId: ids.workspaceId, + knowledgeBaseIds: [ids.knowledgeBaseId], + query, + topK: 15, + searchMode: 'vector', + surface: 'workflow', + tagFilters: options.tagFilters, + }, + }) + return resultSchema.parse({ success: true, data: result }) +} + /** Allow either index or filtered plans, but require successful retrieval within the real surface budget. */ function expectCompleteVectorSearch(diagnostics: z.infer) { const budget = diagnostics.surface === 'dashboard' ? 3000 : 8000 @@ -478,10 +514,12 @@ async function sample( }) saveReport() if (query.query.includes('WITH visible_search_documents')) { - expect(query.query).toContain('"embedding_search"."vector_512"') - expect(diagnostics.vectorCandidateDimensions).toBe(candidateDimensions) + const width = diagnostics.vectorCandidateDimensions! + expect(query.query).toContain( + `"embedding_search"."${width === 1536 ? 'vector' : `vector_${width}`}"` + ) expect(diagnostics.vectorCandidateLimit).toBeGreaterThan(0) - assertIndexedCandidates(parsedPlan[0].Plan, diagnostics.vectorCandidateLimit!) + assertIndexedCandidates(parsedPlan[0].Plan, diagnostics.vectorCandidateLimit!, width) } if (query.query.includes('WITH visible_keyword_documents')) { assertScalarKeywordSorts(parsedPlan[0].Plan) @@ -495,7 +533,7 @@ async function sample( return { result, plans, diagnostics } } -describe.skipIf(!enabled)('Assistant search latency on a realistic indexed corpus', () => { +describe.skipIf(!enabled)('Knowledge search latency on a realistic indexed corpus', () => { beforeAll(async () => { if ( [chunkCount, unrelatedChunkCount].some( @@ -526,7 +564,7 @@ describe.skipIf(!enabled)('Assistant search latency on a realistic indexed corpu .object({ input: z.array(z.string()).length(1), encoding_format: z.literal('base64'), - model: z.literal('text-embedding-3-small'), + model: z.enum(['text-embedding-3-small', 'text-embedding-ada-002']), }) .parse(JSON.parse(String(init?.body))) embeddingCalls += body.input.length @@ -1201,6 +1239,182 @@ describe.skipIf(!enabled)('Assistant search latency on a realistic indexed corpu for (const diagnostics of completed) expectCompleteVectorSearch(diagnostics) }, 180_000) + it('uses compact indexed ranking for workspace KBs with stale estimates and private neighbors', async () => { + const originalAcl = `u:${ids.aliceId}@fixture.test` + await db.execute(sql`ALTER TABLE document SET (autovacuum_enabled = false)`) + try { + /** Analyze a narrow scope, then grow it without updating the planner's ACL histogram. */ + await db.execute(sql`UPDATE document SET acl = CASE WHEN external_id::int % 10 = 1 + THEN ARRAY['pub'] ELSE ARRAY[${originalAcl}] END + WHERE knowledge_base_id = ${ids.knowledgeBaseId}`) + await db.execute(sql`ANALYZE document`) + await db.execute(sql`UPDATE document SET acl = CASE WHEN external_id::int % 5 <> 0 + THEN ARRAY['pub'] ELSE ARRAY[${originalAcl}] END + WHERE knowledge_base_id = ${ids.knowledgeBaseId}`) + for (const topic of [0, 11, 23]) { + const label = `workspace-kb.topic.${topic}` + await prepareOrganizationSample(label) + const { result, plans, diagnostics } = await sample(label, () => + searchWorkspaceKb(`Topic ${topic} deployment`) + ) + expectCompleteVectorSearch(diagnostics) + expect(diagnostics.accessScopeKind).toBe('workspace') + expect(diagnostics.vectorRanking).toBe('candidate-rerank') + expect(result.data.results).toHaveLength(15) + expect(plans.some((plan) => plan.kind === 'vector')).toBe(true) + for (const row of result.data.results) { + expect(row.knowledgeBaseId).toBe(ids.knowledgeBaseId) + expect(Number(row.documentId.split('-doc-')[1]) % 5).not.toBe(0) + } + const expected = await db.execute<{ id: string }>(sql` + SELECT embedding.id FROM embedding JOIN document ON document.id = embedding.document_id + WHERE embedding.knowledge_base_id = ${ids.knowledgeBaseId} AND embedding.enabled + AND document.acl = ARRAY['pub']::text[] + ORDER BY (embedding.embedding <=> ${JSON.stringify(topicVector(topic))}::vector) + 0, embedding.id + LIMIT 15 + `) + const expectedIds = new Set(expected.map(({ id }) => id)) + const recall = + result.data.results.filter((row) => expectedIds.has(row.embeddingId!)).length / + expected.length + expect(recall).toBeGreaterThanOrEqual(0.95) + report[`${label}.recall`] = { neighbors: expected.length, recall } + saveReport() + } + await db + .update(knowledgeBase) + .set({ embeddingModel: 'text-embedding-ada-002' }) + .where(eq(knowledgeBase.id, ids.knowledgeBaseId)) + const fullWidth = await sample('workspace-kb.full-width', () => searchWorkspaceKb()) + expectCompleteVectorSearch(fullWidth.diagnostics) + expect(fullWidth.diagnostics.vectorCandidateDimensions).toBe(dimensions) + expect(fullWidth.result.data.results).toHaveLength(15) + await db + .update(knowledgeBase) + .set({ embeddingModel: 'text-embedding-3-small' }) + .where(eq(knowledgeBase.id, ids.knowledgeBaseId)) + + const workflowId = generateId() + const scheduled: Principal = { + kind: 'delegated', + serviceId: 'executor', + workspaceId: ids.workspaceId, + delegationId: generateId(), + audience: 'sim:knowledge', + issuedAt: new Date(), + expiresAt: new Date(Date.now() + 60_000), + delegationContext: { + kind: 'workflow_execution', + workflowId, + principal: { + kind: 'system', + serviceId: 'schedule', + workspaceId: ids.workspaceId, + workflowId, + }, + currentWorkflow: { workflowId, mode: 'deployment', deploymentVersionId: generateId() }, + }, + } + const scheduledResult = await sample('workspace-kb.scheduled', () => + searchWorkspaceKb('Orion deployment', { principal: scheduled }) + ) + expectCompleteVectorSearch(scheduledResult.diagnostics) + expect(scheduledResult.diagnostics.accessScopeKind).toBe('workspace') + expect(scheduledResult.result.data.results).toHaveLength(15) + + for (const concurrency of [2, 8]) { + const label = `workspace-kb.concurrent.${concurrency}` + await prepareOrganizationSample(label) + diagnosticLog?.mockClear() + const started = performance.now() + const results = await Promise.all( + Array.from({ length: concurrency }, (_, index) => + searchWorkspaceKb(`Topic ${index * 3} deployment`) + ) + ) + const completed = diagnosticLog!.mock.calls + .filter(([message]) => message === 'Knowledge search completed') + .map(([, metadata]) => diagnosticSchema.parse(metadata)) + report[label] = { + milliseconds: performance.now() - started, + resultCounts: results.map((result) => result.data.results.length), + diagnostics: completed, + } + saveReport() + expect(completed).toHaveLength(concurrency) + for (const result of results) expect(result.data.results).toHaveLength(15) + for (const diagnostics of completed) expectCompleteVectorSearch(diagnostics) + } + + await db.insert(knowledgeBaseTagDefinitions).values({ + id: generateId(), + knowledgeBaseId: ids.knowledgeBaseId, + tagSlot: 'tag1', + displayName: 'Fixture group', + }) + await db.execute(sql`UPDATE embedding SET tag1 = 'selected' WHERE knowledge_base_id = ${ids.knowledgeBaseId} + AND document_id IN (SELECT id FROM document WHERE knowledge_base_id = ${ids.knowledgeBaseId} AND external_id::int < 600)`) + const tagged = await sample('workspace-kb.tagged', () => + searchWorkspaceKb('Orion deployment', { + tagFilters: [{ tagName: 'Fixture group', operator: 'eq', value: 'selected' }], + }) + ) + expectCompleteVectorSearch(tagged.diagnostics) + expect(tagged.diagnostics.vectorRanking).toBe('candidate-rerank') + expect(tagged.result.data.results).toHaveLength(15) + for (const row of tagged.result.data.results) { + const ordinal = Number(row.documentId.split('-doc-')[1]) + expect(ordinal).toBeLessThan(600) + expect(ordinal % 5).not.toBe(0) + } + const expectedTagged = await db.execute<{ id: string }>(sql` + SELECT embedding.id FROM embedding JOIN document ON document.id = embedding.document_id + WHERE embedding.knowledge_base_id = ${ids.knowledgeBaseId} AND embedding.enabled + AND document.acl = ARRAY['pub']::text[] AND embedding.tag1 = 'selected' + ORDER BY (embedding.embedding <=> ${JSON.stringify(queryVector)}::vector) + 0, embedding.id LIMIT 15 + `) + const taggedIds = new Set(expectedTagged.map(({ id }) => id)) + const taggedRecall = + tagged.result.data.results.filter((row) => taggedIds.has(row.embeddingId!)).length / + expectedTagged.length + expect(taggedRecall).toBeGreaterThanOrEqual(0.95) + report['workspace-kb.tagged.recall'] = { + neighbors: expectedTagged.length, + recall: taggedRecall, + } + saveReport() + + /** Workspace credentials cannot keep reading a source after its access rewrite begins. */ + await db + .update(knowledgeConnector) + .set({ accessRewritePending: true }) + .where(eq(knowledgeConnector.id, ids.connectorId)) + const denied = await sample('workspace-kb.revoked', () => searchWorkspaceKb()) + expectCompleteVectorSearch(denied.diagnostics) + expect(denied.result.data.results).toEqual([]) + } finally { + await db + .delete(knowledgeBaseTagDefinitions) + .where(eq(knowledgeBaseTagDefinitions.knowledgeBaseId, ids.knowledgeBaseId)) + await db.execute( + sql`UPDATE embedding SET tag1 = NULL WHERE knowledge_base_id = ${ids.knowledgeBaseId} AND tag1 = 'selected'` + ) + await db + .update(knowledgeBase) + .set({ embeddingModel: 'text-embedding-3-small' }) + .where(eq(knowledgeBase.id, ids.knowledgeBaseId)) + await db + .update(knowledgeConnector) + .set({ accessRewritePending: false }) + .where(eq(knowledgeConnector.id, ids.connectorId)) + await db.execute( + sql`UPDATE document SET acl = ARRAY[${originalAcl}] WHERE knowledge_base_id = ${ids.knowledgeBaseId}` + ) + await db.execute(sql`ALTER TABLE document RESET (autovacuum_enabled)`) + await db.execute(sql`ANALYZE document`) + } + }, 180_000) + it('checks live reader access on every search, including after revocation', async () => { await seedSearchReaderFixture(ids) const allowed = await sample('live.allowed', () => search()) diff --git a/apps/sim/lib/knowledge/search/budget.test.ts b/apps/sim/lib/knowledge/search/budget.test.ts index d5e04601737..941fbdbd98e 100644 --- a/apps/sim/lib/knowledge/search/budget.test.ts +++ b/apps/sim/lib/knowledge/search/budget.test.ts @@ -14,7 +14,7 @@ describe('search SQL deadline', () => { now = 101 return ['candidate'] }) - await budget.query('vector.ann', run) + await budget.query('vector.candidate_search', run) await expect(budget.query('vector.exact', run)).rejects.toBeInstanceOf(SearchDeadlineError) expect(run).toHaveBeenCalledTimes(1) }) diff --git a/apps/sim/lib/knowledge/search/diagnostics.ts b/apps/sim/lib/knowledge/search/diagnostics.ts index 9f5f3726575..e308ea56c65 100644 --- a/apps/sim/lib/knowledge/search/diagnostics.ts +++ b/apps/sim/lib/knowledge/search/diagnostics.ts @@ -46,7 +46,6 @@ export type SearchStage = | `${RetrievalLeg}.sql` | 'vector.settings' | 'vector.probe' - | 'vector.ann' | 'vector.rerank' | 'vector.exact' | 'vector.candidate_search' @@ -72,6 +71,7 @@ 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' vectorBudgetMs?: number vectorCandidateLimit?: number diff --git a/apps/sim/lib/knowledge/search/queries.test.ts b/apps/sim/lib/knowledge/search/queries.test.ts index b7a5b9ebabe..abf33e7dcd7 100644 --- a/apps/sim/lib/knowledge/search/queries.test.ts +++ b/apps/sim/lib/knowledge/search/queries.test.ts @@ -316,21 +316,75 @@ describe('getStructuredTagFilters', () => { }) }) -describe('KB block vector retrieval', () => { +describe('workspace-scoped vector retrieval', () => { + const access = { kind: 'workspace' as const, tokens: WORKSPACE_ACCESS_TOKENS } + const getForConnectors = vi.fn() const params: SearchParams = { knowledgeBaseIds: ['kb-small'], topK: 2, - access: { kind: 'workspace', tokens: WORKSPACE_ACCESS_TOKENS }, + access, + accessProvider: { + get: async () => access, + getForConnectors, + getForDocuments: async () => access, + liveSourceConnectorCondition: async () => null, + }, queryVector: { vector: '[0.1,0.2]', dimensions: 1536, model: 'text-embedding-3-small' }, distanceThreshold: 1, } + const probe = Array.from({ length: 400 }, (_, index) => ({ id: `probe-${index}` })) + const candidates = Array.from({ length: 400 }, (_, index) => ({ + id: `candidate-${index}`, + initial_count: 400, + })) + const ranked = [ + { + id: 'near', + documentId: 'near-doc', + connectorId: null, + liveAuthorizationSource: false, + distance: 0.1, + }, + { + id: 'far', + documentId: 'far-doc', + connectorId: null, + liveAuthorizationSource: false, + distance: 0.2, + }, + ] + let probeRows: Array<{ id: string }> + let failSettings: unknown + let failCandidates: unknown - beforeEach(() => resetDbChainMock()) + beforeEach(() => { + resetDbChainMock() + getForConnectors.mockReset() + probeRows = 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('WITH visible_search_documents')) { + if (failCandidates) throw failCandidates + return candidates + } + if (statement.includes('WITH scored_search_candidates')) return ranked + return [] + }) + }) afterEach(() => { vi.restoreAllMocks() 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) => { @@ -347,67 +401,205 @@ describe('KB block vector retrieval', () => { expect(budget.timedOut).toBe(true) expect(dbChainMockFns.transaction).not.toHaveBeenCalled() expect(dbChainMockFns.select).not.toHaveBeenCalled() + expect(dbChainMockFns.execute).not.toHaveBeenCalled() } ) - it('ranks all candidates in a small KB exactly instead of traversing the shared vector index', async () => { - queueTableRows(schemaMock.embedding, [{ id: 'near' }, { id: 'far' }]) - queueTableRows(schemaMock.embedding, [ - { id: 'far', distance: 0.2 }, - { id: 'near', distance: 0.1 }, - ]) - const rows = await handleVectorOnlySearch(params) + it('ranks an exhausted visible scope exactly and rechecks access before returning content', async () => { + probeRows = ranked + queueTableRows(schemaMock.embedding, 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') + } + expect(getForConnectors).not.toHaveBeenCalled() + }) + + 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') + )! + expect(candidate.sql).toContain('CROSS JOIN LATERAL') + expect(candidate.sql).toContain('LIMIT 1') + const serialized = JSON.stringify(candidate) + expect(serialized).toContain('subvector') + expect(serialized).toContain('required_clause') + expect(serialized).toContain('credential') + expect(serialized).toContain('aclVerifiedAt') + expect(serialized).toContain('accessRewritePending') + expect(serialized).toContain('organizationSearchIntegration') + expect(serialized).toContain(String(schemaMock.embeddingSearch.vector512)) + expect(candidate.params).not.toContain(schemaMock.embedding.embedding) + const rerank = statements().find((query) => + query.sql.includes('WITH scored_search_candidates') + )! + expect(rerank.sql).toContain('MATERIALIZED') + expect(JSON.stringify(rerank)).toContain(String(schemaMock.embedding.embedding)) + expect(JSON.stringify(rerank)).toContain('candidate-399') + expect(JSON.stringify(rerank)).not.toContain('probe-399') + expect(getForConnectors).not.toHaveBeenCalled() + }) + + it('keeps workspace-authorized sources eligible when hydration needs another page', async () => { + const initial = ranked.map((row) => ({ + ...row, + connectorId: 'workspace-source', + liveAuthorizationSource: true, + })) + const next = { ...initial[1], id: 'next', distance: 0.3 } + queueTableRows(schemaMock.embedding, initial) + queueTableRows(schemaMock.embedding, [initial[0]]) + queueTableRows(schemaMock.embedding, [next]) + queueTableRows(schemaMock.embedding, [next]) + const rows = await handleVectorOnlySearch({ + ...params, + filters: { documentIds: ['near-doc', 'far-doc'] }, + }) + expect(rows.map((row) => row.id)).toEqual(['near', 'next']) + expect(dbChainMockFns.offset.mock.calls.map(([offset]) => offset)).toEqual([0, 2]) + expect(JSON.stringify(dbChainMockFns.where.mock.calls)).not.toContain('workspace-source') + expect(getForConnectors).not.toHaveBeenCalled() + }) + + 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(dbChainMockFns.execute).not.toHaveBeenCalled() expect(Object.keys(dbChainMockFns.select.mock.calls[0][0])).toEqual(['id']) - expect(dbChainMockFns.limit).toHaveBeenNthCalledWith(1, 201) - expect(render(dbChainMockFns.orderBy.mock.calls[0][0]).sql).toContain('+ 0') - expect( - hasMockCondition( - dbChainMockFns.where.mock.calls[1][0], - (node) => - node.type === 'inArray' && - node.column === schemaMock.embedding.id && - JSON.stringify(node.values) === JSON.stringify(['near', 'far']) - ) - ).toBe(true) + const candidate = statements().find((query) => + query.sql.includes('WITH visible_search_documents') + )! + 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') + expect(getForConnectors).not.toHaveBeenCalled() }) - it.each([1, 200, 201])( - 'reports a %i-candidate SQL timeout as partial, not an empty complete result', - async (count) => { - queueTableRows( - schemaMock.embedding, - Array.from({ length: count }, (_, index) => ({ id: `candidate-${index}` })) - ) - dbChainMockFns.orderBy - .mockImplementationOnce(dbChainMockFns.orderBy.getMockImplementation()!) - .mockRejectedValueOnce(new Error('Statement canceled', { cause: { code: '57014' } })) - const result = await retrieveKnowledgeSearch({ - ...params, - query: 'fixture policy', - searchMode: 'vector', + it('ranks all selected KBs together instead of capping how many results one KB can contribute', async () => { + const knowledgeBaseIds = ['kb-1', 'kb-2', 'kb-3', 'kb-4', 'kb-5'] + queueTableRows( + schemaMock.embedding, + ranked.map((row) => ({ ...row, knowledgeBaseId: 'kb-1' })) + ) + 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') + ) + expect(candidateQueries).toHaveLength(1) + for (const id of knowledgeBaseIds) expect(JSON.stringify(candidateQueries[0])).toContain(id) + expect(dbChainMockFns.transaction).toHaveBeenCalledOnce() + }) + + it.each(['vector.probe', '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 + vi.spyOn(SearchBudget.prototype, 'query').mockImplementation(function ( + this: SearchBudget, + stage: SearchStage, + run: (executor: SearchExecutor) => PromiseLike + ) { + if (stage === failedStage) + return Promise.reject(new Error('Statement canceled', { cause: { code: '57014' } })) + return query.call(this, stage, run) as Promise }) - expect(result).toEqual({ + expect( + await retrieveKnowledgeSearch({ ...params, query: 'fixture policy', searchMode: 'vector' }) + ).toEqual({ rows: [], retrieval: { status: 'partial', timedOutLegs: ['vector'] }, }) - const usesAnn = dbChainMockFns.execute.mock.calls.some(([statement]) => - render(statement).sql.includes('hnsw.iterative_scan') - ) - expect(usesAnn).toBe(count > 200) } ) - it('does not convert an unexpected ranking error into partial retrieval', async () => { - queueTableRows(schemaMock.embedding, [{ id: 'candidate' }]) - const failure = new Error('Connection lost', { cause: { code: '08006' } }) - dbChainMockFns.orderBy - .mockImplementationOnce(dbChainMockFns.orderBy.getMockImplementation()!) - .mockRejectedValueOnce(failure) + it('shares the remaining deadline across candidate selection, reranking and hydration', async () => { + vi.spyOn(performance, 'now').mockReturnValue(0) + const query = SearchBudget.prototype.query + vi.spyOn(SearchBudget.prototype, 'query').mockImplementation(async function ( + this: SearchBudget, + stage: SearchStage, + 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 + }) + queueTableRows(schemaMock.embedding, ranked) + await handleVectorOnlySearch({ ...params, budget: new SearchBudget('vector', 100) }) + expect( + statements() + .filter((query) => query.sql.includes('statement_timeout')) + .map((query) => query.params[0]) + ).toEqual(['100', '70', '70', '40', '20']) + }) + + it('does not convert an unexpected candidate failure into partial retrieval', async () => { + failCandidates = new Error('Connection lost', { cause: { code: '08006' } }) await expect( retrieveKnowledgeSearch({ ...params, query: 'fixture policy', searchMode: 'vector' }) - ).rejects.toBe(failure) + ).rejects.toBe(failCandidates) + }) + + it('does not treat a missing query object as unsupported scan settings', async () => { + failCandidates = new Error('Query object is missing', { cause: { code: '42704' } }) + await expect(handleVectorOnlySearch(params)).rejects.toBe(failCandidates) + failCandidates = undefined + queueTableRows(schemaMock.embedding, ranked) + await handleVectorOnlySearch(params) + expect(statements().filter((query) => query.sql.includes('hnsw.iterative_scan'))).toHaveLength( + 2 + ) + }) + + it('retries unsupported settings after cooldown without changing the candidate query', async () => { + vi.useFakeTimers() + vi.setSystemTime(new Date(0)) + failSettings = new Error('Failed settings query', { cause: { code: '42704' } }) + queueTableRows(schemaMock.embedding, ranked) + await handleVectorOnlySearch(params) + queueTableRows(schemaMock.embedding, ranked) + await handleVectorOnlySearch(params) + expect(statements().filter((query) => query.sql.includes('hnsw.iterative_scan'))).toHaveLength( + 1 + ) + const queries = statements().filter((query) => + query.sql.includes('WITH visible_search_documents') + ) + expect(queries).toHaveLength(2) + expect(JSON.stringify(queries[0])).toBe(JSON.stringify(queries[1])) + await vi.advanceTimersByTimeAsync(10 * 60 * 1000 + 1) + failSettings = undefined + queueTableRows(schemaMock.embedding, ranked) + await handleVectorOnlySearch(params) + expect(statements().filter((query) => query.sql.includes('hnsw.iterative_scan'))).toHaveLength( + 2 + ) }) it('reports incomplete retrieval for 18 expired pool waiters without starting their SQL later', async () => { @@ -452,200 +644,6 @@ describe('KB block vector retrieval', () => { expect(dbChainMockFns.select).not.toHaveBeenCalled() expect(dbChainMockFns.execute).not.toHaveBeenCalled() }) - - it.each([1, 201])( - 'shares the remaining SQL budget between the probe and %i-candidate ranking', - async (count) => { - vi.spyOn(performance, 'now').mockReturnValue(0) - queueTableRows( - schemaMock.embedding, - Array.from({ length: count }, (_, index) => ({ id: `candidate-${index}` })) - ) - const query = SearchBudget.prototype.query - vi.spyOn(SearchBudget.prototype, 'query').mockImplementation(async function ( - this: SearchBudget, - stage: SearchStage, - run: (executor: SearchExecutor) => PromiseLike - ) { - const runQuery: SearchBudget['query'] = query.bind(this) - const result = await runQuery(stage, run) - if (stage === 'vector.probe') vi.spyOn(performance, 'now').mockReturnValue(30) - return result - }) - await handleVectorOnlySearch({ ...params, budget: new SearchBudget('vector', 100) }) - const timeouts = dbChainMockFns.execute.mock.calls - .map(([statement]) => render(statement)) - .filter((statement) => statement.sql.includes('statement_timeout')) - .map((statement) => statement.params[0]) - expect(timeouts).toEqual(count === 1 ? ['100', '70'] : ['100', '70', '70']) - } - ) -}) - -describe('vector scan settings', () => { - const largeProbe = Array.from({ length: 201 }, (_, index) => ({ id: `probe-${index}` })) - const params: SearchParams = { - knowledgeBaseIds: ['kb-small'], - topK: 2, - access: { kind: 'workspace', tokens: WORKSPACE_ACCESS_TOKENS }, - queryVector: { vector: '[0.1,0.2]', dimensions: 1536, model: 'text-embedding-3-small' }, - distanceThreshold: 0.8, - } - - beforeEach(() => { - resetDbChainMock() - queueTableRows(schemaMock.embedding, largeProbe) - }) - - afterEach(() => { - vi.useRealTimers() - }) - - it('tunes an overflowing KB scope without limiting ANN to the probe prefix', async () => { - queueTableRows(schemaMock.embedding, [ - { id: 'far', distance: 0.2 }, - { id: 'near', distance: 0.1 }, - ]) - const rows = await handleVectorOnlySearch(params) - expect(rows.map((row) => row.id)).toEqual(['near', 'far']) - expect(dbChainMockFns.transaction).toHaveBeenCalledOnce() - expect(dbChainMockFns.execute).toHaveBeenCalledOnce() - expect(render(dbChainMockFns.execute.mock.calls[0][0])).toEqual({ - sql: "SELECT set_config('hnsw.iterative_scan', 'relaxed_order', true), set_config('hnsw.max_scan_tuples', ?, true)", - params: ['20000'], - }) - expect(dbChainMockFns.execute.mock.invocationCallOrder[0]).toBeLessThan( - dbChainMockFns.select.mock.invocationCallOrder[1] - ) - expect( - hasMockCondition( - dbChainMockFns.where.mock.calls[1][0], - (node) => - node.type === 'inArray' && - node.column === schemaMock.embedding.knowledgeBaseId && - Array.isArray(node.values) && - node.values.includes('kb-small') - ) - ).toBe(true) - expect(dbChainMockFns.limit).toHaveBeenCalledWith(2) - expect(dbChainMockFns.limit).toHaveBeenCalledTimes(2) - expect( - hasMockCondition( - dbChainMockFns.where.mock.calls[1][0], - (node) => node.type === 'inArray' && node.column === schemaMock.embedding.id - ) - ).toBe(false) - expect(Object.keys(dbChainMockFns.select.mock.calls[1][0])).toEqual(['id', 'distance']) - const ranked = dbChainMockFns.from.mock.calls[2][0] - expect(dbChainMockFns.select.mock.calls[2][0].distance).toBe(ranked.distance) - expect(dbChainMockFns.innerJoin).toHaveBeenCalledWith( - schemaMock.embedding, - expect.objectContaining({ - type: 'eq', - left: schemaMock.embedding.id, - right: ranked.id, - }) - ) - expect(dbChainMockFns.orderBy).toHaveBeenLastCalledWith(ranked.distance) - expect(dbChainMockFns.limit.mock.invocationCallOrder[1]).toBeLessThan( - dbChainMockFns.select.mock.invocationCallOrder[2] - ) - }) - - it('tunes each KB leg and trims their sorted merge', async () => { - const knowledgeBaseIds = ['kb-1', 'kb-2', 'kb-3', 'kb-4', 'kb-5'] - for (let index = 0; index < knowledgeBaseIds.length; index++) { - if (index > 0) queueTableRows(schemaMock.embedding, largeProbe) - queueTableRows(schemaMock.embedding, [{ id: `row-${index}`, distance: (5 - index) / 10 }]) - } - const rows = await handleVectorOnlySearch({ ...params, knowledgeBaseIds }) - expect(rows.map((row) => row.id)).toEqual(['row-4', 'row-3']) - expect(dbChainMockFns.transaction).toHaveBeenCalledTimes(5) - expect(dbChainMockFns.execute).toHaveBeenCalledTimes(5) - expect(dbChainMockFns.select).toHaveBeenCalledTimes(15) - for (const kbId of knowledgeBaseIds) - expect( - dbChainMockFns.where.mock.calls.some(([condition]) => - hasMockCondition( - condition, - (node) => - node.type === 'eq' && - node.left === schemaMock.embedding.knowledgeBaseId && - node.right === kbId - ) - ) - ).toBe(true) - }) - - it('tunes tag vector queries while keeping tag-only reads outside a vector transaction', async () => { - const filtered: SearchParams = { - ...params, - structuredFilters: [{ tagSlot: 'tag1', fieldType: 'text', operator: 'eq', value: 'release' }], - } - queueTableRows(schemaMock.embedding, [ - { id: 'far', distance: 0.2 }, - { id: 'near', distance: 0.1 }, - ]) - expect((await handleTagAndVectorSearch(filtered)).map((row) => row.id)).toEqual(['near', 'far']) - expect(dbChainMockFns.transaction).toHaveBeenCalledOnce() - expect( - hasMockCondition(dbChainMockFns.where.mock.calls[0][0], (node) => { - if (typeof node.toSQL !== 'function') return false - const condition = render(node) - return ( - condition.sql === 'LOWER(?) = LOWER(?)' && - condition.params[0] === schemaMock.embedding.tag1 && - condition.params[1] === 'release' - ) - }) - ).toBe(true) - resetDbChainMock() - await handleTagOnlySearch(filtered) - expect(dbChainMockFns.transaction).not.toHaveBeenCalled() - expect(dbChainMockFns.execute).not.toHaveBeenCalled() - }) - - it('retries unsupported settings after the cooldown without changing the query on fallback', async () => { - vi.useFakeTimers() - vi.setSystemTime(new Date(0)) - dbChainMockFns.execute.mockRejectedValueOnce( - new Error('Failed settings query', { cause: { code: '42704' } }) - ) - queueTableRows(schemaMock.embedding, [{ id: 'fallback', distance: 0.1 }]) - expect((await handleVectorOnlySearch(params)).map((row) => row.id)).toEqual(['fallback']) - expect(dbChainMockFns.execute).toHaveBeenCalledOnce() - expect(dbChainMockFns.select).toHaveBeenCalledTimes(3) - await handleVectorOnlySearch(params) - expect(dbChainMockFns.transaction).toHaveBeenCalledOnce() - await vi.advanceTimersByTimeAsync(10 * 60 * 1000 + 1) - queueTableRows(schemaMock.embedding, largeProbe) - await handleVectorOnlySearch(params) - expect(dbChainMockFns.transaction).toHaveBeenCalledTimes(2) - expect(dbChainMockFns.execute).toHaveBeenCalledTimes(2) - }) - - it('propagates an unrelated settings failure without ranking or disabling later tuning', async () => { - const failure = { code: '08006', message: 'Connection lost' } - dbChainMockFns.execute.mockRejectedValueOnce(failure) - await expect(handleVectorOnlySearch(params)).rejects.toBe(failure) - expect(dbChainMockFns.select).toHaveBeenCalledOnce() - expect(dbChainMockFns.orderBy).not.toHaveBeenCalled() - queueTableRows(schemaMock.embedding, largeProbe) - await handleVectorOnlySearch(params) - expect(dbChainMockFns.execute).toHaveBeenCalledTimes(2) - }) - - it('does not retry a query 42704 or classify it as unsupported scan settings', async () => { - const failure = { code: '42704', message: 'Query object is missing' } - dbChainMockFns.orderBy - .mockImplementationOnce(dbChainMockFns.orderBy.getMockImplementation()!) - .mockRejectedValueOnce(failure) - await expect(handleVectorOnlySearch(params)).rejects.toBe(failure) - expect(dbChainMockFns.select).toHaveBeenCalledTimes(3) - queueTableRows(schemaMock.embedding, largeProbe) - await handleVectorOnlySearch(params) - expect(dbChainMockFns.execute).toHaveBeenCalledTimes(2) - }) }) describe('workspace search filters before ranking', () => { @@ -663,8 +661,8 @@ describe('workspace search filters before ranking', () => { } beforeEach(() => resetDbChainMock()) - function expectScopeOnEveryQuery(skipIdentityProbe = false) { - const queries = dbChainMockFns.where.mock.calls.slice(skipIdentityProbe ? 1 : 0) + function expectScopeOnEveryQuery() { + const queries = dbChainMockFns.where.mock.calls expect(queries.length).toBeGreaterThan(0) for (const [condition] of queries) { expect( @@ -699,15 +697,14 @@ describe('workspace search filters before ranking', () => { it.each([handleVectorOnlySearch, handleTagOnlySearch, handleTagAndVectorSearch])( 'applies the full document scope to vector and tag searches', async (search) => { - const hasIdentityProbe = search !== handleTagOnlySearch - if (hasIdentityProbe) queueTableRows(schemaMock.embedding, [{ id: 'candidate' }]) + queueTableRows(schemaMock.embedding, [{ id: 'candidate' }]) await search({ ...params, structuredFilters: [ { tagSlot: 'tag1', fieldType: 'text', operator: 'eq', value: 'launch' }, ], }) - expectScopeOnEveryQuery(hasIdentityProbe) + expectScopeOnEveryQuery() } ) @@ -996,6 +993,8 @@ 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 === 'keyword') keywordPages.push(candidates) else queueTableRows(schemaMock.embedding, candidates) queueTableRows(schemaMock.embedding, [{ id: 'selected', content: 'verified result' }]) @@ -1020,7 +1019,9 @@ describe('live repository authorization follows ranked candidates', () => { expect(ranking).not.toContain('<=>') expect(ranking).not.toContain('"content"') } else { - expect(Object.keys(dbChainMockFns.select.mock.calls[0][0]).sort()).toEqual( + expect( + Object.keys(dbChainMockFns.select.mock.calls[mode === 'tags' ? 0 : 1][0]).sort() + ).toEqual( [ 'id', 'documentId', @@ -1062,6 +1063,8 @@ 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 === 'keyword') keywordPages.push(candidates) else queueTableRows(schemaMock.embedding, candidates) const hydrated = [{ id: 'gmail', content: 'current permitted content' }] @@ -1091,7 +1094,9 @@ 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 : 2) + expect(dbChainMockFns.select).toHaveBeenCalledTimes( + mode === 'keyword' ? 1 : mode === 'tags' ? 2 : 3 + ) } ) diff --git a/apps/sim/lib/knowledge/search/queries.ts b/apps/sim/lib/knowledge/search/queries.ts index e67a126fa5e..78958d8481d 100644 --- a/apps/sim/lib/knowledge/search/queries.ts +++ b/apps/sim/lib/knowledge/search/queries.ts @@ -47,8 +47,8 @@ const logger = createLogger('KnowledgeSearchQueries') /** SQLSTATE for an unrecognised configuration parameter — pgvector older than 0.8. */ const UNDEFINED_OBJECT_SQLSTATE = '42704' -/** Tuples a relaxed-order scan may visit before giving up on filling the limit. */ -const HNSW_MAX_SCAN_TUPLES = '20000' +/** 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. */ const CANDIDATE_HNSW_MAX_SCAN_TUPLES = '1000' const CANDIDATE_HNSW_EF_SEARCH = '1000' @@ -56,7 +56,6 @@ const CANDIDATE_HNSW_SCAN_MEM_MULTIPLIER = '2' const MIN_VECTOR_RERANK_CANDIDATES = 400 const MAX_VECTOR_RERANK_CANDIDATES = 1600 const VECTOR_RERANK_OVERSAMPLING = 8 -const MAX_EXACT_KB_VECTOR_CANDIDATES = 200 /** How long to stop trying the iterative-scan settings after the server rejected them. */ const HNSW_SETTINGS_UNSUPPORTED_RETRY_MS = 10 * 60 * 1000 @@ -71,10 +70,9 @@ let hnswSettingsUnsupportedUntil = 0 */ async function withVectorScanSettings( run: (executor: SearchExecutor) => Promise, - budget?: SearchBudget, - ranking: 'cosine' | 'candidate' = 'cosine' + budget?: SearchBudget ): Promise { - const stage = ranking === 'candidate' ? 'vector.candidate_search' : 'vector.ann' + const stage = 'vector.candidate_search' const untuned = () => runSearchQuery(budget, stage, run) if (Date.now() < hnswSettingsUnsupportedUntil) return untuned() const acquireStarted = performance.now() @@ -86,9 +84,7 @@ async function withVectorScanSettings( applyingSettings = true await measureSearchStage('vector.settings', () => tx.execute( - ranking === 'candidate' - ? sql`SELECT set_config('hnsw.iterative_scan', 'relaxed_order', true), set_config('hnsw.max_scan_tuples', ${CANDIDATE_HNSW_MAX_SCAN_TUPLES}, true), set_config('hnsw.ef_search', ${CANDIDATE_HNSW_EF_SEARCH}, true), set_config('hnsw.scan_mem_multiplier', ${CANDIDATE_HNSW_SCAN_MEM_MULTIPLIER}, true)` - : sql`SELECT set_config('hnsw.iterative_scan', 'relaxed_order', true), set_config('hnsw.max_scan_tuples', ${HNSW_MAX_SCAN_TUPLES}, true)` + sql`SELECT set_config('hnsw.iterative_scan', 'relaxed_order', true), set_config('hnsw.max_scan_tuples', ${CANDIDATE_HNSW_MAX_SCAN_TUPLES}, true), set_config('hnsw.ef_search', ${CANDIDATE_HNSW_EF_SEARCH}, true), set_config('hnsw.scan_mem_multiplier', ${CANDIDATE_HNSW_SCAN_MEM_MULTIPLIER}, true)` ) ) applyingSettings = false @@ -512,17 +508,18 @@ const SEARCH_READ_CANDIDATE_FIELDS = { )`, } -const LIVE_SEARCH_PAGE_SIZE = 200 -const LIVE_SEARCH_BUDGET_MS = 8000 +const AUTHORIZED_SEARCH_PAGE_SIZE = 200 +const AUTHORIZED_SEARCH_BUDGET_MS = 8000 /** * Verification follows ranked candidates, never the organization's source order. Denied * sources are excluded on refill, so many matches from one revoked source cannot - * consume every result slot. The existing vector tuple budget also bounds candidate work. + * consume every result slot. Candidate pages and the shared deadline bound authorization work. */ async function selectAuthorizedSearchResults(input: { leg: 'vector' | 'keyword' | 'tags' - accessProvider: KnowledgeAccessProvider + access: KnowledgeAccessScope + accessProvider?: KnowledgeAccessProvider filters?: WorkspaceSearchFilters signal?: AbortSignal budget?: SearchBudget @@ -535,8 +532,11 @@ async function selectAuthorizedSearchResults(input: { compareResults?: (a: SearchResult, b: SearchResult) => number hydrate: (ids: string[], access: KnowledgeAccessScope) => Promise }): Promise { - const deadline = Date.now() + LIVE_SEARCH_BUDGET_MS - const pageSize = Math.min(LIVE_SEARCH_PAGE_SIZE, Math.max(input.topK, 20)) + const deadline = Date.now() + AUTHORIZED_SEARCH_BUDGET_MS + const pageSize = Math.min( + AUTHORIZED_SEARCH_PAGE_SIZE, + input.accessProvider ? Math.max(input.topK, 20) : input.topK + ) const results = new Map() const excludedSources = new Set() const considered = new Set() @@ -545,7 +545,7 @@ async function selectAuthorizedSearchResults(input: { try { while ( results.size < input.topK && - scanned < Number(HNSW_MAX_SCAN_TUPLES) && + scanned < MAX_AUTHORIZED_SEARCH_CANDIDATES && (input.budget !== undefined || Date.now() < deadline) ) { input.signal?.throwIfAborted() @@ -576,7 +576,9 @@ async function selectAuthorizedSearchResults(input: { ), ] const access = await measureSearchStage(`${input.leg}.authorization`, () => - input.accessProvider.getForConnectors(connectorIds, input.signal) + input.accessProvider + ? input.accessProvider.getForConnectors(connectorIds, input.signal) + : input.access ) input.signal?.throwIfAborted() const grantedSources = new Set( @@ -590,6 +592,7 @@ async function selectAuthorizedSearchResults(input: { const excludedBefore = excludedSources.size for (const candidate of candidates) { if ( + input.accessProvider && candidate.liveAuthorizationSource && candidate.connectorId && !grantedSources.has(candidate.connectorId) @@ -686,6 +689,7 @@ export async function handleTagOnlySearch(params: SearchParams): Promise { - const { knowledgeBaseIds, topK, queryVector, distanceThreshold, access } = params - + const { queryVector, distanceThreshold } = params if (!queryVector || !distanceThreshold) { throw new Error('Query vector and distance threshold are required for vector-only search') } - - const strategy = getQueryStrategy(knowledgeBaseIds.length, topK) - - const distance = embeddingDistance(queryVector.dimensions, queryVector.vector) - if (params.accessProvider && access.kind === 'user') { - return selectLiveVectorResults(params, params.accessProvider, distance, [ - sql`${distance} < ${distanceThreshold}`, - ]) - } - /** - * A relaxed-order iterative scan may hand rows back slightly out of distance - * order, so both paths re-sort in memory before trimming to `topK`. - */ - if (strategy.useParallel) { - const parallelLimit = Math.ceil(topK / knowledgeBaseIds.length) + 5 - const allResults: SearchResult[] = [] - /** Keep one active KB leg per request so multi-base searches cannot monopolize the pool. */ - for (const kbId of knowledgeBaseIds) { - allResults.push( - ...(await selectScopedVectorResults( - params, - distance, - eq(embedding.knowledgeBaseId, kbId), - parallelLimit - )) - ) - if (params.budget?.timedOut) break - } - return allResults.sort((a, b) => a.distance - b.distance).slice(0, topK) - } - const rows = await selectScopedVectorResults( - params, - distance, - inArray(embedding.knowledgeBaseId, knowledgeBaseIds), - topK - ) - return rows.sort((a, b) => a.distance - b.distance) -} - -/** - * KB runs without a human subject still need bounded small-scope ranking. Probe only chunk - * identities, then reapply every access and visibility predicate before ranking and hydration. - * An overflowing probe selects ANN over the whole scope, never a truncated candidate prefix. - */ -async function selectScopedVectorResults( - params: SearchParams, - distance: SQL, - kbScope: SQL | undefined, - limit: number, - tagConditions: (SQL | undefined)[] = [] -): Promise { - try { - const probe = await runSearchQuery(params.budget, 'vector.probe', (executor) => - executor - .select({ id: embedding.id }) - .from(embedding) - .where(and(kbScope, eq(embedding.enabled, true), ...tagConditions)) - .limit(MAX_EXACT_KB_VECTOR_CANDIDATES + 1) - ) - if (probe.length === 0) return [] - const conditions = [ - kbScope, - ...getVisibilityConditions(params.access, params.filters), - ...tagConditions, - sql`${distance} < ${params.distanceThreshold}`, - ] - if (probe.length <= MAX_EXACT_KB_VECTOR_CANDIDATES) { - annotateSearchDiagnostics({ vectorRanking: 'exact' }) - return await runSearchQuery(params.budget, 'vector.exact', (executor) => - selectRankedVectorResults( - executor, - distance, - [ - ...conditions, - inArray( - embedding.id, - probe.map((candidate) => candidate.id) - ), - ], - limit, - true - ) - ) - } - return await withVectorScanSettings( - (executor) => selectRankedVectorResults(executor, distance, conditions, limit), - params.budget - ) - } catch (error) { - if (!params.budget?.isTimeout(error)) throw error - return [] - } + return selectVectorResults(params) } /** @@ -870,14 +782,27 @@ async function selectScopedVectorResults( * An underfilled index scan expands to a filtered scan within the same statement snapshot. * Live source authorization and content hydration still run after candidate ranking. */ -async function selectLiveVectorResults( - params: SearchParams, - accessProvider: KnowledgeAccessProvider, - distance: SQL, - filters: (SQL | undefined)[] -): Promise { - const conditions = [inArray(embedding.knowledgeBaseId, params.knowledgeBaseIds), ...filters] +async function selectVectorResults(params: SearchParams): Promise { const queryVector = params.queryVector! + const distance = embeddingDistance(queryVector.dimensions, queryVector.vector) + const tagConditions = getStructuredTagFilters(params.structuredFilters ?? [], embedding) + const conditions = [ + inArray(embedding.knowledgeBaseId, params.knowledgeBaseIds), + ...tagConditions, + sql`${distance} < ${params.distanceThreshold!}`, + ] + /** Tags live on chunks; apply them before the candidate limit without fetching full vectors. */ + const candidateTagCondition = tagConditions.length + ? sql`EXISTS ( + SELECT 1 FROM ${embedding} + WHERE ${and(eq(embedding.id, embeddingSearch.id), ...tagConditions)} + )` + : undefined + const accessProvider = params.access.kind === 'user' ? params.accessProvider : undefined + /** Only live-verified readers may defer source authorization until after candidate ranking. */ + const candidateAccess = accessProvider + ? knowledgeMetadataCandidateAccessCondition(params.access) + : knowledgeAccessCondition(params.access) const candidateDistance = embeddingCandidateDistance( queryVector.dimensions, queryVector.vector, @@ -887,8 +812,9 @@ async function selectLiveVectorResults( MAX_VECTOR_RERANK_CANDIDATES, Math.max(MIN_VECTOR_RERANK_CANDIDATES, params.topK * VECTOR_RERANK_OVERSAMPLING) ) - const rows = await selectAuthorizedSearchResults({ + return selectAuthorizedSearchResults({ leg: 'vector', + access: params.access, accessProvider, filters: params.filters, signal: params.signal, @@ -897,23 +823,15 @@ async function selectLiveVectorResults( compareResults: (a, b) => a.distance - b.distance, selectPage: async (limit, offset, excludedSources) => { const visibility = [ - ...getVisibilityConditions( - params.access, - params.filters, - knowledgeMetadataCandidateAccessCondition(params.access) - ), + ...getVisibilityConditions(params.access, params.filters, candidateAccess), excludeSearchSources(excludedSources), ] const candidateDocumentVisibility = [ inArray(document.knowledgeBaseId, params.knowledgeBaseIds), - ...getDocumentVisibilityConditions( - params.access, - params.filters, - knowledgeMetadataCandidateAccessCondition(params.access) - ), + ...getDocumentVisibilityConditions(params.access, params.filters, candidateAccess), excludeSearchSources(excludedSources), ] - /** Explicitly filtered scopes use exact ordering instead of HNSW traversal. */ + /** Exhausted scopes rank exactly; explicit document IDs retain exhaustive passage ordering. */ const exactPage = async (candidateIds?: string[]) => { annotateSearchDiagnostics({ vectorRanking: 'exact' }) const candidates = await runSearchQuery(params.budget, 'vector.exact', (executor) => @@ -934,16 +852,27 @@ async function selectLiveVectorResults( ) return { candidates, nextOffset: offset + candidates.length } } - if (params.filters?.documentIds?.length || params.structuredFilters?.length) { - return exactPage() - } + if (params.filters?.documentIds?.length) return exactPage() /** * Enumerate bounded chunk identities from visible documents. The lateral limit keeps * the probe on document-indexed lookups instead of hashing the entire vector projection. * An exhausted probe fits in the rerank pool and needs only one exact ranking pass. */ const probe = await runSearchQuery(params.budget, 'vector.probe', (executor) => - executor.execute<{ id: string }>(sql` + 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} @@ -988,7 +917,7 @@ async function selectLiveVectorResults( SELECT ${embeddingSearch.id} AS id FROM ${embeddingSearch} CROSS JOIN LATERAL ( SELECT 1 FROM ${document} - WHERE ${and(eq(document.id, embeddingSearch.documentId), ...candidateDocumentVisibility)} + WHERE ${and(eq(document.id, embeddingSearch.documentId), ...candidateDocumentVisibility, candidateTagCondition)} LIMIT 1 ) AS visible WHERE ${and( @@ -1001,7 +930,8 @@ async function selectLiveVectorResults( ${candidateDistance} AS distance FROM ${embeddingSearch} WHERE ${and( inArray(embeddingSearch.knowledgeBaseId, params.knowledgeBaseIds), - eq(embeddingSearch.enabled, true) + eq(embeddingSearch.enabled, true), + candidateTagCondition )} AND (SELECT count(*) FROM initial_candidates) < ${candidateLimit} ), candidates AS ( @@ -1016,8 +946,7 @@ async function selectLiveVectorResults( ) ) SELECT id, (SELECT count(*)::int FROM initial_candidates) AS initial_count FROM candidates `), - params.budget, - 'candidate' + params.budget ) const initialCount = identities[0]?.initial_count ?? 0 annotateSearchDiagnostics({ @@ -1062,37 +991,6 @@ async function selectLiveVectorResults( params.budget ), }) - return rows -} - -/** - * Sort only chunk identities and distances before loading result content. Carrying - * full chunk rows through the vector sort can spill to disk. The bounded subquery - * keeps every visibility predicate before the limit; hydration joins the same - * statement snapshot by primary key, without another distance calculation. - */ -function selectRankedVectorResults( - executor: SearchExecutor, - distance: SQL, - conditions: (SQL | undefined)[], - limit: number, - exact = false -) { - const ranked = executor - .select({ id: embedding.id, distance: distance.as('distance') }) - .from(embedding) - .innerJoin(document, eq(embedding.documentId, document.id)) - .where(and(...conditions)) - .orderBy(exact ? sql`(${distance}) + 0` : distance) - .limit(limit) - .as('ranked_embeddings') - - return executor - .select(getSearchResultFields(ranked.distance)) - .from(ranked) - .innerJoin(embedding, eq(embedding.id, ranked.id)) - .innerJoin(document, eq(document.id, embedding.documentId)) - .orderBy(ranked.distance) } export interface KeywordSearchParams { @@ -1154,6 +1052,7 @@ export async function executeKeywordSearch(params: KeywordSearchParams): Promise /** Keep readable identities and rank scalars separate so sorts never carry full text-search vectors. */ return selectAuthorizedSearchResults({ leg: 'keyword', + access: params.access, accessProvider: params.accessProvider, filters: params.filters, signal: params.signal, @@ -1352,32 +1251,14 @@ export function fuseByReciprocalRank(rankedLists: SearchResult[][], topK: number } export async function handleTagAndVectorSearch(params: SearchParams): Promise { - const { knowledgeBaseIds, topK, structuredFilters, queryVector, distanceThreshold, access } = - params - + const { structuredFilters, queryVector, distanceThreshold } = params if (!structuredFilters || structuredFilters.length === 0) { throw new Error('Tag filters are required for tag and vector search') } if (!queryVector || !distanceThreshold) { throw new Error('Query vector and distance threshold are required for tag and vector search') } - - const tagFilterConditions = getStructuredTagFilters(structuredFilters, embedding) - const distance = embeddingDistance(queryVector.dimensions, queryVector.vector) - if (params.accessProvider && access.kind === 'user') { - return selectLiveVectorResults(params, params.accessProvider, distance, [ - ...tagFilterConditions, - sql`${distance} < ${distanceThreshold}`, - ]) - } - const rows = await selectScopedVectorResults( - params, - distance, - inArray(embedding.knowledgeBaseId, knowledgeBaseIds), - topK, - tagFilterConditions - ) - return rows.sort((a, b) => a.distance - b.distance) + return selectVectorResults(params) } /** From e72dea22e968d5e5f8f625a070bb5aa5ac475a6f Mon Sep 17 00:00:00 2001 From: Vikhyath Mondreti Date: Thu, 17 Sep 2026 14:31:18 -0700 Subject: [PATCH 2/5] fix(knowledge): align search fixtures with shared retrieval --- .../app/api/knowledge/search/utils.test.ts | 32 ++++--- .../search-latency.integration.ts | 83 +++++++++++-------- 2 files changed, 65 insertions(+), 50 deletions(-) diff --git a/apps/sim/app/api/knowledge/search/utils.test.ts b/apps/sim/app/api/knowledge/search/utils.test.ts index fc19016b170..84ae6992e33 100644 --- a/apps/sim/app/api/knowledge/search/utils.test.ts +++ b/apps/sim/app/api/knowledge/search/utils.test.ts @@ -217,6 +217,7 @@ describe('Knowledge Search Utils', () => { Array.from({ length: 201 }, (_, index) => ({ id: `candidate-${index}` })) ) queueTableRows(schemaMock.embedding, [makeResult('second', 0.2), makeResult('first', 0.1)]) + queueTableRows(schemaMock.embedding, [makeResult('second', 0.2), makeResult('first', 0.1)]) const results = await handleTagAndVectorSearch({ knowledgeBaseIds: ['kb-1', 'kb-2'], @@ -231,9 +232,8 @@ describe('Knowledge Search Utils', () => { expect(results.map((row) => row.id)).toEqual(['first', 'second']) expect(dbChainMockFns.select).toHaveBeenCalledTimes(3) - expect(dbChainMockFns.as).toHaveBeenCalledWith('ranked_embeddings') expect(Object.keys(dbChainMockFns.select.mock.calls[0][0])).toEqual(['id']) - expect(dbChainMockFns.limit).toHaveBeenNthCalledWith(1, 201) + expect(dbChainMockFns.limit).toHaveBeenNthCalledWith(1, 400) expect(dbChainMockFns.select.mock.calls[1][0]).toHaveProperty('distance') expect(dbChainMockFns.limit).toHaveBeenCalledWith(2) }) @@ -463,11 +463,7 @@ describe('Knowledge Search Utils', () => { queryVector: JSON.stringify([0.1, 0.2, 0.3]), }) - /** - * A single global LIMIT would let the lexically strongest base consume - * every slot, so an exact-token hit in a smaller base never reaches - * fusion. The vector leg already fans out here; both legs must match. - */ + /** Keyword retrieval preserves its existing per-base lexical candidate selection. */ expect(dbChainMockFns.select).toHaveBeenCalledTimes(knowledgeBaseIds.length) }) @@ -542,6 +538,7 @@ describe('Knowledge Search Utils', () => { }) it('runs a single retrieval leg in vector mode', async () => { + dbChainMockFns.execute.mockResolvedValue([{ id: 'vector-hit' }]) queueTableRows(schemaMock.embedding, [{ id: 'vector-hit' }]) queueTableRows(schemaMock.embedding, [makeResult('vector-hit')]) @@ -551,22 +548,22 @@ describe('Knowledge Search Utils', () => { topK: 10, searchMode: 'vector', query: 'PROJ-1234', - queryVector: JSON.stringify([0.1, 0.2, 0.3]), + queryVector: { vector: JSON.stringify(TEST_EMBEDDING), dimensions: 1536 }, }) expect(results.map((r) => r.id)).toEqual(['vector-hit']) - expect(dbChainMockFns.select).toHaveBeenCalledTimes(3) - expect(dbChainMockFns.as).toHaveBeenCalledWith('ranked_embeddings') + expect(dbChainMockFns.select).toHaveBeenCalledTimes(2) }) it('runs both legs and fuses them in hybrid mode', async () => { /** - * Chains dequeue in creation order: keyword ranking, the budgeted vector - * probe, keyword hydration, then vector ranking and hydration in one query. + * The raw vector probe does not consume a table chain. Keyword ranking and + * hydration complete before vector exact ranking and content hydration. */ + dbChainMockFns.execute.mockResolvedValue([{ id: 'vector-hit' }]) queueTableRows(schemaMock.embedding, [{ id: 'keyword-hit', keywordRank: 0.9 }]) - queueTableRows(schemaMock.embedding, [{ id: 'vector-hit' }]) queueTableRows(schemaMock.embedding, [makeResult('keyword-hit')]) + queueTableRows(schemaMock.embedding, [{ id: 'vector-hit' }]) queueTableRows(schemaMock.embedding, [makeResult('vector-hit')]) const results = await executeKnowledgeSearch({ @@ -575,15 +572,16 @@ describe('Knowledge Search Utils', () => { topK: 10, searchMode: 'hybrid', query: 'PROJ-1234', - queryVector: JSON.stringify([0.1, 0.2, 0.3]), + queryVector: { vector: JSON.stringify(TEST_EMBEDDING), dimensions: 1536 }, }) expect(results.map((r) => r.id).sort()).toEqual(['keyword-hit', 'vector-hit']) - expect(dbChainMockFns.select).toHaveBeenCalledTimes(5) + expect(dbChainMockFns.select).toHaveBeenCalledTimes(4) }) it('propagates unexpected keyword errors after the vector leg finishes', async () => { /** The failing ranking chain is still built first and takes the first queued set. */ + dbChainMockFns.execute.mockResolvedValue([{ id: 'vector-hit' }]) queueTableRows(schemaMock.embedding, [{ id: 'never-ranked', keywordRank: 0 }]) queueTableRows(schemaMock.embedding, [{ id: 'vector-hit' }]) queueTableRows(schemaMock.embedding, [makeResult('vector-hit')]) @@ -600,10 +598,10 @@ describe('Knowledge Search Utils', () => { topK: 10, searchMode: 'hybrid', query: 'PROJ-1234', - queryVector: JSON.stringify([0.1, 0.2, 0.3]), + queryVector: { vector: JSON.stringify(TEST_EMBEDDING), dimensions: 1536 }, }) ).rejects.toBe(failure) - expect(dbChainMockFns.as).toHaveBeenCalledWith('ranked_embeddings') + expect(dbChainMockFns.select).toHaveBeenCalledTimes(3) }) it('skips both query legs when only tag filters are provided', async () => { diff --git a/apps/sim/lib/knowledge/__integration__/search-latency.integration.ts b/apps/sim/lib/knowledge/__integration__/search-latency.integration.ts index b4183e0c4d3..02e80dfad41 100644 --- a/apps/sim/lib/knowledge/__integration__/search-latency.integration.ts +++ b/apps/sim/lib/knowledge/__integration__/search-latency.integration.ts @@ -1,4 +1,4 @@ -/** Real Assistant tool, application authorization, PostgreSQL/pgvector, and result processing. */ +/** Real search adapters, application authorization, PostgreSQL/pgvector, and result processing. */ import { readFileSync, statSync, writeFileSync } from 'node:fs' import type { Principal } from '@sim/auth/principal' import { db } from '@sim/db' @@ -9,7 +9,6 @@ import { document, embedding, knowledgeBase, - knowledgeBaseTagDefinitions, knowledgeConnector, knowledgeConnectorMember, knowledgeDocumentObservation, @@ -103,6 +102,8 @@ function readFixtureReport(file: string) { const reused = reuseFile ? readFixtureReport(reuseFile) : undefined const ids = reused?.fixture ?? createKnowledgeAclFixtureIds() const unrelated = reused?.unrelatedFixture ?? createKnowledgeAclFixtureIds() +const fullWidthFixture = createKnowledgeAclFixtureIds() +const fullWidthChunkCount = 5000 const organizationChatId = generateId() function topicVector(topic = 0) { const vector = Array.from({ length: dimensions }, (_, index) => @@ -124,6 +125,7 @@ const report: Record = { fixtureVersion: 2, chunkCount, unrelatedChunkCount, + fullWidthChunkCount, dimensions, candidateDimensions, chunksPerDocument, @@ -133,7 +135,7 @@ const report: Record = { vectors: 'Normalized 512-dimensional topic/noise geometry with permuted copies across 1536 dimensions; verifies prefix candidate ranking, not semantic embedding quality', cache: evictSharedBuffers - ? 'Organization samples evict PostgreSQL shared buffers before each request; operating-system cache is not cleared' + ? 'Selected workspace and organization samples evict PostgreSQL shared buffers; operating-system cache is not cleared' : 'First and repeated samples; no claim of a cold operating-system cache', layout: reused ? 'Reused fixture; physical layout is inherited from its original report' @@ -383,17 +385,22 @@ async function searchDashboard( async function searchWorkspaceKb( query = 'Orion deployment', - options: { principal?: Principal; tagFilters?: KnowledgeSearchTagFilter[] } = {} + options: { + principal?: Principal + tagFilters?: KnowledgeSearchTagFilter[] + fixture?: typeof ids + } = {} ) { + const fixture = options.fixture ?? ids const result = await searchKnowledge.execute({ principal: options.principal ?? { kind: 'workspace_api_key', - workspaceId: ids.workspaceId, + workspaceId: fixture.workspaceId, keyId: 'fixture-search-key', }, input: { - workspaceId: ids.workspaceId, - knowledgeBaseIds: [ids.knowledgeBaseId], + workspaceId: fixture.workspaceId, + knowledgeBaseIds: [fixture.knowledgeBaseId], query, topK: 15, searchMode: 'vector', @@ -683,6 +690,32 @@ describe.skipIf(!enabled)('Knowledge search latency on a realistic indexed corpu logger.info('Synthetic corpora loaded', { chunkCount, unrelatedChunkCount }) for (const index of indexes) await db.execute(sql.raw(index.indexdef)) } + /** A non-shortenable model must populate its own full-width projection through the write trigger. */ + await seedKnowledgeAclFixture(fullWidthFixture, { connectorType: 'google_drive' }) + await db + .update(knowledgeBase) + .set({ embeddingModel: 'text-embedding-ada-002' }) + .where(eq(knowledgeBase.id, fullWidthFixture.knowledgeBaseId)) + await db.execute(sql` + WITH source AS MATERIALIZED ( + SELECT id, content, embedding FROM embedding + WHERE knowledge_base_id = ${ids.knowledgeBaseId} ORDER BY id LIMIT ${fullWidthChunkCount} + ), documents AS ( + INSERT INTO document + (id, knowledge_base_id, connector_id, external_id, filename, file_url, file_size, + mime_type, processing_status, acl, acl_verified_at) + SELECT ${fullWidthFixture.workspaceId} || '-doc-' || id, ${fullWidthFixture.knowledgeBaseId}, + ${fullWidthFixture.connectorId}, id, 'Full-width deployment guide', + 'https://fixture.invalid/full-width', 12000, 'text/plain', 'completed', + ARRAY['pub']::text[], statement_timestamp() FROM source RETURNING id + ) INSERT INTO embedding + (id, knowledge_base_id, document_id, chunk_index, chunk_hash, content, content_length, + token_count, start_offset, end_offset, embedding) + SELECT ${fullWidthFixture.workspaceId} || '-chunk-' || source.id, + ${fullWidthFixture.knowledgeBaseId}, documents.id, 0, source.id, source.content, + 3000, 750, 0, 3000, source.embedding + FROM source JOIN documents ON documents.id = ${fullWidthFixture.workspaceId} || '-doc-' || source.id + `) await db.execute(sql`ANALYZE document`) await db.execute(sql`ANALYZE embedding`) await db.execute(sql`ANALYZE embedding_search`) @@ -710,9 +743,12 @@ describe.skipIf(!enabled)('Knowledge search latency on a realistic indexed corpu db.$client.options.debug = previousDebug vi.unstubAllGlobals() saveReport() - for (const fixture of process.env.KNOWLEDGE_SEARCH_PERFORMANCE_KEEP_DATABASE === 'true' - ? [] - : [ids, unrelated]) { + for (const fixture of [ + fullWidthFixture, + ...(process.env.KNOWLEDGE_SEARCH_PERFORMANCE_KEEP_DATABASE === 'true' + ? [] + : [ids, unrelated]), + ]) { await db.delete(workspace).where(eq(workspace.id, fixture.workspaceId)) await db.delete(organization).where(eq(organization.id, fixture.organizationId)) await db.delete(user).where(eq(user.id, fixture.aliceId)) @@ -1281,18 +1317,12 @@ describe.skipIf(!enabled)('Knowledge search latency on a realistic indexed corpu report[`${label}.recall`] = { neighbors: expected.length, recall } saveReport() } - await db - .update(knowledgeBase) - .set({ embeddingModel: 'text-embedding-ada-002' }) - .where(eq(knowledgeBase.id, ids.knowledgeBaseId)) - const fullWidth = await sample('workspace-kb.full-width', () => searchWorkspaceKb()) + const fullWidth = await sample('workspace-kb.full-width', () => + searchWorkspaceKb('Orion deployment', { fixture: fullWidthFixture }) + ) expectCompleteVectorSearch(fullWidth.diagnostics) expect(fullWidth.diagnostics.vectorCandidateDimensions).toBe(dimensions) expect(fullWidth.result.data.results).toHaveLength(15) - await db - .update(knowledgeBase) - .set({ embeddingModel: 'text-embedding-3-small' }) - .where(eq(knowledgeBase.id, ids.knowledgeBaseId)) const workflowId = generateId() const scheduled: Principal = { @@ -1346,17 +1376,11 @@ describe.skipIf(!enabled)('Knowledge search latency on a realistic indexed corpu for (const diagnostics of completed) expectCompleteVectorSearch(diagnostics) } - await db.insert(knowledgeBaseTagDefinitions).values({ - id: generateId(), - knowledgeBaseId: ids.knowledgeBaseId, - tagSlot: 'tag1', - displayName: 'Fixture group', - }) await db.execute(sql`UPDATE embedding SET tag1 = 'selected' WHERE knowledge_base_id = ${ids.knowledgeBaseId} AND document_id IN (SELECT id FROM document WHERE knowledge_base_id = ${ids.knowledgeBaseId} AND external_id::int < 600)`) const tagged = await sample('workspace-kb.tagged', () => searchWorkspaceKb('Orion deployment', { - tagFilters: [{ tagName: 'Fixture group', operator: 'eq', value: 'selected' }], + tagFilters: [{ tagName: 'Fixture', operator: 'eq', value: 'selected' }], }) ) expectCompleteVectorSearch(tagged.diagnostics) @@ -1393,16 +1417,9 @@ describe.skipIf(!enabled)('Knowledge search latency on a realistic indexed corpu expectCompleteVectorSearch(denied.diagnostics) expect(denied.result.data.results).toEqual([]) } finally { - await db - .delete(knowledgeBaseTagDefinitions) - .where(eq(knowledgeBaseTagDefinitions.knowledgeBaseId, ids.knowledgeBaseId)) await db.execute( sql`UPDATE embedding SET tag1 = NULL WHERE knowledge_base_id = ${ids.knowledgeBaseId} AND tag1 = 'selected'` ) - await db - .update(knowledgeBase) - .set({ embeddingModel: 'text-embedding-3-small' }) - .where(eq(knowledgeBase.id, ids.knowledgeBaseId)) await db .update(knowledgeConnector) .set({ accessRewritePending: false }) From 7828bdee77cad174b04ffb814e0254225f51f252 Mon Sep 17 00:00:00 2001 From: Vikhyath Mondreti Date: Thu, 17 Sep 2026 14:46:53 -0700 Subject: [PATCH 3/5] fix(knowledge): preserve recall within bounded reranking --- .../search-latency.integration.ts | 86 +++++++++++++------ apps/sim/lib/knowledge/search/queries.ts | 2 +- 2 files changed, 61 insertions(+), 27 deletions(-) diff --git a/apps/sim/lib/knowledge/__integration__/search-latency.integration.ts b/apps/sim/lib/knowledge/__integration__/search-latency.integration.ts index 02e80dfad41..5fe356ae6fe 100644 --- a/apps/sim/lib/knowledge/__integration__/search-latency.integration.ts +++ b/apps/sim/lib/knowledge/__integration__/search-latency.integration.ts @@ -56,6 +56,7 @@ vi.hoisted(() => { if (process.env.KNOWLEDGE_SEARCH_PERFORMANCE_TEST === 'true') { Object.assign(process.env, { OPENAI_API_KEY: 'isolated-embedding-http-fixture', + GEMINI_API_KEY: 'isolated-gemini-http-fixture', CONFLUENCE_CLIENT_ID: 'isolated-confluence-fixture-client', CONFLUENCE_CLIENT_SECRET: 'isolated-confluence-fixture-secret', }) @@ -74,6 +75,7 @@ const unrelatedChunkCount = Number( const evictSharedBuffers = process.env.KNOWLEDGE_SEARCH_PERFORMANCE_EVICT_BUFFERS === 'true' const dimensions = 1536 const candidateDimensions = 512 +const hybridCandidateLimit = 1600 const chunksPerDocument = 4 const logger = createLogger('SearchLatencyIntegration') const fixtureSchema = z.object({ @@ -95,6 +97,7 @@ function readFixtureReport(file: string) { .object({ fixture: fixtureSchema, unrelatedFixture: fixtureSchema, + fullWidthFixture: fixtureSchema.optional(), method: z.object({ fixtureVersion: z.literal(2) }), }) .parse(JSON.parse(readFileSync(file, 'utf8'))) @@ -102,7 +105,7 @@ function readFixtureReport(file: string) { const reused = reuseFile ? readFixtureReport(reuseFile) : undefined const ids = reused?.fixture ?? createKnowledgeAclFixtureIds() const unrelated = reused?.unrelatedFixture ?? createKnowledgeAclFixtureIds() -const fullWidthFixture = createKnowledgeAclFixtureIds() +const fullWidthFixture = reused?.fullWidthFixture ?? createKnowledgeAclFixtureIds() const fullWidthChunkCount = 5000 const organizationChatId = generateId() function topicVector(topic = 0) { @@ -121,6 +124,7 @@ const captured: CapturedQuery[] = [] const report: Record = { fixture: ids, unrelatedFixture: unrelated, + fullWidthFixture, method: { fixtureVersion: 2, chunkCount, @@ -565,13 +569,36 @@ describe.skipIf(!enabled)('Knowledge search latency on a realistic indexed corpu ? new Response(null, { status: 403 }) : Response.json({ type: 'known', accountId: ids.aliceId }) } + if ( + url === + 'https://generativelanguage.googleapis.com/v1beta/models/gemini-embedding-001:batchEmbedContents' + ) { + z.object({ + requests: z + .array( + z.object({ + model: z.literal('models/gemini-embedding-001'), + content: z.object({ + parts: z.array(z.object({ text: z.literal('Orion deployment') })), + }), + outputDimensionality: z.literal(dimensions), + }) + ) + .length(1), + }).parse(JSON.parse(String(init?.body))) + embeddingCalls++ + return Response.json({ + embeddings: [{ values: queryVector }], + usageMetadata: { promptTokenCount: 4 }, + }) + } if (url !== 'https://api.openai.com/v1/embeddings') throw new Error(`Unexpected outbound request in search fixture: ${new URL(url).origin}`) const body = z .object({ input: z.array(z.string()).length(1), encoding_format: z.literal('base64'), - model: z.enum(['text-embedding-3-small', 'text-embedding-ada-002']), + model: z.literal('text-embedding-3-small'), }) .parse(JSON.parse(String(init?.body))) embeddingCalls += body.input.length @@ -691,12 +718,13 @@ describe.skipIf(!enabled)('Knowledge search latency on a realistic indexed corpu for (const index of indexes) await db.execute(sql.raw(index.indexdef)) } /** A non-shortenable model must populate its own full-width projection through the write trigger. */ - await seedKnowledgeAclFixture(fullWidthFixture, { connectorType: 'google_drive' }) - await db - .update(knowledgeBase) - .set({ embeddingModel: 'text-embedding-ada-002' }) - .where(eq(knowledgeBase.id, fullWidthFixture.knowledgeBaseId)) - await db.execute(sql` + if (!reused?.fullWidthFixture) { + await seedKnowledgeAclFixture(fullWidthFixture, { connectorType: 'google_drive' }) + await db + .update(knowledgeBase) + .set({ embeddingModel: 'gemini-embedding-001' }) + .where(eq(knowledgeBase.id, fullWidthFixture.knowledgeBaseId)) + await db.execute(sql` WITH source AS MATERIALIZED ( SELECT id, content, embedding FROM embedding WHERE knowledge_base_id = ${ids.knowledgeBaseId} ORDER BY id LIMIT ${fullWidthChunkCount} @@ -716,6 +744,14 @@ describe.skipIf(!enabled)('Knowledge search latency on a realistic indexed corpu 3000, 750, 0, 3000, source.embedding FROM source JOIN documents ON documents.id = ${fullWidthFixture.workspaceId} || '-doc-' || source.id `) + } + const [fullWidthSize] = await db.execute<{ count: number }>( + sql`SELECT count(*)::int AS count FROM embedding WHERE knowledge_base_id = ${fullWidthFixture.knowledgeBaseId}` + ) + expect(fullWidthSize.count).toBe(fullWidthChunkCount) + await db.execute(sql`UPDATE embedding SET tag1 = 'selected' WHERE knowledge_base_id = ${ids.knowledgeBaseId} + AND tag1 IS DISTINCT FROM 'selected' + AND document_id IN (SELECT id FROM document WHERE knowledge_base_id = ${ids.knowledgeBaseId} AND external_id::int < 600)`) await db.execute(sql`ANALYZE document`) await db.execute(sql`ANALYZE embedding`) await db.execute(sql`ANALYZE embedding_search`) @@ -743,12 +779,9 @@ describe.skipIf(!enabled)('Knowledge search latency on a realistic indexed corpu db.$client.options.debug = previousDebug vi.unstubAllGlobals() saveReport() - for (const fixture of [ - fullWidthFixture, - ...(process.env.KNOWLEDGE_SEARCH_PERFORMANCE_KEEP_DATABASE === 'true' - ? [] - : [ids, unrelated]), - ]) { + for (const fixture of process.env.KNOWLEDGE_SEARCH_PERFORMANCE_KEEP_DATABASE === 'true' + ? [] + : [ids, unrelated, fullWidthFixture]) { await db.delete(workspace).where(eq(workspace.id, fixture.workspaceId)) await db.delete(organization).where(eq(organization.id, fixture.organizationId)) await db.delete(user).where(eq(user.id, fixture.aliceId)) @@ -1089,7 +1122,7 @@ describe.skipIf(!enabled)('Knowledge search latency on a realistic indexed corpu } }, 180_000) - it.each([200, 396, 400, 1000, 2000])( + it.each([200, 1000, 1596, 1600, 2000])( 'keeps a selective scope of %s chunks within both retrieval budgets', async (count) => { const documentCount = count / chunksPerDocument @@ -1116,10 +1149,14 @@ describe.skipIf(!enabled)('Knowledge search latency on a realistic indexed corpu true ) const probe = plans.find((plan) => plan.kind === 'probe')! - expect(probe.plan[0].Plan['Actual Rows']).toBe(Math.min(count, 400)) - expect(assertIndexedChunkProbe(probe.plan[0].Plan)).toBe(Math.min(documentCount, 100)) - expect(plans.filter((plan) => plan.kind === 'vector')).toHaveLength(count < 400 ? 0 : 1) - if (count > 400) { + expect(probe.plan[0].Plan['Actual Rows']).toBe(Math.min(count, hybridCandidateLimit)) + expect(assertIndexedChunkProbe(probe.plan[0].Plan)).toBe( + Math.min(documentCount, hybridCandidateLimit / chunksPerDocument) + ) + expect(plans.filter((plan) => plan.kind === 'vector')).toHaveLength( + count < hybridCandidateLimit ? 0 : 1 + ) + if (count > hybridCandidateLimit) { const rerank = plans.find((plan) => plan.kind === 'rerank')! const actual = await db.$client.unsafe(rerank.query, rerank.parameters).values() const expected = await db.execute<{ id: string }>(sql`SELECT id FROM embedding @@ -1187,8 +1224,10 @@ describe.skipIf(!enabled)('Knowledge search latency on a realistic indexed corpu expectCompleteVectorSearch(broad.diagnostics) expect(broad.result.data.results).toHaveLength(15) const broadProbe = broad.plans.find((plan) => plan.kind === 'probe')! - expect(broadProbe.plan[0].Plan['Actual Rows']).toBe(400) - expect(assertIndexedChunkProbe(broadProbe.plan[0].Plan)).toBe(400 / chunksPerDocument) + expect(broadProbe.plan[0].Plan['Actual Rows']).toBe(hybridCandidateLimit) + expect(assertIndexedChunkProbe(broadProbe.plan[0].Plan)).toBe( + hybridCandidateLimit / chunksPerDocument + ) const { result, plans, diagnostics } = await sample(`member-scope.${surface}`, () => surface === 'copilot' ? search(ids.bobId) : searchDashboard('Orion deployment', ids.bobId) ) @@ -1376,8 +1415,6 @@ describe.skipIf(!enabled)('Knowledge search latency on a realistic indexed corpu for (const diagnostics of completed) expectCompleteVectorSearch(diagnostics) } - await db.execute(sql`UPDATE embedding SET tag1 = 'selected' WHERE knowledge_base_id = ${ids.knowledgeBaseId} - AND document_id IN (SELECT id FROM document WHERE knowledge_base_id = ${ids.knowledgeBaseId} AND external_id::int < 600)`) const tagged = await sample('workspace-kb.tagged', () => searchWorkspaceKb('Orion deployment', { tagFilters: [{ tagName: 'Fixture', operator: 'eq', value: 'selected' }], @@ -1417,9 +1454,6 @@ describe.skipIf(!enabled)('Knowledge search latency on a realistic indexed corpu expectCompleteVectorSearch(denied.diagnostics) expect(denied.result.data.results).toEqual([]) } finally { - await db.execute( - sql`UPDATE embedding SET tag1 = NULL WHERE knowledge_base_id = ${ids.knowledgeBaseId} AND tag1 = 'selected'` - ) await db .update(knowledgeConnector) .set({ accessRewritePending: false }) diff --git a/apps/sim/lib/knowledge/search/queries.ts b/apps/sim/lib/knowledge/search/queries.ts index 78958d8481d..578f71bd02c 100644 --- a/apps/sim/lib/knowledge/search/queries.ts +++ b/apps/sim/lib/knowledge/search/queries.ts @@ -55,7 +55,7 @@ const CANDIDATE_HNSW_EF_SEARCH = '1000' const CANDIDATE_HNSW_SCAN_MEM_MULTIPLIER = '2' const MIN_VECTOR_RERANK_CANDIDATES = 400 const MAX_VECTOR_RERANK_CANDIDATES = 1600 -const VECTOR_RERANK_OVERSAMPLING = 8 +const VECTOR_RERANK_OVERSAMPLING = 32 /** How long to stop trying the iterative-scan settings after the server rejected them. */ const HNSW_SETTINGS_UNSUPPORTED_RETRY_MS = 10 * 60 * 1000 From 454f44c98c10cf19dc978d4b0a2d0131090a39e3 Mon Sep 17 00:00:00 2001 From: Vikhyath Mondreti Date: Thu, 17 Sep 2026 15:09:04 -0700 Subject: [PATCH 4/5] fix(knowledge): validate bounded compact scan plans --- .../search-latency.integration.ts | 51 +++++++++++++------ apps/sim/lib/knowledge/search/queries.ts | 5 +- 2 files changed, 40 insertions(+), 16 deletions(-) diff --git a/apps/sim/lib/knowledge/__integration__/search-latency.integration.ts b/apps/sim/lib/knowledge/__integration__/search-latency.integration.ts index 5fe356ae6fe..a5edda72119 100644 --- a/apps/sim/lib/knowledge/__integration__/search-latency.integration.ts +++ b/apps/sim/lib/knowledge/__integration__/search-latency.integration.ts @@ -75,7 +75,7 @@ const unrelatedChunkCount = Number( const evictSharedBuffers = process.env.KNOWLEDGE_SEARCH_PERFORMANCE_EVICT_BUFFERS === 'true' const dimensions = 1536 const candidateDimensions = 512 -const hybridCandidateLimit = 1600 +const HYBRID_CANDIDATE_LIMIT = 1600 const chunksPerDocument = 4 const logger = createLogger('SearchLatencyIntegration') const fixtureSchema = z.object({ @@ -106,7 +106,7 @@ const reused = reuseFile ? readFixtureReport(reuseFile) : undefined const ids = reused?.fixture ?? createKnowledgeAclFixtureIds() const unrelated = reused?.unrelatedFixture ?? createKnowledgeAclFixtureIds() const fullWidthFixture = reused?.fullWidthFixture ?? createKnowledgeAclFixtureIds() -const fullWidthChunkCount = 5000 +const FULL_WIDTH_CHUNK_COUNT = 5000 const organizationChatId = generateId() function topicVector(topic = 0) { const vector = Array.from({ length: dimensions }, (_, index) => @@ -129,7 +129,7 @@ const report: Record = { fixtureVersion: 2, chunkCount, unrelatedChunkCount, - fullWidthChunkCount, + fullWidthChunkCount: FULL_WIDTH_CHUNK_COUNT, dimensions, candidateDimensions, chunksPerDocument, @@ -162,6 +162,7 @@ interface ExplainNode { 'Node Type': string 'Actual Rows': number 'Actual Loops': number + 'Rows Removed by Filter'?: number 'Plan Rows'?: number 'Shared Hit Blocks'?: number 'Shared Read Blocks'?: number @@ -179,6 +180,7 @@ const explainNodeSchema: z.ZodType = z.lazy(() => 'Node Type': z.string(), 'Actual Rows': z.number(), 'Actual Loops': z.number(), + 'Rows Removed by Filter': z.number().optional(), 'Plan Rows': z.number().optional(), 'Shared Hit Blocks': z.number().optional(), 'Shared Read Blocks': z.number().optional(), @@ -430,7 +432,7 @@ function expectCompleteVectorSearch(diagnostics: z.infer ReturnType, - options: { explain?: boolean } = {} + options: { explain?: boolean; candidateScanRowLimit?: number } = {} ) { captured.length = 0 diagnosticLog?.mockClear() @@ -530,7 +532,24 @@ async function sample( `"embedding_search"."${width === 1536 ? 'vector' : `vector_${width}`}"` ) expect(diagnostics.vectorCandidateLimit).toBeGreaterThan(0) - assertIndexedCandidates(parsedPlan[0].Plan, diagnostics.vectorCandidateLimit!, width) + if (options.candidateScanRowLimit !== undefined) { + /** A small model-specific projection can be cheaper to rank through its KB index. */ + assertCompactCandidates(parsedPlan[0].Plan) + const scans = explainNodes(parsedPlan[0].Plan).filter( + (node) => node['Relation Name'] === 'embedding_search' && node['Actual Loops'] > 0 + ) + expect(scans.length).toBeGreaterThan(0) + for (const node of scans) { + const visited = + (node['Actual Rows'] + (node['Rows Removed by Filter'] ?? 0)) * node['Actual Loops'] + /** EXPLAIN rounds per-worker row averages to integers. */ + expect(visited).toBeLessThanOrEqual( + options.candidateScanRowLimit + node['Actual Loops'] - 1 + ) + } + } else { + assertIndexedCandidates(parsedPlan[0].Plan, diagnostics.vectorCandidateLimit!, width) + } } if (query.query.includes('WITH visible_keyword_documents')) { assertScalarKeywordSorts(parsedPlan[0].Plan) @@ -727,7 +746,7 @@ describe.skipIf(!enabled)('Knowledge search latency on a realistic indexed corpu await db.execute(sql` WITH source AS MATERIALIZED ( SELECT id, content, embedding FROM embedding - WHERE knowledge_base_id = ${ids.knowledgeBaseId} ORDER BY id LIMIT ${fullWidthChunkCount} + WHERE knowledge_base_id = ${ids.knowledgeBaseId} ORDER BY id LIMIT ${FULL_WIDTH_CHUNK_COUNT} ), documents AS ( INSERT INTO document (id, knowledge_base_id, connector_id, external_id, filename, file_url, file_size, @@ -748,7 +767,7 @@ describe.skipIf(!enabled)('Knowledge search latency on a realistic indexed corpu const [fullWidthSize] = await db.execute<{ count: number }>( sql`SELECT count(*)::int AS count FROM embedding WHERE knowledge_base_id = ${fullWidthFixture.knowledgeBaseId}` ) - expect(fullWidthSize.count).toBe(fullWidthChunkCount) + expect(fullWidthSize.count).toBe(FULL_WIDTH_CHUNK_COUNT) await db.execute(sql`UPDATE embedding SET tag1 = 'selected' WHERE knowledge_base_id = ${ids.knowledgeBaseId} AND tag1 IS DISTINCT FROM 'selected' AND document_id IN (SELECT id FROM document WHERE knowledge_base_id = ${ids.knowledgeBaseId} AND external_id::int < 600)`) @@ -1149,14 +1168,14 @@ describe.skipIf(!enabled)('Knowledge search latency on a realistic indexed corpu true ) const probe = plans.find((plan) => plan.kind === 'probe')! - expect(probe.plan[0].Plan['Actual Rows']).toBe(Math.min(count, hybridCandidateLimit)) + expect(probe.plan[0].Plan['Actual Rows']).toBe(Math.min(count, HYBRID_CANDIDATE_LIMIT)) expect(assertIndexedChunkProbe(probe.plan[0].Plan)).toBe( - Math.min(documentCount, hybridCandidateLimit / chunksPerDocument) + Math.min(documentCount, HYBRID_CANDIDATE_LIMIT / chunksPerDocument) ) expect(plans.filter((plan) => plan.kind === 'vector')).toHaveLength( - count < hybridCandidateLimit ? 0 : 1 + count < HYBRID_CANDIDATE_LIMIT ? 0 : 1 ) - if (count > hybridCandidateLimit) { + if (count > HYBRID_CANDIDATE_LIMIT) { const rerank = plans.find((plan) => plan.kind === 'rerank')! const actual = await db.$client.unsafe(rerank.query, rerank.parameters).values() const expected = await db.execute<{ id: string }>(sql`SELECT id FROM embedding @@ -1224,9 +1243,9 @@ describe.skipIf(!enabled)('Knowledge search latency on a realistic indexed corpu expectCompleteVectorSearch(broad.diagnostics) expect(broad.result.data.results).toHaveLength(15) const broadProbe = broad.plans.find((plan) => plan.kind === 'probe')! - expect(broadProbe.plan[0].Plan['Actual Rows']).toBe(hybridCandidateLimit) + expect(broadProbe.plan[0].Plan['Actual Rows']).toBe(HYBRID_CANDIDATE_LIMIT) expect(assertIndexedChunkProbe(broadProbe.plan[0].Plan)).toBe( - hybridCandidateLimit / chunksPerDocument + HYBRID_CANDIDATE_LIMIT / chunksPerDocument ) const { result, plans, diagnostics } = await sample(`member-scope.${surface}`, () => surface === 'copilot' ? search(ids.bobId) : searchDashboard('Orion deployment', ids.bobId) @@ -1356,8 +1375,10 @@ describe.skipIf(!enabled)('Knowledge search latency on a realistic indexed corpu report[`${label}.recall`] = { neighbors: expected.length, recall } saveReport() } - const fullWidth = await sample('workspace-kb.full-width', () => - searchWorkspaceKb('Orion deployment', { fixture: fullWidthFixture }) + const fullWidth = await sample( + 'workspace-kb.full-width', + () => searchWorkspaceKb('Orion deployment', { fixture: fullWidthFixture }), + { candidateScanRowLimit: FULL_WIDTH_CHUNK_COUNT } ) expectCompleteVectorSearch(fullWidth.diagnostics) expect(fullWidth.diagnostics.vectorCandidateDimensions).toBe(dimensions) diff --git a/apps/sim/lib/knowledge/search/queries.ts b/apps/sim/lib/knowledge/search/queries.ts index 578f71bd02c..0c712cf7c63 100644 --- a/apps/sim/lib/knowledge/search/queries.ts +++ b/apps/sim/lib/knowledge/search/queries.ts @@ -49,7 +49,10 @@ const logger = createLogger('KnowledgeSearchQueries') 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. */ +/** + * Stop a permission-starved graph walk early enough to scan the filtered projection instead. + * This approximate iterative-visit threshold excludes pgvector's initial scan; it is not a row limit. + */ const CANDIDATE_HNSW_MAX_SCAN_TUPLES = '1000' const CANDIDATE_HNSW_EF_SEARCH = '1000' const CANDIDATE_HNSW_SCAN_MEM_MULTIPLIER = '2' From e1133d6461942a3757468b5c96f00ca5bb8484c7 Mon Sep 17 00:00:00 2001 From: Vikhyath Mondreti Date: Thu, 17 Sep 2026 15:21:52 -0700 Subject: [PATCH 5/5] fix(knowledge): batch workspace search refills --- .../app/api/knowledge/search/utils.test.ts | 2 +- apps/sim/lib/knowledge/search/queries.test.ts | 36 ++++++++++++++++--- apps/sim/lib/knowledge/search/queries.ts | 5 +-- 3 files changed, 34 insertions(+), 9 deletions(-) diff --git a/apps/sim/app/api/knowledge/search/utils.test.ts b/apps/sim/app/api/knowledge/search/utils.test.ts index 84ae6992e33..7106308b9d5 100644 --- a/apps/sim/app/api/knowledge/search/utils.test.ts +++ b/apps/sim/app/api/knowledge/search/utils.test.ts @@ -235,7 +235,7 @@ describe('Knowledge Search Utils', () => { 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(2) + expect(dbChainMockFns.limit).toHaveBeenCalledWith(20) }) it('should throw error when no filters provided', async () => { diff --git a/apps/sim/lib/knowledge/search/queries.test.ts b/apps/sim/lib/knowledge/search/queries.test.ts index abf33e7dcd7..dc84ba8ffed 100644 --- a/apps/sim/lib/knowledge/search/queries.test.ts +++ b/apps/sim/lib/knowledge/search/queries.test.ts @@ -459,8 +459,10 @@ describe('workspace-scoped vector retrieval', () => { }) it('keeps workspace-authorized sources eligible when hydration needs another page', async () => { - const initial = ranked.map((row) => ({ - ...row, + const initial = Array.from({ length: 20 }, (_, index) => ({ + ...ranked[0], + id: `initial-${index}`, + distance: index / 100, connectorId: 'workspace-source', liveAuthorizationSource: true, })) @@ -473,12 +475,38 @@ describe('workspace-scoped vector retrieval', () => { ...params, filters: { documentIds: ['near-doc', 'far-doc'] }, }) - expect(rows.map((row) => row.id)).toEqual(['near', 'next']) - expect(dbChainMockFns.offset.mock.calls.map(([offset]) => offset)).toEqual([0, 2]) + expect(rows.map((row) => row.id)).toEqual(['initial-0', 'next']) + expect(dbChainMockFns.offset.mock.calls.map(([offset]) => offset)).toEqual([0, 20]) expect(JSON.stringify(dbChainMockFns.where.mock.calls)).not.toContain('workspace-source') expect(getForConnectors).not.toHaveBeenCalled() }) + it('fills a single result from the same candidate page when its nearest row loses access', async () => { + const execute = dbChainMockFns.execute.getMockImplementation()! + dbChainMockFns.execute.mockImplementation(async (query) => { + const statement = render(query) + if (statement.sql.includes('WITH scored_search_candidates')) { + const limit = Number(statement.params.at(-2)) + const offset = Number(statement.params.at(-1)) + const page = ranked.slice(offset, offset + limit) + queueTableRows( + schemaMock.embedding, + page.filter((row) => row.id !== 'near') + ) + return page + } + return execute(query) + }) + + 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(getForConnectors).not.toHaveBeenCalled() + }) + it('does not turn a broad tag filter into exhaustive full-vector ranking', async () => { queueTableRows(schemaMock.embedding, probe) queueTableRows(schemaMock.embedding, ranked) diff --git a/apps/sim/lib/knowledge/search/queries.ts b/apps/sim/lib/knowledge/search/queries.ts index 0c712cf7c63..e2b7dc91c17 100644 --- a/apps/sim/lib/knowledge/search/queries.ts +++ b/apps/sim/lib/knowledge/search/queries.ts @@ -536,10 +536,7 @@ async function selectAuthorizedSearchResults(input: { hydrate: (ids: string[], access: KnowledgeAccessScope) => Promise }): Promise { const deadline = Date.now() + AUTHORIZED_SEARCH_BUDGET_MS - const pageSize = Math.min( - AUTHORIZED_SEARCH_PAGE_SIZE, - input.accessProvider ? Math.max(input.topK, 20) : input.topK - ) + const pageSize = Math.min(AUTHORIZED_SEARCH_PAGE_SIZE, Math.max(input.topK, 20)) const results = new Map() const excludedSources = new Set() const considered = new Set()