Skip to content

Commit 4d3f266

Browse files
committed
improvement(knowledge): choose the vector plan by the caller's reach
- a caller who reaches at least a quarter of the bases' documents walks the whole graph once, as before: pgvector post-filters, so at that share almost every neighbour the graph visits is theirs and one walk is the cheapest exact answer; searching each source on its own only helps a caller whose reach is a small share of the index - the share is counted once against that bound when the probe saturates, and remembered with the saturation - a bounded caller who is a member of an indexed source walks that source instead of ranking it exactly with the rest
1 parent 3ec5f2e commit 4d3f266

3 files changed

Lines changed: 170 additions & 16 deletions

File tree

‎apps/sim/lib/knowledge/search/queries.test.ts‎

Lines changed: 78 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -44,6 +44,7 @@ import {
4444
VECTOR_PROBE_DOCUMENT_LIMIT,
4545
visibleDocumentsQuery,
4646
} from '@/lib/knowledge/search/queries'
47+
import { forgetIndexedVectorSources } from '@/lib/knowledge/search/source-vector-indexes'
4748
import type { StructuredFilter } from '@/lib/knowledge/types'
4849

4950
/**
@@ -1235,6 +1236,7 @@ describe('permitted-document planner', () => {
12351236
rerankRows = []
12361237
sourceExactRows = []
12371238
indexedSourceRows = []
1239+
forgetIndexedVectorSources()
12381240
dbChainMockFns.execute.mockImplementation(async (query) => {
12391241
const statement = render(query).sql
12401242
if (statement.includes('pg_index')) return indexedSourceRows
@@ -1274,13 +1276,61 @@ describe('permitted-document planner', () => {
12741276
traversedRows = [{ id: 'a' }]
12751277
rerankRows = [hit('a', null)]
12761278
queueTableRows(schemaMock.embedding, [hit('a', null)])
1277-
await handleVectorOnlySearch({ ...params, permitted: { kind: 'unbounded' } })
1279+
await handleVectorOnlySearch({ ...params, permitted: { kind: 'unbounded', broad: false } })
12781280
const sqls = statements().map((query) => query.sql)
12791281
expect(sqls.some((sql) => sql.includes('AS visible'))).toBe(true)
12801282
expect(sqls.some(isProbeStatement)).toBe(false)
12811283
expect(sqls.some(isExactRanking)).toBe(false)
12821284
})
12831285

1286+
it('walks the whole graph once for a caller whose reach is broad', async () => {
1287+
const eligibility = { workspace: [], admin: ['other-src'], members: ['member-src'] }
1288+
indexedSourceRows = [{ name: 'idx', connectorId: 'member-src' }]
1289+
traversedRows = [{ id: 'walked-hit', distance: 0.2 }]
1290+
rerankRows = [hit('walked-hit', 'member-src')]
1291+
queueTableRows(schemaMock.embedding, rerankRows)
1292+
await handleVectorOnlySearch({
1293+
...params,
1294+
permitted: { kind: 'unbounded', broad: true },
1295+
accessPlan: {
1296+
connectors: eligibility,
1297+
observers: { confirmed: ['m-1'], observed: [] },
1298+
memberSources: ['member-src'],
1299+
},
1300+
})
1301+
/** One walk over every source, scoped to the bases alone — no source is singled out. */
1302+
const walks = statements().filter((query) => query.sql.includes('AS visible'))
1303+
expect(walks).toHaveLength(1)
1304+
expect(JSON.stringify(walks[0])).not.toContain('"right":"member-src"')
1305+
expect(statements().some((query) => query.sql.includes('WITH readable_documents'))).toBe(false)
1306+
})
1307+
1308+
it('walks an indexed source a bounded caller is a member of instead of ranking it exactly', async () => {
1309+
const eligibility = { workspace: [], admin: ['small-src'], members: ['member-src'] }
1310+
indexedSourceRows = [{ name: 'idx', connectorId: 'member-src' }]
1311+
sourceExactRows = [{ id: 'small-hit', distance: 0.3, saturated: false }]
1312+
traversedRows = [{ id: 'walked-hit', distance: 0.2 }]
1313+
rerankRows = [hit('walked-hit', 'member-src'), hit('small-hit', 'small-src')]
1314+
queueTableRows(schemaMock.embedding, rerankRows)
1315+
await handleVectorOnlySearch({
1316+
...params,
1317+
topK: 2,
1318+
permitted: bounded(
1319+
{ id: 'doc-a', connectorId: 'member-src' },
1320+
{ id: 'doc-b', connectorId: 'small-src' }
1321+
),
1322+
accessPlan: {
1323+
connectors: eligibility,
1324+
observers: { confirmed: ['m-1'], observed: [] },
1325+
memberSources: ['member-src'],
1326+
},
1327+
})
1328+
const walks = statements().filter((query) => query.sql.includes('AS visible'))
1329+
expect(walks).toHaveLength(1)
1330+
expect(JSON.stringify(walks[0])).toContain('"right":"member-src"')
1331+
expect(statements().some((q) => isExactRanking(q.sql))).toBe(false)
1332+
})
1333+
12841334
it('walks a source the caller is a member of and ranks every other source exactly', async () => {
12851335
const eligibility = {
12861336
workspace: [],
@@ -1297,7 +1347,7 @@ describe('permitted-document planner', () => {
12971347
queueTableRows(schemaMock.embedding, rerankRows)
12981348
await handleVectorOnlySearch({
12991349
...params,
1300-
permitted: { kind: 'unbounded' },
1350+
permitted: { kind: 'unbounded', broad: false },
13011351
accessPlan: { connectors: eligibility, observers, memberSources },
13021352
})
13031353
const walks = statements().filter((query) => query.sql.includes('AS visible'))
@@ -1320,7 +1370,7 @@ describe('permitted-document planner', () => {
13201370
queueTableRows(schemaMock.embedding, rerankRows)
13211371
await handleVectorOnlySearch({
13221372
...params,
1323-
permitted: { kind: 'unbounded' },
1373+
permitted: { kind: 'unbounded', broad: false },
13241374
accessPlan: {
13251375
connectors: eligibility,
13261376
observers: { confirmed: [], observed: [] },
@@ -1347,7 +1397,7 @@ describe('permitted-document planner', () => {
13471397
await handleVectorOnlySearch({
13481398
...params,
13491399
topK: 2,
1350-
permitted: { kind: 'unbounded' },
1400+
permitted: { kind: 'unbounded', broad: false },
13511401
accessPlan: {
13521402
connectors: eligibility,
13531403
observers: { confirmed: ['m-1'], observed: [] },
@@ -1368,7 +1418,7 @@ describe('permitted-document planner', () => {
13681418
queueTableRows(schemaMock.embedding, rerankRows)
13691419
await handleVectorOnlySearch({
13701420
...params,
1371-
permitted: { kind: 'unbounded' },
1421+
permitted: { kind: 'unbounded', broad: false },
13721422
accessPlan: {
13731423
connectors: { workspace: [], admin: ['sliced-src'], members: [] },
13741424
observers: { confirmed: [], observed: [] },
@@ -1547,6 +1597,29 @@ describe('permitted-document planner', () => {
15471597
resolvePermittedDocuments({ knowledgeBaseIds, access })
15481598
const probes = () => statements().filter((query) => isProbeStatement(query.sql)).length
15491599

1600+
it('counts a saturated reach against the broad bound once, and remembers the answer', async () => {
1601+
/** The index holds a million documents; the bound is a quarter of them. */
1602+
const counts = { index: 1_000_000, reached: 250_000 }
1603+
dbChainMockFns.execute.mockImplementation(async (query) => {
1604+
const statement = render(query).sql
1605+
if (isProbeStatement(statement)) return [{ id: null, connectorId: null, saturated: true }]
1606+
if (statement.includes(') reached')) return [{ n: counts.reached }]
1607+
if (statement.includes('count(*) AS n')) return [{ n: counts.index }]
1608+
return []
1609+
})
1610+
const reachCounts = () => statements().filter((query) => query.sql.includes(') reached'))
1611+
const broad = await resolve(scope('broad-reach'))
1612+
expect(broad).toEqual({ kind: 'unbounded', broad: true })
1613+
expect(reachCounts()).toHaveLength(1)
1614+
expect(JSON.stringify(reachCounts()[0])).toContain('250000')
1615+
await resolve(scope('broad-reach'))
1616+
expect(reachCounts()).toHaveLength(1)
1617+
counts.reached = 120_000
1618+
const narrow = await resolve(scope('narrow-reach'))
1619+
expect(narrow).toEqual({ kind: 'unbounded', broad: false })
1620+
expect(reachCounts()).toHaveLength(2)
1621+
})
1622+
15501623
it('is remembered, so a broad caller skips the probe on the next search', async () => {
15511624
probeRows = [{ id: null, connectorId: null, saturated: true }]
15521625
const broad = scope('broad')

‎apps/sim/lib/knowledge/search/queries.ts‎

Lines changed: 85 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -1046,7 +1046,16 @@ type PermittedDocument = {
10461046
*/
10471047
export type PermittedDocuments =
10481048
| { kind: 'bounded'; documents: readonly PermittedDocument[] }
1049-
| { kind: 'unbounded' }
1049+
| { kind: 'unbounded'; broad: boolean }
1050+
1051+
/**
1052+
* The share of the index a caller must reach before the whole graph is walked for them. pgvector
1053+
* post-filters, so a walk returns a caller's own neighbours in proportion to their reach: above
1054+
* this share almost every neighbour the graph visits is theirs and one walk is the cheapest exact
1055+
* answer there is; below it the walk spends its budget on chunks they cannot read, and each
1056+
* readable source is searched on its own instead.
1057+
*/
1058+
export const BROAD_REACH_SHARE = 0.25
10501059

10511060
/**
10521061
* How long a caller's saturated reach is remembered. Reach counts the documents a caller's tokens
@@ -1055,7 +1064,51 @@ export type PermittedDocuments =
10551064
*/
10561065
const SATURATED_REACH_TTL_MS = 5 * 60 * 1000
10571066

1058-
const saturatedReach = new LRUCache<string, true>({ max: 10_000, ttl: SATURATED_REACH_TTL_MS })
1067+
/** A saturated reach, and whether it is broad enough to walk the whole graph for. */
1068+
const saturatedReach = new LRUCache<string, { broad: boolean }>({
1069+
max: 10_000,
1070+
ttl: SATURATED_REACH_TTL_MS,
1071+
})
1072+
1073+
/** How many documents the bases hold: the denominator of a reach share, and it moves slowly. */
1074+
const indexDocumentCounts = new LRUCache<string, number>({
1075+
max: 1000,
1076+
ttl: SATURATED_REACH_TTL_MS,
1077+
fetchMethod: async (key) => {
1078+
const [row] = await db.execute<{ n: number }>(sql`
1079+
SELECT count(*) AS n FROM ${document}
1080+
WHERE ${document.knowledgeBaseId} = ANY(${textArrayLiteral(key.split(','))})
1081+
AND ${document.deletedAt} IS NULL`)
1082+
/** An empty answer is not remembered; the bases may simply not have been read yet. */
1083+
return Number(row?.n ?? 0) || undefined
1084+
},
1085+
})
1086+
1087+
/**
1088+
* Whether a saturated reach is broad: the caller reaches at least {@link BROAD_REACH_SHARE} of the
1089+
* bases' documents. Counted once against that bound and remembered with the saturation, so the
1090+
* first search after the window pays for it and the rest do not.
1091+
*/
1092+
async function reachIsBroad(
1093+
knowledgeBaseIds: string[],
1094+
access: KnowledgeAccessScope,
1095+
budget: SearchBudget | undefined
1096+
): Promise<boolean> {
1097+
if (access.kind !== 'user') return true
1098+
const total = (await indexDocumentCounts.fetch([...knowledgeBaseIds].sort().join(','))) ?? 0
1099+
const bound = Math.ceil(total * BROAD_REACH_SHARE)
1100+
if (bound <= VECTOR_PROBE_DOCUMENT_LIMIT) return true
1101+
const [row] = await runSearchQuery(budget, 'permitted_documents', (executor) =>
1102+
executor.execute<{ n: number }>(sql`
1103+
SELECT count(*) AS n FROM (
1104+
SELECT 1 FROM ${document}
1105+
WHERE ${document.deletedAt} IS NULL AND ${knowledgeAclOverlapCondition(access)}
1106+
AND ${inArray(document.knowledgeBaseId, knowledgeBaseIds)}
1107+
LIMIT ${bound}
1108+
) reached`)
1109+
)
1110+
return Number(row?.n ?? 0) >= bound
1111+
}
10591112

10601113
/** Reach depends only on the bases and the caller's tokens; filters narrow the set, not the reach. */
10611114
function reachKey(
@@ -1083,8 +1136,11 @@ export async function resolvePermittedDocuments(params: {
10831136
}): Promise<PermittedDocuments> {
10841137
const key = reachKey(params.knowledgeBaseIds, params.access)
10851138
let probe: ProbeOutcome
1086-
if (key && saturatedReach.get(key)) {
1139+
let broad = true
1140+
const remembered = key ? saturatedReach.get(key) : undefined
1141+
if (remembered) {
10871142
probe = { kind: 'saturated' }
1143+
broad = remembered.broad
10881144
} else {
10891145
try {
10901146
probe = await probeVisibleDocuments(
@@ -1103,12 +1159,19 @@ export async function resolvePermittedDocuments(params: {
11031159
if (!params.budget?.isTimeout(error)) throw error
11041160
probe = { kind: 'timed_out' }
11051161
}
1106-
if (key && probe.kind === 'saturated') saturatedReach.set(key, true)
1162+
if (probe.kind === 'saturated') {
1163+
try {
1164+
broad = await reachIsBroad(params.knowledgeBaseIds, params.access, params.budget)
1165+
} catch (error) {
1166+
if (!params.budget?.isTimeout(error)) throw error
1167+
}
1168+
if (key) saturatedReach.set(key, { broad })
1169+
}
11071170
}
11081171
const permitted: PermittedDocuments =
11091172
probe.kind === 'documents'
11101173
? { kind: 'bounded', documents: probe.documents }
1111-
: { kind: 'unbounded' }
1174+
: { kind: 'unbounded', broad }
11121175
annotateSearchDiagnostics({
11131176
permittedDocuments: permitted.kind,
11141177
...(probe.kind === 'documents' ? { permittedDocumentCount: probe.documents.length } : {}),
@@ -1362,6 +1425,10 @@ async function selectVectorResults(params: SearchParams): Promise<SearchResult[]
13621425
}
13631426
if (params.filters?.documentIds?.length) return exactPage()
13641427
if (candidatePool?.excludedKey !== excludedKey) {
1428+
const plannedIndexedSources =
1429+
params.access.kind === 'user' && params.accessPlan?.memberSources.length
1430+
? await indexedVectorSources()
1431+
: undefined
13651432
annotateSearchDiagnostics({
13661433
vectorRanking: 'candidate-rerank',
13671434
vectorCandidateStorage: 'stored-halfvec',
@@ -1395,24 +1462,33 @@ async function selectVectorResults(params: SearchParams): Promise<SearchResult[]
13951462
)
13961463
}
13971464
let selected: Array<{ id: string }>
1398-
if (params.permitted?.kind === 'bounded') {
1465+
const plan = params.access.kind === 'user' ? params.accessPlan : undefined
1466+
/**
1467+
* A source the caller is a member of that has its own index is walked on its own, which
1468+
* beats ranking it exactly once it is large enough to have earned that index.
1469+
*/
1470+
const walksASource = plan?.memberSources.some(
1471+
(id) => plannedIndexedSources?.has(id) ?? false
1472+
)
1473+
if (params.permitted?.kind === 'bounded' && !walksASource) {
13991474
/**
14001475
* A bounded permitted set is ranked exactly without walking the graph first: the walk
14011476
* post-filters, so when the caller reads a small share of the index it spends its whole
14021477
* uninterruptible tuple budget and still returns almost none of their neighbours.
14031478
*/
14041479
selected = await rankPermittedExactly(params.permitted.documents.map((entry) => entry.id))
1405-
} else if (params.accessPlan && params.access.kind === 'user') {
1480+
} else if (plan && !(params.permitted?.kind === 'unbounded' && params.permitted.broad)) {
14061481
/**
14071482
* Readability follows sources, so each readable source is searched in its own index and
14081483
* the results merged. A member reads a source whole or barely at all: walking one source
14091484
* spends its budget among chunks they can read, where a walk over every source spends it
1410-
* on the sources they cannot.
1485+
* on the sources they cannot. A caller whose reach is broad skips this: for them the
1486+
* whole graph's neighbours are mostly theirs already, and one walk is the cheaper answer.
14111487
*/
14121488
selected = await selectSourceVectorCandidates({
14131489
access: params.access,
14141490
knowledgeBaseIds: params.knowledgeBaseIds,
1415-
plan: params.accessPlan,
1491+
plan,
14161492
documentConditions: candidateDocumentVisibility,
14171493
tagCondition: candidateTagCondition,
14181494
documentTagCondition,

‎apps/sim/lib/knowledge/search/source-vector-indexes.ts‎

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,11 @@ function indexName(connectorId: string): string {
3131
*/
3232
const indexedSources = new LRUCache<'sources', ReadonlySet<string>>({ max: 1, ttl: 60 * 1000 })
3333

34+
/** Forgets the cached answer, after an index was built or dropped. */
35+
export function forgetIndexedVectorSources(): void {
36+
indexedSources.clear()
37+
}
38+
3439
export async function indexedVectorSources(): Promise<ReadonlySet<string>> {
3540
const cached = indexedSources.get('sources')
3641
if (cached) return cached
@@ -120,7 +125,7 @@ export async function ensureSourceVectorIndex(connectorId: string): Promise<bool
120125
documents,
121126
elapsedMs: Date.now() - startedAt,
122127
})
123-
indexedSources.clear()
128+
forgetIndexedVectorSources()
124129
return true
125130
} catch (error) {
126131
logger.error('Source vector index build failed', { connectorId, error: getErrorMessage(error) })
@@ -136,5 +141,5 @@ export async function ensureSourceVectorIndex(connectorId: string): Promise<bool
136141
export async function dropSourceVectorIndex(connectorId: string): Promise<void> {
137142
if (!CONNECTOR_ID.test(connectorId)) return
138143
await db.execute(sql.raw(`DROP INDEX CONCURRENTLY IF EXISTS "${indexName(connectorId)}"`))
139-
indexedSources.clear()
144+
forgetIndexedVectorSources()
140145
}

0 commit comments

Comments
 (0)