Read/write from/to disk for Agent Code. Add overwrite to session file download

sdk/ai/azure-ai-projects/api.md

@@ -260,6 +260,17 @@ namespace azure.ai.projects.aio.operations
                 **kwargs: Any
             ) -> AsyncIterator[bytes]: ...
 
+        @distributed_trace_async
+        async def download_code_to_disk(
+                self, 
+                agent_name: str, 
+                *, 
+                agent_version: Optional[str] = ..., 
+                file_path: Union[str, PathLike[str]], 
+                overwrite: bool = False, 
+                **kwargs: Any
+            ) -> str: ...
+
         @distributed_trace_async
         async def download_session_file(
                 self, 
@@ -277,6 +288,7 @@ namespace azure.ai.projects.aio.operations
                 session_id: str, 
                 *, 
                 file_path: Union[str, PathLike[str]], 
+                overwrite: bool = False, 
                 remote_path: str, 
                 **kwargs: Any
             ) -> None: ...
@@ -9606,6 +9618,17 @@ namespace azure.ai.projects.operations
                 **kwargs: Any
             ) -> Iterator[bytes]: ...
 
+        @distributed_trace
+        def download_code_to_disk(
+                self, 
+                agent_name: str, 
+                *, 
+                agent_version: Optional[str] = ..., 
+                file_path: Union[str, PathLike[str]], 
+                overwrite: bool = False, 
+                **kwargs: Any
+            ) -> str: ...
+
         @distributed_trace
         def download_session_file(
                 self, 
@@ -9623,6 +9646,7 @@ namespace azure.ai.projects.operations
                 session_id: str, 
                 *, 
                 file_path: Union[str, PathLike[str]], 
+                overwrite: bool = False, 
                 remote_path: str, 
                 **kwargs: Any
             ) -> None: ...

sdk/ai/azure-ai-projects/api.metadata.yml

@@ -1,3 +1,3 @@
-apiMdSha256: aaf7da6cfe754ac521a35dde3f23b00242c300533c89e6e1001e434a63bb8e4b
+apiMdSha256: d2a533b0dc76f7cd7128d9bfd30ff9e789edb5641ee2431f9689990c8112bc68
 parserVersion: 0.3.28
 pythonVersion: 3.14.3

sdk/ai/azure-ai-projects/azure/ai/projects/aio/operations/_patch_agents_async.py

@@ -8,6 +8,7 @@
 Follow our quickstart for examples: https://aka.ms/azsdk/python/dpcodegen/python/customize
 """
 
+import hashlib
 import os
 from pathlib import Path
 from typing import Union, Optional, Any, IO, overload
@@ -317,6 +318,7 @@ async def download_session_file_to_disk(
         session_id: str,
         *,
         file_path: Union[str, "os.PathLike[str]"],
+        overwrite: bool = False,
         remote_path: str,
         **kwargs: Any,
     ) -> None:
@@ -331,18 +333,25 @@ async def download_session_file_to_disk(
         :type session_id: str
         :keyword file_path: The full path to the local file where the content should be written. Required.
         :paramtype file_path: str or os.PathLike[str]
+        :keyword overwrite: If True, overwrite the local file if it already exists. If False (default),
+         raise FileExistsError when the file already exists.
+        :paramtype overwrite: bool
         :keyword remote_path: The file path to download from the sandbox, relative to the session home
          directory. Required.
         :paramtype remote_path: str
         :return: None
         :rtype: None
         :raises ~azure.core.exceptions.HttpResponseError:
+        :raises FileExistsError: If *file_path* already exists and *overwrite* is False.
         :raises ValueError: If *file_path* points to a directory.
         :raises OSError: If the file cannot be written.
         """
         p = Path(file_path)
