Merge pull request #15857 from Budibase/budibase-ai-3
Create new global DB for self-host quota.
This commit is contained in:
commit
32fc150fac
|
@ -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)
|
||||
}
|
||||
|
|
|
@ -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)
|
||||
|
|
|
@ -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)
|
||||
}
|
||||
|
||||
|
|
|
@ -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
|
|
@ -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,
|
||||
}
|
||||
)
|
||||
})
|
||||
})
|
||||
})
|
||||
|
|
Loading…
Reference in New Issue