From a462b7dd235ba1c030a3a20b3e6d33c0009df914 Mon Sep 17 00:00:00 2001 From: Waleed Latif Date: Sun, 20 Sep 2026 21:17:46 -0700 Subject: [PATCH 1/5] improvement(knowledge): rank the vector page on the projection and admit searches in parallel --- .../app/api/knowledge/search/utils.test.ts | 5 +- apps/sim/lib/billing/core/usage-gate-cache.ts | 6 +- .../kb-block-search.integration.ts | 2 +- .../search-latency.integration.ts | 103 ++++++----- .../lib/knowledge/access/availability.test.ts | 13 ++ apps/sim/lib/knowledge/access/availability.ts | 34 ++++ apps/sim/lib/knowledge/access/predicate.ts | 14 +- .../lib/knowledge/application/search.test.ts | 2 +- apps/sim/lib/knowledge/application/search.ts | 115 +++++++------ apps/sim/lib/knowledge/search/diagnostics.ts | 4 +- apps/sim/lib/knowledge/search/queries.test.ts | 149 ++++++++++------ apps/sim/lib/knowledge/search/queries.ts | 161 ++++++++++++------ 12 files changed, 406 insertions(+), 202 deletions(-) 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..f36c5c883cf 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,23 @@ 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' + ) + return db.execute<{ id: string }>(sql`SELECT s.id FROM ${embeddingSearch} s + INNER JOIN document d ON d.id = s.document_id + WHERE s.knowledge_base_id = ${ids.knowledgeBaseId} AND s.enabled ${readerClause ?? sql``} + ORDER BY (${distance}) + 0, s.id + LIMIT ${limit}`) +} const captured: CapturedQuery[] = [] const report: Record = { fixture: ids, @@ -213,6 +233,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 +500,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 +540,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 +1027,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 +1048,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 +1099,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 +1144,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 +1192,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 s.document_id 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 +1369,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 +1458,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.test.ts b/apps/sim/lib/knowledge/application/search.test.ts index 2defebd93e1..c72426fb25a 100644 --- a/apps/sim/lib/knowledge/application/search.test.ts +++ b/apps/sim/lib/knowledge/application/search.test.ts @@ -356,8 +356,8 @@ describe('knowledge search application use case', () => { ).rejects.toThrow('Search is not enabled for this organization') expect(mocks.recordActivity).not.toHaveBeenCalled() expect(mocks.requireOrganizationSearch).toHaveBeenCalledExactlyOnceWith('org-canonical') + /** The gate runs beside the embedding call; billing and retrieval still never start after a refusal. */ expect(mocks.resolveBilling).not.toHaveBeenCalled() - expect(mocks.generateEmbedding).not.toHaveBeenCalled() expect(mocks.executeSearch).not.toHaveBeenCalled() }) diff --git a/apps/sim/lib/knowledge/application/search.ts b/apps/sim/lib/knowledge/application/search.ts index 1d355381606..b0c0f95812b 100644 --- a/apps/sim/lib/knowledge/application/search.ts +++ b/apps/sim/lib/knowledge/application/search.ts @@ -11,6 +11,7 @@ import { checkAndBillPayerOverageThreshold } from '@/lib/billing/threshold-billi import { OrchestrationError } from '@/lib/core/orchestration/types' import { resourceScopeFromOwner, resourceScopeKey } from '@/lib/core/resource-scope' import { PlatformEvents } from '@/lib/core/telemetry' +import { runDetached } from '@/lib/core/utils/background' import { generateRequestId } from '@/lib/core/utils/request' import { importDurableSecretProvenance } from '@/lib/execution/durable-secret-provenance' import { reportDurableSecretProvenanceRefusal } from '@/lib/execution/durable-secret-provenance-telemetry' @@ -267,10 +268,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 +283,36 @@ 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 embedding call below instead of ahead of it. + * A refusal still ends the search before any retrieval. + */ + 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,31 +368,40 @@ 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 + const [queryEmbedding, access, searchDefaults, billingAttribution, tagDefinitions] = + await Promise.all([ + hasQuery + ? measureSearchStage('embedding', () => + runWithKnowledgeModelInputProvenance(resultSecretRegistry, () => + generateSearchEmbedding( + input.query!, + embeddingTarget!, + context.workspaceId, + input.signal + ) ) ) - ) - : Promise.resolve(null), - measureSearchStage('access_scope', () => context.access.get()), - measureSearchStage('defaults', () => - resolveKnowledgeSearchDefaults({ - workspaceId: context.workspaceId, - organizationId: context.organizationId, + : Promise.resolve(null), + measureSearchStage('access_scope', () => context.access.get()), + measureSearchStage('defaults', () => + resolveKnowledgeSearchDefaults({ + workspaceId: context.workspaceId, + organizationId: context.organizationId, - /** The signed-in person, if any; never the billing owner or a key's creator. */ - userId: resolvePrincipalSubjectUserId(principal) ?? undefined, - requestedMode: input.searchMode, - }) - ), - ]) + /** The signed-in person, if any; never the billing owner or a key's creator. */ + userId: resolvePrincipalSubjectUserId(principal) ?? undefined, + 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() annotateSearchDiagnostics({ accessScopeKind: access.kind, @@ -601,7 +619,8 @@ const searchKnowledgeUseCase = defineAuthorizedKnowledgeUseCase({ ], }) ) - await measureSearchStage('overage_billing', () => + /** Whether the payer crossed a billing threshold is settled after the results are on their way. */ + runDetached('knowledge search overage billing', () => checkAndBillPayerOverageThreshold(billingAttribution.billingEntity) ) } catch (error) { @@ -609,11 +628,6 @@ const searchKnowledgeUseCase = defineAuthorizedKnowledgeUseCase({ } } - if (filters.length === 0) { - definitionsByKnowledgeBase = await measureSearchStage('tag_definitions', () => - getDocumentTagDefinitionsByKnowledgeBaseIds(knowledgeBaseIds) - ) - } const tagMaps = new Map( [...definitionsByKnowledgeBase].map(([knowledgeBaseId, definitions]) => [ knowledgeBaseId, @@ -723,10 +737,11 @@ const searchKnowledgeUseCase = defineAuthorizedKnowledgeUseCase({ }, afterSuccess: async ({ principal, context, input, result }) => { const actorUserId = resolvePrincipalSubjectUserId(principal) + /** The record is for the organization's activity view; the caller's results never wait on it. */ if (context.organizationId && actorUserId) { - await measureSearchStage('activity_recording', () => + runDetached('organization search activity', () => recordOrganizationSearchActivity({ - organizationId: context.organizationId, + organizationId: context.organizationId!, userId: actorUserId, surface: input.surface ?? 'other', results: result.results, diff --git a/apps/sim/lib/knowledge/search/diagnostics.ts b/apps/sim/lib/knowledge/search/diagnostics.ts index 43d27b0a443..be88897103b 100644 --- a/apps/sim/lib/knowledge/search/diagnostics.ts +++ b/apps/sim/lib/knowledge/search/diagnostics.ts @@ -50,7 +50,7 @@ export type SearchStage = | `${RetrievalLeg}.sql` | 'vector.settings' | 'vector.probe' - | 'vector.rerank' + | 'vector.page' | 'vector.exact_candidates' | 'vector.exact' | 'vector.candidate_search' @@ -84,7 +84,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..4d3d368149c 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,27 @@ 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('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 +581,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 +616,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 +663,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 +696,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 +927,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 +962,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 +975,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 +1049,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 +1084,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 +1285,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 +1431,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 +1458,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 +1556,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 +2026,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 +2051,14 @@ 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('hands back the unread slices of a page a denied source made it rebuild', async () => { @@ -2031,7 +2078,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 +2184,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 +2263,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 +2300,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 +2374,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 +2416,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..d60e32b6656 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,13 @@ 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({ + max: PROJECTION_SOURCE_ACL_TABLES.length, ttl: PROJECTION_FILLED_TTL_MS, - fetchMethod: async () => { + fetchMethod: async (projection) => { + const table = projection === 'embedding_search' ? embeddingSearch : embeddingKeywordTin const [row] = await db.execute<{ unfilled: boolean }>(sql` - SELECT EXISTS (SELECT 1 FROM ${embeddingSearch} WHERE ${embeddingSearch.acl} IS NULL) AS unfilled`) + SELECT EXISTS (SELECT 1 FROM ${table} WHERE ${table.acl} IS NULL) AS unfilled`) return !row?.unfilled }, }) @@ -135,9 +140,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 +903,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 +1492,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 +1512,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 +1522,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 +1631,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 +1653,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 +1706,7 @@ async function selectVectorResults(params: SearchParams): Promise @@ -1781,7 +1828,7 @@ async function selectVectorResults(params: SearchParams): Promise - 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} + const slice = candidatePool.ids.slice(offset, offset + limit) + if (!slice.length) return { candidates: [], nextOffset: 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. + */ + 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))}) `) ) + const page = [...identities].sort((a, b) => (ranked.get(a.id) ?? 0) - (ranked.get(b.id) ?? 0)) return { candidates: page, nextOffset: offset + page.length, @@ -1974,6 +2024,11 @@ 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')) ?? false) + : false /** The projection predicate over the ranked CTE's mirrored columns, plus any excluded source. */ const onRowKeywordVisibility = (excludedSources: readonly string[]) => and( @@ -1984,7 +2039,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 +2241,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 ), From a7bdc757a108c250bda4d8754bb57a0c7ee67768 Mon Sep 17 00:00:00 2001 From: Waleed Latif Date: Sun, 20 Sep 2026 21:26:56 -0700 Subject: [PATCH 2/5] improvement(knowledge): admit before embedding, budget the fill check, pass over emptied slices --- .../search-latency.integration.ts | 12 ++-- .../lib/knowledge/application/search.test.ts | 2 +- apps/sim/lib/knowledge/application/search.ts | 69 ++++++++++--------- apps/sim/lib/knowledge/search/diagnostics.ts | 2 + apps/sim/lib/knowledge/search/queries.test.ts | 15 ++++ apps/sim/lib/knowledge/search/queries.ts | 48 ++++++++----- .../navigation/organization-rollout.test.ts | 7 +- 7 files changed, 98 insertions(+), 57 deletions(-) diff --git a/apps/sim/lib/knowledge/__integration__/search-latency.integration.ts b/apps/sim/lib/knowledge/__integration__/search-latency.integration.ts index f36c5c883cf..e62aa5081c4 100644 --- a/apps/sim/lib/knowledge/__integration__/search-latency.integration.ts +++ b/apps/sim/lib/knowledge/__integration__/search-latency.integration.ts @@ -134,10 +134,12 @@ async function exactProjectionNeighbors(vector: number[], limit: number, readerC JSON.stringify(vector), 'text-embedding-3-small' ) - return db.execute<{ id: string }>(sql`SELECT s.id FROM ${embeddingSearch} s - INNER JOIN document d ON d.id = s.document_id - WHERE s.knowledge_base_id = ${ids.knowledgeBaseId} AND s.enabled ${readerClause ?? sql``} - ORDER BY (${distance}) + 0, s.id + /** 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[] = [] @@ -1197,7 +1199,7 @@ describe.skipIf(!enabled)('Knowledge search latency on a realistic indexed corpu const expected = await exactProjectionNeighbors( queryVector, actual.length, - sql`AND s.document_id IN (${sql.join( + sql`AND ${embeddingSearch.documentId} IN (${sql.join( documentIds.map((id) => sql`${id}`), sql`, ` )})` diff --git a/apps/sim/lib/knowledge/application/search.test.ts b/apps/sim/lib/knowledge/application/search.test.ts index c72426fb25a..2defebd93e1 100644 --- a/apps/sim/lib/knowledge/application/search.test.ts +++ b/apps/sim/lib/knowledge/application/search.test.ts @@ -356,8 +356,8 @@ describe('knowledge search application use case', () => { ).rejects.toThrow('Search is not enabled for this organization') expect(mocks.recordActivity).not.toHaveBeenCalled() expect(mocks.requireOrganizationSearch).toHaveBeenCalledExactlyOnceWith('org-canonical') - /** The gate runs beside the embedding call; billing and retrieval still never start after a refusal. */ expect(mocks.resolveBilling).not.toHaveBeenCalled() + expect(mocks.generateEmbedding).not.toHaveBeenCalled() expect(mocks.executeSearch).not.toHaveBeenCalled() }) diff --git a/apps/sim/lib/knowledge/application/search.ts b/apps/sim/lib/knowledge/application/search.ts index b0c0f95812b..54cab81506d 100644 --- a/apps/sim/lib/knowledge/application/search.ts +++ b/apps/sim/lib/knowledge/application/search.ts @@ -285,8 +285,9 @@ const searchKnowledgeUseCase = defineAuthorizedKnowledgeUseCase({ ) /** * Whether the organization may search at all, and whether this payer still may: neither - * depends on the query, so both run beside the embedding call below instead of ahead of it. - * A refusal still ends the search before any retrieval. + * 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) @@ -368,41 +369,41 @@ const searchKnowledgeUseCase = defineAuthorizedKnowledgeUseCase({ : undefined const resultSecretRegistry = preparedRegistry ?? input.resultSecretRegistry input.signal?.throwIfAborted() - const [queryEmbedding, access, searchDefaults, billingAttribution, tagDefinitions] = - await Promise.all([ - hasQuery - ? measureSearchStage('embedding', () => - runWithKnowledgeModelInputProvenance(resultSecretRegistry, () => - generateSearchEmbedding( - input.query!, - embeddingTarget!, - context.workspaceId, - input.signal - ) - ) - ) - : Promise.resolve(null), - measureSearchStage('access_scope', () => context.access.get()), - measureSearchStage('defaults', () => - resolveKnowledgeSearchDefaults({ - workspaceId: context.workspaceId, - organizationId: context.organizationId, + const [access, searchDefaults, billingAttribution, tagDefinitions] = await Promise.all([ + measureSearchStage('access_scope', () => context.access.get()), + measureSearchStage('defaults', () => + resolveKnowledgeSearchDefaults({ + workspaceId: context.workspaceId, + organizationId: context.organizationId, - /** The signed-in person, if any; never the billing owner or a key's creator. */ - userId: resolvePrincipalSubjectUserId(principal) ?? undefined, - 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), - ]) + /** The signed-in person, if any; never the billing owner or a key's creator. */ + userId: resolvePrincipalSubjectUserId(principal) ?? undefined, + 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, searchMode: searchDefaults.searchMode, diff --git a/apps/sim/lib/knowledge/search/diagnostics.ts b/apps/sim/lib/knowledge/search/diagnostics.ts index be88897103b..e8c753941b6 100644 --- a/apps/sim/lib/knowledge/search/diagnostics.ts +++ b/apps/sim/lib/knowledge/search/diagnostics.ts @@ -51,6 +51,8 @@ export type SearchStage = | 'vector.settings' | 'vector.probe' | 'vector.page' + | 'vector.projection_filled' + | 'keyword.projection_filled' | 'vector.exact_candidates' | 'vector.exact' | 'vector.candidate_search' diff --git a/apps/sim/lib/knowledge/search/queries.test.ts b/apps/sim/lib/knowledge/search/queries.test.ts index 4d3d368149c..3891575fcb3 100644 --- a/apps/sim/lib/knowledge/search/queries.test.ts +++ b/apps/sim/lib/knowledge/search/queries.test.ts @@ -559,6 +559,21 @@ describe('workspace-scoped vector retrieval', () => { 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) diff --git a/apps/sim/lib/knowledge/search/queries.ts b/apps/sim/lib/knowledge/search/queries.ts index d60e32b6656..30e7244e6df 100644 --- a/apps/sim/lib/knowledge/search/queries.ts +++ b/apps/sim/lib/knowledge/search/queries.ts @@ -117,13 +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({ +const projectionFilled = new LRUCache< + ProjectionSourceAclTable, + boolean, + { budget: SearchBudget | undefined; stage: SearchStage } +>({ max: PROJECTION_SOURCE_ACL_TABLES.length, ttl: PROJECTION_FILLED_TTL_MS, - fetchMethod: async (projection) => { + /** 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 db.execute<{ unfilled: boolean }>(sql` + 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 }, }) @@ -1770,7 +1777,11 @@ 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 + const filled = plan + ? ((await projectionFilled.fetch('embedding_search', { + context: { budget: params.budget, stage: 'vector.projection_filled' }, + })) ?? 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. @@ -1891,29 +1902,32 @@ async function selectVectorResults(params: SearchParams): Promise [candidate.id, index])) - const identities = await runSearchQuery(params.budget, 'vector.page', (executor) => - executor.execute(sql` + 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))}) `) - ) - const page = [...identities].sort((a, b) => (ranked.get(a.id) ?? 0) - (ranked.get(b.id) ?? 0)) - 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( @@ -2027,7 +2041,9 @@ export async function executeKeywordSearch(params: KeywordSearchParams): Promise /** 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')) ?? false) + ? ((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[]) => 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) => From 2f3634c3a096ee39646c01e1eccd46b4578afb1d Mon Sep 17 00:00:00 2001 From: Waleed Latif Date: Sun, 20 Sep 2026 21:28:57 -0700 Subject: [PATCH 3/5] improvement(knowledge): keep billing effects awaited and treat a capped pool as exhausted --- apps/sim/lib/knowledge/application/search.ts | 9 +++------ apps/sim/lib/knowledge/search/queries.ts | 3 ++- 2 files changed, 5 insertions(+), 7 deletions(-) diff --git a/apps/sim/lib/knowledge/application/search.ts b/apps/sim/lib/knowledge/application/search.ts index 54cab81506d..6530c6bb64a 100644 --- a/apps/sim/lib/knowledge/application/search.ts +++ b/apps/sim/lib/knowledge/application/search.ts @@ -11,7 +11,6 @@ import { checkAndBillPayerOverageThreshold } from '@/lib/billing/threshold-billi import { OrchestrationError } from '@/lib/core/orchestration/types' import { resourceScopeFromOwner, resourceScopeKey } from '@/lib/core/resource-scope' import { PlatformEvents } from '@/lib/core/telemetry' -import { runDetached } from '@/lib/core/utils/background' import { generateRequestId } from '@/lib/core/utils/request' import { importDurableSecretProvenance } from '@/lib/execution/durable-secret-provenance' import { reportDurableSecretProvenanceRefusal } from '@/lib/execution/durable-secret-provenance-telemetry' @@ -620,8 +619,7 @@ const searchKnowledgeUseCase = defineAuthorizedKnowledgeUseCase({ ], }) ) - /** Whether the payer crossed a billing threshold is settled after the results are on their way. */ - runDetached('knowledge search overage billing', () => + await measureSearchStage('overage_billing', () => checkAndBillPayerOverageThreshold(billingAttribution.billingEntity) ) } catch (error) { @@ -738,11 +736,10 @@ const searchKnowledgeUseCase = defineAuthorizedKnowledgeUseCase({ }, afterSuccess: async ({ principal, context, input, result }) => { const actorUserId = resolvePrincipalSubjectUserId(principal) - /** The record is for the organization's activity view; the caller's results never wait on it. */ if (context.organizationId && actorUserId) { - runDetached('organization search activity', () => + await measureSearchStage('activity_recording', () => recordOrganizationSearchActivity({ - organizationId: context.organizationId!, + organizationId: context.organizationId, userId: actorUserId, surface: input.surface ?? 'other', results: result.results, diff --git a/apps/sim/lib/knowledge/search/queries.ts b/apps/sim/lib/knowledge/search/queries.ts index 30e7244e6df..359ba64f802 100644 --- a/apps/sim/lib/knowledge/search/queries.ts +++ b/apps/sim/lib/knowledge/search/queries.ts @@ -1891,11 +1891,12 @@ async function selectVectorResults(params: SearchParams): Promise= MAX_VECTOR_CANDIDATES, } annotateSearchDiagnostics({ vectorCandidateCount: selected.length, From 50d5770dee8dfdfcbc2c3ed812f74f196dd84db6 Mon Sep 17 00:00:00 2001 From: Waleed Latif Date: Sun, 20 Sep 2026 21:44:28 -0700 Subject: [PATCH 4/5] improvement(knowledge): exclude a denied source through its documents while the projection is unfilled --- apps/sim/lib/knowledge/search/queries.test.ts | 42 +++++++++++++++++++ apps/sim/lib/knowledge/search/queries.ts | 18 ++++---- 2 files changed, 53 insertions(+), 7 deletions(-) diff --git a/apps/sim/lib/knowledge/search/queries.test.ts b/apps/sim/lib/knowledge/search/queries.test.ts index 3891575fcb3..6d9a45a86af 100644 --- a/apps/sim/lib/knowledge/search/queries.test.ts +++ b/apps/sim/lib/knowledge/search/queries.test.ts @@ -2076,6 +2076,48 @@ describe('permitted-document planner', () => { expect(JSON.stringify(pages[1].params)).toContain('"b"') }) + it('excludes a denied source through its documents while the projection is unfilled', async () => { + 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 }] + const rebuilt = JSON.stringify(query).includes('NOT EXISTS') + 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[1])).toContain('NOT EXISTS') + expect(JSON.stringify(walks[1])).toContain('gated-src') + forgetProjectionFilled() + }) + it('hands back the unread slices of a page a denied source made it rebuild', async () => { queueTableRows(schemaMock.knowledgeConnector, [ { diff --git a/apps/sim/lib/knowledge/search/queries.ts b/apps/sim/lib/knowledge/search/queries.ts index 359ba64f802..da80b34ad8e 100644 --- a/apps/sim/lib/knowledge/search/queries.ts +++ b/apps/sim/lib/knowledge/search/queries.ts @@ -1722,12 +1722,22 @@ async function selectVectorResults(params: SearchParams): Promise - const plan = params.access.kind === 'user' ? params.accessPlan : undefined - const filled = plan - ? ((await projectionFilled.fetch('embedding_search', { - context: { budget: params.budget, stage: 'vector.projection_filled' }, - })) ?? 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. From 9ea6ba4483ed69a17b82d017c512f026e321ffaa Mon Sep 17 00:00:00 2001 From: Waleed Latif Date: Sun, 20 Sep 2026 21:47:11 -0700 Subject: [PATCH 5/5] improvement(knowledge): mark the excluded-sources clause and reset the fill memo in its test --- apps/sim/lib/knowledge/search/queries.test.ts | 9 ++++++--- apps/sim/lib/knowledge/search/queries.ts | 4 ++-- 2 files changed, 8 insertions(+), 5 deletions(-) diff --git a/apps/sim/lib/knowledge/search/queries.test.ts b/apps/sim/lib/knowledge/search/queries.test.ts index 6d9a45a86af..134ad41fb9e 100644 --- a/apps/sim/lib/knowledge/search/queries.test.ts +++ b/apps/sim/lib/knowledge/search/queries.test.ts @@ -2077,6 +2077,7 @@ describe('permitted-document planner', () => { }) it('excludes a denied source through its documents while the projection is unfilled', async () => { + forgetProjectionFilled() queueTableRows(schemaMock.knowledgeConnector, [ { id: 'gated-src', @@ -2089,7 +2090,8 @@ describe('permitted-document planner', () => { 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 }] - const rebuilt = JSON.stringify(query).includes('NOT EXISTS') + /** 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')] @@ -2113,8 +2115,9 @@ describe('permitted-document planner', () => { 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[1])).toContain('NOT EXISTS') - expect(JSON.stringify(walks[1])).toContain('gated-src') + 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() }) diff --git a/apps/sim/lib/knowledge/search/queries.ts b/apps/sim/lib/knowledge/search/queries.ts index da80b34ad8e..f7080770bb2 100644 --- a/apps/sim/lib/knowledge/search/queries.ts +++ b/apps/sim/lib/knowledge/search/queries.ts @@ -1736,8 +1736,8 @@ async function selectVectorResults(params: SearchParams): Promise