Merge pull request #12647 from PasanT9/ai-config

Fix AI API issues in distributed setup
wso2 · Oct 9, 2024 · c2f0028 · c2f0028
2 parents 35bd0e9 + 08e2231
commit c2f0028
Show file tree

Hide file tree

Showing 6 changed files with 60 additions and 13 deletions.
diff --git a/...way/src/main/java/org/wso2/carbon/apimgt/gateway/listeners/GatewayJMSMessageListener.java b/...way/src/main/java/org/wso2/carbon/apimgt/gateway/listeners/GatewayJMSMessageListener.java
@@ -423,11 +423,23 @@ public void run() {
                 }
             }
         } else if (EventType.LLM_PROVIDER_CREATE.toString().equals(eventType)) {
-            createLLMProvider(eventJson);
+            try {
+                createLLMProvider(eventJson);
+            } catch (Exception e) {
+                log.error("Error while handling LLM_PROVIDER_CREATE event", e);
+            }
         } else if (EventType.LLM_PROVIDER_DELETE.toString().equals(eventType)) {
-            deleteLLMProvider(eventJson);
+            try {
+                deleteLLMProvider(eventJson);
+            } catch (Exception e) {
+                log.error("Error while handling LLM_PROVIDER_DELETE event", e);
+            }
         } else if (EventType.LLM_PROVIDER_UPDATE.toString().equals(eventType)) {
-            updateLLMProvider(eventJson);
+            try {
+                updateLLMProvider(eventJson);
+            } catch (Exception e) {
+                log.error("Error while handling LLM_PROVIDER_UPDATE event", e);
+            }
         }
     }
 

