adding APK Conf related implementation for AI Provider and Backend AP…

…I Key Security
wso2 · Sep 12, 2024 · d5b4fb6 · d5b4fb6
1 parent b84cd52
commit d5b4fb6
Show file tree

Hide file tree

Showing 10 changed files with 210 additions and 21 deletions.
diff --git a/adapter/internal/operator/controllers/dp/api_controller.go b/adapter/internal/operator/controllers/dp/api_controller.go
@@ -1086,10 +1086,6 @@ func (apiReconciler *APIReconciler) traverseAPIStateAndUpdateOwnerReferences(ctx
 	for _, backendJwt := range apiState.BackendJWTMapping {
 		apiReconciler.retriveParentAPIsAndUpdateOwnerReferene(ctx, &backendJwt)
 	}
-	if apiState.AIProvider != nil && apiState.AIProvider.Name != "" {
-		loggers.LoggerAPKOperator.Infof("Owner Reference .AIProvider override found in API State. AI Provider Name %s", apiState.AIProvider.Name)
-		apiReconciler.retriveParentAPIsAndUpdateOwnerReferene(ctx, apiState.AIProvider)
-	}
 
 }
 

diff --git a/helm-charts/templates/data-plane/gateway-components/AIProviders/ai-provider-azure-ai.yaml b/helm-charts/templates/data-plane/gateway-components/AIProviders/ai-provider-azure-ai.yaml
@@ -0,0 +1,23 @@
+apiVersion: dp.wso2.com/v1alpha3
+kind: AIProvider
+metadata:
+  name: ai-provider-azure-ai
+  namespace: apk
+spec:
+  providerName : "AzureAI"
+  providerAPIVersion : "2024-06-01"
+  organization : "default"
+  model:
+    in: "Body"
+    value: "model"
+  rateLimitFields:
+    promptTokens: 
+      in: "Body"
+      value: "$usage.promptTokens"
+    completionToken:
+      in: "Body"
+      value: "$usage.completionTokens"
+    totalToken:
+      in: "Body"
+      value: "$usage.totalTokens"
+
diff --git a/helm-charts/templates/data-plane/gateway-components/AIProviders/ai-provider-open-ai.yaml b/helm-charts/templates/data-plane/gateway-components/AIProviders/ai-provider-open-ai.yaml
@@ -0,0 +1,23 @@
+apiVersion: dp.wso2.com/v1alpha3
+kind: AIProvider
+metadata:
+  name: ai-provider-open-ai
+  namespace: apk
+spec:
+  providerName : "OpenAI"
+  providerAPIVersion : "v1"
+  organization : "default"
+  model:
+    in: "Body"
+    value: "model"
+  rateLimitFields:
+    promptTokens: 
+      in: "Body"
+      value: "$usage.promptTokens"
+    completionToken:
+      in: "Body"
+      value: "$usage.completionTokens"
+    totalToken:
+      in: "Body"
+      value: "$usage.totalTokens"
+
diff --git a/runtime/config-deployer-service/ballerina/APIClient.bal b/runtime/config-deployer-service/ballerina/APIClient.bal
@@ -336,7 +336,7 @@ public class APIClient {
                 apiArtifact.rateLimitPolicies[rateLimitPolicyCR.metadata.name] = rateLimitPolicyCR;
             }
         }
