From 6b0e12e872a5f75eb23492e5f4f750867b19a4f6 Mon Sep 17 00:00:00 2001 From: Waleed Latif Date: Sun, 20 Sep 2026 22:36:37 -0700 Subject: [PATCH 1/6] improvement(knowledge): score the vector page on the original vectors and pin the clock in the PDF chunk test --- .../lib/knowledge/documents/pdf-ocr-chunking.test.ts | 12 +++++++++++- apps/sim/lib/knowledge/search/queries.test.ts | 10 ++++++---- apps/sim/lib/knowledge/search/queries.ts | 8 +++++++- 3 files changed, 24 insertions(+), 6 deletions(-) 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/search/queries.test.ts b/apps/sim/lib/knowledge/search/queries.test.ts index 6619a842727..1a3e449ee04 100644 --- a/apps/sim/lib/knowledge/search/queries.test.ts +++ b/apps/sim/lib/knowledge/search/queries.test.ts @@ -545,17 +545,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)) diff --git a/apps/sim/lib/knowledge/search/queries.ts b/apps/sim/lib/knowledge/search/queries.ts index 2c6b84804f1..9b7a4231a4b 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') @@ -1943,11 +1944,16 @@ async function selectVectorResults(params: SearchParams): Promise hydrateSearchCandidates( ids, authorized, - distance.as('distance'), + embeddingDistance(queryVector.dimensions, queryVector.vector).as('distance'), params.filters, conditions, 'vector', From f328012fac389b9235ba11739c17b82904e6b03f Mon Sep 17 00:00:00 2001 From: Waleed Latif Date: Sun, 20 Sep 2026 22:44:39 -0700 Subject: [PATCH 2/6] improvement(knowledge): score both pages on the original vectors and say so where the walk is scored --- apps/sim/lib/knowledge/search/queries.ts | 19 ++++++++----------- 1 file changed, 8 insertions(+), 11 deletions(-) diff --git a/apps/sim/lib/knowledge/search/queries.ts b/apps/sim/lib/knowledge/search/queries.ts index 9b7a4231a4b..5f7e1bf60cd 100644 --- a/apps/sim/lib/knowledge/search/queries.ts +++ b/apps/sim/lib/knowledge/search/queries.ts @@ -918,9 +918,10 @@ function hydrateSearchCandidates( ) { 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 @@ -1648,7 +1649,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, @@ -2279,18 +2280,14 @@ export async function executeKeywordSearch(params: KeywordSearchParams): Promise }, /** * 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. + * would detoast one text-search vector per result. The score is the original vector's, like + * the vector leg's, so one response carries one distance scale. */ 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', From 29d4ac8712f3e56467822fd1c801c5ee8d67c63d Mon Sep 17 00:00:00 2001 From: Waleed Latif Date: Sun, 20 Sep 2026 22:59:27 -0700 Subject: [PATCH 3/6] improvement(knowledge): resolve a search's context once and carry candidate identities off the walk --- .../sim/lib/knowledge/application/contexts.ts | 7 +- apps/sim/lib/knowledge/application/search.ts | 115 ++++++++++++------ .../workspace-search.activity.test.ts | 23 +++- .../application/workspace-search.test.ts | 28 ++++- .../knowledge/application/workspace-search.ts | 112 ++++++++++++----- .../lib/knowledge/knowledge-base-reference.ts | 40 ++++++ apps/sim/lib/knowledge/search/queries.test.ts | 74 +++++++---- apps/sim/lib/knowledge/search/queries.ts | 114 +++++++++++------ apps/sim/lib/knowledge/search/search-index.ts | 18 ++- apps/sim/lib/knowledge/service.ts | 33 ++--- packages/testing/src/mocks/schema.mock.ts | 2 + 11 files changed, 392 insertions(+), 174 deletions(-) create mode 100644 apps/sim/lib/knowledge/knowledge-base-reference.ts 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..3f27539435c 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,36 @@ 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[], + signal: AbortSignal | undefined +): KnowledgeSearchContext { + const knowledgeBaseIds = knowledgeBases.map((base) => base.id) + 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 +252,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.signal) } if (!canonicalWorkspaceId) { throw new OrchestrationError('not_found', 'Knowledge base not found') @@ -245,22 +265,25 @@ 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.signal) } -const searchKnowledgeUseCase = defineAuthorizedKnowledgeUseCase({ - operation: knowledgeOperations.search, - resolveContext: ({ principal, input }: { principal: Principal; input: SearchKnowledgeInput }) => - measureSearchStage('knowledge_context', () => resolveKnowledgeSearchContext(input, principal)), - async execute({ principal, input, context }) { +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, @@ -723,8 +746,22 @@ const searchKnowledgeUseCase = defineAuthorizedKnowledgeUseCase({ accessScopeKind: access.kind, resultSecretRegistry: registry, } - }, - afterSuccess: async ({ principal, context, input, result }) => { + } +} + +/** 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', () => @@ -756,7 +793,17 @@ const searchKnowledgeUseCase = defineAuthorizedKnowledgeUseCase({ 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..6dccf80c58a 100644 --- a/apps/sim/lib/knowledge/application/workspace-search.activity.test.ts +++ b/apps/sim/lib/knowledge/application/workspace-search.activity.test.ts @@ -33,7 +33,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 ?? (async () => undefined), })) import { @@ -82,10 +93,12 @@ 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' }), + }) + ) 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..5908a1b1a4c 100644 --- a/apps/sim/lib/knowledge/application/workspace-search.test.ts +++ b/apps/sim/lib/knowledge/application/workspace-search.test.ts @@ -21,7 +21,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 ?? (async () => undefined), })) import { searchWorkspaceKnowledge } from '@/lib/knowledge/application/workspace-search' @@ -47,10 +58,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..d7f1c2a8341 100644 --- a/apps/sim/lib/knowledge/application/workspace-search.ts +++ b/apps/sim/lib/knowledge/application/workspace-search.ts @@ -1,19 +1,79 @@ +import type { Principal } from '@sim/auth/principal' import { resolvePrincipalSubjectUserId } from '@sim/auth/principal' 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' +/** A search's result when the owner has an index; an owner without one answers empty. */ +type ScopedSearchResult = + | SearchKnowledgeResult + | { + results: [] + query: string + knowledgeBases: [] + retrieval: { status: 'complete'; timedOutLegs: [] } + } + +function emptySearch(query: string | undefined): ScopedSearchResult { + return { + results: [], + query: query ?? '', + knowledgeBases: [], + retrieval: { status: 'complete' as const, timedOutLegs: [] }, + } +} + +/** + * Runs the search over the owner's index under the context this use case already resolved and + * authorized: the index is the one base, and nothing about it is read twice. + */ +function searchOwnersIndex( + principal: Principal, + input: Omit, + context: KnowledgeResourceContext, + index: ActiveKnowledgeBaseReference +): Promise { + const searchInput: SearchKnowledgeInput = { ...input, knowledgeBaseIds: [index.id] } + validateKnowledgeSearchInput(searchInput) + return runKnowledgeSearch({ + principal, + input: searchInput, + context: buildKnowledgeSearchContext(principal, context, [index], input.signal), + }) +} + +/** The shared follow-up applies to a search that ran; an empty answer recorded its own. */ +function afterScopedSearch(execution: { + principal: Principal + context: KnowledgeResourceContext + input: Pick + result: ScopedSearchResult +}): Promise | undefined { + return 'userId' in execution.result + ? afterKnowledgeSearch({ ...execution, result: execution.result }) + : undefined +} + export type SearchWorkspaceKnowledgeInput = Omit< SearchKnowledgeInput, 'knowledgeBaseIds' | 'workspaceId' @@ -31,18 +91,16 @@ const searchWorkspaceKnowledgeUseCase = defineAuthorizedKnowledgeUseCase({ 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({ + if (!index) return emptySearch(input.query) + return searchOwnersIndex( principal, - input: { ...input, workspaceId: context.workspaceId, knowledgeBaseIds: [index.id] }, - }) + { ...input, workspaceId: context.workspaceId }, + context, + index + ) }, + afterSuccess: ({ principal, context, input, result }) => + afterScopedSearch({ principal, context, input, result }), }) export const searchWorkspaceKnowledge = instrumentSearchUseCase( @@ -81,18 +139,12 @@ const searchOrganizationKnowledgeUseCase = defineAuthorizedKnowledgeUseCase({ results: [], }) } - return { - results: [], - query: input.query ?? '', - knowledgeBases: [], - retrieval: { status: 'complete' as const, timedOutLegs: [] }, - } + return emptySearch(input.query) } - return searchKnowledge.execute({ - principal, - input: { ...input, knowledgeBaseIds: [index.id] }, - }) + return searchOwnersIndex(principal, input, context, index) }, + afterSuccess: ({ principal, context, input, result }) => + afterScopedSearch({ principal, context, input, result }), }) export const searchOrganizationKnowledge = instrumentSearchUseCase( @@ -129,23 +181,21 @@ const searchScopedKnowledgeUseCase = defineAuthorizedKnowledgeUseCase({ results: [], }) } - return { - results: [], - query: input.query ?? '', - knowledgeBases: [], - retrieval: { status: 'complete' as const, timedOutLegs: [] }, - } + return emptySearch(input.query) } - return searchKnowledge.execute({ + return searchOwnersIndex( principal, - input: { + { ...input, workspaceId: input.workspaceId ?? undefined, organizationId: input.organizationId ?? undefined, - knowledgeBaseIds: [index.id], }, - }) + context, + index + ) }, + afterSuccess: ({ principal, context, input, result }) => + afterScopedSearch({ principal, context, input, result }), }) export const searchScopedKnowledge = instrumentSearchUseCase( 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 1a3e449ee04..c407f1f27ce 100644 --- a/apps/sim/lib/knowledge/search/queries.test.ts +++ b/apps/sim/lib/knowledge/search/queries.test.ts @@ -71,6 +71,7 @@ describe('retrieval leg budgets', () => { 'applies vector budget %s without shortening keyword or tag retrieval', async (vectorBudgetMs) => { resetDbChainMock() + forgetProjectionFilled() vi.spyOn(performance, 'now').mockReturnValue(1000) const remaining = SearchBudget.prototype.remaining const deadlines = new Map() @@ -405,6 +406,7 @@ describe('workspace-scoped vector retrieval', () => { beforeEach(() => { resetDbChainMock() + forgetProjectionFilled() getForConnectors.mockReset() probeRows = probe traversedRows = candidates @@ -413,6 +415,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 +534,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' }, @@ -637,6 +643,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, @@ -726,7 +733,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 () => { @@ -840,7 +847,10 @@ describe('workspace search filters before ranking', () => { modifiedAfter: '2026-09-01T00:00:00Z', }, } - beforeEach(() => resetDbChainMock()) + beforeEach(() => { + resetDbChainMock() + forgetProjectionFilled() + }) function expectScopeOnEveryQuery() { const queries = dbChainMockFns.where.mock.calls @@ -940,6 +950,7 @@ describe('hydration follows ranked candidates', () => { beforeEach(() => { resetDbChainMock() + forgetProjectionFilled() probePages.length = 0 exactPages.length = 0 candidatePages.length = 0 @@ -947,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() ?? [] @@ -1157,7 +1170,7 @@ describe('hydration follows ranked candidates', () => { ) } else { const ranking = statements().find((query) => isExactRanking(query.sql))! - expect(ranking.sql).toContain('AS id FROM') + expect(ranking.sql).toContain('AS "connectorId"') expect(ranking.sql).not.toContain('"content"') } const rankingOrder = @@ -1293,6 +1306,7 @@ describe('permitted-document planner', () => { beforeEach(() => { resetDbChainMock() + forgetProjectionFilled() probeRows = [] exactRows = [] traversedRows = [] @@ -1304,6 +1318,8 @@ describe('permitted-document planner', () => { 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 @@ -1980,6 +1996,7 @@ describe('permitted-document planner', () => { expect(statements().some((query) => isWalk(query.sql))).toBe(true) resetDbChainMock() + forgetProjectionFilled() dbChainMockFns.execute.mockImplementation(async () => []) await retrieveKnowledgeSearch({ ...search, @@ -2030,7 +2047,8 @@ 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 () => { + forgetProjectionFilled() queueTableRows(schemaMock.knowledgeConnector, [ { id: 'gated-src', @@ -2041,23 +2059,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, []) @@ -2072,14 +2088,12 @@ 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('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) + forgetProjectionFilled() }) it('excludes a denied source through its documents while the projection is unfilled', async () => { @@ -2143,6 +2157,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,6 +2249,7 @@ describe('filters on a resolved scope', () => { beforeEach(() => { resetDbChainMock() + forgetProjectionFilled() forgetIndexedVectorSources() forgetSearchReach() forgetProjectionFilled() @@ -2243,6 +2260,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 @@ -2321,6 +2340,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', @@ -2438,6 +2459,8 @@ 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 (statement.includes('AS unfilled')) return [{ unfilled: true }] if (isWalk(statement)) return traversedRows if (isPageStatement(statement)) return rerankRows @@ -2478,11 +2501,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({ @@ -2492,6 +2517,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 5f7e1bf60cd..e2f70c6f6d0 100644 --- a/apps/sim/lib/knowledge/search/queries.ts +++ b/apps/sim/lib/knowledge/search/queries.ts @@ -117,23 +117,38 @@ const PROJECTION_FILLED_TTL_MS = 60_000 * Whether the ranking projection still holds rows the backfill has not filled. Read off the * unfilled-rows index in microseconds and remembered briefly: the answer only ever changes once. */ -const projectionFilled = new LRUCache< - ProjectionSourceAclTable, - boolean, - { budget: SearchBudget | undefined; stage: SearchStage } ->({ +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. */ - fetchMethod: async (projection, _stale, { context }) => { - const table = projection === 'embedding_search' ? embeddingSearch : embeddingKeywordTin - const [row] = await runSearchQuery(context.budget, context.stage, (executor) => +}) + +/** + * Whether the projection's rows all carry their mirrored columns. The read that misses the cache + * is the search's own, under its budget like every other read of the leg, so one search's + * deadline never fails another's; it is one index page, once a minute. A read that fails for a + * reason other than the deadline is answered as unfilled, the slower and safe form. + */ +async function isProjectionFilled( + projection: ProjectionSourceAclTable, + stage: SearchStage, + budget: SearchBudget | undefined +): Promise { + const cached = projectionFilled.get(projection) + if (cached !== undefined) return cached + const table = projection === 'embedding_search' ? embeddingSearch : embeddingKeywordTin + try { + const [row] = await runSearchQuery(budget, stage, (executor) => executor.execute<{ unfilled: boolean }>(sql` SELECT EXISTS (SELECT 1 FROM ${table} WHERE ${table.acl} IS NULL) AS unfilled`) ) - return !row?.unfilled - }, -}) + const filled = !row?.unfilled + projectionFilled.set(projection, filled) + return filled + } catch (error) { + if (budget?.isTimeout(error)) throw error + return false + } +} /** Forgets whether the projection was filled, after its rows changed. */ export function forgetProjectionFilled(): void { @@ -1517,7 +1532,7 @@ async function selectSourceVectorCandidates(input: { candidateDistance: SQL candidateLimit: number budget?: SearchBudget -}): Promise> { +}): Promise { const sources = planSourceVectorCandidates({ plan: input.plan, indexedSources: await indexedVectorSources(), @@ -1533,7 +1548,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 — @@ -1548,8 +1563,9 @@ 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 ${embeddingSearch.id} AS id, ${embeddingSearch.documentId} AS "documentId", ${embeddingSearch.connectorId} AS "connectorId", + ${input.candidateDistance} AS distance FROM ${embeddingSearch} /* on-row visibility */ WHERE ${and( base, @@ -1593,14 +1609,15 @@ 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 ${embeddingSearch.id} AS id, ${embeddingSearch.documentId} AS "documentId", ${embeddingSearch.connectorId} AS "connectorId", + ${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}`) @@ -1628,11 +1645,11 @@ 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 })) + .map(({ id, documentId, connectorId }) => ({ id, documentId, connectorId })) } /** Sources ranked at once; each holds a connection for its own statement. */ @@ -1690,7 +1707,7 @@ async function selectVectorResults(params: SearchParams): Promise; limit: number; exhausted: boolean } + | { excludedKey: string; ids: SearchReadCandidate[]; limit: number; exhausted: boolean } | undefined return selectAuthorizedSearchResults({ leg: 'vector', @@ -1733,12 +1750,14 @@ async function selectVectorResults(params: SearchParams): Promise - executor.execute<{ id: string }>(sql` - SELECT ${embeddingSearch.id} AS id FROM ${embeddingSearch} + executor.execute(sql` + SELECT ${embeddingSearch.id} AS id, ${embeddingSearch.documentId} AS "documentId", ${embeddingSearch.connectorId} AS "connectorId" + FROM ${embeddingSearch} WHERE ${and( inArray(embeddingSearch.knowledgeBaseId, params.knowledgeBaseIds), eq(embeddingSearch.enabled, true), @@ -1796,7 +1816,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. @@ -1848,10 +1868,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 ${embeddingSearch.id} AS id, ${embeddingSearch.documentId} AS "documentId", ${embeddingSearch.connectorId} AS "connectorId" + FROM ${embeddingSearch} /* on-row visibility */ WHERE ${and( scopeOfWalk, projectionCandidateAccessCondition(embeddingSearch, params.access, plan, { filled }), @@ -1862,9 +1883,11 @@ async function selectVectorResults(params: SearchParams): Promise 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', From f138017ec6280ba02854cdafe6190a6b7fd61d0d Mon Sep 17 00:00:00 2001 From: Waleed Latif Date: Sun, 20 Sep 2026 23:12:34 -0700 Subject: [PATCH 4/6] improvement(knowledge): one scoped search shape, page and pool cleanups, admission ahead of the embedding beside the scope reads --- .../lib/billing/core/billing-attribution.ts | 15 +- apps/sim/lib/knowledge/access/availability.ts | 15 +- apps/sim/lib/knowledge/application/search.ts | 908 +++++++++--------- .../workspace-search.activity.test.ts | 11 +- .../application/workspace-search.test.ts | 3 +- .../knowledge/application/workspace-search.ts | 262 +++-- apps/sim/lib/knowledge/search/queries.test.ts | 27 +- apps/sim/lib/knowledge/search/queries.ts | 167 ++-- 8 files changed, 690 insertions(+), 718 deletions(-) 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/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/search.ts b/apps/sim/lib/knowledge/application/search.ts index 3f27539435c..b68ba56076d 100644 --- a/apps/sim/lib/knowledge/application/search.ts +++ b/apps/sim/lib/knowledge/application/search.ts @@ -201,9 +201,10 @@ export function buildKnowledgeSearchContext( principal: Principal, context: KnowledgeResourceContext, knowledgeBases: ActiveKnowledgeBaseReference[], - signal: AbortSignal | undefined + input: Pick ): KnowledgeSearchContext { const knowledgeBaseIds = knowledgeBases.map((base) => base.id) + const signal = input.signal return { ...context, knowledgeBases, @@ -211,7 +212,7 @@ export function buildKnowledgeSearchContext( principal, context.organizationId ? { ...context, knowledgeBaseIds, signal } - : { workspaceId: context.workspaceId!, knowledgeBaseIds, signal } + : { workspaceId: context.workspaceId, knowledgeBaseIds, signal } ), } } @@ -257,7 +258,7 @@ async function resolveKnowledgeSearchContext( const context = await resolveKnowledgeOrganizationContext({ organizationId: canonicalOrganizationId, }) - return buildKnowledgeSearchContext(principal, context, resolved, input.signal) + return buildKnowledgeSearchContext(principal, context, resolved, input) } if (!canonicalWorkspaceId) { throw new OrchestrationError('not_found', 'Knowledge base not found') @@ -265,7 +266,7 @@ async function resolveKnowledgeSearchContext( const workspaceContext = await resolveKnowledgeWorkspaceContext({ workspaceId: canonicalWorkspaceId, }) - return buildKnowledgeSearchContext(principal, workspaceContext, resolved, input.signal) + return buildKnowledgeSearchContext(principal, workspaceContext, resolved, input) } export interface KnowledgeSearchExecution { @@ -283,134 +284,115 @@ export async function runKnowledgeSearch({ 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' - ) - } - const userId = resolveKnowledgeAttributedUserId(principal, context) - const shouldMeter = !( - input.skipUsageBilling && - principal.kind === 'delegated' && - principal.serviceId === 'executor' + 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!) ) - } - } - 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) + : await measureSearchStage('billing_attribution', () => + resolveKnowledgeBillingAttribution(principal, context) + ) + : undefined + if (shouldMeter && billingAttribution) { + const usage = await measureSearchStage('usage_admission', () => + checkSearchUsageLimits(billingAttribution) ) - structuredFilters = built.structuredFilters - definitionsByKnowledgeBase = built.definitionsByKnowledgeBase + if (usage.isExceeded) { + throw new KnowledgeUsageLimitExceededError( + usage.message || 'Usage limit exceeded. Please upgrade your plan to continue.' + ) + } } + return billingAttribution + } - /** - * 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 }, - ]) + 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) ) - 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 }) - ) - : undefined - const resultSecretRegistry = preparedRegistry ?? input.resultSecretRegistry - input.signal?.throwIfAborted() - const [access, searchDefaults, billingAttribution, tagDefinitions] = await Promise.all([ - measureSearchStage('access_scope', () => context.access.get()), - measureSearchStage('defaults', () => - resolveKnowledgeSearchDefaults({ - workspaceId: context.workspaceId, - organizationId: context.organizationId, + structuredFilters = built.structuredFilters + definitionsByKnowledgeBase = built.definitionsByKnowledgeBase + } - /** The signed-in person, if any; never the billing owner or a key's creator. */ - userId: resolvePrincipalSubjectUserId(principal) ?? undefined, - requestedMode: input.searchMode, - }) - ), - admit(), - /** The tag names the results are labelled with depend on the bases alone. */ - filters.length === 0 - ? measureSearchStage('tag_definitions', () => - getDocumentTagDefinitionsByKnowledgeBaseIds(knowledgeBaseIds) - ) - : Promise.resolve(definitionsByKnowledgeBase), + /** + * 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 }, ]) - definitionsByKnowledgeBase = tagDefinitions + ) + 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 }) + ) + : undefined + const resultSecretRegistry = preparedRegistry ?? input.resultSecretRegistry + input.signal?.throwIfAborted() + /** The embedding is requested the moment admission passes, beside the scope and defaults reads. */ + const admittedEmbedding = async () => { + const billingAttribution = await admit() input.signal?.throwIfAborted() const queryEmbedding = hasQuery ? await measureSearchStage('embedding', () => @@ -424,329 +406,351 @@ export async function runKnowledgeSearch({ ) ) : null - input.signal?.throwIfAborted() - annotateSearchDiagnostics({ - accessScopeKind: access.kind, + return { billingAttribution, queryEmbedding } + } + const [access, searchDefaults, admitted, tagDefinitions, rerankerCredential] = await Promise.all([ + measureSearchStage('access_scope', () => context.access.get()), + measureSearchStage('defaults', () => + resolveKnowledgeSearchDefaults({ + workspaceId: context.workspaceId, + organizationId: context.organizationId, + + /** The signed-in person, if any; never the billing owner or a key's creator. */ + userId: resolvePrincipalSubjectUserId(principal) ?? undefined, + requestedMode: input.searchMode, + }) + ), + admittedEmbedding(), + /** The tag names the results are labelled with depend on the bases alone. */ + filters.length === 0 + ? measureSearchStage('tag_definitions', () => + 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, + ]) + const { billingAttribution, queryEmbedding } = admitted + definitionsByKnowledgeBase = tagDefinitions + 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) + ) + : 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 } : {}), + } + }) + 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') + } } } + 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. */ @@ -761,39 +765,37 @@ export async function afterKnowledgeSearch({ 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, - }) + 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. */ 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 6dccf80c58a..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(), @@ -44,7 +45,7 @@ vi.mock('@/lib/knowledge/application/search', () => ({ access: {}, }), validateKnowledgeSearchInput: () => undefined, - afterKnowledgeSearch: mocks.afterSearch ?? (async () => undefined), + afterKnowledgeSearch: mocks.afterSearch, })) import { @@ -63,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([ @@ -99,6 +104,8 @@ describe.each([ 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 5908a1b1a4c..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(), @@ -32,7 +33,7 @@ vi.mock('@/lib/knowledge/application/search', () => ({ access: {}, }), validateKnowledgeSearchInput: () => undefined, - afterKnowledgeSearch: mocks.afterSearch ?? (async () => undefined), + afterKnowledgeSearch: mocks.afterSearch, })) import { searchWorkspaceKnowledge } from '@/lib/knowledge/application/workspace-search' diff --git a/apps/sim/lib/knowledge/application/workspace-search.ts b/apps/sim/lib/knowledge/application/workspace-search.ts index d7f1c2a8341..99ad8878978 100644 --- a/apps/sim/lib/knowledge/application/workspace-search.ts +++ b/apps/sim/lib/knowledge/application/workspace-search.ts @@ -24,181 +24,139 @@ import { recordOrganizationSearchActivity } from '@/lib/knowledge/search/activit import { measureSearchStage } from '@/lib/knowledge/search/diagnostics' import { findSearchIndex, findWorkspaceSearchIndex } from '@/lib/knowledge/search/search-index' -/** A search's result when the owner has an index; an owner without one answers empty. */ -type ScopedSearchResult = - | SearchKnowledgeResult - | { - results: [] - query: string - knowledgeBases: [] - retrieval: { status: 'complete'; timedOutLegs: [] } - } +export type SearchWorkspaceKnowledgeInput = Omit< + SearchKnowledgeInput, + 'knowledgeBaseIds' | 'workspaceId' +> & { + workspaceId: string +} -function emptySearch(query: string | undefined): ScopedSearchResult { - return { - results: [], - query: query ?? '', - knowledgeBases: [], - retrieval: { status: 'complete' as const, timedOutLegs: [] }, - } +export type SearchOrganizationKnowledgeInput = Omit< + SearchWorkspaceKnowledgeInput, + 'workspaceId' +> & { organizationId: string } + +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 } /** - * Runs the search over the owner's index under the context this use case already resolved and - * authorized: the index is the one base, and nothing about it is read twice. + * 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. */ -function searchOwnersIndex( +async function searchWithoutIndex( principal: Principal, - input: Omit, context: KnowledgeResourceContext, - index: ActiveKnowledgeBaseReference -): Promise { - const searchInput: SearchKnowledgeInput = { ...input, knowledgeBaseIds: [index.id] } - validateKnowledgeSearchInput(searchInput) - return runKnowledgeSearch({ - principal, - input: searchInput, - context: buildKnowledgeSearchContext(principal, context, [index], input.signal), - }) + input: Pick +): Promise { + 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: [] }, + } } -/** The shared follow-up applies to a search that ran; an empty answer recorded its own. */ -function afterScopedSearch(execution: { - principal: Principal - context: KnowledgeResourceContext - input: Pick - result: ScopedSearchResult -}): Promise | undefined { - return 'userId' in execution.result - ? afterKnowledgeSearch({ ...execution, result: execution.result }) - : undefined -} +type ScopedSearchInput = Omit -export type SearchWorkspaceKnowledgeInput = Omit< - SearchKnowledgeInput, - 'knowledgeBaseIds' | 'workspaceId' -> & { - workspaceId: string +/** + * 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], + } + 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. */ -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 emptySearch(input.query) - return searchOwnersIndex( - principal, - { ...input, workspaceId: context.workspaceId }, - context, - index - ) - }, - afterSuccess: ({ principal, context, input, result }) => - afterScopedSearch({ principal, context, input, result }), -}) - export const searchWorkspaceKnowledge = instrumentSearchUseCase( 'workspace_application', - searchWorkspaceKnowledgeUseCase + defineScopedSearchUseCase({ + resolveContext: (input) => resolveKnowledgeWorkspaceContext(input), + findIndex: (context) => findWorkspaceSearchIndex(context.workspaceId!), + searchInput: (input, context) => ({ ...input, workspaceId: context.workspaceId }), + }) ) -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 }) { - input.signal?.throwIfAborted() - const index = await measureSearchStage('index_resolution', () => - findSearchIndex({ - kind: 'organization', - organizationId: context.organizationId, - }) - ) - 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 emptySearch(input.query) - } - return searchOwnersIndex(principal, input, context, index) - }, - afterSuccess: ({ principal, context, input, result }) => - afterScopedSearch({ principal, context, input, result }), -}) - 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 emptySearch(input.query) - } - return searchOwnersIndex( - principal, - { - ...input, - workspaceId: input.workspaceId ?? undefined, - organizationId: input.organizationId ?? undefined, - }, - context, - index - ) - }, - afterSuccess: ({ principal, context, input, result }) => - afterScopedSearch({ principal, context, input, result }), -}) - 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/search/queries.test.ts b/apps/sim/lib/knowledge/search/queries.test.ts index c407f1f27ce..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()) @@ -71,7 +74,6 @@ describe('retrieval leg budgets', () => { 'applies vector budget %s without shortening keyword or tag retrieval', async (vectorBudgetMs) => { resetDbChainMock() - forgetProjectionFilled() vi.spyOn(performance, 'now').mockReturnValue(1000) const remaining = SearchBudget.prototype.remaining const deadlines = new Map() @@ -406,7 +408,6 @@ describe('workspace-scoped vector retrieval', () => { beforeEach(() => { resetDbChainMock() - forgetProjectionFilled() getForConnectors.mockReset() probeRows = probe traversedRows = candidates @@ -823,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') @@ -849,7 +851,6 @@ describe('workspace search filters before ranking', () => { } beforeEach(() => { resetDbChainMock() - forgetProjectionFilled() }) function expectScopeOnEveryQuery() { @@ -950,7 +951,6 @@ describe('hydration follows ranked candidates', () => { beforeEach(() => { resetDbChainMock() - forgetProjectionFilled() probePages.length = 0 exactPages.length = 0 candidatePages.length = 0 @@ -1170,7 +1170,8 @@ describe('hydration follows ranked candidates', () => { ) } else { const ranking = statements().find((query) => isExactRanking(query.sql))! - expect(ranking.sql).toContain('AS "connectorId"') + /** 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 = @@ -1306,7 +1307,6 @@ describe('permitted-document planner', () => { beforeEach(() => { resetDbChainMock() - forgetProjectionFilled() probeRows = [] exactRows = [] traversedRows = [] @@ -1315,7 +1315,6 @@ 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. */ @@ -1600,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' }], @@ -1622,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 () => { @@ -1996,7 +1993,6 @@ describe('permitted-document planner', () => { expect(statements().some((query) => isWalk(query.sql))).toBe(true) resetDbChainMock() - forgetProjectionFilled() dbChainMockFns.execute.mockImplementation(async () => []) await retrieveKnowledgeSearch({ ...search, @@ -2048,7 +2044,6 @@ describe('permitted-document planner', () => { }) it('rebuilds the pool without a gated source the caller turns out not to hold', async () => { - forgetProjectionFilled() queueTableRows(schemaMock.knowledgeConnector, [ { id: 'gated-src', @@ -2091,13 +2086,12 @@ describe('permitted-document planner', () => { const walks = statements().filter((query) => isWalk(query.sql)) expect(walks).toHaveLength(2) expect(JSON.stringify(walks[0])).not.toContain('/* excluded sources */') + expect(JSON.stringify(walks[1])).toContain('/* excluded sources */') expect(JSON.stringify(walks[1])).toContain('OR NOT (') expect(statements().some((query) => isPageStatement(query.sql))).toBe(false) - forgetProjectionFilled() }) it('excludes a denied source through its documents while the projection is unfilled', async () => { - forgetProjectionFilled() queueTableRows(schemaMock.knowledgeConnector, [ { id: 'gated-src', @@ -2138,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 () => { @@ -2249,10 +2242,8 @@ describe('filters on a resolved scope', () => { beforeEach(() => { resetDbChainMock() - forgetProjectionFilled() forgetIndexedVectorSources() forgetSearchReach() - forgetProjectionFilled() probeRows = [] traversedRows = [] rerankRows = [] @@ -2403,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) @@ -2461,7 +2451,6 @@ describe('filters on a resolved scope', () => { 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 unfilled')) return [{ unfilled: true }] if (isWalk(statement)) return traversedRows if (isPageStatement(statement)) return rerankRows return [] diff --git a/apps/sim/lib/knowledge/search/queries.ts b/apps/sim/lib/knowledge/search/queries.ts index e2f70c6f6d0..75cf97dc86f 100644 --- a/apps/sim/lib/knowledge/search/queries.ts +++ b/apps/sim/lib/knowledge/search/queries.ts @@ -117,43 +117,46 @@ const PROJECTION_FILLED_TTL_MS = 60_000 * Whether the ranking projection still holds rows the backfill has not filled. Read off the * unfilled-rows index in microseconds and remembered briefly: the answer only ever changes once. */ -const projectionFilled = new LRUCache({ +const projectionFilled = new LRUCache< + ProjectionSourceAclTable, + boolean, + { budget: SearchBudget | undefined; stage: SearchStage } +>({ max: PROJECTION_SOURCE_ACL_TABLES.length, ttl: PROJECTION_FILLED_TTL_MS, + /** + * 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 + 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 + } + }, }) -/** - * Whether the projection's rows all carry their mirrored columns. The read that misses the cache - * is the search's own, under its budget like every other read of the leg, so one search's - * deadline never fails another's; it is one index page, once a minute. A read that fails for a - * reason other than the deadline is answered as unfilled, the slower and safe form. - */ +/** 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 { - const cached = projectionFilled.get(projection) - if (cached !== undefined) return cached - const table = projection === 'embedding_search' ? embeddingSearch : embeddingKeywordTin - try { - const [row] = await runSearchQuery(budget, stage, (executor) => - executor.execute<{ unfilled: boolean }>(sql` - SELECT EXISTS (SELECT 1 FROM ${table} WHERE ${table.acl} IS NULL) AS unfilled`) - ) - const filled = !row?.unfilled - projectionFilled.set(projection, filled) - return filled - } catch (error) { - if (budget?.isTimeout(error)) throw error - return false - } + return (await projectionFilled.fetch(projection, { context: { budget, stage } })) ?? false } -/** Forgets whether the projection was filled, after its rows changed. */ +/** 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 @@ -708,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, @@ -929,7 +938,9 @@ 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) /** @@ -938,21 +949,22 @@ function hydrateSearchCandidates( * 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`. */ @@ -1564,8 +1576,7 @@ async function selectSourceVectorCandidates(input: { withVectorScanSettings( (executor) => executor.execute(sql` - SELECT ${embeddingSearch.id} AS id, ${embeddingSearch.documentId} AS "documentId", ${embeddingSearch.connectorId} AS "connectorId", - ${input.candidateDistance} AS distance + SELECT ${PROJECTION_CANDIDATE_COLUMNS}, ${input.candidateDistance} AS distance FROM ${embeddingSearch} /* on-row visibility */ WHERE ${and( base, @@ -1611,8 +1622,7 @@ async function selectSourceVectorCandidates(input: { const rows = await runSearchQuery(input.budget, 'vector.source_exact', (executor) => executor.execute(sql` WITH readable_chunks AS MATERIALIZED ( - SELECT ${embeddingSearch.id} AS id, ${embeddingSearch.documentId} AS "documentId", ${embeddingSearch.connectorId} AS "connectorId", - ${input.candidateDistance} AS distance + SELECT ${PROJECTION_CANDIDATE_COLUMNS}, ${input.candidateDistance} AS distance FROM ${embeddingSearch} WHERE ${readableChunks} LIMIT ${SOURCE_EXACT_CHUNK_LIMIT + 1} @@ -1649,7 +1659,6 @@ async function selectSourceVectorCandidates(input: { return ranked .sort((a, b) => Number(a.distance) - Number(b.distance)) .slice(0, input.candidateLimit) - .map(({ id, documentId, connectorId }) => ({ id, documentId, connectorId })) } /** Sources ranked at once; each holds a connection for its own statement. */ @@ -1707,7 +1716,14 @@ async function selectVectorResults(params: SearchParams): Promise executor.execute(sql` - SELECT ${embeddingSearch.id} AS id, ${embeddingSearch.documentId} AS "documentId", ${embeddingSearch.connectorId} AS "connectorId" + SELECT ${PROJECTION_CANDIDATE_COLUMNS} FROM ${embeddingSearch} WHERE ${and( inArray(embeddingSearch.knowledgeBaseId, params.knowledgeBaseIds), @@ -1871,7 +1887,7 @@ async function selectVectorResults(params: SearchParams): Promise( plan ? sql` - SELECT ${embeddingSearch.id} AS id, ${embeddingSearch.documentId} AS "documentId", ${embeddingSearch.connectorId} AS "connectorId" + SELECT ${PROJECTION_CANDIDATE_COLUMNS} FROM ${embeddingSearch} /* on-row visibility */ WHERE ${and( scopeOfWalk, @@ -1935,6 +1951,7 @@ async function selectVectorResults(params: SearchParams): Promise= MAX_VECTOR_CANDIDATES, + filled, } annotateSearchDiagnostics({ vectorCandidateCount: selected.length, @@ -1942,20 +1959,20 @@ async function selectVectorResults(params: SearchParams): Promise hydrateSearchCandidates( ids, @@ -1990,7 +2002,8 @@ async function selectVectorResults(params: SearchParams): Promise hydrateSearchCandidates( ids, From 22fed447eb0a55c27351978673a3e2c95a4df838 Mon Sep 17 00:00:00 2001 From: Waleed Latif Date: Sun, 20 Sep 2026 23:22:49 -0700 Subject: [PATCH 5/6] improvement(knowledge): embed only after every prerequisite holds, refuse a contradicting owner, count the shared fill read --- .../kb-block-search.integration.ts | 10 ++- apps/sim/lib/knowledge/application/search.ts | 75 ++++++++----------- .../knowledge/application/workspace-search.ts | 8 ++ 3 files changed, 47 insertions(+), 46 deletions(-) diff --git a/apps/sim/lib/knowledge/__integration__/kb-block-search.integration.ts b/apps/sim/lib/knowledge/__integration__/kb-block-search.integration.ts index 3b61286d0fa..5ec095222bd 100644 --- a/apps/sim/lib/knowledge/__integration__/kb-block-search.integration.ts +++ b/apps/sim/lib/knowledge/__integration__/kb-block-search.integration.ts @@ -132,10 +132,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 +144,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/application/search.ts b/apps/sim/lib/knowledge/application/search.ts index b68ba56076d..41fe47443e7 100644 --- a/apps/sim/lib/knowledge/application/search.ts +++ b/apps/sim/lib/knowledge/application/search.ts @@ -390,51 +390,42 @@ export async function runKnowledgeSearch({ : undefined const resultSecretRegistry = preparedRegistry ?? input.resultSecretRegistry input.signal?.throwIfAborted() - /** The embedding is requested the moment admission passes, beside the scope and defaults reads. */ - const admittedEmbedding = async () => { - const billingAttribution = await admit() - input.signal?.throwIfAborted() - const queryEmbedding = hasQuery - ? await measureSearchStage('embedding', () => - runWithKnowledgeModelInputProvenance(resultSecretRegistry, () => - generateSearchEmbedding( - input.query!, - embeddingTarget!, - context.workspaceId, - input.signal - ) - ) - ) - : null - return { billingAttribution, queryEmbedding } - } - const [access, searchDefaults, admitted, tagDefinitions, rerankerCredential] = await Promise.all([ - measureSearchStage('access_scope', () => context.access.get()), - measureSearchStage('defaults', () => - resolveKnowledgeSearchDefaults({ - workspaceId: context.workspaceId, - organizationId: context.organizationId, + const [access, searchDefaults, billingAttribution, tagDefinitions, rerankerCredential] = + await Promise.all([ + measureSearchStage('access_scope', () => context.access.get()), + measureSearchStage('defaults', () => + resolveKnowledgeSearchDefaults({ + workspaceId: context.workspaceId, + organizationId: context.organizationId, - /** The signed-in person, if any; never the billing owner or a key's creator. */ - userId: resolvePrincipalSubjectUserId(principal) ?? undefined, - requestedMode: input.searchMode, - }) - ), - admittedEmbedding(), - /** The tag names the results are labelled with depend on the bases alone. */ - filters.length === 0 - ? measureSearchStage('tag_definitions', () => - 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, - ]) - const { billingAttribution, queryEmbedding } = admitted + /** The signed-in person, if any; never the billing owner or a key's creator. */ + userId: resolvePrincipalSubjectUserId(principal) ?? undefined, + requestedMode: input.searchMode, + }) + ), + admit(), + /** The tag names the results are labelled with depend on the bases alone. */ + filters.length === 0 + ? measureSearchStage('tag_definitions', () => + getDocumentTagDefinitionsByKnowledgeBaseIds(knowledgeBaseIds) + ) + : Promise.resolve(definitionsByKnowledgeBase), + /** 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() + /** 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, diff --git a/apps/sim/lib/knowledge/application/workspace-search.ts b/apps/sim/lib/knowledge/application/workspace-search.ts index 99ad8878978..d69fe6bca21 100644 --- a/apps/sim/lib/knowledge/application/workspace-search.ts +++ b/apps/sim/lib/knowledge/application/workspace-search.ts @@ -1,5 +1,6 @@ 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' @@ -112,6 +113,13 @@ function defineScopedSearchUseCase< ...surface.searchInput(input, context), knowledgeBaseIds: [index.id], } + /** 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') + } validateKnowledgeSearchInput(searchInput) return runKnowledgeSearch({ principal, From 28c0359ccd6b07b9466ce60d52e9e174301f9963 Mon Sep 17 00:00:00 2001 From: Waleed Latif Date: Sun, 20 Sep 2026 23:31:46 -0700 Subject: [PATCH 6/6] improvement(knowledge): reset the projection-fill memo per integration iteration --- .../knowledge/__integration__/kb-block-search.integration.ts | 3 +++ 1 file changed, 3 insertions(+) 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 5ec095222bd..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) => {