Skip to content

feat(knowledge): add Live sync option to KB connectors + fix embedding billing #3959

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
waleedlatif1 merged 12 commits into from
Apr 4, 2026
Merged

feat(knowledge): add Live sync option to KB connectors + fix embedding billing #3959

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
12 commits
Select commit Hold shift + click to select a range
8cd90e7
feat(knowledge): add Live sync option to KB connector modal for Max/E…
waleedlatif1 Apr 3, 2026
987411c
fix(knowledge): record embedding usage cost for KB document processing
waleedlatif1 Apr 3, 2026
3347481
fix(knowledge): simplify embedding billing — use calculateCost, retur…
waleedlatif1 Apr 3, 2026
446c506
fix(knowledge): remove dev bypass from hasLiveSyncAccess
waleedlatif1 Apr 3, 2026
11d9f83
chore(knowledge): rename sync-intervals to consts, fix stale TSDoc co…
waleedlatif1 Apr 3, 2026
26682c1
improvement(knowledge): extract MaxBadge component, capture billing c…
waleedlatif1 Apr 3, 2026
47dbe1c
fix(knowledge): add knowledge-base to usage_log_source enum, fix docs…
waleedlatif1 Apr 3, 2026
0809ea1
fix(knowledge): generate migration for knowledge-base usage_log_sourc…
waleedlatif1 Apr 3, 2026
57f190a
fix(knowledge): add knowledge-base to usage_log_source enum via drizz…
waleedlatif1 Apr 3, 2026
4e6cb8f
fix(knowledge): fix search embedding test mocks, parallelize billing …
waleedlatif1 Apr 3, 2026
d3c7bff
fix(knowledge): warn when embedding model has no pricing entry
waleedlatif1 Apr 3, 2026
9e5d32d
fix(knowledge): call checkAndBillOverageThreshold after embedding usage
waleedlatif1 Apr 3, 2026
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
Prev Previous commit
Next Next commit
fix(knowledge): fix search embedding test mocks, parallelize billing … 
…lookups
  • Loading branch information
@waleedlatif1
waleedlatif1 committed Apr 4, 2026
commit 4e6cb8f46a709ad12f24277a08e44486b44dfa5c
103 changes: 43 additions & 60 deletions apps/sim/app/api/knowledge/search/utils.test.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@
* @vitest-environment node
*/
import { createEnvMock, databaseMock, loggerMock } from '@sim/testing'
import { mockNextFetchResponse } from '@sim/testing/mocks'
import { beforeEach, describe, expect, it, vi } from 'vitest'

vi.mock('drizzle-orm')
Expand All @@ -14,16 +15,6 @@ vi.mock('@/lib/knowledge/documents/utils', () => ({
retryWithExponentialBackoff: (fn: any) => fn(),
}))

vi.stubGlobal(
'fetch',
vi.fn().mockResolvedValue({
ok: true,
json: async () => ({
data: [{ embedding: [0.1, 0.2, 0.3] }],
}),
})
)

vi.mock('@/lib/core/config/env', () => createEnvMock())

