Addressing review comments

Author: kumar-shivam-ranjan

ads/aqua/config/container_config.py

@@ -203,12 +203,13 @@ def from_service_config(
                     {
                         "PORT": container.workload_configuration_details_list[
                             0
-                        ].additional_configurations.get("PORT", ""),
+                        ].additional_configurations.get("PORT", "")
+                    },
+                    {
                         "HEALTH_CHECK_PORT": container.workload_configuration_details_list[
                             0
                         ].additional_configurations.get("HEALTH_CHECK_PORT", UNKNOWN),
                     },
-                    {},
                 ]
                 container_spec = AquaContainerConfigSpec(
                     cli_param=container.workload_configuration_details_list[0].cmd,

ads/aqua/extension/common_ws_msg_handler.py

@@ -5,13 +5,12 @@
 
 import json
 from importlib import metadata
-from typing import List, Union
+from typing import List, Optional, 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 (
     AdsVersionResponse,
-    CompatibilityCheckResponse,
     RequestResponseType,
 )
 
@@ -25,7 +24,7 @@ def __init__(self, message: Union[str, bytes]):
         super().__init__(message)
 
     @handle_exceptions
-    def process(self) -> Union[AdsVersionResponse, CompatibilityCheckResponse]:
+    def process(self) -> Optional[AdsVersionResponse]:
         request = json.loads(self.message)
         if request.get("kind") == "AdsVersion":
             version = metadata.version("oracle_ads")
@@ -35,9 +34,3 @@ def process(self) -> Union[AdsVersionResponse, CompatibilityCheckResponse]:
                 data=version,
             )
             return response
-        if request.get("kind") == "CompatibilityCheck":
-            return CompatibilityCheckResponse(
-                message_id=request.get("message_id"),
-                kind=RequestResponseType.CompatibilityCheck,
-                data={"status": "ok"},
-            )

ads/aqua/extension/model_handler.py

@@ -15,7 +15,7 @@
 from ads.aqua.extension.errors import Errors
 from ads.aqua.model import AquaModelApp
 from ads.aqua.model.entities import AquaModelSummary, HFModelSummary
-from ads.config import USER
+from ads.config import SERVICE
 from ads.model.common.utils import MetadataArtifactPathType
 
 
@@ -82,7 +82,7 @@ def list(self):
         # project_id is no needed.
         project_id = self.get_argument("project_id", default=None)
         model_type = self.get_argument("model_type", default=None)
