flag ai with feature flags correctly

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)