Skip to content

Commit 85449a1

Browse files
authored
Merge pull request #2478 from Tharsanan1/airl-apk-conf
Add apk conf for airl
2 parents 95c0ffb + bcacdcd commit 85449a1

File tree

14 files changed

+481
-11
lines changed

14 files changed

+481
-11
lines changed

common-controller/internal/operator/controllers/dp/airatelimitpolicy_controller.go

+3
Original file line numberDiff line numberDiff line change
@@ -105,6 +105,9 @@ func (r *AIRateLimitPolicyReconciler) Reconcile(ctx context.Context, req ctrl.Re
105105
xds.UpdateRateLimiterPolicies(conf.CommonController.Server.Label)
106106
} else {
107107
loggers.LoggerAPKOperator.Infof("ratelimits found")
108+
if ratelimitPolicy.Spec.Override == nil {
109+
ratelimitPolicy.Spec.Override = ratelimitPolicy.Spec.Default
110+
}
108111
if ratelimitPolicy.Spec.TargetRef.Name != "" {
109112
r.ods.AddorUpdateAIRatelimitToStore(ratelimitKey, ratelimitPolicy.Spec)
110113
xds.UpdateRateLimitXDSCacheForAIRatelimitPolicies(r.ods.GetAIRatelimitPolicySpecs())

runtime/config-deployer-service/ballerina/APIClient.bal

+37
Original file line numberDiff line numberDiff line change
@@ -232,6 +232,11 @@ public class APIClient {
232232
serviceEntry: false,
233233
url: self.constructURlFromService(sandboxEndpointConfig.endpoint)
234234
};
235+
AIRatelimit? aiRatelimit = sandboxEndpointConfig.aiRatelimit;
236+
if aiRatelimit is AIRatelimit && aiRatelimit.enabled {
237+
model:AIRateLimitPolicy airl = self.generateAIRateLimitPolicyCR(apkConf, aiRatelimit.token, aiRatelimit.request, backendService.metadata.name, organization);
238+
apiArtifact.aiRatelimitPolicies[airl.metadata.name] = airl;
239+
}
235240
}
236241
}
237242
if (endpointType == () || endpointType == PRODUCTION_TYPE) {
@@ -246,6 +251,11 @@ public class APIClient {
246251
serviceEntry: false,
247252
url: self.constructURlFromService(productionEndpointConfig.endpoint)
248253
};
254+
AIRatelimit? aiRatelimit = productionEndpointConfig.aiRatelimit;
255+
if aiRatelimit is AIRatelimit && aiRatelimit.enabled {
256+
model:AIRateLimitPolicy airl = self.generateAIRateLimitPolicyCR(apkConf, aiRatelimit.token, aiRatelimit.request, backendService.metadata.name, organization);
257+
apiArtifact.aiRatelimitPolicies[airl.metadata.name] = airl;
258+
}
249259
}
250260
}
251261
return endpointIdMap;
@@ -1506,6 +1516,29 @@ public class APIClient {
15061516
return rateLimitPolicyCR;
15071517
}
15081518