-        category = self.get_argument("category", default=USER)
+        category = self.get_argument("category", default=SERVICE)
         return self.finish(
             AquaModelApp().list(
                 compartment_id=compartment_id,

ads/aqua/model/model.py

@@ -185,7 +185,7 @@ def create(
         target_compartment = compartment_id or COMPARTMENT_OCID
 
         # Skip model copying if it is registered model
-        if service_model.freeform_tags.get(Tags.BASE_MODEL_CUSTOM) is not None:
+        if service_model.freeform_tags.get(Tags.BASE_MODEL_CUSTOM, None) is not None:
             logger.info(
                 f"Aqua Model {model_id} already exists in the user's compartment."
                 "Skipped copying."
@@ -919,12 +919,13 @@ def _process_model(
     def list(
         self,
         compartment_id: str = None,
+        category: str = None,
         project_id: str = None,
         model_type: str = None,
         **kwargs,
     ) -> List["AquaModelSummary"]:
         """Lists all Aqua models within a specified compartment and/or project.
-        If `compartment_id` is not specified, the method defaults to returning
+        If `category` is not specified, the method defaults to returning
         the service models within the pre-configured default compartment. By default, the list
         of models in the service compartment are cached. Use clear_model_list_cache() to invalidate
         the cache.
@@ -933,6 +934,8 @@ def list(
         ----------
         compartment_id: (str, optional). Defaults to `None`.
             The compartment OCID.
+        category: (str,optional). Defaults to `SERVICE`
+            The category of the models to fetch. Can be either `USER` or `SERVICE`
         project_id: (str, optional). Defaults to `None`.
             The project OCID.
         model_type: (str, optional). Defaults to `None`.
@@ -946,9 +949,9 @@ def list(
             The list of the `ads.aqua.model.AquaModelSummary`.
         """
 
-        models = []
-        category = kwargs.pop("category", USER)
-        if compartment_id and category != SERVICE:
+        category = category or kwargs.pop("category", SERVICE)
+        compartment_id = compartment_id or COMPARTMENT_OCID
+        if category == USER:
             # tracks number of times custom model listing was called
             self.telemetry.record_event_async(
                 category="aqua/custom/model", action="list"
@@ -957,33 +960,32 @@ def list(
             logger.info(f"Fetching custom models from compartment_id={compartment_id}.")
             model_type = model_type.upper() if model_type else ModelType.FT
             models = self._rqs(compartment_id, model_type=model_type)
+            logger.info(
+                f"Fetched {len(models)} models from {compartment_id or COMPARTMENT_OCID}."
+            )
         else:
             # tracks number of times service model listing was called
             self.telemetry.record_event_async(
                 category="aqua/service/model", action="list"
             )
 
             if AQUA_SERVICE_MODELS in self._service_models_cache:
-                logger.info(
-                    f"Returning service models list in {AQUA_SERVICE_MODELS} from cache."
-                )
+                logger.info("Returning service models list from cache.")
                 return self._service_models_cache.get(AQUA_SERVICE_MODELS)
-            logger.info("Fetching service models.")
+            logger.info("Fetching service models from cache.")
             lifecycle_state = kwargs.pop(
                 "lifecycle_state", Model.LIFECYCLE_STATE_ACTIVE
             )
 
             models = self.list_resource(
                 self.ds_client.list_models,
-                compartment_id=compartment_id or COMPARTMENT_OCID,
+                compartment_id=compartment_id,
                 lifecycle_state=lifecycle_state,
                 category=category,
                 **kwargs,
             )
+            logger.info(f"Fetched {len(models)} service models.")
 
-        logger.info(
-            f"Fetched {len(models)} models from {AQUA_SERVICE_MODELS if category==SERVICE else compartment_id}."
-        )
         aqua_models = []
         inference_containers = self.get_container_config().to_dict().get("inference")
         for model in models:
@@ -997,7 +999,6 @@ def list(
                     project_id=project_id or UNKNOWN,
                 )
             )
-
         if category == SERVICE:
             self._service_models_cache.__setitem__(
                 key=AQUA_SERVICE_MODELS, value=aqua_models

tests/unitary/with_extras/aqua/test_model.py

@@ -1483,7 +1483,7 @@ def test_list_custom_models(self, mock_get_container_config):
             ]
         )
 
-        results = self.app.list(TestDataset.COMPARTMENT_ID)
+        results = self.app.list(TestDataset.COMPARTMENT_ID, category=ads.config.USER)
 
         self.app._rqs.assert_called_with(TestDataset.COMPARTMENT_ID, model_type="FT")
 

tests/unitary/with_extras/aqua/test_model_handler.py

@@ -27,7 +27,7 @@
 )
 from ads.aqua.model import AquaModelApp
 from ads.aqua.model.entities import AquaModel, AquaModelSummary, HFModelSummary
-from ads.config import USER
+from ads.config import USER, SERVICE
 
 
 class ModelHandlerTestCase(TestCase):
@@ -134,7 +134,7 @@ def test_list(self, mock_list):
             mock_finish.side_effect = lambda x: x
             self.model_handler.list()
             mock_list.assert_called_with(
-                compartment_id=None, project_id=None, model_type=None, category=USER
+                compartment_id=None, project_id=None, model_type=None, category=SERVICE
             )
 
     @parameterized.expand(

tests/unitary/with_extras/aqua/test_ui.py

@@ -541,7 +541,6 @@ def test_list_containers(self, mock_list_service_containers):
         mock_list_service_containers.return_value = TestDataset.CONTAINERS_LIST
 
         test_result = self.app.list_containers()
-        print("test_result: ", test_result)
         expected_result = {
             "evaluate": [
                 {
@@ -566,13 +565,11 @@ def test_list_containers(self, mock_list_service_containers):
                     "spec": {
                         "cli_param": "--served-model-name odsc-llm --disable-custom-all-reduce --seed 42 ",
                         "env_vars": [
-                            {
-                                "HEALTH_CHECK_PORT": "8080",
-                                "MODEL_DEPLOY_ENABLE_STREAMING": "true",
-                                "MODEL_DEPLOY_HEALTH_ENDPOINT": "",
-                                "MODEL_DEPLOY_PREDICT_ENDPOINT": "/v1/completions",
-                                "PORT": "8080",
-                            }
+                            {"MODEL_DEPLOY_PREDICT_ENDPOINT": "/v1/completions"},
+                            {"MODEL_DEPLOY_HEALTH_ENDPOINT": ""},
+                            {"MODEL_DEPLOY_ENABLE_STREAMING": "true"},
+                            {"PORT": "8080"},
+                            {"HEALTH_CHECK_PORT": "8080"},
                         ],
                         "health_check_port": "8080",
                         "restricted_params": [