From 20c863433d9b47820e4a3cbbe969f6e5d0e43620 Mon Sep 17 00:00:00 2001 From: Savindu Dimal Date: Wed, 18 Sep 2024 17:33:20 +0530 Subject: [PATCH] Add REST API changes to filter AI API subscription rate limiting policies --- .../wso2/carbon/apimgt/api/model/Tier.java | 28 ++++++++ .../carbon/apimgt/impl/utils/APIUtil.java | 4 +- .../src/main/resources/publisher-api.yaml | 29 ++++++++ .../publisher/v1/dto/ThrottlingPolicyDTO.java | 68 ++++++++++++++++++- .../mappings/ThrottlingPolicyMappingUtil.java | 3 + .../apimgt/rest/api/publisher/v1/ApisApi.java | 4 +- .../rest/api/publisher/v1/ApisApiService.java | 2 +- .../publisher/v1/ThrottlingPoliciesApi.java | 4 +- .../v1/ThrottlingPoliciesApiService.java | 2 +- .../publisher/v1/impl/ApisApiServiceImpl.java | 13 ++-- .../ThrottlingPoliciesApiServiceImpl.java | 20 +++++- .../src/main/resources/publisher-api.yaml | 29 ++++++++ 12 files changed, 188 insertions(+), 18 deletions(-) diff --git a/components/apimgt/org.wso2.carbon.apimgt.api/src/main/java/org/wso2/carbon/apimgt/api/model/Tier.java b/components/apimgt/org.wso2.carbon.apimgt.api/src/main/java/org/wso2/carbon/apimgt/api/model/Tier.java index 85ab2b2ed29e..4d3317d2cb99 100644 --- a/components/apimgt/org.wso2.carbon.apimgt.api/src/main/java/org/wso2/carbon/apimgt/api/model/Tier.java +++ b/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{ private int rateLimitCount; private String rateLimitTimeUnit; private String bandwidthDataUnit; + private long totalTokenCount; + private long promptTokenCount; + private long completionTokenCount; public Map 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; 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 99f7059f3344..ab764ab491c1 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 @@ -6485,10 +6485,12 @@ public static Map 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()); diff --git a/components/apimgt/org.wso2.carbon.apimgt.rest.api.common/src/main/resources/publisher-api.yaml b/components/apimgt/org.wso2.carbon.apimgt.rest.api.common/src/main/resources/publisher-api.yaml index 51ec67f66186..69f7db689074 100644 --- a/components/apimgt/org.wso2.carbon.apimgt.rest.api.common/src/main/resources/publisher-api.yaml +++ b/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 diff --git a/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 b/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 index c4e820f48802..12b8227c845a 100644 --- a/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 +++ b/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 attributes = new HashMap(); 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"); diff --git a/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 b/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 index 0f6d1c0bf631..6bc044f24fb5 100644 --- a/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 +++ b/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())); } diff --git a/components/apimgt/org.wso2.carbon.apimgt.rest.api.publisher.v1/src/gen/java/org/wso2/carbon/apimgt/rest/api/publisher/v1/ApisApi.java b/components/apimgt/org.wso2.carbon.apimgt.rest.api.publisher.v1/src/gen/java/org/wso2/carbon/apimgt/rest/api/publisher/v1/ApisApi.java index 1b1b92debca2..dbafa7734150 100644 --- a/components/apimgt/org.wso2.carbon.apimgt.rest.api.publisher.v1/src/gen/java/org/wso2/carbon/apimgt/rest/api/publisher/v1/ApisApi.java +++ b/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 diff --git a/components/apimgt/org.wso2.carbon.apimgt.rest.api.publisher.v1/src/gen/java/org/wso2/carbon/apimgt/rest/api/publisher/v1/ApisApiService.java b/components/apimgt/org.wso2.carbon.apimgt.rest.api.publisher.v1/src/gen/java/org/wso2/carbon/apimgt/rest/api/publisher/v1/ApisApiService.java index 109eed40030d..5d47c2da35bb 100644 --- a/components/apimgt/org.wso2.carbon.apimgt.rest.api.publisher.v1/src/gen/java/org/wso2/carbon/apimgt/rest/api/publisher/v1/ApisApiService.java +++ b/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; diff --git a/components/apimgt/org.wso2.carbon.apimgt.rest.api.publisher.v1/src/gen/java/org/wso2/carbon/apimgt/rest/api/publisher/v1/ThrottlingPoliciesApi.java b/components/apimgt/org.wso2.carbon.apimgt.rest.api.publisher.v1/src/gen/java/org/wso2/carbon/apimgt/rest/api/publisher/v1/ThrottlingPoliciesApi.java index b20196aa1298..5582ec551c58 100644 --- a/components/apimgt/org.wso2.carbon.apimgt.rest.api.publisher.v1/src/gen/java/org/wso2/carbon/apimgt/rest/api/publisher/v1/ThrottlingPoliciesApi.java +++ b/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 diff --git a/components/apimgt/org.wso2.carbon.apimgt.rest.api.publisher.v1/src/gen/java/org/wso2/carbon/apimgt/rest/api/publisher/v1/ThrottlingPoliciesApiService.java b/components/apimgt/org.wso2.carbon.apimgt.rest.api.publisher.v1/src/gen/java/org/wso2/carbon/apimgt/rest/api/publisher/v1/ThrottlingPoliciesApiService.java index f61fd3f800d6..fc804d780c54 100644 --- a/components/apimgt/org.wso2.carbon.apimgt.rest.api.publisher.v1/src/gen/java/org/wso2/carbon/apimgt/rest/api/publisher/v1/ThrottlingPoliciesApiService.java +++ b/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; } diff --git a/components/apimgt/org.wso2.carbon.apimgt.rest.api.publisher.v1/src/main/java/org/wso2/carbon/apimgt/rest/api/publisher/v1/impl/ApisApiServiceImpl.java b/components/apimgt/org.wso2.carbon.apimgt.rest.api.publisher.v1/src/main/java/org/wso2/carbon/apimgt/rest/api/publisher/v1/impl/ApisApiServiceImpl.java index b110333d5885..dd19e3d3b8d1 100644 --- a/components/apimgt/org.wso2.carbon.apimgt.rest.api.publisher.v1/src/main/java/org/wso2/carbon/apimgt/rest/api/publisher/v1/impl/ApisApiServiceImpl.java +++ b/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 availableThrottlingPolicyList = new ThrottlingPoliciesApiServiceImpl() - .getThrottlingPolicyList(ThrottlingPolicyDTO.PolicyLevelEnum.SUBSCRIPTION.toString(), true); + List availableThrottlingPolicyList = new ThrottlingPoliciesApiServiceImpl().getThrottlingPolicyList( + ThrottlingPolicyDTO.PolicyLevelEnum.SUBSCRIPTION.toString(), true, isAiApi); - if (apiInfo != null ) { + if (apiInfo != null) { List apiPolicies = apiInfo.getPolicies(); - List apiThrottlingPolicies = ApisApiServiceImplUtils.filterAPIThrottlingPolicies(apiPolicies, availableThrottlingPolicyList); + List apiThrottlingPolicies = ApisApiServiceImplUtils.filterAPIThrottlingPolicies(apiPolicies, + availableThrottlingPolicyList); return Response.ok().entity(apiThrottlingPolicies).build(); } return null; diff --git a/components/apimgt/org.wso2.carbon.apimgt.rest.api.publisher.v1/src/main/java/org/wso2/carbon/apimgt/rest/api/publisher/v1/impl/ThrottlingPoliciesApiServiceImpl.java b/components/apimgt/org.wso2.carbon.apimgt.rest.api.publisher.v1/src/main/java/org/wso2/carbon/apimgt/rest/api/publisher/v1/impl/ThrottlingPoliciesApiServiceImpl.java index 4cc00dffb79a..e3a71fe8d771 100644 --- a/components/apimgt/org.wso2.carbon.apimgt.rest.api.publisher.v1/src/main/java/org/wso2/carbon/apimgt/rest/api/publisher/v1/impl/ThrottlingPoliciesApiServiceImpl.java +++ b/components/apimgt/org.wso2.carbon.apimgt.rest.api.publisher.v1/src/main/java/org/wso2/carbon/apimgt/rest/api/publisher/v1/impl/ThrottlingPoliciesApiServiceImpl.java @@ -68,12 +68,12 @@ public class ThrottlingPoliciesApiServiceImpl implements ThrottlingPoliciesApiSe */ @Override public Response getAllThrottlingPolicies(String policyLevel, Integer limit, Integer offset, - String ifNoneMatch, MessageContext messageContext) { + String ifNoneMatch, Boolean isAiApi, MessageContext messageContext) { //pre-processing //setting default limit and offset if they are null limit = limit != null ? limit : RestApiConstants.PAGINATION_LIMIT_DEFAULT; offset = offset != null ? offset : RestApiConstants.PAGINATION_OFFSET_DEFAULT; - List tierList = getThrottlingPolicyList(policyLevel, false); + List tierList = getThrottlingPolicyList(policyLevel, false, isAiApi); ThrottlingPolicyListDTO policyListDTO = ThrottlingPolicyMappingUtil .fromTierListToDTO(tierList, policyLevel, limit, offset); //todo: set total counts properly @@ -179,7 +179,8 @@ public Response getThrottlingPolicyByName(String policyName, String policyLevel, * @param policyLevel * @return list of throttling policies */ - public List getThrottlingPolicyList(String policyLevel, boolean includeAsyncPolicies) { + public List getThrottlingPolicyList(String policyLevel, boolean includeAsyncPolicies, + boolean filterAIAPIPolicies) { try { List tierList = new ArrayList<>(); String tenantDomain = RestApiCommonUtil.getLoggedInUserTenantDomain(); @@ -197,6 +198,19 @@ public List getThrottlingPolicyList(String policyLevel, boolean includeAsy tierList = tierList.stream().filter(tier -> !PolicyConstants.EVENT_COUNT_TYPE.equals( tier.getQuotaPolicyType())).collect(Collectors.toList()); } + + // Handle filtering for AI API policies + if (filterAIAPIPolicies) { + // Keep tiers where QuotaPolicyType is AI_API_QUOTA_TYPE or is null + tierList = tierList.stream() + .filter(tier -> PolicyConstants.AI_API_QUOTA_TYPE.equals(tier.getQuotaPolicyType()) + || tier.getQuotaPolicyType() == null).collect(Collectors.toList()); + } else { + // Remove tiers where QuotaPolicyType is AI_API_QUOTA_TYPE + tierList = tierList.stream() + .filter(tier -> !PolicyConstants.AI_API_QUOTA_TYPE.equals(tier.getQuotaPolicyType())) + .collect(Collectors.toList()); + } } else if (ThrottlingPolicyDTO.PolicyLevelEnum.API.toString().equals(policyLevel)) { Map resourceTiersMap = APIUtil.getTiers(APIConstants.TIER_RESOURCE_TYPE, tenantDomain); diff --git a/components/apimgt/org.wso2.carbon.apimgt.rest.api.publisher.v1/src/main/resources/publisher-api.yaml b/components/apimgt/org.wso2.carbon.apimgt.rest.api.publisher.v1/src/main/resources/publisher-api.yaml index 51ec67f66186..69f7db689074 100644 --- a/components/apimgt/org.wso2.carbon.apimgt.rest.api.publisher.v1/src/main/resources/publisher-api.yaml +++ b/components/apimgt/org.wso2.carbon.apimgt.rest.api.publisher.v1/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