Foundry change (Azure#40679)

* adding target storage for capability host

* Update sdk/ml/azure-ai-ml/CHANGELOG.md

Co-authored-by: Copilot <[email protected]>

* unit test fix

* correcting change log

---------

Co-authored-by: Copilot <[email protected]>

Author: achauhan-scc

sdk/ml/azure-ai-ml/CHANGELOG.md

@@ -2,6 +2,9 @@
 
 ## 1.27.0 (unreleased)
 
+### Features Added
+- Added Target storage connection for capability host.
+
 ### Bugs Fixed
 - Fix for compute Instance, disableLocalAuth property should be depend on ssh public access enabled.
 - Removing Git-related properties from job properties if a PAT token is detected in the repository URL.

sdk/ml/azure-ai-ml/azure/ai/ml/_ml_client.py

@@ -543,7 +543,7 @@ def __init__(
         self._capability_hosts = CapabilityHostsOperations(
             self._operation_scope,
             self._operation_config,
-            self._service_client_10_2024_preview,
+            self._service_client_01_2025_preview,
             self._operation_container,
             self._credential,
             **kwargs,

sdk/ml/azure-ai-ml/azure/ai/ml/_schema/workspace/ai_workspaces/capability_host.py

@@ -16,3 +16,4 @@ class CapabilityHostSchema(PathAwareSchema):
     vector_store_connections = fields.List(fields.Str(), required=False)
     ai_services_connections = fields.List(fields.Str(), required=False)
     storage_connections = fields.List(fields.Str(), required=False)
+    thread_storage_connections = fields.List(fields.Str(), required=False)

sdk/ml/azure-ai-ml/azure/ai/ml/entities/_autogen_entities/_model_base.py

@@ -6,23 +6,25 @@
 # --------------------------------------------------------------------------
 # pylint: disable=protected-access, broad-except, no-member
 
+import base64
 import calendar
+import copy
 import decimal
+import email.utils
+import enum
 import functools
-import sys
 import logging
-import base64
 import re
-import copy
+import sys
 import typing
-import enum
-import email.utils
-from datetime import datetime, date, time, timedelta, timezone
+from datetime import date, datetime, time, timedelta, timezone
 from json import JSONEncoder
-from typing_extensions import Self
+
 import isodate
-from azure.core.exceptions import DeserializationError
+from typing_extensions import Self
+
 from azure.core import CaseInsensitiveEnumMeta
+from azure.core.exceptions import DeserializationError
 from azure.core.pipeline import PipelineResponse
 from azure.core.serialization import NULL
 

sdk/ml/azure-ai-ml/azure/ai/ml/entities/_workspace/_ai_workspaces/capability_host.py

@@ -4,35 +4,20 @@
 
 import os
 from os import PathLike
-from typing import (
-    List,
-    Optional,
-    Union,
-    IO,
-    Any,
-    AnyStr,
-    Dict,
-)
 from pathlib import Path
-from azure.ai.ml._utils._experimental import experimental
-from azure.ai.ml.entities._resource import Resource
-from azure.ai.ml.constants._workspace import CapabilityHostKind
-from azure.ai.ml.constants._common import (
-    BASE_PATH_CONTEXT_KEY,
-    PARAMS_OVERRIDE_KEY,
-)
+from typing import IO, Any, AnyStr, Dict, List, Optional, Union
 
-from azure.ai.ml._schema.workspace.ai_workspaces.capability_host import (
-    CapabilityHostSchema,
+from azure.ai.ml._restclient.v2025_01_01_preview.models._models_py3 import CapabilityHost as RestCapabilityHost
+from azure.ai.ml._restclient.v2025_01_01_preview.models._models_py3 import (
+    CapabilityHostProperties as RestCapabilityHostProperties,
 )
+from azure.ai.ml._schema.workspace.ai_workspaces.capability_host import CapabilityHostSchema
+from azure.ai.ml._utils._experimental import experimental
 from azure.ai.ml._utils.utils import dump_yaml_to_file
+from azure.ai.ml.constants._common import BASE_PATH_CONTEXT_KEY, PARAMS_OVERRIDE_KEY
+from azure.ai.ml.constants._workspace import CapabilityHostKind
+from azure.ai.ml.entities._resource import Resource
 from azure.ai.ml.entities._util import load_from_dict
-from azure.ai.ml._restclient.v2024_10_01_preview.models._models_py3 import (
-    CapabilityHost as RestCapabilityHost,
-)
-from azure.ai.ml._restclient.v2024_10_01_preview.models._models_py3 import (
-    CapabilityHostProperties as RestCapabilityHostProperties,
-)
 
 
 @experimental
@@ -75,6 +60,7 @@ def __init__(
         vector_store_connections: Optional[List[str]] = None,
         ai_services_connections: Optional[List[str]] = None,
         storage_connections: Optional[List[str]] = None,
+        thread_storage_connections: Optional[List[str]] = None,
         capability_host_kind: Union[str, CapabilityHostKind] = CapabilityHostKind.AGENTS,
         **kwargs: Any,
     ):
@@ -83,6 +69,7 @@ def __init__(
         self.ai_services_connections = ai_services_connections
         self.storage_connections = storage_connections
         self.vector_store_connections = vector_store_connections
+        self.thread_storage_connections = thread_storage_connections
 
     def dump(
         self,
@@ -160,6 +147,9 @@ def _from_rest_object(cls, rest_obj: RestCapabilityHost) -> "CapabilityHost":
             ai_services_connections=(rest_obj.properties.ai_services_connections if rest_obj.properties else None),
             storage_connections=(rest_obj.properties.storage_connections if rest_obj.properties else None),
             vector_store_connections=(rest_obj.properties.vector_store_connections if rest_obj.properties else None),
+            thread_storage_connections=(
+                rest_obj.properties.thread_storage_connections if rest_obj.properties else None
+            ),
             capability_host_kind=(
                 rest_obj.properties.capability_host_kind if rest_obj.properties else CapabilityHostKind.AGENTS
             ),
@@ -178,6 +168,7 @@ def _to_rest_object(self) -> RestCapabilityHost:
             ai_services_connections=self.ai_services_connections,
             storage_connections=self.storage_connections,
             vector_store_connections=self.vector_store_connections,
+            thread_storage_connections=self.thread_storage_connections,
             description=self.description,
             capability_host_kind=self.capability_host_kind,
         )

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

@@ -8,7 +8,7 @@
 from marshmallow.exceptions import ValidationError as SchemaValidationError
 
 from azure.ai.ml._exception_helper import log_and_raise_error
-from azure.ai.ml._restclient.v2024_10_01_preview import AzureMachineLearningWorkspaces as ServiceClient102024Preview
+from azure.ai.ml._restclient.v2025_01_01_preview import AzureMachineLearningWorkspaces as ServiceClient012025Preview
 from azure.ai.ml._scope_dependent_operations import (
     OperationConfig,
     OperationsContainer,
@@ -40,9 +40,9 @@ class CapabilityHostsOperations(_ScopeDependentOperations):
     :type operation_scope: ~azure.ai.ml._scope_dependent_operations.OperationScope
     :param operation_config: Common configuration for operations classes of an MLClient object.
     :type operation_config: ~azure.ai.ml._scope_dependent_operations.OperationConfig
-    :param service_client_10_2024: Service client to allow end users to operate on Azure Machine Learning Workspace
-        resources (ServiceClient102024Preview).
-    :type service_client_10_2024: ~azure.ai.ml._restclient.v2024_10_01_preview._azure_machine_learning_workspaces.AzureMachineLearningWorkspaces    # pylint: disable=line-too-long
+    :param service_client_01_2025: Service client to allow end users to operate on Azure Machine Learning Workspace
+        resources (ServiceClient012025Preview).
+    :type service_client_01_2025: ~azure.ai.ml._restclient.v2025_01_01_preview._azure_machine_learning_workspaces.AzureMachineLearningWorkspaces    # pylint: disable=line-too-long
     :param all_operations: All operations classes of an MLClient object.
     :type all_operations: ~azure.ai.ml._scope_dependent_operations.OperationsContainer
     :param credentials: Credential to use for authentication.
@@ -55,7 +55,7 @@ def __init__(
         self,
         operation_scope: OperationScope,
         operation_config: OperationConfig,
-        service_client_10_2024: ServiceClient102024Preview,
+        service_client_01_2025: ServiceClient012025Preview,
         all_operations: OperationsContainer,
         credentials: TokenCredential,
         **kwargs: Any,
@@ -66,9 +66,9 @@ def __init__(
         :type operation_scope: ~azure.ai.ml._scope_dependent_operations.OperationScope
         :param operation_config: Common configuration for operations classes of an MLClient object.
         :type operation_config: ~azure.ai.ml._scope_dependent_operations.OperationConfig
-        :param service_client_10_2024: Service client to allow end users to operate on Azure Machine Learning Workspace
-            resources (ServiceClient102024Preview).
-        :type service_client_10_2024: ~azure.ai.ml._restclient.v2024_10_01_preview._azure_machine_learning_workspaces.AzureMachineLearningWorkspaces    # pylint: disable=line-too-long
+        :param service_client_01_2025: Service client to allow end users to operate on Azure Machine Learning Workspace
+            resources (ServiceClient012025Preview).
+        :type service_client_01_2025: ~azure.ai.ml._restclient.v2025_01_01_preview._azure_machine_learning_workspaces.AzureMachineLearningWorkspaces    # pylint: disable=line-too-long
         :param all_operations: All operations classes of an MLClient object.
         :type all_operations: ~azure.ai.ml._scope_dependent_operations.OperationsContainer
         :param credentials: Credential to use for authentication.
@@ -80,8 +80,8 @@ def __init__(
         super(CapabilityHostsOperations, self).__init__(operation_scope, operation_config)
         ops_logger.update_filter()
         self._all_operations = all_operations
-        self._capability_hosts_operations = service_client_10_2024.capability_hosts
-        self._workspace_operations = service_client_10_2024.workspaces
+        self._capability_hosts_operations = service_client_01_2025.capability_hosts
+        self._workspace_operations = service_client_01_2025.workspaces
         self._credentials = credentials
         self._init_kwargs = kwargs
 

sdk/ml/azure-ai-ml/tests/test_configs/workspace/ai_workspaces/test_capability_host_project.yml

@@ -6,4 +6,6 @@ ai_services_connections:
   - "aiservice_connection_1"
 storage_connections:
   - "storage_connection_1"
+thread_storage_connections:
+  - "thread_storage_connection_1"
 capability_host_kind: "Agents"

sdk/ml/azure-ai-ml/tests/workspace/ai_workspaces/unittests/test_capability_host_entity.py

@@ -1,17 +1,14 @@
+import os
+from datetime import datetime
+
 import pytest
 
 from azure.ai.ml import load_capability_host
-from azure.ai.ml.entities._workspace._ai_workspaces.capability_host import (
-    CapabilityHost,
-)
-from azure.ai.ml._restclient.v2024_10_01_preview.models._models_py3 import (
-    CapabilityHost as RestCapabilityHost,
-)
-from azure.ai.ml._restclient.v2024_10_01_preview.models._models_py3 import (
+from azure.ai.ml._restclient.v2025_01_01_preview.models._models_py3 import CapabilityHost as RestCapabilityHost
+from azure.ai.ml._restclient.v2025_01_01_preview.models._models_py3 import (
     CapabilityHostProperties as RestCapabilityHostProperties,
 )
-from datetime import datetime
-import os
+from azure.ai.ml.entities._workspace._ai_workspaces.capability_host import CapabilityHost
 
 
 @pytest.mark.unittest
@@ -54,6 +51,11 @@ def test_capability_host_project_schema(self) -> None:
             and len(capability_host.vector_store_connections) == 1
             and capability_host.vector_store_connections[0] == "vector_store_connection_1"
         )
+        assert (
+            capability_host.thread_storage_connections is not None
+            and len(capability_host.thread_storage_connections) == 1
+            and capability_host.thread_storage_connections[0] == "thread_storage_connection_1"
+        )
 
     def test_capability_host_constructor(self) -> None:
         capability_host = CapabilityHost(

sdk/ml/azure-ai-ml/tests/workspace/ai_workspaces/unittests/test_capability_host_operations.py

@@ -1,13 +1,11 @@
 from unittest.mock import Mock
 
 import pytest
-from azure.ai.ml import load_capability_host
 
-from azure.ai.ml._scope_dependent_operations import OperationScope, OperationConfig
-from azure.ai.ml.operations._capability_hosts_operations import (
-    CapabilityHostsOperations,
-)
+from azure.ai.ml import load_capability_host
+from azure.ai.ml._scope_dependent_operations import OperationConfig, OperationScope
 from azure.ai.ml.exceptions import MlException
+from azure.ai.ml.operations._capability_hosts_operations import CapabilityHostsOperations
 
 
 @pytest.fixture
@@ -19,14 +17,14 @@ def mock_credential() -> Mock:
 def mock_capability_hosts_operation(
     mock_workspace_scope: OperationScope,
     mock_operation_config: OperationConfig,
-    mock_aml_services_2024_10_01_preview: Mock,
+    mock_aml_services_2025_01_01_preview: Mock,
     mock_machinelearning_client: Mock,
     mock_credential: Mock,
 ) -> CapabilityHostsOperations:
     yield CapabilityHostsOperations(
         operation_scope=mock_workspace_scope,
         operation_config=mock_operation_config,
-        service_client_10_2024=mock_aml_services_2024_10_01_preview,
+        service_client_01_2025=mock_aml_services_2025_01_01_preview,
         all_operations=mock_machinelearning_client._operation_container,
         credentials=mock_credential,
     )
@@ -36,14 +34,14 @@ def mock_capability_hosts_operation(
 def mock_invalid_capability_hosts_operation(
     mock_registry_scope: OperationScope,
     mock_operation_config: OperationConfig,
-    mock_aml_services_2024_10_01_preview: Mock,
+    mock_aml_services_2025_01_01_preview: Mock,
     mock_machinelearning_client: Mock,
     mock_credential: Mock,
 ) -> CapabilityHostsOperations:
     yield CapabilityHostsOperations(
         operation_scope=mock_registry_scope,
         operation_config=mock_operation_config,
-        service_client_10_2024=mock_aml_services_2024_10_01_preview,
+        service_client_01_2025=mock_aml_services_2025_01_01_preview,
         all_operations=mock_machinelearning_client._operation_container,
         credentials=mock_credential,
     )