-        if p.exists() and p.is_dir():
-            raise ValueError(f"Provide a valid file path, not a folder path `{file_path}`.")
+        if p.exists():
+            if p.is_dir():
+                raise ValueError(f"Provide a valid file path, not a folder path `{file_path}`.")
+            if not overwrite:
+                raise FileExistsError(f"The file `{file_path}` already exists. Set overwrite=True to replace it.")
 
         # Download the file content using the existing method
         content_iterator = await self.download_session_file(
@@ -356,3 +365,60 @@ async def download_session_file_to_disk(
         with open(file_path, "wb") as f:
             async for chunk in content_iterator:
                 f.write(chunk)
+
+    @distributed_trace_async
+    async def download_code_to_disk(
+        self,
+        agent_name: str,
+        *,
+        file_path: Union[str, "os.PathLike[str]"],
+        overwrite: bool = False,
+        agent_version: Optional[str] = None,
+        **kwargs: Any,
+    ) -> str:
+        """Download agent code directly to disk.
+
+        Downloads the code zip for a code-based hosted agent and writes it to a local file.
+
+        If ``agent_version`` is supplied, downloads that version's code zip; otherwise
+        downloads the latest version's code zip.
+
+        :param agent_name: The name of the agent. Required.
+        :type agent_name: str
+        :keyword file_path: The full path to the local file where the code zip should be written. Required.
+        :paramtype file_path: str or os.PathLike[str]
+        :keyword overwrite: If True, overwrite the local file if it already exists. If False (default),
+         raise FileExistsError when the file already exists.
+        :paramtype overwrite: bool
+        :keyword agent_version: The version of the agent whose code zip should be downloaded.
+         If omitted, the latest version's code zip is downloaded. Default value is None.
+        :paramtype agent_version: str
+        :return: The SHA-256 hex digest of the downloaded file.
+        :rtype: str
+        :raises ~azure.core.exceptions.HttpResponseError:
+        :raises FileExistsError: If *file_path* already exists and *overwrite* is False.
+        :raises ValueError: If *file_path* points to a directory.
+        :raises OSError: If the file cannot be written.
+        """
+        p = Path(file_path)
+        if p.exists():
+            if p.is_dir():
+                raise ValueError(f"Provide a valid file path, not a folder path `{file_path}`.")
+            if not overwrite:
+                raise FileExistsError(f"The file `{file_path}` already exists. Set overwrite=True to replace it.")
+
+        # Download the code content using the existing method
+        content_iterator = await self.download_code(
+            agent_name=agent_name,
+            agent_version=agent_version,
+            **kwargs,
+        )
+
+        # Write the content to disk and calculate SHA-256
+        sha = hashlib.sha256()
+        with open(file_path, "wb") as f:
+            async for chunk in content_iterator:
+                f.write(chunk)
+                sha.update(chunk)
+
+        return sha.hexdigest()

sdk/ai/azure-ai-projects/azure/ai/projects/operations/_patch_agents.py

@@ -8,6 +8,7 @@
 Follow our quickstart for examples: https://aka.ms/azsdk/python/dpcodegen/python/customize
 """
 
+import hashlib
 import os
 from pathlib import Path
 from typing import Union, Optional, Any, IO, overload
@@ -333,6 +334,7 @@ def download_session_file_to_disk(
         session_id: str,
         *,
         file_path: Union[str, "os.PathLike[str]"],
+        overwrite: bool = False,
         remote_path: str,
         **kwargs: Any,
     ) -> None:
@@ -350,15 +352,22 @@ def download_session_file_to_disk(
         :keyword remote_path: The file path to download from the sandbox, relative to the session home
          directory. Required.
         :paramtype remote_path: str
+        :keyword overwrite: If True, overwrite the local file if it already exists. If False (default),
+         raise FileExistsError when the file already exists.
+        :paramtype overwrite: bool
         :return: None
         :rtype: None
         :raises ~azure.core.exceptions.HttpResponseError:
+        :raises FileExistsError: If *file_path* already exists and *overwrite* is False.
         :raises ValueError: If *file_path* points to a directory.
         :raises OSError: If the file cannot be written.
         """
         p = Path(file_path)
-        if p.exists() and p.is_dir():
-            raise ValueError(f"Provide a valid file path, not a folder path `{file_path}`.")
+        if p.exists():
+            if p.is_dir():
+                raise ValueError(f"Provide a valid file path, not a folder path `{file_path}`.")
+            if not overwrite:
+                raise FileExistsError(f"The file `{file_path}` already exists. Set overwrite=True to replace it.")
 
         # Download the file content using the existing method
         content_iterator = self.download_session_file(
@@ -372,3 +381,60 @@ def download_session_file_to_disk(
         with open(file_path, "wb") as f:
             for chunk in content_iterator:
                 f.write(chunk)
+
+    @distributed_trace
+    def download_code_to_disk(
+        self,
+        agent_name: str,
+        *,
+        file_path: Union[str, "os.PathLike[str]"],
+        overwrite: bool = False,
+        agent_version: Optional[str] = None,
+        **kwargs: Any,
+    ) -> str:
+        """Download agent code directly to disk.
+
+        Downloads the code zip for a code-based hosted agent and writes it to a local file.
+
+        If ``agent_version`` is supplied, downloads that version's code zip; otherwise
+        downloads the latest version's code zip.
+
+        :param agent_name: The name of the agent. Required.
+        :type agent_name: str
+        :keyword file_path: The full path to the local file where the code zip should be written. Required.
+        :paramtype file_path: str or os.PathLike[str]
+        :keyword overwrite: If True, overwrite the local file if it already exists. If False (default),
+         raise FileExistsError when the file already exists.
+        :paramtype overwrite: bool
+        :keyword agent_version: The version of the agent whose code zip should be downloaded.
+         If omitted, the latest version's code zip is downloaded. Default value is None.
+        :paramtype agent_version: str
+        :return: The SHA-256 hex digest of the downloaded file.
+        :rtype: str
+        :raises ~azure.core.exceptions.HttpResponseError:
+        :raises FileExistsError: If *file_path* already exists and *overwrite* is False.
+        :raises ValueError: If *file_path* points to a directory.
+        :raises OSError: If the file cannot be written.
+        """
+        p = Path(file_path)
+        if p.exists():
+            if p.is_dir():
+                raise ValueError(f"Provide a valid file path, not a folder path `{file_path}`.")
+            if not overwrite:
+                raise FileExistsError(f"The file `{file_path}` already exists. Set overwrite=True to replace it.")
+
+        # Download the code content using the existing method
+        content_iterator = self.download_code(
+            agent_name=agent_name,
+            agent_version=agent_version,
+            **kwargs,
+        )
+
+        # Write the content to disk and calculate SHA-256
+        sha = hashlib.sha256()
+        with open(file_path, "wb") as f:
+            for chunk in content_iterator:
+                f.write(chunk)
+                sha.update(chunk)
+
+        return sha.hexdigest()

sdk/ai/azure-ai-projects/docs/public-methods.md

@@ -4,16 +4,16 @@ This document lists all public methods available on `AIProjectClient` and its su
 
 ## Summary
 
-There are a total of 142 unique public methods:
+There are a total of 143 unique public methods:
 - 5 stable methods on the client
-- 56 stable methods on top-level sub-clients
+- 57 stable methods on top-level sub-clients
 - 81 beta methods on nested beta sub-clients
 
 ### Top-level sub-clients (stable operations)
 
 | Subclient | Class Name | Methods Count |
 |-----------|------------|----------------|
-| `agents` | AgentsOperations | 24 |
+| `agents` | AgentsOperations | 25 |
 | `connections` | ConnectionsOperations | 3 |
 | `datasets` | DatasetsOperations | 9 |
 | `deployments` | DeploymentsOperations | 2 |
@@ -66,6 +66,7 @@ Alphabetically sorted. An asterisk at the end of the method name means is a hand
 .agents.delete_version
 .agents.disable
 .agents.download_code
+.agents.download_code_to_disk*
 .agents.download_session_file
 .agents.download_session_file_to_disk*
 .agents.enable

sdk/ai/azure-ai-projects/samples/hosted_agents/sample_create_hosted_agent_from_code.py

@@ -39,15 +39,11 @@
        REMOTE_BUILD; defaults to `false` (BUNDLED).
 """
 
-import hashlib
 import os
 import tempfile
 from pathlib import Path
-
 from dotenv import load_dotenv
-
 from azure.identity import DefaultAzureCredential
-
 from azure.ai.projects import AIProjectClient
 from azure.ai.projects.models import (
     CodeConfiguration,
@@ -71,7 +67,7 @@
 
 with (
     DefaultAzureCredential() as credential,
-    AIProjectClient(endpoint=endpoint, credential=credential, allow_preview=True) as project_client,
+    AIProjectClient(endpoint=endpoint, credential=credential) as project_client,
 ):
     content = CreateAgentVersionFromCodeContent(
         metadata=CreateAgentVersionFromCodeMetadata(
@@ -112,15 +108,14 @@
 
     # Download the zip for the version we just created, streaming to a temp file.
     version_zip_path = Path(tempfile.gettempdir()) / f"{agent_name}-{created.version}.zip"
-    sha = hashlib.sha256()
-    with open(version_zip_path, "wb") as f:
-        for chunk in project_client.agents.download_code(
-            agent_name=agent_name,
-            agent_version=created.version,
-        ):
-            f.write(chunk)
-            sha.update(chunk)
-    downloaded_version_sha256 = sha.hexdigest()
+
+    downloaded_version_sha256 = project_client.agents.download_code_to_disk(
+        agent_name=agent_name,
+        agent_version=created.version,
+        file_path=version_zip_path,
+        overwrite=True,
+    )
+
     print(
         f"Downloaded version code zip to {version_zip_path}: {version_zip_path.stat().st_size} bytes, "
         f"sha256={downloaded_version_sha256} (matches uploaded: {downloaded_version_sha256 == code_zip_sha256})"

sdk/ai/azure-ai-projects/samples/hosted_agents/sample_create_hosted_agent_from_code_async.py

@@ -41,15 +41,11 @@
 """
 
 import asyncio
-import hashlib
 import os
 import tempfile
 from pathlib import Path
-
 from dotenv import load_dotenv
-
 from azure.identity.aio import DefaultAzureCredential
-
 from azure.ai.projects.aio import AIProjectClient
 from azure.ai.projects.models import (
     CodeConfiguration,
@@ -62,7 +58,6 @@
 from hosted_agents_util import select_echo_agent_code_zip, wait_for_agent_version_active_async
 from rbac_util import ensure_agent_identity_rbac_async
 
-
 async def main() -> None:
     load_dotenv()
 
@@ -75,7 +70,7 @@ async def main() -> None:
 
     async with (
         DefaultAzureCredential() as credential,
-        AIProjectClient(endpoint=endpoint, credential=credential, allow_preview=True) as project_client,
+        AIProjectClient(endpoint=endpoint, credential=credential) as project_client,
     ):
         content = CreateAgentVersionFromCodeContent(
             metadata=CreateAgentVersionFromCodeMetadata(
@@ -118,16 +113,14 @@ async def main() -> None:
 
         # Download the zip for the version we just created, streaming to a temp file.
         version_zip_path = Path(tempfile.gettempdir()) / f"{agent_name}-{created.version}.zip"
-        sha = hashlib.sha256()
-        version_stream = await project_client.agents.download_code(
+
+        downloaded_version_sha256 = await project_client.agents.download_code_to_disk(
             agent_name=agent_name,
             agent_version=created.version,
+            file_path=version_zip_path,
+            overwrite=True,
         )
-        with open(version_zip_path, "wb") as f:
-            async for chunk in version_stream:
-                f.write(chunk)
-                sha.update(chunk)
-        downloaded_version_sha256 = sha.hexdigest()
+
         print(
             f"Downloaded version code zip to {version_zip_path}: {version_zip_path.stat().st_size} bytes, "
             f"sha256={downloaded_version_sha256} (matches uploaded: {downloaded_version_sha256 == code_zip_sha256})"