Introduce generic property name to enforce the target endpoint

Author: ashera96

components/apimgt/org.wso2.carbon.apimgt.api/src/main/java/org/wso2/carbon/apimgt/api/APIConstants.java

@@ -182,6 +182,9 @@ public static class AIAPIConstants {
         public static final String TRAFFIC_FLOW_DIRECTION_IN = "IN";
         public static final String TRAFFIC_FLOW_DIRECTION_OUT = "OUT";
         public static final String API_LLM_ENDPOINT = "_API_LLMEndpoint_";
+        // New generic property name for passing target model configurations between mediators.
+        public static final String TARGET_MODEL_CONFIGS = "TARGET_MODEL_CONFIGS";
+        // Backward compatible constant name. New code should use TARGET_MODEL_CONFIGS.
         public static final String ROUND_ROBIN_CONFIGS = "ROUND_ROBIN_CONFIGS";
         public static final String FAILOVER_CONFIGS = "FAILOVER_CONFIGS";
         public static final String TARGET_MODEL_ENDPOINT = "TARGET_MODEL_ENDPOINT";

components/apimgt/org.wso2.carbon.apimgt.gateway/src/main/java/org/wso2/carbon/apimgt/gateway/mediators/AIAPIMediator.java

@@ -175,11 +175,17 @@ private void processInboundRequest(MessageContext messageContext,
             return;
         }
 
-        Map<String, Object> roundRobinConfigs;
-        if (messageContext.getProperty(APIConstants.AIAPIConstants.ROUND_ROBIN_CONFIGS) != null) {
-            roundRobinConfigs =
-                    (Map<String, Object>) messageContext.getProperty(APIConstants.AIAPIConstants.ROUND_ROBIN_CONFIGS);
-            handleLoadBalancing(messageContext, providerConfiguration, roundRobinConfigs, provider);
+        // Prefer the new generic TARGET_MODEL_CONFIGS; fall back to ROUND_ROBIN_CONFIGS for compatibility
+        Map<String, Object> targetModelConfigs = null;
+        if (messageContext.getProperty(APIConstants.AIAPIConstants.TARGET_MODEL_CONFIGS) != null) {
+            targetModelConfigs = (Map<String, Object>) messageContext
+                    .getProperty(APIConstants.AIAPIConstants.TARGET_MODEL_CONFIGS);
+        } else if (messageContext.getProperty(APIConstants.AIAPIConstants.ROUND_ROBIN_CONFIGS) != null) {
+            targetModelConfigs = (Map<String, Object>) messageContext
+                    .getProperty(APIConstants.AIAPIConstants.ROUND_ROBIN_CONFIGS);
+        }
+        if (targetModelConfigs != null) {
+            handleLoadBalancing(messageContext, providerConfiguration, targetModelConfigs, provider);
             if (log.isDebugEnabled()) {
                 log.debug("Load balancing configured, processing with round-robin configurations");
             }
@@ -649,10 +655,14 @@ private void processOutboundResponse(MessageContext messageContext,
         llmProviderService.getResponseMetadata(llmResponseMetaData, providerConfigs.getMetadata(), metadataMap);
         messageContext.setProperty(APIConstants.AIAPIConstants.AI_API_RESPONSE_METADATA, metadataMap);
 
+        // Read target model configs from the new generic property first, then fall back for compatibility
         Map<String, Object> roundRobinConfigs = null;
-        if (messageContext.getProperty(APIConstants.AIAPIConstants.ROUND_ROBIN_CONFIGS) != null) {
-            roundRobinConfigs =
-                    (Map<String, Object>) messageContext.getProperty(APIConstants.AIAPIConstants.ROUND_ROBIN_CONFIGS);
+        if (messageContext.getProperty(APIConstants.AIAPIConstants.TARGET_MODEL_CONFIGS) != null) {
+            roundRobinConfigs = (Map<String, Object>) messageContext
+                    .getProperty(APIConstants.AIAPIConstants.TARGET_MODEL_CONFIGS);
+        } else if (messageContext.getProperty(APIConstants.AIAPIConstants.ROUND_ROBIN_CONFIGS) != null) {
+            roundRobinConfigs = (Map<String, Object>) messageContext
+                    .getProperty(APIConstants.AIAPIConstants.ROUND_ROBIN_CONFIGS);
         }
 
         Map<String, Object> failoverConfigs = null;