-        if apkConf.apiPolicies != () || apkConf.corsConfiguration != () || apkConf.subscriptionValidation == true {
+        if apkConf.apiPolicies != () || apkConf.corsConfiguration != () || apkConf.aiProvider != () || apkConf.subscriptionValidation == true {
             model:APIPolicy? apiPolicyCR = check self.generateAPIPolicyAndBackendCR(apiArtifact, apkConf, (), apkConf.apiPolicies, organization, apiArtifact.uniqueId);
             if apiPolicyCR != () {
                 apiArtifact.apiPolicies[apiPolicyCR.metadata.name] = apiPolicyCR;
@@ -721,6 +721,21 @@ public class APIClient {
                 defaultSpecData.cORSPolicy = cORSPolicy;
             }
         }
+        AIProvider? aiProvider = apkConf.aiProvider;
+        if aiProvider is AIProvider {
+            string aiProviderRef;
+            if aiProvider.name == "AzureAI" {
+                aiProviderRef = "ai-provider-azure-ai";
+            } else if aiProvider.name == "OpenAI" {
+                aiProviderRef = "ai-provider-open-ai";
+            } else if aiProvider.name == "GoogleAI" {
+                aiProviderRef = "ai-provider-google-ai";
+            } else {
+                aiProviderRef = aiProvider.name;
+            }
+            model:AIProviderReference aIProviderPolicy = {name: aiProviderRef};
+            defaultSpecData.aiProvider = aIProviderPolicy;
+        }
         if defaultSpecData != {} {
             model:APIPolicy? apiPolicyCR = self.generateAPIPolicyCR(apkConf, targetRefName, operations, organization, defaultSpecData);
             if apiPolicyCR != () {
@@ -1278,8 +1293,10 @@ public class APIClient {
         }
         if endpointSecurity is EndpointSecurity {
             if endpointSecurity?.enabled ?: false {
-                // When user adds basic auth endpoint security username and password
-                BasicEndpointSecurity? securityType = endpointSecurity.securityType;
+                // When user adds basic auth endpoint security username and password should be provided.
+                // When user adds api key endpoint security api key name and api key value should be provided.
+                BasicEndpointSecurity|APIKeyEndpointSecurity? securityType = endpointSecurity.securityType;
+                log:printDebug("Security Type: "+ securityType.toString());
 
                 if securityType is BasicEndpointSecurity {
                     securityConfig = {
@@ -1292,6 +1309,18 @@ public class APIClient {
                         }
                     };
                 }
+                if securityType is APIKeyEndpointSecurity {
+                    securityConfig = {
+                        apiKey: {
+                            name: <string>securityType.apiKeyNameKey,
+                            'in: <string>securityType.'in,
+                            valueFrom: {
+                                name: <string>securityType.secretName, 
+                                valueKey: <string>securityType.apiKeyValueKey
+                            }
+                        }
+                    };
+                }
             }
             backendService.spec.security = securityConfig;
         }

diff --git a/runtime/config-deployer-service/ballerina/Dependencies.toml b/runtime/config-deployer-service/ballerina/Dependencies.toml
@@ -44,7 +44,7 @@ modules = [
 [[package]]
 org = "ballerina"
 name = "crypto"
-version = "2.6.2"
+version = "2.6.3"
 dependencies = [
 	{org = "ballerina", name = "jballerina.java"},
 	{org = "ballerina", name = "time"}
@@ -70,7 +70,7 @@ modules = [
 [[package]]
 org = "ballerina"
 name = "http"
-version = "2.10.12"
+version = "2.10.15"
 dependencies = [
 	{org = "ballerina", name = "auth"},
 	{org = "ballerina", name = "cache"},
@@ -102,7 +102,7 @@ modules = [
 [[package]]
 org = "ballerina"
 name = "io"
-version = "1.6.0"
+version = "1.6.2"
 dependencies = [
 	{org = "ballerina", name = "jballerina.java"},
 	{org = "ballerina", name = "lang.value"}
@@ -273,7 +273,7 @@ dependencies = [
 [[package]]
 org = "ballerina"
 name = "observe"
-version = "1.2.2"
+version = "1.2.3"
 dependencies = [
 	{org = "ballerina", name = "jballerina.java"}
 ]
@@ -302,7 +302,7 @@ modules = [
 [[package]]
 org = "ballerina"
 name = "sql"
-version = "1.11.1"
+version = "1.11.2"
 dependencies = [
 	{org = "ballerina", name = "io"},
 	{org = "ballerina", name = "jballerina.java"},
@@ -335,15 +335,15 @@ modules = [
 [[package]]
 org = "ballerina"
 name = "time"
-version = "2.4.0"
+version = "2.4.1"
 dependencies = [
 	{org = "ballerina", name = "jballerina.java"}
 ]
 
 [[package]]
 org = "ballerina"
 name = "url"
-version = "2.4.0"
+version = "2.4.1"
 dependencies = [
 	{org = "ballerina", name = "jballerina.java"}
 ]
@@ -424,8 +424,7 @@ dependencies = [
 modules = [
 	{org = "wso2", packageName = "apk_common_lib", moduleName = "apk_common_lib"},
 	{org = "wso2", packageName = "apk_common_lib", moduleName = "apk_common_lib.java.io"},
-	{org = "wso2", packageName = "apk_common_lib", moduleName = "apk_common_lib.java.lang"},
-	{org = "wso2", packageName = "apk_common_lib", moduleName = "apk_common_lib.org.wso2.apk.common"}
+	{org = "wso2", packageName = "apk_common_lib", moduleName = "apk_common_lib.java.lang"}
 ]
 
 [[package]]

diff --git a/runtime/config-deployer-service/ballerina/modules/model/APIPolicy.bal b/runtime/config-deployer-service/ballerina/modules/model/APIPolicy.bal
@@ -34,6 +34,7 @@ public type APIPolicyData record {
     CORSPolicy cORSPolicy?;
     BackendJwtReference backendJwtPolicy?;
     boolean subscriptionValidation?;
+    AIProviderReference aiProvider?;
 };
 
 public type InterceptorReference record {
@@ -44,6 +45,10 @@ public type BackendJwtReference record {
     string name?;
 };
 
+public type AIProviderReference record {
+    string name?;
+};
+
 public type APIPolicyList record {
     string apiVersion = "dp.wso2.com/v1alpha3";
     string kind = "APIPolicyList";

diff --git a/runtime/config-deployer-service/ballerina/modules/model/Backend.bal b/runtime/config-deployer-service/ballerina/modules/model/Backend.bal
@@ -44,6 +44,18 @@ public type BasicSecurityConfig record {
 
 public type SecurityConfig record {
     BasicSecurityConfig basic?;
+    APIKeySecurityConfig apiKey?;
+};
+
+public type APIKeySecurityConfig record {
+    string name;
+    string 'in;
+    ValueFrom valueFrom;
+};
+
+public type ValueFrom record {
+    string name;
+    string valueKey;
 };
 
 public type SecretRefConfig record {
@@ -52,7 +64,6 @@ public type SecretRefConfig record {
    string passwordKey = "password"; 
 };
 
-
 public type RefConfig record {
     string key;
     string name;

diff --git a/runtime/config-deployer-service/ballerina/resources/apk-conf-schema.yaml b/runtime/config-deployer-service/ballerina/resources/apk-conf-schema.yaml
@@ -55,6 +55,8 @@ components:
             Environment of the API
         endpointConfigurations:
           $ref: "#/components/schemas/EndpointConfigurations"
+        aiProvider:
+          $ref: "#/components/schemas/AIProvider"
         operations:
           type: array
           items:
@@ -325,6 +327,14 @@ components:
             - Day
           example: Minute
       additionalProperties: false
+    AIProvider:
+      title: AI Provider
+      properties:
+        name:
+          type: string
+        apiVersion:
+          type: string
+      additionalProperties: false
     EndpointConfigurations:
       title: Endpoint configs
       properties:
@@ -368,6 +378,7 @@ components:
         securityType:
           oneOf:
             - $ref: "#/components/schemas/BasicEndpointSecurity"
+            - $ref: "#/components/schemas/APIKeyEndpointSecurity"
       additionalProperties: false
     BasicEndpointSecurity:
       type: object
@@ -379,6 +390,22 @@ components:
         passwordKey:
           type: string
       additionalProperties: false
+    APIKeyEndpointSecurity:
+      type: object
+      properties:
+        secretName:
+          type: string
+        in:
+          type: string
+          enum:
+            - Header
+            - Query
+          default: Header
+        apiKeyNameKey:
+          type: string
+        apiKeyValueKey:
+          type: string
+      additionalProperties: false
     Resiliency:
       type: object
       properties:

diff --git a/runtime/config-deployer-service/ballerina/types.bal b/runtime/config-deployer-service/ballerina/types.bal
@@ -60,7 +60,7 @@ public type ErrorListItem record {
 # + securityType - Configuration for the basic endpoint security.
 public type EndpointSecurity record {
     boolean enabled?;
-    BasicEndpointSecurity securityType?;
+    BasicEndpointSecurity|APIKeyEndpointSecurity securityType?;
 };
 
 # Configuration for Custom Claims.
@@ -102,9 +102,22 @@ public type RateLimit record {
 # + userNameKey - The key to retrieve the username from the secret.
 # + passwordKey - The key to retrieve the password from the secret.
 public type BasicEndpointSecurity record {
-    string secretName?;
-    string userNameKey?;
-    string passwordKey?;
+    string secretName;
+    string userNameKey;
+    string passwordKey;
+};
+
+# Configuration for API Key Endpoint Security.
+# 
+# + secretName - The name of the secret containing the API key.
+# + in - The mode of sending Ex: header/query.
+# + apiKeyNameKey - API key Header name.
+# + apiKeyValueKey - Key to retrieve API key Header value from the secret.
+public type APIKeyEndpointSecurity record {
+    string secretName;
+    string 'in;
+    string apiKeyNameKey;
+    string apiKeyValueKey;
 };
 
 # Configuration for APK Operations.
@@ -407,6 +420,7 @@ public type InterceptorPolicy record {
 # + subscriptionValidation - Is subscription validation enabled for the API.
 # + environment - Environment of the API.
 # + endpointConfigurations - Sandbox and production endpoint configurations of the API
+# + aiProvider - AI provider configuration for the API.
 # + operations - List of operations for this API.
 # + apiPolicies - Policies like interceptor to be added to the entire API.
 # + rateLimit - Rate limiting configuration for the API.
@@ -427,6 +441,7 @@ public type APKConf record {
     boolean subscriptionValidation = false;
     string environment?;
     EndpointConfigurations endpointConfigurations?;
+    AIProvider aiProvider?;
     APKOperations[] operations?;
     APIOperationPolicies apiPolicies?;
     RateLimit rateLimit?;
@@ -435,6 +450,15 @@ public type APKConf record {
     CORSConfiguration corsConfiguration?;
 };
 
+# Configuration for an AI provider.
+# 
+# + name - The name of the AI provider.
+# + apiVersion - The version of the AI provider.
+public type AIProvider record {
+    string name;
+    string apiVersion;
+};
+
 # Configuration for Interceptor Policy parameters.
 #
 # + backendUrl - Backend URL of the interceptor service.