Merge pull request #15857 from Budibase/budibase-ai-3

Create new global DB for self-host quota.
This commit is contained in:
Sam Rose 2025-04-01 11:38:22 +01:00 committed by GitHub
commit 32fc150fac
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
6 changed files with 170 additions and 86 deletions

View File

@ -96,6 +96,24 @@ async function get<T extends Document>(db: Database, id: string): Promise<T> {
return cacheItem.doc
}
async function tryGet<T extends Document>(
db: Database,
id: string
): Promise<T | null> {
const cache = await getCache()
const cacheKey = makeCacheKey(db, id)
let cacheItem: CacheItem<T> | null = await cache.get(cacheKey)
if (!cacheItem) {
const doc = await db.tryGet<T>(id)
if (!doc) {
return null
}
cacheItem = makeCacheItem(doc)
await cache.store(cacheKey, cacheItem)
}
return cacheItem.doc
}
async function remove(db: Database, docOrId: any, rev?: any): Promise<void> {
const cache = await getCache()
if (!docOrId) {
@ -123,10 +141,17 @@ export class Writethrough {
return put(this.db, doc, writeRateMs)
}
/**
* @deprecated use `tryGet` instead
*/
async get<T extends Document>(id: string) {
return get<T>(this.db, id)
}
async tryGet<T extends Document>(id: string) {
return tryGet<T>(this.db, id)
}
async remove(docOrId: any, rev?: any) {
return remove(this.db, docOrId, rev)
}

View File

@ -60,6 +60,11 @@ export const StaticDatabases = {
SCIM_LOGS: {
name: "scim-logs",
},
// Used by self-host users making use of Budicloud resources. Introduced when
// we started letting self-host users use Budibase AI in the cloud.
SELF_HOST_CLOUD: {
name: "self-host-cloud",
},
}
export const APP_PREFIX = prefixed(DocumentType.APP)

View File

@ -157,6 +157,33 @@ export async function doInTenant<T>(
return newContext(updates, task)
}
// We allow self-host licensed users to make use of some Budicloud services
// (e.g. Budibase AI). When they do this, they use their license key as an API
// key. We use that license key to identify the tenant ID, and we set the
// context to be self-host using cloud. This affects things like where their
// quota documents get stored (because we want to avoid creating a new global
// DB for each self-host tenant).
export async function doInSelfHostTenantUsingCloud<T>(
tenantId: string,
task: () => T
): Promise<T> {
const updates = { tenantId, isSelfHostUsingCloud: true }
return newContext(updates, task)
}
export function isSelfHostUsingCloud() {
const context = Context.get()
return !!context?.isSelfHostUsingCloud
}
export function getSelfHostCloudDB() {
const context = Context.get()
if (!context || !context.isSelfHostUsingCloud) {
throw new Error("Self-host cloud DB not found")
}
return getDB(StaticDatabases.SELF_HOST_CLOUD.name)
}
export async function doInAppContext<T>(
appId: string,
task: () => T
@ -325,6 +352,11 @@ export function getGlobalDB(): Database {
if (!context || (env.MULTI_TENANCY && !context.tenantId)) {
throw new Error("Global DB not found")
}
if (context.isSelfHostUsingCloud) {
throw new Error(
"Global DB not found - self-host users using cloud don't have a global DB"
)
}
return getDB(baseGlobalDBName(context?.tenantId))
}
@ -344,6 +376,11 @@ export function getAppDB(opts?: any): Database {
if (!appId) {
throw new Error("Unable to retrieve app DB - no app ID.")
}
if (isSelfHostUsingCloud()) {
throw new Error(
"App DB not found - self-host users using cloud don't have app DBs"
)
}
return getDB(appId, opts)
}

View File

@ -5,6 +5,7 @@ import { GoogleSpreadsheet } from "google-spreadsheet"
// keep this out of Budibase types, don't want to expose context info
export type ContextMap = {
tenantId?: string
isSelfHostUsingCloud?: boolean
appId?: string
identity?: IdentityContext
environmentVariables?: Record<string, string>

@ -1 +1 @@
Subproject commit 4417bceb24eabdd9a8c1615fb83c4e6fe8c0c914
Subproject commit 8eb981cf01151261697a8f26c08c4c28f66b8e15

View File

@ -29,7 +29,6 @@ interface TestSetup {
name: string
setup: SetupFn
mockLLMResponse: MockLLMResponseFn
selfHostOnly?: boolean
}
function budibaseAI(): SetupFn {
@ -80,7 +79,7 @@ function customAIConfig(providerConfig: Partial<ProviderConfig>): SetupFn {
}
}
const providers: TestSetup[] = [
const allProviders: TestSetup[] = [
{
name: "OpenAI API key",
setup: async () => {
@ -89,7 +88,6 @@ const providers: TestSetup[] = [
})
},
mockLLMResponse: mockChatGPTResponse,
selfHostOnly: true,
},
{
name: "OpenAI API key with custom config",
@ -126,9 +124,9 @@ describe("AI", () => {
nock.cleanAll()
})
describe.each(providers)(
describe.each(allProviders)(
"provider: $name",
({ setup, mockLLMResponse, selfHostOnly }: TestSetup) => {
({ setup, mockLLMResponse }: TestSetup) => {
let cleanup: () => Promise<void> | void
beforeAll(async () => {
cleanup = await setup(config)
@ -243,86 +241,104 @@ describe("AI", () => {
)
})
})
!selfHostOnly &&
describe("POST /api/ai/chat", () => {
let envCleanup: () => void
let featureCleanup: () => void
beforeAll(() => {
envCleanup = setEnv({ SELF_HOSTED: false })
featureCleanup = features.testutils.setFeatureFlags("*", {
AI_JS_GENERATION: true,
})
})
afterAll(() => {
featureCleanup()
envCleanup()
})
beforeEach(() => {
const license: License = {
plan: {
type: PlanType.FREE,
model: PlanModel.PER_USER,
usesInvoicing: false,
},
features: [],
quotas: {} as any,
tenantId: config.tenantId,
}
nock(env.ACCOUNT_PORTAL_URL).get("/api/license").reply(200, license)
})
it("handles correct chat response", async () => {
mockLLMResponse("Hi there!")
const { message } = await config.api.ai.chat({
messages: [{ role: "user", content: "Hello!" }],
licenseKey: "test-key",
})
expect(message).toBe("Hi there!")
})
it("handles chat response error", async () => {
mockLLMResponse(() => {
throw new Error("LLM error")
})
await config.api.ai.chat(
{
messages: [{ role: "user", content: "Hello!" }],
licenseKey: "test-key",
},
{ status: 500 }
)
})
it("handles no license", async () => {
nock.cleanAll()
nock(env.ACCOUNT_PORTAL_URL).get("/api/license").reply(404)
await config.api.ai.chat(
{
messages: [{ role: "user", content: "Hello!" }],
licenseKey: "test-key",
},
{
status: 403,
}
)
})
it("handles no license key", async () => {
await config.api.ai.chat(
{
messages: [{ role: "user", content: "Hello!" }],
// @ts-expect-error - intentionally wrong
licenseKey: undefined,
},
{
status: 403,
}
)
})
})
}
)
})
describe("BudibaseAI", () => {
const config = new TestConfiguration()
let cleanup: () => void | Promise<void>
beforeAll(async () => {
await config.init()
cleanup = await budibaseAI()(config)
})
afterAll(async () => {
if ("then" in cleanup) {
await cleanup()
} else {
cleanup()
}
config.end()
})
describe("POST /api/ai/chat", () => {
let envCleanup: () => void
let featureCleanup: () => void
beforeAll(() => {
envCleanup = setEnv({ SELF_HOSTED: false })
featureCleanup = features.testutils.setFeatureFlags("*", {
AI_JS_GENERATION: true,
})
})
afterAll(() => {
featureCleanup()
envCleanup()
})
beforeEach(() => {
nock.cleanAll()
const license: License = {
plan: {
type: PlanType.FREE,
model: PlanModel.PER_USER,
usesInvoicing: false,
},
features: [],
quotas: {} as any,
tenantId: config.tenantId,
}
nock(env.ACCOUNT_PORTAL_URL).get("/api/license").reply(200, license)
})
it("handles correct chat response", async () => {
mockChatGPTResponse("Hi there!")
const { message } = await config.api.ai.chat({
messages: [{ role: "user", content: "Hello!" }],
licenseKey: "test-key",
})
expect(message).toBe("Hi there!")
})
it("handles chat response error", async () => {
mockChatGPTResponse(() => {
throw new Error("LLM error")
})
await config.api.ai.chat(
{
messages: [{ role: "user", content: "Hello!" }],
licenseKey: "test-key",
},
{ status: 500 }
)
})
it("handles no license", async () => {
nock.cleanAll()
nock(env.ACCOUNT_PORTAL_URL).get("/api/license").reply(404)
await config.api.ai.chat(
{
messages: [{ role: "user", content: "Hello!" }],
licenseKey: "test-key",
},
{
status: 403,
}
)
})
it("handles no license key", async () => {
await config.api.ai.chat(
{
messages: [{ role: "user", content: "Hello!" }],
// @ts-expect-error - intentionally wrong
licenseKey: undefined,
},
{
status: 403,
}
)
})
})
})