Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -84,6 +84,29 @@ describe('Assistant retrieval tools', () => {
next: null,
})
})
it('returns empty incomplete retrieval as a recoverable search outcome and logs coverage', async () => {
mocks.search.mockResolvedValue({
retrieval: { status: 'partial', timedOutLegs: ['vector', 'keyword'] },
knowledgeBases: [{ id: 'index', name: 'Enterprise Search' }],
results: [],
})

const result = await searchWorkspaceServerTool.execute({ query: 'canaries' }, context)

expect(result).toMatchObject({
success: true,
message: expect.stringContaining('cannot establish absence or completeness'),
data: {
retrieval: { status: 'partial', timedOutLegs: ['vector', 'keyword'] },
results: [],
},
})
expect(result).not.toHaveProperty('error')
expect(mocks.info).toHaveBeenCalledWith(
'Knowledge search completed',
expect.objectContaining({ passageBytes: 0, originalPassageBytes: 0, outcome: 'success' })
)
})
it('pins organization and private chat while reusing the canonical search index and citations', async () => {
const orgContext = {
...context,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -93,7 +93,7 @@ export const searchWorkspaceServerTool: BaseServerTool = {
const names = new Map(result.knowledgeBases.map((base) => [base.id, base.name]))
const output = {
success: true,
message: `${result.retrieval.status === 'partial' ? 'Partial search: a retrieval branch reached its deadline. These results cannot establish absence or completeness. ' : ''}Found ${result.results.length} passage previews. Read a document at its chunkIndex for more context. ${CITATION_INSTRUCTION}`,
message: `${result.retrieval.status === 'partial' ? 'Search coverage is incomplete. Continue with a more specific query or source filter; these results cannot establish absence or completeness. ' : ''}Found ${result.results.length} passage previews. Read a document at its chunkIndex for more context. ${CITATION_INSTRUCTION}`,
data: {
query: safeQuery,
retrieval: result.retrieval,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -83,7 +83,11 @@ const reused = reuseFile ? readFixtureReport(reuseFile) : undefined
const ids = reused?.fixture ?? createKnowledgeAclFixtureIds()
const unrelated = reused?.unrelatedFixture ?? createKnowledgeAclFixtureIds()
const organizationChatId = generateId()
const queryVector = Array.from({ length: dimensions }, (_, index) => (index === 0 ? 1 : 0))
const queryVector = Array.from({ length: dimensions }, (_, index) =>
Math.sin((index + 1) * 12.9898)
)
const queryMagnitude = Math.hypot(...queryVector)
for (let index = 0; index < queryVector.length; index++) queryVector[index] /= queryMagnitude
const captured: CapturedQuery[] = []
const report: Record<string, unknown> = {
fixture: ids,
Expand Down Expand Up @@ -133,6 +137,7 @@ const explainSchema = z.array(z.object({ Plan: explainNodeSchema }).passthrough(

function usesVectorIndex(node: ExplainNode): boolean {
return (
node['Index Name'] === 'embedding_binary_hnsw_idx' ||
node['Index Name'] === 'embedding_vector_hnsw_idx' ||
(node.Plans?.some(usesVectorIndex) ?? false)
)
Expand Down Expand Up @@ -237,14 +242,21 @@ async function sample(label: string, run: () => ReturnType<typeof search>) {
const plan = await db.$client.begin(async (tx) => {
await tx.unsafe("SET LOCAL hnsw.iterative_scan = 'relaxed_order'")
await tx.unsafe('SET LOCAL hnsw.max_scan_tuples = 20000')
if (query.query.includes('binary_quantize')) {
await tx.unsafe('SET LOCAL hnsw.max_scan_tuples = 100000')
await tx.unsafe('SET LOCAL hnsw.ef_search = 200')
await tx.unsafe('SET LOCAL hnsw.scan_mem_multiplier = 4')
}
return tx.unsafe(`EXPLAIN (ANALYZE, BUFFERS, FORMAT JSON) ${query.query}`, query.parameters)
})
plans.push({
kind: query.query.includes('keyword_rank')
? 'keyword'
: query.query.includes('order by')
: query.query.includes('binary_quantize')
? 'vector'
: 'probe',
: query.query.includes('order by')
? 'rerank'
: 'probe',
query: query.query,
parameters: query.parameters,
plan: explainSchema.parse(plan[0]['QUERY PLAN']),
Expand Down Expand Up @@ -378,8 +390,8 @@ describe.skipIf(!enabled)('Assistant search latency on a realistic indexed corpu
CASE WHEN n % 8 = 0 THEN 'Orion deployment reference. ' ELSE 'Engineering operations reference. ' END ||
(SELECT string_agg(md5(n::text || ':' || paragraph::text), ' ') FROM generate_series(1, 90) paragraph),
3000, 750, 0, 3000,
l2_normalize(ARRAY(SELECT (CASE WHEN coordinate = n % 32 + 1 THEN 1 ELSE 0 END +
0.025 * sin(n::double precision * coordinate * 12.9898 + coordinate * 78.233))::real
l2_normalize(ARRAY(SELECT (sin(coordinate * (n % 32 + 1) * 12.9898) +
0.25 * sin(n::double precision * coordinate * 12.9898 + coordinate * 78.233))::real
FROM generate_series(1, ${dimensions}) coordinate)::vector(1536))
FROM generate_series(${first}::int, ${last}::int) n`)
})
Expand Down Expand Up @@ -494,7 +506,13 @@ describe.skipIf(!enabled)('Assistant search latency on a realistic indexed corpu
}
)
if (delayedLegs === 'both') {
expect(result).toMatchObject({ success: false, retryable: true })
expect(result).toMatchObject({
success: true,
data: {
retrieval: { status: 'partial', timedOutLegs: ['vector', 'keyword'] },
results: [],
},
})
return
}
expect(result).toMatchObject({
Expand Down Expand Up @@ -526,6 +544,18 @@ describe.skipIf(!enabled)('Assistant search latency on a realistic indexed corpu
const vectorPlans = plans.filter((plan) => plan.kind === 'vector')
expect(vectorPlans).toHaveLength(1)
expect(usesVectorIndex(vectorPlans[0].plan[0].Plan)).toBe(true)
expect(plans.some((plan) => plan.kind === 'rerank')).toBe(true)
const rerank = plans.find((plan) => plan.kind === 'rerank')!
const actual = await db.$client.unsafe(rerank.query, rerank.parameters).values()
const expected = await db.execute<{ id: string }>(sql`SELECT id FROM embedding
WHERE knowledge_base_id = ${ids.knowledgeBaseId} AND enabled
ORDER BY (embedding <=> ${JSON.stringify(queryVector)}::vector) + 0, id
LIMIT ${actual.length}`)
const expectedIds = new Set(expected.map(({ id }) => id))
const recall = actual.filter(([id]) => expectedIds.has(id)).length / expected.length
expect(recall).toBeGreaterThanOrEqual(0.95)
report[`recall.${iteration}`] = { neighbors: expected.length, recall }
saveReport()
}
expect(embeddingCalls - before).toBe(2)
}, 180_000)
Expand Down Expand Up @@ -578,7 +608,7 @@ describe.skipIf(!enabled)('Assistant search latency on a realistic indexed corpu
expect(probe).toHaveLength(1)
expect(probe[0].query).not.toContain('<=>')
expect(probe[0].plan[0].Plan['Actual Rows']).toBe(12)
const vector = plans.filter((plan) => plan.kind === 'vector')
const vector = plans.filter((plan) => plan.kind === 'rerank')
expect(vector).toHaveLength(1)
expect(vector[0].query).toContain('"embedding"."id" in')
} finally {
Expand Down Expand Up @@ -672,7 +702,7 @@ describe.skipIf(!enabled)('Assistant search latency on a realistic indexed corpu
}, 180_000)
/** Opt in with local Sim and Go URLs; uses the real configured provider, billing adapter, and async resume protocol. */
it.skipIf(!process.env.KNOWLEDGE_SEARCH_ASSISTANT_URL)(
'answers through local Go Assistant with progressive reads and citations',
'recovers quietly from incomplete search through local Go Assistant, then reads and cites evidence',
async () => {
const assistantUrl = new URL(process.env.KNOWLEDGE_SEARCH_ASSISTANT_URL!)
const simUrl = new URL(process.env.KNOWLEDGE_SEARCH_SIM_URL!)
Expand Down Expand Up @@ -708,6 +738,18 @@ describe.skipIf(!enabled)('Assistant search latency on a realistic indexed corpu
bytes: number
}> = []
let answer = ''
let incompleteSearch = true
const query = SearchBudget.prototype.query
const delayed = vi.spyOn(SearchBudget.prototype, 'query').mockImplementation(function <T>(
this: SearchBudget,
stage: SearchStage,
run: (executor: SearchExecutor) => PromiseLike<T>
): Promise<T> {
return query.call(this, stage, async (tx) => {
if (incompleteSearch) await tx.execute(sql`SELECT pg_sleep(9)`)
return run(tx)
}) as Promise<T>
})
const started = performance.now()
try {
await db
Expand Down Expand Up @@ -808,9 +850,20 @@ describe.skipIf(!enabled)('Assistant search latency on a realistic indexed corpu
bytes: Buffer.byteLength(JSON.stringify(result)),
})
const { success } = z.object({ success: z.boolean() }).parse(result)
expect(success).toBe(true)
if (incompleteSearch) {
expect(call.toolName).toBe('search_workspace')
expect(result).toMatchObject({
data: {
retrieval: { status: 'partial', timedOutLegs: ['vector', 'keyword'] },
results: [],
},
})
}
return { callId, name: call.toolName, success, data: result }
})
)
incompleteSearch = false
path = '/api/tools/resume'
body = {
checkpointId: checkpoint.checkpointId,
Expand All @@ -828,10 +881,12 @@ describe.skipIf(!enabled)('Assistant search latency on a realistic indexed corpu
expect(answer).toContain('SILVER COMET')
expect(answer).toContain('K7M2-84')
expect(answer).toContain('<source>')
expect(answer).not.toMatch(/timed?\s*out|timeout|internal retr(?:y|ies)/i)
expect(calls.some((call) => call.name === 'read_document')).toBe(true)
expect(calls.some((call) => call.name === 'search_workspace')).toBe(true)
expect(calls.filter((call) => call.name === 'search_workspace').length).toBeGreaterThan(1)
expect(calls.every((call) => call.bytes < 40000)).toBe(true)
} finally {
delayed.mockRestore()
await db.update(embedding).set(original).where(eq(embedding.id, chunkId))
}
},
Expand Down
32 changes: 31 additions & 1 deletion apps/sim/lib/knowledge/application/search.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,7 @@ const mocks = vi.hoisted(() => ({
checkActorUsage: vi.fn(),
generateEmbedding: vi.fn(),
executeSearch: vi.fn(),
retrieval: vi.fn(),
getDocumentMetadata: vi.fn(),
getTagDefinitions: vi.fn(),
getTagDefinitionsBatch: vi.fn(),
Expand Down Expand Up @@ -91,7 +92,7 @@ vi.mock('@/lib/knowledge/search/queries', () => ({
generateSearchEmbedding: mocks.generateEmbedding,
retrieveKnowledgeSearch: async (...args: unknown[]) => ({
rows: await mocks.executeSearch(...args),
retrieval: { status: 'complete', timedOutLegs: [] },
retrieval: mocks.retrieval(),
}),
getDocumentMetadataByIds: mocks.getDocumentMetadata,
}))
Expand Down Expand Up @@ -129,6 +130,7 @@ const knowledgeBase = {

describe('knowledge search application use case', () => {
beforeEach(() => {
mocks.retrieval.mockReturnValue({ status: 'complete', timedOutLegs: [] })
vi.clearAllMocks()
mocks.rerank.mockReset()
resetDbChainMock()
Expand Down Expand Up @@ -204,6 +206,34 @@ describe('knowledge search application use case', () => {
expect(result.totalResults).toBe(0)
})

it.each([false, true])(
'requires explicit partial-result support for empty incomplete searches (allowPartialResults=%s)',
async (allowPartialResults) => {
mocks.retrieval.mockReturnValue({ status: 'partial', timedOutLegs: ['vector', 'keyword'] })
mocks.executeSearch.mockResolvedValue([])
const result = searchKnowledge.execute({
principal: { kind: 'session', userId: 'user-1', sessionId: 'session-1' },
input: {
workspaceId: 'workspace-1',
knowledgeBaseIds: ['knowledge-1'],
query: 'canaries',
topK: 20,
allowPartialResults,
},
})

if (!allowPartialResults) {
await expect(result).rejects.toThrow('retrieval deadline')
return
}
await expect(result).resolves.toMatchObject({
results: [],
totalResults: 0,
retrieval: { status: 'partial', timedOutLegs: ['vector', 'keyword'] },
})
}
)

describe.each(['workspace', 'organization'] as const)('%s ranking policy', (scope) => {
beforeEach(() => {
if (scope === 'organization') {
Expand Down
3 changes: 0 additions & 3 deletions apps/sim/lib/knowledge/application/search.ts
Original file line number Diff line number Diff line change
Expand Up @@ -709,9 +709,6 @@ const searchKnowledgeUseCase = defineAuthorizedKnowledgeUseCase({
}
}
annotateSearchDiagnostics({ resultCount: results.length })
if (retrieved.retrieval.status === 'partial' && results.length === 0) {
throw new SearchDeadlineError()
}
const cost = baseCost
? {
input: baseCost.input,
Expand Down
3 changes: 3 additions & 0 deletions apps/sim/lib/knowledge/search/diagnostics.ts
Original file line number Diff line number Diff line change
Expand Up @@ -47,6 +47,7 @@ export type SearchStage =
| 'vector.settings'
| 'vector.probe'
| 'vector.ann'
| 'vector.rerank'
| 'vector.exact'

/** Fixed, content-free fields. Never pass queries, filters, document identities, SQL, or errors. */
Expand All @@ -68,6 +69,8 @@ export interface SearchDiagnosticMetadata {
searchMode?: 'hybrid' | 'vector'
boostRecency?: boolean
embeddingDimensions?: number
vectorRanking?: 'exact' | 'binary-rerank'
vectorCandidateCount?: number
resultCount?: number
/** Tool output before the executor's final egress projection; counts only, never content. */
toolResultBytes?: number
Expand Down
19 changes: 15 additions & 4 deletions apps/sim/lib/knowledge/search/queries.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -528,6 +528,7 @@ describe('live repository authorization follows ranked candidates', () => {
schemaMock.embedding,
Array.from({ length: 200 }, (_, index) => candidate(`probe-${index}`, 'allowed-source'))
)
queueTableRows(schemaMock.embedding, [{ id: 'far' }, { id: 'near' }])
queueTableRows(schemaMock.embedding, [
{ ...candidate('far', 'allowed-source'), distance: 0.3 },
{ ...candidate('near', 'allowed-source'), distance: 0.1 },
Expand All @@ -544,13 +545,17 @@ describe('live repository authorization follows ranked candidates', () => {
})
expect(rows.map((row) => row.id)).toEqual(['near'])
expect(dbChainMockFns.orderBy.mock.calls[0]).toHaveLength(1)
expect(Object.keys(dbChainMockFns.select.mock.calls[1][0])).toEqual(['id', 'distance'])
expect(Object.keys(dbChainMockFns.select.mock.calls[1][0])).toEqual(['id'])
expect(render(dbChainMockFns.orderBy.mock.calls[0][0]).sql).toContain('binary_quantize')
expect(JSON.stringify(dbChainMockFns.orderBy.mock.calls[1][0])).toContain('<=>')
expect(dbChainMockFns.limit.mock.calls[1]).toEqual([4000])
expect(JSON.stringify(dbChainMockFns.where.mock.calls[1][0])).not.toContain('<=>')
expect(dbChainMockFns.limit.mock.invocationCallOrder[1]).toBeLessThan(
dbChainMockFns.select.mock.invocationCallOrder[2]
)
expect(JSON.stringify(dbChainMockFns.where.mock.calls[1][0])).toContain('OFFSET 0')
expect(getForConnectors).toHaveBeenCalledExactlyOnceWith(['allowed-source'], undefined)
expect(JSON.stringify(dbChainMockFns.where.mock.calls[2][0])).toContain('github_read_grant')
expect(JSON.stringify(dbChainMockFns.where.mock.calls[3][0])).toContain('github_read_grant')
})

it('finishes empty scopes after the bounded probe without scanning HNSW or calling providers', async () => {
Expand Down Expand Up @@ -592,6 +597,7 @@ describe('live repository authorization follows ranked candidates', () => {
schemaMock.embedding,
Array.from({ length: 200 }, (_, index) => candidate(`probe-${index}`, 'allowed-source'))
)
queueTableRows(schemaMock.embedding, [{ id: 'partial' }])
queueTableRows(schemaMock.embedding, [candidate('partial', 'allowed-source')])
queueTableRows(schemaMock.embedding, [candidate('selected', 'allowed-source')])
queueTableRows(schemaMock.embedding, [
Expand All @@ -614,6 +620,10 @@ describe('live repository authorization follows ranked candidates', () => {
candidate(`approximate-${index}`, 'allowed-source')
)
queueTableRows(schemaMock.embedding, probe)
queueTableRows(
schemaMock.embedding,
approximate.map(({ id }) => ({ id }))
)
queueTableRows(schemaMock.embedding, approximate)
queueTableRows(schemaMock.embedding, [])
queueTableRows(schemaMock.embedding, probe)
Expand All @@ -628,7 +638,7 @@ describe('live repository authorization follows ranked candidates', () => {
const rows = await handleVectorOnlySearch({ ...params, structuredFilters: undefined })

expect(rows.map((row) => row.id)).toEqual(['selected'])
expect(dbChainMockFns.offset.mock.calls).toEqual([[0], [20], [0], [20]])
expect(dbChainMockFns.offset.mock.calls).toEqual([[0], [0], [20]])
expect(getForConnectors).toHaveBeenCalledTimes(2)
expect(dbChainMockFns.transaction).toHaveBeenCalledTimes(2)
})
Expand All @@ -638,6 +648,7 @@ describe('live repository authorization follows ranked candidates', () => {
candidate(`probe-${index}`, 'allowed-source')
)
queueTableRows(schemaMock.embedding, probe)
queueTableRows(schemaMock.embedding, [{ id: 'far' }])
queueTableRows(schemaMock.embedding, [
{ ...candidate('far', 'allowed-source'), distance: 0.7 },
...Array.from({ length: 19 }, (_, index) => candidate(`hidden-${index}`, 'allowed-source')),
Expand All @@ -662,7 +673,7 @@ describe('live repository authorization follows ranked candidates', () => {
})

expect(rows.map((row) => row.id)).toEqual(['nearer', 'near'])
expect(dbChainMockFns.offset.mock.calls).toEqual([[0], [20], [0]])
expect(dbChainMockFns.offset.mock.calls).toEqual([[0], [0]])
expect(
hasMockCondition(
dbChainMockFns.where.mock.calls.at(-1)![0],
Expand Down
Loading
Loading