Merge pull request #14959 from Budibase/flag-ai-correctly-for-premium

Flag ai correctly for premium

packages/builder/src/components/automation/SetupPanel/CronBuilder.svelte

@@ -9,7 +9,7 @@
   } from "@budibase/bbui"
   import { onMount, createEventDispatcher } from "svelte"
   import { flags } from "stores/builder"
-  import { featureFlags } from "stores/portal"
+  import { featureFlags, licensing } from "stores/portal"
   import { API } from "api"
   import MagicWand from "../../../../assets/MagicWand.svelte"
 
@@ -26,7 +26,9 @@
   let aiCronPrompt = ""
   let loadingAICronExpression = false
 
-  $: aiEnabled = $featureFlags.AI_CUSTOM_CONFIGS || $featureFlags.BUDIBASE_AI
+  $: aiEnabled =
+    ($featureFlags.AI_CUSTOM_CONFIGS && $licensing.customAIConfigsEnabled) ||
+    ($featureFlags.BUDIBASE_AI && $licensing.budibaseAIEnabled)
   $: {
     if (cronExpression) {
       try {

packages/builder/src/pages/builder/portal/settings/ai/AISettings.spec.js

@@ -1,7 +1,7 @@
 import { it, expect, describe, vi } from "vitest"
 import AISettings from "./index.svelte"
 import { render, fireEvent } from "@testing-library/svelte"
-import { admin, licensing } from "stores/portal"
+import { admin, licensing, featureFlags } from "stores/portal"
 import { notifications } from "@budibase/bbui"
 
 vi.spyOn(notifications, "error").mockImplementation(vi.fn)
@@ -12,12 +12,17 @@ const Hosting = {
   Self: "self",
 }
 
-function setupEnv(hosting, features = {}) {
+function setupEnv(hosting, features = {}, flags = {}) {
   const defaultFeatures = {
     budibaseAIEnabled: false,
     customAIConfigsEnabled: false,
     ...features,
   }
+  const defaultFlags = {
+    BUDIBASE_AI: false,
+    AI_CUSTOM_CONFIGS: false,
+    ...flags,
+  }
   admin.subscribe = vi.fn().mockImplementation(callback => {
     callback({ cloud: hosting === Hosting.Cloud })
     return () => {}
@@ -26,6 +31,10 @@ function setupEnv(hosting, features = {}) {
     callback(defaultFeatures)
     return () => {}
   })
+  featureFlags.subscribe = vi.fn().mockImplementation(callback => {
+    callback(defaultFlags)
+    return () => {}
+  })
 }
 
 describe("AISettings", () => {
@@ -72,7 +81,11 @@ describe("AISettings", () => {
       let addConfigurationButton
       let configModal
 
-      setupEnv(Hosting.Cloud, { customAIConfigsEnabled: true })
+      setupEnv(
+        Hosting.Cloud,
+        { customAIConfigsEnabled: true },
+        { AI_CUSTOM_CONFIGS: true }
+      )
       instance = render(AISettings)
       addConfigurationButton = instance.queryByText("Add configuration")
       expect(addConfigurationButton).toBeInTheDocument()
@@ -85,7 +98,11 @@ describe("AISettings", () => {
       let addConfigurationButton
       let configModal
 
-      setupEnv(Hosting.Self, { customAIConfigsEnabled: true })
+      setupEnv(
+        Hosting.Self,
+        { customAIConfigsEnabled: true },
+        { AI_CUSTOM_CONFIGS: true }
+      )
       instance = render(AISettings)
       addConfigurationButton = instance.queryByText("Add configuration")
       expect(addConfigurationButton).toBeInTheDocument()

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

@@ -12,7 +12,7 @@
     Tags,
     Tag,
   } from "@budibase/bbui"
-  import { admin, licensing } from "stores/portal"
+  import { admin, licensing, featureFlags } from "stores/portal"
   import { API } from "api"
   import AIConfigModal from "./ConfigModal.svelte"
   import AIConfigTile from "./AIConfigTile.svelte"
@@ -27,7 +27,8 @@
   let editingUuid
 
   $: isCloud = $admin.cloud
-  $: customAIConfigsEnabled = $licensing.customAIConfigsEnabled
+  $: customAIConfigsEnabled =
+    $featureFlags.AI_CUSTOM_CONFIGS && $licensing.customAIConfigsEnabled
 
   async function fetchAIConfig() {
     try {

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

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

packages/server/src/api/controllers/row/staticFormula.ts

@@ -4,8 +4,15 @@ import {
   processAIColumns,
   processFormulas,
 } from "../../../utilities/rowProcessor"
-import { context } from "@budibase/backend-core"
-import { Table, Row, FormulaType, FieldType, ViewV2 } from "@budibase/types"
+import { context, features } from "@budibase/backend-core"
+import {
+  Table,
+  Row,
+  FeatureFlag,
+  FormulaType,
+  FieldType,
+  ViewV2,
+} from "@budibase/types"
 import * as linkRows from "../../../db/linkedRows"
 import isEqual from "lodash/isEqual"
 import { cloneDeep } from "lodash/fp"
@@ -145,8 +152,10 @@ export async function finaliseRow(
     contextRows: [enrichedRow],
   })
   const aiEnabled =
-    (await pro.features.isBudibaseAIEnabled()) ||
-    (await pro.features.isAICustomConfigsEnabled())
+    ((await features.flags.isEnabled(FeatureFlag.BUDIBASE_AI)) &&
+      (await pro.features.isBudibaseAIEnabled())) ||
+    ((await features.flags.isEnabled(FeatureFlag.AI_CUSTOM_CONFIGS)) &&
+      (await pro.features.isAICustomConfigsEnabled()))
   if (aiEnabled) {
     row = await processAIColumns(table, row, {
       contextRows: [enrichedRow],

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

@@ -7,17 +7,17 @@ import {
   AutomationIOType,
   OpenAIStepInputs,
   OpenAIStepOutputs,
+  FeatureFlag,
 } from "@budibase/types"
-import { env } from "@budibase/backend-core"
+import { env, features } from "@budibase/backend-core"
 import * as automationUtils from "../automationUtils"
 import * as pro from "@budibase/pro"
 
 enum Model {
-  GPT_35_TURBO = "gpt-3.5-turbo",
-  // will only work with api keys that have access to the GPT4 API
-  GPT_4 = "gpt-4",
-  GPT_4O = "gpt-4o",
   GPT_4O_MINI = "gpt-4o-mini",
+  GPT_4O = "gpt-4o",
+  GPT_4 = "gpt-4",
+  GPT_35_TURBO = "gpt-3.5-turbo",
 }
 
 export const definition: AutomationStepDefinition = {
@@ -99,8 +99,12 @@ export async function run({
 
   try {
     let response
-    const customConfigsEnabled = await pro.features.isAICustomConfigsEnabled()
-    const budibaseAIEnabled = await pro.features.isBudibaseAIEnabled()
+    const customConfigsEnabled =
+      (await features.flags.isEnabled(FeatureFlag.AI_CUSTOM_CONFIGS)) &&
+      (await pro.features.isAICustomConfigsEnabled())
+    const budibaseAIEnabled =
+      (await features.flags.isEnabled(FeatureFlag.BUDIBASE_AI)) &&
+      (await pro.features.isBudibaseAIEnabled())
 
     if (budibaseAIEnabled || customConfigsEnabled) {
       const llm = await pro.ai.LargeLanguageModel.forCurrentTenant(inputs.model)