From 7b867abd19bb17cacae455f4ed8411f9f1b222b7 Mon Sep 17 00:00:00 2001
From: Martin McKeaveney <martinmckeaveney@gmail.com>
Date: Mon, 4 Nov 2024 22:15:41 +0000
Subject: [PATCH] flag ai with feature flags correctly

---
 .../src/api/controllers/row/staticFormula.ts   | 17 +++++++++++++----
 .../server/src/automations/steps/openai.ts     | 18 +++++++++++-------
 2 files changed, 24 insertions(+), 11 deletions(-)

diff --git a/packages/server/src/api/controllers/row/staticFormula.ts b/packages/server/src/api/controllers/row/staticFormula.ts
index 28853a8faf..31a6678012 100644
--- a/packages/server/src/api/controllers/row/staticFormula.ts
+++ b/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],
diff --git a/packages/server/src/automations/steps/openai.ts b/packages/server/src/automations/steps/openai.ts
index b1dfa3df5b..b3db0da9c5 100644
--- a/packages/server/src/automations/steps/openai.ts
+++ b/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)