diff --git a/...t/org.wso2.carbon.apimgt.impl/src/main/java/org/wso2/carbon/apimgt/impl/APIConstants.java b/...t/org.wso2.carbon.apimgt.impl/src/main/java/org/wso2/carbon/apimgt/impl/APIConstants.java
@@ -519,7 +519,7 @@ public final class APIConstants {
     public static class AI {
 
         public static final String API_CHAT = "APIChat";
-        public static final String API_CHAT_ENABLED = "Enabled";
+        public static final String ENABLED = "Enabled";
         public static final String API_CHAT_AUTH_TOKEN = "AuthToken";
         public static final String API_CHAT_KEY = "Key";
         public static final String API_CHAT_ENDPOINT = "Endpoint";
@@ -541,6 +541,7 @@ public static class AI {
         public static final String MARKETPLACE_ASSISTANT_PUBLISH_API_RESOURCE = "ApiPublishResource";
         public static final String MARKETPLACE_ASSISTANT_DELETE_API_RESOURCE = "ApiDeleteResource";
         public static final String MARKETPLACE_ASSISTANT_API_COUNT_RESOURCE = "ApiCountResource";
+        public static final String AI_CONFIGURATION = "AiConfiguration";
 
         private AI() {
 

diff --git a/...carbon.apimgt.impl/src/main/java/org/wso2/carbon/apimgt/impl/APIManagerConfiguration.java b/...carbon.apimgt.impl/src/main/java/org/wso2/carbon/apimgt/impl/APIManagerConfiguration.java
@@ -135,6 +135,7 @@ public class APIManagerConfiguration {
     private JSONObject subscriberAttributes = new JSONObject();
     private static Map<String, String> analyticsMaskProps;
     private TokenValidationDto tokenValidationDto = new TokenValidationDto();
+    private boolean enableAiConfiguration;
 
     public Map<String, List<String>> getRestApiJWTAuthAudiences() {
         return restApiJWTAuthAudiences;
@@ -658,6 +659,8 @@ private void readChildElements(OMElement serverConfig,
                 setMarketplaceAssistantConfiguration(element);
             } else if (APIConstants.AI.API_CHAT.equals(localName)) {
                 setApiChatConfiguration(element);
+            } else if (APIConstants.AI.AI_CONFIGURATION.equals(localName)){
+                setAiConfiguration(element);
             } else if (APIConstants.TokenValidationConstants.TOKEN_VALIDATION_CONFIG.equals(localName)) {
                 setTokenValidation(element);
             }
@@ -2489,9 +2492,28 @@ public void setMarketplaceAssistantConfiguration(OMElement omElement){
         }
     }
 
+    private void setAiConfiguration(OMElement omElement) {
+
+        OMElement aiConfigurationEnabled =
+                omElement.getFirstChildWithName(new QName(APIConstants.AI.ENABLED));
+        if (aiConfigurationEnabled != null) {
+            setEnableAiConfiguration(Boolean.parseBoolean(aiConfigurationEnabled.getText()));
+        }
+    }
+
+    public boolean isEnableAiConfiguration() {
+
+        return enableAiConfiguration;
+    }
+
+    public void setEnableAiConfiguration(boolean enableAiConfiguration) {
+
+        this.enableAiConfiguration = enableAiConfiguration;
+    }
+
     public void setApiChatConfiguration(OMElement omElement){
         OMElement apiChatEnableElement =
-                omElement.getFirstChildWithName(new QName(APIConstants.AI.API_CHAT_ENABLED));
+                omElement.getFirstChildWithName(new QName(APIConstants.AI.ENABLED));
         if (apiChatEnableElement != null) {
             apiChatConfigurationDto.setEnabled(Boolean.parseBoolean(apiChatEnableElement.getText()));
         }

diff --git a/...pimgt.impl/src/main/java/org/wso2/carbon/apimgt/impl/listeners/ServerStartupListener.java b/...pimgt.impl/src/main/java/org/wso2/carbon/apimgt/impl/listeners/ServerStartupListener.java
@@ -61,10 +61,14 @@ public void completedServerStartup() {
                 Map<String, TokenIssuerDto> tokenIssuerDtoMap =
                         apiManagerConfiguration.getJwtConfigurationDto().getTokenIssuerDtoMap();
                 tokenIssuerDtoMap.forEach((issuer, tokenIssuer) -> KeyManagerHolder.addGlobalJWTValidators(tokenIssuer));
-                try {
-                    LLMProviderRegistrationService.registerDefaultLLMProviders(MultitenantConstants.SUPER_TENANT_DOMAIN_NAME);
-                } catch (APIManagementException e) {
-                    log.error("Error occurred during onboarding pre defined LLM Providers", e);
+                if (apiManagerConfiguration.isEnableAiConfiguration()) {
+                    try {
+                        LLMProviderRegistrationService
+                                .registerDefaultLLMProviders(MultitenantConstants.SUPER_TENANT_DOMAIN_NAME);
+                    } catch (APIManagementException e) {
+                        log.error("Error occurred during onboarding pre defined LLM Providers in tenant "
+                                + MultitenantConstants.SUPER_TENANT_DOMAIN_NAME, e);
+                    }
                 }
             }
         } else {

diff --git a/...apimgt.impl/src/main/java/org/wso2/carbon/apimgt/impl/observers/CommonConfigDeployer.java b/...apimgt.impl/src/main/java/org/wso2/carbon/apimgt/impl/observers/CommonConfigDeployer.java
@@ -121,10 +121,13 @@ public void run() {
         KeyManagerConfigurationDataRetriever keyManagerConfigurationDataRetriever =
                 new KeyManagerConfigurationDataRetriever(tenantDomain);
         keyManagerConfigurationDataRetriever.startLoadKeyManagerConfigurations();
-        try {
-            LLMProviderRegistrationService.registerDefaultLLMProviders(tenantDomain);
-        } catch (APIManagementException e) {
-            log.error("Error occurred during onboarding pre defined LLM Providers");
+        if (configuration.isEnableAiConfiguration()) {
+            try {
+                LLMProviderRegistrationService.registerDefaultLLMProviders(tenantDomain);
+            } catch (APIManagementException e) {
+                log.error("Error occurred during onboarding pre defined LLM Providers in tenant domain "
+                        + tenantDomain, e);
+            }
         }
     }
 }
diff --git a/...re.feature/src/main/resources/conf_templates/templates/repository/conf/api-manager.xml.j2 b/...re.feature/src/main/resources/conf_templates/templates/repository/conf/api-manager.xml.j2
@@ -389,6 +389,11 @@
 
     </APIChat>
 
+
+    <AiConfiguration>
+        <Enabled>{{apim.ai.llm_provider_events_enabled}}</Enabled>
+    </AiConfiguration>
+
     <!--
         Configurations relevant to enable rate limit feature.
         This enables to use resources considering subscription levels.