[ODSC-55513] Show Aqua extension for the default region (#778)

Author: VipulMascarenhas

ads/aqua/extension/common_handler.py

@@ -7,8 +7,9 @@
 from importlib import metadata
 
 from ads.aqua import ODSC_MODEL_COMPARTMENT_OCID
-from ads.aqua.decorator import handle_exceptions
 from ads.aqua.exception import AquaResourceAccessError
+from ads.aqua.utils import known_realm
+from ads.aqua.decorator import handle_exceptions
 from ads.aqua.extension.base_handler import AquaAPIhandler
 
 
@@ -25,8 +26,23 @@ class CompatibilityCheckHandler(AquaAPIhandler):
 
     @handle_exceptions
     def get(self):
+        """This method provides the availability status of Aqua. If ODSC_MODEL_COMPARTMENT_OCID environment variable
+        is set, then status `ok` is returned. For regions where Aqua is available but the environment variable is not
+        set due to accesses/policies, we return the `compatible` status to indicate that the extension can be enabled
+        for the selected notebook session.
+
+        Returns
+        -------
+            status dict:
+                ok or compatible
+        Raises:
+            AquaResourceAccessError: raised when aqua is not accessible in the given session/region.
+
+        """
         if ODSC_MODEL_COMPARTMENT_OCID:
             return self.finish(dict(status="ok"))
+        elif known_realm():
+            return self.finish(dict(status="compatible"))
         else:
             raise AquaResourceAccessError(
                 f"The AI Quick actions extension is not compatible in the given region."

ads/aqua/utils.py

@@ -79,6 +79,7 @@
 LIFECYCLE_DETAILS_MISSING_JOBRUN = "The asscociated JobRun resource has been deleted."
 READY_TO_DEPLOY_STATUS = "ACTIVE"
 READY_TO_FINE_TUNE_STATUS = "TRUE"
+AQUA_GA_LIST = ["id19sfcrra6z"]
 
 
 class LifecycleStatus(Enum):
@@ -733,3 +734,14 @@ def _is_valid_mvs(mvs: ModelVersionSet, target_tag: str) -> bool:
         return False
 
     return target_tag in mvs.freeform_tags
+
+
+def known_realm():
+    """This helper function returns True if the Aqua service is available by default in the given namespace.
+    Returns
+    -------
+    bool:
+        Return True if aqua service is available.
+
+    """
+    return os.environ.get("CONDA_BUCKET_NS") in AQUA_GA_LIST

tests/unitary/with_extras/aqua/test_common_handler.py

@@ -0,0 +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 os
+import json
+from importlib import reload
+from tornado.web import Application
+from tornado.testing import AsyncHTTPTestCase
+
+import ads.config
+import ads.aqua
+from ads.aqua.utils import AQUA_GA_LIST
+from ads.aqua.extension.common_handler import CompatibilityCheckHandler
+
+
+class TestDataset:
+    SERVICE_COMPARTMENT_ID = "ocid1.compartment.oc1..<OCID>"
+
+
+class TestYourHandler(AsyncHTTPTestCase):
+    def get_app(self):
+        return Application([(r"/hello", CompatibilityCheckHandler)])
+
+    def setUp(self):
+        super().setUp()
+        os.environ["ODSC_MODEL_COMPARTMENT_OCID"] = TestDataset.SERVICE_COMPARTMENT_ID
+        reload(ads.config)
+        reload(ads.aqua)
+        reload(ads.aqua.extension.common_handler)
+
+    def tearDown(self):
+        super().tearDown()
+        os.environ.pop("ODSC_MODEL_COMPARTMENT_OCID", None)
+        reload(ads.config)
+        reload(ads.aqua)
+        reload(ads.aqua.extension.common_handler)
+
+    def test_get_ok(self):
+        response = self.fetch("/hello", method="GET")
+        assert json.loads(response.body)["status"] == "ok"
+
+    def test_get_compatible_status(self):
+        os.environ["ODSC_MODEL_COMPARTMENT_OCID"] = ""
+        os.environ["CONDA_BUCKET_NS"] = AQUA_GA_LIST[0]
+        reload(ads.common)
+        reload(ads.aqua)
+        reload(ads.aqua.extension.common_handler)
+        response = self.fetch("/hello", method="GET")
+        assert json.loads(response.body)["status"] == "compatible"
+
+    def test_raise_not_compatible_error(self):
+        os.environ["ODSC_MODEL_COMPARTMENT_OCID"] = ""
+        os.environ["CONDA_BUCKET_NS"] = "test-namespace"
+        reload(ads.common)
+        reload(ads.aqua)
+        reload(ads.aqua.extension.common_handler)
+        response = self.fetch("/hello", method="GET")
+        body = json.loads(response.body)
+        assert body["status"] == 404
+        assert (
+            body["reason"]
+            == "The AI Quick actions extension is not compatible in the given region."
+        )