Add default AI subscription policies

SavinduDimal · Oct 9, 2024 · 52e0d77 · 52e0d77
1 parent 3dbb09f
commit 52e0d77
Show file tree

Hide file tree

Showing 2 changed files with 61 additions and 0 deletions.
diff --git a/...t/org.wso2.carbon.apimgt.impl/src/main/java/org/wso2/carbon/apimgt/impl/APIConstants.java b/...t/org.wso2.carbon.apimgt.impl/src/main/java/org/wso2/carbon/apimgt/impl/APIConstants.java
@@ -1963,6 +1963,10 @@ public enum RegistryResourceTypesForUI {
     public static final String DEFAULT_SUB_POLICY_ASYNC_WH_BRONZE = "AsyncWHBronze";
     public static final String DEFAULT_SUB_POLICY_ASYNC_WH_UNLIMITED = "AsyncWHUnlimited";
 
+    public static final String DEFAULT_SUB_POLICY_AI_GOLD = "AIGold";
+    public static final String DEFAULT_SUB_POLICY_AI_SILVER = "AISilver";
+    public static final String DEFAULT_SUB_POLICY_AI_BRONZE = "AIBronze";
+
     public static final String DEFAULT_SUB_POLICY_GOLD_DESC = "Allows 5000 requests per minute";
     public static final String DEFAULT_SUB_POLICY_SILVER_DESC = "Allows 2000 requests per minute";
     public static final String DEFAULT_SUB_POLICY_BRONZE_DESC = "Allows 1000 requests per minute";
@@ -1987,6 +1991,10 @@ public enum RegistryResourceTypesForUI {
     public static final String DEFAULT_SUB_POLICY_ASYNC_WH_UNLIMITED_DESC = "Allows unlimited events and " +
             "unlimited active subscriptions";
 
+    public static final String DEFAULT_SUB_POLICY_AI_GOLD_DESC = "Allows 50000 total tokens and 500 requests per minute";
+    public static final String DEFAULT_SUB_POLICY_AI_SILVER_DESC = "Allows 10000 total tokens and 100 requests per minute";
+    public static final String DEFAULT_SUB_POLICY_AI_BRONZE_DESC = "Allows 1000 total tokens and 10 requests per minute";
+
     public static final String DEFAULT_API_POLICY_FIFTY_THOUSAND_REQ_PER_MIN = "50KPerMin";
     public static final String DEFAULT_API_POLICY_TWENTY_THOUSAND_REQ_PER_MIN = "20KPerMin";
     public static final String DEFAULT_API_POLICY_TEN_THOUSAND_REQ_PER_MIN = "10KPerMin";

diff --git a/.../org.wso2.carbon.apimgt.impl/src/main/java/org/wso2/carbon/apimgt/impl/utils/APIUtil.java b/.../org.wso2.carbon.apimgt.impl/src/main/java/org/wso2/carbon/apimgt/impl/utils/APIUtil.java
@@ -6316,6 +6316,59 @@ public static void addDefaultTenantAdvancedThrottlePolicies(String tenantDomain,
             SubscriptionPolicy retrievedPolicy = apiMgtDAO.getSubscriptionPolicy(policyName, tenantId);
             deployRetrievedSubscriptionPolicy(tenantId, retrievedPolicy);
         }
+
+        //Adding AI API Quota based subscription level policies
+        long[] totalTokenCountValues = new long[]{50000, 10000, 1000};
+        long[] aiPolicyRequestCount =  new long[]{500, 100, 10};
+        String[] aiApiQuotaSubPolicyNames = new String[]{APIConstants.DEFAULT_SUB_POLICY_AI_GOLD,
+                APIConstants.DEFAULT_SUB_POLICY_AI_SILVER, APIConstants.DEFAULT_SUB_POLICY_AI_BRONZE};
+        String[] aiApiQuotaSubPolicyDescriptions = new String[]{
+                APIConstants.DEFAULT_SUB_POLICY_AI_GOLD_DESC, APIConstants.DEFAULT_SUB_POLICY_AI_SILVER_DESC,
+                APIConstants.DEFAULT_SUB_POLICY_AI_BRONZE_DESC};
+
+        for (int i = 0; i < aiApiQuotaSubPolicyNames.length; i++) {
+            policyName = aiApiQuotaSubPolicyNames[i];
+            boolean needDeployment = false;
+            SubscriptionPolicy subscriptionPolicy = new SubscriptionPolicy(policyName);
+            subscriptionPolicy.setDisplayName(policyName);
+            subscriptionPolicy.setDescription(aiApiQuotaSubPolicyDescriptions[i]);
+            subscriptionPolicy.setTenantId(tenantId);
+            subscriptionPolicy.setDeployed(true);
+            QuotaPolicy defaultQuotaPolicy = new QuotaPolicy();
+            AIAPIQuotaLimit aiApiQuotaLimit = new AIAPIQuotaLimit();
+            aiApiQuotaLimit.setRequestCount(aiPolicyRequestCount[i]);
+            aiApiQuotaLimit.setTotalTokenCount(totalTokenCountValues[i]);
+            aiApiQuotaLimit.setUnitTime(1);
+            aiApiQuotaLimit.setTimeUnit(APIConstants.TIME_UNIT_MINUTE);
+            defaultQuotaPolicy.setType(PolicyConstants.AI_API_QUOTA_TYPE);
+            defaultQuotaPolicy.setLimit(aiApiQuotaLimit);
+            subscriptionPolicy.setDefaultQuotaPolicy(defaultQuotaPolicy);
+            subscriptionPolicy.setStopOnQuotaReach(true);
+            subscriptionPolicy.setBillingPlan(APIConstants.BILLING_PLAN_FREE);
+
+            if (!apiMgtDAO.isPolicyExist(PolicyConstants.POLICY_LEVEL_SUB, tenantId, policyName)) {
+                apiMgtDAO.addSubscriptionPolicy(subscriptionPolicy);
+                needDeployment = true;
+            }
+
+            if (!needDeployment) {
+                continue;
+            }
+            String superTenantDomain = null;
+            try {
+                superTenantDomain = ServiceReferenceHolder.getInstance().getRealmService().getTenantManager().
+                        getSuperTenantDomain();
+            } catch (UserStoreException e) {
+                handleInternalException("Error in getting the super tenant domain", e);
+            }
+
+            boolean isSuperTenant = tenantDomain.equals(superTenantDomain);
+            if (isSuperTenant) {
+                continue;
+            }
+            SubscriptionPolicy retrievedPolicy = apiMgtDAO.getSubscriptionPolicy(policyName, tenantId);
+            deployRetrievedSubscriptionPolicy(tenantId, retrievedPolicy);
+        }
 
         long tenThousandPerMinTier = defualtLimits.containsKey(APIConstants.DEFAULT_API_POLICY_TEN_THOUSAND_REQ_PER_MIN) ?
                 defualtLimits.get(APIConstants.DEFAULT_API_POLICY_TEN_THOUSAND_REQ_PER_MIN) : 10000;