import {
Expand Down Expand Up @@ -178,17 +169,16 @@ describe('Knowledge Search Utils', () => {
OPENAI_API_KEY: 'test-openai-key',
})

const fetchSpy = vi.mocked(fetch)
fetchSpy.mockResolvedValueOnce({
ok: true,
json: async () => ({
mockNextFetchResponse({
json: {
data: [{ embedding: [0.1, 0.2, 0.3] }],
}),
} as any)
usage: { prompt_tokens: 1, total_tokens: 1 },
},
})

const result = await generateSearchEmbedding('test query')

expect(fetchSpy).toHaveBeenCalledWith(
expect(vi.mocked(fetch)).toHaveBeenCalledWith(
'https://test.openai.azure.com/openai/deployments/text-embedding-ada-002/embeddings?api-version=2024-12-01-preview',
expect.objectContaining({
headers: expect.objectContaining({
Expand All @@ -209,17 +199,16 @@ describe('Knowledge Search Utils', () => {
OPENAI_API_KEY: 'test-openai-key',
})

const fetchSpy = vi.mocked(fetch)
fetchSpy.mockResolvedValueOnce({
ok: true,
json: async () => ({
mockNextFetchResponse({
json: {
data: [{ embedding: [0.1, 0.2, 0.3] }],
}),
} as any)
usage: { prompt_tokens: 1, total_tokens: 1 },
},
})

const result = await generateSearchEmbedding('test query')

expect(fetchSpy).toHaveBeenCalledWith(
expect(vi.mocked(fetch)).toHaveBeenCalledWith(
'https://api.openai.com/v1/embeddings',
expect.objectContaining({
headers: expect.objectContaining({
Expand All @@ -243,17 +232,16 @@ describe('Knowledge Search Utils', () => {
OPENAI_API_KEY: 'test-openai-key',
})

const fetchSpy = vi.mocked(fetch)
fetchSpy.mockResolvedValueOnce({
ok: true,
json: async () => ({
mockNextFetchResponse({
json: {
data: [{ embedding: [0.1, 0.2, 0.3] }],
}),
} as any)
usage: { prompt_tokens: 1, total_tokens: 1 },
},
})

await generateSearchEmbedding('test query')

expect(fetchSpy).toHaveBeenCalledWith(
expect(vi.mocked(fetch)).toHaveBeenCalledWith(
expect.stringContaining('api-version='),
expect.any(Object)
)
Expand All @@ -273,17 +261,16 @@ describe('Knowledge Search Utils', () => {
OPENAI_API_KEY: 'test-openai-key',
})

const fetchSpy = vi.mocked(fetch)
fetchSpy.mockResolvedValueOnce({
ok: true,
json: async () => ({
mockNextFetchResponse({
json: {
data: [{ embedding: [0.1, 0.2, 0.3] }],
}),
} as any)
usage: { prompt_tokens: 1, total_tokens: 1 },
},
})

await generateSearchEmbedding('test query', 'text-embedding-3-small')

expect(fetchSpy).toHaveBeenCalledWith(
expect(vi.mocked(fetch)).toHaveBeenCalledWith(
'https://test.openai.azure.com/openai/deployments/custom-embedding-model/embeddings?api-version=2024-12-01-preview',
expect.any(Object)
)
Expand Down Expand Up @@ -311,13 +298,12 @@ describe('Knowledge Search Utils', () => {
KB_OPENAI_MODEL_NAME: 'text-embedding-ada-002',
})

const fetchSpy = vi.mocked(fetch)
fetchSpy.mockResolvedValueOnce({
mockNextFetchResponse({
ok: false,
status: 404,
statusText: 'Not Found',
text: async () => 'Deployment not found',
} as any)
text: 'Deployment not found',
})

await expect(generateSearchEmbedding('test query')).rejects.toThrow('Embedding API failed')

Expand All @@ -332,13 +318,12 @@ describe('Knowledge Search Utils', () => {
OPENAI_API_KEY: 'test-openai-key',
})

const fetchSpy = vi.mocked(fetch)
fetchSpy.mockResolvedValueOnce({
mockNextFetchResponse({
ok: false,
status: 429,
statusText: 'Too Many Requests',
text: async () => 'Rate limit exceeded',
} as any)
text: 'Rate limit exceeded',
})

await expect(generateSearchEmbedding('test query')).rejects.toThrow('Embedding API failed')

Expand All @@ -356,17 +341,16 @@ describe('Knowledge Search Utils', () => {
KB_OPENAI_MODEL_NAME: 'text-embedding-ada-002',
})

const fetchSpy = vi.mocked(fetch)
fetchSpy.mockResolvedValueOnce({
ok: true,
json: async () => ({
mockNextFetchResponse({
json: {
data: [{ embedding: [0.1, 0.2, 0.3] }],
}),
} as any)
usage: { prompt_tokens: 1, total_tokens: 1 },
},
})

await generateSearchEmbedding('test query')

expect(fetchSpy).toHaveBeenCalledWith(
expect(vi.mocked(fetch)).toHaveBeenCalledWith(
expect.any(String),
expect.objectContaining({
body: JSON.stringify({
Expand All @@ -387,17 +371,16 @@ describe('Knowledge Search Utils', () => {
OPENAI_API_KEY: 'test-openai-key',
})

const fetchSpy = vi.mocked(fetch)
fetchSpy.mockResolvedValueOnce({
ok: true,
json: async () => ({
mockNextFetchResponse({
json: {
data: [{ embedding: [0.1, 0.2, 0.3] }],
}),
} as any)
usage: { prompt_tokens: 1, total_tokens: 1 },
},
})

await generateSearchEmbedding('test query', 'text-embedding-3-small')

expect(fetchSpy).toHaveBeenCalledWith(
expect(vi.mocked(fetch)).toHaveBeenCalledWith(
expect.any(String),
expect.objectContaining({
body: JSON.stringify({
Expand Down
12 changes: 8 additions & 4 deletions apps/sim/lib/billing/core/subscription.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -448,9 +448,11 @@ export async function hasInboxAccess(userId: string): Promise<boolean> {
if (!isProd) {
return true
}
const sub = await getHighestPrioritySubscription(userId)
const [sub, billingStatus] = await Promise.all([
getHighestPrioritySubscription(userId),
getEffectiveBillingStatus(userId),
])
if (!sub) return false
const billingStatus = await getEffectiveBillingStatus(userId)
if (!hasUsableSubscriptionAccess(sub.status, billingStatus.billingBlocked)) return false
return getPlanTierCredits(sub.plan) >= 25000 || checkEnterprisePlan(sub)
} catch (error) {
Expand All @@ -470,9 +472,11 @@ export async function hasLiveSyncAccess(userId: string): Promise<boolean> {
if (!isHosted) {
return true
}
const sub = await getHighestPrioritySubscription(userId)
const [sub, billingStatus] = await Promise.all([
getHighestPrioritySubscription(userId),
getEffectiveBillingStatus(userId),
])
if (!sub) return false
const billingStatus = await getEffectiveBillingStatus(userId)
if (!hasUsableSubscriptionAccess(sub.status, billingStatus.billingBlocked)) return false
return getPlanTierCredits(sub.plan) >= 25000 || checkEnterprisePlan(sub)
} catch (error) {
Expand Down