oracle-ads 2.11.15__py3-none-any.whl → 2.11.17__py3-none-any.whl

ads/aqua/extension/common_ws_msg_handler.py

@@ -0,0 +1,57 @@
+#!/usr/bin/env python
+
+# Copyright (c) 2024 Oracle and/or its affiliates.
+# Licensed under the Universal Permissive License v 1.0 as shown at https://oss.oracle.com/licenses/upl/
+
+import json
+from importlib import metadata
+from typing import List, Union
+
+from ads.aqua import ODSC_MODEL_COMPARTMENT_OCID, fetch_service_compartment
+from ads.aqua.common.decorator import handle_exceptions
+from ads.aqua.common.errors import AquaResourceAccessError
+from ads.aqua.common.utils import known_realm
+from ads.aqua.extension.aqua_ws_msg_handler import AquaWSMsgHandler
+from ads.aqua.extension.models.ws_models import (
+    AdsVersionResponse,
+    CompatibilityCheckResponse,
+    RequestResponseType,
+)
+
+
+class AquaCommonWsMsgHandler(AquaWSMsgHandler):
+    @staticmethod
+    def get_message_types() -> List[RequestResponseType]:
+        return [RequestResponseType.AdsVersion, RequestResponseType.CompatibilityCheck]
+
+    def __init__(self, message: Union[str, bytes]):
+        super().__init__(message)
+
+    @handle_exceptions
+    def process(self) -> Union[AdsVersionResponse, CompatibilityCheckResponse]:
+        request = json.loads(self.message)
+        if request.get("kind") == "AdsVersion":
+            version = metadata.version("oracle_ads")
+            response = AdsVersionResponse(
+                message_id=request.get("message_id"),
+                kind=RequestResponseType.AdsVersion,
+                data=version,
+            )
+            return response
+        if request.get("kind") == "CompatibilityCheck":
+            if ODSC_MODEL_COMPARTMENT_OCID or fetch_service_compartment():
+                return CompatibilityCheckResponse(
+                    message_id=request.get("message_id"),
+                    kind=RequestResponseType.CompatibilityCheck,
+                    data={"status": "ok"},
+                )
+            elif known_realm():
+                return CompatibilityCheckResponse(
+                    message_id=request.get("message_id"),
+                    kind=RequestResponseType.CompatibilityCheck,
+                    data={"status": "compatible"},
+                )
+            else:
+                raise AquaResourceAccessError(
+                    "The AI Quick actions extension is not compatible in the given region."
+                )

ads/aqua/extension/deployment_handler.py

@@ -1,5 +1,4 @@
 #!/usr/bin/env python
-# -*- coding: utf-8 -*-
 # Copyright (c) 2024 Oracle and/or its affiliates.
 # Licensed under the Universal Permissive License v 1.0 as shown at https://oss.oracle.com/licenses/upl/
 
@@ -8,8 +7,8 @@ from urllib.parse import urlparse
 from tornado.web import HTTPError
 
 from ads.aqua.common.decorator import handle_exceptions
-from ads.aqua.extension.errors import Errors
 from ads.aqua.extension.base_handler import AquaAPIhandler
+from ads.aqua.extension.errors import Errors
 from ads.aqua.modeldeployment import AquaDeploymentApp, MDInferenceResponse
 from ads.aqua.modeldeployment.entities import ModelParams
 from ads.config import COMPARTMENT_OCID, PROJECT_OCID
@@ -66,8 +65,8 @@ class AquaDeploymentHandler(AquaAPIhandler):
         """
         try:
             input_data = self.get_json_body()
-        except Exception:
-            raise HTTPError(400, Errors.INVALID_INPUT_DATA_FORMAT)
+        except Exception as ex:
+            raise HTTPError(400, Errors.INVALID_INPUT_DATA_FORMAT) from ex
 
         if not input_data:
             raise HTTPError(400, Errors.NO_INPUT_DATA)
@@ -100,6 +99,9 @@ class AquaDeploymentHandler(AquaAPIhandler):
         health_check_port = input_data.get("health_check_port")
         env_var = input_data.get("env_var")
         container_family = input_data.get("container_family")
+        ocpus = input_data.get("ocpus")
+        memory_in_gbs = input_data.get("memory_in_gbs")
+        model_file = input_data.get("model_file")
 
         self.finish(
             AquaDeploymentApp().create(
@@ -119,6 +121,9 @@ class AquaDeploymentHandler(AquaAPIhandler):
                 health_check_port=health_check_port,
                 env_var=env_var,
                 container_family=container_family,
+                ocpus=ocpus,
+                memory_in_gbs=memory_in_gbs,
+                model_file=model_file,
             )
         )
 
@@ -153,9 +158,7 @@ class AquaDeploymentInferenceHandler(AquaAPIhandler):
                 return False
             if not url.netloc:
                 return False
-            if not url.path.endswith("/predict"):
-                return False
-            return True
+            return url.path.endswith("/predict")
         except Exception:
             return False
 
@@ -170,8 +173,8 @@ class AquaDeploymentInferenceHandler(AquaAPIhandler):
         """
         try:
             input_data = self.get_json_body()
