Merge pull request #14965 from Budibase/ai-config-fixes-2

further AI config updates for case where there's no configuration

packages/builder/src/pages/builder/portal/settings/index.svelte

@@ -1,8 +1,8 @@
 <script>
   import { redirect } from "@roxi/routify"
-  import { licensing, featureFlags } from "stores/portal"
+  import { featureFlags } from "stores/portal"
 
-  if ($featureFlags.AI_CUSTOM_CONFIGS && $licensing.customAIConfigsEnabled) {
+  if ($featureFlags.AI_CUSTOM_CONFIGS) {
     $redirect("./ai")
   } else {
     $redirect("./auth")

packages/pro

@@ -1 +1 @@
-Subproject commit 2ab8536b6005576684810d774f1ac22239218546
+Subproject commit 04bee88597edb1edb88ed299d0597b587f0362ec

packages/server/src/api/routes/tests/row.spec.ts

@@ -56,6 +56,7 @@ jest.mock("@budibase/pro", () => ({
   ai: {
     LargeLanguageModel: {
       forCurrentTenant: async () => ({
+        initialised: true,
         run: jest.fn(() => `Mock LLM Response`),
         buildPromptFromAIOperation: jest.fn(),
       }),

packages/server/src/api/routes/tests/search.spec.ts

@@ -54,6 +54,7 @@ jest.mock("@budibase/pro", () => ({
   ai: {
     LargeLanguageModel: {
       forCurrentTenant: async () => ({
+        initialised: true,
         run: jest.fn(() => `Mock LLM Response`),
         buildPromptFromAIOperation: jest.fn(),
       }),

packages/server/src/automations/steps/openai.ts

@@ -106,21 +106,15 @@ export async function run({
       (await features.flags.isEnabled(FeatureFlag.BUDIBASE_AI)) &&
       (await pro.features.isBudibaseAIEnabled())
 
+    let llm
     if (budibaseAIEnabled || customConfigsEnabled) {
-      const llm = await pro.ai.LargeLanguageModel.forCurrentTenant(inputs.model)
+      llm = await pro.ai.LargeLanguageModel.forCurrentTenant(inputs.model)
-      response = await llm.run(inputs.prompt)
-    } else {
-      // fallback to the default that uses the environment variable for backwards compat
-      if (!env.OPENAI_API_KEY) {
-        return {
-          success: false,
-          response:
-            "OpenAI API Key not configured - please add the OPENAI_API_KEY environment variable.",
-        }
-      }
-      response = await legacyOpenAIPrompt(inputs)
     }
 
+    response = llm?.initialised
+      ? await llm.run(inputs.prompt)
+      : await legacyOpenAIPrompt(inputs)
+
     return {
       response,
       success: true,

packages/server/src/automations/tests/openai.spec.ts

@@ -1,9 +1,6 @@
 import { getConfig, runStep, afterAll as _afterAll } from "./utilities"
 import { OpenAI } from "openai"
-import {
+import { setEnv as setCoreEnv } from "@budibase/backend-core"
-  withEnv as withCoreEnv,
-  setEnv as setCoreEnv,
-} from "@budibase/backend-core"
 import * as pro from "@budibase/pro"
 
 jest.mock("openai", () => ({
@@ -28,6 +25,7 @@ jest.mock("@budibase/pro", () => ({
   ai: {
     LargeLanguageModel: {
       forCurrentTenant: jest.fn().mockImplementation(() => ({
+        initialised: true,
         init: jest.fn(),
         run: jest.fn(),
       })),
@@ -63,16 +61,6 @@ describe("test the openai action", () => {
 
   afterAll(_afterAll)
 
-  it("should present the correct error message when the OPENAI_API_KEY variable isn't set", async () => {
-    await withCoreEnv({ OPENAI_API_KEY: "" }, async () => {
-      let res = await runStep("OPENAI", { prompt: OPENAI_PROMPT })
-      expect(res.response).toEqual(
-        "OpenAI API Key not configured - please add the OPENAI_API_KEY environment variable."
-      )
-      expect(res.success).toBeFalsy()
-    })
-  })
-
   it("should be able to receive a response from ChatGPT given a prompt", async () => {
     const res = await runStep("OPENAI", { prompt: OPENAI_PROMPT })
     expect(res.response).toEqual("This is a test")

packages/server/src/utilities/rowProcessor/tests/utils.spec.ts

@@ -18,6 +18,7 @@ jest.mock("@budibase/pro", () => ({
   ai: {
     LargeLanguageModel: {
       forCurrentTenant: async () => ({
+        initialised: true,
         run: jest.fn(() => "response from LLM"),
         buildPromptFromAIOperation: buildPromptMock,
       }),

packages/server/src/utilities/rowProcessor/utils.ts

@@ -108,7 +108,7 @@ export async function processAIColumns<T extends Row | Row[]>(
     span?.addTags({ table_id: table._id, numRows })
     const rows = Array.isArray(inputRows) ? inputRows : [inputRows]
     const llm = await pro.ai.LargeLanguageModel.forCurrentTenant("gpt-4o-mini")
-    if (rows) {
+    if (rows && llm.initialised) {
       // Ensure we have snippet context
       await context.ensureSnippetContext()
 

packages/worker/src/api/controllers/global/configs.ts

@@ -12,27 +12,29 @@ import {
 } from "@budibase/backend-core"
 import { checkAnyUserExists } from "../../../utilities/users"
 import {
+  AIConfig,
+  AIInnerConfig,
   Config,
   ConfigType,
   Ctx,
   GetPublicOIDCConfigResponse,
   GetPublicSettingsResponse,
   GoogleInnerConfig,
+  isAIConfig,
   isGoogleConfig,
   isOIDCConfig,
   isSettingsConfig,
   isSMTPConfig,
   OIDCConfigs,
+  OIDCLogosConfig,
+  PASSWORD_REPLACEMENT,
+  QuotaUsageType,
   SettingsBrandingConfig,
   SettingsInnerConfig,
   SSOConfig,
   SSOConfigType,
+  StaticQuotaName,
   UserCtx,
-  OIDCLogosConfig,
-  AIConfig,
-  PASSWORD_REPLACEMENT,
-  isAIConfig,
-  AIInnerConfig,
 } from "@budibase/types"
 import * as pro from "@budibase/pro"
 
@@ -83,6 +85,12 @@ const getEventFns = async (config: Config, existing?: Config) => {
       fns.push(events.email.SMTPUpdated)
     } else if (isAIConfig(config)) {
       fns.push(() => events.ai.AIConfigUpdated)
+      if (
+        Object.keys(existing.config).length > Object.keys(config.config).length
+      ) {
+        fns.push(() => pro.quotas.removeCustomAIConfig())
+      }
+      fns.push(() => pro.quotas.addCustomAIConfig())
     } else if (isGoogleConfig(config)) {
       fns.push(() => events.auth.SSOUpdated(ConfigType.GOOGLE))
       if (!existing.config.activated && config.config.activated) {
@@ -248,7 +256,6 @@ export async function save(ctx: UserCtx<Config>) {
         if (existingConfig) {
           await verifyAIConfig(config, existingConfig)
         }
-        await pro.quotas.addCustomAIConfig()
         break
     }
   } catch (err: any) {
@@ -518,7 +525,11 @@ export async function destroy(ctx: UserCtx) {
     await db.remove(id, rev)
     await cache.destroy(cache.CacheKey.CHECKLIST)
     if (id === configs.generateConfigID(ConfigType.AI)) {
-      await pro.quotas.removeCustomAIConfig()
+      await pro.quotas.set(
+        StaticQuotaName.AI_CUSTOM_CONFIGS,
+        QuotaUsageType.STATIC,
+        0
+      )
     }
     ctx.body = { message: "Config deleted successfully" }
   } catch (err: any) {