diff --git a/apps/sim/app/api/knowledge/search/utils.test.ts b/apps/sim/app/api/knowledge/search/utils.test.ts index 83217cbd6b2..e0eac74f0a6 100644 --- a/apps/sim/app/api/knowledge/search/utils.test.ts +++ b/apps/sim/app/api/knowledge/search/utils.test.ts @@ -216,7 +216,8 @@ describe('Knowledge Search Utils', () => { const statement = (query as { toSQL: () => { sql: string } }).toSQL().sql if (statement.includes('AS visible')) return [] if (statement.includes(') + 0 LIMIT')) return [{ id: 'first' }, { id: 'second' }] - if (statement.includes('WITH scored_search_candidates')) + /** The page reads the pool slice's identities; the walk's order is kept client-side. */ + if (statement.includes('AS "connectorId"') && statement.includes('= ANY(')) return [makeResult('second', 0.2), makeResult('first', 0.1)] return [{ id: 'doc-first' }, { id: 'doc-second' }] }) @@ -240,7 +241,7 @@ describe('Knowledge Search Utils', () => { const exact = dbChainMockFns.execute.mock.calls .map(([query]) => (query as { toSQL: () => { sql: string; params: unknown[] } }).toSQL()) .find((statement) => statement.sql.includes(') + 0 LIMIT'))! - expect(exact.params).toContain(400) + expect(exact.params).toContain(200) }) it('should throw error when no filters provided', async () => { diff --git a/apps/sim/lib/billing/core/usage-gate-cache.ts b/apps/sim/lib/billing/core/usage-gate-cache.ts index 982e5a406d9..00508a9e20d 100644 --- a/apps/sim/lib/billing/core/usage-gate-cache.ts +++ b/apps/sim/lib/billing/core/usage-gate-cache.ts @@ -14,9 +14,11 @@ import { coalesceLocally } from '@/lib/concurrency/singleflight' * every uncached call. Bulk ingestion re-checks per document and knowledge * search checks per query. Staleness is bounded by this TTL and fails in the * harmless direction: a payer who crosses their limit keeps going for at most - * this long, which charges nobody wrongly. + * this long, which charges nobody wrongly. Five minutes: the sum is a few + * hundred milliseconds for a busy payer, and a minute made every search after + * a pause pay it. */ -export const USAGE_GATE_TTL_MS = 60 * 1000 +export const USAGE_GATE_TTL_MS = 5 * 60 * 1000 /** * Recent gate answers, admitted and refused, with `LRUCache` supplying the TTL 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 05abb8e22df..3b61286d0fa 100644 --- a/apps/sim/lib/knowledge/__integration__/kb-block-search.integration.ts +++ b/apps/sim/lib/knowledge/__integration__/kb-block-search.integration.ts @@ -143,7 +143,7 @@ describe('API-key KB block fan-out', () => { expect(matching('hnsw.iterative_scan')).toHaveLength(bases.length) expect(matching('AS visible')).toHaveLength(bases.length) expect(matching(') + 0 LIMIT')).toHaveLength(bases.length) - expect(matching('scored_search_candidates')).toHaveLength(bases.length) + expect(matching('"embedding_search"."id" = ANY(')).toHaveLength(bases.length) /** The probe enumerates visible documents and reports saturation; it never ranks them. */ expect( statements.filter( diff --git a/apps/sim/lib/knowledge/__integration__/search-latency.integration.ts b/apps/sim/lib/knowledge/__integration__/search-latency.integration.ts index 68981209d69..e62aa5081c4 100644 --- a/apps/sim/lib/knowledge/__integration__/search-latency.integration.ts +++ b/apps/sim/lib/knowledge/__integration__/search-latency.integration.ts @@ -8,6 +8,7 @@ import { credentialGroup, document, embedding, + embeddingSearch, knowledgeBase, knowledgeConnector, knowledgeConnectorMember, @@ -19,7 +20,7 @@ import { } from '@sim/db/schema' import { createLogger, Logger } from '@sim/logger' import { generateId } from '@sim/utils/id' -import { and, eq, inArray, sql } from 'drizzle-orm' +import { and, eq, inArray, type SQL, sql } from 'drizzle-orm' import { NextRequest } from 'next/server' import { afterAll, beforeAll, describe, expect, it, type MockInstance, vi } from 'vitest' import { z } from 'zod' @@ -41,6 +42,7 @@ import { seedKnowledgeMemberFixture, } from '@/lib/knowledge/__integration__/seed-source-access-fixture' import { type KnowledgeSearchTagFilter, searchKnowledge } from '@/lib/knowledge/application/search' +import type { KbEmbeddingDimensions } from '@/lib/knowledge/embedding-models' import { SearchBudget, SearchDeadlineError, @@ -48,6 +50,7 @@ import { } from '@/lib/knowledge/search/budget' import type { SearchStage } from '@/lib/knowledge/search/diagnostics' import type { WorkspaceSearchFilters } from '@/lib/knowledge/search/filters' +import { embeddingCandidateDistance } from '@/lib/knowledge/vector-columns' import { POST as searchRoute } from '@/app/api/knowledge/search/route' import { ResolvedSecretTraceRegistry } from '@/executor/utils/resolved-secret-trace-registry' @@ -120,6 +123,25 @@ function topicVector(topic = 0) { return vector.map((value) => value / magnitude) } const queryVector = topicVector() + +/** + * The exact nearest chunks on the projection's stored halfvec, which is what the page's order + * is measured against: the walk ranks on that column, and nothing rescores it. + */ +async function exactProjectionNeighbors(vector: number[], limit: number, readerClause?: SQL) { + const distance = embeddingCandidateDistance( + dimensions as KbEmbeddingDimensions, + JSON.stringify(vector), + 'text-embedding-3-small' + ) + /** Unaliased: the distance expression qualifies its column with the table's own name. */ + return db.execute<{ id: string }>(sql`SELECT ${embeddingSearch.id} AS id FROM ${embeddingSearch} + INNER JOIN document d ON d.id = ${embeddingSearch.documentId} + WHERE ${embeddingSearch.knowledgeBaseId} = ${ids.knowledgeBaseId} + AND ${embeddingSearch.enabled} ${readerClause ?? sql``} + ORDER BY (${distance}) + 0, ${embeddingSearch.id} + LIMIT ${limit}`) +} const captured: CapturedQuery[] = [] const report: Record = { fixture: ids, @@ -213,6 +235,10 @@ function explainNodes(node: ExplainNode): ExplainNode[] { * aliases its own lateral `scoped_chunk`, so this cannot match it, and matching on the rendered * clause casing would silently stop these assertions from running at all. */ +/** The vector page reads a pool slice's identities from the projection and its documents. */ +const VECTOR_PAGE_JOIN = + 'INNER JOIN "document" ON "document"."id" = "embedding_search"."document_id"' + function isVectorCandidateQuery(statement: string) { return statement.toLowerCase().includes(') as visible') } @@ -476,12 +502,12 @@ async function sample( item.query.includes('limit') || item.query.includes('CROSS JOIN LATERAL') || isVectorCandidateQuery(item.query) || - item.query.includes('WITH scored_search_candidates') || + item.query.includes(VECTOR_PAGE_JOIN) || item.query.includes('WITH matched_keyword_chunks')) ) const plans: Array< CapturedQuery & { - kind: 'keyword' | 'vector' | 'rerank' | 'probe' + kind: 'keyword' | 'vector' | 'page' | 'probe' plan: z.infer } > = [] @@ -516,9 +542,8 @@ async function sample( ? 'keyword' : isVectorCandidateQuery(query.query) ? 'vector' - : query.query.includes('order by') || - query.query.includes('WITH scored_search_candidates') - ? 'rerank' + : query.query.includes('order by') || query.query.includes(VECTOR_PAGE_JOIN) + ? 'page' : 'probe', query: query.query, parameters: query.parameters, @@ -1004,13 +1029,10 @@ describe.skipIf(!enabled)('Knowledge search latency on a realistic indexed corpu expect(vectorPlans).toHaveLength(1) expect(vectorPlans[0].plan[0].Plan['Actual Rows']).toBeGreaterThan(0) assertCompactCandidates(vectorPlans[0].plan[0].Plan) - expect(plans.some((plan) => plan.kind === 'rerank')).toBe(true) - 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 - WHERE knowledge_base_id = ${ids.knowledgeBaseId} AND enabled - ORDER BY (embedding <=> ${JSON.stringify(queryVector)}::vector) + 0, id - LIMIT ${actual.length}`) + expect(plans.some((plan) => plan.kind === 'page')).toBe(true) + const page = plans.find((plan) => plan.kind === 'page')! + const actual = await db.$client.unsafe(page.query, page.parameters).values() + const expected = await exactProjectionNeighbors(queryVector, actual.length) const expectedIds = new Set(expected.map(({ id }) => id)) const recall = actual.filter(([id]) => expectedIds.has(id)).length / expected.length expect(recall).toBeGreaterThanOrEqual(0.95) @@ -1028,12 +1050,9 @@ describe.skipIf(!enabled)('Knowledge search latency on a realistic indexed corpu expectCompleteVectorSearch(diagnostics) const candidates = plans.find((plan) => plan.kind === 'vector')! assertCompactCandidates(candidates.plan[0].Plan) - 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 - WHERE knowledge_base_id = ${ids.knowledgeBaseId} AND enabled - ORDER BY (embedding <=> ${JSON.stringify(topicVector(topic))}::vector) + 0, id - LIMIT ${actual.length}`) + const page = plans.find((plan) => plan.kind === 'page')! + const actual = await db.$client.unsafe(page.query, page.parameters).values() + const expected = await exactProjectionNeighbors(topicVector(topic), actual.length) const expectedIds = new Set(expected.map(({ id }) => id)) const recall = actual.filter(([id]) => expectedIds.has(id)).length / expected.length expect(recall).toBeGreaterThanOrEqual(0.95) @@ -1082,15 +1101,14 @@ describe.skipIf(!enabled)('Knowledge search latency on a realistic indexed corpu ) expectCompleteVectorSearch(diagnostics) expect(result.data.results).toHaveLength(15) - const rerank = plans.find((plan) => plan.kind === 'rerank')! - expect(rerank).toBeDefined() - const actual = await db.$client.unsafe(rerank.query, rerank.parameters).values() - const expected = await db.execute<{ id: string }>(sql`SELECT e.id FROM embedding e - INNER JOIN document d ON d.id = e.document_id - WHERE e.knowledge_base_id = ${ids.knowledgeBaseId} AND e.enabled - AND d.acl @> ARRAY[${reader}]::text[] - ORDER BY (e.embedding <=> ${JSON.stringify(queryVector)}::vector) + 0, e.id - LIMIT ${actual.length}`) + const page = plans.find((plan) => plan.kind === 'page')! + expect(page).toBeDefined() + const actual = await db.$client.unsafe(page.query, page.parameters).values() + const expected = await exactProjectionNeighbors( + queryVector, + actual.length, + sql`AND d.acl @> ARRAY[${reader}]::text[]` + ) expect(expected.length).toBeGreaterThan(0) const expectedIds = new Set(expected.map(({ id }) => id)) const recall = actual.filter(([id]) => expectedIds.has(id)).length / expected.length @@ -1128,9 +1146,10 @@ describe.skipIf(!enabled)('Knowledge search latency on a realistic indexed corpu expect(probe[0].query).not.toContain('<=>') expect(probe[0].plan[0].Plan['Actual Rows']).toBe(12) expect(assertIndexedChunkProbe(probe[0].plan[0].Plan)).toBe(documentIds.length) - const vector = plans.filter((plan) => plan.kind === 'rerank') - expect(vector).toHaveLength(1) - expect(vector[0].query).toContain('"embedding"."id" in') + /** The page reads the bounded ranking's identities from the projection, never the original vectors. */ + const page = plans.filter((plan) => plan.kind === 'page') + expect(page).toHaveLength(1) + expect(page[0].query).not.toContain('"embedding"."embedding"') } } finally { await db @@ -1175,16 +1194,16 @@ describe.skipIf(!enabled)('Knowledge search latency on a realistic indexed corpu count < HYBRID_CANDIDATE_LIMIT ? 0 : 1 ) 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 - WHERE knowledge_base_id = ${ids.knowledgeBaseId} AND enabled - AND document_id IN (${sql.join( - documentIds.map((id) => sql`${id}`), - sql`, ` - )}) - ORDER BY (embedding <=> ${JSON.stringify(queryVector)}::vector) + 0, id - LIMIT ${actual.length}`) + const page = plans.find((plan) => plan.kind === 'page')! + const actual = await db.$client.unsafe(page.query, page.parameters).values() + const expected = await exactProjectionNeighbors( + queryVector, + actual.length, + sql`AND ${embeddingSearch.documentId} IN (${sql.join( + documentIds.map((id) => sql`${id}`), + sql`, ` + )})` + ) const expectedIds = new Set(expected.map(({ id }) => id)) const recall = actual.filter(([id]) => expectedIds.has(id)).length / expected.length expect(recall).toBeGreaterThanOrEqual(0.95) @@ -1352,7 +1371,7 @@ describe.skipIf(!enabled)('Knowledge search latency on a realistic indexed corpu ) expectCompleteVectorSearch(diagnostics) expect(diagnostics.accessScopeKind).toBe('workspace') - expect(diagnostics.vectorRanking).toBe('candidate-rerank') + expect(diagnostics.vectorRanking).toBe('projection-walk') expect(result.data.results).toHaveLength(15) expect(plans.some((plan) => plan.kind === 'vector')).toBe(true) for (const row of result.data.results) { @@ -1441,7 +1460,7 @@ describe.skipIf(!enabled)('Knowledge search latency on a realistic indexed corpu }) ) expectCompleteVectorSearch(tagged.diagnostics) - expect(tagged.diagnostics.vectorRanking).toBe('candidate-rerank') + expect(tagged.diagnostics.vectorRanking).toBe('projection-walk') expect(tagged.result.data.results).toHaveLength(15) for (const row of tagged.result.data.results) { const ordinal = Number(row.documentId.split('-doc-')[1]) diff --git a/apps/sim/lib/knowledge/access/availability.test.ts b/apps/sim/lib/knowledge/access/availability.test.ts index 576966fdb67..f4eaa5dc090 100644 --- a/apps/sim/lib/knowledge/access/availability.test.ts +++ b/apps/sim/lib/knowledge/access/availability.test.ts @@ -29,6 +29,7 @@ vi.mock('@/lib/credential-groups/scoped-availability', () => ({ })) import { + forgetKnowledgeAccessAvailability, requireOrganizationSearchAvailable, resolveKnowledgeAccessAvailability, } from '@/lib/knowledge/access/availability' @@ -36,6 +37,7 @@ import { describe('knowledge access availability ownership', () => { beforeEach(() => { vi.clearAllMocks() + forgetKnowledgeAccessAvailability() mocks.featureEnabled.mockResolvedValue(true) mocks.enterprise.mockResolvedValue(true) mocks.scopedGroups.mockResolvedValue(true) @@ -59,6 +61,17 @@ describe('knowledge access availability ownership', () => { expect(mocks.workspaceGroups).not.toHaveBeenCalled() }) + it('answers the same owner from one read for a minute', async () => { + await resolveKnowledgeAccessAvailability({ organizationId: 'org-1' }) + await resolveKnowledgeAccessAvailability({ organizationId: 'org-1' }) + expect(mocks.enterprise).toHaveBeenCalledTimes(1) + await resolveKnowledgeAccessAvailability({ organizationId: 'org-2' }) + expect(mocks.enterprise).toHaveBeenCalledTimes(2) + forgetKnowledgeAccessAvailability() + await resolveKnowledgeAccessAvailability({ organizationId: 'org-1' }) + expect(mocks.enterprise).toHaveBeenCalledTimes(3) + }) + it('keeps source mirroring independent from managed identity availability', async () => { mocks.scopedGroups.mockResolvedValue(false) await expect(resolveKnowledgeAccessAvailability({ organizationId: 'org-1' })).resolves.toEqual({ diff --git a/apps/sim/lib/knowledge/access/availability.ts b/apps/sim/lib/knowledge/access/availability.ts index 31c9d9722b6..6a99f6873c6 100644 --- a/apps/sim/lib/knowledge/access/availability.ts +++ b/apps/sim/lib/knowledge/access/availability.ts @@ -1,3 +1,4 @@ +import { LRUCache } from 'lru-cache' import { isOrganizationOnEnterprisePlan } from '@/lib/billing/core/subscription' import { getWorkspaceOwnerSubscriptionAccess, @@ -39,11 +40,44 @@ export interface KnowledgeAccessAvailability { memberScoped: boolean } +/** + * How long a resolved availability holds. A search resolves it three times over — the gate, the + * defaults, the reader's scope — each a subscription and a billing read; one read per owner per + * minute answers all of them, and a plan or flag change lands within the minute. + */ +const AVAILABILITY_TTL_MS = 60 * 1000 + +const availabilityCache = new LRUCache< + string, + KnowledgeAccessAvailability, + KnowledgeMemberAccessContext +>({ + max: 10_000, + ttl: AVAILABILITY_TTL_MS, + fetchMethod: (_key, _stale, { context }) => readKnowledgeAccessAvailability(context), +}) + export async function resolveKnowledgeAccessAvailability( context: KnowledgeMemberAccessContext ): Promise { if (context.organizationId && context.workspaceId) throw new Error('Knowledge access requires one resource owner') + /** A caller that brings its own billing snapshot is answered from that snapshot, uncached. */ + if (context.ownerBilling) return readKnowledgeAccessAvailability(context) + const key = `${context.organizationId ?? ''}|${context.workspaceId ?? ''}|${context.userId ?? ''}` + const availability = await availabilityCache.fetch(key, { context }) + if (!availability) throw new Error('Knowledge access availability could not be resolved') + return availability +} + +/** Forgets every resolved availability, for tests and for a settings change that must land now. */ +export function forgetKnowledgeAccessAvailability(): void { + availabilityCache.clear() +} + +async function readKnowledgeAccessAvailability( + context: KnowledgeMemberAccessContext +): Promise { if ( !(await isFeatureEnabled( 'knowledge-member-access', diff --git a/apps/sim/lib/knowledge/access/predicate.ts b/apps/sim/lib/knowledge/access/predicate.ts index 28ad3f9f49c..b10c6276432 100644 --- a/apps/sim/lib/knowledge/access/predicate.ts +++ b/apps/sim/lib/knowledge/access/predicate.ts @@ -357,7 +357,15 @@ export function projectionCandidateAccessCondition( documentId: AnyPgColumn | SQL }, scope: KnowledgeAccessScope | SystemAccessScope, - plan: SearchAccessPlan + plan: SearchAccessPlan, + options: { + /** + * Whether every row of the projection carries its mirrored source and ACL. While the fill + * is under way, a row it has not reached is decided on its document; once it is complete no + * such row exists, and the predicate is the array test alone. + */ + filled?: boolean + } = {} ): SQL { if (scope.kind === 'system') return sql`true` if (scope.tokens.length === 0) return sql`false` @@ -374,12 +382,14 @@ export function projectionCandidateAccessCondition( const owned = plan.uploads ? sql`(${projection.connectorId} IS NULL OR ${inSources(mirrored)})` : inSources(mirrored) + const onRow = sql`(${projection.acl} && ${tokens} AND ${owned})` + if (options.filled) return onRow const unfilled = sql`(${projection.acl} IS NULL AND EXISTS ( SELECT 1 FROM ${document} WHERE ${document.id} = ${projection.documentId} AND ${knowledgeCandidateAccessConditionForConnectors(scope, plan)} ))` - return sql`(${unfilled} OR (${projection.acl} && ${tokens} AND ${owned}))` + return sql`(${unfilled} OR ${onRow})` } /** diff --git a/apps/sim/lib/knowledge/application/search.ts b/apps/sim/lib/knowledge/application/search.ts index 1d355381606..6530c6bb64a 100644 --- a/apps/sim/lib/knowledge/application/search.ts +++ b/apps/sim/lib/knowledge/application/search.ts @@ -267,10 +267,6 @@ const searchKnowledgeUseCase = defineAuthorizedKnowledgeUseCase({ knowledgeBaseCount: context.knowledgeBases.length, }) input.signal?.throwIfAborted() - if (context.organizationId) - await measureSearchStage('availability', () => - requireOrganizationSearchAvailable(context.organizationId!) - ) const requestId = generateRequestId() const hasQuery = Boolean(input.query?.trim()) const filters = input.tagFilters ?? [] @@ -286,24 +282,37 @@ const searchKnowledgeUseCase = defineAuthorizedKnowledgeUseCase({ principal.kind === 'delegated' && principal.serviceId === 'executor' ) - const billingAttribution = hasQuery - ? input.resolveBillingAttribution && context.workspaceId - ? await measureSearchStage('billing_attribution', () => - input.resolveBillingAttribution!(context.workspaceId!) - ) - : await measureSearchStage('billing_attribution', () => - resolveKnowledgeBillingAttribution(principal, context) - ) - : undefined - if (shouldMeter && billingAttribution) { - const usage = await measureSearchStage('usage_admission', () => - checkSearchUsageLimits(billingAttribution) - ) - if (usage.isExceeded) { - throw new KnowledgeUsageLimitExceededError( - usage.message || 'Usage limit exceeded. Please upgrade your plan to continue.' + /** + * Whether the organization may search at all, and whether this payer still may: neither + * depends on the query, so both run beside the scope and defaults reads below instead of + * ahead of them. Admission stays ahead of the embedding call, which a refused search must + * never make. + */ + const admit = async (): Promise => { + if (context.organizationId) + await measureSearchStage('availability', () => + requireOrganizationSearchAvailable(context.organizationId!) ) + const billingAttribution = hasQuery + ? input.resolveBillingAttribution && context.workspaceId + ? await measureSearchStage('billing_attribution', () => + input.resolveBillingAttribution!(context.workspaceId!) + ) + : await measureSearchStage('billing_attribution', () => + resolveKnowledgeBillingAttribution(principal, context) + ) + : undefined + if (shouldMeter && billingAttribution) { + const usage = await measureSearchStage('usage_admission', () => + checkSearchUsageLimits(billingAttribution) + ) + if (usage.isExceeded) { + throw new KnowledgeUsageLimitExceededError( + usage.message || 'Usage limit exceeded. Please upgrade your plan to continue.' + ) + } } + return billingAttribution } const knowledgeBaseIds = context.knowledgeBases.map((knowledgeBase) => knowledgeBase.id) @@ -359,19 +368,7 @@ const searchKnowledgeUseCase = defineAuthorizedKnowledgeUseCase({ : undefined const resultSecretRegistry = preparedRegistry ?? input.resultSecretRegistry input.signal?.throwIfAborted() - const [queryEmbedding, access, searchDefaults] = await Promise.all([ - hasQuery - ? measureSearchStage('embedding', () => - runWithKnowledgeModelInputProvenance(resultSecretRegistry, () => - generateSearchEmbedding( - input.query!, - embeddingTarget!, - context.workspaceId, - input.signal - ) - ) - ) - : Promise.resolve(null), + const [access, searchDefaults, billingAttribution, tagDefinitions] = await Promise.all([ measureSearchStage('access_scope', () => context.access.get()), measureSearchStage('defaults', () => resolveKnowledgeSearchDefaults({ @@ -383,7 +380,28 @@ const searchKnowledgeUseCase = defineAuthorizedKnowledgeUseCase({ requestedMode: input.searchMode, }) ), + admit(), + /** The tag names the results are labelled with depend on the bases alone. */ + filters.length === 0 + ? measureSearchStage('tag_definitions', () => + getDocumentTagDefinitionsByKnowledgeBaseIds(knowledgeBaseIds) + ) + : Promise.resolve(definitionsByKnowledgeBase), ]) + definitionsByKnowledgeBase = tagDefinitions + input.signal?.throwIfAborted() + const queryEmbedding = hasQuery + ? await measureSearchStage('embedding', () => + runWithKnowledgeModelInputProvenance(resultSecretRegistry, () => + generateSearchEmbedding( + input.query!, + embeddingTarget!, + context.workspaceId, + input.signal + ) + ) + ) + : null input.signal?.throwIfAborted() annotateSearchDiagnostics({ accessScopeKind: access.kind, @@ -609,11 +627,6 @@ const searchKnowledgeUseCase = defineAuthorizedKnowledgeUseCase({ } } - if (filters.length === 0) { - definitionsByKnowledgeBase = await measureSearchStage('tag_definitions', () => - getDocumentTagDefinitionsByKnowledgeBaseIds(knowledgeBaseIds) - ) - } const tagMaps = new Map( [...definitionsByKnowledgeBase].map(([knowledgeBaseId, definitions]) => [ knowledgeBaseId, diff --git a/apps/sim/lib/knowledge/search/diagnostics.ts b/apps/sim/lib/knowledge/search/diagnostics.ts index 43d27b0a443..e8c753941b6 100644 --- a/apps/sim/lib/knowledge/search/diagnostics.ts +++ b/apps/sim/lib/knowledge/search/diagnostics.ts @@ -50,7 +50,9 @@ export type SearchStage = | `${RetrievalLeg}.sql` | 'vector.settings' | 'vector.probe' - | 'vector.rerank' + | 'vector.page' + | 'vector.projection_filled' + | 'keyword.projection_filled' | 'vector.exact_candidates' | 'vector.exact' | 'vector.candidate_search' @@ -84,7 +86,7 @@ export interface SearchDiagnosticMetadata { searchMode?: 'hybrid' | 'vector' boostRecency?: boolean embeddingDimensions?: number - vectorRanking?: 'exact' | 'exact-candidates' | 'candidate-rerank' | 'per-source' + vectorRanking?: 'exact' | 'exact-candidates' | 'projection-walk' | 'per-source' vectorCandidateStorage?: 'stored-halfvec' /** * Whether the bounded traversal filled its candidate limit. `underfilled` means visibility diff --git a/apps/sim/lib/knowledge/search/queries.test.ts b/apps/sim/lib/knowledge/search/queries.test.ts index fbb5c368c68..134ad41fb9e 100644 --- a/apps/sim/lib/knowledge/search/queries.test.ts +++ b/apps/sim/lib/knowledge/search/queries.test.ts @@ -47,6 +47,7 @@ import { retrieveKnowledgeSearch, type SearchParams, VECTOR_PROBE_DOCUMENT_LIMIT, + vectorCandidatePoolLimit, visibleDocumentsQuery, } from '@/lib/knowledge/search/queries' import { forgetIndexedVectorSources } from '@/lib/knowledge/search/source-vector-indexes' @@ -142,6 +143,13 @@ function isExactRanking(sql: string) { return sql.includes(') + 0 LIMIT') } +/** The page read: a slice of the pool's identities, from the projection and its documents. */ +function isPageStatement(sql: string) { + return ( + sql.includes('AS "connectorId"') && sql.includes('= ANY(') && !sql.includes('ranked_tin_chunks') + ) +} + const statements = () => dbChainMockFns.execute.mock.calls.map(([query]) => render(query)) function renderOne(filters: StructuredFilter[]) { @@ -413,7 +421,7 @@ describe('workspace-scoped vector retrieval', () => { if (failCandidates) throw failCandidates return traversedRows } - if (statement.includes('WITH scored_search_candidates')) return ranked + if (isPageStatement(statement)) return ranked if (isExactRanking(statement)) return exactRows if (isProbeStatement(statement)) return probeRows return [] @@ -512,9 +520,7 @@ describe('workspace-scoped vector retrieval', () => { probeRows = [] expect(await handleVectorOnlySearch(params)).toEqual([]) expect(statements().filter((query) => isExactRanking(query.sql))).toHaveLength(0) - expect( - statements().filter((query) => query.sql.includes('WITH scored_search_candidates')) - ).toHaveLength(0) + expect(statements().filter((query) => isPageStatement(query.sql))).toHaveLength(0) }) it('spends only its own share of the leg on a probe that runs long', async () => { @@ -539,6 +545,42 @@ describe('workspace-scoped vector retrieval', () => { expect(statements().filter((query) => isExactRanking(query.sql))).toHaveLength(0) }) + it('walks for a pool sized to the page, and scores results on the projection', async () => { + queueTableRows(schemaMock.embedding, [...ranked].reverse()) + await handleVectorOnlySearch(params) + const walk = statements().find((query) => isWalk(query.sql))! + /** The page is the walk's order, so the walk ends at a page's worth of candidates, not a rerank's. */ + expect(walk.params).toContain(200) + expect(walk.params).not.toContain(1600) + /** Hydration scores each result on the stored halfvec; the original vector is never read. */ + const fields = JSON.stringify(dbChainMockFns.select.mock.calls[0][0]) + expect(fields).toContain(String(schemaMock.embeddingSearch.vector512)) + expect(fields).not.toContain(String(schemaMock.embedding.embedding)) + expect(JSON.stringify(dbChainMockFns.leftJoin.mock.calls)).toContain('embeddingSearch') + }) + + it('passes over a slice whose documents went away instead of ending the pool there', async () => { + const execute = dbChainMockFns.execute.getMockImplementation()! + let pages = 0 + dbChainMockFns.execute.mockImplementation(async (query) => { + const statement = render(query) + /** The first slice's documents are gone; the next slice still has the readable rows. */ + if (isPageStatement(statement.sql)) return pages++ === 0 ? [] : ranked + return execute(query) + }) + queueTableRows(schemaMock.embedding, [...ranked].reverse()) + expect((await handleVectorOnlySearch(params)).map((row) => row.id)).toEqual(['near', 'far']) + expect(pages).toBe(2) + expect(statements().filter((query) => isWalk(query.sql))).toHaveLength(1) + }) + + it('sizes the pool to the pages asked for, doubling a pool the pages outran', () => { + expect(vectorCandidatePoolLimit(20, undefined)).toBe(200) + expect(vectorCandidatePoolLimit(150, undefined)).toBe(300) + expect(vectorCandidatePoolLimit(210, 200)).toBe(420) + expect(vectorCandidatePoolLimit(5000, 1600)).toBe(1600) + }) + 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']) @@ -554,13 +596,12 @@ describe('workspace-scoped vector retrieval', () => { 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') + /** The page is the ranking's own order; nothing rescores it against the original vectors. */ + const page = statements().find((query) => isPageStatement(query.sql))! + expect(page.sql).not.toContain('MATERIALIZED') + expect(JSON.stringify(page)).not.toContain(String(schemaMock.embedding.embedding)) + expect(JSON.stringify(page)).toContain('candidate-0') + expect(JSON.stringify(page)).not.toContain('probe-') expect(getForConnectors).not.toHaveBeenCalled() }) @@ -590,15 +631,12 @@ describe('workspace-scoped vector retrieval', () => { 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) + if (isPageStatement(statement.sql)) { queueTableRows( schemaMock.embedding, - page.filter((row) => row.id !== 'near') + ranked.filter((row) => row.id !== 'near') ) - return page + return ranked } return execute(query) }) @@ -640,7 +678,7 @@ describe('workspace-scoped vector retrieval', () => { expect(dbChainMockFns.transaction).toHaveBeenCalledOnce() }) - it.each(['vector.candidate_search', 'vector.rerank', 'vector.sql'] as const)( + it.each(['vector.candidate_search', 'vector.page', 'vector.sql'] as const)( 'reports a %s timeout as partial, not a complete empty search', async (failedStage) => { const query = SearchBudget.prototype.query @@ -673,7 +711,7 @@ describe('workspace-scoped vector retrieval', () => { ) { const result = await (query.bind(this) as SearchBudget['query'])(stage, run) if (stage === 'vector.candidate_search') vi.spyOn(performance, 'now').mockReturnValue(60) - if (stage === 'vector.rerank') vi.spyOn(performance, 'now').mockReturnValue(80) + if (stage === 'vector.page') vi.spyOn(performance, 'now').mockReturnValue(80) return result }) queueTableRows(schemaMock.embedding, ranked) @@ -904,7 +942,7 @@ describe('hydration follows ranked candidates', () => { dbChainMockFns.execute.mockImplementation(async (query) => { const statement = render(query).sql if (statement.includes('AS visible')) return candidatePages.shift() ?? [] - if (statement.includes('WITH scored_search_candidates')) return rerankPages.shift() ?? [] + if (isPageStatement(statement)) return rerankPages.shift() ?? [] if (statement.includes('WITH matched_keyword_chunks')) return keywordPages.shift() ?? [] if (isExactRanking(statement)) return exactPages.shift() ?? [] if (isProbeStatement(statement)) return probePages.shift() ?? [] @@ -939,10 +977,10 @@ describe('hydration follows ranked candidates', () => { expect(render(candidateQuery).sql).toContain('LIMIT 1') expect(JSON.stringify(candidateQuery)).toContain('required_clause') expect(JSON.stringify(candidateQuery)).toContain('subvector') - const rankQuery = dbChainMockFns.execute.mock.calls.find(([query]) => - render(query).sql.includes('WITH scored_search_candidates') + const pageQuery = dbChainMockFns.execute.mock.calls.find(([query]) => + isPageStatement(render(query).sql) )![0] - expect(render(rankQuery).sql).toContain('MATERIALIZED') + expect(JSON.stringify(pageQuery)).not.toContain(String(schemaMock.embedding.embedding)) }) it('finishes a scope the probe finds nothing in without ranking it', async () => { @@ -952,9 +990,7 @@ describe('hydration follows ranked candidates', () => { expect(probe.params).toContain(VECTOR_PROBE_DOCUMENT_LIMIT + 1) expect(probe.sql).not.toContain('<=>') expect(JSON.stringify(probe)).toContain('required_clause') - expect( - statements().filter((query) => query.sql.includes('WITH scored_search_candidates')) - ).toHaveLength(0) + expect(statements().filter((query) => isPageStatement(query.sql))).toHaveLength(0) }) it('ranks the permitted set exactly when the traversal comes back underfilled', async () => { @@ -1028,9 +1064,7 @@ describe('hydration follows ranked candidates', () => { const rows = await handleVectorOnlySearch({ ...params, structuredFilters: undefined }) expect(rows.map((row) => row.id)).toEqual(['selected']) expect( - dbChainMockFns.execute.mock.calls.filter(([query]) => - render(query).sql.includes('WITH scored_search_candidates') - ) + dbChainMockFns.execute.mock.calls.filter(([query]) => isPageStatement(render(query).sql)) ).toHaveLength(2) }) @@ -1065,9 +1099,7 @@ describe('hydration follows ranked candidates', () => { expect(rows.map((row) => row.id)).toEqual(['nearer', 'near']) expect( - dbChainMockFns.execute.mock.calls.filter(([query]) => - render(query).sql.includes('WITH scored_search_candidates') - ) + dbChainMockFns.execute.mock.calls.filter(([query]) => isPageStatement(render(query).sql)) ).toHaveLength(2) expect( hasMockCondition( @@ -1268,7 +1300,7 @@ describe('permitted-document planner', () => { const statement = render(query).sql if (statement.includes('pg_index')) return indexedSourceRows if (isWalk(statement)) return traversedRows - if (statement.includes('WITH scored_search_candidates')) return rerankRows + if (isPageStatement(statement)) return rerankRows if (statement.includes('WITH readable_chunks')) return sourceExactRows if (isExactRanking(statement)) return exactRows if (isProbeStatement(statement)) return probeRows @@ -1414,9 +1446,7 @@ describe('permitted-document planner', () => { const walks = statements().filter((query) => isWalk(query.sql)) expect(walks).toHaveLength(1) expect(JSON.stringify(walks[0])).toContain('sliced-src') - const reranked = JSON.stringify( - statements().find((query) => query.sql.includes('scored_search_candidates')) - ) + const reranked = JSON.stringify(statements().find((query) => isPageStatement(query.sql))) expect(reranked).toContain('walked-hit') expect(reranked).not.toContain('arbitrary-hit') }) @@ -1443,9 +1473,7 @@ describe('permitted-document planner', () => { /** Uploads carry no connector, so their slice runs even with no sliced source beside them. */ const exact = statements().filter((query) => query.sql.includes('WITH readable_chunks')) expect(exact).toHaveLength(1) - expect( - JSON.stringify(statements().find((q) => q.sql.includes('scored_search_candidates'))) - ).toContain('upload-hit') + expect(JSON.stringify(statements().find((q) => isPageStatement(q.sql)))).toContain('upload-hit') }) it('ranks every source exactly when the caller is a member of none', async () => { @@ -1543,9 +1571,36 @@ describe('permitted-document planner', () => { /** The ranked CTE carries the mirrored source and ACL the predicate tests. */ expect(statement).toContain('AS connector_id') expect(statement).toContain('ranked_tin_chunks.acl') + /** Every row is filled, so none is decided on its document. */ + expect(statement).not.toContain('IS NULL AND EXISTS (') + /** The candidates matched where they were ranked; hydration does not match them again. */ + expect(JSON.stringify(dbChainMockFns.where.mock.calls)).not.toContain('@@') + }) + + it('decides a row the backfill has not reached on its document while the fill runs', async () => { + forgetProjectionFilled() + tinPages.push({ + ranked: 1, + candidates: [{ id: 'a', documentId: 'doc-a', connectorId: 'src-a' }], + }) + const execute = dbChainMockFns.execute.getMockImplementation()! + dbChainMockFns.execute.mockImplementation(async (query) => + render(query).sql.includes('AS unfilled') ? [{ unfilled: true }] : execute(query) + ) + await keyword({ + accessPlan: { + connectors: { workspace: [], admin: ['src-a'], members: [] }, + observers: { confirmed: [], observed: [] }, + memberSources: [], + connectorTypes: new Map(), + uploads: true, + }, + }) + const statement = JSON.stringify(tinStatements()[0]) /** A row the backfill has not filled (`acl IS NULL`) is decided on its document instead. */ expect(statement).toContain('IS NULL AND EXISTS (') expect(statement).toContain('ranked_tin_chunks.document_id') + forgetProjectionFilled() }) it('widens the window for a broad resolved scope whose first page came back short', async () => { @@ -1986,8 +2041,11 @@ describe('permitted-document planner', () => { const statement = render(query).sql /** The exclusion is the only clause that negates a connector membership. */ const rebuilt = JSON.stringify(query).includes('OR NOT (') - if (statement.includes('WITH scored_search_candidates')) - return rebuilt ? [hit('b', 'other-src')] : [hit('a', 'gated-src')] + /** The page carries no exclusion of its own; the rebuilt pool is what asks for 'b'. */ + if (isPageStatement(statement)) + return JSON.stringify(render(query).params).includes('"b"') + ? [hit('b', 'other-src')] + : [hit('a', 'gated-src')] /** The first walk's pool is the gated source's; the rebuilt one reaches the accessible chunk. */ if (isWalk(statement)) return Array.from({ length: 400 }, (_, i) => ({ @@ -2008,10 +2066,59 @@ describe('permitted-document planner', () => { }) expect(getForConnectors).toHaveBeenCalledOnce() expect(result.rows.map((row) => row.id)).toEqual(['b']) - const reranks = statements().filter((query) => query.sql.includes('scored_search_candidates')) - expect(reranks).toHaveLength(2) - expect(JSON.stringify(reranks[0])).not.toContain('OR NOT (') - expect(JSON.stringify(reranks[1])).toContain('OR NOT (') + /** The exclusion lives in the walk that rebuilds the pool, so the rebuilt page asks for 'b'. */ + const walks = statements().filter((query) => isWalk(query.sql)) + expect(walks).toHaveLength(2) + expect(JSON.stringify(walks[0])).not.toContain('OR NOT (') + expect(JSON.stringify(walks[1])).toContain('OR NOT (') + const pages = statements().filter((query) => isPageStatement(query.sql)) + expect(pages).toHaveLength(2) + expect(JSON.stringify(pages[1].params)).toContain('"b"') + }) + + it('excludes a denied source through its documents while the projection is unfilled', async () => { + forgetProjectionFilled() + queueTableRows(schemaMock.knowledgeConnector, [ + { + id: 'gated-src', + accessMode: 'admin', + connectorType: 'confluence', + githubRepository: false, + }, + ]) + dbChainMockFns.execute.mockImplementation(async (query) => { + const statement = render(query).sql + /** The fill has not reached every row, so a denied source cannot be read off the row. */ + if (statement.includes('AS unfilled')) return [{ unfilled: true }] + /** The mock renders nested fragments as parameters, so the marker is found in the whole query. */ + const rebuilt = JSON.stringify(query).includes('/* excluded sources */') + if (isPageStatement(statement)) + return JSON.stringify(render(query).params).includes('"b"') + ? [hit('b', 'other-src')] + : [hit('a', 'gated-src')] + if (isWalk(statement)) + return Array.from({ length: 400 }, (_, i) => ({ + id: i === 0 ? (rebuilt ? 'b' : 'a') : `w-${i}`, + distance: 0.1, + })) + return [] + }) + queueTableRows(schemaMock.embedding, []) + queueTableRows(schemaMock.embedding, [hit('b', 'other-src')]) + const getForConnectors = vi.fn(async () => reader) + const result = await retrieveKnowledgeSearch({ + ...liveSearch, + searchMode: 'vector', + access: reader, + accessProvider: { ...provider, getForConnectors }, + }) + expect(result.rows.map((row) => row.id)).toEqual(['b']) + const walks = statements().filter((query) => isWalk(query.sql)) + expect(walks).toHaveLength(2) + expect(JSON.stringify(walks[0])).not.toContain('/* excluded sources */') + expect(JSON.stringify(walks[1])).toContain('NOT EXISTS (SELECT 1 FROM') + expect(JSON.stringify(walks[1])).toContain('/* excluded sources */') + forgetProjectionFilled() }) it('hands back the unread slices of a page a denied source made it rebuild', async () => { @@ -2031,7 +2138,7 @@ describe('permitted-document planner', () => { dbChainMockFns.execute.mockImplementation(async (query) => { const statement = render(query).sql const rebuilt = JSON.stringify(query).includes('OR NOT (') - if (statement.includes('WITH scored_search_candidates')) + if (isPageStatement(statement)) return rebuilt ? [hit('b', 'other-src')] : [ @@ -2137,7 +2244,7 @@ describe('filters on a resolved scope', () => { if (statement.includes(') reached')) return [{ n: 250_000 }] if (isProbeStatement(statement)) return probeRows if (isWalk(statement)) return traversedRows - if (statement.includes('WITH scored_search_candidates')) return rerankRows + if (isPageStatement(statement)) return rerankRows if (statement.includes('ranked_tin_chunks')) return [{ ranked: 0, candidates: [] }] return [] }) @@ -2216,7 +2323,7 @@ describe('filters on a resolved scope', () => { return [{ 'QUERY PLAN': [{ Plan: { 'Plan Rows': 1_000_000 } }] }] if (statement.includes(') reached')) return [{ n: 250_000 }] if (isWalk(statement)) return traversedRows - if (statement.includes('WITH scored_search_candidates')) return rerankRows + if (isPageStatement(statement)) return rerankRows return [] }) const result = await retrieveKnowledgeSearch({ @@ -2253,7 +2360,7 @@ describe('filters on a resolved scope', () => { if (statement.includes('AS saturated')) return [{ id: 'doc-recent', connectorId: null, saturated: false }] if (isWalk(statement)) return traversedRows - if (statement.includes('WITH scored_search_candidates')) return rerankRows + if (isPageStatement(statement)) return rerankRows return [] }) const search = () => @@ -2327,7 +2434,7 @@ describe('filters on a resolved scope', () => { const statement = render(query).sql if (statement.includes('AS unfilled')) return [{ unfilled: true }] if (isWalk(statement)) return traversedRows - if (statement.includes('WITH scored_search_candidates')) return rerankRows + if (isPageStatement(statement)) return rerankRows return [] }) await handleVectorOnlySearch({ @@ -2369,7 +2476,7 @@ describe('filters on a resolved scope', () => { dbChainMockFns.execute.mockImplementation(async (query) => { const statement = render(query).sql if (isWalk(statement)) return traversedRows - if (statement.includes('WITH scored_search_candidates')) return rerankRows + if (isPageStatement(statement)) return rerankRows return [] }) await handleVectorOnlySearch({ diff --git a/apps/sim/lib/knowledge/search/queries.ts b/apps/sim/lib/knowledge/search/queries.ts index b5a101c94ec..f7080770bb2 100644 --- a/apps/sim/lib/knowledge/search/queries.ts +++ b/apps/sim/lib/knowledge/search/queries.ts @@ -7,6 +7,10 @@ import { embeddingSearch, knowledgeConnector, } from '@sim/db/schema' +import { + PROJECTION_SOURCE_ACL_TABLES, + type ProjectionSourceAclTable, +} from '@sim/db/script-migrations/0021_embedding_search_connector' import { createLogger } from '@sim/logger' import { sha256Hex } from '@sim/security/hash' import { getErrorMessage, getPostgresErrorCode } from '@sim/utils/errors' @@ -113,12 +117,20 @@ const PROJECTION_FILLED_TTL_MS = 60_000 * Whether the ranking projection still holds rows the backfill has not filled. Read off the * unfilled-rows index in microseconds and remembered briefly: the answer only ever changes once. */ -const projectionFilled = new LRUCache({ - max: 1, +const projectionFilled = new LRUCache< + ProjectionSourceAclTable, + boolean, + { budget: SearchBudget | undefined; stage: SearchStage } +>({ + max: PROJECTION_SOURCE_ACL_TABLES.length, ttl: PROJECTION_FILLED_TTL_MS, - fetchMethod: async () => { - const [row] = await db.execute<{ unfilled: boolean }>(sql` - SELECT EXISTS (SELECT 1 FROM ${embeddingSearch} WHERE ${embeddingSearch.acl} IS NULL) AS unfilled`) + /** The read that misses the cache spends the leg's own budget, like every other read of the leg. */ + fetchMethod: async (projection, _stale, { context }) => { + const table = projection === 'embedding_search' ? embeddingSearch : embeddingKeywordTin + const [row] = await runSearchQuery(context.budget, context.stage, (executor) => + executor.execute<{ unfilled: boolean }>(sql` + SELECT EXISTS (SELECT 1 FROM ${table} WHERE ${table.acl} IS NULL) AS unfilled`) + ) return !row?.unfilled }, }) @@ -135,9 +147,22 @@ export function forgetProjectionFilled(): void { */ const CANDIDATE_HNSW_EF_SEARCH = '200' const CANDIDATE_HNSW_SCAN_MEM_MULTIPLIER = '2' -const MIN_VECTOR_RERANK_CANDIDATES = 400 -const MAX_VECTOR_RERANK_CANDIDATES = 1600 -const VECTOR_RERANK_OVERSAMPLING = 32 +/** + * Candidates one walk gathers: the pages the search has asked for so far and as many again, in + * case the full read predicate refuses some. The walk ends as soon as it has them, so a pool the + * size of a page ends long before one sized for a rerank; a pool the pages outrun is walked again, + * wider. The ceiling bounds the widest walk. + */ +const VECTOR_CANDIDATE_POOL_MIN = 200 +const MAX_VECTOR_CANDIDATES = 1600 + +/** The pool a search needs to serve `needed` candidates, at least twice the last pool. */ +export function vectorCandidatePoolLimit(needed: number, previous: number | undefined): number { + return Math.min( + MAX_VECTOR_CANDIDATES, + Math.max(VECTOR_CANDIDATE_POOL_MIN, needed * 2, (previous ?? 0) * 2) + ) +} /** * The probe's share of the leg. It ranks nothing, so it must never be why the leg misses its * own deadline. @@ -885,11 +910,17 @@ function hydrateSearchCandidates( budget?: SearchBudget ) { const accessCondition = knowledgeAccessCondition(access) + /** + * The score comes from the projection's stored halfvec, the column the walk ranked on: the + * original vector lives out of line in toast storage that no cache holds, and reading it back + * for every hydrated row was a random page read per result on every novel query. + */ return runSearchQuery(budget, `${leg}.sql`, (executor) => executor .select(getSearchResultFields(distance)) .from(embedding) .innerJoin(document, eq(embedding.documentId, document.id)) + .leftJoin(embeddingSearch, eq(embeddingSearch.id, embedding.id)) .where( and( inArray(embedding.id, ids), @@ -1468,6 +1499,8 @@ async function selectSourceVectorCandidates(input: { plan: SearchAccessPlan tagCondition: SQL | undefined documentCondition: SQL | undefined + /** Sources the caller turned out not to hold, kept out of every source's ranking. */ + exclusion: SQL | undefined /** Whether every projection row carries its mirrored columns, so a walk needs no document. */ projectionFilled: boolean candidateDistance: SQL @@ -1486,7 +1519,8 @@ async function selectSourceVectorCandidates(input: { const base = and( inArray(embeddingSearch.knowledgeBaseId, input.knowledgeBaseIds), eq(embeddingSearch.enabled, true), - input.tagCondition + input.tagCondition, + input.exclusion ) type RankedChunks = Promise> /** @@ -1495,7 +1529,9 @@ async function selectSourceVectorCandidates(input: { * so the graph is not stalled by a document lookup per candidate; the tag filter, which lives on * the chunk, still joins. */ - const onRow = projectionCandidateAccessCondition(embeddingSearch, input.access, input.plan) + const onRow = projectionCandidateAccessCondition(embeddingSearch, input.access, input.plan, { + filled: input.projectionFilled, + }) const walk = (scope: SQL): (() => RankedChunks) => () => @@ -1602,7 +1638,12 @@ const SOURCE_RANKING_CONCURRENCY = 3 */ async function selectVectorResults(params: SearchParams): Promise { const queryVector = params.queryVector! - const distance = embeddingDistance(queryVector.dimensions, queryVector.vector) + /** One score for ranking, threshold and results alike: the projection's, which stays in cache. */ + const distance = embeddingCandidateDistance( + queryVector.dimensions, + queryVector.vector, + queryVector.model + ) const tagConditions = getStructuredTagFilters(params.structuredFilters ?? [], embedding) const conditions = [ inArray(embedding.knowledgeBaseId, params.knowledgeBaseIds), @@ -1619,15 +1660,7 @@ async function selectVectorResults(params: SearchParams): Promise } | undefined + let candidatePool: + | { excludedKey: string; ids: Array<{ id: string }>; limit: number; exhausted: boolean } + | undefined return selectAuthorizedSearchResults({ leg: 'vector', access: params.access, @@ -1678,6 +1713,7 @@ async function selectVectorResults(params: SearchParams): Promise - const plan = params.access.kind === 'user' ? params.accessPlan : undefined - const filled = plan ? ((await projectionFilled.fetch('embedding_search')) ?? false) : false /** * A source the caller is a member of that has its own index is walked on its own, which * beats ranking it exactly once it is large enough to have earned that index. @@ -1757,6 +1817,7 @@ async function selectVectorResults(params: SearchParams): Promise @@ -1781,7 +1843,7 @@ async function selectVectorResults(params: SearchParams): Promise= MAX_VECTOR_CANDIDATES, + } annotateSearchDiagnostics({ vectorCandidateCount: selected.length, vectorCandidateScan: selected.length < candidateLimit ? 'underfilled' : 'planned', }) } - const identities = candidatePool.ids - if (!identities.length) return { candidates: [], nextOffset: offset } - /** Score each bounded candidate once; sorting the materialized scalar cannot invoke HNSW again. */ - const page = await runSearchQuery(params.budget, 'vector.rerank', (executor) => - executor.execute(sql` - WITH scored_search_candidates AS MATERIALIZED ( - SELECT ${embedding.id} AS id, ${document.id} AS "documentId", - ${document.connectorId} AS "connectorId", - ${distance} AS distance - FROM ${embedding} INNER JOIN ${document} ON ${document.id} = ${embedding.documentId} - WHERE ${and( - inArray( - embedding.id, - identities.map(({ id }) => id) - ), - ...conditions, - ...visibility - )} - ) SELECT * FROM scored_search_candidates ORDER BY distance, id LIMIT ${limit} OFFSET ${offset} + /** + * The walk's order is the page's order: it ranked on the stored halfvec, and rescoring the + * pool against the original vectors read one out-of-line vector per candidate from storage + * no cache holds, seconds on a query nobody had run before. Only the page's identities are + * read here; the full read predicate follows at hydration, as before. A slice whose + * documents all went away since the walk is passed over, not mistaken for the pool's end. + */ + for (let start = offset; start < candidatePool.ids.length; start += limit) { + const slice = candidatePool.ids.slice(start, start + limit) + const ranked = new Map(slice.map((candidate, index) => [candidate.id, index])) + const identities = await runSearchQuery(params.budget, 'vector.page', (executor) => + executor.execute(sql` + SELECT ${embeddingSearch.id} AS id, ${document.id} AS "documentId", + ${document.connectorId} AS "connectorId" + FROM ${embeddingSearch} + INNER JOIN ${document} ON ${document.id} = ${embeddingSearch.documentId} + WHERE ${embeddingSearch.id} = ANY(${textArrayLiteral(slice.map((candidate) => candidate.id))}) `) - ) - return { - candidates: page, - nextOffset: offset + page.length, + ) + if (!identities.length) continue + const page = [...identities].sort( + (a, b) => (ranked.get(a.id) ?? 0) - (ranked.get(b.id) ?? 0) + ) + return { candidates: page, nextOffset: start + slice.length } } + return { candidates: [], nextOffset: candidatePool.ids.length } }, hydrate: (ids, authorized) => hydrateSearchCandidates( @@ -1974,6 +2043,13 @@ export async function executeKeywordSearch(params: KeywordSearchParams): Promise : {}), }) const accessPlan = access.kind === 'user' ? params.accessPlan : undefined + /** A filled projection decides readability on the ranked row alone; none of its rows needs the document. */ + const tinFilled = + accessPlan && tinQuery + ? ((await projectionFilled.fetch('embedding_keyword_tin', { + context: { budget: params.budget, stage: 'keyword.projection_filled' }, + })) ?? false) + : false /** The projection predicate over the ranked CTE's mirrored columns, plus any excluded source. */ const onRowKeywordVisibility = (excludedSources: readonly string[]) => and( @@ -1984,7 +2060,8 @@ export async function executeKeywordSearch(params: KeywordSearchParams): Promise documentId: sql`ranked_tin_chunks.document_id`, }, access, - accessPlan! + accessPlan!, + { filled: tinFilled } ), dateFilterCondition(params.filters) ? sql`EXISTS (SELECT 1 FROM ${document} WHERE ${and(sql`${document.id} = ranked_tin_chunks.document_id`, dateFilterCondition(params.filters))})` @@ -2185,13 +2262,22 @@ export async function executeKeywordSearch(params: KeywordSearchParams): Promise ) return { candidates, nextOffset: offset + candidates.length } }, + /** + * Every candidate already matched the query where it was ranked; matching it again here + * would detoast one text-search vector per result. The score is the projection's, like the + * vector leg's, so the two legs fuse on the same distance. + */ hydrate: (ids, authorized) => hydrateSearchCandidates( ids, authorized, - embeddingDistance(queryVector.dimensions, queryVector.vector).as('distance'), + embeddingCandidateDistance( + queryVector.dimensions, + queryVector.vector, + queryVector.model + ).as('distance'), params.filters, - conditions, + [inArray(embedding.knowledgeBaseId, knowledgeBaseIds), ...tagFilterConditions], 'keyword', params.budget ), diff --git a/apps/sim/lib/navigation/organization-rollout.test.ts b/apps/sim/lib/navigation/organization-rollout.test.ts index dee375993b4..e0f4aaf644c 100644 --- a/apps/sim/lib/navigation/organization-rollout.test.ts +++ b/apps/sim/lib/navigation/organization-rollout.test.ts @@ -22,13 +22,18 @@ vi.mock('@/lib/billing/core/access', () => ({ isOrganizationBillingBlocked: vi.fn().mockResolvedValue(false), })) -import { requireOrganizationSearchAvailable } from '@/lib/knowledge/access/availability' +import { + forgetKnowledgeAccessAvailability, + requireOrganizationSearchAvailable, +} from '@/lib/knowledge/access/availability' import { resolveAppEntryPath } from '@/lib/navigation/resolve-app-entry' afterAll(resetEnvFlagsMock) describe('organization rollout during impersonation', () => { beforeEach(() => { + /** Each case answers the same organization differently; the memo must not carry one across. */ + forgetKnowledgeAccessAvailability() vi.clearAllMocks() setEnvFlags({ isAppConfigEnabled: true, isHosted: true }) mocks.landing.mockImplementation(async (userId: string) =>