Merge branch 'feature/aquav1.0.3' into ODSC-59697/llama-cpp-deployment

Author: VipulMascarenhas

ads/aqua/extension/common_ws_msg_handler.py

@@ -0,0 +1,45 @@
+import json
+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 RequestResponseType, AdsVersionResponse, AdsVersionRequest, \
+    CompatibilityCheckResponse
+from importlib import metadata
+
+
+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) -> 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(
+                    f"The AI Quick actions extension is not compatible in the given region."
+                )

ads/aqua/extension/deployment_ws_msg_handler.py

@@ -0,0 +1,40 @@
+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 RequestResponseType, ListDeploymentResponse, ListDeploymentRequest, \
+    ModelDeploymentDetailsResponse
+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) -> 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/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 @@
 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) -> 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/models/ws_models.py

@@ -7,23 +7,30 @@
 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"
 
 
 @dataclass
 class BaseResponse(DataClassSerializable):
     message_id: str
     kind: RequestResponseType
-    data: object
+    data: Optional[object]
 
 
 @dataclass
@@ -40,21 +47,89 @@ 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
 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,41 @@
+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 RequestResponseType, ListModelsResponse, ListModelsRequest, \
+    ModelDetailsResponse
+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) -> 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 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 @@
     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