Skip to content

Commit c03ce29

Browse files
waleedlatif1waleedlatif1
authored
improvement(search): reduce retrieval overhead (#7653)
* Improve organization search ranking and bound rerank admission * improvement(search): reduce embedding and retrieval overhead * fix(search): align boundary test fixtures with search adapters * improvement(search): preserve existing ranking policy
1 parent 4d5e4aa commit c03ce29

25 files changed

Lines changed: 26880 additions & 260 deletions

‎apps/sim/app/api/knowledge/search/utils.test.ts‎

Lines changed: 18 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@ import { ResolvedSecretTraceRegistry } from '@/executor/utils/resolved-secret-tr
2121
vi.mock('@/lib/core/rate-limiter/provider-admission', () => ({
2222
PROVIDER_QUOTA_COOLDOWN_MS: 300_000,
2323
ProviderQuotaExhaustedError: class ProviderQuotaExhaustedError extends Error {},
24+
ProviderAdmissionTimeoutError: class ProviderAdmissionTimeoutError extends Error {},
2425
isProviderQuotaExhausted: vi.fn().mockResolvedValue(false),
2526
recordProviderCooldown: vi.fn().mockResolvedValue(undefined),
2627
waitForProviderAdmission: vi.fn().mockResolvedValue(undefined),
@@ -105,18 +106,23 @@ function makeResult(id: string, distance = 0.1): SearchResult {
105106
}
106107
}
107108

108-
const TEST_EMBEDDING = [0.1, 0.2, 0.3, ...Array.from({ length: 1533 }, () => 0)]
109+
const TEST_EMBEDDING = [0.1, 0.2, 0.3, ...Array.from({ length: 1533 }, () => 0)].map(Math.fround)
109110

110111
function mockNextEmbeddingResponse(): void {
111-
vi.mocked(fetch).mockResolvedValueOnce(
112-
new Response(
112+
vi.mocked(fetch).mockImplementationOnce(async (_url, init) => {
113+
const request = JSON.parse(String(init?.body))
114+
const embedding =
115+
request.encoding_format === 'base64'
116+
? Buffer.from(new Float32Array(TEST_EMBEDDING).buffer).toString('base64')
117+
: TEST_EMBEDDING
118+
return new Response(
113119
JSON.stringify({
114-
data: [{ embedding: TEST_EMBEDDING, index: 0 }],
120+
data: [{ embedding, index: 0 }],
115121
usage: { prompt_tokens: 1, total_tokens: 1 },
116122
}),
117123
{ status: 200, headers: { 'Content-Type': 'application/json' } }
118124
)
119-
)
125+
})
120126
}
121127

122128
describe('Knowledge Search Utils', () => {
@@ -220,7 +226,8 @@ describe('Knowledge Search Utils', () => {
220226
})
221227

222228
expect(results.map((row) => row.id)).toEqual(['first', 'second'])
223-
expect(dbChainMockFns.select).toHaveBeenCalledTimes(1)
229+
expect(dbChainMockFns.select).toHaveBeenCalledTimes(2)
230+
expect(dbChainMockFns.as).toHaveBeenCalledWith('ranked_embeddings')
224231
expect(dbChainMockFns.select.mock.calls[0][0]).toHaveProperty('distance')
225232
expect(dbChainMockFns.limit).toHaveBeenCalledWith(2)
226233
})
@@ -541,7 +548,8 @@ describe('Knowledge Search Utils', () => {
541548
})
542549

543550
expect(results.map((r) => r.id)).toEqual(['vector-hit'])
544-
expect(dbChainMockFns.select).toHaveBeenCalledTimes(1)
551+
expect(dbChainMockFns.select).toHaveBeenCalledTimes(2)
552+
expect(dbChainMockFns.as).toHaveBeenCalledWith('ranked_embeddings')
545553
})
546554

547555
it('runs both legs and fuses them in hybrid mode', async () => {
@@ -565,7 +573,7 @@ describe('Knowledge Search Utils', () => {
565573
})
566574

567575
expect(results.map((r) => r.id).sort()).toEqual(['keyword-hit', 'vector-hit'])
568-
expect(dbChainMockFns.select).toHaveBeenCalledTimes(3)
576+
expect(dbChainMockFns.select).toHaveBeenCalledTimes(4)
569577
})
570578

571579
it('falls back to vector results when the keyword leg fails', async () => {
@@ -830,7 +838,7 @@ describe('Knowledge Search Utils', () => {
830838
body: JSON.stringify({
831839
input: ['test query'],
832840
model: 'text-embedding-3-small',
833-
encoding_format: 'float',
841+
encoding_format: 'base64',
834842
dimensions: 1536,
835843
}),
836844
})
@@ -860,7 +868,7 @@ describe('Knowledge Search Utils', () => {
860868
body: JSON.stringify({
861869
input: ['prefix {{TOKEN}} suffix'],
862870
model: 'text-embedding-3-small',
863-
encoding_format: 'float',
871+
encoding_format: 'base64',
864872
dimensions: 1536,
865873
}),
866874
})

‎apps/sim/app/api/knowledge/utils.test.ts‎

Lines changed: 12 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@ import * as workspacesUtilsModule from '@/lib/workspaces/utils'
2323
vi.mock('@/lib/core/rate-limiter/provider-admission', () => ({
2424
PROVIDER_QUOTA_COOLDOWN_MS: 300_000,
2525
ProviderQuotaExhaustedError: class ProviderQuotaExhaustedError extends Error {},
26+
ProviderAdmissionTimeoutError: class ProviderAdmissionTimeoutError extends Error {},
2627
isProviderQuotaExhausted: vi.fn().mockResolvedValue(false),
2728
recordProviderCooldown: vi.fn().mockResolvedValue(undefined),
2829
waitForProviderAdmission: vi.fn().mockResolvedValue(undefined),
@@ -125,18 +126,24 @@ function createTestEmbedding(value: number): number[] {
125126
return Array.from({ length: TEST_EMBEDDING_DIMENSION }, () => value)
126127
}
127128

128-
function createEmbeddingResponse(values: number[]): Response {
129+
function createEmbeddingResponse(values: number[], encoding: 'base64' | 'float'): Response {
129130
return new Response(
130131
JSON.stringify({
131-
data: values.map((value, index) => ({ embedding: createTestEmbedding(value), index })),
132+
data: values.map((value, index) => ({
133+
embedding:
134+
encoding === 'base64'
135+
? Buffer.from(new Float32Array(createTestEmbedding(value)).buffer).toString('base64')
136+
: createTestEmbedding(value),
137+
index,
138+
})),
132139
usage: { prompt_tokens: values.length, total_tokens: values.length },
133140
}),
134141
{ status: 200, headers: { 'Content-Type': 'application/json' } }
135142
)
136143
}
137144

138145
function createEmbeddingFetchMock() {
139-
return vi.fn().mockResolvedValue(createEmbeddingResponse([0.1, 0.3]))
146+
return vi.fn().mockResolvedValue(createEmbeddingResponse([0.1, 0.3], 'base64'))
140147
}
141148

142149
vi.stubGlobal('fetch', createEmbeddingFetchMock())
@@ -286,7 +293,7 @@ describe('Knowledge Utils', () => {
286293
})
287294

288295
const fetchSpy = vi.mocked(fetch)
289-
fetchSpy.mockResolvedValueOnce(createEmbeddingResponse([0.1]))
296+
fetchSpy.mockResolvedValueOnce(createEmbeddingResponse([0.1], 'float'))
290297

291298
await generateEmbeddings(['test text'], DEFAULT_EMBEDDING_TARGET)
292299

@@ -310,7 +317,7 @@ describe('Knowledge Utils', () => {
310317
})
311318

312319
const fetchSpy = vi.mocked(fetch)
313-
fetchSpy.mockResolvedValueOnce(createEmbeddingResponse([0.1]))
320+
fetchSpy.mockResolvedValueOnce(createEmbeddingResponse([0.1], 'base64'))
314321

315322
await generateEmbeddings(['test text'], DEFAULT_EMBEDDING_TARGET)
316323

‎apps/sim/app/api/v2/knowledge/search/route.provenance.test.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -48,7 +48,7 @@ vi.mock('@/lib/knowledge/application/contexts', () => ({
4848
}))
4949

5050
vi.mock('@/lib/knowledge/service', () => ({
51-
getKnowledgeBaseById: mocks.getKnowledgeBase,
51+
getActiveKnowledgeBaseReference: mocks.getKnowledgeBase,
5252
}))
5353

5454
vi.mock('@/lib/knowledge/embeddings', () => ({

‎apps/sim/lib/core/rate-limiter/storage/db-token-bucket.test.ts‎

Lines changed: 48 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -47,4 +47,52 @@ describe('PostgreSQL token bucket', () => {
4747
})
4848
expect(dbChainMockFns.set).toHaveBeenCalledWith(expect.objectContaining({ tokens: '0' }))
4949
})
50+
51+
it('initializes request and cooldown buckets in a consistent order without duplicate keys', async () => {
52+
const now = new Date()
53+
dbChainMockFns.limit.mockResolvedValue([
54+
{ key: 'cooldown', tokens: '0', lastRefillAt: now, blockedUntil: null },
55+
{ key: 'requests', tokens: '10', lastRefillAt: now, blockedUntil: null },
56+
{ key: 'tokens', tokens: '10', lastRefillAt: now, blockedUntil: null },
57+
])
58+
59+
expect(
60+
await new DbTokenBucket().consumeTokensAtomically(
61+
[
62+
{ key: 'tokens', cost: 2, config: CONFIG },
63+
{ key: 'requests', cost: 1, config: CONFIG },
64+
],
65+
{ cooldownKeys: ['cooldown', 'cooldown'], deadlineAt: now.getTime() + 1000 }
66+
)
67+
).toEqual({ allowed: true, retryAfterMs: 0 })
68+
69+
expect(dbChainMockFns.values).toHaveBeenCalledExactlyOnceWith([
70+
{ key: 'cooldown', tokens: '0', lastRefillAt: now, updatedAt: now },
71+
{ key: 'requests', tokens: '10', lastRefillAt: now, updatedAt: now },
72+
{ key: 'tokens', tokens: '10', lastRefillAt: now, updatedAt: now },
73+
])
74+
expect(dbChainMockFns.set).toHaveBeenCalledWith({
75+
tokens: '8',
76+
lastRefillAt: now,
77+
updatedAt: now,
78+
})
79+
expect(dbChainMockFns.set).toHaveBeenCalledWith({
80+
tokens: '9',
81+
lastRefillAt: now,
82+
updatedAt: now,
83+
})
84+
})
85+
86+
it('accepts an empty reservation without attempting an empty insert', async () => {
87+
dbChainMockFns.limit.mockResolvedValue([])
88+
89+
expect(
90+
await new DbTokenBucket().consumeTokensAtomically([], {
91+
cooldownKeys: [],
92+
deadlineAt: Date.now() + 1000,
93+
})
94+
).toEqual({ allowed: true, retryAfterMs: 0 })
95+
expect(dbChainMockFns.insert).not.toHaveBeenCalled()
96+
expect(dbChainMockFns.update).not.toHaveBeenCalled()
97+
})
5098
})

‎apps/sim/lib/core/rate-limiter/storage/db-token-bucket.ts‎

Lines changed: 12 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -94,16 +94,20 @@ export class DbTokenBucket implements RateLimitStorageAdapter {
9494
...new Set([...options.cooldownKeys, ...reservations.map((item) => item.key)]),
9595
].sort()
9696
const createdAt = new Date()
97-
for (const key of keys) {
98-
const reservation = reservations.find((item) => item.key === key)
97+
if (keys.length > 0) {
9998
await tx
10099
.insert(rateLimitBucket)
101-
.values({
102-
key,
103-
tokens: String(reservation?.config.maxTokens ?? 0),
104-
lastRefillAt: createdAt,
105-
updatedAt: createdAt,
106-
})
100+
.values(
101+
keys.map((key) => {
102+
const reservation = reservations.find((item) => item.key === key)
103+
return {
104+
key,
105+
tokens: String(reservation?.config.maxTokens ?? 0),
106+
lastRefillAt: createdAt,
107+
updatedAt: createdAt,
108+
}
109+
})
110+
)
107111
.onConflictDoNothing()
108112
}
109113
const rows = await tx

0 commit comments

Comments
 (0)