diff --git a/common-controller/internal/operator/controllers/dp/airatelimitpolicy_controller.go b/common-controller/internal/operator/controllers/dp/airatelimitpolicy_controller.go index 4791376a3..98789f89a 100644 --- a/common-controller/internal/operator/controllers/dp/airatelimitpolicy_controller.go +++ b/common-controller/internal/operator/controllers/dp/airatelimitpolicy_controller.go @@ -105,6 +105,9 @@ func (r *AIRateLimitPolicyReconciler) Reconcile(ctx context.Context, req ctrl.Re xds.UpdateRateLimiterPolicies(conf.CommonController.Server.Label) } else { loggers.LoggerAPKOperator.Infof("ratelimits found") + if ratelimitPolicy.Spec.Override == nil { + ratelimitPolicy.Spec.Override = ratelimitPolicy.Spec.Default + } if ratelimitPolicy.Spec.TargetRef.Name != "" { r.ods.AddorUpdateAIRatelimitToStore(ratelimitKey, ratelimitPolicy.Spec) xds.UpdateRateLimitXDSCacheForAIRatelimitPolicies(r.ods.GetAIRatelimitPolicySpecs()) diff --git a/runtime/config-deployer-service/ballerina/APIClient.bal b/runtime/config-deployer-service/ballerina/APIClient.bal index 1469775e0..0dda6fdb1 100644 --- a/runtime/config-deployer-service/ballerina/APIClient.bal +++ b/runtime/config-deployer-service/ballerina/APIClient.bal @@ -232,6 +232,11 @@ public class APIClient { serviceEntry: false, url: self.constructURlFromService(sandboxEndpointConfig.endpoint) }; + AIRatelimit? aiRatelimit = sandboxEndpointConfig.aiRatelimit; + if aiRatelimit is AIRatelimit && aiRatelimit.enabled { + model:AIRateLimitPolicy airl = self.generateAIRateLimitPolicyCR(apkConf, aiRatelimit.token, aiRatelimit.request, backendService.metadata.name, organization); + apiArtifact.aiRatelimitPolicies[airl.metadata.name] = airl; + } } } if (endpointType == () || endpointType == PRODUCTION_TYPE) { @@ -246,6 +251,11 @@ public class APIClient { serviceEntry: false, url: self.constructURlFromService(productionEndpointConfig.endpoint) }; + AIRatelimit? aiRatelimit = productionEndpointConfig.aiRatelimit; + if aiRatelimit is AIRatelimit && aiRatelimit.enabled { + model:AIRateLimitPolicy airl = self.generateAIRateLimitPolicyCR(apkConf, aiRatelimit.token, aiRatelimit.request, backendService.metadata.name, organization); + apiArtifact.aiRatelimitPolicies[airl.metadata.name] = airl; + } } } return endpointIdMap; @@ -1367,6 +1377,29 @@ public class APIClient { return rateLimitPolicyCR; } + public isolated function generateAIRateLimitPolicyCR(APKConf apkConf, TokenAIRL tokenAIRL, RequestAIRL requestAIRL, string targetRefName, commons:Organization organization) returns model:AIRateLimitPolicy { + string apiIdentifierHash = crypto:hashSha1((apkConf.name + apkConf.version).toBytes()).toBase16(); + model:AIRateLimitPolicy aiRateLimitPolicyCR = { + metadata: { + name: self.retrieveAIRateLimitPolicyName(apiIdentifierHash, targetRefName), + labels: self.getLabels(apkConf, organization) + }, + spec: { + default: { + organization: organization.name, + requestCount: {unit: requestAIRL.unit, requestsPerUnit: requestAIRL.requestLimit}, + tokenCount: {unit: tokenAIRL.unit, requestTokenCount: tokenAIRL.promptLimit, responseTokenCount: tokenAIRL.completionLimit, totalTokenCount: tokenAIRL.totalLimit} + }, + targetRef: { + group: "dp.wso2.com", + kind: "Backend", + name: targetRefName + } + } + }; + return aiRateLimitPolicyCR; + } + isolated function retrieveRateLimitData(RateLimit rateLimit, commons:Organization organization) returns model:RateLimitData { model:RateLimitData rateLimitData = { api: { @@ -1794,6 +1827,10 @@ public class APIClient { } } + public isolated function retrieveAIRateLimitPolicyName(string apiID, string targetRef) returns string { + return "airl-" + apiID + "-" + targetRef; + } + private isolated function validateAndRetrieveAPKConfiguration(json apkconfJson) returns APKConf|commons:APKError? { do { runtimeapi:APKConfValidationResponse validationResponse = check apkConfValidator.validate(apkconfJson.toJsonString()); diff --git a/runtime/config-deployer-service/ballerina/ConfigGenreatorClient.bal b/runtime/config-deployer-service/ballerina/ConfigGenreatorClient.bal index e410cdb51..9209208e5 100644 --- a/runtime/config-deployer-service/ballerina/ConfigGenreatorClient.bal +++ b/runtime/config-deployer-service/ballerina/ConfigGenreatorClient.bal @@ -199,6 +199,10 @@ public class ConfigGeneratorClient { string yamlString = check self.convertJsonToYaml(rateLimitPolicy.toJsonString()); _ = check self.storeFile(yamlString, rateLimitPolicy.metadata.name, zipDir); } + foreach model:AIRateLimitPolicy airateLimitPolicy in apiArtifact.aiRatelimitPolicies { + string yamlString = check self.convertJsonToYaml(airateLimitPolicy.toJsonString()); + _ = check self.storeFile(yamlString, airateLimitPolicy.metadata.name, zipDir); + } foreach model:APIPolicy apiPolicy in apiArtifact.apiPolicies { string yamlString = check self.convertJsonToYaml(apiPolicy.toJsonString()); _ = check self.storeFile(yamlString, apiPolicy.metadata.name, zipDir); diff --git a/runtime/config-deployer-service/ballerina/DeployerClient.bal b/runtime/config-deployer-service/ballerina/DeployerClient.bal index e5beb5d2d..318329010 100644 --- a/runtime/config-deployer-service/ballerina/DeployerClient.bal +++ b/runtime/config-deployer-service/ballerina/DeployerClient.bal @@ -99,6 +99,7 @@ public class DeployerClient { _ = check self.deleteScopeCrsForAPI(existingAPI, apiArtifact?.organization); check self.deleteBackends(existingAPI, apiArtifact?.organization); check self.deleteRateLimitPolicyCRs(existingAPI, apiArtifact?.organization); + check self.deleteAIRateLimitPolicyCRs(existingAPI, apiArtifact?.organization); check self.deleteAPIPolicyCRs(existingAPI, apiArtifact?.organization); check self.deleteInterceptorServiceCRs(existingAPI, apiArtifact?.organization); check self.deleteBackendJWTConfig(existingAPI, apiArtifact?.organization); @@ -118,6 +119,7 @@ public class DeployerClient { check self.deployBackendServices(apiArtifact, ownerReference); check self.deployAuthenticationCRs(apiArtifact, ownerReference); check self.deployRateLimitPolicyCRs(apiArtifact, ownerReference); + check self.deployAIRateLimitPolicyCRs(apiArtifact, ownerReference); check self.deployInterceptorServiceCRs(apiArtifact, ownerReference); check self.deployBackendJWTConfigs(apiArtifact, ownerReference); check self.deployAPIPolicyCRs(apiArtifact, ownerReference); @@ -586,6 +588,30 @@ public class DeployerClient { } } + private isolated function deployAIRateLimitPolicyCRs(model:APIArtifact apiArtifact, model:OwnerReference ownerReference) returns error? { + foreach model:AIRateLimitPolicy rateLimitPolicy in apiArtifact.aiRatelimitPolicies { + rateLimitPolicy.metadata.ownerReferences = [ownerReference]; + http:Response deployRateLimitPolicyResult = check deployAIRateLimitPolicyCR(rateLimitPolicy, apiArtifact?.namespace); + if deployRateLimitPolicyResult.statusCode == http:STATUS_CREATED { + log:printDebug("Deployed AIRateLimitPolicy Successfully" + rateLimitPolicy.toString()); + } else if deployRateLimitPolicyResult.statusCode == http:STATUS_CONFLICT { + log:printDebug("AIRateLimitPolicy already exists" + rateLimitPolicy.toString()); + model:AIRateLimitPolicy rateLimitPolicyFromK8s = check getAIRateLimitPolicyCR(rateLimitPolicy.metadata.name, apiArtifact?.namespace); + rateLimitPolicy.metadata.resourceVersion = rateLimitPolicyFromK8s.metadata.resourceVersion; + http:Response rateLimitPolicyCR = check updateAIRateLimitPolicyCR(rateLimitPolicy, apiArtifact?.namespace); + if rateLimitPolicyCR.statusCode != http:STATUS_OK { + json responsePayLoad = check rateLimitPolicyCR.getJsonPayload(); + model:Status statusResponse = check responsePayLoad.cloneWithType(model:Status); + check self.handleK8sTimeout(statusResponse); + } + } else { + json responsePayLoad = check deployRateLimitPolicyResult.getJsonPayload(); + model:Status statusResponse = check responsePayLoad.cloneWithType(model:Status); + check self.handleK8sTimeout(statusResponse); + } + } + } + private isolated function deleteRateLimitPolicyCRs(model:API api, string organization) returns commons:APKError? { do { model:RateLimitPolicyList|http:ClientError rateLimitPolicyCrListResponse = check getRateLimitPolicyCRsForAPI(api.spec.apiName, api.spec.apiVersion, api.metadata?.namespace, organization); @@ -610,6 +636,29 @@ public class DeployerClient { } } + private isolated function deleteAIRateLimitPolicyCRs(model:API api, string organization) returns commons:APKError? { + do { + model:AIRateLimitPolicyList|http:ClientError aiRateLimitPolicyCrListResponse = check getAIRateLimitPolicyCRsForAPI(api.spec.apiName, api.spec.apiVersion, api.metadata?.namespace, organization); + if aiRateLimitPolicyCrListResponse is model:AIRateLimitPolicyList { + foreach model:AIRateLimitPolicy item in aiRateLimitPolicyCrListResponse.items { + http:Response|http:ClientError rateLimitPolicyCRDeletionResponse = deleteAIRateLimitPolicyCR(item.metadata.name, item.metadata?.namespace); + if rateLimitPolicyCRDeletionResponse is http:Response { + if rateLimitPolicyCRDeletionResponse.statusCode != http:STATUS_OK { + json responsePayLoad = check rateLimitPolicyCRDeletionResponse.getJsonPayload(); + model:Status statusResponse = check responsePayLoad.cloneWithType(model:Status); + check self.handleK8sTimeout(statusResponse); + } + } else { + log:printError("Error occured while deleting AI rate limit policy"); + } + } + return; + } + } on fail var e { + log:printError("Error occured deleting AI rate limit policy", e); + return e909022("Error occured deleting AI rate limit policy", e); + } + } private isolated function deleteAPIPolicyCRs(model:API api, string organization) returns commons:APKError? { do { model:APIPolicyList|http:ClientError apiPolicyCrListResponse = check getAPIPolicyCRsForAPI(api.spec.apiName, api.spec.apiVersion, api.metadata?.namespace, organization); diff --git a/runtime/config-deployer-service/ballerina/K8sClient.bal b/runtime/config-deployer-service/ballerina/K8sClient.bal index 448e8f0eb..624e92931 100644 --- a/runtime/config-deployer-service/ballerina/K8sClient.bal +++ b/runtime/config-deployer-service/ballerina/K8sClient.bal @@ -250,26 +250,51 @@ isolated function deployRateLimitPolicyCR(model:RateLimitPolicy rateLimitPolicy, return k8sApiServerEp->post(endpoint, rateLimitPolicy, targetType = http:Response); } +isolated function deployAIRateLimitPolicyCR(model:AIRateLimitPolicy rateLimitPolicy, string namespace) returns http:Response|http:ClientError { + string endpoint = "/apis/dp.wso2.com/v1alpha3/namespaces/" + namespace + "/airatelimitpolicies"; + return k8sApiServerEp->post(endpoint, rateLimitPolicy, targetType = http:Response); +} + isolated function updateRateLimitPolicyCR(model:RateLimitPolicy rateLimitPolicy, string namespace) returns http:Response|http:ClientError { string endpoint = "/apis/dp.wso2.com/v1alpha1/namespaces/" + namespace + "/ratelimitpolicies/" + rateLimitPolicy.metadata.name; return k8sApiServerEp->put(endpoint, rateLimitPolicy, targetType = http:Response); } +isolated function updateAIRateLimitPolicyCR(model:AIRateLimitPolicy rateLimitPolicy, string namespace) returns http:Response|http:ClientError { + string endpoint = "/apis/dp.wso2.com/v1alpha3/namespaces/" + namespace + "/airatelimitpolicies/" + rateLimitPolicy.metadata.name; + return k8sApiServerEp->put(endpoint, rateLimitPolicy, targetType = http:Response); +} + isolated function getRateLimitPolicyCR(string name, string namespace) returns model:RateLimitPolicy|http:ClientError { string endpoint = "/apis/dp.wso2.com/v1alpha1/namespaces/" + namespace + "/ratelimitpolicies/" + name; return k8sApiServerEp->get(endpoint, targetType = model:RateLimitPolicy); } +isolated function getAIRateLimitPolicyCR(string name, string namespace) returns model:AIRateLimitPolicy|http:ClientError { + string endpoint = "/apis/dp.wso2.com/v1alpha3/namespaces/" + namespace + "/airatelimitpolicies/" + name; + return k8sApiServerEp->get(endpoint, targetType = model:AIRateLimitPolicy); +} + isolated function deleteRateLimitPolicyCR(string name, string namespace) returns http:Response|http:ClientError { string endpoint = "/apis/dp.wso2.com/v1alpha1/namespaces/" + namespace + "/ratelimitpolicies/" + name; return k8sApiServerEp->delete(endpoint, targetType = http:Response); } +isolated function deleteAIRateLimitPolicyCR(string name, string namespace) returns http:Response|http:ClientError { + string endpoint = "/apis/dp.wso2.com/v1alpha3/namespaces/" + namespace + "/airatelimitpolicies/" + name; + return k8sApiServerEp->delete(endpoint, targetType = http:Response); +} + isolated function getRateLimitPolicyCRsForAPI(string apiName, string apiVersion, string namespace, string organization) returns model:RateLimitPolicyList|http:ClientError|error { string endpoint = "/apis/dp.wso2.com/v1alpha1/namespaces/" + namespace + "/ratelimitpolicies?labelSelector=" + check generateUrlEncodedLabelSelector(apiName, apiVersion, organization); return k8sApiServerEp->get(endpoint, targetType = model:RateLimitPolicyList); } +isolated function getAIRateLimitPolicyCRsForAPI(string apiName, string apiVersion, string namespace, string organization) returns model:AIRateLimitPolicyList|http:ClientError|error { + string endpoint = "/apis/dp.wso2.com/v1alpha3/namespaces/" + namespace + "/airatelimitpolicies?labelSelector=" + check generateUrlEncodedLabelSelector(apiName, apiVersion, organization); + return k8sApiServerEp->get(endpoint, targetType = model:AIRateLimitPolicyList); +} + isolated function deployAPIPolicyCR(model:APIPolicy apiPolicy, string namespace) returns http:Response|http:ClientError { string endpoint = "/apis/dp.wso2.com/v1alpha3/namespaces/" + namespace + "/apipolicies"; return k8sApiServerEp->post(endpoint, apiPolicy, targetType = http:Response); diff --git a/runtime/config-deployer-service/ballerina/modules/model/AIRatelimitPolicy.bal b/runtime/config-deployer-service/ballerina/modules/model/AIRatelimitPolicy.bal new file mode 100644 index 000000000..6c7a0f8f4 --- /dev/null +++ b/runtime/config-deployer-service/ballerina/modules/model/AIRatelimitPolicy.bal @@ -0,0 +1,47 @@ +// +// Copyright (c) 2024, WSO2 LLC. (http://www.wso2.com). +// +// WSO2 LLC. licenses this file to you under the Apache License, +// Version 2.0 (the "License"); you may not use this file except +// in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. +// +public type AIRateLimitPolicy record { + string apiVersion = "dp.wso2.com/v1alpha3"; + string kind = "AIRateLimitPolicy"; + Metadata metadata; + AIRateLimitPolicySpec spec; +}; + +public type AIRateLimitPolicySpec record {| + AIRateLimitPolicyData override?; + AIRateLimitPolicyData default?; + TargetRef targetRef; +|}; + +public type AIRateLimitPolicyData record { + string organization; + TokenAIRL tokenCount; + RequestAIRL requestCount; +}; + +public type TokenAIRL record { + string unit; + int requestTokenCount; + int responseTokenCount; + int totalTokenCount; +}; + +public type RequestAIRL record { + string unit; + int requestsPerUnit; +}; diff --git a/runtime/config-deployer-service/ballerina/modules/model/APIArtifact.bal b/runtime/config-deployer-service/ballerina/modules/model/APIArtifact.bal index 7f272176a..82c2e1727 100644 --- a/runtime/config-deployer-service/ballerina/modules/model/APIArtifact.bal +++ b/runtime/config-deployer-service/ballerina/modules/model/APIArtifact.bal @@ -13,6 +13,7 @@ public type APIArtifact record {| map authenticationMap = {}; map scopes = {}; map rateLimitPolicies = {}; + map aiRatelimitPolicies = {}; map apiPolicies = {}; map interceptorServices = {}; boolean sandboxEndpointAvailable = false; diff --git a/runtime/config-deployer-service/ballerina/modules/model/RateLimit.bal b/runtime/config-deployer-service/ballerina/modules/model/RateLimit.bal index a24cffbf2..8cf969c21 100644 --- a/runtime/config-deployer-service/ballerina/modules/model/RateLimit.bal +++ b/runtime/config-deployer-service/ballerina/modules/model/RateLimit.bal @@ -43,3 +43,10 @@ public type RateLimitPolicyList record { ListMeta metadata; RateLimitPolicy[] items; }; + +public type AIRateLimitPolicyList record { + string apiVersion = "dp.wso2.com/v1alpha3"; + string kind = "AIRateLimitPolicyList"; + ListMeta metadata; + AIRateLimitPolicy[] items; +}; diff --git a/runtime/config-deployer-service/ballerina/resources/apk-conf-schema.yaml b/runtime/config-deployer-service/ballerina/resources/apk-conf-schema.yaml index afe8fc75d..248af0f7c 100644 --- a/runtime/config-deployer-service/ballerina/resources/apk-conf-schema.yaml +++ b/runtime/config-deployer-service/ballerina/resources/apk-conf-schema.yaml @@ -360,6 +360,8 @@ components: $ref: "#/components/schemas/Certificate" resiliency: $ref: "#/components/schemas/Resiliency" + aiRatelimit: + $ref: "#/components/schemas/AIRatelimit" additionalProperties: false Certificate: type: object @@ -599,3 +601,57 @@ components: type: string default: string additionalProperties: false + AIRatelimit: + type: object + required: + - enabled + - token + - request + properties: + enabled: + type: boolean + default: true + token: + $ref: "#/components/schemas/TokenAIRL" + request: + $ref: "#/components/schemas/RequestAIRL" + TokenAIRL: + type: object + required: + - promptLimit + - completionLimit + - totalLimit + - unit + properties: + promptLimit: + type: integer + default: 0 + completionLimit: + type: integer + default: 0 + totalLimit: + type: integer + default": 0 + unit: + type: string + default: Minute + enum: + - Minute + - Hour + - Day + RequestAIRL: + type: object + required: + - requestLimit + - unit + properties: + requestLimit: + type: integer + default: 0 + unit: + type: string + default: Minute + enum: + - Minute + - Hour + - Day diff --git a/runtime/config-deployer-service/ballerina/types.bal b/runtime/config-deployer-service/ballerina/types.bal index dee63aa89..ed2abbccd 100644 --- a/runtime/config-deployer-service/ballerina/types.bal +++ b/runtime/config-deployer-service/ballerina/types.bal @@ -234,6 +234,38 @@ public type Resiliency record { RetryPolicy retryPolicy?; }; +# Configuration of AIRatelimit settings. +# +# + token - Configuration for the CircuitBreaker. +# + request - Configuration for the Timeout. +public type AIRatelimit record { + boolean enabled; + TokenAIRL token; + RequestAIRL request; +}; + +# Configuration for Token AI rate limit settings. +# +# + promptLimit - Limit for prompts within the specified unit. +# + completionLimit - Limit for completions within the specified unit. +# + totalLimit - Total limit combining prompt and completion counts. +# + unit - The time unit for the rate limits (Minute, Hour, Day). +public type TokenAIRL record { + int promptLimit; + int completionLimit; + int totalLimit; + string unit; +}; + +# Configuration for Request AI rate limit settings. +# +# + requestLimit - Limit for requests within the specified unit. +# + unit - The time unit for the request limits (Minute, Hour, Day). +public type RequestAIRL record { + int requestLimit; + string unit; +}; + # Configuration of CircuitBreaker settings. # # + maxConnectionPools - The maximum number of connection pools allowed. @@ -267,11 +299,13 @@ public type EndpointConfigurations record { # + endpointSecurity - The security configuration for the endpoint. # + certificate - The certificate configuration for the endpoint. # + resiliency - The resiliency configuration for the endpoint. +# + AIRatelimit - The AIRatelimit configuration for the AI ratelimit. public type EndpointConfiguration record { string|K8sService endpoint; EndpointSecurity endpointSecurity?; Certificate certificate?; Resiliency resiliency?; + AIRatelimit aiRatelimit?; }; # Configuration of OAuth2 Authentication type. diff --git a/runtime/config-deployer-service/docker/config-deployer/conf/apk-schema.json b/runtime/config-deployer-service/docker/config-deployer/conf/apk-schema.json index 19f723ee2..9ab31107d 100644 --- a/runtime/config-deployer-service/docker/config-deployer/conf/apk-schema.json +++ b/runtime/config-deployer-service/docker/config-deployer/conf/apk-schema.json @@ -532,6 +532,10 @@ "resiliency": { "$ref": "#/schemas/Resiliency", "description": "Resiliency configuration for the API endpoint." + }, + "aiRatelimit": { + "$ref": "#/schemas/AIRatelimit", + "description": "AI ratelimit configuration for the API endpoint." } }, "additionalProperties": false @@ -632,6 +636,93 @@ }, "additionalProperties": false }, + "AIRatelimit": { + "type": "object", + "required": [ + "enabled", + "token", + "request" + ], + "description": "Endpoint AI ratelimit related configurations of the API", + "properties": { + "enabled" : { + "type" : "boolean", + "default": true, + "description": "States whether the AI ratelimit is turned on or not" + }, + "token": { + "$ref": "#/schemas/TokenAIRL" + }, + "request": { + "$ref": "#/schemas/RequestAIRL" + } + }, + "additionalProperties": false + }, + "TokenAIRL": { + "type": "object", + "required": [ + "promptLimit", + "completionLimit", + "totalLimit", + "unit" + ], + "description": "Token limits configuration for AI rate limiting", + "properties": { + "promptLimit": { + "type": "integer", + "default": 0, + "description": "Limit for prompts within the specified unit" + }, + "completionLimit": { + "type": "integer", + "default": 0, + "description": "Limit for completions within the specified unit" + }, + "totalLimit": { + "type": "integer", + "default": 0, + "description": "Total limit combining prompt and completion counts" + }, + "unit": { + "type": "string", + "default": "Minute", + "enum": [ + "Minute", + "Hour", + "Day" + ], + "description": "The time unit for the rate limits" + } + }, + "additionalProperties": false + }, + "RequestAIRL": { + "type": "object", + "required": [ + "requestLimit", + "unit" + ], + "description": "Request limits configuration for AI rate limiting", + "properties": { + "requestLimit": { + "type": "integer", + "default": 0, + "description": "Limit for requests within the specified unit" + }, + "unit": { + "type": "string", + "default": "Minute", + "enum": [ + "Minute", + "Hour", + "Day" + ], + "description": "The time unit for the request limits" + } + }, + "additionalProperties": false + }, "CircuitBreaker": { "type": "object", "properties": {