Merge pull request #12605 from AnuGayan/master-AI-API-temp

Fix unit test failures due to backend throttling
wso2 · Sep 26, 2024 · 687fb64 · 687fb64
2 parents 1a64965 + 2c559f7
commit 687fb64
Show file tree

Hide file tree

Showing 4 changed files with 11 additions and 43 deletions.
diff --git a/...way/src/main/java/org/wso2/carbon/apimgt/gateway/handlers/throttling/ThrottleHandler.java b/...way/src/main/java/org/wso2/carbon/apimgt/gateway/handlers/throttling/ThrottleHandler.java
@@ -184,15 +184,13 @@ private boolean doRoleBasedAccessThrottlingWithCEP(MessageContext synCtx, Config
         //Throttled decisions
         boolean isThrottled = false;
         boolean isResourceLevelThrottled = false;
-        boolean isOperationLevelThrottled = false;
         boolean isApplicationLevelThrottled;
         boolean isSubscriptionLevelThrottled;
         boolean isSubscriptionLevelSpikeThrottled = false;
         boolean isApiLevelThrottled = false;
         boolean isBlockedRequest = false;
         boolean apiLevelThrottledTriggered = false;
         boolean policyLevelUserTriggered = false;
-        String ipLevelBlockingKey;
         String appLevelBlockingKey = "";
         String subscriptionLevelBlockingKey = "";
         boolean stopOnQuotaReach = true;

diff --git a/...src/main/java/org/wso2/carbon/apimgt/rest/api/admin/v1/impl/ThrottlingApiServiceImpl.java b/...src/main/java/org/wso2/carbon/apimgt/rest/api/admin/v1/impl/ThrottlingApiServiceImpl.java
@@ -811,10 +811,6 @@ public Response throttlingPoliciesSubscriptionPolicyIdPut(String policyId, Strin
             //overridden properties
             body.setPolicyId(policyId);
             body.setPolicyName(existingPolicy.getPolicyName());
-            if (PolicyConstants.AI_API_QUOTA_TYPE_ENUM_VALUE.equals(body.getDefaultLimit().getType().toString())) {
-                body.setDefaultLimit(RestApiAdminUtils.overrideTokenBasedQuotaLimits(body.getDefaultLimit(),
-                        existingPolicy.getDefaultQuotaPolicy()));
-            }
 
             // validate if permission info exists and halt the execution in case of an error
             validatePolicyPermissions(body);

diff --git a/...in.v1/src/main/java/org/wso2/carbon/apimgt/rest/api/admin/v1/utils/RestApiAdminUtils.java b/...in.v1/src/main/java/org/wso2/carbon/apimgt/rest/api/admin/v1/utils/RestApiAdminUtils.java
@@ -383,30 +383,4 @@ public static void restoreTenantTheme(int tenantId, File tenantThemeDirectory, F
         FileUtils.deleteDirectory(backupDirectory);
         apiAdmin.updateTenantTheme(tenantId, existingTenantTheme);
     }
-
-    /**
-     * Override token based quota limits if they are null in the TokenCountLimitDTO.
-     *
-     * @param throttleLimitDTO TokenCountLimitDTO object
-     * @param quotaPolicy    Token count limits of the existing subscription policy
-     * @return Overridden TokenCountLimitDTO object
-     */
-    public static ThrottleLimitDTO overrideTokenBasedQuotaLimits(ThrottleLimitDTO throttleLimitDTO,
-            QuotaPolicy quotaPolicy) {
-
-        AIAPIQuotaLimitDTO aiApiQuotaLimitDTO = throttleLimitDTO.getAiApiQuota();
-        AIAPIQuotaLimit AIAPIQuotaLimit = (AIAPIQuotaLimit) quotaPolicy.getLimit();
-        throttleLimitDTO.setType(throttleLimitDTO.getType());
-        if (throttleLimitDTO.getAiApiQuota().getTotalTokenCount() == null) {
-            aiApiQuotaLimitDTO.setTotalTokenCount(AIAPIQuotaLimit.getTotalTokenCount());
-        }
-        if (throttleLimitDTO.getAiApiQuota().getPromptTokenCount() == null) {
-            aiApiQuotaLimitDTO.setPromptTokenCount(AIAPIQuotaLimit.getPromptTokenCount());
-        }
-        if (throttleLimitDTO.getAiApiQuota().getCompletionTokenCount() == null) {
-            aiApiQuotaLimitDTO.setCompletionTokenCount(AIAPIQuotaLimit.getCompletionTokenCount());
-        }
-        throttleLimitDTO.setAiApiQuota(aiApiQuotaLimitDTO);
-        return throttleLimitDTO;
-    }
 }
diff --git a/...imgt.core.feature/src/main/resources/policy_templates/throttle_policy_template_ai_sub.xml b/...imgt.core.feature/src/main/resources/policy_templates/throttle_policy_template_ai_sub.xml
@@ -13,25 +13,25 @@
 		#if($quotaPolicy != "")
 		FROM EligibilityStream[isEligible==true]#throttler:timeBatch($quotaPolicy.getLimit().getUnitTime() $quotaPolicy.getLimit().getTimeUnit(), 0)
 		select throttleKey,
-		#if($quotaPolicy.getLimit().getRequestCount() != "")
-		ifThenElse((count(messageID) >= 1000), true,
+		#if($quotaPolicy.getLimit().getRequestCount() != 0)
+		ifThenElse((count(messageID) >= $quotaPolicy.getLimit().getRequestCount()), true,
 		#end
-		#if($quotaPolicy.getLimit().getPromptTokenCount() != "")
-		ifThenElse((sum(cast(map:get(propertiesMap,'promptTokenCount'),'long')) >= 30000L), true,
+		#if($quotaPolicy.getAiApiQuota().getPromptTokenCount() != 0)
+		ifThenElse((sum(cast(map:get(propertiesMap,'PROMPT_TOKENS'),'long')) >= $quotaPolicy.getAiApiQuota().getPromptTokenCount()L), true,
 		#end
-		#if($quotaPolicy.getLimit().getCompletionTokenCount() != "")
-		ifThenElse((sum(cast(map:get(propertiesMap,'completionTokenCount'),'long')) >= 30000L), true,
+		#if($quotaPolicy.getAiApiQuota().getCompletionTokenCount() != 0)
+		ifThenElse((sum(cast(map:get(propertiesMap,'COMPLETION_TOKENS'),'long')) >= $quotaPolicy.getAiApiQuota().getCompletionTokenCount()L), true,
 		#end
-		#if($quotaPolicy.getLimit().getTotalTokenCount() != "")
-		ifThenElse((sum(cast(map:get(propertiesMap,'totalTokenCount'),'long')) >= 30000L), true, false)
+		#if($quotaPolicy.getAiApiQuota().getTotalTokenCount() != 0)
+		ifThenElse((sum(cast(map:get(propertiesMap,'TOTAL_TOKENS'),'long')) >= $quotaPolicy.getAiApiQuota().getTotalTokenCount()L), true, false)
 		#end
-		#if($quotaPolicy.getLimit().getPromptTokenCount() != "")
+		#if($quotaPolicy.getLimit().getPromptTokenCount() != 0)
 		)
 		#end
-		#if($quotaPolicy.getLimit().getCompletionTokenCount() != "")
+		#if($quotaPolicy.getLimit().getCompletionTokenCount() != 0)
 		)
 		#end
-		#if($quotaPolicy.getRequestCountLimit().getRequestCount() != "")
+		#if($quotaPolicy.getRequestCountLimit().getRequestCount() != 0)
 		)
 		#end
 		as isThrottled, expiryTimeStamp