Merge pull request #14846 from Budibase/ai-fixes-3.0

BB AI QA testing fixes

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

@@ -31,10 +31,7 @@
 
   async function fetchAIConfig() {
     try {
-      const aiDoc = await API.getConfig(ConfigTypes.AI)
-      if (aiDoc._id) {
-        fullAIConfig = aiDoc
-      }
+      fullAIConfig = await API.getConfig(ConfigTypes.AI)
     } catch (error) {
       notifications.error("Error fetching AI config")
     }
@@ -66,6 +63,7 @@
       }
       // Add new or update existing custom AI Config
       fullAIConfig.config[id] = editingAIConfig
+      fullAIConfig.type = ConfigTypes.AI
     }
 
     try {

packages/pro

@@ -1 +1 @@
-Subproject commit 297fdc937e9c650b4964fc1a942b60022b195865
+Subproject commit 60411963053b98e0415a1e9285b721fc26c628c8

packages/server/src/automations/actions.ts

@@ -18,6 +18,7 @@ import * as loop from "./steps/loop"
 import * as collect from "./steps/collect"
 import * as branch from "./steps/branch"
 import * as triggerAutomationRun from "./steps/triggerAutomationRun"
+import * as openai from "./steps/openai"
 import env from "../environment"
 import {
   PluginType,
@@ -50,6 +51,7 @@ const ACTION_IMPLS: ActionImplType = {
   QUERY_ROWS: queryRow.run,
   COLLECT: collect.run,
   TRIGGER_AUTOMATION_RUN: triggerAutomationRun.run,
+  OPENAI: openai.run,
   // these used to be lowercase step IDs, maintain for backwards compat
   discord: discord.run,
   slack: slack.run,
@@ -89,21 +91,25 @@ export const BUILTIN_ACTION_DEFINITIONS: Record<
 // ran at all
 if (env.SELF_HOSTED) {
   const bash = require("./steps/bash")
-  const openai = require("./steps/openai")
 
   // @ts-ignore
   ACTION_IMPLS["EXECUTE_BASH"] = bash.run
   // @ts-ignore
   BUILTIN_ACTION_DEFINITIONS["EXECUTE_BASH"] = bash.definition
-  // @ts-ignore
-  ACTION_IMPLS.OPENAI = openai.run
-  BUILTIN_ACTION_DEFINITIONS.OPENAI = openai.definition
 }
 
 export async function getActionDefinitions() {
   if (await features.flags.isEnabled(FeatureFlag.AUTOMATION_BRANCHING)) {
     BUILTIN_ACTION_DEFINITIONS["BRANCH"] = branch.definition
   }
+  if (
+    env.SELF_HOSTED ||
+    (await features.flags.isEnabled(FeatureFlag.BUDIBASE_AI)) ||
+    (await features.flags.isEnabled(FeatureFlag.AI_CUSTOM_CONFIGS))
+  ) {
+    BUILTIN_ACTION_DEFINITIONS["OPENAI"] = openai.definition
+  }
+
   const actionDefinitions = BUILTIN_ACTION_DEFINITIONS
   if (env.SELF_HOSTED) {
     const plugins = await sdk.plugins.fetch(PluginType.AUTOMATION)

packages/types/src/documents/app/automation/schema.ts

@@ -126,16 +126,16 @@ export type ActionImplementations<T extends Hosting> = {
     n8nStepInputs,
     ExternalAppStepOutputs
   >
+  [AutomationActionStepId.OPENAI]: ActionImplementation<
+    OpenAIStepInputs,
+    OpenAIStepOutputs
+  >
 } & (T extends "self"
   ? {
       [AutomationActionStepId.EXECUTE_BASH]: ActionImplementation<
         BashStepInputs,
         BashStepOutputs
       >
-      [AutomationActionStepId.OPENAI]: ActionImplementation<
-        OpenAIStepInputs,
-        OpenAIStepOutputs
-      >
     }
   : {})
 

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

@@ -44,9 +44,7 @@ const getEventFns = async (config: Config, existing?: Config) => {
       fns.push(events.email.SMTPCreated)
     } else if (isAIConfig(config)) {
       fns.push(() => events.ai.AIConfigCreated)
-      fns.push(() =>
-        pro.quotas.updateCustomAIConfigCount(Object.keys(config.config).length)
-      )
+      fns.push(() => pro.quotas.addCustomAIConfig())
     } else if (isGoogleConfig(config)) {
       fns.push(() => events.auth.SSOCreated(ConfigType.GOOGLE))
       if (config.config.activated) {
@@ -85,9 +83,6 @@ const getEventFns = async (config: Config, existing?: Config) => {
       fns.push(events.email.SMTPUpdated)
     } else if (isAIConfig(config)) {
       fns.push(() => events.ai.AIConfigUpdated)
-      fns.push(() =>
-        pro.quotas.updateCustomAIConfigCount(Object.keys(config.config).length)
-      )
     } else if (isGoogleConfig(config)) {
       fns.push(() => events.auth.SSOUpdated(ConfigType.GOOGLE))
       if (!existing.config.activated && config.config.activated) {
@@ -253,7 +248,7 @@ export async function save(ctx: UserCtx<Config>) {
         if (existingConfig) {
           await verifyAIConfig(config, existingConfig)
         }
-        await pro.quotas.updateCustomAIConfigCount(Object.keys(config).length)
+        await pro.quotas.addCustomAIConfig()
         break
     }
   } catch (err: any) {
@@ -342,29 +337,43 @@ export async function find(ctx: UserCtx) {
     let scopedConfig = await configs.getConfig(type)
 
     if (scopedConfig) {
-      if (type === ConfigType.OIDC_LOGOS) {
-        enrichOIDCLogos(scopedConfig)
-      }
-
-      if (type === ConfigType.AI) {
-        await pro.sdk.ai.enrichAIConfig(scopedConfig)
-        // Strip out the API Keys from the response so they don't show in the UI
-        for (const key in scopedConfig.config) {
-          if (scopedConfig.config[key].apiKey) {
-            scopedConfig.config[key].apiKey = PASSWORD_REPLACEMENT
-          }
-        }
-      }
-      ctx.body = scopedConfig
+      await handleConfigType(type, scopedConfig)
+    } else if (type === ConfigType.AI) {
+      scopedConfig = { config: {} } as AIConfig
+      await handleAIConfig(scopedConfig)
     } else {
-      // don't throw an error, there simply is nothing to return
+      // If no config found and not AI type, just return an empty body
       ctx.body = {}
+      return
     }
+
+    ctx.body = scopedConfig
   } catch (err: any) {
     ctx.throw(err?.status || 400, err)
   }
 }
 
+async function handleConfigType(type: ConfigType, config: Config) {
+  if (type === ConfigType.OIDC_LOGOS) {
+    enrichOIDCLogos(config)
+  } else if (type === ConfigType.AI) {
+    await handleAIConfig(config)
+  }
+}
+
+async function handleAIConfig(config: AIConfig) {
+  await pro.sdk.ai.enrichAIConfig(config)
+  stripApiKeys(config)
+}
+
+function stripApiKeys(config: AIConfig) {
+  for (const key in config?.config) {
+    if (config.config[key].apiKey) {
+      config.config[key].apiKey = PASSWORD_REPLACEMENT
+    }
+  }
+}
+
 export async function publicOidc(ctx: Ctx<void, GetPublicOIDCConfigResponse>) {
   try {
     // Find the config with the most granular scope based on context
@@ -508,6 +517,9 @@ export async function destroy(ctx: UserCtx) {
   try {
     await db.remove(id, rev)
     await cache.destroy(cache.CacheKey.CHECKLIST)
+    if (id === configs.generateConfigID(ConfigType.AI)) {
+      await pro.quotas.removeCustomAIConfig()
+    }
     ctx.body = { message: "Config deleted successfully" }
   } catch (err: any) {
     ctx.throw(err.status, err)

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

@@ -13,10 +13,6 @@ describe("Global configs controller", () => {
     await config.afterAll()
   })
 
-  afterEach(() => {
-    jest.resetAllMocks()
-  })
-
   it("Should strip secrets when pulling AI config", async () => {
     const data = structures.configs.ai()
     await config.api.configs.saveConfig(data)