chore: check Fabric Tenant

cognitive/src/main/scala/com/microsoft/azure/synapse/ml/services/CognitiveServiceBase.scala

@@ -6,7 +6,7 @@ package com.microsoft.azure.synapse.ml.services
 import com.microsoft.azure.synapse.ml.codegen.Wrappable
 import com.microsoft.azure.synapse.ml.core.contracts.HasOutputCol
 import com.microsoft.azure.synapse.ml.core.schema.DatasetExtensions
-import com.microsoft.azure.synapse.ml.fabric.{FabricClient, TokenLibrary}
+import com.microsoft.azure.synapse.ml.fabric.{FabricClient, TokenLibrary, FabricTenantSetting}
 import com.microsoft.azure.synapse.ml.io.http._
 import com.microsoft.azure.synapse.ml.logging.SynapseMLLogging
 import com.microsoft.azure.synapse.ml.logging.common.PlatformDetails
@@ -219,7 +219,18 @@ trait HasCustomCogServiceDomain extends Wrappable with HasURL with HasUrlPath {
       |    self._java_obj = self._java_obj.setDefaultInternalEndpoint(value)
       |    return self
       |
+      |def checkUsingDefaultModelEndpoint(self):
+      |    from synapse.ml.fabric.service_discovery import get_fabric_env_config
+      |    url = get_fabric_env_config().fabric_env_config.get_mlflow_workload_endpoint()  + "/cognitive/openai/"
+      |    return url == self._java_obj.getUrl()
+      |
       |def _transform(self, dataset: DataFrame) -> DataFrame:
+      |    if running_on_synapse_internal():
+      |       if self.checkUsingDefaultModelEndpoint():
+      |          deployment_name = self._java_obj.getDeploymentName()
+      |          if not self._java_obj.getModelStatus(deployment_name):
+      |              raise PermissionError("Please check fabric LLM setting or switch to Azure OpenAI endpoint")
+      |
       |    return super()._transform(dataset)
       |""".stripMargin
   }
@@ -497,7 +508,7 @@ abstract class CognitiveServicesBaseNoHandler(val uid: String) extends Transform
   with HasURL with ComplexParamsWritable
   with HasSubscriptionKey with HasErrorCol
   with HasAADToken with HasCustomCogServiceDomain
-  with SynapseMLLogging {
+  with SynapseMLLogging with FabricTenantSetting {
 
   setDefault(
     outputCol -> (this.uid + "_output"),

core/src/main/scala/com/microsoft/azure/synapse/ml/fabric/OpenAIFabricSetting.scala

@@ -0,0 +1,48 @@
+package com.microsoft.azure.synapse.ml.fabric
+
+import spray.json.{JsValue, JsString}
+
+trait FabricTenantSetting extends RESTUtils {
+
+  private def getHeaders: Map[String, String] = {
+    Map(
+      "Authorization" -> s"Bearer ${TokenLibrary.getAccessToken}",
+      "Content-Type" -> "application/json"
+    )
+  }
+
+  def usagePost(url: String, body: String): JsValue = {
+    usagePost(url, body, getHeaders);
+  }
+
+  def getModelStatus(modelName: String): Boolean = {
+
+    val payload =
+      s"""["${modelName}"]"""
+
+    val mlWorkloadEndpointML = FabricClient.MLWorkloadEndpointML
+    val url = mlWorkloadEndpointML + "cognitive/openai/tenantsetting"
+    val modelStatus = usagePost(url, payload).asJsObject.fields.get(modelName.toLowerCase).get
+
+    // Allowed, Disallowed, DisallowedForCrossGeo, ModelNotFound, InvalidResult
+    val resultString: String = modelStatus match {
+      case JsString(value) => value
+    }
+
+    resultString match {
+      case "Disallowed" => throw new Exception(s"Default OpenAI model ${modelName} is Disallowed, " +
+        s"please contact your admin if you want to use default fabric LLM model. " +
+        s"Or you can set your Azure OpenAI credentials.")
+      case "DisallowedForCrossGeo" => throw new Exception(s"Default OpenAI model ${modelName} is Disallowed " +
+        s"for Cross Geo, please contact your admin if you want to use default fabric LLM model. " +
+        s"Or you can set your Azure OpenAI credentials.")
+      case "ModelNotFound" => throw new Exception(s"Default OpenAI model ${modelName} not found, " +
+        s"please check your deployment name.")
+      case "InvalidResult" => throw new Exception("Cannot get tenant admin setting status correctly")
+      case _ => ()
+    }
+
+    resultString == "Allowed"
+  }
+
+}