Merge pull request #1194 from Tharsanan1/airl

Add dp to cp airl
wso2 · Sep 26, 2024 · 0b1d90f · 0b1d90f
2 parents 25e2be9 + 921b48d
commit 0b1d90f
Show file tree

Hide file tree

Showing 2 changed files with 61 additions and 0 deletions.
diff --git a/apim-apk-agent/pkg/managementserver/rest_server.go b/apim-apk-agent/pkg/managementserver/rest_server.go
@@ -216,6 +216,52 @@ func createAPIYaml(apiCPEvent *APICPEvent) (string, string) {
 			"accessControlExposeHeaders":    apiCPEvent.API.CORSPolicy.AccessControlExposeHeaders,
 		}
 	}
+	maxTps := make(map[string]interface{})
+
+	// Handle Production fields
+	if apiCPEvent.API.ProdAIRL != nil {
+		maxTps["production"] = apiCPEvent.API.ProdAIRL.RequestCount
+		maxTps["productionTimeUnit"] = apiCPEvent.API.ProdAIRL.TimeUnit
+
+		tokenConfig := make(map[string]interface{})
+		if apiCPEvent.API.ProdAIRL.PromptTokenCount != nil {
+			tokenConfig["productionMaxPromptTokenCount"] = *apiCPEvent.API.ProdAIRL.PromptTokenCount
+		}
+		if apiCPEvent.API.ProdAIRL.CompletionTokenCount != nil {
+			tokenConfig["productionMaxCompletionTokenCount"] = *apiCPEvent.API.ProdAIRL.CompletionTokenCount
+		}
+		if apiCPEvent.API.ProdAIRL.TotalTokenCount != nil {
+			tokenConfig["productionMaxTotalTokenCount"] = *apiCPEvent.API.ProdAIRL.TotalTokenCount
+		}
+		if len(tokenConfig) > 0 {
+			tokenConfig["isTokenBasedThrottlingEnabled"] = true
+			maxTps["tokenBasedThrottlingConfiguration"] = tokenConfig
+		}
+	}
+
+	// Handle Sandbox fields
+	if apiCPEvent.API.SandAIRL != nil {
+		maxTps["sandbox"] = apiCPEvent.API.SandAIRL.RequestCount
+		maxTps["sandboxTimeUnit"] = apiCPEvent.API.SandAIRL.TimeUnit
+
+		// Add sandbox token-based throttling config
+		if maxTps["tokenBasedThrottlingConfiguration"] == nil {
+			maxTps["tokenBasedThrottlingConfiguration"] = make(map[string]interface{})
+		}
+		tokenConfig := maxTps["tokenBasedThrottlingConfiguration"].(map[string]interface{})
+		if apiCPEvent.API.SandAIRL.PromptTokenCount != nil {
+			tokenConfig["sandboxMaxPromptTokenCount"] = *apiCPEvent.API.SandAIRL.PromptTokenCount
+		}
+		if apiCPEvent.API.SandAIRL.CompletionTokenCount != nil {
+			tokenConfig["sandboxMaxCompletionTokenCount"] = *apiCPEvent.API.SandAIRL.CompletionTokenCount
+		}
+		if apiCPEvent.API.SandAIRL.TotalTokenCount != nil {
+			tokenConfig["sandboxMaxTotalTokenCount"] = *apiCPEvent.API.SandAIRL.TotalTokenCount
+		}
+	}
+	if len(maxTps) > 0 {
+		data["data"].(map[string]interface{})["maxTps"] = maxTps
+	}
 	logger.LoggerMgtServer.Debugf("Prepared yaml : %+v", data)
 	definition := apiCPEvent.API.Definition
 	if strings.EqualFold(apiCPEvent.API.APIType, "rest") {
@@ -287,6 +333,10 @@ func createAPIYaml(apiCPEvent *APICPEvent) (string, string) {
 	return string(yamlBytes), definition
 }
 
+func getUint32OrNil(i *uint32) {
+
+}
+
 func createDeployementYaml(vhost string) string {
 	config, err := config.ReadConfigs()
 	envLabel := []string{"Default"}

diff --git a/apim-apk-agent/pkg/managementserver/types.go b/apim-apk-agent/pkg/managementserver/types.go
@@ -146,6 +146,17 @@ type API struct {
 	AuthHeader       string            `json:"authHeader"`
 	APIKeyHeader     string            `json:"apiKeyHeader"`
 	Operations       []OperationFromDP `json:"operations"`
+	SandAIRL         *AIRL              `json:"sandAIRL"`
+	ProdAIRL         *AIRL              `json:"prodAIRL"`
+}
+
+// AIRL holds AI ratelimit related data
+type AIRL struct {
+	PromptTokenCount     *uint32   `json:"promptTokenCount"`
+	CompletionTokenCount *uint32   `json:"CompletionTokenCount"`
+	TotalTokenCount      *uint32   `json:"totalTokenCount"`
+	TimeUnit             string `json:"timeUnit"`
+	RequestCount         *uint32   `json:"requestCount"`
 }
 
 // APKHeaders contains the request and response header modifier information