1519+
public isolated function generateAIRateLimitPolicyCR(APKConf apkConf, TokenAIRL tokenAIRL, RequestAIRL requestAIRL, string targetRefName, commons:Organization organization) returns model:AIRateLimitPolicy {
1520+
string apiIdentifierHash = crypto:hashSha1((apkConf.name + apkConf.version).toBytes()).toBase16();
1521+
model:AIRateLimitPolicy aiRateLimitPolicyCR = {
1522+
metadata: {
1523+
name: self.retrieveAIRateLimitPolicyName(apiIdentifierHash, targetRefName),
1524+
labels: self.getLabels(apkConf, organization)
1525+
},
1526+
spec: {
1527+
default: {
1528+
organization: organization.name,
1529+
requestCount: {unit: requestAIRL.unit, requestsPerUnit: requestAIRL.requestLimit},
1530+
tokenCount: {unit: tokenAIRL.unit, requestTokenCount: tokenAIRL.promptLimit, responseTokenCount: tokenAIRL.completionLimit, totalTokenCount: tokenAIRL.totalLimit}
1531+
},
1532+
targetRef: {
1533+
group: "dp.wso2.com",
1534+
kind: "Backend",
1535+
name: targetRefName
1536+
}
1537+
}
1538+
};
1539+
return aiRateLimitPolicyCR;
1540+
}
1541+
15091542
isolated function retrieveRateLimitData(RateLimit rateLimit, commons:Organization organization) returns model:RateLimitData {
15101543
model:RateLimitData rateLimitData = {
15111544
api: {
@@ -1933,6 +1966,10 @@ public class APIClient {
19331966
}
19341967
}
19351968

1969+
public isolated function retrieveAIRateLimitPolicyName(string apiID, string targetRef) returns string {
1970+
return "airl-" + apiID + "-" + targetRef;
1971+
}
1972+
19361973
private isolated function validateAndRetrieveAPKConfiguration(json apkconfJson) returns APKConf|commons:APKError? {
19371974
do {
19381975
runtimeapi:APKConfValidationResponse validationResponse = check apkConfValidator.validate(apkconfJson.toJsonString());

runtime/config-deployer-service/ballerina/ConfigGenreatorClient.bal

+4
Original file line numberDiff line numberDiff line change
@@ -234,6 +234,10 @@ on fail var e {
234234
string yamlString = check self.convertJsonToYaml(rateLimitPolicy.toJsonString());
235235
_ = check self.storeFile(yamlString, rateLimitPolicy.metadata.name, zipDir);
236236
}
237+
foreach model:AIRateLimitPolicy airateLimitPolicy in apiArtifact.aiRatelimitPolicies {
238+
string yamlString = check self.convertJsonToYaml(airateLimitPolicy.toJsonString());
239+
_ = check self.storeFile(yamlString, airateLimitPolicy.metadata.name, zipDir);
240+
}
237241
foreach model:APIPolicy apiPolicy in apiArtifact.apiPolicies {
238242
string yamlString = check self.convertJsonToYaml(apiPolicy.toJsonString());
239243
_ = check self.storeFile(yamlString, apiPolicy.metadata.name, zipDir);

runtime/config-deployer-service/ballerina/DeployerClient.bal

+49
Original file line numberDiff line numberDiff line change
@@ -102,6 +102,7 @@ public class DeployerClient {
102102
_ = check self.deleteScopeCrsForAPI(existingAPI, <string>apiArtifact?.organization);
103103
check self.deleteBackends(existingAPI, <string>apiArtifact?.organization);
104104
check self.deleteRateLimitPolicyCRs(existingAPI, <string>apiArtifact?.organization);
105+
check self.deleteAIRateLimitPolicyCRs(existingAPI, <string>apiArtifact?.organization);
105106
check self.deleteAPIPolicyCRs(existingAPI, <string>apiArtifact?.organization);
106107
check self.deleteInterceptorServiceCRs(existingAPI, <string>apiArtifact?.organization);
107108
check self.deleteBackendJWTConfig(existingAPI, <string>apiArtifact?.organization);
@@ -121,6 +122,7 @@ public class DeployerClient {
121122
check self.deployBackendServices(apiArtifact, ownerReference);
122123
check self.deployAuthenticationCRs(apiArtifact, ownerReference);
123124
check self.deployRateLimitPolicyCRs(apiArtifact, ownerReference);
125+
check self.deployAIRateLimitPolicyCRs(apiArtifact, ownerReference);
124126
check self.deployInterceptorServiceCRs(apiArtifact, ownerReference);
125127
check self.deployBackendJWTConfigs(apiArtifact, ownerReference);
126128
check self.deployAPIPolicyCRs(apiArtifact, ownerReference);
@@ -660,6 +662,30 @@ public class DeployerClient {
660662
}
661663
}
662664

665+
private isolated function deployAIRateLimitPolicyCRs(model:APIArtifact apiArtifact, model:OwnerReference ownerReference) returns error? {
666+
foreach model:AIRateLimitPolicy rateLimitPolicy in apiArtifact.aiRatelimitPolicies {
667+
rateLimitPolicy.metadata.ownerReferences = [ownerReference];
668+
http:Response deployRateLimitPolicyResult = check deployAIRateLimitPolicyCR(rateLimitPolicy, <string>apiArtifact?.namespace);
669+
if deployRateLimitPolicyResult.statusCode == http:STATUS_CREATED {
670+
log:printDebug("Deployed AIRateLimitPolicy Successfully" + rateLimitPolicy.toString());
671+
} else if deployRateLimitPolicyResult.statusCode == http:STATUS_CONFLICT {
672+
log:printDebug("AIRateLimitPolicy already exists" + rateLimitPolicy.toString());
673+
model:AIRateLimitPolicy rateLimitPolicyFromK8s = check getAIRateLimitPolicyCR(rateLimitPolicy.metadata.name, <string>apiArtifact?.namespace);
674+
rateLimitPolicy.metadata.resourceVersion = rateLimitPolicyFromK8s.metadata.resourceVersion;
675+
http:Response rateLimitPolicyCR = check updateAIRateLimitPolicyCR(rateLimitPolicy, <string>apiArtifact?.namespace);
676+
if rateLimitPolicyCR.statusCode != http:STATUS_OK {
677+
json responsePayLoad = check rateLimitPolicyCR.getJsonPayload();
678+
model:Status statusResponse = check responsePayLoad.cloneWithType(model:Status);
679+
check self.handleK8sTimeout(statusResponse);
680+
}
681+
} else {
682+
json responsePayLoad = check deployRateLimitPolicyResult.getJsonPayload();
683+
model:Status statusResponse = check responsePayLoad.cloneWithType(model:Status);
684+
check self.handleK8sTimeout(statusResponse);
685+
}
686+
}
687+
}
688+
663689
private isolated function deleteRateLimitPolicyCRs(model:API api, string organization) returns commons:APKError? {
664690
do {
665691
model:RateLimitPolicyList|http:ClientError rateLimitPolicyCrListResponse = check getRateLimitPolicyCRsForAPI(api.spec.apiName, api.spec.apiVersion, <string>api.metadata?.namespace, organization);
@@ -684,6 +710,29 @@ public class DeployerClient {
684710
}
685711
}
686712

713+
private isolated function deleteAIRateLimitPolicyCRs(model:API api, string organization) returns commons:APKError? {
714+
do {
715+
model:AIRateLimitPolicyList|http:ClientError aiRateLimitPolicyCrListResponse = check getAIRateLimitPolicyCRsForAPI(api.spec.apiName, api.spec.apiVersion, <string>api.metadata?.namespace, organization);
716+
if aiRateLimitPolicyCrListResponse is model:AIRateLimitPolicyList {
717+
foreach model:AIRateLimitPolicy item in aiRateLimitPolicyCrListResponse.items {
718+
http:Response|http:ClientError rateLimitPolicyCRDeletionResponse = deleteAIRateLimitPolicyCR(item.metadata.name, <string>item.metadata?.namespace);
719+
if rateLimitPolicyCRDeletionResponse is http:Response {
720+
if rateLimitPolicyCRDeletionResponse.statusCode != http:STATUS_OK {
721+
json responsePayLoad = check rateLimitPolicyCRDeletionResponse.getJsonPayload();
722+
model:Status statusResponse = check responsePayLoad.cloneWithType(model:Status);
723+
check self.handleK8sTimeout(statusResponse);
724+
}
725+
} else {
726+
log:printError("Error occured while deleting AI rate limit policy");
727+
}
728+
}
729+
return;
730+
}
731+
} on fail var e {
732+
log:printError("Error occured deleting AI rate limit policy", e);
733+
return e909022("Error occured deleting AI rate limit policy", e);
734+
}
735+
}
687736
private isolated function deleteAPIPolicyCRs(model:API api, string organization) returns commons:APKError? {
688737
do {
689738
model:APIPolicyList|http:ClientError apiPolicyCrListResponse = check getAPIPolicyCRsForAPI(api.spec.apiName, api.spec.apiVersion, <string>api.metadata?.namespace, organization);

runtime/config-deployer-service/ballerina/K8sClient.bal

+25
Original file line numberDiff line numberDiff line change
@@ -251,26 +251,51 @@ isolated function deployRateLimitPolicyCR(model:RateLimitPolicy rateLimitPolicy,
251251
return k8sApiServerEp->post(endpoint, rateLimitPolicy, targetType = http:Response);
252252
}
253253

254+
isolated function deployAIRateLimitPolicyCR(model:AIRateLimitPolicy rateLimitPolicy, string namespace) returns http:Response|http:ClientError {
255+
string endpoint = "/apis/dp.wso2.com/v1alpha3/namespaces/" + namespace + "/airatelimitpolicies";
256+
return k8sApiServerEp->post(endpoint, rateLimitPolicy, targetType = http:Response);
257+
}
258+
254259
isolated function updateRateLimitPolicyCR(model:RateLimitPolicy rateLimitPolicy, string namespace) returns http:Response|http:ClientError {
255260
string endpoint = "/apis/dp.wso2.com/v1alpha1/namespaces/" + namespace + "/ratelimitpolicies/" + rateLimitPolicy.metadata.name;
256261
return k8sApiServerEp->put(endpoint, rateLimitPolicy, targetType = http:Response);
257262
}
258263

264+
isolated function updateAIRateLimitPolicyCR(model:AIRateLimitPolicy rateLimitPolicy, string namespace) returns http:Response|http:ClientError {
265+
string endpoint = "/apis/dp.wso2.com/v1alpha3/namespaces/" + namespace + "/airatelimitpolicies/" + rateLimitPolicy.metadata.name;
266+
return k8sApiServerEp->put(endpoint, rateLimitPolicy, targetType = http:Response);
267+
}
268+
259269
isolated function getRateLimitPolicyCR(string name, string namespace) returns model:RateLimitPolicy|http:ClientError {
260270
string endpoint = "/apis/dp.wso2.com/v1alpha1/namespaces/" + namespace + "/ratelimitpolicies/" + name;
261271
return k8sApiServerEp->get(endpoint, targetType = model:RateLimitPolicy);
262272
}
263273

274+
isolated function getAIRateLimitPolicyCR(string name, string namespace) returns model:AIRateLimitPolicy|http:ClientError {
275+
string endpoint = "/apis/dp.wso2.com/v1alpha3/namespaces/" + namespace + "/airatelimitpolicies/" + name;
276+
return k8sApiServerEp->get(endpoint, targetType = model:AIRateLimitPolicy);
277+
}
278+
264279
isolated function deleteRateLimitPolicyCR(string name, string namespace) returns http:Response|http:ClientError {
265280
string endpoint = "/apis/dp.wso2.com/v1alpha1/namespaces/" + namespace + "/ratelimitpolicies/" + name;
266281
return k8sApiServerEp->delete(endpoint, targetType = http:Response);
267282
}
268283

284+
isolated function deleteAIRateLimitPolicyCR(string name, string namespace) returns http:Response|http:ClientError {
285+
string endpoint = "/apis/dp.wso2.com/v1alpha3/namespaces/" + namespace + "/airatelimitpolicies/" + name;
286+
return k8sApiServerEp->delete(endpoint, targetType = http:Response);
287+
}
288+
269289
isolated function getRateLimitPolicyCRsForAPI(string apiName, string apiVersion, string namespace, string organization) returns model:RateLimitPolicyList|http:ClientError|error {
270290
string endpoint = "/apis/dp.wso2.com/v1alpha1/namespaces/" + namespace + "/ratelimitpolicies?labelSelector=" + check generateUrlEncodedLabelSelector(apiName, apiVersion, organization);
271291
return k8sApiServerEp->get(endpoint, targetType = model:RateLimitPolicyList);
272292
}
273293

294+
isolated function getAIRateLimitPolicyCRsForAPI(string apiName, string apiVersion, string namespace, string organization) returns model:AIRateLimitPolicyList|http:ClientError|error {
295+
string endpoint = "/apis/dp.wso2.com/v1alpha3/namespaces/" + namespace + "/airatelimitpolicies?labelSelector=" + check generateUrlEncodedLabelSelector(apiName, apiVersion, organization);
296+
return k8sApiServerEp->get(endpoint, targetType = model:AIRateLimitPolicyList);
297+
}
298+
274299
isolated function deployAPIPolicyCR(model:APIPolicy apiPolicy, string namespace) returns http:Response|http:ClientError {
275300
string endpoint = "/apis/dp.wso2.com/v1alpha3/namespaces/" + namespace + "/apipolicies";
276301
return k8sApiServerEp->post(endpoint, apiPolicy, targetType = http:Response);
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,47 @@
1+
//
2+
// Copyright (c) 2024, WSO2 LLC. (http://www.wso2.com).
3+
//
4+
// WSO2 LLC. licenses this file to you under the Apache License,
5+
// Version 2.0 (the "License"); you may not use this file except
6+
// in compliance with the License.
7+
// You may obtain a copy of the License at
8+
//
9+
// http://www.apache.org/licenses/LICENSE-2.0
10+
//
11+
// Unless required by applicable law or agreed to in writing,
12+
// software distributed under the License is distributed on an
13+
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14+
// KIND, either express or implied. See the License for the
15+
// specific language governing permissions and limitations
16+
// under the License.
17+
//
18+
public type AIRateLimitPolicy record {
19+
string apiVersion = "dp.wso2.com/v1alpha3";
20+
string kind = "AIRateLimitPolicy";
21+
Metadata metadata;
22+
AIRateLimitPolicySpec spec;
23+
};
24+
25+
public type AIRateLimitPolicySpec record {|
26+
AIRateLimitPolicyData override?;
27+
AIRateLimitPolicyData default?;
28+
TargetRef targetRef;
29+
|};
30+
31+
public type AIRateLimitPolicyData record {
32+
string organization;
33+
TokenAIRL tokenCount;
34+
RequestAIRL requestCount;
35+
};
36+
37+
public type TokenAIRL record {
38+
string unit;
39+
int requestTokenCount;
40+
int responseTokenCount;
41+
int totalTokenCount;
42+
};
43+
44+
public type RequestAIRL record {
45+
string unit;
46+
int requestsPerUnit;
47+
};

runtime/config-deployer-service/ballerina/modules/model/APIArtifact.bal

+1
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@ public type APIArtifact record {|
1515
map<Authentication> authenticationMap = {};
1616
map<Scope> scopes = {};
1717
map<RateLimitPolicy> rateLimitPolicies = {};
18+
map<AIRateLimitPolicy> aiRatelimitPolicies = {};
1819
map<APIPolicy> apiPolicies = {};
1920
map<InterceptorService> interceptorServices = {};
2021
boolean sandboxEndpointAvailable = false;

runtime/config-deployer-service/ballerina/modules/model/RateLimit.bal

+7
Original file line numberDiff line numberDiff line change
@@ -43,3 +43,10 @@ public type RateLimitPolicyList record {
4343
ListMeta metadata;
4444
RateLimitPolicy[] items;
4545
};
46+
47+
public type AIRateLimitPolicyList record {
48+
string apiVersion = "dp.wso2.com/v1alpha3";
49+
string kind = "AIRateLimitPolicyList";
50+
ListMeta metadata;
51+
AIRateLimitPolicy[] items;
52+
};

runtime/config-deployer-service/ballerina/resources/apk-conf-schema.yaml

+56
Original file line numberDiff line numberDiff line change
@@ -360,6 +360,8 @@ components:
360360
$ref: "#/components/schemas/Certificate"
361361
resiliency:
362362
$ref: "#/components/schemas/Resiliency"
363+
aiRatelimit:
364+
$ref: "#/components/schemas/AIRatelimit"
363365
additionalProperties: false
364366
Certificate:
365367
type: object
@@ -599,3 +601,57 @@ components:
599601
type: string
600602
default: string
601603
additionalProperties: false
604+
AIRatelimit:
605+
type: object
606+
required:
607+
- enabled
608+
- token
609+
- request
610+
properties:
611+
enabled:
612+
type: boolean
613+
default: true
614+
token:
615+
$ref: "#/components/schemas/TokenAIRL"
616+
request:
617+
$ref: "#/components/schemas/RequestAIRL"
618+
TokenAIRL:
619+
type: object
620+
required:
621+
- promptLimit
622+
- completionLimit
623+
- totalLimit
624+
- unit
625+
properties:
626+
promptLimit:
627+
type: integer
628+
default: 0
629+
completionLimit:
630+
type: integer
631+
default: 0
632+
totalLimit:
633+
type: integer
634+
default": 0
635+
unit:
636+
type: string
637+
default: Minute
638+
enum:
639+
- Minute
640+
- Hour
641+
- Day
642+
RequestAIRL:
643+
type: object
644+
required:
645+
- requestLimit
646+
- unit
647+
properties:
648+
requestLimit:
649+
type: integer
650+
default: 0
651+
unit:
652+
type: string
653+
default: Minute
654+
enum:
655+
- Minute
656+
- Hour
657+
- Day

0 commit comments

Comments
 (0)