diff --git a/apps/sim/app/api/knowledge/search/route.test.ts b/apps/sim/app/api/knowledge/search/route.test.ts index 46e4a113eef..445fe171c63 100644 --- a/apps/sim/app/api/knowledge/search/route.test.ts +++ b/apps/sim/app/api/knowledge/search/route.test.ts @@ -46,6 +46,7 @@ describe('workspace search route', () => { expect(call.input.filters).toEqual({ source: 'slack', documentIds: ['doc-1'] }) expect(call.input.signal).toBe(request.signal) expect(call.input.allowPartialResults).toBe(true) + expect(call.input.vectorBudgetMs).toBe(3000) controller.abort() expect(call.input.signal.aborted).toBe(true) await expect(response.json()).resolves.toEqual({ diff --git a/apps/sim/app/api/knowledge/search/route.ts b/apps/sim/app/api/knowledge/search/route.ts index 178027a04b4..b1db5e04d39 100644 --- a/apps/sim/app/api/knowledge/search/route.ts +++ b/apps/sim/app/api/knowledge/search/route.ts @@ -9,6 +9,8 @@ import { knowledgeOperations } from '@/lib/knowledge/application/operations' import { searchScopedKnowledge } from '@/lib/knowledge/application/workspace-search' import { sourceAuthor } from '@/lib/knowledge/search/author' +const DIRECT_SEARCH_VECTOR_BUDGET_MS = 3000 + export const POST = defineInternalJsonRoute({ contract: searchWorkspaceKnowledgeContract, auth: internalSessionAuth, @@ -25,6 +27,7 @@ export const POST = defineInternalJsonRoute({ query: body.query, topK: body.topK, allowPartialResults: true, + vectorBudgetMs: DIRECT_SEARCH_VECTOR_BUDGET_MS, surface: 'dashboard' as const, signal: request.signal, }), diff --git a/apps/sim/app/workspace/[workspaceId]/home/components/knowledge-search-results/knowledge-search-results.test.tsx b/apps/sim/app/workspace/[workspaceId]/home/components/knowledge-search-results/knowledge-search-results.test.tsx index 7825d523178..85a3e336b60 100644 --- a/apps/sim/app/workspace/[workspaceId]/home/components/knowledge-search-results/knowledge-search-results.test.tsx +++ b/apps/sim/app/workspace/[workspaceId]/home/components/knowledge-search-results/knowledge-search-results.test.tsx @@ -68,13 +68,13 @@ describe('source indexing context in search results', () => { mocks.overview.mockReturnValue({ data: undefined }) await render() expect(container.textContent).not.toContain('Still indexing') - expect(container.textContent).toContain('No documents you can read match') + expect(container.textContent).toContain('Search found no results.') }) }) describe('incomplete search coverage', () => { it.each([false, true])( - 'keeps matches and offers a retry without claiming absence (hasResults=%s)', + 'shows matches without timeout copy or retry controls (hasResults=%s)', async (hasResults) => { mocks.search.mockReturnValue({ data: { @@ -106,17 +106,17 @@ describe('incomplete search coverage', () => { await render() expect(container.textContent).not.toContain('Search couldn’t run') expect(container.textContent).not.toContain('No documents') - expect(container.textContent).not.toContain('0 documents') + expect(container.textContent).not.toContain('Some results may be missing.') + expect(container.textContent).not.toContain('Search is incomplete.') expect(container.textContent).toContain( - hasResults ? 'Some results may be missing.' : 'Search is incomplete.' + hasResults ? '1 document' : 'Search found no results.' ) if (hasResults) expect(container.textContent).toContain('Release plan') const retry = [...container.querySelectorAll('button')].find( (button) => button.textContent === 'Try again' ) - expect(retry).toBeDefined() - await act(async () => retry!.click()) - expect(mocks.retry).toHaveBeenCalledOnce() + expect(retry).toBeUndefined() + expect(mocks.retry).not.toHaveBeenCalled() } ) }) diff --git a/apps/sim/app/workspace/[workspaceId]/home/components/knowledge-search-results/knowledge-search-results.tsx b/apps/sim/app/workspace/[workspaceId]/home/components/knowledge-search-results/knowledge-search-results.tsx index 853d3f5f7e0..bb51a95f74c 100644 --- a/apps/sim/app/workspace/[workspaceId]/home/components/knowledge-search-results/knowledge-search-results.tsx +++ b/apps/sim/app/workspace/[workspaceId]/home/components/knowledge-search-results/knowledge-search-results.tsx @@ -144,7 +144,6 @@ export function KnowledgeSearchResults({ .filter((provider) => provider.isSyncing) .map((provider) => connectorDisplayName(provider.connectorType)) const documents = useMemo(() => groupResultsByDocument(search?.results ?? []), [search?.results]) - const incomplete = search?.retrieval.status === 'partial' const sourceTypes = [ ...new Set([ ...(filters.source ? [filters.source] : []), @@ -205,24 +204,18 @@ export function KnowledgeSearchResults({
- {incomplete && documents.length === 0 ? ( - 'Search is incomplete.' + {documents.length === 0 ? ( + 'Search found no results.' ) : ( <> {documents.length === 1 ? '1 document' : `${documents.length} documents`} {' · searched as you'} - {incomplete && Some results may be missing.} )} {indexingNote && {indexingNote}} - {incomplete && ( - void refetchSearch()}> - {isFetching ? 'Retrying…' : 'Try again'} - - )}
{showFilters && (
@@ -256,15 +249,7 @@ export function KnowledgeSearchResults({ ))}
)} - {documents.length === 0 ? ( - !incomplete && ( -

- {filtersActive - ? 'No documents match these filters.' - : `No documents you can read match “${query}”.`} -

- ) - ) : ( + {documents.length > 0 && (
{documents.map((result) => { const source = toSource(result, query, scope) diff --git a/apps/sim/lib/knowledge/__integration__/search-latency.integration.ts b/apps/sim/lib/knowledge/__integration__/search-latency.integration.ts index 919cce00ac4..df209d20d1b 100644 --- a/apps/sim/lib/knowledge/__integration__/search-latency.integration.ts +++ b/apps/sim/lib/knowledge/__integration__/search-latency.integration.ts @@ -480,9 +480,11 @@ describe.skipIf(!enabled)('Assistant search latency on a realistic indexed corpu expect(run).not.toHaveBeenCalled() }) - it.each(['keyword', 'both'] as const)( - 'handles %s SQL branches exceeding the deadline explicitly', + it.each(['vector', 'keyword', 'both'] as const)( + 'keeps the Assistant budget when %s SQL branches are delayed', async (delayedLegs) => { + diagnosticLog?.mockClear() + let vectorDelayed = false const query = SearchBudget.prototype.query const delayed = vi.spyOn(SearchBudget.prototype, 'query').mockImplementation(function ( this: SearchBudget, @@ -490,7 +492,13 @@ describe.skipIf(!enabled)('Assistant search latency on a realistic indexed corpu run: (executor: SearchExecutor) => PromiseLike ): Promise { return query.call(this, stage, async (tx) => { - if (delayedLegs === 'both' || this.leg === 'keyword') + if (delayedLegs === 'vector' && this.leg === 'vector' && !vectorDelayed) { + vectorDelayed = true + await tx.execute(sql`SELECT pg_sleep(4)`) + } else if ( + delayedLegs === 'both' || + (delayedLegs === 'keyword' && this.leg === 'keyword') + ) await tx.execute(sql`SELECT pg_sleep(9)`) return run(tx) }) as Promise @@ -509,6 +517,10 @@ describe.skipIf(!enabled)('Assistant search latency on a realistic indexed corpu }), } ) + const completed = diagnosticLog?.mock.calls.find( + ([message]) => message === 'Knowledge search completed' + ) + expect(diagnosticSchema.parse(completed?.[1]).vectorBudgetMs).toBe(8000) if (delayedLegs === 'both') { expect(result).toMatchObject({ success: true, @@ -521,14 +533,19 @@ describe.skipIf(!enabled)('Assistant search latency on a realistic indexed corpu } expect(result).toMatchObject({ success: true, - data: { retrieval: { status: 'partial', timedOutLegs: ['keyword'] } }, + data: { + retrieval: + delayedLegs === 'vector' + ? { status: 'complete', timedOutLegs: [] } + : { status: 'partial', timedOutLegs: ['keyword'] }, + }, }) const parsed = resultSchema.parse(result) expect(parsed.data.results.length).toBeGreaterThan(0) expect( parsed.data.results.every((row) => row.knowledgeBaseId === ids.knowledgeBaseId) ).toBe(true) - report['deadline.partial'] = { resultCount: parsed.data.results.length } + report[`assistant.deadline.${delayedLegs}`] = { resultCount: parsed.data.results.length } } finally { delayed.mockRestore() } @@ -539,6 +556,7 @@ describe.skipIf(!enabled)('Assistant search latency on a realistic indexed corpu it.each(['vector', 'both'] as const)( 'returns incomplete dashboard coverage when %s SQL branches exceed their deadline', async (delayedLegs) => { + diagnosticLog?.mockClear() const authenticate = vi.spyOn(internalSessionAuth, 'authenticate').mockResolvedValue({ kind: 'session', userId: ids.aliceId, @@ -552,7 +570,7 @@ describe.skipIf(!enabled)('Assistant search latency on a realistic indexed corpu ): Promise { return query.call(this, stage, async (tx) => { if (delayedLegs === 'both' || this.leg === 'vector') - await tx.execute(sql`SELECT pg_sleep(9)`) + await tx.execute(sql`SELECT pg_sleep(${delayedLegs === 'both' ? 9 : 4})`) return run(tx) }) as Promise }) @@ -569,6 +587,13 @@ describe.skipIf(!enabled)('Assistant search latency on a realistic indexed corpu }) ) expect(response.status).toBe(200) + const completed = diagnosticLog?.mock.calls.find( + ([message]) => message === 'Knowledge search completed' + ) + const diagnostics = diagnosticSchema.parse(completed?.[1]) + expect(diagnostics.vectorBudgetMs).toBe(3000) + expect(diagnostics.stages.vector.totalMs).toBeGreaterThan(2500) + expect(diagnostics.stages.vector.totalMs).toBeLessThan(4000) const data = workspaceKnowledgeSearchDataSchema.parse((await response.json()).data) expect(data.retrieval).toEqual({ status: 'partial', diff --git a/apps/sim/lib/knowledge/application/search.test.ts b/apps/sim/lib/knowledge/application/search.test.ts index d34cc59cfaf..e42e26018c4 100644 --- a/apps/sim/lib/knowledge/application/search.test.ts +++ b/apps/sim/lib/knowledge/application/search.test.ts @@ -234,6 +234,25 @@ describe('knowledge search application use case', () => { } ) + it.each([ + { surface: 'dashboard' as const, vectorBudgetMs: 3000 }, + { surface: 'copilot' as const, vectorBudgetMs: undefined }, + { surface: 'workflow' as const, vectorBudgetMs: undefined }, + ])('forwards only the configured vector budget for $surface', async (options) => { + await searchKnowledge.execute({ + principal: { kind: 'session', userId: 'user-1', sessionId: 'session-1' }, + input: { + knowledgeBaseIds: ['knowledge-1'], + query: 'release', + topK: 10, + ...options, + }, + }) + expect(mocks.executeSearch).toHaveBeenCalledWith( + expect.objectContaining({ vectorBudgetMs: options.vectorBudgetMs }) + ) + }) + describe.each(['workspace', 'organization'] as const)('%s ranking policy', (scope) => { beforeEach(() => { if (scope === 'organization') { diff --git a/apps/sim/lib/knowledge/application/search.ts b/apps/sim/lib/knowledge/application/search.ts index cdf61bfc5fd..1f8871529cb 100644 --- a/apps/sim/lib/knowledge/application/search.ts +++ b/apps/sim/lib/knowledge/application/search.ts @@ -90,8 +90,10 @@ export class KnowledgeSearchProvenanceUnavailableError extends Error { export type KnowledgeSearchTagFilter = KnowledgeTagNameFilter export interface SearchKnowledgeInput { - /** Only surfaces displaying retrieval status may accept incomplete evidence. */ + /** Allows returning available results when a retrieval leg times out. */ allowPartialResults?: boolean + /** Trusted adapter's vector retrieval budget; omitted callers use the shared default. */ + vectorBudgetMs?: number /** Optional assertion from a trusted adapter or public contract. */ workspaceId?: string organizationId?: string @@ -404,6 +406,7 @@ const searchKnowledgeUseCase = defineAuthorizedKnowledgeUseCase({ : input.topK const retrieved = await measureSearchStage('retrieval', () => retrieveKnowledgeSearch({ + vectorBudgetMs: input.vectorBudgetMs, knowledgeBaseIds, topK: candidateTopK, filters: input.filters, diff --git a/apps/sim/lib/knowledge/search/budget.ts b/apps/sim/lib/knowledge/search/budget.ts index 7f8a0eaeecf..8011733d926 100644 --- a/apps/sim/lib/knowledge/search/budget.ts +++ b/apps/sim/lib/knowledge/search/budget.ts @@ -19,7 +19,7 @@ export class SearchDeadlineError extends Error { } } -/** One leg's state, using the hybrid request's shared deadline across every refill and fallback. */ +/** One leg's absolute deadline, reused across every refill and fallback. */ export class SearchBudget { timedOut = false diff --git a/apps/sim/lib/knowledge/search/diagnostics.ts b/apps/sim/lib/knowledge/search/diagnostics.ts index 4265db0c886..04c618d9721 100644 --- a/apps/sim/lib/knowledge/search/diagnostics.ts +++ b/apps/sim/lib/knowledge/search/diagnostics.ts @@ -70,6 +70,7 @@ export interface SearchDiagnosticMetadata { boostRecency?: boolean embeddingDimensions?: number vectorRanking?: 'exact' | 'binary-rerank' + vectorBudgetMs?: number vectorCandidateLimit?: number vectorCandidateCount?: number resultCount?: number diff --git a/apps/sim/lib/knowledge/search/queries.test.ts b/apps/sim/lib/knowledge/search/queries.test.ts index f73f148ff63..86c5bdf8458 100644 --- a/apps/sim/lib/knowledge/search/queries.test.ts +++ b/apps/sim/lib/knowledge/search/queries.test.ts @@ -15,12 +15,14 @@ import { WORKSPACE_ACCESS_TOKENS, } from '@/lib/knowledge/access/types' import { buildTagFilterCondition } from '@/lib/knowledge/documents/tag-filter' +import { SearchBudget } from '@/lib/knowledge/search/budget' import { executeKeywordSearch, getStructuredTagFilters, handleTagAndVectorSearch, handleTagOnlySearch, handleVectorOnlySearch, + retrieveKnowledgeSearch, type SearchParams, } from '@/lib/knowledge/search/queries' import type { StructuredFilter } from '@/lib/knowledge/types' @@ -36,6 +38,59 @@ const embeddingTable = { boolean1: 'boolean1', } +describe('retrieval leg budgets', () => { + afterEach(() => vi.restoreAllMocks()) + + it.each([undefined, 3000])( + 'applies vector budget %s without shortening keyword or tag retrieval', + async (vectorBudgetMs) => { + resetDbChainMock() + vi.spyOn(performance, 'now').mockReturnValue(1000) + const remaining = SearchBudget.prototype.remaining + const deadlines = new Map() + vi.spyOn(SearchBudget.prototype, 'remaining').mockImplementation(function ( + this: SearchBudget + ) { + deadlines.set(this.leg, this.deadline) + return remaining.call(this) + }) + const access: UserAccessScope = { + kind: 'user', + userId: 'user-1', + tokens: WORKSPACE_ACCESS_TOKENS, + } + const params = { + knowledgeBaseIds: ['knowledge-1'], + topK: 10, + access, + accessProvider: { + get: async () => access, + getForConnectors: async () => access, + getForDocuments: async () => access, + }, + searchMode: 'hybrid' as const, + vectorBudgetMs, + } + await retrieveKnowledgeSearch({ + ...params, + query: 'release', + queryVector: { vector: '[1,0]', dimensions: 1536 }, + }) + await retrieveKnowledgeSearch({ + ...params, + structuredFilters: [ + { tagSlot: 'tag1', fieldType: 'text', operator: 'eq', value: 'release' }, + ], + }) + expect(Object.fromEntries(deadlines)).toEqual({ + vector: 1000 + (vectorBudgetMs ?? 8000), + keyword: 9000, + tags: 9000, + }) + } + ) +}) + /** * The global `drizzle-orm` mock renders `sql` fragments to a `?`-placeholder * string via `toSQL()`, so we can assert the exact predicate each filter builds. diff --git a/apps/sim/lib/knowledge/search/queries.ts b/apps/sim/lib/knowledge/search/queries.ts index af35de09073..cb0f82b20cf 100644 --- a/apps/sim/lib/knowledge/search/queries.ts +++ b/apps/sim/lib/knowledge/search/queries.ts @@ -1241,6 +1241,8 @@ export async function handleTagAndVectorSearch(params: SearchParams): Promise