two version of R2R are hereHEADmaster

1 files changed, 483 insertions, 0 deletions

diff --git a/.venv/lib/python3.12/site-packages/azure/ai/ml/operations/_index_operations.py b/.venv/lib/python3.12/site-packages/azure/ai/ml/operations/_index_operations.py
new file mode 100644
index 00000000..28e409c7
--- /dev/null
+++ b/.venv/lib/python3.12/site-packages/azure/ai/ml/operations/_index_operations.py
@@ -0,0 +1,483 @@
+# ---------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# ---------------------------------------------------------
+
+# pylint: disable=protected-access
+import json
+from typing import Any, Callable, Dict, Iterable, List, Optional, Union
+
+from azure.ai.ml._artifacts._artifact_utilities import _check_and_upload_path
+
+# cspell:disable-next-line
+from azure.ai.ml._restclient.azure_ai_assets_v2024_04_01.azureaiassetsv20240401 import (
+    MachineLearningServicesClient as AzureAiAssetsClient042024,
+)
+
+# cspell:disable-next-line
+from azure.ai.ml._restclient.azure_ai_assets_v2024_04_01.azureaiassetsv20240401.models import Index as RestIndex
+from azure.ai.ml._restclient.v2023_04_01_preview.models import ListViewType
+from azure.ai.ml._scope_dependent_operations import (
+    OperationConfig,
+    OperationsContainer,
+    OperationScope,
+    _ScopeDependentOperations,
+)
+from azure.ai.ml._telemetry import ActivityType, monitor_with_activity
+from azure.ai.ml._utils._asset_utils import (
+    _resolve_label_to_asset,
+    _validate_auto_delete_setting_in_data_output,
+    _validate_workspace_managed_datastore,
+)
+from azure.ai.ml._utils._http_utils import HttpPipeline
+from azure.ai.ml._utils._logger_utils import OpsLogger
+from azure.ai.ml._utils.utils import _get_base_urls_from_discovery_service
+from azure.ai.ml.constants._common import AssetTypes, AzureMLResourceType, WorkspaceDiscoveryUrlKey
+from azure.ai.ml.dsl import pipeline
+from azure.ai.ml.entities import PipelineJob, PipelineJobSettings
+from azure.ai.ml.entities._assets import Index
+from azure.ai.ml.entities._credentials import ManagedIdentityConfiguration, UserIdentityConfiguration
+from azure.ai.ml.entities._indexes import (
+    AzureAISearchConfig,
+    GitSource,
+    IndexDataSource,
+    LocalSource,
+    ModelConfiguration,
+)
+from azure.ai.ml.entities._indexes.data_index_func import index_data as index_data_func
+from azure.ai.ml.entities._indexes.entities.data_index import (
+    CitationRegex,
+    Data,
+    DataIndex,
+    Embedding,
+    IndexSource,
+    IndexStore,
+)
+from azure.ai.ml.entities._indexes.utils._open_ai_utils import build_connection_id, build_open_ai_protocol
+from azure.ai.ml.entities._inputs_outputs import Input
+from azure.ai.ml.exceptions import ErrorCategory, ErrorTarget, ValidationErrorType, ValidationException
+from azure.ai.ml.operations._datastore_operations import DatastoreOperations
+from azure.core.credentials import TokenCredential
+
+ops_logger = OpsLogger(__name__)
+module_logger = ops_logger.module_logger
+
+
+class IndexOperations(_ScopeDependentOperations):
+    """Represents a client for performing operations on index assets.
+
+    You should not instantiate this class directly. Instead, you should create MLClient and use this client via the
+    property MLClient.index
+    """
+
+    def __init__(
+        self,
+        *,
+        operation_scope: OperationScope,
+        operation_config: OperationConfig,
+        credential: TokenCredential,
+        datastore_operations: DatastoreOperations,
+        all_operations: OperationsContainer,
+        **kwargs: Any,
+    ):
+        super().__init__(operation_scope, operation_config)
+        ops_logger.update_filter()
+        self._credential = credential
+        # Dataplane service clients are lazily created as they are needed
+        self.__azure_ai_assets_client: Optional[AzureAiAssetsClient042024] = None
+        # Kwargs to propagate to dataplane service clients
+        self._service_client_kwargs: Dict[str, Any] = kwargs.pop("_service_client_kwargs", {})
+        self._all_operations = all_operations
+
+        self._datastore_operation: DatastoreOperations = datastore_operations
+        self._requests_pipeline: HttpPipeline = kwargs.pop("requests_pipeline")
+
+        # Maps a label to a function which given an asset name,
+        # returns the asset associated with the label
+        self._managed_label_resolver: Dict[str, Callable[[str], Index]] = {"latest": self._get_latest_version}
+
+    @property
+    def _azure_ai_assets(self) -> AzureAiAssetsClient042024:
+        """Lazily instantiated client for azure ai assets api.
+
+        .. note::
+
+            Property is lazily instantiated since the api's base url depends on the user's workspace, and
+            must be fetched remotely.
+        """
+        if self.__azure_ai_assets_client is None:
+            workspace_operations = self._all_operations.all_operations[AzureMLResourceType.WORKSPACE]
+
+            endpoint = _get_base_urls_from_discovery_service(
+                workspace_operations, self._operation_scope.workspace_name, self._requests_pipeline
+            )[WorkspaceDiscoveryUrlKey.API]
+
+            self.__azure_ai_assets_client = AzureAiAssetsClient042024(
+                endpoint=endpoint,
+                subscription_id=self._operation_scope.subscription_id,
+                resource_group_name=self._operation_scope.resource_group_name,
+                workspace_name=self._operation_scope.workspace_name,
+                credential=self._credential,
+                **self._service_client_kwargs,
+            )
+
+        return self.__azure_ai_assets_client
+
+    @monitor_with_activity(ops_logger, "Index.CreateOrUpdate", ActivityType.PUBLICAPI)
+    def create_or_update(self, index: Index, **kwargs) -> Index:
+        """Returns created or updated index asset.
+
+        If not already in storage, asset will be uploaded to the workspace's default datastore.
+
+        :param index: Index asset object.
+        :type index: Index
+        :return: Index asset object.
+        :rtype: ~azure.ai.ml.entities.Index
+        """
+
+        if not index.name:
+            msg = "Must specify a name."
+            raise ValidationException(
+                message=msg,
+                target=ErrorTarget.INDEX,
+                no_personal_data_message=msg,
+                error_category=ErrorCategory.USER_ERROR,
+                error_type=ValidationErrorType.MISSING_FIELD,
+            )
+
+        if not index.version:
+            if not index._auto_increment_version:
+                msg = "Must specify a version."
+                raise ValidationException(
+                    message=msg,
+                    target=ErrorTarget.INDEX,
+                    no_personal_data_message=msg,
+                    error_category=ErrorCategory.USER_ERROR,
+                    error_type=ValidationErrorType.MISSING_FIELD,
+                )
+
+            next_version = self._azure_ai_assets.indexes.get_next_version(index.name).next_version
+
+            if next_version is None:
+                msg = "Version not specified, could not automatically increment version. Set a version to resolve."
+                raise ValidationException(
+                    message=msg,
+                    target=ErrorTarget.INDEX,
+                    no_personal_data_message=msg,
+                    error_category=ErrorCategory.SYSTEM_ERROR,
+                    error_type=ValidationErrorType.MISSING_FIELD,
+                )
+
+            index.version = str(next_version)
+
+        _ = _check_and_upload_path(
+            artifact=index,
+            asset_operations=self,
+            datastore_name=index.datastore,
+            artifact_type=ErrorTarget.INDEX,
+            show_progress=self._show_progress,
+        )
+
+        return Index._from_rest_object(
+            self._azure_ai_assets.indexes.create_or_update(
+                name=index.name, version=index.version, body=index._to_rest_object(), **kwargs
+            )
+        )
+
+    @monitor_with_activity(ops_logger, "Index.Get", ActivityType.PUBLICAPI)
+    def get(self, name: str, *, version: Optional[str] = None, label: Optional[str] = None, **kwargs) -> Index:
+        """Returns information about the specified index asset.
+
+        :param str name: Name of the index asset.
+        :keyword Optional[str] version: Version of the index asset. Mutually exclusive with `label`.
+        :keyword Optional[str] label: The label of the index asset. Mutually exclusive with  `version`.
+        :raises ~azure.ai.ml.exceptions.ValidationException: Raised if Index cannot be successfully validated.
+            Details will be provided in the error message.
+        :return: Index asset object.
+        :rtype: ~azure.ai.ml.entities.Index
+        """
+        if version and label:
+            msg = "Cannot specify both version and label."
+            raise ValidationException(
+                message=msg,
+                target=ErrorTarget.INDEX,
+                no_personal_data_message=msg,
+                error_category=ErrorCategory.USER_ERROR,
+                error_type=ValidationErrorType.INVALID_VALUE,
+            )
+
+        if label:
+            return _resolve_label_to_asset(self, name, label)
+
+        if not version:
+            msg = "Must provide either version or label."
+            raise ValidationException(
+                message=msg,
+                target=ErrorTarget.INDEX,
+                no_personal_data_message=msg,
+                error_category=ErrorCategory.USER_ERROR,
+                error_type=ValidationErrorType.MISSING_FIELD,
+            )
+
+        index_version_resource = self._azure_ai_assets.indexes.get(name=name, version=version, **kwargs)
+
+        return Index._from_rest_object(index_version_resource)
+
+    def _get_latest_version(self, name: str) -> Index:
+        return Index._from_rest_object(self._azure_ai_assets.indexes.get_latest(name))
+
+    @monitor_with_activity(ops_logger, "Index.List", ActivityType.PUBLICAPI)
+    def list(
+        self, name: Optional[str] = None, *, list_view_type: ListViewType = ListViewType.ACTIVE_ONLY, **kwargs
+    ) -> Iterable[Index]:
+        """List all Index assets in workspace.
+
+        If an Index is specified by name, list all version of that Index.
+
+        :param name: Name of the model.
+        :type name: Optional[str]
+        :keyword list_view_type: View type for including/excluding (for example) archived models.
+            Defaults to :attr:`ListViewType.ACTIVE_ONLY`.
+        :paramtype list_view_type: ListViewType
+        :return: An iterator like instance of Index objects
+        :rtype: ~azure.core.paging.ItemPaged[~azure.ai.ml.entities.Index]
+        """
+
+        def cls(rest_indexes: Iterable[RestIndex]) -> List[Index]:
+            return [Index._from_rest_object(i) for i in rest_indexes]
+
+        if name is None:
+            return self._azure_ai_assets.indexes.list_latest(cls=cls, **kwargs)
+
+        return self._azure_ai_assets.indexes.list(name, list_view_type=list_view_type, cls=cls, **kwargs)
+
+    def build_index(
+        self,
+        *,
+        ######## required args ##########
+        name: str,
+        embeddings_model_config: ModelConfiguration,
+        ######## chunking information ##########
+        data_source_citation_url: Optional[str] = None,
+        tokens_per_chunk: Optional[int] = None,
+        token_overlap_across_chunks: Optional[int] = None,
+        input_glob: Optional[str] = None,
+        ######## other generic args ########
+        document_path_replacement_regex: Optional[str] = None,
+        ######## Azure AI Search index info ########
+        index_config: Optional[AzureAISearchConfig] = None,  # todo better name?
+        ######## data source info ########
+        input_source: Union[IndexDataSource, str],
+        input_source_credential: Optional[Union[ManagedIdentityConfiguration, UserIdentityConfiguration]] = None,
+    ) -> 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.
+        :paramtype name: str
+        :keyword embeddings_model_config: Model config for the embedding model.
+        :paramtype embeddings_model_config: ~azure.ai.ml.entities._indexes.ModelConfiguration
+        :keyword data_source_citation_url: The URL of the data source.
+        :paramtype data_source_citation_url: Optional[str]
+        :keyword tokens_per_chunk: The size of chunks to be used for indexing.
+        :paramtype tokens_per_chunk: Optional[int]
+        :keyword token_overlap_across_chunks: The amount of overlap between chunks.
+        :paramtype token_overlap_across_chunks: Optional[int]
+        :keyword input_glob: The glob pattern to be used for indexing.
+        :paramtype input_glob: Optional[str]
+        :keyword document_path_replacement_regex: The regex pattern for replacing document paths.
+        :paramtype document_path_replacement_regex: Optional[str]
+        :keyword index_config: The configuration for the ACS output.
+        :paramtype index_config: Optional[~azure.ai.ml.entities._indexes.AzureAISearchConfig]
+        :keyword input_source: The input source for the index.
+        :paramtype input_source: Union[~azure.ai.ml.entities._indexes.IndexDataSource, str]
+        :keyword input_source_credential: The identity to be used for the 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.Index, ~azure.ai.ml.entities.Job]
+        :raises ValueError: If the `source_input` is not type ~typing.Str or
+            ~azure.ai.ml.entities._indexes.LocalSource.
+        """
+        if document_path_replacement_regex:
+            document_path_replacement_regex = json.loads(document_path_replacement_regex)
+
+        data_index = DataIndex(
+            name=name,
+            source=IndexSource(
+                input_data=Data(
+                    type="uri_folder",
+                    path=".",
+                ),
+                input_glob=input_glob,
+                chunk_size=tokens_per_chunk,
+                chunk_overlap=token_overlap_across_chunks,
+                citation_url=data_source_citation_url,
+                citation_url_replacement_regex=(
+                    CitationRegex(
+                        match_pattern=document_path_replacement_regex["match_pattern"],  # type: ignore[index]
+                        replacement_pattern=document_path_replacement_regex[
+                            "replacement_pattern"  # type: ignore[index]
+                        ],
+                    )
+                    if document_path_replacement_regex
+                    else None
+                ),
+            ),
+            embedding=Embedding(
+                model=build_open_ai_protocol(
+                    model=embeddings_model_config.model_name, deployment=embeddings_model_config.deployment_name
+                ),
+                connection=build_connection_id(embeddings_model_config.connection_name, self._operation_scope),
+            ),
+            index=(
+                IndexStore(
+                    type="acs",
+                    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")
+            ),
+            # name is replaced with a unique value each time the job is run
+            path=f"azureml://datastores/workspaceblobstore/paths/indexes/{name}/{{name}}",
+        )
+
+        if isinstance(input_source, LocalSource):
+            data_index.source.input_data = Data(
+                type="uri_folder",
+                path=input_source.input_data.path,
+            )
+
+            return self._create_data_indexing_job(data_index=data_index, identity=input_source_credential)
+
+        if isinstance(input_source, str):
+            data_index.source.input_data = Data(
+                type="uri_folder",
+                path=input_source,
+            )
+
+            return self._create_data_indexing_job(data_index=data_index, identity=input_source_credential)
+
+        if isinstance(input_source, GitSource):
+            from azure.ai.ml import MLClient
+
+            ml_registry = MLClient(credential=self._credential, registry_name="azureml")
+            git_clone_component = ml_registry.components.get("llm_rag_git_clone", label="latest")
+
+            # Clone Git Repo and use as input to index_job
+            @pipeline(default_compute="serverless")  # type: ignore[call-overload]
+            def git_to_index(
+                git_url,
+                branch_name="",
+                git_connection_id="",
+            ):
+                git_clone = git_clone_component(git_repository=git_url, branch_name=branch_name)
+                git_clone.environment_variables["AZUREML_WORKSPACE_CONNECTION_ID_GIT"] = git_connection_id
+
+                index_job = index_data_func(
+                    description=data_index.description,
+                    data_index=data_index,
+                    input_data_override=git_clone.outputs.output_data,
+                    ml_client=MLClient(
+                        subscription_id=self._subscription_id,
+                        resource_group_name=self._resource_group_name,
+                        workspace_name=self._workspace_name,
+                        credential=self._credential,
+                    ),
+                )
+                # pylint: disable=no-member
+                return index_job.outputs
+
+            git_index_job = git_to_index(
+                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
+
+            # Submit the DataIndex Job
+            return self._all_operations.all_operations[AzureMLResourceType.JOB].create_or_update(git_index_job)
+        raise ValueError(f"Unsupported input source type {type(input_source)}")
+
+    def _create_data_indexing_job(
+        self,
+        data_index: DataIndex,
+        identity: Optional[Union[ManagedIdentityConfiguration, UserIdentityConfiguration]] = None,
+        compute: str = "serverless",
+        serverless_instance_type: Optional[str] = None,
+        input_data_override: Optional[Input] = None,
+        submit_job: bool = True,
+        **kwargs,
+    ) -> PipelineJob:
+        """
+        Returns the data import job that is creating the data asset.
+
+        :param data_index: DataIndex object.
+        :type data_index: azure.ai.ml.entities._dataindex
+        :param identity: Identity configuration for the job.
+        :type identity: Optional[Union[ManagedIdentityConfiguration, UserIdentityConfiguration]]
+        :param compute: The compute target to use for the job. Default: "serverless".
+        :type compute: str
+        :param serverless_instance_type: The instance type to use for serverless compute.
+        :type serverless_instance_type: Optional[str]
+        :param input_data_override: Input data override for the job.
+            Used to pipe output of step into DataIndex Job in a pipeline.
+        :type input_data_override: Optional[Input]
+        :param submit_job: Whether to submit the job to the service. Default: True.
+        :type submit_job: bool
+        :return: data import job object.
+        :rtype: ~azure.ai.ml.entities.PipelineJob.
+        """
+        # pylint: disable=no-member
+        from azure.ai.ml import MLClient
+
+        default_name = "data_index_" + data_index.name if data_index.name is not None else ""
+        experiment_name = kwargs.pop("experiment_name", None) or default_name
+        data_index.type = AssetTypes.URI_FOLDER
+
+        # avoid specifying auto_delete_setting in job output now
+        _validate_auto_delete_setting_in_data_output(data_index.auto_delete_setting)
+
+        # block customer specified path on managed datastore
+        data_index.path = _validate_workspace_managed_datastore(data_index.path)
+
+        if "${{name}}" not in str(data_index.path) and "{name}" not in str(data_index.path):
+            data_index.path = str(data_index.path).rstrip("/") + "/${{name}}"
+
+        index_job = index_data_func(
+            description=data_index.description or kwargs.pop("description", None) or default_name,
+            name=data_index.name or kwargs.pop("name", None),
+            display_name=kwargs.pop("display_name", None) or default_name,
+            experiment_name=experiment_name,
+            compute=compute,
+            serverless_instance_type=serverless_instance_type,
+            data_index=data_index,
+            ml_client=MLClient(
+                subscription_id=self._subscription_id,
+                resource_group_name=self._resource_group_name,
+                workspace_name=self._workspace_name,
+                credential=self._credential,
+            ),
+            identity=identity,
+            input_data_override=input_data_override,
+            **kwargs,
+        )
+        index_pipeline = PipelineJob(
+            description=index_job.description,
+            tags=index_job.tags,
+            name=index_job.name,
+            display_name=index_job.display_name,
+            experiment_name=experiment_name,
+            properties=index_job.properties or {},
+            settings=PipelineJobSettings(force_rerun=True, default_compute=compute),
+            jobs={default_name: index_job},
+        )
+        index_pipeline.properties["azureml.mlIndexAssetName"] = data_index.name
+        index_pipeline.properties["azureml.mlIndexAssetKind"] = data_index.index.type
+        index_pipeline.properties["azureml.mlIndexAssetSource"] = kwargs.pop("mlindex_asset_source", "Data Asset")
+
+        if submit_job:
+            return self._all_operations.all_operations[AzureMLResourceType.JOB].create_or_update(
+                job=index_pipeline, skip_validation=True, **kwargs
+            )
+        return index_pipeline