From e194a2feb8f64e163862ff6ff64ed65f535a6bd1 Mon Sep 17 00:00:00 2001
From: Dhruwang Jariwala <67850763+Dhruwang@users.noreply.github.com>
Date: Tue, 31 Dec 2024 14:31:08 +0530
Subject: [PATCH] chore: make ai feature available for self hosting (#4543)

---
 .../web/modules/ee/license-check/lib/utils.ts | 23 ++++++++++++-------
 .../license-check/types/enterprise-license.ts |  1 +
 2 files changed, 16 insertions(+), 8 deletions(-)

diff --git a/apps/web/modules/ee/license-check/lib/utils.ts b/apps/web/modules/ee/license-check/lib/utils.ts
index 17ea780d0e..6d1d3579d6 100644
--- a/apps/web/modules/ee/license-check/lib/utils.ts
+++ b/apps/web/modules/ee/license-check/lib/utils.ts
@@ -89,6 +89,7 @@ const fetchLicenseForE2ETesting = async (): Promise<{
           projects: 3,
           whitelabel: true,
           removeBranding: true,
+          ai: true,
         },
         lastChecked: currentTime,
       };
@@ -154,6 +155,7 @@ export const getEnterpriseLicense = async (): Promise<{
           whitelabel: false,
           removeBranding: false,
           contacts: false,
+          ai: false,
         },
         lastChecked: new Date(),
       };
@@ -350,18 +352,23 @@ export const getIsSSOEnabled = async (): Promise<boolean> => {
 };
 
 export const getIsOrganizationAIReady = async (billingPlan: TOrganizationBillingPlan) => {
-  // TODO: We'll remove the IS_FORMBRICKS_CLOUD check once we have the AI feature available for self-hosted customers
+  if (!IS_AI_CONFIGURED) return false;
+  if (E2E_TESTING) {
+    const previousResult = await fetchLicenseForE2ETesting();
+    return previousResult && previousResult.features ? previousResult.features.ai : false;
+  }
+  const license = await getEnterpriseLicense();
+
   if (IS_FORMBRICKS_CLOUD) {
-    return (
-      IS_AI_CONFIGURED &&
-      (await getEnterpriseLicense()).active &&
-      (billingPlan === PROJECT_FEATURE_KEYS.STARTUP ||
-        billingPlan === PROJECT_FEATURE_KEYS.SCALE ||
-        billingPlan === PROJECT_FEATURE_KEYS.ENTERPRISE)
+    return Boolean(
+      license.features?.ai &&
+        (billingPlan === PROJECT_FEATURE_KEYS.STARTUP ||
+          billingPlan === PROJECT_FEATURE_KEYS.SCALE ||
+          billingPlan === PROJECT_FEATURE_KEYS.ENTERPRISE)
     );
   }
 
-  return false;
+  return Boolean(license.features?.ai);
 };
 
 export const getIsAIEnabled = async (organization: TOrganization) => {
diff --git a/apps/web/modules/ee/license-check/types/enterprise-license.ts b/apps/web/modules/ee/license-check/types/enterprise-license.ts
index a384144e82..987da12f39 100644
--- a/apps/web/modules/ee/license-check/types/enterprise-license.ts
+++ b/apps/web/modules/ee/license-check/types/enterprise-license.ts
@@ -12,6 +12,7 @@ const ZEnterpriseLicenseFeatures = z.object({
   removeBranding: z.boolean(),
   twoFactorAuth: z.boolean(),
   sso: z.boolean(),
+  ai: z.boolean(),
 });
 
 export type TEnterpriseLicenseFeatures = z.infer<typeof ZEnterpriseLicenseFeatures>;