Addressing review comments

Author: kumar-shivam-ranjan

ads/aqua/common/enums.py

@@ -52,8 +52,9 @@ class InferenceContainerTypeFamily(str, metaclass=ExtendedEnumMeta):
     AQUA_VLLM_CONTAINER_FAMILY = "odsc-vllm-serving"
     AQUA_TGI_CONTAINER_FAMILY = "odsc-tgi-serving"
     AQUA_LLAMA_CPP_CONTAINER_FAMILY = "odsc-llama-cpp-serving"
-    AQUA_TEI_CONTAINER_FAMILY = "odsc-tei-serving"
 
+class CustomInferenceContainerTypeFamily(str,metaclass=ExtendedEnumMeta):
+    AQUA_TEI_CONTAINER_FAMILY="odsc-tei-serving"
 
 class InferenceContainerParamType(str, metaclass=ExtendedEnumMeta):
     PARAM_TYPE_VLLM = "VLLM_PARAMS"

ads/aqua/extension/model_handler.py

@@ -8,7 +8,7 @@
 from tornado.web import HTTPError
 
 from ads.aqua.common.decorator import handle_exceptions
-from ads.aqua.common.enums import InferenceContainerTypeFamily
+from ads.aqua.common.enums import InferenceContainerTypeFamily,CustomInferenceContainerTypeFamily
 from ads.aqua.common.errors import AquaRuntimeError, AquaValueError
 from ads.aqua.common.utils import (
     get_hf_model_info,
@@ -166,11 +166,10 @@ def put(self, id):
         inference_container = input_data.get("inference_container")
         inference_container_uri = input_data.get("inference_container_uri")
         inference_containers = AquaModelApp.list_valid_inference_containers()
+        inference_containers.extend(CustomInferenceContainerTypeFamily.values())
         if (
             inference_container is not None
             and inference_container not in inference_containers
-            and inference_container
-            != InferenceContainerTypeFamily.AQUA_TEI_CONTAINER_FAMILY
         ):
             raise HTTPError(
                 400, Errors.INVALID_VALUE_OF_PARAMETER.format("inference_container")

ads/aqua/model/model.py

@@ -17,7 +17,7 @@
 from ads.aqua.common.enums import (
     FineTuningContainerTypeFamily,
     InferenceContainerTypeFamily,
-    Tags,
+    Tags, CustomInferenceContainerTypeFamily,
 )
 from ads.aqua.common.errors import AquaRuntimeError, AquaValueError
 from ads.aqua.common.utils import (
@@ -405,19 +405,18 @@ def edit_registered_model(
         if ds_model.freeform_tags.get(Tags.BASE_MODEL_CUSTOM, None):
             if ds_model.freeform_tags.get(Tags.AQUA_SERVICE_MODEL_TAG, None):
                 raise AquaRuntimeError(
-                    f"Failed to edit model:{id}. Only registered unverified models can be edited."
+                    f"Only registered unverified models can be edited."
                 )
             else:
                 custom_metadata_list = ds_model.custom_metadata_list
                 freeform_tags = ds_model.freeform_tags
                 if inference_container:
                     if (
-                        inference_container
-                        == InferenceContainerTypeFamily.AQUA_TEI_CONTAINER_FAMILY
+                            inference_container in CustomInferenceContainerTypeFamily.values()
                         and inference_container_uri is None
                     ):
                         raise AquaRuntimeError(
-                            f"Failed to edit model:{id}. Inference container URI must be provided."
+                            f"Inference container URI must be provided."
                         )
                     else:
                         custom_metadata_list.add(
@@ -429,8 +428,7 @@ def edit_registered_model(
                         )
                 if inference_container_uri:
                     if (
-                        inference_container
-                        == InferenceContainerTypeFamily.AQUA_TEI_CONTAINER_FAMILY
+                            inference_container in CustomInferenceContainerTypeFamily.values()
                         or inference_container is None
                     ):
                         custom_metadata_list.add(
@@ -442,7 +440,7 @@ def edit_registered_model(
                         )
                     else:
                         raise AquaRuntimeError(
-                            f"Failed to edit model:{id}. Inference container URI can be edited only with TEI container."
+                            f"Inference container URI can be edited only with container values: {CustomInferenceContainerTypeFamily.values()}"
                         )
 
                 if enable_finetuning is not None:
@@ -480,7 +478,7 @@ def edit_registered_model(
                 AquaApp().update_model(id, update_model_details)
         else:
             raise AquaRuntimeError(
-                f"Failed to edit model:{id}. Only registered unverified models can be edited."
+                f"Only registered unverified models can be edited."
             )
 
     def _fetch_metric_from_metadata(
@@ -900,8 +898,7 @@ def _create_model_catalog_entry(
             # only add cmd vars if inference container is not an SMC
             if (
                 inference_container not in smc_container_set
-                and inference_container
-                == InferenceContainerTypeFamily.AQUA_TEI_CONTAINER_FAMILY
+                and inference_container in CustomInferenceContainerTypeFamily.values()
             ):
                 cmd_vars = generate_tei_cmd_var(os_path)
                 metadata.add(