-        except Exception:
-            raise HTTPError(400, Errors.INVALID_INPUT_DATA_FORMAT)
+        except Exception as ex:
+            raise HTTPError(400, Errors.INVALID_INPUT_DATA_FORMAT) from ex
 
         if not input_data:
             raise HTTPError(400, Errors.NO_INPUT_DATA)
@@ -192,10 +195,10 @@ class AquaDeploymentInferenceHandler(AquaAPIhandler):
         )
         try:
             model_params_obj = ModelParams(**model_params)
-        except:
+        except Exception as ex:
             raise HTTPError(
                 400, Errors.INVALID_INPUT_DATA_FORMAT.format("model_params")
-            )
+            ) from ex
 
         return self.finish(
             MDInferenceResponse(prompt, model_params_obj).get_model_deployment_response(
@@ -236,8 +239,8 @@ class AquaDeploymentParamsHandler(AquaAPIhandler):
         """
         try:
             input_data = self.get_json_body()
-        except Exception:
-            raise HTTPError(400, Errors.INVALID_INPUT_DATA_FORMAT)
+        except Exception as ex:
+            raise HTTPError(400, Errors.INVALID_INPUT_DATA_FORMAT) from ex
 
         if not input_data:
             raise HTTPError(400, Errors.NO_INPUT_DATA)

ads/aqua/extension/deployment_ws_msg_handler.py

@@ -0,0 +1,54 @@
+#!/usr/bin/env python
+
+# Copyright (c) 2024 Oracle and/or its affiliates.
+# Licensed under the Universal Permissive License v 1.0 as shown at https://oss.oracle.com/licenses/upl/
+
+import json
+from typing import List, Union
+
+from ads.aqua.common.decorator import handle_exceptions
+from ads.aqua.extension.aqua_ws_msg_handler import AquaWSMsgHandler
+from ads.aqua.extension.models.ws_models import (
+    ListDeploymentResponse,
+    ModelDeploymentDetailsResponse,
+    RequestResponseType,
+)
+from ads.aqua.modeldeployment import AquaDeploymentApp
+from ads.config import COMPARTMENT_OCID
+
+
+class AquaDeploymentWSMsgHandler(AquaWSMsgHandler):
+    def __init__(self, message: Union[str, bytes]):
+        super().__init__(message)
+
+    @staticmethod
+    def get_message_types() -> List[RequestResponseType]:
+        return [
+            RequestResponseType.ListDeployments,
+            RequestResponseType.DeploymentDetails,
+        ]
+
+    @handle_exceptions
+    def process(self) -> Union[ListDeploymentResponse, ModelDeploymentDetailsResponse]:
+        request = json.loads(self.message)
+        if request.get("kind") == "ListDeployments":
+            deployment_list = AquaDeploymentApp().list(
+                compartment_id=request.get("compartment_id") or COMPARTMENT_OCID,
+                project_id=request.get("project_id"),
+            )
+            response = ListDeploymentResponse(
+                message_id=request.get("message_id"),
+                kind=RequestResponseType.ListDeployments,
+                data=deployment_list,
+            )
+            return response
+        elif request.get("kind") == "DeploymentDetails":
+            deployment_details = AquaDeploymentApp().get(
+                request.get("model_deployment_id")
+            )
+            response = ModelDeploymentDetailsResponse(
+                message_id=request.get("message_id"),
+                kind=RequestResponseType.DeploymentDetails,
+                data=deployment_details,
+            )
+            return response

ads/aqua/extension/errors.py

@@ -1,5 +1,4 @@
 #!/usr/bin/env python
-# -*- coding: utf-8 -*-
 # Copyright (c) 2024 Oracle and/or its affiliates.
 # Licensed under the Universal Permissive License v 1.0 as shown at https://oss.oracle.com/licenses/upl/
 
@@ -8,3 +7,4 @@ class Errors(str):
     INVALID_INPUT_DATA_FORMAT = "Invalid format of input data."
     NO_INPUT_DATA = "No input data provided."
     MISSING_REQUIRED_PARAMETER = "Missing required parameter: '{}'"
+    MISSING_ONEOF_REQUIRED_PARAMETER = "Either '{}' or '{}' is required."

ads/aqua/extension/evaluation_ws_msg_handler.py

@@ -3,13 +3,14 @@
 # Copyright (c) 2024 Oracle and/or its affiliates.
 # Licensed under the Universal Permissive License v 1.0 as shown at https://oss.oracle.com/licenses/upl/
 
+import json
 from typing import List, Union
 
 from ads.aqua.common.decorator import handle_exceptions
 from ads.aqua.evaluation import AquaEvaluationApp
 from ads.aqua.extension.aqua_ws_msg_handler import AquaWSMsgHandler
 from ads.aqua.extension.models.ws_models import (
-    ListEvaluationsRequest,
+    EvaluationDetailsResponse,
     ListEvaluationsResponse,
     RequestResponseType,
 )
@@ -19,21 +20,42 @@ from ads.config import COMPARTMENT_OCID
 class AquaEvaluationWSMsgHandler(AquaWSMsgHandler):
     @staticmethod
     def get_message_types() -> List[RequestResponseType]:
-        return [RequestResponseType.ListEvaluations]
+        return [
+            RequestResponseType.ListEvaluations,
+            RequestResponseType.EvaluationDetails,
+        ]
 
     def __init__(self, message: Union[str, bytes]):
         super().__init__(message)
 
     @handle_exceptions
-    def process(self) -> ListEvaluationsResponse:
-        list_eval_request = ListEvaluationsRequest.from_json(self.message)
+    def process(self) -> Union[ListEvaluationsResponse, EvaluationDetailsResponse]:
+        request = json.loads(self.message)
+        if request["kind"] == "ListEvaluations":
+            return self.list_evaluations(request)
+        if request["kind"] == "EvaluationDetails":
+            return self.evaluation_details(request)
 
+    @staticmethod
+    def list_evaluations(request) -> ListEvaluationsResponse:
         eval_list = AquaEvaluationApp().list(
-            list_eval_request.compartment_id or COMPARTMENT_OCID,
+            request.get("compartment_id") or COMPARTMENT_OCID
         )
         response = ListEvaluationsResponse(
-            message_id=list_eval_request.message_id,
+            message_id=request["message_id"],
             kind=RequestResponseType.ListEvaluations,
             data=eval_list,
         )
         return response
+
+    @staticmethod
+    def evaluation_details(request) -> EvaluationDetailsResponse:
+        evaluation_details = AquaEvaluationApp().get(
+            eval_id=request.get("evaluation_id")
+        )
+        response = EvaluationDetailsResponse(
+            message_id=request.get("message_id"),
+            kind=RequestResponseType.EvaluationDetails,
+            data=evaluation_details,
+        )
+        return response

ads/aqua/extension/model_handler.py

@@ -1,27 +1,66 @@
 #!/usr/bin/env python
-# -*- coding: utf-8 -*-
 # Copyright (c) 2024 Oracle and/or its affiliates.
 # Licensed under the Universal Permissive License v 1.0 as shown at https://oss.oracle.com/licenses/upl/
 
-import re
 from typing import Optional
 from urllib.parse import urlparse
 
 from tornado.web import HTTPError
-from ads.aqua.extension.errors import Errors
+
 from ads.aqua.common.decorator import handle_exceptions
+from ads.aqua.common.errors import AquaRuntimeError, AquaValueError
+from ads.aqua.common.utils import get_hf_model_info
 from ads.aqua.extension.base_handler import AquaAPIhandler
+from ads.aqua.extension.errors import Errors
 from ads.aqua.model import AquaModelApp
+from ads.aqua.model.constants import ModelTask
+from ads.aqua.model.entities import AquaModelSummary, HFModelSummary
+from ads.aqua.ui import ModelFormat
 
 
 class AquaModelHandler(AquaAPIhandler):
     """Handler for Aqua Model REST APIs."""
 
     @handle_exceptions
-    def get(self, model_id=""):
+    def get(
+        self,
+        model_id="",
+    ):
         """Handle GET request."""
-        if not model_id:
+        url_parse = urlparse(self.request.path)
+        paths = url_parse.path.strip("/")
+        if paths.startswith("aqua/model/files"):
+            os_path = self.get_argument("os_path", None)
+            model_name = self.get_argument("model_name", None)
+
+            model_format = self.get_argument("model_format")
+            if not model_format:
+                raise HTTPError(
+                    400, Errors.MISSING_REQUIRED_PARAMETER.format("model_format")
+                )
+            try:
+                model_format = ModelFormat(model_format.upper())
+            except ValueError as err:
+                raise AquaValueError(f"Invalid model format: {model_format}") from err
+            else:
+                if os_path:
+                    return self.finish(
+                        AquaModelApp.get_model_files(os_path, model_format)
+                    )
+                elif model_name:
+                    return self.finish(
+                        AquaModelApp.get_hf_model_files(model_name, model_format)
+                    )
+                else:
+                    raise HTTPError(
+                        400,
+                        Errors.MISSING_ONEOF_REQUIRED_PARAMETER.format(
+                            "os_path", "model_name"
+                        ),
+                    )
+        elif not model_id:
             return self.list()
+
         return self.read(model_id)
 
     def read(self, model_id):
@@ -29,7 +68,7 @@ class AquaModelHandler(AquaAPIhandler):
         return self.finish(AquaModelApp().get(model_id))
 
     @handle_exceptions
-    def delete(self, id=""):
+    def delete(self):
         """Handles DELETE request for clearing cache"""
         url_parse = urlparse(self.request.path)
         paths = url_parse.path.strip("/")
@@ -63,8 +102,8 @@ class AquaModelHandler(AquaAPIhandler):
         """
         try:
             input_data = self.get_json_body()
-        except Exception:
-            raise HTTPError(400, Errors.INVALID_INPUT_DATA_FORMAT)
+        except Exception as ex:
+            raise HTTPError(400, Errors.INVALID_INPUT_DATA_FORMAT) from ex
 
         if not input_data:
             raise HTTPError(400, Errors.NO_INPUT_DATA)
@@ -81,15 +120,21 @@ class AquaModelHandler(AquaAPIhandler):
         finetuning_container = input_data.get("finetuning_container")
         compartment_id = input_data.get("compartment_id")
         project_id = input_data.get("project_id")
+        model_file = input_data.get("model_file")
+        download_from_hf = (
+            str(input_data.get("download_from_hf", "false")).lower() == "true"
+        )
 
         return self.finish(
             AquaModelApp().register(
                 model=model,
                 os_path=os_path,
+                download_from_hf=download_from_hf,
                 inference_container=inference_container,
                 finetuning_container=finetuning_container,
                 compartment_id=compartment_id,
                 project_id=project_id,
+                model_file=model_file,
             )
         )
 
@@ -105,7 +150,88 @@ class AquaModelLicenseHandler(AquaAPIhandler):
         return self.finish(AquaModelApp().load_license(model_id))
 
 
+class AquaHuggingFaceHandler(AquaAPIhandler):
+    """Handler for Aqua Hugging Face REST APIs."""
+
+    @staticmethod
+    def _find_matching_aqua_model(model_id: str) -> Optional[AquaModelSummary]:
+        """
+        Finds a matching model in AQUA based on the model ID from Hugging Face.
+
+        Parameters
+        ----------
+        model_id (str): The Hugging Face model ID to match.
+
+        Returns
+        -------
+        Optional[AquaModelSummary]
+            Returns the matching AquaModelSummary object if found, else None.
+        """
+        # Convert the Hugging Face model ID to lowercase once
+        model_id_lower = model_id.lower()
+
+        aqua_model_app = AquaModelApp()
+        model_ocid = aqua_model_app._find_matching_aqua_model(model_id=model_id_lower)
+        if model_ocid:
+            return aqua_model_app.get(model_ocid, load_model_card=False)
+
+        return None
+
+    @handle_exceptions
+    def post(self, *args, **kwargs):
+        """Handles post request for the HF Models APIs
+
+        Raises
+        ------
+        HTTPError
+            Raises HTTPError if inputs are missing or are invalid.
+        """
+        try:
+            input_data = self.get_json_body()
+        except Exception as ex:
+            raise HTTPError(400, Errors.INVALID_INPUT_DATA_FORMAT) from ex
+
+        if not input_data:
+            raise HTTPError(400, Errors.NO_INPUT_DATA)
+
+        model_id = input_data.get("model_id")
+
+        if not model_id:
+            raise HTTPError(400, Errors.MISSING_REQUIRED_PARAMETER.format("model_id"))
+
+        # Get model info from the HF
+        hf_model_info = get_hf_model_info(repo_id=model_id)
+
+        # Check if model is not disabled
+        if hf_model_info.disabled:
+            raise AquaRuntimeError(
+                f"The chosen model '{hf_model_info.id}' is currently disabled and cannot be imported into AQUA. "
+                "Please verify the model's status on the Hugging Face Model Hub or select a different model."
+            )
+
+        # Check pipeline_tag, it should be `text-generation`
+        if (
+            not hf_model_info.pipeline_tag
+            or hf_model_info.pipeline_tag.lower() != ModelTask.TEXT_GENERATION
+        ):
+            raise AquaRuntimeError(
+                f"Unsupported pipeline tag for the chosen model: '{hf_model_info.pipeline_tag}'. "
+                f"AQUA currently supports the following tasks only: {', '.join(ModelTask.values())}. "
+                "Please select a model with a compatible pipeline tag."
+            )
+
+        # Check if it is a service/verified model
+        aqua_model_info: AquaModelSummary = self._find_matching_aqua_model(
+            model_id=hf_model_info.id
+        )
+
+        return self.finish(
+            HFModelSummary(model_info=hf_model_info, aqua_model_info=aqua_model_info)
+        )
+
+
 __handlers__ = [
     ("model/?([^/]*)", AquaModelHandler),
     ("model/?([^/]*)/license", AquaModelLicenseHandler),
+    ("model/hf/search/?([^/]*)", AquaHuggingFaceHandler),
 ]

ads/aqua/extension/models/ws_models.py

@@ -7,15 +7,22 @@
 from dataclasses import dataclass
 from typing import List, Optional
 
-from ads.aqua.evaluation.entities import AquaEvaluationSummary
-from ads.aqua.model.entities import AquaModelSummary
+from ads.aqua.evaluation.entities import AquaEvaluationSummary, AquaEvaluationDetail
+from ads.aqua.model.entities import AquaModelSummary, AquaModel
+from ads.aqua.modeldeployment.entities import AquaDeployment, AquaDeploymentDetail
 from ads.common.extended_enum import ExtendedEnumMeta
 from ads.common.serializer import DataClassSerializable
 
 
 class RequestResponseType(str, metaclass=ExtendedEnumMeta):
     ListEvaluations = "ListEvaluations"
+    EvaluationDetails = "EvaluationDetails"
+    ListDeployments = "ListDeployments"
+    DeploymentDetails = "DeploymentDetails"
     ListModels = "ListModels"
+    ModelDetails = "ModelDetails"
+    AdsVersion = "AdsVersion"
+    CompatibilityCheck = "CompatibilityCheck"
     Error = "Error"
 
 
@@ -23,7 +30,7 @@ class RequestResponseType(str, metaclass=ExtendedEnumMeta):
 class BaseResponse(DataClassSerializable):
     message_id: str
     kind: RequestResponseType
-    data: object
+    data: Optional[object]
 
 
 @dataclass
@@ -40,9 +47,37 @@ class ListEvaluationsRequest(BaseRequest):
     kind = RequestResponseType.ListEvaluations
 
 
+@dataclass
+class EvaluationDetailsRequest(BaseRequest):
+    kind = RequestResponseType.EvaluationDetails
+    evaluation_id: str
+
+
 @dataclass
 class ListModelsRequest(BaseRequest):
     compartment_id: Optional[str] = None
+    project_id: Optional[str] = None
+    model_type: Optional[str] = None
+    kind = RequestResponseType.ListDeployments
+
+
+@dataclass
+class ModelDetailsRequest(BaseRequest):
+    kind = RequestResponseType.ModelDetails
+    model_id: str
+
+
+@dataclass
+class ListDeploymentRequest(BaseRequest):
+    compartment_id: str
+    project_id: Optional[str] = None
+    kind = RequestResponseType.ListDeployments
+
+
+@dataclass
+class DeploymentDetailsRequest(BaseRequest):
+    model_deployment_id: str
+    kind = RequestResponseType.DeploymentDetails
 
 
 @dataclass
@@ -50,11 +85,51 @@ class ListEvaluationsResponse(BaseResponse):
     data: List[AquaEvaluationSummary]
 
 
+@dataclass
+class EvaluationDetailsResponse(BaseResponse):
+    data: AquaEvaluationDetail
+
+
+@dataclass
+class ListDeploymentResponse(BaseResponse):
+    data: List[AquaDeployment]
+
+
+@dataclass
+class ModelDeploymentDetailsResponse(BaseResponse):
+    data: AquaDeploymentDetail
+
+
 @dataclass
 class ListModelsResponse(BaseResponse):
     data: List[AquaModelSummary]
 
 
+@dataclass
+class ModelDetailsResponse(BaseResponse):
+    data: AquaModel
+
+
+@dataclass
+class AdsVersionRequest(BaseRequest):
+    kind: RequestResponseType.AdsVersion
+
+
+@dataclass
+class AdsVersionResponse(BaseResponse):
+    data: str
+
+
+@dataclass
+class CompatibilityCheckRequest(BaseRequest):
+    kind: RequestResponseType.CompatibilityCheck
+
+
+@dataclass
+class CompatibilityCheckResponse(BaseResponse):
+    data: object
+
+
 @dataclass
 class AquaWsError(DataClassSerializable):
     status: str

ads/aqua/extension/models_ws_msg_handler.py

@@ -0,0 +1,49 @@
+#!/usr/bin/env python
+
+# Copyright (c) 2024 Oracle and/or its affiliates.
+# Licensed under the Universal Permissive License v 1.0 as shown at https://oss.oracle.com/licenses/upl/
+
+import json
+from typing import List, Union
+
+from ads.aqua.common.decorator import handle_exceptions
+from ads.aqua.extension.aqua_ws_msg_handler import AquaWSMsgHandler
+from ads.aqua.extension.models.ws_models import (
+    ListModelsResponse,
+    ModelDetailsResponse,
+    RequestResponseType,
+)
+from ads.aqua.model import AquaModelApp
+
+
+class AquaModelWSMsgHandler(AquaWSMsgHandler):
+    def __init__(self, message: Union[str, bytes]):
+        super().__init__(message)
+
+    @staticmethod
+    def get_message_types() -> List[RequestResponseType]:
+        return [RequestResponseType.ListModels, RequestResponseType.ModelDetails]
+
+    @handle_exceptions
+    def process(self) -> Union[ListModelsResponse, ModelDetailsResponse]:
+        request = json.loads(self.message)
+        if request.get("kind") == "ListModels":
+            models_list = AquaModelApp().list(
+                compartment_id=request.get("compartment_id"),
+                project_id=request.get("project_id"),
+                model_type=request.get("model_type"),
+            )
+            response = ListModelsResponse(
+                message_id=request.get("message_id"),
+                kind=RequestResponseType.ListModels,
+                data=models_list,
+            )
+            return response
+        elif request.get("kind") == "ModelDetails":
+            model_id = request.get("model_id")
+            response = AquaModelApp().get(model_id)
+            return ModelDetailsResponse(
+                message_id=request.get("message_id"),
+                kind=RequestResponseType.ModelDetails,
+                data=response,
+            )

ads/aqua/extension/ui_websocket_handler.py

@@ -14,6 +14,8 @@ from tornado.websocket import WebSocketHandler
 
 from ads.aqua import logger
 from ads.aqua.extension.aqua_ws_msg_handler import AquaWSMsgHandler
+from ads.aqua.extension.common_ws_msg_handler import AquaCommonWsMsgHandler
+from ads.aqua.extension.deployment_ws_msg_handler import AquaDeploymentWSMsgHandler
 from ads.aqua.extension.evaluation_ws_msg_handler import AquaEvaluationWSMsgHandler
 from ads.aqua.extension.models.ws_models import (
     AquaWsError,
@@ -22,6 +24,7 @@ from ads.aqua.extension.models.ws_models import (
     ErrorResponse,
     RequestResponseType,
 )
+from ads.aqua.extension.models_ws_msg_handler import AquaModelWSMsgHandler
 
 MAX_WORKERS = 20
 
@@ -43,7 +46,10 @@ def get_aqua_internal_error_response(message_id: str) -> ErrorResponse:
 class AquaUIWebSocketHandler(WebSocketHandler):
     """Handler for Aqua Websocket."""
 
-    _handlers_: List[Type[AquaWSMsgHandler]] = [AquaEvaluationWSMsgHandler]
+    _handlers_: List[Type[AquaWSMsgHandler]] = [AquaEvaluationWSMsgHandler,
+                                                AquaDeploymentWSMsgHandler,
+                                                AquaModelWSMsgHandler,
+                                                AquaCommonWsMsgHandler]
 
     thread_pool: ThreadPoolExecutor