Provide the support to work with common AI component

wso2 · Jan 17, 2025 · 0abe1ef · 0abe1ef
1 parent 716699e
commit 0abe1ef
Show file tree

Hide file tree

Showing 3 changed files with 30 additions and 19 deletions.
diff --git a/....management/org.wso2.carbon.identity.api.server.branding.preference.management.v1/pom.xml b/....management/org.wso2.carbon.identity.api.server.branding.preference.management.v1/pom.xml
@@ -171,5 +171,10 @@
             <artifactId>org.wso2.carbon.identity.application.authentication.framework</artifactId>
             <scope>provided</scope>
         </dependency>
+        <dependency>
+            <groupId>org.wso2.carbon.identity.framework</groupId>
+            <artifactId>org.wso2.carbon.identity.ai.service.mgt</artifactId>
+            <scope>provided</scope>
+        </dependency>
     </dependencies>
 </project>
diff --git a/.../server/branding/preference/management/v1/core/BrandingAIPreferenceManagementService.java b/.../server/branding/preference/management/v1/core/BrandingAIPreferenceManagementService.java
@@ -20,15 +20,15 @@
 
 import org.apache.commons.logging.Log;
 import org.apache.commons.logging.LogFactory;
+import org.wso2.carbon.identity.ai.service.mgt.exceptions.AIClientException;
+import org.wso2.carbon.identity.ai.service.mgt.exceptions.AIServerException;
 import org.wso2.carbon.identity.api.server.branding.preference.management.common.BrandingPreferenceServiceHolder;
 import org.wso2.carbon.identity.api.server.branding.preference.management.v1.model.BrandingGenerationRequestModel;
 import org.wso2.carbon.identity.api.server.branding.preference.management.v1.model.BrandingGenerationResponseModel;
 import org.wso2.carbon.identity.api.server.branding.preference.management.v1.model.BrandingGenerationResultModel;
 import org.wso2.carbon.identity.api.server.branding.preference.management.v1.model.BrandingGenerationStatusModel;
 import org.wso2.carbon.identity.api.server.common.error.APIError;
 import org.wso2.carbon.identity.api.server.common.error.ErrorResponse;
-import org.wso2.carbon.identity.branding.preference.management.core.exception.BrandingAIClientException;
-import org.wso2.carbon.identity.branding.preference.management.core.exception.BrandingAIServerException;
 
 import java.util.HashMap;
 import java.util.List;
@@ -61,9 +61,9 @@ public BrandingGenerationResponseModel generateBrandingPreference(
             BrandingGenerationResponseModel response = new BrandingGenerationResponseModel();
             response.setOperationId(operationId);
             return response;
-        } catch (BrandingAIServerException e) {
+        } catch (AIServerException e) {
             throw handleServerException(e);
-        } catch (BrandingAIClientException e) {
+        } catch (AIClientException e) {
             throw handleClientException(e);
         }
     }
@@ -82,9 +82,9 @@ public BrandingGenerationStatusModel getBrandingPreferenceGenerationStatus(Strin
             BrandingGenerationStatusModel response = new BrandingGenerationStatusModel();
             response.setStatus(convertObjectToMap(generationStatus));
             return response;
-        } catch (BrandingAIServerException e) {
+        } catch (AIServerException e) {
             throw handleServerException(e);
-        } catch (BrandingAIClientException e) {
+        } catch (AIClientException e) {
             throw handleClientException(e);
         }
     }
