Add REST API changes to filter AI API subscription rate limiting poli…

…cies

components/apimgt/org.wso2.carbon.apimgt.api/src/main/java/org/wso2/carbon/apimgt/api/model/Tier.java

@@ -47,6 +47,9 @@ public class Tier implements Serializable, Comparable<Tier>{
     private int rateLimitCount;
     private String rateLimitTimeUnit;
     private String bandwidthDataUnit;
+    private long totalTokenCount;
+    private long promptTokenCount;
+    private long completionTokenCount;
 
     public Map<String, String> getMonetizationAttributes() {
         return monetizationAttributes;
@@ -183,6 +186,31 @@ public String getBandwidthDataUnit() {
         return bandwidthDataUnit;
     }
 
+    public long getTotalTokenCount() {
+        return totalTokenCount;
+    }
+
+    public void setTotalTokenCount(long totalTokenCount) {
+        this.totalTokenCount = totalTokenCount;
+    }
+
+    public long getPromptTokenCount() {
+        return promptTokenCount;
+    }
+
+    public void setPromptTokenCount(long promptTokenCount) {
+        this.promptTokenCount = promptTokenCount;
+    }
+
+    public long getCompletionTokenCount() {
+        return completionTokenCount;
+    }
+
+    public void setCompletionTokenCount(long completionTokenCount) {
+        this.completionTokenCount = completionTokenCount;
+    }
+
+
     @Override
     public boolean equals(Object o) {
         if (this == o) return true;

components/apimgt/org.wso2.carbon.apimgt.impl/src/main/java/org/wso2/carbon/apimgt/impl/utils/APIUtil.java

@@ -6485,10 +6485,12 @@ public static Map<String, Tier> getTiersFromPolicies(String policyLevel, int ten
                     tier.setRequestCount(bandwidthLimit.getDataAmount());
                     tier.setBandwidthDataUnit(bandwidthLimit.getDataUnit());
                 } else if (limit instanceof AIAPIQuotaLimit){
-                    // Todo: Need to implement this according to publisher and developer portals' requirements
                     AIAPIQuotaLimit AIAPIQuotaLimit = (AIAPIQuotaLimit) limit;
                     tier.setRequestsPerMin(AIAPIQuotaLimit.getRequestCount());
                     tier.setRequestCount(AIAPIQuotaLimit.getRequestCount());
+                    tier.setTotalTokenCount(AIAPIQuotaLimit.getTotalTokenCount());
+                    tier.setPromptTokenCount(AIAPIQuotaLimit.getPromptTokenCount());
+                    tier.setCompletionTokenCount(AIAPIQuotaLimit.getCompletionTokenCount());
                 } else {
                     EventCountLimit eventCountLimit = (EventCountLimit) limit;
                     tier.setRequestCount(eventCountLimit.getEventCount());

components/apimgt/org.wso2.carbon.apimgt.rest.api.common/src/main/resources/publisher-api.yaml

@@ -1105,6 +1105,7 @@ paths:
         - $ref: '#/components/parameters/apiId'
         - $ref: '#/components/parameters/requestedTenant'
         - $ref: '#/components/parameters/If-None-Match'
+        - $ref: '#/components/parameters/isAiApi'
       responses:
         200:
           description: |
@@ -4442,6 +4443,7 @@ paths:
         - $ref: '#/components/parameters/offset'
         - $ref: '#/components/parameters/policyLevel'
         - $ref: '#/components/parameters/If-None-Match'
+        - $ref: '#/components/parameters/isAiApi'
       responses:
         200:
           description: |
@@ -11058,6 +11060,24 @@ components:
             Unit of data allowed to be transfered. Allowed values are "KB", "MB" and "GB"
           type: string
           example: KB
+        totalTokenCount:
+          type: integer
+          description: |
+            Maximum number of total tokens which can be used within a provided unit time
+          format: int64
+          example: 1000
+        promptTokenCount:
+          type: integer
+          description: |
+            Maximum number of prompt tokens which can be used within a provided unit time
+          format: int64
+          example: 500
+        completionTokenCount:
+          type: integer
+          description: |
+            Maximum number of completion tokens which can be used within a provided unit time
+          format: int64
+          example: 600
         unitTime:
           type: integer
           format: int64
@@ -11080,6 +11100,7 @@ components:
           enum:
             - REQUESTCOUNT
             - BANDWIDTHVOLUME
+            - AIAPIQUOTA
           example: REQUESTCOUNT
         tierPlan:
           type: string
@@ -13774,6 +13795,14 @@ components:
       schema:
         type: string
         default: application/json
+    isAiApi:
+      name: isAiApi
+      in: query
+      description: |
+        Indicates the quota policy type to be AI API quota or not.
+      schema:
+        type: boolean
+        default: false
     offset:
       name: offset
       in: query

components/apimgt/org.wso2.carbon.apimgt.rest.api.publisher.v1.common/src/gen/java/org/wso2/carbon/apimgt/rest/api/publisher/v1/dto/ThrottlingPolicyDTO.java

@@ -61,6 +61,9 @@ public static PolicyLevelEnum fromValue(String v) {
     private Map<String, String> attributes = new HashMap<String, String>();
     private Long requestCount = null;
     private String dataUnit = null;
+    private Long totalTokenCount = null;
+    private Long promptTokenCount = null;
+    private Long completionTokenCount = null;
     private Long unitTime = null;
     private String timeUnit = null;
     private Integer rateLimitCount = 0;
@@ -70,7 +73,8 @@ public static PolicyLevelEnum fromValue(String v) {
     @XmlEnum(String.class)
     public enum QuotaPolicyTypeEnum {
         REQUESTCOUNT("REQUESTCOUNT"),
-        BANDWIDTHVOLUME("BANDWIDTHVOLUME");
+        BANDWIDTHVOLUME("BANDWIDTHVOLUME"),
+        AIAPIQUOTA("AIAPIQUOTA");
         private String value;
 
         QuotaPolicyTypeEnum (String v) {
@@ -256,6 +260,60 @@ public void setDataUnit(String dataUnit) {
     this.dataUnit = dataUnit;
   }
 
+  /**
+   * Maximum number of total tokens which can be used within a provided unit time 
+   **/
+  public ThrottlingPolicyDTO totalTokenCount(Long totalTokenCount) {
+    this.totalTokenCount = totalTokenCount;
+    return this;
+  }
+
+
+  @ApiModelProperty(example = "1000", value = "Maximum number of total tokens which can be used within a provided unit time ")
+  @JsonProperty("totalTokenCount")
+  public Long getTotalTokenCount() {
+    return totalTokenCount;
+  }
+  public void setTotalTokenCount(Long totalTokenCount) {
+    this.totalTokenCount = totalTokenCount;
+  }
+
+  /**
+   * Maximum number of prompt tokens which can be used within a provided unit time 
+   **/
+  public ThrottlingPolicyDTO promptTokenCount(Long promptTokenCount) {
+    this.promptTokenCount = promptTokenCount;
+    return this;
+  }
+
+
+  @ApiModelProperty(example = "500", value = "Maximum number of prompt tokens which can be used within a provided unit time ")
+  @JsonProperty("promptTokenCount")
+  public Long getPromptTokenCount() {
+    return promptTokenCount;
+  }
+  public void setPromptTokenCount(Long promptTokenCount) {
+    this.promptTokenCount = promptTokenCount;
+  }
+
+  /**
+   * Maximum number of completion tokens which can be used within a provided unit time 
+   **/
+  public ThrottlingPolicyDTO completionTokenCount(Long completionTokenCount) {
+    this.completionTokenCount = completionTokenCount;
+    return this;
+  }
+
+
+  @ApiModelProperty(example = "600", value = "Maximum number of completion tokens which can be used within a provided unit time ")
+  @JsonProperty("completionTokenCount")
+  public Long getCompletionTokenCount() {
+    return completionTokenCount;
+  }
+  public void setCompletionTokenCount(Long completionTokenCount) {
+    this.completionTokenCount = completionTokenCount;
+  }
+
   /**
    **/
   public ThrottlingPolicyDTO unitTime(Long unitTime) {
@@ -418,6 +476,9 @@ public boolean equals(java.lang.Object o) {
         Objects.equals(attributes, throttlingPolicy.attributes) &&
         Objects.equals(requestCount, throttlingPolicy.requestCount) &&
         Objects.equals(dataUnit, throttlingPolicy.dataUnit) &&
+        Objects.equals(totalTokenCount, throttlingPolicy.totalTokenCount) &&
+        Objects.equals(promptTokenCount, throttlingPolicy.promptTokenCount) &&
+        Objects.equals(completionTokenCount, throttlingPolicy.completionTokenCount) &&
         Objects.equals(unitTime, throttlingPolicy.unitTime) &&
         Objects.equals(timeUnit, throttlingPolicy.timeUnit) &&
         Objects.equals(rateLimitCount, throttlingPolicy.rateLimitCount) &&
@@ -430,7 +491,7 @@ public boolean equals(java.lang.Object o) {
 
   @Override
   public int hashCode() {
-    return Objects.hash(name, description, policyLevel, displayName, attributes, requestCount, dataUnit, unitTime, timeUnit, rateLimitCount, rateLimitTimeUnit, quotaPolicyType, tierPlan, stopOnQuotaReach, monetizationProperties);
+    return Objects.hash(name, description, policyLevel, displayName, attributes, requestCount, dataUnit, totalTokenCount, promptTokenCount, completionTokenCount, unitTime, timeUnit, rateLimitCount, rateLimitTimeUnit, quotaPolicyType, tierPlan, stopOnQuotaReach, monetizationProperties);
   }
 
   @Override
@@ -445,6 +506,9 @@ public String toString() {
     sb.append("    attributes: ").append(toIndentedString(attributes)).append("\n");
     sb.append("    requestCount: ").append(toIndentedString(requestCount)).append("\n");
     sb.append("    dataUnit: ").append(toIndentedString(dataUnit)).append("\n");
+    sb.append("    totalTokenCount: ").append(toIndentedString(totalTokenCount)).append("\n");
+    sb.append("    promptTokenCount: ").append(toIndentedString(promptTokenCount)).append("\n");
+    sb.append("    completionTokenCount: ").append(toIndentedString(completionTokenCount)).append("\n");
     sb.append("    unitTime: ").append(toIndentedString(unitTime)).append("\n");
     sb.append("    timeUnit: ").append(toIndentedString(timeUnit)).append("\n");
     sb.append("    rateLimitCount: ").append(toIndentedString(rateLimitCount)).append("\n");

components/apimgt/org.wso2.carbon.apimgt.rest.api.publisher.v1.common/src/main/java/org/wso2/carbon/apimgt/rest/api/publisher/v1/common/mappings/ThrottlingPolicyMappingUtil.java

@@ -133,6 +133,9 @@ public static ThrottlingPolicyDTO fromTierToDTO(Tier tier, String tierLevel) {
         dto.setRateLimitCount(tier.getRateLimitCount());
         dto.setRateLimitTimeUnit(tier.getRateLimitTimeUnit());
         dto.setDataUnit(tier.getBandwidthDataUnit());
+        dto.setTotalTokenCount(tier.getTotalTokenCount());
+        dto.setPromptTokenCount(tier.getPromptTokenCount());
+        dto.setCompletionTokenCount(tier.getCompletionTokenCount());
         if (tier.getQuotaPolicyType() != null) {
             dto.setQuotaPolicyType(mapQuotaPolicyTypeFromModeltoDTO(tier.getQuotaPolicyType()));
         }

components/apimgt/org.wso2.carbon.apimgt.rest.api.publisher.v1/src/gen/java/org/wso2/carbon/apimgt/rest/api/publisher/v1/ApisApi.java

@@ -1112,8 +1112,8 @@ public Response getAPISpecificOperationPolicyContentByPolicyId(@ApiParam(value =
         @ApiResponse(code = 304, message = "Not Modified. Empty body because the client has already the latest version of the requested resource. ", response = Void.class),
         @ApiResponse(code = 404, message = "Not Found. The specified resource does not exist.", response = ErrorDTO.class),
         @ApiResponse(code = 406, message = "Not Acceptable. The requested media type is not supported.", response = ErrorDTO.class) })
-    public Response getAPISubscriptionPolicies(@ApiParam(value = "**API ID** consisting of the **UUID** of the API. ",required=true) @PathParam("apiId") String apiId,  @ApiParam(value = "For cross-tenant invocations, this is used to specify the tenant domain, where the resource need to be   retirieved from. " )@HeaderParam("X-WSO2-Tenant") String xWSO2Tenant,  @ApiParam(value = "Validator for conditional requests; based on the ETag of the formerly retrieved variant of the resource. " )@HeaderParam("If-None-Match") String ifNoneMatch) throws APIManagementException{
-        return delegate.getAPISubscriptionPolicies(apiId, xWSO2Tenant, ifNoneMatch, securityContext);
+    public Response getAPISubscriptionPolicies(@ApiParam(value = "**API ID** consisting of the **UUID** of the API. ",required=true) @PathParam("apiId") String apiId,  @ApiParam(value = "For cross-tenant invocations, this is used to specify the tenant domain, where the resource need to be   retirieved from. " )@HeaderParam("X-WSO2-Tenant") String xWSO2Tenant,  @ApiParam(value = "Validator for conditional requests; based on the ETag of the formerly retrieved variant of the resource. " )@HeaderParam("If-None-Match") String ifNoneMatch,  @ApiParam(value = "Indicates the quota policy type to be AI API quota or not. ", defaultValue="false") @DefaultValue("false") @QueryParam("isAiApi") Boolean isAiApi) throws APIManagementException{
+        return delegate.getAPISubscriptionPolicies(apiId, xWSO2Tenant, ifNoneMatch, isAiApi, securityContext);
     }
 
     @GET

components/apimgt/org.wso2.carbon.apimgt.rest.api.publisher.v1/src/gen/java/org/wso2/carbon/apimgt/rest/api/publisher/v1/ApisApiService.java

@@ -115,7 +115,7 @@ public interface ApisApiService {
       public Response getAPIRevisionDeployments(String apiId, MessageContext messageContext) throws APIManagementException;
       public Response getAPIRevisions(String apiId, String query, MessageContext messageContext) throws APIManagementException;
       public Response getAPISpecificOperationPolicyContentByPolicyId(String apiId, String operationPolicyId, MessageContext messageContext) throws APIManagementException;
-      public Response getAPISubscriptionPolicies(String apiId, String xWSO2Tenant, String ifNoneMatch, MessageContext messageContext) throws APIManagementException;
+      public Response getAPISubscriptionPolicies(String apiId, String xWSO2Tenant, String ifNoneMatch, Boolean isAiApi, MessageContext messageContext) throws APIManagementException;
       public Response getAPISwagger(String apiId, String ifNoneMatch, MessageContext messageContext) throws APIManagementException;
       public Response getAPIThumbnail(String apiId, String ifNoneMatch, MessageContext messageContext) throws APIManagementException;
       public Response getAllAPISpecificOperationPolicies(String apiId, Integer limit, Integer offset, String query, MessageContext messageContext) throws APIManagementException;

components/apimgt/org.wso2.carbon.apimgt.rest.api.publisher.v1/src/gen/java/org/wso2/carbon/apimgt/rest/api/publisher/v1/ThrottlingPoliciesApi.java

@@ -53,8 +53,8 @@ public class ThrottlingPoliciesApi  {
         @ApiResponse(code = 200, message = "OK. List of policies returned. ", response = ThrottlingPolicyListDTO.class),
         @ApiResponse(code = 304, message = "Not Modified. Empty body because the client has already the latest version of the requested resource (Will be supported in future). ", response = Void.class),
         @ApiResponse(code = 406, message = "Not Acceptable. The requested media type is not supported.", response = ErrorDTO.class) })
-    public Response getAllThrottlingPolicies(@ApiParam(value = "List API or Application or Resource type policies. ",required=true, allowableValues="api, subcription") @PathParam("policyLevel") String policyLevel,  @ApiParam(value = "Maximum size of resource array to return. ", defaultValue="25") @DefaultValue("25") @QueryParam("limit") Integer limit,  @ApiParam(value = "Starting point within the complete list of items qualified. ", defaultValue="0") @DefaultValue("0") @QueryParam("offset") Integer offset,  @ApiParam(value = "Validator for conditional requests; based on the ETag of the formerly retrieved variant of the resource. " )@HeaderParam("If-None-Match") String ifNoneMatch) throws APIManagementException{
-        return delegate.getAllThrottlingPolicies(policyLevel, limit, offset, ifNoneMatch, securityContext);
+    public Response getAllThrottlingPolicies(@ApiParam(value = "List API or Application or Resource type policies. ",required=true, allowableValues="api, subcription") @PathParam("policyLevel") String policyLevel,  @ApiParam(value = "Maximum size of resource array to return. ", defaultValue="25") @DefaultValue("25") @QueryParam("limit") Integer limit,  @ApiParam(value = "Starting point within the complete list of items qualified. ", defaultValue="0") @DefaultValue("0") @QueryParam("offset") Integer offset,  @ApiParam(value = "Validator for conditional requests; based on the ETag of the formerly retrieved variant of the resource. " )@HeaderParam("If-None-Match") String ifNoneMatch,  @ApiParam(value = "Indicates the quota policy type to be AI API quota or not. ", defaultValue="false") @DefaultValue("false") @QueryParam("isAiApi") Boolean isAiApi) throws APIManagementException{
+        return delegate.getAllThrottlingPolicies(policyLevel, limit, offset, ifNoneMatch, isAiApi, securityContext);
     }
 
     @GET

components/apimgt/org.wso2.carbon.apimgt.rest.api.publisher.v1/src/gen/java/org/wso2/carbon/apimgt/rest/api/publisher/v1/ThrottlingPoliciesApiService.java

@@ -23,7 +23,7 @@
 
 
 public interface ThrottlingPoliciesApiService {
-      public Response getAllThrottlingPolicies(String policyLevel, Integer limit, Integer offset, String ifNoneMatch, MessageContext messageContext) throws APIManagementException;
+      public Response getAllThrottlingPolicies(String policyLevel, Integer limit, Integer offset, String ifNoneMatch, Boolean isAiApi, MessageContext messageContext) throws APIManagementException;
       public Response getSubscriptionThrottlingPolicies(Integer limit, Integer offset, String ifNoneMatch, MessageContext messageContext) throws APIManagementException;
       public Response getThrottlingPolicyByName(String policyName, String policyLevel, String ifNoneMatch, MessageContext messageContext) throws APIManagementException;
 }

components/apimgt/org.wso2.carbon.apimgt.rest.api.publisher.v1/src/main/java/org/wso2/carbon/apimgt/rest/api/publisher/v1/impl/ApisApiServiceImpl.java

@@ -3575,17 +3575,18 @@ public Response generateMockScripts(String apiId, String ifNoneMatch, MessageCon
     }
 
     @Override
-    public Response getAPISubscriptionPolicies(String apiId, String ifNoneMatch, String xWSO2Tenant,
-                                                     MessageContext messageContext) throws APIManagementException {
+    public Response getAPISubscriptionPolicies(String apiId, String xWSO2Tenant, String ifNoneMatch, Boolean isAiApi,
+            MessageContext messageContext) throws APIManagementException {
         APIProvider apiProvider = RestApiCommonUtil.getLoggedInUserProvider();
         String organization = RestApiUtil.getValidatedOrganization(messageContext);
         APIDTO apiInfo = getAPIByID(apiId, apiProvider, organization);
-        List<Tier> availableThrottlingPolicyList = new ThrottlingPoliciesApiServiceImpl()
-                .getThrottlingPolicyList(ThrottlingPolicyDTO.PolicyLevelEnum.SUBSCRIPTION.toString(), true);
+        List<Tier> availableThrottlingPolicyList = new ThrottlingPoliciesApiServiceImpl().getThrottlingPolicyList(
+                ThrottlingPolicyDTO.PolicyLevelEnum.SUBSCRIPTION.toString(), true, isAiApi);
 
-        if (apiInfo != null ) {
+        if (apiInfo != null) {
             List<String> apiPolicies = apiInfo.getPolicies();
-            List<Tier> apiThrottlingPolicies = ApisApiServiceImplUtils.filterAPIThrottlingPolicies(apiPolicies, availableThrottlingPolicyList);
+            List<Tier> apiThrottlingPolicies = ApisApiServiceImplUtils.filterAPIThrottlingPolicies(apiPolicies,
+                    availableThrottlingPolicyList);
             return Response.ok().entity(apiThrottlingPolicies).build();
         }
         return null;