websocket changes for 1.0.3

Author: kumar-shivam-ranjan

ads/aqua/extension/deployment_ws_msg_handler.py

@@ -0,0 +1,31 @@
+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 RequestResponseType, ListDeploymentResponse, ListDeploymentRequest
+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]
+
+    @handle_exceptions
+    def process(self) -> ListDeploymentResponse:
+        list_deployment_request = ListDeploymentRequest.from_json(self.message)
+        deployment_list = AquaDeploymentApp().list(
+            compartment_id=list_deployment_request.compartment_id or COMPARTMENT_OCID,
+            project_id=list_deployment_request.project_id,
+        )
+        response = ListDeploymentResponse(
+            message_id=list_deployment_request.message_id,
+            kind=RequestResponseType.ListDeployments,
+            data=deployment_list,
+        )
+        return response

ads/aqua/extension/models/ws_models.py

@@ -9,12 +9,14 @@
 
 from ads.aqua.evaluation.entities import AquaEvaluationSummary
 from ads.aqua.model.entities import AquaModelSummary
+from ads.aqua.modeldeployment.entities import AquaDeployment
 from ads.common.extended_enum import ExtendedEnumMeta
 from ads.common.serializer import DataClassSerializable
 
 
 class RequestResponseType(str, metaclass=ExtendedEnumMeta):
     ListEvaluations = "ListEvaluations"
+    ListDeployments = "ListDeployments"
     ListModels = "ListModels"
     Error = "Error"
 
@@ -43,13 +45,26 @@ class ListEvaluationsRequest(BaseRequest):
 @dataclass
 class ListModelsRequest(BaseRequest):
     compartment_id: Optional[str] = None
+    project_id: Optional[str] = None
+    model_type: Optional[str] = None
+    kind = RequestResponseType.ListDeployments
 
 
 @dataclass
 class ListEvaluationsResponse(BaseResponse):
     data: List[AquaEvaluationSummary]
 
+@dataclass
+class ListDeploymentRequest(BaseRequest):
+    compartment_id: str
+    project_id: Optional[str] = None
+    kind = RequestResponseType.ListDeployments
 
+
+@dataclass
+class ListDeploymentResponse(BaseResponse):
+    data: List[AquaDeployment]
+    
 @dataclass
 class ListModelsResponse(BaseResponse):
     data: List[AquaModelSummary]

ads/aqua/extension/models_ws_msg_handler.py

@@ -0,0 +1,33 @@
+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 RequestResponseType,ListModelsResponse, ListModelsRequest
+from ads.aqua.model import AquaModelApp
+from ads.config import COMPARTMENT_OCID
+
+
+class AquaModelWSMsgHandler(AquaWSMsgHandler):
+
+    def __init__(self, message: Union[str, bytes]):
+        super().__init__(message)
+
+    @staticmethod
+    def get_message_types() -> List[RequestResponseType]:
+        return [RequestResponseType.ListModels]
+
+    @handle_exceptions
+    def process(self) -> ListModelsResponse:
+        list_models_request = ListModelsRequest.from_json(self.message)
+        print(list_models_request)
+        models_list = AquaModelApp().list(
+            compartment_id=list_models_request.compartment_id or COMPARTMENT_OCID,
+            project_id=list_models_request.project_id,
+            model_type=list_models_request.model_type
+        )
+        response = ListModelsResponse(
+            message_id=list_models_request.message_id,
+            kind=RequestResponseType.ListModels,
+            data=models_list,
+        )
+        return response

ads/aqua/extension/ui_websocket_handler.py

@@ -14,6 +14,7 @@
 
 from ads.aqua import logger
 from ads.aqua.extension.aqua_ws_msg_handler import AquaWSMsgHandler
+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 +23,7 @@
     ErrorResponse,
     RequestResponseType,
 )
+from ads.aqua.extension.models_ws_msg_handler import AquaModelWSMsgHandler
 
 MAX_WORKERS = 20
 
@@ -43,7 +45,7 @@ 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]
 
     thread_pool: ThreadPoolExecutor