@@ -106,22 +106,22 @@ public BrandingGenerationResultModel getBrandingPreferenceGenerationResult(Strin
             response.setStatus(getStatusFromResult(resultMap));
 
             if (!resultMap.containsKey("data")) {
-                throw new BrandingAIServerException(ERROR_CODE_ERROR_GETTING_BRANDING_RESULT_STATUS.getMessage(),
+                throw new AIServerException(ERROR_CODE_ERROR_GETTING_BRANDING_RESULT_STATUS.getMessage(),
                         ERROR_CODE_ERROR_GETTING_BRANDING_RESULT_STATUS.getCode());
             }
 
             Map<String, Object> dataMap = (Map<String, Object>) resultMap.get("data");
             response.setData(dataMap);
             return response;
-        } catch (BrandingAIServerException e) {
+        } catch (AIServerException e) {
             throw handleServerException(e);
-        } catch (BrandingAIClientException e) {
+        } catch (AIClientException e) {
             throw handleClientException(e);
         }
     }
 
     private BrandingGenerationResultModel.StatusEnum getStatusFromResult(Map<String, Object> resultMap)
-            throws BrandingAIServerException {
+            throws AIServerException {
 
         if (resultMap.containsKey("status")) {
             String status = (String) resultMap.get("status");
@@ -133,32 +133,32 @@ private BrandingGenerationResultModel.StatusEnum getStatusFromResult(Map<String,
                 return BrandingGenerationResultModel.StatusEnum.FAILED;
             }
         }
-        throw new BrandingAIServerException(ERROR_CODE_ERROR_GETTING_BRANDING_RESULT.getMessage(),
+        throw new AIServerException(ERROR_CODE_ERROR_GETTING_BRANDING_RESULT.getMessage(),
                 ERROR_CODE_ERROR_GETTING_BRANDING_RESULT.getCode());
     }
 
 
-    private APIError handleClientException(BrandingAIClientException error) {
+    private APIError handleClientException(AIClientException error) {
 
         ErrorResponse.Builder errorResponseBuilder = new ErrorResponse.Builder()
                 .withCode(error.getErrorCode())
                 .withMessage(error.getMessage());
-        if (error.getBrandingAIResponse() != null) {
-            Response.Status status = Response.Status.fromStatusCode(error.getBrandingAIResponse().getStatusCode());
-            errorResponseBuilder.withDescription(error.getBrandingAIResponse().getResponseBody());
+        if (error.getServerMessage() != null) {
+            Response.Status status = Response.Status.fromStatusCode(error.getServerStatusCode());
+            errorResponseBuilder.withDescription(error.getServerMessage());
             return new APIError(status, errorResponseBuilder.build());
         }
         return new APIError(Response.Status.BAD_REQUEST, errorResponseBuilder.build());
     }
 
-    private APIError handleServerException(BrandingAIServerException error) {
+    private APIError handleServerException(AIServerException error) {
 
         ErrorResponse.Builder errorResponseBuilder = new ErrorResponse.Builder()
                 .withCode(error.getErrorCode())
                 .withMessage(error.getMessage());
-        if (error.getBrandingAIResponse() != null) {
-            Response.Status status = Response.Status.fromStatusCode(error.getBrandingAIResponse().getStatusCode());
-            errorResponseBuilder.withDescription(error.getBrandingAIResponse().getResponseBody());
+        if (error.getServerMessage() != null) {
+            Response.Status status = Response.Status.fromStatusCode(error.getServerStatusCode());
+            errorResponseBuilder.withDescription(error.getServerMessage());
             return new APIError(status, errorResponseBuilder.build());
         }
         return new APIError(Response.Status.INTERNAL_SERVER_ERROR, errorResponseBuilder.build());

diff --git a/pom.xml b/pom.xml
@@ -550,6 +550,12 @@
                 <version>${carbon.identity.framework.version}</version>
                 <scope>provided</scope>
             </dependency>
+            <dependency>
+                <groupId>org.wso2.carbon.identity.framework</groupId>
+                <artifactId>org.wso2.carbon.identity.ai.service.mgt</artifactId>
+                <version>${carbon.identity.framework.version}</version>
+                <scope>provided</scope>
+            </dependency>
             <dependency>
                 <groupId>org.wso2.carbon</groupId>
                 <artifactId>org.wso2.carbon.admin.advisory.mgt</artifactId>