diff --git a/components/apimgt/org.wso2.carbon.apimgt.impl/src/main/java/org/wso2/carbon/apimgt/impl/APIConstants.java b/components/apimgt/org.wso2.carbon.apimgt.impl/src/main/java/org/wso2/carbon/apimgt/impl/APIConstants.java index 618f3277f8d5..77623e4768e4 100755 --- a/components/apimgt/org.wso2.carbon.apimgt.impl/src/main/java/org/wso2/carbon/apimgt/impl/APIConstants.java +++ b/components/apimgt/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/components/apimgt/org.wso2.carbon.apimgt.impl/src/main/java/org/wso2/carbon/apimgt/impl/utils/APIUtil.java b/components/apimgt/org.wso2.carbon.apimgt.impl/src/main/java/org/wso2/carbon/apimgt/impl/utils/APIUtil.java index 6fedcef51c7c..06fdd2d8f95e 100644 --- a/components/apimgt/org.wso2.carbon.apimgt.impl/src/main/java/org/wso2/carbon/apimgt/impl/utils/APIUtil.java +++ b/components/apimgt/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;