diff --git a/apps/sim/lib/billing/core/billing-attribution.ts b/apps/sim/lib/billing/core/billing-attribution.ts index 9b996b6ad83..2dfdfaa21ce 100644 --- a/apps/sim/lib/billing/core/billing-attribution.ts +++ b/apps/sim/lib/billing/core/billing-attribution.ts @@ -742,13 +742,16 @@ export async function resolveBillingAttribution({ /** The organization payer is independent of the person making the request. */ export async function resolveOrganizationBillingPayer(organizationId: string) { - const [owner] = await db - .select({ userId: member.userId }) - .from(member) - .where(and(eq(member.organizationId, organizationId), eq(member.role, 'owner'))) - .limit(1) + /** The owner and the subscription are independent reads; neither waits on the other. */ + const [[owner], payerSubscription] = await Promise.all([ + db + .select({ userId: member.userId }) + .from(member) + .where(and(eq(member.organizationId, organizationId), eq(member.role, 'owner'))) + .limit(1), + getOrganizationSubscription(organizationId, { onError: 'throw' }), + ]) if (!owner) throw new Error('Organization billing owner is unavailable') - const payerSubscription = await getOrganizationSubscription(organizationId, { onError: 'throw' }) if (payerSubscription && payerSubscription.referenceId !== organizationId) throw new Error('Organization subscription belongs to a different payer') return { organizationId, billedAccountUserId: owner.userId, payerSubscription } 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 3b61286d0fa..9427b171efc 100644 --- a/apps/sim/lib/knowledge/__integration__/kb-block-search.integration.ts +++ b/apps/sim/lib/knowledge/__integration__/kb-block-search.integration.ts @@ -11,6 +11,7 @@ import { } from '@/lib/knowledge/__integration__/seed-source-access-fixture' import { createKnowledgeAccessProvider } from '@/lib/knowledge/access/scope' import { + forgetProjectionFilled, resolvePermittedDocuments, retrieveKnowledgeSearch, VECTOR_PROBE_DOCUMENT_LIMIT, @@ -92,6 +93,8 @@ describe('API-key KB block fan-out', () => { it.each([false, true])( 'completes 18 concurrent KB searches with access checks intact (tag filter: %s)', async (withTags) => { + /** The projection-fill memo outlives an iteration; each one must read it once, like a cold process. */ + forgetProjectionFilled() const previousDebug = db.$client.options.debug const statements: string[] = [] db.$client.options.debug = (_connection, query) => { @@ -132,10 +135,11 @@ describe('API-key KB block fan-out', () => { statements.filter((query) => query.includes(fragment)) /** * Every statement runs under the leg's deadline: the candidate search applies it with the - * scan settings in one statement, and the probe, the exact ranking, the rerank and - * hydration each open with one of their own. + * scan settings in one statement, and the probe, the exact ranking and hydration each + * open with one of their own. The projection-fill read is shared by the searches that + * miss its memo together, so it appears once. */ - expect(matching('statement_timeout')).toHaveLength(bases.length * 5) + expect(matching('statement_timeout')).toHaveLength(bases.length * 4 + 1) /** * A scope this small leaves the bounded traversal short of its candidate limit, so every * search probes once and rescues once — never a widening retry loop. @@ -143,7 +147,8 @@ 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('"embedding_search"."id" = ANY(')).toHaveLength(bases.length) + /** The walk carries each candidate's identities, so a filled projection reads no page. */ + expect(matching('"embedding_search"."id" = ANY(')).toHaveLength(0) /** The probe enumerates visible documents and reports saturation; it never ranks them. */ expect( statements.filter( diff --git a/apps/sim/lib/knowledge/access/availability.ts b/apps/sim/lib/knowledge/access/availability.ts index 6a99f6873c6..8128e4dd3a1 100644 --- a/apps/sim/lib/knowledge/access/availability.ts +++ b/apps/sim/lib/knowledge/access/availability.ts @@ -64,7 +64,10 @@ export async function resolveKnowledgeAccessAvailability( 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 ?? ''}` + /** An organization's answer depends on the organization alone; a workspace's on its viewer too. */ + const key = context.organizationId + ? `${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 @@ -89,14 +92,14 @@ async function readKnowledgeAccessAvailability( return { sourceMirrored: false, memberScoped: false } } if (context.organizationId) { - return { - sourceMirrored: - !isHosted || (await isOrganizationOnEnterprisePlan(context.organizationId, 'throw')), - memberScoped: await isScopedCredentialGroupsAvailable({ + const [enterprise, memberScoped] = await Promise.all([ + isHosted ? isOrganizationOnEnterprisePlan(context.organizationId, 'throw') : true, + isScopedCredentialGroupsAvailable({ kind: 'organization', organizationId: context.organizationId, }), - } + ]) + return { sourceMirrored: enterprise, memberScoped } } if (!context.workspaceId) throw new Error('Knowledge access requires a resource owner') const ownerBilling = diff --git a/apps/sim/lib/knowledge/application/contexts.ts b/apps/sim/lib/knowledge/application/contexts.ts index 3d2be05e2cb..085f8b47447 100644 --- a/apps/sim/lib/knowledge/application/contexts.ts +++ b/apps/sim/lib/knowledge/application/contexts.ts @@ -18,15 +18,12 @@ import { } from '@/lib/knowledge/connectors/service' import type { ActiveKnowledgeDocument } from '@/lib/knowledge/documents/service' import { getKnowledgeDocument, getKnowledgeDocumentById } from '@/lib/knowledge/documents/service' +import type { ActiveKnowledgeBaseReference } from '@/lib/knowledge/knowledge-base-reference' import { getRestorableKnowledgeBase, type RestorableKnowledgeBase, } from '@/lib/knowledge/orchestration/restore' -import { - type ActiveKnowledgeBaseReference, - getActiveKnowledgeBaseReference, - getKnowledgeBaseById, -} from '@/lib/knowledge/service' +import { getActiveKnowledgeBaseReference, getKnowledgeBaseById } from '@/lib/knowledge/service' import { getTagDefinitionById } from '@/lib/knowledge/tags/service' import type { DocumentTagDefinition } from '@/lib/knowledge/tags/types' import type { KnowledgeBaseWithCounts } from '@/lib/knowledge/types' diff --git a/apps/sim/lib/knowledge/application/search.ts b/apps/sim/lib/knowledge/application/search.ts index 43cc7ad22d9..41fe47443e7 100644 --- a/apps/sim/lib/knowledge/application/search.ts +++ b/apps/sim/lib/knowledge/application/search.ts @@ -33,6 +33,7 @@ import { instrumentSearchUseCase } from '@/lib/knowledge/application/search-diag import { ALL_TAG_SLOTS } from '@/lib/knowledge/constants' import { getEmbeddingModelInfo, toKbEmbeddingDimensions } from '@/lib/knowledge/embedding-models' import { generateSearchEmbedding, type KbEmbeddingTarget } from '@/lib/knowledge/embeddings' +import type { ActiveKnowledgeBaseReference } from '@/lib/knowledge/knowledge-base-reference' import { runWithKnowledgeModelInputProvenance } from '@/lib/knowledge/model-input-provenance' import { hasRerankerCredential, rerank } from '@/lib/knowledge/reranker' import type { RerankerStatus } from '@/lib/knowledge/reranker-models' @@ -47,10 +48,7 @@ import { type SearchResult, } from '@/lib/knowledge/search/queries' import { importKnowledgeSearchResultSecretProvenance } from '@/lib/knowledge/secret-provenance' -import { - type ActiveKnowledgeBaseReference, - getActiveKnowledgeBaseReferences, -} from '@/lib/knowledge/service' +import { getActiveKnowledgeBaseReferences } from '@/lib/knowledge/service' import { type KnowledgeTagNameFilter, resolveKnowledgeTagFilters, @@ -163,6 +161,7 @@ export interface SearchKnowledgeResult { knowledgeBaseId: string topK: number totalResults: number + rerankerStatus: RerankerStatus cost?: KnowledgeSearchCost workspaceId?: string userId: string @@ -171,10 +170,8 @@ export interface SearchKnowledgeResult { resultSecretRegistry?: ResolvedSecretTraceRegistry } -async function resolveKnowledgeSearchContext( - input: SearchKnowledgeInput, - principal: Principal -): Promise { +/** The request's shape, checked before anything is read for it. */ +export function validateKnowledgeSearchInput(input: SearchKnowledgeInput): void { if ( input.knowledgeBaseIds.length < 1 || input.knowledgeBaseIds.length > KNOWLEDGE_SEARCH_COST_POLICY.maxKnowledgeBases @@ -194,6 +191,37 @@ async function resolveKnowledgeSearchContext( `topK must be an integer between 1 and ${KNOWLEDGE_SEARCH_COST_POLICY.maxTopK}` ) } +} + +/** + * The search context over bases already resolved and authorized under `context`: what the + * caller may read across them comes from the principal, never from the input. + */ +export function buildKnowledgeSearchContext( + principal: Principal, + context: KnowledgeResourceContext, + knowledgeBases: ActiveKnowledgeBaseReference[], + input: Pick +): KnowledgeSearchContext { + const knowledgeBaseIds = knowledgeBases.map((base) => base.id) + const signal = input.signal + return { + ...context, + knowledgeBases, + access: createKnowledgeAccessProvider( + principal, + context.organizationId + ? { ...context, knowledgeBaseIds, signal } + : { workspaceId: context.workspaceId, knowledgeBaseIds, signal } + ), + } +} + +async function resolveKnowledgeSearchContext( + input: SearchKnowledgeInput, + principal: Principal +): Promise { + validateKnowledgeSearchInput(input) const knowledgeBases = await getActiveKnowledgeBaseReferences(input.knowledgeBaseIds) const missingIds = input.knowledgeBaseIds.filter((_, index) => { const knowledgeBase = knowledgeBases[index] @@ -225,19 +253,12 @@ async function resolveKnowledgeSearchContext( `Knowledge bases not found or access denied: ${input.knowledgeBaseIds.join(', ')}` ) } + const resolved = knowledgeBases as ActiveKnowledgeBaseReference[] if (canonicalOrganizationId) { const context = await resolveKnowledgeOrganizationContext({ organizationId: canonicalOrganizationId, }) - return { - ...context, - knowledgeBases: knowledgeBases as ActiveKnowledgeBaseReference[], - access: createKnowledgeAccessProvider(principal, { - ...context, - knowledgeBaseIds: knowledgeBases.map((base) => base!.id), - signal: input.signal, - }), - } + return buildKnowledgeSearchContext(principal, context, resolved, input) } if (!canonicalWorkspaceId) { throw new OrchestrationError('not_found', 'Knowledge base not found') @@ -245,129 +266,132 @@ async function resolveKnowledgeSearchContext( const workspaceContext = await resolveKnowledgeWorkspaceContext({ workspaceId: canonicalWorkspaceId, }) - return { - ...workspaceContext, - knowledgeBases: knowledgeBases as ActiveKnowledgeBaseReference[], - access: createKnowledgeAccessProvider(principal, { - workspaceId: canonicalWorkspaceId, - knowledgeBaseIds: knowledgeBases.map((base) => base!.id), - signal: input.signal, - }), - } + return buildKnowledgeSearchContext(principal, workspaceContext, resolved, input) } -const searchKnowledgeUseCase = defineAuthorizedKnowledgeUseCase({ - operation: knowledgeOperations.search, - resolveContext: ({ principal, input }: { principal: Principal; input: SearchKnowledgeInput }) => - measureSearchStage('knowledge_context', () => resolveKnowledgeSearchContext(input, principal)), - async execute({ principal, input, context }) { - annotateSearchDiagnostics({ - scopeKind: context.organizationId ? 'organization' : 'workspace', - knowledgeBaseCount: context.knowledgeBases.length, - }) - input.signal?.throwIfAborted() - const requestId = generateRequestId() - const hasQuery = Boolean(input.query?.trim()) - const filters = input.tagFilters ?? [] - if (!hasQuery && filters.length === 0) { - throw new OrchestrationError( - 'validation', - 'Please provide either a search query or tag filters to search your knowledge base' - ) - } - const userId = resolveKnowledgeAttributedUserId(principal, context) - const shouldMeter = !( - input.skipUsageBilling && - principal.kind === 'delegated' && - principal.serviceId === 'executor' +export interface KnowledgeSearchExecution { + principal: Principal + input: SearchKnowledgeInput + context: KnowledgeSearchContext +} + +/** + * The search itself, over a context its caller has already resolved and authorized: the + * operation each search surface shares once it has decided which bases the request may read. + */ +export async function runKnowledgeSearch({ + principal, + input, + context, +}: KnowledgeSearchExecution): Promise { + annotateSearchDiagnostics({ + scopeKind: context.organizationId ? 'organization' : 'workspace', + knowledgeBaseCount: context.knowledgeBases.length, + }) + input.signal?.throwIfAborted() + const requestId = generateRequestId() + const hasQuery = Boolean(input.query?.trim()) + const filters = input.tagFilters ?? [] + if (!hasQuery && filters.length === 0) { + throw new OrchestrationError( + 'validation', + 'Please provide either a search query or tag filters to search your knowledge base' ) - /** - * 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.' + } + const userId = resolveKnowledgeAttributedUserId(principal, context) + const shouldMeter = !( + input.skipUsageBilling && + principal.kind === 'delegated' && + principal.serviceId === 'executor' + ) + /** + * 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 } + return billingAttribution + } - const knowledgeBaseIds = context.knowledgeBases.map((knowledgeBase) => knowledgeBase.id) - let structuredFilters: StructuredFilter[] = [] - let definitionsByKnowledgeBase = new Map() - if (filters.length > 0) { - const built = await measureSearchStage('tag_filters', () => - resolveKnowledgeTagFilters(filters, knowledgeBaseIds) - ) - structuredFilters = built.structuredFilters - definitionsByKnowledgeBase = built.definitionsByKnowledgeBase - } + const knowledgeBaseIds = context.knowledgeBases.map((knowledgeBase) => knowledgeBase.id) + let structuredFilters: StructuredFilter[] = [] + let definitionsByKnowledgeBase = new Map() + if (filters.length > 0) { + const built = await measureSearchStage('tag_filters', () => + resolveKnowledgeTagFilters(filters, knowledgeBaseIds) + ) + structuredFilters = built.structuredFilters + definitionsByKnowledgeBase = built.definitionsByKnowledgeBase + } - /** - * One query embedding serves every leg, so every base in the request has to - * be indexed the same way. The width is part of that: it selects the - * pgvector column each comparison reads, and two bases on the same model at - * different widths still live in different columns. - * - * Built only for a query search. A tag-only request never embeds anything, - * so resolving a width it will not use would let one base recorded at an - * unstorable width fail a request that does not depend on it. - */ - const embeddingTargets = new Map( - context.knowledgeBases.map((kb) => [ - `${kb.embeddingModel}:${kb.embeddingDimension}`, - { model: kb.embeddingModel, dimensions: kb.embeddingDimension }, - ]) + /** + * One query embedding serves every leg, so every base in the request has to + * be indexed the same way. The width is part of that: it selects the + * pgvector column each comparison reads, and two bases on the same model at + * different widths still live in different columns. + * + * Built only for a query search. A tag-only request never embeds anything, + * so resolving a width it will not use would let one base recorded at an + * unstorable width fail a request that does not depend on it. + */ + const embeddingTargets = new Map( + context.knowledgeBases.map((kb) => [ + `${kb.embeddingModel}:${kb.embeddingDimension}`, + { model: kb.embeddingModel, dimensions: kb.embeddingDimension }, + ]) + ) + if (hasQuery && embeddingTargets.size > 1) { + throw new OrchestrationError( + 'validation', + 'Selected knowledge bases use different embedding models or vector widths and cannot be searched together. Search them separately.' ) - if (hasQuery && embeddingTargets.size > 1) { - throw new OrchestrationError( - 'validation', - 'Selected knowledge bases use different embedding models or vector widths and cannot be searched together. Search them separately.' + } + const selectedTarget = [...embeddingTargets.values()][0] + const embeddingModel = selectedTarget.model + /** + * The width is narrowed to a storable one only for a query search, which is + * the only kind that reads a vector column. A tag-only search must not fail + * on a width it never uses. + */ + const embeddingTarget: KbEmbeddingTarget | undefined = hasQuery + ? { + model: selectedTarget.model, + dimensions: toKbEmbeddingDimensions(selectedTarget.dimensions), + } + : undefined + const preparedRegistry = input.prepareModelInputProvenance + ? await measureSearchStage('input_provenance', () => + input.prepareModelInputProvenance!({ userId, workspaceId: context.workspaceId }) ) - } - const selectedTarget = [...embeddingTargets.values()][0] - const embeddingModel = selectedTarget.model - /** - * The width is narrowed to a storable one only for a query search, which is - * the only kind that reads a vector column. A tag-only search must not fail - * on a width it never uses. - */ - const embeddingTarget: KbEmbeddingTarget | undefined = hasQuery - ? { - model: selectedTarget.model, - dimensions: toKbEmbeddingDimensions(selectedTarget.dimensions), - } - : undefined - const preparedRegistry = input.prepareModelInputProvenance - ? await measureSearchStage('input_provenance', () => - input.prepareModelInputProvenance!({ userId, workspaceId: context.workspaceId }) - ) - : undefined - const resultSecretRegistry = preparedRegistry ?? input.resultSecretRegistry - input.signal?.throwIfAborted() - const [access, searchDefaults, billingAttribution, tagDefinitions] = await Promise.all([ + : undefined + const resultSecretRegistry = preparedRegistry ?? input.resultSecretRegistry + input.signal?.throwIfAborted() + const [access, searchDefaults, billingAttribution, tagDefinitions, rerankerCredential] = + await Promise.all([ measureSearchStage('access_scope', () => context.access.get()), measureSearchStage('defaults', () => resolveKnowledgeSearchDefaults({ @@ -386,377 +410,393 @@ const searchKnowledgeUseCase = defineAuthorizedKnowledgeUseCase({ getDocumentTagDefinitionsByKnowledgeBaseIds(knowledgeBaseIds) ) : Promise.resolve(definitionsByKnowledgeBase), + /** A surface may ask to rerank; without a key for the workspace or the platform there is nothing to ask. */ + input.rerankerEnabled && hasQuery + ? hasRerankerCredential(context.workspaceId, input.rerankerApiKey) + : false, ]) - definitionsByKnowledgeBase = tagDefinitions - input.signal?.throwIfAborted() - const queryEmbedding = hasQuery - ? await measureSearchStage('embedding', () => - runWithKnowledgeModelInputProvenance(resultSecretRegistry, () => - generateSearchEmbedding( - input.query!, - embeddingTarget!, - context.workspaceId, - input.signal - ) - ) + definitionsByKnowledgeBase = tagDefinitions + input.signal?.throwIfAborted() + /** Requested only once every prerequisite held: a search refused for any reason spends no model call. */ + 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, + boostRecency: searchDefaults.boostRecency, + embeddingDimensions: embeddingTarget?.dimensions, + }) + const useReranker = rerankerCredential + const candidateTopK = useReranker + ? input.rerankerInputCount !== undefined + ? Math.min( + KNOWLEDGE_SEARCH_COST_POLICY.maxTopK, + Math.max(input.topK, input.rerankerInputCount) ) - : null - input.signal?.throwIfAborted() - annotateSearchDiagnostics({ - accessScopeKind: access.kind, + : Math.min(KNOWLEDGE_SEARCH_COST_POLICY.maxTopK, input.topK * 4) + : input.topK + const retrieved = await measureSearchStage('retrieval', () => + retrieveKnowledgeSearch({ + vectorBudgetMs: input.vectorBudgetMs, + knowledgeBaseIds, + topK: candidateTopK, + filters: input.filters, + access, + accessProvider: context.access, + signal: input.signal, searchMode: searchDefaults.searchMode, boostRecency: searchDefaults.boostRecency, - embeddingDimensions: embeddingTarget?.dimensions, + query: input.query, + queryVector: hasQuery + ? { + vector: JSON.stringify(queryEmbedding?.embedding ?? null), + dimensions: embeddingTarget!.dimensions, + model: embeddingTarget!.model, + } + : undefined, + structuredFilters: structuredFilters.length > 0 ? structuredFilters : undefined, }) - /** A surface may ask to rerank; without a key for the workspace or the platform there is nothing to ask. */ - const useReranker = - Boolean(input.rerankerEnabled && hasQuery) && - (await hasRerankerCredential(context.workspaceId, input.rerankerApiKey)) - const candidateTopK = useReranker - ? input.rerankerInputCount !== undefined - ? Math.min( - KNOWLEDGE_SEARCH_COST_POLICY.maxTopK, - Math.max(input.topK, input.rerankerInputCount) - ) - : Math.min(KNOWLEDGE_SEARCH_COST_POLICY.maxTopK, input.topK * 4) - : input.topK - const retrieved = await measureSearchStage('retrieval', () => - retrieveKnowledgeSearch({ - vectorBudgetMs: input.vectorBudgetMs, - knowledgeBaseIds, - topK: candidateTopK, - filters: input.filters, - access, - accessProvider: context.access, - signal: input.signal, - searchMode: searchDefaults.searchMode, - boostRecency: searchDefaults.boostRecency, - query: input.query, - queryVector: hasQuery - ? { - vector: JSON.stringify(queryEmbedding?.embedding ?? null), - dimensions: embeddingTarget!.dimensions, - model: embeddingTarget!.model, - } - : undefined, - structuredFilters: structuredFilters.length > 0 ? structuredFilters : undefined, + ) + + annotateSearchDiagnostics({ + retrievalStatus: retrieved.retrieval.status, + timedOutLegs: retrieved.retrieval.timedOutLegs, + }) + if (retrieved.retrieval.status === 'partial' && !input.allowPartialResults) + throw new SearchDeadlineError() + let rows = retrieved.rows + input.signal?.throwIfAborted() + /** Public callers have no input envelope, but persisted reranker inputs still need provenance. */ + const registrySubjectUserId = resolvePrincipalSubjectUserId(principal) + const registry = + resultSecretRegistry ?? + (input.prepareModelInputProvenance || useReranker + ? new ResolvedSecretTraceRegistry( + [], + registrySubjectUserId + ? { userId: registrySubjectUserId, workspaceId: context.workspaceId } + : undefined + ) + : undefined) + let provenanceSnapshot: Awaited< + ReturnType + > | null = null + if (registry) { + provenanceSnapshot = await measureSearchStage('result_provenance', () => + importKnowledgeSearchResultSecretProvenance({ + registry, + results: rows, }) ) - - annotateSearchDiagnostics({ - retrievalStatus: retrieved.retrieval.status, - timedOutLegs: retrieved.retrieval.timedOutLegs, - }) - if (retrieved.retrieval.status === 'partial' && !input.allowPartialResults) - throw new SearchDeadlineError() - let rows = retrieved.rows - input.signal?.throwIfAborted() - /** Public callers have no input envelope, but persisted reranker inputs still need provenance. */ - const registrySubjectUserId = resolvePrincipalSubjectUserId(principal) - const registry = - resultSecretRegistry ?? - (input.prepareModelInputProvenance || useReranker - ? new ResolvedSecretTraceRegistry( - [], - registrySubjectUserId - ? { userId: registrySubjectUserId, workspaceId: context.workspaceId } - : undefined - ) - : undefined) - let provenanceSnapshot: Awaited< - ReturnType - > | null = null - if (registry) { - provenanceSnapshot = await measureSearchStage('result_provenance', () => - importKnowledgeSearchResultSecretProvenance({ - registry, - results: rows, + if (!provenanceSnapshot.imported) { + registry.markIncomplete('knowledge-result-provenance-unavailable') + if (useReranker) { + reportDurableSecretProvenanceRefusal({ + surface: 'knowledge', + cause: 'knowledge-result-provenance-unavailable', + workspaceId: context.workspaceId, }) - ) - if (!provenanceSnapshot.imported) { - registry.markIncomplete('knowledge-result-provenance-unavailable') - if (useReranker) { - reportDurableSecretProvenanceRefusal({ - surface: 'knowledge', - cause: 'knowledge-result-provenance-unavailable', - workspaceId: context.workspaceId, - }) - throw new KnowledgeSearchProvenanceUnavailableError() - } + throw new KnowledgeSearchProvenanceUnavailableError() } } + } - const rerankerScores = new Map() - let rerankerBilled = false - let rerankerIsBYOK = false - /** - * Returned on every search. The fallback to vector ordering is deliberate — a - * Cohere outage should not take knowledge search down with it — but until this - * was reported the fallback was also invisible: a 200 whose results were - * byte-identical to an unreranked search, with no `rerankerScore` anywhere and - * nothing to say why. - * - * It starts at the outcome that holds if the rerank call below never happens or - * never completes, so only the success path has to move it. A request with - * nothing to rank — no query text, or no candidate rows — is `skipped` rather - * than `unavailable`: the reranker was never the obstacle. Anything else that - * was asked for and did not produce a usable ordering is `unavailable`, - * including a request that reaches here with no model, which no HTTP contract - * can now produce. - * - * A call that returns without raising but hands back an empty ordering counts - * as `unavailable` too, and it is not the reranker "matching nothing": - * `rerank` asks for `top_n` over a non-empty document list, so a provider that - * ranked them returns one entry per document. Empty means the response carried - * nothing usable — no results, or only indices outside the batch, which - * `rerank` drops. The caller is left in vector order with no `rerankerScore`, - * which is exactly what `unavailable` promises, and retrying is exactly the - * right advice. - */ - let rerankerStatus: RerankerStatus = !input.rerankerEnabled - ? 'not_requested' - : !hasQuery || rows.length === 0 - ? 'skipped' - : 'unavailable' - if (useReranker && input.rerankerModel && rows.length > 0) { - const candidateCount = rows.length - try { - const reranked = await measureSearchStage('reranking', () => - runWithKnowledgeModelInputProvenance(registry, () => - rerank( - input.query!, - rows.map((row) => ({ id: row.id, text: row.content })), - { - model: input.rerankerModel!, - topN: input.topK, - workspaceId: context.workspaceId, + const rerankerScores = new Map() + let rerankerBilled = false + let rerankerIsBYOK = false + /** + * Returned on every search. The fallback to vector ordering is deliberate — a + * Cohere outage should not take knowledge search down with it — but until this + * was reported the fallback was also invisible: a 200 whose results were + * byte-identical to an unreranked search, with no `rerankerScore` anywhere and + * nothing to say why. + * + * It starts at the outcome that holds if the rerank call below never happens or + * never completes, so only the success path has to move it. A request with + * nothing to rank — no query text, or no candidate rows — is `skipped` rather + * than `unavailable`: the reranker was never the obstacle. Anything else that + * was asked for and did not produce a usable ordering is `unavailable`, + * including a request that reaches here with no model, which no HTTP contract + * can now produce. + * + * A call that returns without raising but hands back an empty ordering counts + * as `unavailable` too, and it is not the reranker "matching nothing": + * `rerank` asks for `top_n` over a non-empty document list, so a provider that + * ranked them returns one entry per document. Empty means the response carried + * nothing usable — no results, or only indices outside the batch, which + * `rerank` drops. The caller is left in vector order with no `rerankerScore`, + * which is exactly what `unavailable` promises, and retrying is exactly the + * right advice. + */ + let rerankerStatus: RerankerStatus = !input.rerankerEnabled + ? 'not_requested' + : !hasQuery || rows.length === 0 + ? 'skipped' + : 'unavailable' + if (useReranker && input.rerankerModel && rows.length > 0) { + const candidateCount = rows.length + try { + const reranked = await measureSearchStage('reranking', () => + runWithKnowledgeModelInputProvenance(registry, () => + rerank( + input.query!, + rows.map((row) => ({ id: row.id, text: row.content })), + { + model: input.rerankerModel!, + topN: input.topK, + workspaceId: context.workspaceId, - apiKey: input.rerankerApiKey, - signal: input.signal, - } - ) + apiKey: input.rerankerApiKey, + signal: input.signal, + } ) ) - rerankerBilled = true - rerankerIsBYOK = reranked.isBYOK - if (reranked.results.length === 0) { - rows = rows.slice(0, input.topK) - } else { - const byId = new Map(rows.map((row) => [row.id, row])) - rows = reranked.results - .map((ranked) => byId.get(ranked.item.id)) - .filter((row): row is SearchResult => Boolean(row)) - for (const ranked of reranked.results) { - rerankerScores.set(ranked.item.id, ranked.relevanceScore) - } - rerankerStatus = 'applied' - } - } catch (error) { - input.signal?.throwIfAborted() - if (registry?.isPermanentlyIncomplete()) throw error - logger.warn('Knowledge reranker failed; using vector ordering', { - error: getErrorMessage(error), - model: input.rerankerModel, - candidateCount, - }) + ) + rerankerBilled = true + rerankerIsBYOK = reranked.isBYOK + if (reranked.results.length === 0) { rows = rows.slice(0, input.topK) - rerankerStatus = 'unavailable' + } else { + const byId = new Map(rows.map((row) => [row.id, row])) + rows = reranked.results + .map((ranked) => byId.get(ranked.item.id)) + .filter((row): row is SearchResult => Boolean(row)) + for (const ranked of reranked.results) { + rerankerScores.set(ranked.item.id, ranked.relevanceScore) + } + rerankerStatus = 'applied' } - logger.info('Knowledge reranker completed', { - status: rerankerStatus, + } catch (error) { + input.signal?.throwIfAborted() + if (registry?.isPermanentlyIncomplete()) throw error + logger.warn('Knowledge reranker failed; using vector ordering', { + error: getErrorMessage(error), + model: input.rerankerModel, candidateCount, - resultCount: rows.length, - workspaceId: context.workspaceId, }) - } else if (useReranker) { rows = rows.slice(0, input.topK) + rerankerStatus = 'unavailable' } + logger.info('Knowledge reranker completed', { + status: rerankerStatus, + candidateCount, + resultCount: rows.length, + workspaceId: context.workspaceId, + }) + } else if (useReranker) { + rows = rows.slice(0, input.topK) + } - let tokenCount = 0 - let baseCost: ReturnType | null = null - if (hasQuery) { - tokenCount = estimateTokenCount( - input.query!, - getEmbeddingModelInfo(embeddingModel).tokenizerProvider - ).count - if (!queryEmbedding?.isBYOK) baseCost = calculateCost(embeddingModel, tokenCount, 0, false) - } - let rerankerCost = 0 - if (rerankerBilled && input.rerankerModel && !rerankerIsBYOK) { - const pricing = getRerankModelPricing(input.rerankerModel) - if (pricing) { - rerankerCost = pricing.perSearchUnit - baseCost = baseCost - ? { - ...baseCost, - input: baseCost.input + rerankerCost, - total: baseCost.total + rerankerCost, - } - : { - input: rerankerCost, - output: 0, - total: rerankerCost, - pricing: { input: 0, output: 0, updatedAt: pricing.updatedAt }, - } - } + let tokenCount = 0 + let baseCost: ReturnType | null = null + if (hasQuery) { + tokenCount = estimateTokenCount( + input.query!, + getEmbeddingModelInfo(embeddingModel).tokenizerProvider + ).count + if (!queryEmbedding?.isBYOK) baseCost = calculateCost(embeddingModel, tokenCount, 0, false) + } + let rerankerCost = 0 + if (rerankerBilled && input.rerankerModel && !rerankerIsBYOK) { + const pricing = getRerankModelPricing(input.rerankerModel) + if (pricing) { + rerankerCost = pricing.perSearchUnit + baseCost = baseCost + ? { + ...baseCost, + input: baseCost.input + rerankerCost, + total: baseCost.total + rerankerCost, + } + : { + input: rerankerCost, + output: 0, + total: rerankerCost, + pricing: { input: 0, output: 0, updatedAt: pricing.updatedAt }, + } } - if (shouldMeter && billingAttribution && baseCost && baseCost.total > 0) { - try { - await measureSearchStage('usage_recording', () => - recordUsage({ - userId, - ...(context.workspaceId ? { workspaceId: context.workspaceId } : {}), - ...toBillingContext(billingAttribution), - entries: [ - { - category: 'model', - source: 'knowledge-base', - description: embeddingModel, - cost: baseCost.total, - sourceReference: `kb-search:${requestId}`, - }, - ], - }) - ) - await measureSearchStage('overage_billing', () => - checkAndBillPayerOverageThreshold(billingAttribution.billingEntity) - ) - } catch (error) { - logger.error('Failed to record Knowledge search usage', { error }) - } + } + if (shouldMeter && billingAttribution && baseCost && baseCost.total > 0) { + try { + await measureSearchStage('usage_recording', () => + recordUsage({ + userId, + ...(context.workspaceId ? { workspaceId: context.workspaceId } : {}), + ...toBillingContext(billingAttribution), + entries: [ + { + category: 'model', + source: 'knowledge-base', + description: embeddingModel, + cost: baseCost.total, + sourceReference: `kb-search:${requestId}`, + }, + ], + }) + ) + await measureSearchStage('overage_billing', () => + checkAndBillPayerOverageThreshold(billingAttribution.billingEntity) + ) + } catch (error) { + logger.error('Failed to record Knowledge search usage', { error }) } + } - const tagMaps = new Map( - [...definitionsByKnowledgeBase].map(([knowledgeBaseId, definitions]) => [ - knowledgeBaseId, - new Map(definitions.map((definition) => [definition.tagSlot, definition.displayName])), - ]) - ) - /** - * The provenance snapshot vouches for the name, URL, and tags a model may see; the source - * card's modified time and connector type ride on the hydrated row, read under the same - * predicate as the content. - */ - const results = rows.map((row): KnowledgeSearchItem => { - const metadata: Record = {} - const tagMap = tagMaps.get(row.knowledgeBaseId) - const provenanceDocument = provenanceSnapshot?.documentMetadata[row.documentId] - const document = provenanceDocument ?? row - for (const slot of ALL_TAG_SLOTS) { - const value = - provenanceDocument && slot.startsWith('tag') - ? provenanceDocument[ - slot as 'tag1' | 'tag2' | 'tag3' | 'tag4' | 'tag5' | 'tag6' | 'tag7' - ] - : row[slot] - if (value !== null && value !== undefined) metadata[tagMap?.get(slot) ?? slot] = value - } - const rerankerScore = rerankerScores.get(row.id) - return { - embeddingId: row.id, - knowledgeBaseId: row.knowledgeBaseId, - documentId: row.documentId, - documentName: document?.filename ?? null, - sourceUrl: document?.sourceUrl ?? null, - sourceModifiedAt: row.sourceModifiedAt ?? null, - connectorType: row.connectorType ?? null, - content: row.content, - chunkIndex: row.chunkIndex, - metadata, - similarity: hasQuery ? 1 - row.distance : 1, - ...(rerankerScore !== undefined ? { rerankerScore } : {}), - } - }) - if (registry && provenanceSnapshot) { - for (const [documentId, document] of Object.entries(provenanceSnapshot.documentMetadata)) { - const renderedMetadata = results - .filter((result) => result.documentId === documentId) - .map((result) => ({ - documentName: result.documentName, - sourceUrl: result.sourceUrl, - metadata: result.metadata, - })) - if (renderedMetadata.length === 0) continue - if ( - !(await measureSearchStage('metadata_provenance', () => - importDurableSecretProvenance(registry, document.provenance, renderedMetadata) - )) - ) { - registry.markIncomplete('knowledge-result-provenance-unavailable') - } - } + const tagMaps = new Map( + [...definitionsByKnowledgeBase].map(([knowledgeBaseId, definitions]) => [ + knowledgeBaseId, + new Map(definitions.map((definition) => [definition.tagSlot, definition.displayName])), + ]) + ) + /** + * The provenance snapshot vouches for the name, URL, and tags a model may see; the source + * card's modified time and connector type ride on the hydrated row, read under the same + * predicate as the content. + */ + const results = rows.map((row): KnowledgeSearchItem => { + const metadata: Record = {} + const tagMap = tagMaps.get(row.knowledgeBaseId) + const provenanceDocument = provenanceSnapshot?.documentMetadata[row.documentId] + const document = provenanceDocument ?? row + for (const slot of ALL_TAG_SLOTS) { + const value = + provenanceDocument && slot.startsWith('tag') + ? provenanceDocument[slot as 'tag1' | 'tag2' | 'tag3' | 'tag4' | 'tag5' | 'tag6' | 'tag7'] + : row[slot] + if (value !== null && value !== undefined) metadata[tagMap?.get(slot) ?? slot] = value } - annotateSearchDiagnostics({ resultCount: results.length }) - const cost = baseCost - ? { - input: baseCost.input, - output: baseCost.output, - total: baseCost.total, - tokens: { prompt: tokenCount, completion: 0, total: tokenCount }, - model: embeddingModel, - pricing: baseCost.pricing, - ...(rerankerBilled && !rerankerIsBYOK - ? { - rerankerCost, - rerankerModel: input.rerankerModel, - rerankerSearchUnits: 1, - } - : {}), - } - : undefined + const rerankerScore = rerankerScores.get(row.id) return { - retrieval: retrieved.retrieval, - results, - query: input.query ?? '', - knowledgeBaseIds, - knowledgeBases: context.knowledgeBases.map((knowledgeBase) => ({ - id: knowledgeBase.id, - name: knowledgeBase.name, - })), - knowledgeBaseId: knowledgeBaseIds[0], - topK: input.topK, - totalResults: results.length, - rerankerStatus, - cost, - ...(context.workspaceId ? { workspaceId: context.workspaceId } : {}), - userId, - accessScopeKind: access.kind, - resultSecretRegistry: registry, + embeddingId: row.id, + knowledgeBaseId: row.knowledgeBaseId, + documentId: row.documentId, + documentName: document?.filename ?? null, + sourceUrl: document?.sourceUrl ?? null, + sourceModifiedAt: row.sourceModifiedAt ?? null, + connectorType: row.connectorType ?? null, + content: row.content, + chunkIndex: row.chunkIndex, + metadata, + similarity: hasQuery ? 1 - row.distance : 1, + ...(rerankerScore !== undefined ? { rerankerScore } : {}), } - }, - afterSuccess: async ({ principal, context, input, result }) => { - const actorUserId = resolvePrincipalSubjectUserId(principal) - if (context.organizationId && actorUserId) { - await measureSearchStage('activity_recording', () => - recordOrganizationSearchActivity({ - organizationId: context.organizationId, - userId: actorUserId, - surface: input.surface ?? 'other', - results: result.results, - }) - ) + }) + if (registry && provenanceSnapshot) { + for (const [documentId, document] of Object.entries(provenanceSnapshot.documentMetadata)) { + const renderedMetadata = results + .filter((result) => result.documentId === documentId) + .map((result) => ({ + documentName: result.documentName, + sourceUrl: result.sourceUrl, + metadata: result.metadata, + })) + if (renderedMetadata.length === 0) continue + if ( + !(await measureSearchStage('metadata_provenance', () => + importDurableSecretProvenance(registry, document.provenance, renderedMetadata) + )) + ) { + registry.markIncomplete('knowledge-result-provenance-unavailable') + } } - PlatformEvents.knowledgeBaseSearched({ - knowledgeBaseId: result.knowledgeBaseId, - knowledgeBaseIds: result.knowledgeBaseIds, - documentIds: [...new Set(result.results.map((item) => item.documentId))], - connectorTypes: [ - ...new Set( - result.results.flatMap((item) => (item.connectorType ? [item.connectorType] : [])) - ), - ], - resultsCount: result.totalResults, - workspaceId: context.workspaceId, - actorUserId: resolvePrincipalSubjectUserId(principal) ?? undefined, - principalKind: principal.kind, - delegatedServiceId: - principal.kind === 'delegated' || principal.kind === 'organization_delegated' - ? principal.serviceId - : undefined, - accessScopeKind: result.accessScopeKind, - surface: input.surface, - }) - }, + } + annotateSearchDiagnostics({ resultCount: results.length }) + const cost = baseCost + ? { + input: baseCost.input, + output: baseCost.output, + total: baseCost.total, + tokens: { prompt: tokenCount, completion: 0, total: tokenCount }, + model: embeddingModel, + pricing: baseCost.pricing, + ...(rerankerBilled && !rerankerIsBYOK + ? { + rerankerCost, + rerankerModel: input.rerankerModel, + rerankerSearchUnits: 1, + } + : {}), + } + : undefined + return { + retrieval: retrieved.retrieval, + results, + query: input.query ?? '', + knowledgeBaseIds, + knowledgeBases: context.knowledgeBases.map((knowledgeBase) => ({ + id: knowledgeBase.id, + name: knowledgeBase.name, + })), + knowledgeBaseId: knowledgeBaseIds[0], + topK: input.topK, + totalResults: results.length, + rerankerStatus, + cost, + ...(context.workspaceId ? { workspaceId: context.workspaceId } : {}), + userId, + accessScopeKind: access.kind, + resultSecretRegistry: registry, + } +} + +/** What follows a completed search on every surface: the organization's activity record and the platform event. */ +export async function afterKnowledgeSearch({ + principal, + context, + input, + result, +}: { + principal: Principal + context: KnowledgeResourceContext + input: Pick + result: SearchKnowledgeResult +}): Promise { + const actorUserId = resolvePrincipalSubjectUserId(principal) + if (context.organizationId && actorUserId) { + await measureSearchStage('activity_recording', () => + recordOrganizationSearchActivity({ + organizationId: context.organizationId, + userId: actorUserId, + surface: input.surface ?? 'other', + results: result.results, + }) + ) + } + PlatformEvents.knowledgeBaseSearched({ + knowledgeBaseId: result.knowledgeBaseId, + knowledgeBaseIds: result.knowledgeBaseIds, + documentIds: [...new Set(result.results.map((item) => item.documentId))], + connectorTypes: [ + ...new Set( + result.results.flatMap((item) => (item.connectorType ? [item.connectorType] : [])) + ), + ], + resultsCount: result.totalResults, + workspaceId: context.workspaceId, + actorUserId: resolvePrincipalSubjectUserId(principal) ?? undefined, + principalKind: principal.kind, + delegatedServiceId: + principal.kind === 'delegated' || principal.kind === 'organization_delegated' + ? principal.serviceId + : undefined, + accessScopeKind: result.accessScopeKind, + surface: input.surface, + }) +} + +/** Search over explicitly named bases: resolves and authorizes them, then runs the shared search. */ +const searchKnowledgeUseCase = defineAuthorizedKnowledgeUseCase({ + operation: knowledgeOperations.search, + resolveContext: ({ principal, input }: { principal: Principal; input: SearchKnowledgeInput }) => + measureSearchStage('knowledge_context', () => resolveKnowledgeSearchContext(input, principal)), + execute: ({ principal, input, context }) => runKnowledgeSearch({ principal, input, context }), + afterSuccess: ({ principal, context, input, result }) => + afterKnowledgeSearch({ principal, context, input, result }), }) export const searchKnowledge = instrumentSearchUseCase( diff --git a/apps/sim/lib/knowledge/application/workspace-search.activity.test.ts b/apps/sim/lib/knowledge/application/workspace-search.activity.test.ts index 301766ff25f..dc933d694fa 100644 --- a/apps/sim/lib/knowledge/application/workspace-search.activity.test.ts +++ b/apps/sim/lib/knowledge/application/workspace-search.activity.test.ts @@ -4,6 +4,7 @@ import { queueTableRows, resetDbChainMock } from '@sim/testing' import { beforeEach, describe, expect, it, vi } from 'vitest' const mocks = vi.hoisted(() => ({ + afterSearch: vi.fn(async () => undefined), context: vi.fn(), policy: vi.fn(), findIndex: vi.fn(), @@ -33,7 +34,18 @@ vi.mock('@/lib/knowledge/search/activity', () => ({ recordOrganizationSearchActivity: mocks.activity, })) vi.mock('@/lib/knowledge/application/search', () => ({ - searchKnowledge: { execute: mocks.search }, + runKnowledgeSearch: mocks.search, + buildKnowledgeSearchContext: ( + _principal: unknown, + context: unknown, + knowledgeBases: unknown + ) => ({ + ...(context as object), + knowledgeBases, + access: {}, + }), + validateKnowledgeSearchInput: () => undefined, + afterKnowledgeSearch: mocks.afterSearch, })) import { @@ -52,7 +64,11 @@ beforeEach(() => { mocks.findIndex.mockResolvedValue(null) mocks.available.mockResolvedValue(undefined) mocks.activity.mockResolvedValue(undefined) - mocks.search.mockResolvedValue({ results: [], knowledgeBases: [{ id: 'index' }] }) + mocks.search.mockResolvedValue({ + results: [], + knowledgeBases: [{ id: 'index' }], + knowledgeBaseId: 'index', + }) }) describe.each([ @@ -82,10 +98,14 @@ describe.each([ mocks.findIndex.mockResolvedValueOnce({ id: 'index' }) await operation.execute({ principal, input }) expect(mocks.search).toHaveBeenCalledOnce() - expect(mocks.search).toHaveBeenCalledWith({ - principal, - input: expect.objectContaining({ knowledgeBaseIds: ['index'], surface: 'slack' }), - }) + expect(mocks.search).toHaveBeenCalledWith( + expect.objectContaining({ + principal, + input: expect.objectContaining({ knowledgeBaseIds: ['index'], surface: 'slack' }), + }) + ) + /** A searched index is followed up once, by the shared hook; the empty path records nothing here. */ + expect(mocks.afterSearch).toHaveBeenCalledOnce() expect(mocks.activity).not.toHaveBeenCalled() }) diff --git a/apps/sim/lib/knowledge/application/workspace-search.test.ts b/apps/sim/lib/knowledge/application/workspace-search.test.ts index 53f26f7404d..1d632f6535c 100644 --- a/apps/sim/lib/knowledge/application/workspace-search.test.ts +++ b/apps/sim/lib/knowledge/application/workspace-search.test.ts @@ -9,6 +9,7 @@ import { import { beforeEach, describe, expect, it, vi } from 'vitest' const mocks = vi.hoisted(() => ({ + afterSearch: vi.fn(async () => undefined), resolveWorkspace: vi.fn(), permission: vi.fn(), search: vi.fn(), @@ -21,7 +22,18 @@ vi.mock('@/lib/knowledge/application/contexts', () => ({ resolveKnowledgeWorkspaceContext: mocks.resolveWorkspace, })) vi.mock('@/lib/knowledge/application/search', () => ({ - searchKnowledge: { execute: mocks.search }, + runKnowledgeSearch: mocks.search, + buildKnowledgeSearchContext: ( + _principal: unknown, + context: unknown, + knowledgeBases: unknown + ) => ({ + ...(context as object), + knowledgeBases, + access: {}, + }), + validateKnowledgeSearchInput: () => undefined, + afterKnowledgeSearch: mocks.afterSearch, })) import { searchWorkspaceKnowledge } from '@/lib/knowledge/application/workspace-search' @@ -47,10 +59,17 @@ describe('canonical workspace search', () => { it('authorizes the person before selecting the canonical active index and passes the same principal and filters', async () => { queueTableRows(schemaMock.knowledgeBase, [{ id: 'index' }]) await searchWorkspaceKnowledge.execute({ principal, input }) - expect(mocks.search).toHaveBeenCalledWith({ - principal, - input: { ...input, knowledgeBaseIds: ['index'] }, - }) + /** The search runs under the context this use case resolved; the index is its one base. */ + expect(mocks.search).toHaveBeenCalledWith( + expect.objectContaining({ + principal, + input: { ...input, knowledgeBaseIds: ['index'] }, + context: expect.objectContaining({ + workspaceId: 'workspace', + knowledgeBases: [expect.objectContaining({ id: 'index' })], + }), + }) + ) expect( hasMockCondition( dbChainMockFns.where.mock.calls[0][0], diff --git a/apps/sim/lib/knowledge/application/workspace-search.ts b/apps/sim/lib/knowledge/application/workspace-search.ts index 81924891d5a..d69fe6bca21 100644 --- a/apps/sim/lib/knowledge/application/workspace-search.ts +++ b/apps/sim/lib/knowledge/application/workspace-search.ts @@ -1,15 +1,26 @@ +import type { Principal } from '@sim/auth/principal' import { resolvePrincipalSubjectUserId } from '@sim/auth/principal' +import { OrchestrationError } from '@/lib/core/orchestration/types' import { type ResourceOwner, resourceScopeFromOwner } from '@/lib/core/resource-scope' import { requireOrganizationSearchAvailable } from '@/lib/knowledge/access/availability' import { defineAuthorizedKnowledgeUseCase } from '@/lib/knowledge/application/authorized-knowledge-use-case' import { + type KnowledgeResourceContext, resolveKnowledgeOrganizationContext, resolveKnowledgeOwnerContext, resolveKnowledgeWorkspaceContext, } from '@/lib/knowledge/application/contexts' import { knowledgeOperations } from '@/lib/knowledge/application/operations' -import { type SearchKnowledgeInput, searchKnowledge } from '@/lib/knowledge/application/search' +import { + afterKnowledgeSearch, + buildKnowledgeSearchContext, + runKnowledgeSearch, + type SearchKnowledgeInput, + type SearchKnowledgeResult, + validateKnowledgeSearchInput, +} from '@/lib/knowledge/application/search' import { instrumentSearchUseCase } from '@/lib/knowledge/application/search-diagnostics' +import type { ActiveKnowledgeBaseReference } from '@/lib/knowledge/knowledge-base-reference' import { recordOrganizationSearchActivity } from '@/lib/knowledge/search/activity' import { measureSearchStage } from '@/lib/knowledge/search/diagnostics' import { findSearchIndex, findWorkspaceSearchIndex } from '@/lib/knowledge/search/search-index' @@ -21,134 +32,139 @@ export type SearchWorkspaceKnowledgeInput = Omit< workspaceId: string } -/** Search and Assistant share the workspace's canonical Enterprise Search index. */ -const searchWorkspaceKnowledgeUseCase = defineAuthorizedKnowledgeUseCase({ - operation: knowledgeOperations.search, - resolveContext: ({ input }: { input: SearchWorkspaceKnowledgeInput }) => - measureSearchStage('scope_resolution', () => resolveKnowledgeWorkspaceContext(input)), - async execute({ principal, input, context }) { - input.signal?.throwIfAborted() - const index = await measureSearchStage('index_resolution', () => - findWorkspaceSearchIndex(context.workspaceId) - ) - if (!index) - return { - results: [], - query: input.query ?? '', - knowledgeBases: [], - retrieval: { status: 'complete' as const, timedOutLegs: [] }, - } - return searchKnowledge.execute({ - principal, - input: { ...input, workspaceId: context.workspaceId, knowledgeBaseIds: [index.id] }, - }) - }, -}) - -export const searchWorkspaceKnowledge = instrumentSearchUseCase( - 'workspace_application', - searchWorkspaceKnowledgeUseCase -) - export type SearchOrganizationKnowledgeInput = Omit< SearchWorkspaceKnowledgeInput, 'workspaceId' > & { organizationId: string } -/** Organization Search and Assistant resolve the same index and provider ACLs. */ -const searchOrganizationKnowledgeUseCase = defineAuthorizedKnowledgeUseCase({ - operation: knowledgeOperations.search, - resolveContext: ({ input }: { input: SearchOrganizationKnowledgeInput }) => - measureSearchStage('scope_resolution', () => resolveKnowledgeOrganizationContext(input)), - async execute({ principal, input, context }) { +export type SearchScopedKnowledgeInput = Omit< + SearchKnowledgeInput, + 'knowledgeBaseIds' | 'workspaceId' | 'organizationId' +> & + ResourceOwner + +/** What an owner without an index answers: nothing, completely. */ +interface SearchWithoutIndex { + results: [] + query: string + knowledgeBases: [] + retrieval: { status: 'complete'; timedOutLegs: [] } +} + +type ScopedSearchResult = SearchKnowledgeResult | SearchWithoutIndex + +/** Whether an index was searched, which is what the follow-up to a search is for. */ +function searchedAnIndex(result: ScopedSearchResult): result is SearchKnowledgeResult { + return 'knowledgeBaseId' in result +} + +/** + * An owner without an index answers empty. An organization still has to be allowed to search, + * and its empty search is recorded like any other, so the activity view shows the attempt. + */ +async function searchWithoutIndex( + principal: Principal, + context: KnowledgeResourceContext, + input: Pick +): Promise { + if (context.organizationId) { + await requireOrganizationSearchAvailable(context.organizationId) input.signal?.throwIfAborted() - const index = await measureSearchStage('index_resolution', () => - findSearchIndex({ - kind: 'organization', + const userId = resolvePrincipalSubjectUserId(principal) + if (userId) + await recordOrganizationSearchActivity({ organizationId: context.organizationId, + userId, + surface: input.surface ?? 'other', + results: [], }) - ) - if (!index) { - if (context.organizationId) { - await requireOrganizationSearchAvailable(context.organizationId) - input.signal?.throwIfAborted() - const userId = resolvePrincipalSubjectUserId(principal) - if (userId) - await recordOrganizationSearchActivity({ - organizationId: context.organizationId, - userId, - surface: input.surface ?? 'other', - results: [], - }) + } + return { + results: [], + query: input.query ?? '', + knowledgeBases: [], + retrieval: { status: 'complete', timedOutLegs: [] }, + } +} + +type ScopedSearchInput = Omit + +/** + * A search surface that resolves an owner, finds the owner's index and searches it. The owner is + * resolved and authorized once, here; the search runs under that context, and nothing about the + * index is read twice. Surfaces differ only in how they name the owner and find the index. + */ +function defineScopedSearchUseCase< + I extends Pick, +>(surface: { + resolveContext: (input: I) => Promise + findIndex: (context: KnowledgeResourceContext) => Promise + searchInput: (input: I, context: KnowledgeResourceContext) => ScopedSearchInput +}) { + return defineAuthorizedKnowledgeUseCase({ + operation: knowledgeOperations.search, + resolveContext: ({ input }: { input: I }) => + measureSearchStage('scope_resolution', () => surface.resolveContext(input)), + async execute({ principal, input, context }): Promise { + input.signal?.throwIfAborted() + const index = await measureSearchStage('index_resolution', () => surface.findIndex(context)) + if (!index) return searchWithoutIndex(principal, context, input) + const searchInput: SearchKnowledgeInput = { + ...surface.searchInput(input, context), + knowledgeBaseIds: [index.id], } - return { - results: [], - query: input.query ?? '', - knowledgeBases: [], - retrieval: { status: 'complete' as const, timedOutLegs: [] }, + /** An owner the request asserts is the one that was resolved, or the request names none. */ + if ( + (searchInput.organizationId && searchInput.organizationId !== context.organizationId) || + (searchInput.workspaceId && searchInput.workspaceId !== context.workspaceId) + ) { + throw new OrchestrationError('not_found', 'Knowledge base not found') } - } - return searchKnowledge.execute({ - principal, - input: { ...input, knowledgeBaseIds: [index.id] }, - }) - }, -}) + validateKnowledgeSearchInput(searchInput) + return runKnowledgeSearch({ + principal, + input: searchInput, + context: buildKnowledgeSearchContext(principal, context, [index], searchInput), + }) + }, + afterSuccess: ({ principal, context, input, result }) => + searchedAnIndex(result) + ? afterKnowledgeSearch({ principal, context, input, result }) + : undefined, + }) +} +/** Search and Assistant share the workspace's canonical Enterprise Search index. */ +export const searchWorkspaceKnowledge = instrumentSearchUseCase( + 'workspace_application', + defineScopedSearchUseCase({ + resolveContext: (input) => resolveKnowledgeWorkspaceContext(input), + findIndex: (context) => findWorkspaceSearchIndex(context.workspaceId!), + searchInput: (input, context) => ({ ...input, workspaceId: context.workspaceId }), + }) +) + +/** Organization Search and Assistant resolve the same index and provider ACLs. */ export const searchOrganizationKnowledge = instrumentSearchUseCase( 'organization_application', - searchOrganizationKnowledgeUseCase + defineScopedSearchUseCase({ + resolveContext: (input) => resolveKnowledgeOrganizationContext(input), + findIndex: (context) => + findSearchIndex({ kind: 'organization', organizationId: context.organizationId! }), + searchInput: (input) => input, + }) ) -export type SearchScopedKnowledgeInput = Omit< - SearchKnowledgeInput, - 'knowledgeBaseIds' | 'workspaceId' | 'organizationId' -> & - ResourceOwner - /** The routed owner selects the index; current membership and provider ACLs select its documents. */ -const searchScopedKnowledgeUseCase = defineAuthorizedKnowledgeUseCase({ - operation: knowledgeOperations.search, - resolveContext: ({ input }: { input: SearchScopedKnowledgeInput }) => - measureSearchStage('scope_resolution', () => resolveKnowledgeOwnerContext(input)), - async execute({ principal, input, context }) { - input.signal?.throwIfAborted() - const index = await measureSearchStage('index_resolution', () => - findSearchIndex(resourceScopeFromOwner(context)) - ) - if (!index) { - if (context.organizationId) { - await requireOrganizationSearchAvailable(context.organizationId) - input.signal?.throwIfAborted() - const userId = resolvePrincipalSubjectUserId(principal) - if (userId) - await recordOrganizationSearchActivity({ - organizationId: context.organizationId, - userId, - surface: input.surface ?? 'other', - results: [], - }) - } - return { - results: [], - query: input.query ?? '', - knowledgeBases: [], - retrieval: { status: 'complete' as const, timedOutLegs: [] }, - } - } - return searchKnowledge.execute({ - principal, - input: { - ...input, - workspaceId: input.workspaceId ?? undefined, - organizationId: input.organizationId ?? undefined, - knowledgeBaseIds: [index.id], - }, - }) - }, -}) - export const searchScopedKnowledge = instrumentSearchUseCase( 'scoped_application', - searchScopedKnowledgeUseCase + defineScopedSearchUseCase({ + resolveContext: (input) => resolveKnowledgeOwnerContext(input), + findIndex: (context) => findSearchIndex(resourceScopeFromOwner(context)), + searchInput: (input) => ({ + ...input, + workspaceId: input.workspaceId ?? undefined, + organizationId: input.organizationId ?? undefined, + }), + }) ) diff --git a/apps/sim/lib/knowledge/documents/pdf-ocr-chunking.test.ts b/apps/sim/lib/knowledge/documents/pdf-ocr-chunking.test.ts index 67ab2b4a992..a24c672550b 100644 --- a/apps/sim/lib/knowledge/documents/pdf-ocr-chunking.test.ts +++ b/apps/sim/lib/knowledge/documents/pdf-ocr-chunking.test.ts @@ -2,7 +2,7 @@ * @vitest-environment node */ import { PDFDocument, StandardFonts } from 'pdf-lib' -import { afterEach, describe, expect, it, vi } from 'vitest' +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest' import { PermanentDocumentProcessingError } from '@/lib/knowledge/documents/document-processing-error' import type { OcrRequestPolicy } from '@/lib/knowledge/documents/ocr-request-policy' import { buildLargestFittingPdfChunk } from '@/lib/knowledge/documents/pdf-ocr-chunking' @@ -27,7 +27,17 @@ function policy(overrides: Partial = {}): OcrRequestPolicy { } describe('buildLargestFittingPdfChunk', () => { + /** + * pdf-lib stamps the modification date into every save and the object streams are deflated, + * so two builds of the same pages a second apart can differ by a byte. The byte budgets below + * are derived from one build and applied to another; a pinned clock keeps them comparable. + */ + beforeEach(() => { + vi.useFakeTimers({ toFake: ['Date'] }) + vi.setSystemTime(new Date('2026-01-01T00:00:00Z')) + }) afterEach(() => { + vi.useRealTimers() vi.restoreAllMocks() }) diff --git a/apps/sim/lib/knowledge/knowledge-base-reference.ts b/apps/sim/lib/knowledge/knowledge-base-reference.ts new file mode 100644 index 00000000000..08a4a1d7e1f --- /dev/null +++ b/apps/sim/lib/knowledge/knowledge-base-reference.ts @@ -0,0 +1,40 @@ +import { knowledgeBase } from '@sim/db/schema' +import type { ChunkingConfig, KnowledgeBaseWithCounts } from '@/lib/knowledge/types' + +/** + * Canonical identity and configuration for application authorization and retrieval: what a + * search or a document read needs to know about a base, and nothing it counts. + */ +export type ActiveKnowledgeBaseReference = Omit< + KnowledgeBaseWithCounts, + 'tokenCount' | 'docCount' | 'connectorTypes' | 'hasPermissionScopedConnector' +> + +/** The columns a reference carries, so every lookup that returns one reads the same shape. */ +export const ACTIVE_KNOWLEDGE_BASE_REFERENCE_FIELDS = { + id: knowledgeBase.id, + userId: knowledgeBase.userId, + name: knowledgeBase.name, + isSearchIndex: knowledgeBase.isSearchIndex, + description: knowledgeBase.description, + embeddingModel: knowledgeBase.embeddingModel, + embeddingDimension: knowledgeBase.embeddingDimension, + chunkingConfig: knowledgeBase.chunkingConfig, + createdAt: knowledgeBase.createdAt, + updatedAt: knowledgeBase.updatedAt, + deletedAt: knowledgeBase.deletedAt, + workspaceId: knowledgeBase.workspaceId, + organizationId: knowledgeBase.organizationId, + folderId: knowledgeBase.folderId, +} + +type ActiveKnowledgeBaseRow = Omit & { + chunkingConfig: unknown +} + +/** The stored row as a reference; the chunking configuration is JSON the schema does not type. */ +export function toActiveKnowledgeBaseReference( + row: ActiveKnowledgeBaseRow +): ActiveKnowledgeBaseReference { + return { ...row, chunkingConfig: row.chunkingConfig as ChunkingConfig } +} diff --git a/apps/sim/lib/knowledge/search/queries.test.ts b/apps/sim/lib/knowledge/search/queries.test.ts index 6619a842727..52266d4a23c 100644 --- a/apps/sim/lib/knowledge/search/queries.test.ts +++ b/apps/sim/lib/knowledge/search/queries.test.ts @@ -64,6 +64,9 @@ const embeddingTable = { boolean1: 'boolean1', } +/** The projection-fill memo outlives a test; every case starts without one. */ +beforeEach(() => forgetProjectionFilled()) + describe('retrieval leg budgets', () => { afterEach(() => vi.restoreAllMocks()) @@ -413,6 +416,8 @@ describe('workspace-scoped vector retrieval', () => { failCandidates = undefined dbChainMockFns.execute.mockImplementation(async (query) => { const statement = render(query).sql + /** The fixtures model the page read, which only an unfilled projection makes. */ + if (statement.includes('AS unfilled')) return [{ unfilled: true }] if (statement.includes('hnsw.iterative_scan')) { if (failSettings) throw failSettings return [] @@ -530,6 +535,8 @@ describe('workspace-scoped vector retrieval', () => { const execute = dbChainMockFns.execute.getMockImplementation()! dbChainMockFns.execute.mockImplementation(async (query) => { const statement = render(query).sql + /** The fixtures model the page read, which only an unfilled projection makes. */ + if (statement.includes('AS unfilled')) return [{ unfilled: true }] if (isProbeStatement(statement)) throw new Error('canceling statement due to statement timeout', { cause: { code: '57014' }, @@ -545,17 +552,19 @@ 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 () => { + it('walks for a pool sized to the page, and scores the page on the original vectors', 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. */ + /** The hydrated page is scored on the original vector; the threshold stays on the projection. */ 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(fields).toContain(String(schemaMock.embedding.embedding)) + expect(JSON.stringify(dbChainMockFns.where.mock.calls)).toContain( + String(schemaMock.embeddingSearch.vector512) + ) expect(JSON.stringify(dbChainMockFns.leftJoin.mock.calls)).toContain('embeddingSearch') /** The source card's name, URL and connector type ride on the same read; no second pass. */ expect(fields).toContain(String(schemaMock.document.filename)) @@ -635,6 +644,7 @@ describe('workspace-scoped vector retrieval', () => { const execute = dbChainMockFns.execute.getMockImplementation()! dbChainMockFns.execute.mockImplementation(async (query) => { const statement = render(query) + if (statement.sql.includes('AS unfilled')) return [{ unfilled: true }] if (isPageStatement(statement.sql)) { queueTableRows( schemaMock.embedding, @@ -724,7 +734,7 @@ describe('workspace-scoped vector retrieval', () => { statements() .filter((query) => query.sql.includes('statement_timeout')) .map((query) => query.params[0]) - ).toEqual(['100', '40', '20']) + ).toEqual(['100', '100', '40', '20']) }) it('applies the scan settings in the deadline statement rather than one of their own', async () => { @@ -814,7 +824,8 @@ describe('workspace-scoped vector retrieval', () => { } for (const resume of release) resume() const settled = await Promise.allSettled(transactions) - expect(settled).toHaveLength(18) + /** The searches that miss the projection-fill memo together share one read; each search's own read is refused at its deadline before it starts. */ + expect(settled).toHaveLength(1) for (const transaction of settled) { expect(transaction.status).toBe('rejected') if (transaction.status === 'rejected') @@ -838,7 +849,9 @@ describe('workspace search filters before ranking', () => { modifiedAfter: '2026-09-01T00:00:00Z', }, } - beforeEach(() => resetDbChainMock()) + beforeEach(() => { + resetDbChainMock() + }) function expectScopeOnEveryQuery() { const queries = dbChainMockFns.where.mock.calls @@ -945,6 +958,8 @@ describe('hydration follows ranked candidates', () => { keywordPages.length = 0 dbChainMockFns.execute.mockImplementation(async (query) => { const statement = render(query).sql + /** The fixtures model the page read, which only an unfilled projection makes. */ + if (statement.includes('AS unfilled')) return [{ unfilled: true }] if (statement.includes('AS visible')) return candidatePages.shift() ?? [] if (isPageStatement(statement)) return rerankPages.shift() ?? [] if (statement.includes('WITH matched_keyword_chunks')) return keywordPages.shift() ?? [] @@ -1155,7 +1170,8 @@ describe('hydration follows ranked candidates', () => { ) } else { const ranking = statements().find((query) => isExactRanking(query.sql))! - expect(ranking.sql).toContain('AS id FROM') + /** The identities are one nested fragment; the mock renders it into the parameters. */ + expect(JSON.stringify(ranking)).toContain('connectorId') expect(ranking.sql).not.toContain('"content"') } const rankingOrder = @@ -1299,9 +1315,10 @@ describe('permitted-document planner', () => { indexedSourceRows = [] forgetIndexedVectorSources() forgetSearchReach() - forgetProjectionFilled() dbChainMockFns.execute.mockImplementation(async (query) => { const statement = render(query).sql + /** The fixtures model the page read, which only an unfilled projection makes. */ + if (statement.includes('AS unfilled')) return [{ unfilled: true }] if (statement.includes('pg_index')) return indexedSourceRows if (isWalk(statement)) return traversedRows if (isPageStatement(statement)) return rerankRows @@ -1582,7 +1599,6 @@ describe('permitted-document planner', () => { }) 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' }], @@ -1604,7 +1620,6 @@ describe('permitted-document planner', () => { /** 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 () => { @@ -2028,7 +2043,7 @@ describe('permitted-document planner', () => { expect(getForConnectors).not.toHaveBeenCalled() }) - it('rebuilds the pages without a gated source the caller turns out not to hold', async () => { + it('rebuilds the pool without a gated source the caller turns out not to hold', async () => { queueTableRows(schemaMock.knowledgeConnector, [ { id: 'gated-src', @@ -2039,23 +2054,21 @@ describe('permitted-document planner', () => { ]) /** * The first pool is filled by the gated source alone; only a pool built without it — the - * exclusion carries the source id into the statement — reaches the accessible candidate. + * exclusion carries the source id into the walk — reaches the accessible candidate. The + * projection is filled, so the walk carries each candidate's source and no page is read. */ dbChainMockFns.execute.mockImplementation(async (query) => { const statement = render(query).sql - /** The exclusion is the only clause that negates a connector membership. */ - const rebuilt = JSON.stringify(query).includes('OR NOT (') - /** 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 (statement.includes('AS unfilled')) return [{ unfilled: false }] + const rebuilt = JSON.stringify(query).includes('/* excluded sources */') if (isWalk(statement)) - return Array.from({ length: 400 }, (_, i) => ({ - id: i === 0 ? (rebuilt ? 'b' : 'a') : `w-${i}`, - distance: 0.1, - })) + return Array.from({ length: 400 }, (_, i) => + i === 0 + ? rebuilt + ? hit('b', 'other-src') + : hit('a', 'gated-src') + : hit(`w-${i}`, rebuilt ? 'other-src' : 'gated-src') + ) return [] }) queueTableRows(schemaMock.embedding, []) @@ -2070,18 +2083,15 @@ describe('permitted-document planner', () => { }) expect(getForConnectors).toHaveBeenCalledOnce() expect(result.rows.map((row) => row.id)).toEqual(['b']) - /** 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[0])).not.toContain('/* excluded sources */') + expect(JSON.stringify(walks[1])).toContain('/* excluded sources */') 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"') + expect(statements().some((query) => isPageStatement(query.sql))).toBe(false) }) it('excludes a denied source through its documents while the projection is unfilled', async () => { - forgetProjectionFilled() queueTableRows(schemaMock.knowledgeConnector, [ { id: 'gated-src', @@ -2122,7 +2132,6 @@ describe('permitted-document planner', () => { 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 () => { @@ -2141,6 +2150,8 @@ describe('permitted-document planner', () => { */ dbChainMockFns.execute.mockImplementation(async (query) => { const statement = render(query).sql + /** The fixtures model the page read, which only an unfilled projection makes. */ + if (statement.includes('AS unfilled')) return [{ unfilled: true }] const rebuilt = JSON.stringify(query).includes('OR NOT (') if (isPageStatement(statement)) return rebuilt @@ -2233,7 +2244,6 @@ describe('filters on a resolved scope', () => { resetDbChainMock() forgetIndexedVectorSources() forgetSearchReach() - forgetProjectionFilled() probeRows = [] traversedRows = [] rerankRows = [] @@ -2241,6 +2251,8 @@ describe('filters on a resolved scope', () => { indexedSourceRows = [] dbChainMockFns.execute.mockImplementation(async (query) => { const statement = render(query).sql + /** The fixtures model the page read, which only an unfilled projection makes. */ + if (statement.includes('AS unfilled')) return [{ unfilled: true }] if (statement.includes('EXPLAIN')) return [{ 'QUERY PLAN': [{ Plan: { 'Plan Rows': 1_000_000 } }] }] if (statement.includes('pg_index')) return indexedSourceRows @@ -2319,6 +2331,8 @@ describe('filters on a resolved scope', () => { it("estimates a filter under the leg's deadline and walks when the estimate runs out of time", async () => { dbChainMockFns.execute.mockImplementation(async (query) => { const statement = render(query).sql + /** The fixtures model the page read, which only an unfilled projection makes. */ + if (statement.includes('AS unfilled')) return [{ unfilled: true }] if (statement.includes('EXPLAIN') && JSON.stringify(render(query)).includes('"type":"gte"')) throw Object.assign(new Error('canceling statement due to statement timeout'), { code: '57014', @@ -2380,7 +2394,6 @@ describe('filters on a resolved scope', () => { expect(probes()).toHaveLength(1) estimated = 1_000_000 forgetSearchReach() - forgetProjectionFilled() dbChainMockFns.execute.mockClear() await search() expect(probes()).toHaveLength(0) @@ -2436,6 +2449,7 @@ describe('filters on a resolved scope', () => { queueTableRows(schemaMock.embedding, rerankRows) dbChainMockFns.execute.mockImplementation(async (query) => { const statement = render(query).sql + /** The fixtures model the page read, which only an unfilled projection makes. */ if (statement.includes('AS unfilled')) return [{ unfilled: true }] if (isWalk(statement)) return traversedRows if (isPageStatement(statement)) return rerankRows @@ -2476,11 +2490,13 @@ describe('filters on a resolved scope', () => { .map((query) => query.params.find((param) => param === '20000' || param === '100000')) expect(scanCaps().at(-1)).toBe('20000') resetDbChainMock() + forgetProjectionFilled() queueTableRows(schemaMock.embedding, rerankRows) dbChainMockFns.execute.mockImplementation(async (query) => { const statement = render(query).sql - if (isWalk(statement)) return traversedRows - if (isPageStatement(statement)) return rerankRows + /** A filled projection: the walk carries the identities and no page is read. */ + if (statement.includes('AS unfilled')) return [{ unfilled: false }] + if (isWalk(statement)) return rerankRows return [] }) await handleVectorOnlySearch({ @@ -2490,6 +2506,7 @@ describe('filters on a resolved scope', () => { }) expect(datesDocument(statements().filter((query) => isWalk(query.sql))[0])).toBe(false) expect(scanCaps().at(-1)).toBe('100000') + expect(statements().some((query) => isPageStatement(query.sql))).toBe(false) }) it('leaves the keyword leg short when its deadline passes before the ranking is resolved', async () => { diff --git a/apps/sim/lib/knowledge/search/queries.ts b/apps/sim/lib/knowledge/search/queries.ts index 2c6b84804f1..75cf97dc86f 100644 --- a/apps/sim/lib/knowledge/search/queries.ts +++ b/apps/sim/lib/knowledge/search/queries.ts @@ -65,6 +65,7 @@ import type { StructuredFilter } from '@/lib/knowledge/types' import { embeddingCandidateDimensions, embeddingCandidateDistance, + embeddingDistance, } from '@/lib/knowledge/vector-columns' const logger = createLogger('KnowledgeSearchQueries') @@ -123,21 +124,39 @@ const projectionFilled = new LRUCache< >({ max: PROJECTION_SOURCE_ACL_TABLES.length, ttl: PROJECTION_FILLED_TTL_MS, - /** The read that misses the cache spends the leg's own budget, like every other read of the leg. */ + /** + * The read that misses the cache is the search's own, under its budget like every other read + * of the leg, and the searches that miss together share it. A read that fails is not + * remembered: it answers unfilled, the slower and safe form, and the next search reads again. + */ 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 + try { + 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 + } catch { + return undefined + } }, }) -/** Forgets whether the projection was filled, after its rows changed. */ +/** Whether every row of the projection carries its mirrored source and ACL; unknown counts as not yet. */ +async function isProjectionFilled( + projection: ProjectionSourceAclTable, + stage: SearchStage, + budget: SearchBudget | undefined +): Promise { + return (await projectionFilled.fetch(projection, { context: { budget, stage } })) ?? false +} + +/** Forgets whether the projections were filled, after their rows changed. */ export function forgetProjectionFilled(): void { projectionFilled.clear() } + /** * Beam width per iteration. A beam is the granularity of cancellation: pgvector calls * `CHECK_FOR_INTERRUPTS` only while building an index, never inside `hnswgettuple`, so neither @@ -692,6 +711,12 @@ type SearchReadCandidate = { connectorId: string | null } +/** + * The same identities read off a projection row in raw SQL: the aliases are what + * `SearchReadCandidate` deserializes, so every walk reads them from one place. + */ +const PROJECTION_CANDIDATE_COLUMNS = sql`${embeddingSearch.id} AS id, ${embeddingSearch.documentId} AS "documentId", ${embeddingSearch.connectorId} AS "connectorId"` + /** Only opaque identifiers leave candidate ranking; content stays behind the full read predicate. */ const SEARCH_READ_CANDIDATE_FIELDS = { id: embedding.id, @@ -913,29 +938,33 @@ function hydrateSearchCandidates( filters: WorkspaceSearchFilters | undefined, conditions: (SQL | undefined)[], leg: RetrievalLeg, - budget?: SearchBudget + budget?: SearchBudget, + /** Whether a condition reads the projection's stored halfvec, which only the vector leg's threshold does. */ + joinProjection = false ) { 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. + * The projection joins so a condition on its stored halfvec — the candidate threshold — can be + * tested here; the returned score is whatever the leg passes as `distance`. Both legs pass the + * original vector's cosine distance: one out-of-line read per hydrated row, the page's size, + * where scoring the whole candidate pool that way read one per candidate on every novel query. */ - return runSearchQuery(budget, `${leg}.sql`, (executor) => - executor + return runSearchQuery(budget, `${leg}.sql`, (executor) => { + const read = executor .select(getSearchResultFields(distance)) .from(embedding) .innerJoin(document, eq(embedding.documentId, document.id)) - .leftJoin(embeddingSearch, eq(embeddingSearch.id, embedding.id)) .leftJoin(knowledgeConnector, eq(knowledgeConnector.id, document.connectorId)) - .where( - and( - inArray(embedding.id, ids), - ...getVisibilityConditions(access, filters, accessCondition), - ...conditions - ) + return ( + joinProjection ? read.leftJoin(embeddingSearch, eq(embeddingSearch.id, embedding.id)) : read + ).where( + and( + inArray(embedding.id, ids), + ...getVisibilityConditions(access, filters, accessCondition), + ...conditions ) - ) + ) + }) } /** Candidates each hybrid leg retrieves before the fused list is trimmed to `topK`. */ @@ -1515,7 +1544,7 @@ async function selectSourceVectorCandidates(input: { candidateDistance: SQL candidateLimit: number budget?: SearchBudget -}): Promise> { +}): Promise { const sources = planSourceVectorCandidates({ plan: input.plan, indexedSources: await indexedVectorSources(), @@ -1531,7 +1560,7 @@ async function selectSourceVectorCandidates(input: { input.tagCondition, input.exclusion ) - type RankedChunks = Promise> + type RankedChunks = Promise> /** * Walks one source's own index, or the sliced sources together when their slice saturated. * Readability is decided on the row the walk visits — the source and ACL are mirrored there — @@ -1546,8 +1575,8 @@ async function selectSourceVectorCandidates(input: { () => withVectorScanSettings( (executor) => - executor.execute<{ id: string; distance: number }>(sql` - SELECT ${embeddingSearch.id} AS id, ${input.candidateDistance} AS distance + executor.execute(sql` + SELECT ${PROJECTION_CANDIDATE_COLUMNS}, ${input.candidateDistance} AS distance FROM ${embeddingSearch} /* on-row visibility */ WHERE ${and( base, @@ -1591,14 +1620,14 @@ async function selectSourceVectorCandidates(input: { : sql`EXISTS (SELECT 1 FROM ${document} WHERE ${and(eq(document.id, embeddingSearch.documentId), input.documentCondition)})` ) const rows = await runSearchQuery(input.budget, 'vector.source_exact', (executor) => - executor.execute<{ id: string; distance: number; saturated: boolean }>(sql` + executor.execute(sql` WITH readable_chunks AS MATERIALIZED ( - SELECT ${embeddingSearch.id} AS id, ${input.candidateDistance} AS distance + SELECT ${PROJECTION_CANDIDATE_COLUMNS}, ${input.candidateDistance} AS distance FROM ${embeddingSearch} WHERE ${readableChunks} LIMIT ${SOURCE_EXACT_CHUNK_LIMIT + 1} ) - SELECT id, distance + 0 AS distance, + SELECT id, "documentId", "connectorId", distance + 0 AS distance, (SELECT count(*) FROM readable_chunks) > ${SOURCE_EXACT_CHUNK_LIMIT} AS saturated FROM readable_chunks ORDER BY distance LIMIT ${input.candidateLimit}`) @@ -1626,11 +1655,10 @@ async function selectSourceVectorCandidates(input: { } } ) - const ranked: Array<{ id: string; distance: number }> = scored.flat() + const ranked: Array = scored.flat() return ranked .sort((a, b) => Number(a.distance) - Number(b.distance)) .slice(0, input.candidateLimit) - .map((row) => ({ id: row.id })) } /** Sources ranked at once; each holds a connection for its own statement. */ @@ -1647,7 +1675,7 @@ const SOURCE_RANKING_CONCURRENCY = 3 */ async function selectVectorResults(params: SearchParams): Promise { const queryVector = params.queryVector! - /** One score for ranking, threshold and results alike: the projection's, which stays in cache. */ + /** The walk and the candidate threshold use the projection's score, which stays in cache; the page is scored on the original vectors at hydration. */ const distance = embeddingCandidateDistance( queryVector.dimensions, queryVector.vector, @@ -1688,7 +1716,14 @@ async function selectVectorResults(params: SearchParams): Promise; limit: number; exhausted: boolean } + | { + excludedKey: string + ids: SearchReadCandidate[] + limit: number + exhausted: boolean + /** Whether the pool's rows carry their source, so a page needs no read of its own. */ + filled: boolean + } | undefined return selectAuthorizedSearchResults({ leg: 'vector', @@ -1731,23 +1766,8 @@ async function selectVectorResults(params: SearchParams): Promise - executor.execute<{ id: string }>(sql` - SELECT ${embeddingSearch.id} AS id FROM ${embeddingSearch} + executor.execute(sql` + SELECT ${PROJECTION_CANDIDATE_COLUMNS} + FROM ${embeddingSearch} WHERE ${and( inArray(embeddingSearch.knowledgeBaseId, params.knowledgeBaseIds), eq(embeddingSearch.enabled, true), @@ -1794,7 +1832,7 @@ async function selectVectorResults(params: SearchParams): Promise + let selected: SearchReadCandidate[] /** * 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. @@ -1846,10 +1884,11 @@ async function selectVectorResults(params: SearchParams): Promise - executor.execute<{ id: string }>( + executor.execute( plan ? sql` - SELECT ${embeddingSearch.id} AS id FROM ${embeddingSearch} /* on-row visibility */ + SELECT ${PROJECTION_CANDIDATE_COLUMNS} + FROM ${embeddingSearch} /* on-row visibility */ WHERE ${and( scopeOfWalk, projectionCandidateAccessCondition(embeddingSearch, params.access, plan, { filled }), @@ -1860,9 +1899,11 @@ async function selectVectorResults(params: SearchParams): Promise= MAX_VECTOR_CANDIDATES, + filled, } annotateSearchDiagnostics({ vectorCandidateCount: selected.length, @@ -1917,11 +1959,20 @@ async function selectVectorResults(params: SearchParams): Promise @@ -2271,20 +2325,12 @@ 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. - */ + /** Every candidate already matched the query where it was ranked; matching it again here would detoast one text-search vector per result. */ hydrate: (ids, authorized) => hydrateSearchCandidates( ids, authorized, - embeddingCandidateDistance( - queryVector.dimensions, - queryVector.vector, - queryVector.model - ).as('distance'), + embeddingDistance(queryVector.dimensions, queryVector.vector).as('distance'), params.filters, [inArray(embedding.knowledgeBaseId, knowledgeBaseIds), ...tagFilterConditions], 'keyword', diff --git a/apps/sim/lib/knowledge/search/search-index.ts b/apps/sim/lib/knowledge/search/search-index.ts index 80e58a77b63..25d4aea5998 100644 --- a/apps/sim/lib/knowledge/search/search-index.ts +++ b/apps/sim/lib/knowledge/search/search-index.ts @@ -3,11 +3,21 @@ import { knowledgeBase } from '@sim/db/schema' import { and, eq, isNull } from 'drizzle-orm' import type { ResourceScope } from '@/lib/core/resource-scope' import { resourceScopeCondition } from '@/lib/core/resource-scope.server' +import { + ACTIVE_KNOWLEDGE_BASE_REFERENCE_FIELDS, + type ActiveKnowledgeBaseReference, + toActiveKnowledgeBaseReference, +} from '@/lib/knowledge/knowledge-base-reference' -/** Resolves the single active Enterprise Search index owned by a workspace. */ -export async function findSearchIndex(scope: ResourceScope) { +/** + * Resolves the single active Enterprise Search index owned by a workspace or organization, as + * the reference a search runs over, so the caller that found it need not read it again. + */ +export async function findSearchIndex( + scope: ResourceScope +): Promise { const [index] = await db - .select({ id: knowledgeBase.id }) + .select(ACTIVE_KNOWLEDGE_BASE_REFERENCE_FIELDS) .from(knowledgeBase) .where( and( @@ -17,7 +27,7 @@ export async function findSearchIndex(scope: ResourceScope) { ) ) .limit(1) - return index ?? null + return index ? toActiveKnowledgeBaseReference(index) : null } export function findWorkspaceSearchIndex(workspaceId: string) { diff --git a/apps/sim/lib/knowledge/service.ts b/apps/sim/lib/knowledge/service.ts index 91352dbdef0..23159f496f1 100644 --- a/apps/sim/lib/knowledge/service.ts +++ b/apps/sim/lib/knowledge/service.ts @@ -36,6 +36,11 @@ import { isKnowledgeMemberAccessAvailable } from '@/lib/knowledge/access/availab import { knowledgeAccessCondition } from '@/lib/knowledge/access/predicate' import type { KnowledgeAccessProvider } from '@/lib/knowledge/access/types' import { mirrorsSourceAcls } from '@/lib/knowledge/connectors/access-modes' +import { + ACTIVE_KNOWLEDGE_BASE_REFERENCE_FIELDS, + type ActiveKnowledgeBaseReference, + toActiveKnowledgeBaseReference, +} from '@/lib/knowledge/knowledge-base-reference' import { type KnowledgeReadAccess, knowledgeReadAccessBatches } from '@/lib/knowledge/read-access' import type { ChunkingConfig, @@ -949,28 +954,6 @@ export async function getKnowledgeBaseNames( return new Map(rows.map((row) => [row.id, row.name])) } -export type ActiveKnowledgeBaseReference = Omit< - KnowledgeBaseWithCounts, - 'tokenCount' | 'docCount' | 'connectorTypes' | 'hasPermissionScopedConnector' -> - -const ACTIVE_KNOWLEDGE_BASE_REFERENCE_FIELDS = { - id: knowledgeBase.id, - userId: knowledgeBase.userId, - name: knowledgeBase.name, - isSearchIndex: knowledgeBase.isSearchIndex, - description: knowledgeBase.description, - embeddingModel: knowledgeBase.embeddingModel, - embeddingDimension: knowledgeBase.embeddingDimension, - chunkingConfig: knowledgeBase.chunkingConfig, - createdAt: knowledgeBase.createdAt, - updatedAt: knowledgeBase.updatedAt, - deletedAt: knowledgeBase.deletedAt, - workspaceId: knowledgeBase.workspaceId, - organizationId: knowledgeBase.organizationId, - folderId: knowledgeBase.folderId, -} - /** * Canonical identity and configuration for application authorization and retrieval. * Reading a reference never scans the base's documents to compute display counts. @@ -984,7 +967,7 @@ export async function getActiveKnowledgeBaseReference( .where(and(eq(knowledgeBase.id, knowledgeBaseId), isNull(knowledgeBase.deletedAt))) .limit(1) - return row ? { ...row, chunkingConfig: row.chunkingConfig as ChunkingConfig } : null + return row ? toActiveKnowledgeBaseReference(row) : null } /** Loads active references in one statement while preserving requested order and missing entries. */ @@ -1004,9 +987,7 @@ export async function getActiveKnowledgeBaseReferences( isNull(knowledgeBase.deletedAt) ) ) - const byId = new Map( - rows.map((row) => [row.id, { ...row, chunkingConfig: row.chunkingConfig as ChunkingConfig }]) - ) + const byId = new Map(rows.map((row) => [row.id, toActiveKnowledgeBaseReference(row)])) return knowledgeBaseIds.map((id) => byId.get(id) ?? null) } diff --git a/packages/testing/src/mocks/schema.mock.ts b/packages/testing/src/mocks/schema.mock.ts index 66a9138aef9..e2efd629801 100644 --- a/packages/testing/src/mocks/schema.mock.ts +++ b/packages/testing/src/mocks/schema.mock.ts @@ -990,6 +990,8 @@ export const schemaMock = { id: 'embeddingSearch.id', knowledgeBaseId: 'embeddingSearch.knowledgeBaseId', documentId: 'embeddingSearch.documentId', + connectorId: 'embeddingSearch.connectorId', + acl: 'embeddingSearch.acl', enabled: 'embeddingSearch.enabled', binary: 'embeddingSearch.binary', binary384: 'embeddingSearch.binary384',