address API review comments (#36058)

* address comments

* reintroduce index

* retrigger checks

* lint fixes

* retrigger checks

* remove experimental tag

* investigating build issues

* revert experimental changes

* make dataclass kw-only

* commit docstring

* make index config experimental

* remove unused experimental import

* correct build_index docstring/typehint

* kevin's suggestion - remove redundant experimental warning

* make all init args keyword only

* re-order annotations

* run black

---------

Co-authored-by: Ubuntu <azureuser@adramadev0.1u2n51k150xetd4yrig4jsoeod.xx.internal.cloudapp.net>
Co-authored-by: Miles Holland <[email protected]>

Author: 3 people

sdk/ml/azure-ai-ml/azure/ai/ml/entities/_indexes/input/_ai_search_config.py

@@ -8,22 +8,24 @@
 # Defines stuff related to the resulting created index, like the index type.
 
 from typing import Optional
+from azure.ai.ml._utils._experimental import experimental
 
 
+@experimental
 class AzureAISearchConfig:
     """Config class for creating an Azure AI Search index.
 
-    :param ai_search_index_name: The name of the Azure AI Search index.
-    :type ai_search_index_name: Optional[str]
-    :param ai_search_connection_id: The Azure AI Search connection ID.
-    :type ai_search_connection_id: Optional[str]
+    :param index_name: The name of the Azure AI Search index.
+    :type index_name: Optional[str]
+    :param connection_id: The Azure AI Search connection ID.
+    :type connection_id: Optional[str]
     """
 
     def __init__(
         self,
         *,
-        ai_search_index_name: Optional[str] = None,
-        ai_search_connection_id: Optional[str] = None,
+        index_name: Optional[str] = None,
+        connection_id: Optional[str] = None,
     ) -> None:
-        self.ai_search_index_name = ai_search_index_name
-        self.ai_search_connection_id = ai_search_connection_id
+        self.index_name = index_name
+        self.connection_id = connection_id

sdk/ml/azure-ai-ml/azure/ai/ml/entities/_indexes/input/_index_data_source.py

@@ -3,6 +3,7 @@
 # ---------------------------------------------------------
 from typing import Union
 
+from azure.ai.ml._utils._experimental import experimental
 from azure.ai.ml.entities._inputs_outputs import Input
 from azure.ai.ml.constants._common import IndexInputType
 
@@ -12,6 +13,7 @@
 
 
 # Defines stuff related to supplying inputs for an index AKA the base data.
+@experimental
 class IndexDataSource:
     """Base class for configs that define data that will be processed into an ML index.
     This class should not be instantiated directly. Use one of its child classes instead.
@@ -28,24 +30,26 @@ def __init__(self, *, input_type: Union[str, IndexInputType]):
 # Field bundle for creating an index from files located in a Git repo.
 # TODO Does git_url need to specifically be an SSH or HTTPS style link?
 # TODO What is git connection id?
+@experimental
 class GitSource(IndexDataSource):
     """Config class for creating an ML index from files located in a git repository.
 
-    :param git_url: A link to the repository to use.
-    :type git_url: str
-    :param git_branch_name: The name of the branch to use from the target repository.
-    :type git_branch_name: str
-    :param git_connection_id: The connection ID for GitHub
-    :type git_connection_id: str
+    :param url: A link to the repository to use.
+    :type url: str
+    :param branch_name: The name of the branch to use from the target repository.
+    :type branch_name: str
+    :param connection_id: The connection ID for GitHub
+    :type connection_id: str
     """
 
-    def __init__(self, *, git_url: str, git_branch_name: str, git_connection_id: str):
-        self.git_url = git_url
-        self.git_branch_name = git_branch_name
-        self.git_connection_id = git_connection_id
+    def __init__(self, *, url: str, branch_name: str, connection_id: str):
+        self.url = url
+        self.branch_name = branch_name
+        self.connection_id = connection_id
         super().__init__(input_type=IndexInputType.GIT)
 
 
+@experimental
 class LocalSource(IndexDataSource):
     """Config class for creating an ML index from a collection of local files.
 

sdk/ml/azure-ai-ml/azure/ai/ml/entities/_indexes/model_config.py

@@ -3,6 +3,7 @@
 # ---------------------------------------------------------
 from dataclasses import dataclass
 from typing import Any, Dict, Optional
+from azure.ai.ml._utils._experimental import experimental
 from azure.ai.ml._utils.utils import camel_to_snake
 from azure.ai.ml.entities._workspace.connections.workspace_connection import WorkspaceConnection
 from azure.ai.ml.entities._workspace.connections.connection_subtypes import (
@@ -11,6 +12,7 @@
 )
 
 
+@experimental
 @dataclass
 class ModelConfiguration:
     """Configuration for a embedding model.
@@ -42,12 +44,33 @@ class ModelConfiguration:
     deployment_name: Optional[str]
     model_kwargs: Dict[str, Any]
 
+    def __init__(
+        self,
+        *,
+        api_base: Optional[str],
+        api_key: Optional[str],
+        api_version: Optional[str],
+        connection_name: Optional[str],
+        connection_type: Optional[str],
+        model_name: Optional[str],
+        deployment_name: Optional[str],
+        model_kwargs: Dict[str, Any]
+    ):
+        self.api_base = api_base
+        self.api_key = api_key
+        self.api_version = api_version
+        self.connection_name = connection_name
+        self.connection_type = connection_type
+        self.model_name = model_name
+        self.deployment_name = deployment_name
+        self.model_kwargs = model_kwargs
+
     @staticmethod
     def from_connection(
         connection: WorkspaceConnection,
         model_name: Optional[str] = None,
         deployment_name: Optional[str] = None,
-        **model_kwargs
+        **kwargs
     ) -> "ModelConfiguration":
         """Create an model configuration from a Connection.
 
@@ -57,8 +80,8 @@ def from_connection(
         :type model_name: Optional[str]
         :param deployment_name: The name of the deployment.
         :type deployment_name: Optional[str]
-        :keyword model_kwargs: Additional keyword arguments for the model.
-        :paramtype model_kwargs: Dict[str, Any]
+        :keyword kwargs: Additional keyword arguments for the model.
+        :paramtype kwargs: Dict[str, Any]
         :return: The model configuration.
         :rtype: ~azure.ai.ml.entities._indexes.entities.ModelConfiguration
         :raises TypeError: If the connection is not an AzureOpenAIConnection.
@@ -97,5 +120,5 @@ def from_connection(
             connection_type=connection_type,
             model_name=model_name,
             deployment_name=deployment_name,
-            model_kwargs=model_kwargs,
+            model_kwargs=kwargs,
         )

sdk/ml/azure-ai-ml/azure/ai/ml/operations/_index_operations.py

@@ -270,7 +270,7 @@ def build_index(
         ######## data source info ########
         input_source: Union[IndexDataSource, str],
         input_source_credential: Optional[Union[ManagedIdentityConfiguration, UserIdentityConfiguration]] = None,
-    ) -> Union["MLIndex", "Job"]:  # type: ignore[name-defined]
+    ) -> Union["Index", "Job"]:  # type: ignore[name-defined]
         """Builds an index on the cloud using the Azure AI Resources service.
 
         :keyword name: The name of the index to be created.
@@ -295,7 +295,7 @@ def build_index(
         :paramtype input_source_credential: Optional[Union[~azure.ai.ml.entities.ManagedIdentityConfiguration,
             ~azure.ai.ml.entities.UserIdentityConfiguration]]
         :return: If the `source_input` is a GitSource, returns a created DataIndex Job object.
-        :rtype: Union[~azure.ai.ml.entities._indexes.MLIndex, ~azure.ai.ml.entities.Job]
+        :rtype: Union[~azure.ai.ml.entities.Index, ~azure.ai.ml.entities.Job]
         :raises ValueError: If the `source_input` is not type ~typing.Str or
             ~azure.ai.ml.entities._indexes.LocalSource.
         """
@@ -333,8 +333,8 @@ def build_index(
             index=(
                 IndexStore(
                     type="acs",
-                    connection=build_connection_id(index_config.ai_search_connection_id, self._operation_scope),
-                    name=index_config.ai_search_index_name,
+                    connection=build_connection_id(index_config.connection_id, self._operation_scope),
+                    name=index_config.index_name,
                 )
                 if index_config is not None
                 else IndexStore(type="faiss")
@@ -390,9 +390,9 @@ def git_to_index(
                 return index_job.outputs
 
             git_index_job = git_to_index(
-                git_url=input_source.git_url,
-                branch_name=input_source.git_branch_name,
-                git_connection_id=input_source.git_connection_id,
+                git_url=input_source.url,
+                branch_name=input_source.branch_name,
+                git_connection_id=input_source.connection_id,
             )
             # Ensure repo cloned each run to get latest, comment out to have first clone reused.
             git_index_job.settings.force_rerun = True