From 5e1c90152ea1bb881f8532e710cb07a10ca5c448 Mon Sep 17 00:00:00 2001 From: Darren Cohen <39422044+dargilco@users.noreply.github.com> Date: Wed, 5 Aug 2026 15:48:32 -0700 Subject: [PATCH 1/8] polling sample --- ...mpleqna_for_finetuning_with_app_polling.py | 253 ++++++++++++++++++ 1 file changed, 253 insertions(+) create mode 100644 sdk/ai/azure-ai-projects/samples/datasets/sample_dataset_generation_job_simpleqna_for_finetuning_with_app_polling.py diff --git a/sdk/ai/azure-ai-projects/samples/datasets/sample_dataset_generation_job_simpleqna_for_finetuning_with_app_polling.py b/sdk/ai/azure-ai-projects/samples/datasets/sample_dataset_generation_job_simpleqna_for_finetuning_with_app_polling.py new file mode 100644 index 000000000000..87501b01b221 --- /dev/null +++ b/sdk/ai/azure-ai-projects/samples/datasets/sample_dataset_generation_job_simpleqna_for_finetuning_with_app_polling.py @@ -0,0 +1,253 @@ +# pylint: disable=line-too-long,useless-suppression +# ------------------------------------ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. +# ------------------------------------ + +""" +DESCRIPTION: + Generates supervised fine-tuning data from a Markdown reference document + uploaded as an Azure OpenAI File. The sample: + + 1. Uploads a short reference document via the Azure OpenAI Files API + (`purpose=user_data`) so it can be referenced by file id. + 2. Creates a `DataGenerationJob` (scenario=SUPERVISED_FINETUNING, + type=simple_qna) without SDK polling. + 3. Polls the job from application code until it reaches a terminal state, + then prints every generated file output. + 4. Cleans up the generated fine-tuning files and the Azure OpenAI input file. + + `simple_qna` REQUIRES `model_options` — the service uses the configured LLM + to synthesize the QnA pairs. Setting `train_split` triggers a split of + the generated samples into two Azure OpenAI output files. + +USAGE: + python sample_dataset_generation_job_simpleqna_for_finetuning_with_app_polling.py + + Before running the sample: + + pip install "azure-ai-projects>=2.4.0" azure-identity openai python-dotenv + + Set these environment variables with your own values: + 1) FOUNDRY_PROJECT_ENDPOINT - Required. The Azure AI Project endpoint, as found + in the overview page of your Microsoft Foundry project. + 2) FOUNDRY_MODEL_NAME - Required. The name of an Azure OpenAI model + deployment used to synthesize the QnA samples. For `simple_qna` fine-tuning, + the deployment must support the chat completions API (e.g. `gpt-4o`, `gpt-4.1`). + 3) DATASET_NAME - Optional. Name to assign to the generated output files + (used as the file name prefix). Defaults to `simpleqna-finetuning-sample`. + The service caps the rendered output name at 50 characters, so keep + custom values short — the sample appends a unique run id suffix. + 4) POLL_INTERVAL_SECONDS - Optional. Number of seconds to sleep between status + polls for the data generation job. Defaults to 10. +""" + +import io +import os +import time +import uuid +from datetime import datetime, timezone + +from dotenv import load_dotenv + +from azure.identity import DefaultAzureCredential +from azure.ai.projects import AIProjectClient +from azure.ai.projects.models import ( + DataGenerationJob, + DataGenerationJobInputs, + DataGenerationJobOutputOptions, + DataGenerationJobScenario, + DataGenerationModelOptions, + FileDataGenerationJobOutput, + FileDataGenerationJobSource, + JobStatus, + SimpleQnADataGenerationJobOptions, + SimpleQnAFineTuningQuestionType, +) + +load_dotenv() + +endpoint = os.environ["FOUNDRY_PROJECT_ENDPOINT"] +model_name = os.environ["FOUNDRY_MODEL_NAME"] +dataset_name = os.environ.get("DATASET_NAME", "simpleqna-finetuning-sample") +poll_interval_seconds = int(os.environ.get("POLL_INTERVAL_SECONDS", "10")) + +TERMINAL_STATUSES = {JobStatus.SUCCEEDED, JobStatus.FAILED, JobStatus.CANCELLED} + +# Unique per-run output name so repeated runs do not collide. +# Output names are capped at 50 characters by the service. +run_id = ( + f"{datetime.now(tz=timezone.utc).strftime('%y%m%d%H%M%S')}-{uuid.uuid4().hex[:4]}" +) +output_name = f"{dataset_name}-{run_id}" +if len(output_name) > 50: + raise ValueError( + f"Output name `{output_name}` exceeds the 50-character service limit. " + f"Lower DATASET_NAME (currently `{dataset_name}`) so that `-` fits within 50 characters." + ) + +# Reference document the sample uploads as an Azure OpenAI file. The service +# requires the file to contain at least 1 KB of content to generate QnA from. +SEED_REFERENCE_DOCUMENT = """# Widgets and Gizmos Reference + +## Products +- Widget: blue, manufactured at Factory 7 in Acme, carbon-fiber, rated to 80 C, sold in packs of 4, 250 g each. +- Gizmo: red, manufactured at Factory 12 in Bedrock, carbon-fiber, rated to 80 C, sold individually, 1.2 kg each. +- Sprocket: green, manufactured at Factory 3 in Acme, stainless steel, rated to 200 C, sold individually, 500 g each. + +## Operations +- Factory operates weekdays 0700-1900 local time. +- Closed on public holidays, except for the annual maintenance run on December 27. +- ISO 9001 certified; audited annually by an independent third party. +- Quality control samples every 100th unit and runs full destructive testing on every 5000th unit. + +## Customer support +- Warranty claims: email support@example.com with the serial number printed on the underside of the product. +- Returns: accepted within 30 days if unopened; opened items are eligible for repair only. +- Bulk orders (50+ units): contact sales@example.com for volume pricing and an extended 90-day return window. +- Replacement parts: orderable directly from the support portal using the original order number. + +## Pricing and SLAs +- Widget pack: USD 24.99 per 4-pack; free shipping on orders over USD 75. +- Gizmo unit: USD 49.99; free shipping on orders over USD 75. +- Sprocket unit: USD 14.99; ships from regional warehouses in 1-2 business days. +- Standard support response: within one business day. Priority support response: within four hours. +""" + +with ( + DefaultAzureCredential() as credential, + AIProjectClient(endpoint=endpoint, credential=credential) as project_client, + project_client.get_openai_client() as openai_client, +): + + # ------------------------------------------------------------------ + # 1. Upload the seed reference document as an Azure OpenAI file. + # ------------------------------------------------------------------ + seed_filename = f"widgets-gizmos-seed-{run_id}.md" + print(f"Upload the seed reference document as Azure OpenAI file `{seed_filename}`.") + seed_file = openai_client.files.create( + file=(seed_filename, io.BytesIO(SEED_REFERENCE_DOCUMENT.encode("utf-8"))), + purpose="user_data", + ) + print(f"Uploaded Azure OpenAI file (id: {seed_file.id}).") + + # Wait for the file to finish processing — the data generation service + # rejects references to files that are not yet in the `processed` state. + print("Wait for the Azure OpenAI file to be processed.", end="", flush=True) + while seed_file.status not in ("processed", "error"): + time.sleep(2) + seed_file = openai_client.files.retrieve(file_id=seed_file.id) + print(".", end="", flush=True) + print() + if seed_file.status != "processed": + raise RuntimeError( + f"Azure OpenAI file `{seed_file.id}` failed to process: status=`{seed_file.status}`." + ) + + # ------------------------------------------------------------------ + # 2. Submit a fine-tuning data generation job without SDK polling. + # ------------------------------------------------------------------ + job = DataGenerationJob( + inputs=DataGenerationJobInputs( + name=f"simpleqna-finetuning-{run_id}", + scenario=DataGenerationJobScenario.SUPERVISED_FINETUNING, + sources=[ + FileDataGenerationJobSource( + description="Widgets & Gizmos product / operations reference (Azure OpenAI file).", + id=seed_file.id, + ), + ], + options=SimpleQnADataGenerationJobOptions( + # Service requires max_samples to be between 15 and 1000. + max_samples=15, + # `simple_qna` REQUIRES model_options. + model_options=DataGenerationModelOptions(model=model_name), + # Split generated samples 80% training / 20% validation. + train_split=0.8, + # Ask for both short-answer and long-answer questions. + question_types=[ + SimpleQnAFineTuningQuestionType.SHORT_ANSWER, + SimpleQnAFineTuningQuestionType.LONG_ANSWER, + ], + ), + output_options=DataGenerationJobOutputOptions(name=output_name), + ), + ) + + print("Create a dataset generation job without SDK polling.") + created_jobs: list[DataGenerationJob] = [] + + def raw_response_hook(response): + # With polling disabled, this hook receives the initial response containing the created job. + response.http_response.read() + created_jobs.append(DataGenerationJob(response.http_response.json())) + + project_client.beta.datasets.begin_create_generation_job( + job=job, + polling=False, + raw_response_hook=raw_response_hook, + ) + if not created_jobs: + raise RuntimeError("The create operation did not return a data generation job.") + job = created_jobs[0] + print(f"Created job: id={job.id}, status={job.status}") + + # ------------------------------------------------------------------ + # 3. Poll from application code until the job reaches a terminal state. + # ------------------------------------------------------------------ + print(f"Polling job `{job.id}` to completion...", end="", flush=True) + while job.status not in TERMINAL_STATUSES: + time.sleep(poll_interval_seconds) + job = project_client.beta.datasets.get_generation_job(job_id=job.id) + print(".", end="", flush=True) + print() + print(f"Final job status: `{job.status}`.") + + if job.status == JobStatus.FAILED: + message = job.error.message if job.error else "" + raise RuntimeError(f"Data generation job `{job.id}` failed: {message}") + if job.status == JobStatus.CANCELLED: + raise RuntimeError(f"Data generation job `{job.id}` was cancelled.") + if job.result is None: + raise RuntimeError( + f"Data generation job `{job.id}` completed without a result." + ) + + job_result = job.result + print(f"Data generation result: {job_result}") + + # ------------------------------------------------------------------ + # 4. Inspect the generated fine-tuning file outputs. + # ------------------------------------------------------------------ + # `train_split=0.8` produces two Azure OpenAI files: a training partition + # and a validation partition. Both are emitted as FileDataGenerationJobOutput + # entries in `job_result.outputs`. + file_outputs = [ + output + for output in (job_result.outputs or []) + if isinstance(output, FileDataGenerationJobOutput) + ] + if not file_outputs: + raise RuntimeError("The data generation job did not produce any file outputs.") + + print(f"Generated {len(file_outputs)} fine-tuning file(s):") + for output in file_outputs: + if not output.id: + raise RuntimeError("A file output was returned without an id.") + # Resolve the Azure OpenAI file to surface its real filename and size. + file_info = openai_client.files.retrieve(file_id=output.id) + print( + f" - filename=`{file_info.filename}` id=`{output.id}` bytes={file_info.bytes}" + ) + if job_result.generated_samples is not None: + print(f"Generated samples: {job_result.generated_samples}") + + # ------------------------------------------------------------------ + # 5. Clean up. + # ------------------------------------------------------------------ + for output in file_outputs: + print(f"Delete the generated Azure OpenAI file `{output.id}`.") + openai_client.files.delete(file_id=output.id) + + print(f"Delete the Azure OpenAI input file `{seed_file.id}`.") + openai_client.files.delete(file_id=seed_file.id) From fbd25faff5c6a2ba1c558767df168446f975e5ae Mon Sep 17 00:00:00 2001 From: Darren Cohen <39422044+dargilco@users.noreply.github.com> Date: Wed, 5 Aug 2026 16:29:21 -0700 Subject: [PATCH 2/8] Custom LROPoller for dataset generation jobs --- .../ai/projects/aio/operations/_patch.py | 5 +- .../aio/operations/_patch_datasets_async.py | 129 +++++++++++++++++- .../azure/ai/projects/models/_patch.py | 84 +++++++++++- .../azure/ai/projects/operations/_patch.py | 5 +- .../ai/projects/operations/_patch_datasets.py | 126 ++++++++++++++++- ...mpleqna_for_finetuning_with_app_polling.py | 17 +-- .../tests/datasets/test_datasets.py | 26 +++- .../tests/datasets/test_datasets_async.py | 28 +++- 8 files changed, 397 insertions(+), 23 deletions(-) diff --git a/sdk/ai/azure-ai-projects/azure/ai/projects/aio/operations/_patch.py b/sdk/ai/azure-ai-projects/azure/ai/projects/aio/operations/_patch.py index 5d8893177cae..d734c8c691da 100644 --- a/sdk/ai/azure-ai-projects/azure/ai/projects/aio/operations/_patch.py +++ b/sdk/ai/azure-ai-projects/azure/ai/projects/aio/operations/_patch.py @@ -10,7 +10,7 @@ from typing import Any, List from ._patch_agents_async import AgentsOperations -from ._patch_datasets_async import DatasetsOperations +from ._patch_datasets_async import BetaDatasetsOperations, DatasetsOperations from ._patch_evaluation_rules_async import EvaluationRulesOperations from ._patch_telemetry_async import TelemetryOperations from ._patch_connections_async import ConnectionsOperations @@ -19,7 +19,6 @@ from ...operations._patch import _BETA_OPERATION_FEATURE_HEADERS, _OperationMethodHeaderProxy from ._operations import ( BetaAgentsOperations, - BetaDatasetsOperations, BetaEvaluationTaxonomiesOperations, BetaEvaluatorsOperations, BetaInsightsOperations, @@ -74,6 +73,8 @@ def __init__(self, *args: Any, **kwargs: Any) -> None: self.memory_stores = BetaMemoryStoresOperations(self._client, self._config, self._serialize, self._deserialize) # Replace with patched class that includes create (3-step upload helper) self.models = BetaModelsOperations(self._client, self._config, self._serialize, self._deserialize) + # Replace with patched class that returns AsyncDatasetGenerationLROPoller + self.datasets = BetaDatasetsOperations(self._client, self._config, self._serialize, self._deserialize) for property_name, foundry_features_value in _BETA_OPERATION_FEATURE_HEADERS.items(): setattr( diff --git a/sdk/ai/azure-ai-projects/azure/ai/projects/aio/operations/_patch_datasets_async.py b/sdk/ai/azure-ai-projects/azure/ai/projects/aio/operations/_patch_datasets_async.py index dc7095c827ea..6612e31eacad 100644 --- a/sdk/ai/azure-ai-projects/azure/ai/projects/aio/operations/_patch_datasets_async.py +++ b/sdk/ai/azure-ai-projects/azure/ai/projects/aio/operations/_patch_datasets_async.py @@ -11,13 +11,23 @@ import os import re import logging -from typing import Any, Tuple, Optional +from typing import Any, IO, Tuple, Optional, Union, cast, overload +from collections.abc import MutableMapping from pathlib import Path from urllib.parse import urlsplit from azure.storage.blob.aio import ContainerClient +from azure.core.polling import AsyncNoPolling, AsyncPollingMethod +from azure.core.polling.async_base_polling import AsyncLROBasePolling from azure.core.tracing.decorator_async import distributed_trace_async +from azure.core.utils import case_insensitive_dict -from ._operations import DatasetsOperations as DatasetsOperationsGenerated +from ._operations import ( + BetaDatasetsOperations as BetaDatasetsOperationsGenerated, + DatasetsOperations as DatasetsOperationsGenerated, +) +from ... import models as _models +from ..._utils.model_base import _deserialize +from ...models import AsyncDatasetGenerationLROPoller from ...models._models import ( FileDatasetVersion, FolderDatasetVersion, @@ -28,6 +38,121 @@ logger = logging.getLogger(__name__) +JSON = MutableMapping[str, Any] + + +class BetaDatasetsOperations(BetaDatasetsOperationsGenerated): + """Custom async operations for beta data generation jobs.""" + + @overload + async def begin_create_generation_job( + self, + job: _models.DataGenerationJob, + *, + operation_id: Optional[str] = None, + content_type: str = "application/json", + **kwargs: Any, + ) -> AsyncDatasetGenerationLROPoller: ... + + @overload + async def begin_create_generation_job( + self, + job: JSON, + *, + operation_id: Optional[str] = None, + content_type: str = "application/json", + **kwargs: Any, + ) -> AsyncDatasetGenerationLROPoller: ... + + @overload + async def begin_create_generation_job( + self, + job: IO[bytes], + *, + operation_id: Optional[str] = None, + content_type: str = "application/json", + **kwargs: Any, + ) -> AsyncDatasetGenerationLROPoller: ... + + @distributed_trace_async + async def begin_create_generation_job( + self, + job: Union[_models.DataGenerationJob, JSON, IO[bytes]], + *, + operation_id: Optional[str] = None, + **kwargs: Any, + ) -> AsyncDatasetGenerationLROPoller: + """Create a data generation job. + + :param job: The job to create. Required. + :type job: ~azure.ai.projects.models.DataGenerationJob or JSON or IO[bytes] + :keyword operation_id: Client-generated unique ID for idempotent retries. When absent, the + server creates the job unconditionally. Default value is None. + :paramtype operation_id: str + :return: A poller that returns DataGenerationJobResult and exposes the job ID in ``details``. + :rtype: ~azure.ai.projects.models.AsyncDatasetGenerationLROPoller + :raises ~azure.core.exceptions.HttpResponseError: + """ + headers = case_insensitive_dict(kwargs.pop("headers", {}) or {}) + params = kwargs.pop("params", {}) or {} + + content_type: Optional[str] = kwargs.pop("content_type", headers.pop("Content-Type", None)) + cls = kwargs.pop("cls", None) + polling: Union[bool, AsyncPollingMethod] = kwargs.pop("polling", True) + lro_delay = kwargs.pop("polling_interval", self._config.polling_interval) + continuation_token: Optional[str] = kwargs.pop("continuation_token", None) + raw_result = None + if continuation_token is None: + raw_result = await self._create_generation_job_initial( + job=job, + operation_id=operation_id, + content_type=content_type, + cls=lambda x, y, z: x, + headers=headers, + params=params, + **kwargs, + ) + await raw_result.http_response.read() # type: ignore + kwargs.pop("error_map", None) + + def get_long_running_output(pipeline_response): + response_headers = {} + response = pipeline_response.http_response + response_headers["Operation-Location"] = self._deserialize( + "str", response.headers.get("Operation-Location") + ) + response_headers["Location"] = self._deserialize("str", response.headers.get("Location")) + + deserialized = _deserialize(_models.DataGenerationJobResult, response.json().get("result", {})) + if cls: + return cls(pipeline_response, deserialized, response_headers) + return deserialized + + path_format_arguments = { + "endpoint": self._serialize.url("self._config.endpoint", self._config.endpoint, "str", skip_quote=True), + } + + if polling is True: + polling_method: AsyncPollingMethod = cast( + AsyncPollingMethod, + AsyncLROBasePolling(lro_delay, path_format_arguments=path_format_arguments, **kwargs), + ) + elif polling is False: + polling_method = cast(AsyncPollingMethod, AsyncNoPolling()) + else: + polling_method = polling + if continuation_token: + return AsyncDatasetGenerationLROPoller.from_continuation_token( + polling_method=polling_method, + continuation_token=continuation_token, + client=self._client, + deserialization_callback=get_long_running_output, + ) + assert raw_result is not None + return AsyncDatasetGenerationLROPoller( # type: ignore + self._client, raw_result, get_long_running_output, polling_method + ) + class DatasetsOperations(DatasetsOperationsGenerated): """ diff --git a/sdk/ai/azure-ai-projects/azure/ai/projects/models/_patch.py b/sdk/ai/azure-ai-projects/azure/ai/projects/models/_patch.py index 257e53dded78..bba5cad64ab2 100644 --- a/sdk/ai/azure-ai-projects/azure/ai/projects/models/_patch.py +++ b/sdk/ai/azure-ai-projects/azure/ai/projects/models/_patch.py @@ -33,7 +33,7 @@ TracesPreviewEvalRunDataSource, ) from ._models import CustomCredential as CustomCredentialGenerated -from ..models import MemoryStoreUpdateCompletedResult, MemoryStoreUpdateResult +from ..models import DataGenerationJobResult, MemoryStoreUpdateCompletedResult, MemoryStoreUpdateResult from ._enums import _FoundryFeaturesOptInKeys, _AgentDefinitionOptInKeys _FOUNDRY_FEATURES_HEADER_NAME: Final[str] = "Foundry-Features" @@ -380,7 +380,88 @@ def from_continuation_token( return cls(client, initial_response, deserialization_callback, polling_method) +class DatasetGenerationLROPoller(LROPoller[DataGenerationJobResult]): + """Custom LROPoller for data generation job operations.""" + + def __init__(self, client: Any, initial_response: Any, deserialization_callback: Any, polling_method: Any) -> None: + self._job_id = self._get_job_id(initial_response) + super().__init__(client, initial_response, deserialization_callback, polling_method) + + @staticmethod + def _get_job_id(initial_response: Any) -> Optional[str]: + try: + return initial_response.http_response.json().get("id") + except (AttributeError, TypeError, ValueError): + return None + + @property + def details(self) -> Mapping[str, Any]: + """Returns metadata associated with the data generation job operation. + + :return: A mapping containing the created data generation job ID. + :rtype: Mapping[str, Any] + """ + return {"job_id": self._job_id} + + @classmethod + def from_continuation_token( + cls, polling_method: PollingMethod[DataGenerationJobResult], continuation_token: str, **kwargs: Any + ) -> "DatasetGenerationLROPoller": + """Create a poller from a continuation token. + + :param polling_method: The polling strategy to adopt. + :type polling_method: ~azure.core.polling.PollingMethod + :param continuation_token: An opaque continuation token. + :type continuation_token: str + :return: An instance of DatasetGenerationLROPoller. + :rtype: DatasetGenerationLROPoller + """ + client, initial_response, deserialization_callback = polling_method.from_continuation_token( + continuation_token, **kwargs + ) + return cls(client, initial_response, deserialization_callback, polling_method) + + +class AsyncDatasetGenerationLROPoller(AsyncLROPoller[DataGenerationJobResult]): + """Custom AsyncLROPoller for data generation job operations.""" + + def __init__(self, client: Any, initial_response: Any, deserialization_callback: Any, polling_method: Any) -> None: + super().__init__(client, initial_response, deserialization_callback, polling_method) + self._job_id = DatasetGenerationLROPoller._get_job_id(initial_response) + + @property + def details(self) -> Mapping[str, Any]: + """Returns metadata associated with the data generation job operation. + + :return: A mapping containing the created data generation job ID. + :rtype: Mapping[str, Any] + """ + return {"job_id": self._job_id} + + @classmethod + def from_continuation_token( + cls, + polling_method: AsyncPollingMethod[DataGenerationJobResult], + continuation_token: str, + **kwargs: Any, + ) -> "AsyncDatasetGenerationLROPoller": + """Create a poller from a continuation token. + + :param polling_method: The polling strategy to adopt. + :type polling_method: ~azure.core.polling.AsyncPollingMethod + :param continuation_token: An opaque continuation token. + :type continuation_token: str + :return: An instance of AsyncDatasetGenerationLROPoller. + :rtype: AsyncDatasetGenerationLROPoller + """ + client, initial_response, deserialization_callback = polling_method.from_continuation_token( + continuation_token, **kwargs + ) + return cls(client, initial_response, deserialization_callback, polling_method) + + __all__: List[str] = [ + "AsyncDatasetGenerationLROPoller", "AsyncUpdateMemoriesLROPoller", "AzureAIAgentTargetParam", "AzureAIBenchmarkPreviewEvalRunDataSource", @@ -388,6 +469,7 @@ def from_continuation_token( "AzureAIModelTargetParam", "AzureAIResponsesEvalRunDataSource", "CustomCredential", + "DatasetGenerationLROPoller", "EvalCsvFileIdSource", "EvalCsvRunDataSource", "TestingCriterionAzureAIEvaluator", diff --git a/sdk/ai/azure-ai-projects/azure/ai/projects/operations/_patch.py b/sdk/ai/azure-ai-projects/azure/ai/projects/operations/_patch.py index 283443056bf4..49559c1b2e1d 100644 --- a/sdk/ai/azure-ai-projects/azure/ai/projects/operations/_patch.py +++ b/sdk/ai/azure-ai-projects/azure/ai/projects/operations/_patch.py @@ -13,7 +13,7 @@ from typing import Any, Callable, List from ..models._patch import _FOUNDRY_FEATURES_HEADER_NAME, _BETA_OPERATION_FEATURE_HEADERS, _has_header_case_insensitive from ._patch_agents import AgentsOperations -from ._patch_datasets import DatasetsOperations +from ._patch_datasets import BetaDatasetsOperations, DatasetsOperations from ._patch_evaluation_rules import EvaluationRulesOperations from ._patch_telemetry import TelemetryOperations from ._patch_connections import ConnectionsOperations @@ -21,7 +21,6 @@ from ._patch_models import BetaModelsOperations from ._operations import ( BetaAgentsOperations, - BetaDatasetsOperations, BetaEvaluationTaxonomiesOperations, BetaEvaluatorsOperations, BetaInsightsOperations, @@ -129,6 +128,8 @@ def __init__(self, *args: Any, **kwargs: Any) -> None: self.memory_stores = BetaMemoryStoresOperations(self._client, self._config, self._serialize, self._deserialize) # Replace with patched class that includes create (3-step upload helper) self.models = BetaModelsOperations(self._client, self._config, self._serialize, self._deserialize) + # Replace with patched class that returns DatasetGenerationLROPoller + self.datasets = BetaDatasetsOperations(self._client, self._config, self._serialize, self._deserialize) for property_name, foundry_features_value in _BETA_OPERATION_FEATURE_HEADERS.items(): setattr( diff --git a/sdk/ai/azure-ai-projects/azure/ai/projects/operations/_patch_datasets.py b/sdk/ai/azure-ai-projects/azure/ai/projects/operations/_patch_datasets.py index bf2c0db51271..be33b5a2763d 100644 --- a/sdk/ai/azure-ai-projects/azure/ai/projects/operations/_patch_datasets.py +++ b/sdk/ai/azure-ai-projects/azure/ai/projects/operations/_patch_datasets.py @@ -11,12 +11,22 @@ import os import re import logging -from typing import Any, Tuple, Optional +from typing import Any, IO, Tuple, Optional, Union, cast, overload +from collections.abc import MutableMapping from pathlib import Path from urllib.parse import urlsplit from azure.storage.blob import ContainerClient +from azure.core.polling import NoPolling, PollingMethod +from azure.core.polling.base_polling import LROBasePolling from azure.core.tracing.decorator import distributed_trace -from ._operations import DatasetsOperations as DatasetsOperationsGenerated +from azure.core.utils import case_insensitive_dict +from ._operations import ( + BetaDatasetsOperations as BetaDatasetsOperationsGenerated, + DatasetsOperations as DatasetsOperationsGenerated, +) +from .. import models as _models +from .._utils.model_base import _deserialize +from ..models import DatasetGenerationLROPoller from ..models._models import ( FileDatasetVersion, FolderDatasetVersion, @@ -27,6 +37,118 @@ logger = logging.getLogger(__name__) +JSON = MutableMapping[str, Any] + + +class BetaDatasetsOperations(BetaDatasetsOperationsGenerated): + """Custom operations for beta data generation jobs.""" + + @overload + def begin_create_generation_job( + self, + job: _models.DataGenerationJob, + *, + operation_id: Optional[str] = None, + content_type: str = "application/json", + **kwargs: Any, + ) -> DatasetGenerationLROPoller: ... + + @overload + def begin_create_generation_job( + self, + job: JSON, + *, + operation_id: Optional[str] = None, + content_type: str = "application/json", + **kwargs: Any, + ) -> DatasetGenerationLROPoller: ... + + @overload + def begin_create_generation_job( + self, + job: IO[bytes], + *, + operation_id: Optional[str] = None, + content_type: str = "application/json", + **kwargs: Any, + ) -> DatasetGenerationLROPoller: ... + + @distributed_trace + def begin_create_generation_job( + self, + job: Union[_models.DataGenerationJob, JSON, IO[bytes]], + *, + operation_id: Optional[str] = None, + **kwargs: Any, + ) -> DatasetGenerationLROPoller: + """Create a data generation job. + + :param job: The job to create. Required. + :type job: ~azure.ai.projects.models.DataGenerationJob or JSON or IO[bytes] + :keyword operation_id: Client-generated unique ID for idempotent retries. When absent, the + server creates the job unconditionally. Default value is None. + :paramtype operation_id: str + :return: A poller that returns DataGenerationJobResult and exposes the job ID in ``details``. + :rtype: ~azure.ai.projects.models.DatasetGenerationLROPoller + :raises ~azure.core.exceptions.HttpResponseError: + """ + headers = case_insensitive_dict(kwargs.pop("headers", {}) or {}) + params = kwargs.pop("params", {}) or {} + + content_type: Optional[str] = kwargs.pop("content_type", headers.pop("Content-Type", None)) + cls = kwargs.pop("cls", None) + polling: Union[bool, PollingMethod] = kwargs.pop("polling", True) + lro_delay = kwargs.pop("polling_interval", self._config.polling_interval) + continuation_token: Optional[str] = kwargs.pop("continuation_token", None) + raw_result = None + if continuation_token is None: + raw_result = self._create_generation_job_initial( + job=job, + operation_id=operation_id, + content_type=content_type, + cls=lambda x, y, z: x, + headers=headers, + params=params, + **kwargs, + ) + raw_result.http_response.read() # type: ignore + kwargs.pop("error_map", None) + + def get_long_running_output(pipeline_response): + response_headers = {} + response = pipeline_response.http_response + response_headers["Operation-Location"] = self._deserialize( + "str", response.headers.get("Operation-Location") + ) + response_headers["Location"] = self._deserialize("str", response.headers.get("Location")) + + deserialized = _deserialize(_models.DataGenerationJobResult, response.json().get("result", {})) + if cls: + return cls(pipeline_response, deserialized, response_headers) + return deserialized + + path_format_arguments = { + "endpoint": self._serialize.url("self._config.endpoint", self._config.endpoint, "str", skip_quote=True), + } + + if polling is True: + polling_method: PollingMethod = cast( + PollingMethod, LROBasePolling(lro_delay, path_format_arguments=path_format_arguments, **kwargs) + ) + elif polling is False: + polling_method = cast(PollingMethod, NoPolling()) + else: + polling_method = polling + if continuation_token: + return DatasetGenerationLROPoller.from_continuation_token( + polling_method=polling_method, + continuation_token=continuation_token, + client=self._client, + deserialization_callback=get_long_running_output, + ) + assert raw_result is not None + return DatasetGenerationLROPoller(self._client, raw_result, get_long_running_output, polling_method) # type: ignore + class DatasetsOperations(DatasetsOperationsGenerated): """ diff --git a/sdk/ai/azure-ai-projects/samples/datasets/sample_dataset_generation_job_simpleqna_for_finetuning_with_app_polling.py b/sdk/ai/azure-ai-projects/samples/datasets/sample_dataset_generation_job_simpleqna_for_finetuning_with_app_polling.py index 87501b01b221..3844c28e3c62 100644 --- a/sdk/ai/azure-ai-projects/samples/datasets/sample_dataset_generation_job_simpleqna_for_finetuning_with_app_polling.py +++ b/sdk/ai/azure-ai-projects/samples/datasets/sample_dataset_generation_job_simpleqna_for_finetuning_with_app_polling.py @@ -175,21 +175,14 @@ ) print("Create a dataset generation job without SDK polling.") - created_jobs: list[DataGenerationJob] = [] - - def raw_response_hook(response): - # With polling disabled, this hook receives the initial response containing the created job. - response.http_response.read() - created_jobs.append(DataGenerationJob(response.http_response.json())) - - project_client.beta.datasets.begin_create_generation_job( + poller = project_client.beta.datasets.begin_create_generation_job( job=job, polling=False, - raw_response_hook=raw_response_hook, ) - if not created_jobs: - raise RuntimeError("The create operation did not return a data generation job.") - job = created_jobs[0] + job_id = poller.details["job_id"] + if not job_id: + raise RuntimeError("The create operation did not return a data generation job ID.") + job = project_client.beta.datasets.get_generation_job(job_id=job_id) print(f"Created job: id={job.id}, status={job.status}") # ------------------------------------------------------------------ diff --git a/sdk/ai/azure-ai-projects/tests/datasets/test_datasets.py b/sdk/ai/azure-ai-projects/tests/datasets/test_datasets.py index cf7df74c2ffe..cf68496df52f 100644 --- a/sdk/ai/azure-ai-projects/tests/datasets/test_datasets.py +++ b/sdk/ai/azure-ai-projects/tests/datasets/test_datasets.py @@ -5,12 +5,15 @@ # ------------------------------------ import os import re +from unittest.mock import MagicMock + import pytest from test_base import TestBase, servicePreparer from devtools_testutils import recorded_by_proxy, is_live, is_live_and_not_recording, add_general_regex_sanitizer from azure.ai.projects import AIProjectClient -from azure.ai.projects.models import DatasetVersion, DatasetType +from azure.ai.projects.models import DatasetGenerationLROPoller, DatasetVersion, DatasetType from azure.ai.projects.models._enums import ConnectionType +from azure.ai.projects.operations._patch_datasets import BetaDatasetsOperations from azure.core.exceptions import HttpResponseError # Construct the paths to the data folder and data file used in this test @@ -20,6 +23,27 @@ data_file2 = os.path.join(data_folder, "data_file2.txt") +def test_begin_create_generation_job_exposes_job_id(): + """The sync create operation exposes its job ID without SDK polling.""" + operation = BetaDatasetsOperations.__new__(BetaDatasetsOperations) + operation._client = MagicMock() # pylint: disable=protected-access + operation._config = MagicMock(polling_interval=0) # pylint: disable=protected-access + operation._serialize = MagicMock() # pylint: disable=protected-access + operation._serialize.url.return_value = "https://example.test" # pylint: disable=protected-access + operation._deserialize = MagicMock() # pylint: disable=protected-access + + initial_response = MagicMock() + initial_response.http_response.json.return_value = {"id": "job-sync"} + operation._create_generation_job_initial = MagicMock( # pylint: disable=protected-access + return_value=initial_response + ) + + poller = operation.begin_create_generation_job(job={}, polling=False) + + assert isinstance(poller, DatasetGenerationLROPoller) + assert poller.details["job_id"] == "job-sync" + + @pytest.mark.skipif( not is_live_and_not_recording(), reason="Skipped when using recordings due to flakiness of recording blob storage calls", diff --git a/sdk/ai/azure-ai-projects/tests/datasets/test_datasets_async.py b/sdk/ai/azure-ai-projects/tests/datasets/test_datasets_async.py index ac2770ddb2a5..72d0b992d7b4 100644 --- a/sdk/ai/azure-ai-projects/tests/datasets/test_datasets_async.py +++ b/sdk/ai/azure-ai-projects/tests/datasets/test_datasets_async.py @@ -5,12 +5,15 @@ # ------------------------------------ import os import re +from unittest.mock import AsyncMock, MagicMock + import pytest from test_base import TestBase, servicePreparer from devtools_testutils.aio import recorded_by_proxy_async from devtools_testutils import is_live, is_live_and_not_recording, add_general_regex_sanitizer from azure.ai.projects.aio import AIProjectClient -from azure.ai.projects.models import DatasetVersion, DatasetType +from azure.ai.projects.aio.operations._patch_datasets_async import BetaDatasetsOperations +from azure.ai.projects.models import AsyncDatasetGenerationLROPoller, DatasetVersion, DatasetType from azure.ai.projects.models._enums import ConnectionType from azure.core.exceptions import HttpResponseError @@ -21,6 +24,29 @@ data_file2 = os.path.join(data_folder, "data_file2.txt") +@pytest.mark.asyncio +async def test_begin_create_generation_job_exposes_job_id_async(): + """The async create operation exposes its job ID without SDK polling.""" + operation = BetaDatasetsOperations.__new__(BetaDatasetsOperations) + operation._client = MagicMock() # pylint: disable=protected-access + operation._config = MagicMock(polling_interval=0) # pylint: disable=protected-access + operation._serialize = MagicMock() # pylint: disable=protected-access + operation._serialize.url.return_value = "https://example.test" # pylint: disable=protected-access + operation._deserialize = MagicMock() # pylint: disable=protected-access + + initial_response = MagicMock() + initial_response.http_response.json.return_value = {"id": "job-async"} + initial_response.http_response.read = AsyncMock() + operation._create_generation_job_initial = AsyncMock( # pylint: disable=protected-access + return_value=initial_response + ) + + poller = await operation.begin_create_generation_job(job={}, polling=False) + + assert isinstance(poller, AsyncDatasetGenerationLROPoller) + assert poller.details["job_id"] == "job-async" + + @pytest.mark.skipif( not is_live_and_not_recording(), reason="Skipped when using recordings due to flakiness of recording blob storage calls", From 2b3bc97da0fb5ea2c665b8e556711af33b5534da Mon Sep 17 00:00:00 2001 From: Darren Cohen <39422044+dargilco@users.noreply.github.com> Date: Wed, 5 Aug 2026 16:41:47 -0700 Subject: [PATCH 3/8] Updates --- ...mpleqna_for_finetuning_with_app_polling.py | 2 +- ...a_for_finetuning_with_app_polling_async.py | 241 ++++++++++++++++++ 2 files changed, 242 insertions(+), 1 deletion(-) create mode 100644 sdk/ai/azure-ai-projects/samples/datasets/sample_dataset_generation_job_simpleqna_for_finetuning_with_app_polling_async.py diff --git a/sdk/ai/azure-ai-projects/samples/datasets/sample_dataset_generation_job_simpleqna_for_finetuning_with_app_polling.py b/sdk/ai/azure-ai-projects/samples/datasets/sample_dataset_generation_job_simpleqna_for_finetuning_with_app_polling.py index 3844c28e3c62..d7d00abe9bff 100644 --- a/sdk/ai/azure-ai-projects/samples/datasets/sample_dataset_generation_job_simpleqna_for_finetuning_with_app_polling.py +++ b/sdk/ai/azure-ai-projects/samples/datasets/sample_dataset_generation_job_simpleqna_for_finetuning_with_app_polling.py @@ -26,7 +26,7 @@ Before running the sample: - pip install "azure-ai-projects>=2.4.0" azure-identity openai python-dotenv + pip install "azure-ai-projects>=2.5.0" azure-identity openai python-dotenv Set these environment variables with your own values: 1) FOUNDRY_PROJECT_ENDPOINT - Required. The Azure AI Project endpoint, as found diff --git a/sdk/ai/azure-ai-projects/samples/datasets/sample_dataset_generation_job_simpleqna_for_finetuning_with_app_polling_async.py b/sdk/ai/azure-ai-projects/samples/datasets/sample_dataset_generation_job_simpleqna_for_finetuning_with_app_polling_async.py new file mode 100644 index 000000000000..52af02fed04b --- /dev/null +++ b/sdk/ai/azure-ai-projects/samples/datasets/sample_dataset_generation_job_simpleqna_for_finetuning_with_app_polling_async.py @@ -0,0 +1,241 @@ +# pylint: disable=line-too-long,useless-suppression +# ------------------------------------ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. +# ------------------------------------ + +""" +DESCRIPTION: + Generates supervised fine-tuning data from a Markdown reference document + uploaded as an Azure OpenAI File. The sample: + + 1. Uploads a short reference document via the Azure OpenAI Files API + (`purpose=user_data`) so it can be referenced by file id. + 2. Creates a `DataGenerationJob` (scenario=SUPERVISED_FINETUNING, + type=simple_qna) without SDK polling. + 3. Polls the job asynchronously from application code until it reaches a + terminal state, then prints every generated file output. + 4. Cleans up the generated fine-tuning files and the Azure OpenAI input file. + + `simple_qna` REQUIRES `model_options` — the service uses the configured LLM + to synthesize the QnA pairs. Setting `train_split` triggers a split of + the generated samples into two Azure OpenAI output files. + +USAGE: + python sample_dataset_generation_job_simpleqna_for_finetuning_with_app_polling_async.py + + Before running the sample: + + pip install "azure-ai-projects>=2.5.0" azure-identity openai python-dotenv aiohttp + + Set these environment variables with your own values: + 1) FOUNDRY_PROJECT_ENDPOINT - Required. The Azure AI Project endpoint, as found + in the overview page of your Microsoft Foundry project. + 2) FOUNDRY_MODEL_NAME - Required. The name of an Azure OpenAI model + deployment used to synthesize the QnA samples. For `simple_qna` fine-tuning, + the deployment must support the chat completions API (e.g. `gpt-4o`, `gpt-4.1`). + 3) DATASET_NAME - Optional. Name to assign to the generated output files + (used as the file name prefix). Defaults to `simpleqna-finetuning-sample`. + The service caps the rendered output name at 50 characters, so keep + custom values short — the sample appends a unique run id suffix. + 4) POLL_INTERVAL_SECONDS - Optional. Number of seconds to sleep between status + polls for the data generation job. Defaults to 10. +""" + +import asyncio +import os +import uuid +from datetime import datetime, timezone + +from dotenv import load_dotenv + +from azure.identity.aio import DefaultAzureCredential +from azure.ai.projects.aio import AIProjectClient +from azure.ai.projects.models import ( + DataGenerationJob, + DataGenerationJobInputs, + DataGenerationJobOutputOptions, + DataGenerationJobScenario, + DataGenerationModelOptions, + FileDataGenerationJobOutput, + FileDataGenerationJobSource, + JobStatus, + SimpleQnADataGenerationJobOptions, + SimpleQnAFineTuningQuestionType, +) + +load_dotenv() + +endpoint = os.environ["FOUNDRY_PROJECT_ENDPOINT"] +model_name = os.environ["FOUNDRY_MODEL_NAME"] +dataset_name = os.environ.get("DATASET_NAME", "simpleqna-finetuning-sample") +poll_interval_seconds = int(os.environ.get("POLL_INTERVAL_SECONDS", "10")) + +TERMINAL_STATUSES = {JobStatus.SUCCEEDED, JobStatus.FAILED, JobStatus.CANCELLED} + +# Unique per-run output name so repeated runs do not collide. +# Output names are capped at 50 characters by the service. +run_id = f"{datetime.now(tz=timezone.utc).strftime('%y%m%d%H%M%S')}-{uuid.uuid4().hex[:4]}" +output_name = f"{dataset_name}-{run_id}" +if len(output_name) > 50: + raise ValueError( + f"Output name `{output_name}` exceeds the 50-character service limit. " + f"Lower DATASET_NAME (currently `{dataset_name}`) so that `-` fits within 50 characters." + ) + +# Reference document the sample uploads as an Azure OpenAI file. The service +# requires the file to contain at least 1 KB of content to generate QnA from. +SEED_REFERENCE_DOCUMENT = """# Widgets and Gizmos Reference + +## Products +- Widget: blue, manufactured at Factory 7 in Acme, carbon-fiber, rated to 80 C, sold in packs of 4, 250 g each. +- Gizmo: red, manufactured at Factory 12 in Bedrock, carbon-fiber, rated to 80 C, sold individually, 1.2 kg each. +- Sprocket: green, manufactured at Factory 3 in Acme, stainless steel, rated to 200 C, sold individually, 500 g each. + +## Operations +- Factory operates weekdays 0700-1900 local time. +- Closed on public holidays, except for the annual maintenance run on December 27. +- ISO 9001 certified; audited annually by an independent third party. +- Quality control samples every 100th unit and runs full destructive testing on every 5000th unit. + +## Customer support +- Warranty claims: email support@example.com with the serial number printed on the underside of the product. +- Returns: accepted within 30 days if unopened; opened items are eligible for repair only. +- Bulk orders (50+ units): contact sales@example.com for volume pricing and an extended 90-day return window. +- Replacement parts: orderable directly from the support portal using the original order number. + +## Pricing and SLAs +- Widget pack: USD 24.99 per 4-pack; free shipping on orders over USD 75. +- Gizmo unit: USD 49.99; free shipping on orders over USD 75. +- Sprocket unit: USD 14.99; ships from regional warehouses in 1-2 business days. +- Standard support response: within one business day. Priority support response: within four hours. +""" + + +async def main() -> None: + async with ( + DefaultAzureCredential() as credential, + AIProjectClient(endpoint=endpoint, credential=credential) as project_client, + project_client.get_openai_client() as openai_client, + ): + + # ------------------------------------------------------------------ + # 1. Upload the seed reference document as an Azure OpenAI file. + # ------------------------------------------------------------------ + seed_filename = f"widgets-gizmos-seed-{run_id}.md" + print(f"Upload the seed reference document as Azure OpenAI file `{seed_filename}`.") + seed_file = await openai_client.files.create( + file=(seed_filename, SEED_REFERENCE_DOCUMENT.encode("utf-8"), "text/markdown"), + purpose="user_data", + ) + print(f"Uploaded Azure OpenAI file (id: {seed_file.id}).") + + # Wait for the file to finish processing — the data generation service + # rejects references to files that are not yet in the `processed` state. + print("Wait for the Azure OpenAI file to be processed.", end="", flush=True) + while seed_file.status not in ("processed", "error"): + await asyncio.sleep(2) + seed_file = await openai_client.files.retrieve(file_id=seed_file.id) + print(".", end="", flush=True) + print() + if seed_file.status != "processed": + raise RuntimeError(f"Azure OpenAI file `{seed_file.id}` failed to process: status=`{seed_file.status}`.") + + # ------------------------------------------------------------------ + # 2. Submit a fine-tuning data generation job without SDK polling. + # ------------------------------------------------------------------ + job = DataGenerationJob( + inputs=DataGenerationJobInputs( + name=f"simpleqna-finetuning-{run_id}", + scenario=DataGenerationJobScenario.SUPERVISED_FINETUNING, + sources=[ + FileDataGenerationJobSource( + description="Widgets & Gizmos product / operations reference (Azure OpenAI file).", + id=seed_file.id, + ), + ], + options=SimpleQnADataGenerationJobOptions( + # Service requires max_samples to be between 15 and 1000. + max_samples=15, + # `simple_qna` REQUIRES model_options. + model_options=DataGenerationModelOptions(model=model_name), + # Split generated samples 80% training / 20% validation. + train_split=0.8, + # Ask for both short-answer and long-answer questions. + question_types=[ + SimpleQnAFineTuningQuestionType.SHORT_ANSWER, + SimpleQnAFineTuningQuestionType.LONG_ANSWER, + ], + ), + output_options=DataGenerationJobOutputOptions(name=output_name), + ), + ) + + print("Create a dataset generation job without SDK polling.") + poller = await project_client.beta.datasets.begin_create_generation_job( + job=job, + polling=False, + ) + job_id = poller.details["job_id"] + if not job_id: + raise RuntimeError("The create operation did not return a data generation job ID.") + job = await project_client.beta.datasets.get_generation_job(job_id=job_id) + print(f"Created job: id={job.id}, status={job.status}") + + # ------------------------------------------------------------------ + # 3. Poll from application code until the job reaches a terminal state. + # ------------------------------------------------------------------ + print(f"Polling job `{job.id}` to completion...", end="", flush=True) + while job.status not in TERMINAL_STATUSES: + await asyncio.sleep(poll_interval_seconds) + job = await project_client.beta.datasets.get_generation_job(job_id=job.id) + print(".", end="", flush=True) + print() + print(f"Final job status: `{job.status}`.") + + if job.status == JobStatus.FAILED: + message = job.error.message if job.error else "" + raise RuntimeError(f"Data generation job `{job.id}` failed: {message}") + if job.status == JobStatus.CANCELLED: + raise RuntimeError(f"Data generation job `{job.id}` was cancelled.") + if job.result is None: + raise RuntimeError(f"Data generation job `{job.id}` completed without a result.") + + job_result = job.result + print(f"Data generation result: {job_result}") + + # ------------------------------------------------------------------ + # 4. Inspect the generated fine-tuning file outputs. + # ------------------------------------------------------------------ + # `train_split=0.8` produces two Azure OpenAI files: a training partition + # and a validation partition. Both are emitted as FileDataGenerationJobOutput + # entries in `job_result.outputs`. + file_outputs = [ + output for output in (job_result.outputs or []) if isinstance(output, FileDataGenerationJobOutput) + ] + if not file_outputs: + raise RuntimeError("The data generation job did not produce any file outputs.") + + print(f"Generated {len(file_outputs)} fine-tuning file(s):") + for output in file_outputs: + if not output.id: + raise RuntimeError("A file output was returned without an id.") + # Resolve the Azure OpenAI file to surface its real filename and size. + file_info = await openai_client.files.retrieve(file_id=output.id) + print(f" - filename=`{file_info.filename}` id=`{output.id}` bytes={file_info.bytes}") + if job_result.generated_samples is not None: + print(f"Generated samples: {job_result.generated_samples}") + + # ------------------------------------------------------------------ + # 5. Clean up. + # ------------------------------------------------------------------ + for output in file_outputs: + print(f"Delete the generated Azure OpenAI file `{output.id}`.") + await openai_client.files.delete(file_id=output.id) + + print(f"Delete the Azure OpenAI input file `{seed_file.id}`.") + await openai_client.files.delete(file_id=seed_file.id) + + +if __name__ == "__main__": + asyncio.run(main()) From 309b8f56422229189c122b1868943f9b9868e1f9 Mon Sep 17 00:00:00 2001 From: Darren Cohen <39422044+dargilco@users.noreply.github.com> Date: Wed, 5 Aug 2026 17:13:55 -0700 Subject: [PATCH 4/8] Evaluator poller --- .../ai/projects/aio/operations/_patch.py | 4 +- .../aio/operations/_patch_evaluators_async.py | 144 ++++++++++++++++++ .../azure/ai/projects/models/_patch.py | 77 +++++++++- .../azure/ai/projects/operations/_patch.py | 4 +- .../projects/operations/_patch_evaluators.py | 142 +++++++++++++++++ .../tests/evaluators/test_evaluators.py | 31 ++++ .../tests/evaluators/test_evaluators_async.py | 35 +++++ 7 files changed, 432 insertions(+), 5 deletions(-) create mode 100644 sdk/ai/azure-ai-projects/azure/ai/projects/aio/operations/_patch_evaluators_async.py create mode 100644 sdk/ai/azure-ai-projects/azure/ai/projects/operations/_patch_evaluators.py create mode 100644 sdk/ai/azure-ai-projects/tests/evaluators/test_evaluators.py create mode 100644 sdk/ai/azure-ai-projects/tests/evaluators/test_evaluators_async.py diff --git a/sdk/ai/azure-ai-projects/azure/ai/projects/aio/operations/_patch.py b/sdk/ai/azure-ai-projects/azure/ai/projects/aio/operations/_patch.py index d734c8c691da..336547835829 100644 --- a/sdk/ai/azure-ai-projects/azure/ai/projects/aio/operations/_patch.py +++ b/sdk/ai/azure-ai-projects/azure/ai/projects/aio/operations/_patch.py @@ -11,6 +11,7 @@ from typing import Any, List from ._patch_agents_async import AgentsOperations from ._patch_datasets_async import BetaDatasetsOperations, DatasetsOperations +from ._patch_evaluators_async import BetaEvaluatorsOperations from ._patch_evaluation_rules_async import EvaluationRulesOperations from ._patch_telemetry_async import TelemetryOperations from ._patch_connections_async import ConnectionsOperations @@ -20,7 +21,6 @@ from ._operations import ( BetaAgentsOperations, BetaEvaluationTaxonomiesOperations, - BetaEvaluatorsOperations, BetaInsightsOperations, BetaOperations as GeneratedBetaOperations, BetaRedTeamsOperations, @@ -65,7 +65,7 @@ class BetaOperations(GeneratedBetaOperations): def __init__(self, *args: Any, **kwargs: Any) -> None: super().__init__(*args, **kwargs) - # Replace with patched class that includes upload() + # Replace with patched class that returns AsyncEvaluatorGenerationLROPoller self.evaluators = BetaEvaluatorsOperations(self._client, self._config, self._serialize, self._deserialize) # Replace with patched class that adds file-path overload to upload_session_file self.agents = BetaAgentsOperations(self._client, self._config, self._serialize, self._deserialize) diff --git a/sdk/ai/azure-ai-projects/azure/ai/projects/aio/operations/_patch_evaluators_async.py b/sdk/ai/azure-ai-projects/azure/ai/projects/aio/operations/_patch_evaluators_async.py new file mode 100644 index 000000000000..f23b8d13ac72 --- /dev/null +++ b/sdk/ai/azure-ai-projects/azure/ai/projects/aio/operations/_patch_evaluators_async.py @@ -0,0 +1,144 @@ +# pylint: disable=line-too-long,useless-suppression +# ------------------------------------ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. +# ------------------------------------ +"""Custom async evaluator operations.""" + +from collections.abc import MutableMapping +from typing import Any, IO, Optional, Union, cast, overload + +from azure.core.polling import AsyncNoPolling, AsyncPollingMethod +from azure.core.polling.async_base_polling import AsyncLROBasePolling +from azure.core.tracing.decorator_async import distributed_trace_async +from azure.core.utils import case_insensitive_dict + +from ._operations import BetaEvaluatorsOperations as BetaEvaluatorsOperationsGenerated +from ... import models as _models +from ..._utils.model_base import _deserialize +from ...models import AsyncEvaluatorGenerationLROPoller + +JSON = MutableMapping[str, Any] + + +class BetaEvaluatorsOperations(BetaEvaluatorsOperationsGenerated): + """Custom async operations for beta evaluator generation jobs.""" + + @overload + async def begin_create_generation_job( + self, + job: _models.EvaluatorGenerationJob, + *, + operation_id: Optional[str] = None, + content_type: str = "application/json", + **kwargs: Any, + ) -> AsyncEvaluatorGenerationLROPoller: ... + + @overload + async def begin_create_generation_job( + self, + job: JSON, + *, + operation_id: Optional[str] = None, + content_type: str = "application/json", + **kwargs: Any, + ) -> AsyncEvaluatorGenerationLROPoller: ... + + @overload + async def begin_create_generation_job( + self, + job: IO[bytes], + *, + operation_id: Optional[str] = None, + content_type: str = "application/json", + **kwargs: Any, + ) -> AsyncEvaluatorGenerationLROPoller: ... + + @distributed_trace_async + async def begin_create_generation_job( + self, + job: Union[_models.EvaluatorGenerationJob, JSON, IO[bytes]], + *, + operation_id: Optional[str] = None, + **kwargs: Any, + ) -> AsyncEvaluatorGenerationLROPoller: + """Create an evaluator generation job. + + :param job: The job to create. Required. + :type job: ~azure.ai.projects.models.EvaluatorGenerationJob or JSON or IO[bytes] + :keyword operation_id: Client-generated unique ID for idempotent retries. When absent, the + server creates the job unconditionally. Default value is None. + :paramtype operation_id: str + :return: A poller that returns EvaluatorVersion and exposes the job ID in ``details``. + :rtype: ~azure.ai.projects.models.AsyncEvaluatorGenerationLROPoller + :raises ~azure.core.exceptions.HttpResponseError: + """ + headers = case_insensitive_dict(kwargs.pop("headers", {}) or {}) + params = kwargs.pop("params", {}) or {} + + content_type: Optional[str] = kwargs.pop( + "content_type", headers.pop("Content-Type", None) + ) + cls = kwargs.pop("cls", None) + polling: Union[bool, AsyncPollingMethod] = kwargs.pop("polling", True) + lro_delay = kwargs.pop("polling_interval", self._config.polling_interval) + continuation_token: Optional[str] = kwargs.pop("continuation_token", None) + raw_result = None + if continuation_token is None: + raw_result = await self._create_generation_job_initial( + job=job, + operation_id=operation_id, + content_type=content_type, + cls=lambda x, y, z: x, + headers=headers, + params=params, + **kwargs, + ) + await raw_result.http_response.read() # type: ignore + kwargs.pop("error_map", None) + + def get_long_running_output(pipeline_response): + response_headers = {} + response = pipeline_response.http_response + response_headers["Operation-Location"] = self._deserialize( + "str", response.headers.get("Operation-Location") + ) + response_headers["Location"] = self._deserialize( + "str", response.headers.get("Location") + ) + + deserialized = _deserialize( + _models.EvaluatorVersion, response.json().get("result", {}) + ) + if cls: + return cls(pipeline_response, deserialized, response_headers) + return deserialized + + path_format_arguments = { + "endpoint": self._serialize.url( + "self._config.endpoint", self._config.endpoint, "str", skip_quote=True + ), + } + + if polling is True: + polling_method: AsyncPollingMethod = cast( + AsyncPollingMethod, + AsyncLROBasePolling( + lro_delay, path_format_arguments=path_format_arguments, **kwargs + ), + ) + elif polling is False: + polling_method = cast(AsyncPollingMethod, AsyncNoPolling()) + else: + polling_method = polling + if continuation_token: + return AsyncEvaluatorGenerationLROPoller.from_continuation_token( + polling_method=polling_method, + continuation_token=continuation_token, + client=self._client, + deserialization_callback=get_long_running_output, + ) + assert raw_result is not None + return AsyncEvaluatorGenerationLROPoller( # type: ignore + self._client, raw_result, get_long_running_output, polling_method + ) diff --git a/sdk/ai/azure-ai-projects/azure/ai/projects/models/_patch.py b/sdk/ai/azure-ai-projects/azure/ai/projects/models/_patch.py index bba5cad64ab2..a6e3b0993605 100644 --- a/sdk/ai/azure-ai-projects/azure/ai/projects/models/_patch.py +++ b/sdk/ai/azure-ai-projects/azure/ai/projects/models/_patch.py @@ -33,7 +33,7 @@ TracesPreviewEvalRunDataSource, ) from ._models import CustomCredential as CustomCredentialGenerated -from ..models import DataGenerationJobResult, MemoryStoreUpdateCompletedResult, MemoryStoreUpdateResult +from ..models import DataGenerationJobResult, EvaluatorVersion, MemoryStoreUpdateCompletedResult, MemoryStoreUpdateResult from ._enums import _FoundryFeaturesOptInKeys, _AgentDefinitionOptInKeys _FOUNDRY_FEATURES_HEADER_NAME: Final[str] = "Foundry-Features" @@ -460,8 +460,82 @@ def from_continuation_token( return cls(client, initial_response, deserialization_callback, polling_method) +class EvaluatorGenerationLROPoller(LROPoller[EvaluatorVersion]): + """Custom LROPoller for evaluator generation job operations.""" + + def __init__(self, client: Any, initial_response: Any, deserialization_callback: Any, polling_method: Any) -> None: + self._job_id = DatasetGenerationLROPoller._get_job_id(initial_response) + super().__init__(client, initial_response, deserialization_callback, polling_method) + + @property + def details(self) -> Mapping[str, Any]: + """Returns metadata associated with the evaluator generation job operation. + + :return: A mapping containing the created evaluator generation job ID. + :rtype: Mapping[str, Any] + """ + return {"job_id": self._job_id} + + @classmethod + def from_continuation_token( + cls, polling_method: PollingMethod[EvaluatorVersion], continuation_token: str, **kwargs: Any + ) -> "EvaluatorGenerationLROPoller": + """Create a poller from a continuation token. + + :param polling_method: The polling strategy to adopt. + :type polling_method: ~azure.core.polling.PollingMethod + :param continuation_token: An opaque continuation token. + :type continuation_token: str + :return: An instance of EvaluatorGenerationLROPoller. + :rtype: EvaluatorGenerationLROPoller + """ + client, initial_response, deserialization_callback = polling_method.from_continuation_token( + continuation_token, **kwargs + ) + return cls(client, initial_response, deserialization_callback, polling_method) + + +class AsyncEvaluatorGenerationLROPoller(AsyncLROPoller[EvaluatorVersion]): + """Custom AsyncLROPoller for evaluator generation job operations.""" + + def __init__(self, client: Any, initial_response: Any, deserialization_callback: Any, polling_method: Any) -> None: + super().__init__(client, initial_response, deserialization_callback, polling_method) + self._job_id = DatasetGenerationLROPoller._get_job_id(initial_response) + + @property + def details(self) -> Mapping[str, Any]: + """Returns metadata associated with the evaluator generation job operation. + + :return: A mapping containing the created evaluator generation job ID. + :rtype: Mapping[str, Any] + """ + return {"job_id": self._job_id} + + @classmethod + def from_continuation_token( + cls, + polling_method: AsyncPollingMethod[EvaluatorVersion], + continuation_token: str, + **kwargs: Any, + ) -> "AsyncEvaluatorGenerationLROPoller": + """Create a poller from a continuation token. + + :param polling_method: The polling strategy to adopt. + :type polling_method: ~azure.core.polling.AsyncPollingMethod + :param continuation_token: An opaque continuation token. + :type continuation_token: str + :return: An instance of AsyncEvaluatorGenerationLROPoller. + :rtype: AsyncEvaluatorGenerationLROPoller + """ + client, initial_response, deserialization_callback = polling_method.from_continuation_token( + continuation_token, **kwargs + ) + return cls(client, initial_response, deserialization_callback, polling_method) + + __all__: List[str] = [ "AsyncDatasetGenerationLROPoller", + "AsyncEvaluatorGenerationLROPoller", "AsyncUpdateMemoriesLROPoller", "AzureAIAgentTargetParam", "AzureAIBenchmarkPreviewEvalRunDataSource", @@ -470,6 +544,7 @@ def from_continuation_token( "AzureAIResponsesEvalRunDataSource", "CustomCredential", "DatasetGenerationLROPoller", + "EvaluatorGenerationLROPoller", "EvalCsvFileIdSource", "EvalCsvRunDataSource", "TestingCriterionAzureAIEvaluator", diff --git a/sdk/ai/azure-ai-projects/azure/ai/projects/operations/_patch.py b/sdk/ai/azure-ai-projects/azure/ai/projects/operations/_patch.py index 49559c1b2e1d..8e88a3e7eafa 100644 --- a/sdk/ai/azure-ai-projects/azure/ai/projects/operations/_patch.py +++ b/sdk/ai/azure-ai-projects/azure/ai/projects/operations/_patch.py @@ -14,6 +14,7 @@ from ..models._patch import _FOUNDRY_FEATURES_HEADER_NAME, _BETA_OPERATION_FEATURE_HEADERS, _has_header_case_insensitive from ._patch_agents import AgentsOperations from ._patch_datasets import BetaDatasetsOperations, DatasetsOperations +from ._patch_evaluators import BetaEvaluatorsOperations from ._patch_evaluation_rules import EvaluationRulesOperations from ._patch_telemetry import TelemetryOperations from ._patch_connections import ConnectionsOperations @@ -22,7 +23,6 @@ from ._operations import ( BetaAgentsOperations, BetaEvaluationTaxonomiesOperations, - BetaEvaluatorsOperations, BetaInsightsOperations, BetaOperations as GeneratedBetaOperations, BetaRedTeamsOperations, @@ -120,7 +120,7 @@ class BetaOperations(GeneratedBetaOperations): def __init__(self, *args: Any, **kwargs: Any) -> None: super().__init__(*args, **kwargs) - # Replace with patched class that includes upload() + # Replace with patched class that returns EvaluatorGenerationLROPoller self.evaluators = BetaEvaluatorsOperations(self._client, self._config, self._serialize, self._deserialize) # Replace with patched class that adds file-path overload to upload_session_file self.agents = BetaAgentsOperations(self._client, self._config, self._serialize, self._deserialize) diff --git a/sdk/ai/azure-ai-projects/azure/ai/projects/operations/_patch_evaluators.py b/sdk/ai/azure-ai-projects/azure/ai/projects/operations/_patch_evaluators.py new file mode 100644 index 000000000000..143b1dc66283 --- /dev/null +++ b/sdk/ai/azure-ai-projects/azure/ai/projects/operations/_patch_evaluators.py @@ -0,0 +1,142 @@ +# pylint: disable=line-too-long,useless-suppression +# ------------------------------------ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. +# ------------------------------------ +"""Custom evaluator operations.""" + +from collections.abc import MutableMapping +from typing import Any, IO, Optional, Union, cast, overload + +from azure.core.polling import NoPolling, PollingMethod +from azure.core.polling.base_polling import LROBasePolling +from azure.core.tracing.decorator import distributed_trace +from azure.core.utils import case_insensitive_dict + +from ._operations import BetaEvaluatorsOperations as BetaEvaluatorsOperationsGenerated +from .. import models as _models +from .._utils.model_base import _deserialize +from ..models import EvaluatorGenerationLROPoller + +JSON = MutableMapping[str, Any] + + +class BetaEvaluatorsOperations(BetaEvaluatorsOperationsGenerated): + """Custom operations for beta evaluator generation jobs.""" + + @overload + def begin_create_generation_job( + self, + job: _models.EvaluatorGenerationJob, + *, + operation_id: Optional[str] = None, + content_type: str = "application/json", + **kwargs: Any, + ) -> EvaluatorGenerationLROPoller: ... + + @overload + def begin_create_generation_job( + self, + job: JSON, + *, + operation_id: Optional[str] = None, + content_type: str = "application/json", + **kwargs: Any, + ) -> EvaluatorGenerationLROPoller: ... + + @overload + def begin_create_generation_job( + self, + job: IO[bytes], + *, + operation_id: Optional[str] = None, + content_type: str = "application/json", + **kwargs: Any, + ) -> EvaluatorGenerationLROPoller: ... + + @distributed_trace + def begin_create_generation_job( + self, + job: Union[_models.EvaluatorGenerationJob, JSON, IO[bytes]], + *, + operation_id: Optional[str] = None, + **kwargs: Any, + ) -> EvaluatorGenerationLROPoller: + """Create an evaluator generation job. + + :param job: The job to create. Required. + :type job: ~azure.ai.projects.models.EvaluatorGenerationJob or JSON or IO[bytes] + :keyword operation_id: Client-generated unique ID for idempotent retries. When absent, the + server creates the job unconditionally. Default value is None. + :paramtype operation_id: str + :return: A poller that returns EvaluatorVersion and exposes the job ID in ``details``. + :rtype: ~azure.ai.projects.models.EvaluatorGenerationLROPoller + :raises ~azure.core.exceptions.HttpResponseError: + """ + headers = case_insensitive_dict(kwargs.pop("headers", {}) or {}) + params = kwargs.pop("params", {}) or {} + + content_type: Optional[str] = kwargs.pop( + "content_type", headers.pop("Content-Type", None) + ) + cls = kwargs.pop("cls", None) + polling: Union[bool, PollingMethod] = kwargs.pop("polling", True) + lro_delay = kwargs.pop("polling_interval", self._config.polling_interval) + continuation_token: Optional[str] = kwargs.pop("continuation_token", None) + raw_result = None + if continuation_token is None: + raw_result = self._create_generation_job_initial( + job=job, + operation_id=operation_id, + content_type=content_type, + cls=lambda x, y, z: x, + headers=headers, + params=params, + **kwargs, + ) + raw_result.http_response.read() # type: ignore + kwargs.pop("error_map", None) + + def get_long_running_output(pipeline_response): + response_headers = {} + response = pipeline_response.http_response + response_headers["Operation-Location"] = self._deserialize( + "str", response.headers.get("Operation-Location") + ) + response_headers["Location"] = self._deserialize( + "str", response.headers.get("Location") + ) + + deserialized = _deserialize( + _models.EvaluatorVersion, response.json().get("result", {}) + ) + if cls: + return cls(pipeline_response, deserialized, response_headers) + return deserialized + + path_format_arguments = { + "endpoint": self._serialize.url( + "self._config.endpoint", self._config.endpoint, "str", skip_quote=True + ), + } + + if polling is True: + polling_method: PollingMethod = cast( + PollingMethod, + LROBasePolling( + lro_delay, path_format_arguments=path_format_arguments, **kwargs + ), + ) + elif polling is False: + polling_method = cast(PollingMethod, NoPolling()) + else: + polling_method = polling + if continuation_token: + return EvaluatorGenerationLROPoller.from_continuation_token( + polling_method=polling_method, + continuation_token=continuation_token, + client=self._client, + deserialization_callback=get_long_running_output, + ) + assert raw_result is not None + return EvaluatorGenerationLROPoller(self._client, raw_result, get_long_running_output, polling_method) # type: ignore diff --git a/sdk/ai/azure-ai-projects/tests/evaluators/test_evaluators.py b/sdk/ai/azure-ai-projects/tests/evaluators/test_evaluators.py new file mode 100644 index 000000000000..de4bb5ffc2c6 --- /dev/null +++ b/sdk/ai/azure-ai-projects/tests/evaluators/test_evaluators.py @@ -0,0 +1,31 @@ +# ------------------------------------ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. +# ------------------------------------ +"""Unit tests for sync evaluator generation pollers.""" + +from unittest.mock import MagicMock + +from azure.ai.projects.models import EvaluatorGenerationLROPoller +from azure.ai.projects.operations._patch_evaluators import BetaEvaluatorsOperations + + +def test_begin_create_generation_job_exposes_job_id(): + """The sync create operation exposes its job ID without SDK polling.""" + operation = BetaEvaluatorsOperations.__new__(BetaEvaluatorsOperations) + operation._client = MagicMock() # pylint: disable=protected-access + operation._config = MagicMock(polling_interval=0) # pylint: disable=protected-access + operation._serialize = MagicMock() # pylint: disable=protected-access + operation._serialize.url.return_value = "https://example.test" # pylint: disable=protected-access + operation._deserialize = MagicMock() # pylint: disable=protected-access + + initial_response = MagicMock() + initial_response.http_response.json.return_value = {"id": "evaluator-job-sync"} + operation._create_generation_job_initial = MagicMock( # pylint: disable=protected-access + return_value=initial_response + ) + + poller = operation.begin_create_generation_job(job={}, polling=False) + + assert isinstance(poller, EvaluatorGenerationLROPoller) + assert poller.details["job_id"] == "evaluator-job-sync" \ No newline at end of file diff --git a/sdk/ai/azure-ai-projects/tests/evaluators/test_evaluators_async.py b/sdk/ai/azure-ai-projects/tests/evaluators/test_evaluators_async.py new file mode 100644 index 000000000000..2e74a90110f3 --- /dev/null +++ b/sdk/ai/azure-ai-projects/tests/evaluators/test_evaluators_async.py @@ -0,0 +1,35 @@ +# ------------------------------------ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. +# ------------------------------------ +"""Unit tests for async evaluator generation pollers.""" + +from unittest.mock import AsyncMock, MagicMock + +import pytest + +from azure.ai.projects.aio.operations._patch_evaluators_async import BetaEvaluatorsOperations +from azure.ai.projects.models import AsyncEvaluatorGenerationLROPoller + + +@pytest.mark.asyncio +async def test_begin_create_generation_job_exposes_job_id_async(): + """The async create operation exposes its job ID without SDK polling.""" + operation = BetaEvaluatorsOperations.__new__(BetaEvaluatorsOperations) + operation._client = MagicMock() # pylint: disable=protected-access + operation._config = MagicMock(polling_interval=0) # pylint: disable=protected-access + operation._serialize = MagicMock() # pylint: disable=protected-access + operation._serialize.url.return_value = "https://example.test" # pylint: disable=protected-access + operation._deserialize = MagicMock() # pylint: disable=protected-access + + initial_response = MagicMock() + initial_response.http_response.json.return_value = {"id": "evaluator-job-async"} + initial_response.http_response.read = AsyncMock() + operation._create_generation_job_initial = AsyncMock( # pylint: disable=protected-access + return_value=initial_response + ) + + poller = await operation.begin_create_generation_job(job={}, polling=False) + + assert isinstance(poller, AsyncEvaluatorGenerationLROPoller) + assert poller.details["job_id"] == "evaluator-job-async" \ No newline at end of file From 8e57421ecdff962e2635c0549c3337652363c8da Mon Sep 17 00:00:00 2001 From: Darren Cohen <39422044+dargilco@users.noreply.github.com> Date: Wed, 5 Aug 2026 17:35:34 -0700 Subject: [PATCH 5/8] agent optimization custom lro poller --- .../ai/projects/aio/operations/_patch.py | 5 +- .../aio/operations/_patch_agents_async.py | 127 +++++++++++++++++- .../azure/ai/projects/models/_patch.py | 83 +++++++++++- .../azure/ai/projects/operations/_patch.py | 5 +- .../ai/projects/operations/_patch_agents.py | 125 ++++++++++++++++- ...e_optimization_job_advanced_app_polling.py | 19 +-- ...mization_job_advanced_app_polling_async.py | 25 ++-- .../agents/test_agent_optimization_poller.py | 35 +++++ .../test_agent_optimization_poller_async.py | 39 ++++++ 9 files changed, 424 insertions(+), 39 deletions(-) create mode 100644 sdk/ai/azure-ai-projects/tests/agents/test_agent_optimization_poller.py create mode 100644 sdk/ai/azure-ai-projects/tests/agents/test_agent_optimization_poller_async.py diff --git a/sdk/ai/azure-ai-projects/azure/ai/projects/aio/operations/_patch.py b/sdk/ai/azure-ai-projects/azure/ai/projects/aio/operations/_patch.py index 336547835829..5bb74cf4fe6d 100644 --- a/sdk/ai/azure-ai-projects/azure/ai/projects/aio/operations/_patch.py +++ b/sdk/ai/azure-ai-projects/azure/ai/projects/aio/operations/_patch.py @@ -9,7 +9,7 @@ """ from typing import Any, List -from ._patch_agents_async import AgentsOperations +from ._patch_agents_async import AgentsOperations, BetaAgentsOperations from ._patch_datasets_async import BetaDatasetsOperations, DatasetsOperations from ._patch_evaluators_async import BetaEvaluatorsOperations from ._patch_evaluation_rules_async import EvaluationRulesOperations @@ -19,7 +19,6 @@ from ._patch_models_async import BetaModelsOperations from ...operations._patch import _BETA_OPERATION_FEATURE_HEADERS, _OperationMethodHeaderProxy from ._operations import ( - BetaAgentsOperations, BetaEvaluationTaxonomiesOperations, BetaInsightsOperations, BetaOperations as GeneratedBetaOperations, @@ -67,7 +66,7 @@ def __init__(self, *args: Any, **kwargs: Any) -> None: super().__init__(*args, **kwargs) # Replace with patched class that returns AsyncEvaluatorGenerationLROPoller self.evaluators = BetaEvaluatorsOperations(self._client, self._config, self._serialize, self._deserialize) - # Replace with patched class that adds file-path overload to upload_session_file + # Replace with patched class that returns AsyncAgentOptimizationLROPoller self.agents = BetaAgentsOperations(self._client, self._config, self._serialize, self._deserialize) # Replace with patched class that includes begin_update_memories self.memory_stores = BetaMemoryStoresOperations(self._client, self._config, self._serialize, self._deserialize) diff --git a/sdk/ai/azure-ai-projects/azure/ai/projects/aio/operations/_patch_agents_async.py b/sdk/ai/azure-ai-projects/azure/ai/projects/aio/operations/_patch_agents_async.py index cd906a8d8498..8d8dcb180454 100644 --- a/sdk/ai/azure-ai-projects/azure/ai/projects/aio/operations/_patch_agents_async.py +++ b/sdk/ai/azure-ai-projects/azure/ai/projects/aio/operations/_patch_agents_async.py @@ -8,11 +8,21 @@ Follow our quickstart for examples: https://aka.ms/azsdk/python/dpcodegen/python/customize """ -from typing import Union, Optional, Any, IO, overload +from typing import Union, Optional, Any, IO, cast, overload from azure.core.exceptions import HttpResponseError +from azure.core.polling import AsyncNoPolling, AsyncPollingMethod +from azure.core.polling.async_base_polling import AsyncLROBasePolling from azure.core.tracing.decorator_async import distributed_trace_async -from ._operations import AgentsOperations as GeneratedAgentsOperations, JSON, _Unset +from azure.core.utils import case_insensitive_dict +from ._operations import ( + AgentsOperations as GeneratedAgentsOperations, + BetaAgentsOperations as BetaAgentsOperationsGenerated, + JSON, + _Unset, +) from ... import models as _models +from ..._utils.model_base import _deserialize +from ...models import AsyncAgentOptimizationLROPoller from ...operations._patch_agents import _compute_sha256_from_stream from ...models._patch import ( _FOUNDRY_FEATURES_HEADER_NAME, @@ -314,3 +324,116 @@ async def create_version_from_code( new_exc.model = exc.model raise new_exc from exc raise + + +class BetaAgentsOperations(BetaAgentsOperationsGenerated): + """Custom async operations for beta agent optimization jobs.""" + + @overload + async def begin_create_optimization_job( + self, + job: _models.OptimizationJob, + *, + operation_id: Optional[str] = None, + content_type: str = "application/json", + **kwargs: Any, + ) -> AsyncAgentOptimizationLROPoller: ... + + @overload + async def begin_create_optimization_job( + self, + job: JSON, + *, + operation_id: Optional[str] = None, + content_type: str = "application/json", + **kwargs: Any, + ) -> AsyncAgentOptimizationLROPoller: ... + + @overload + async def begin_create_optimization_job( + self, + job: IO[bytes], + *, + operation_id: Optional[str] = None, + content_type: str = "application/json", + **kwargs: Any, + ) -> AsyncAgentOptimizationLROPoller: ... + + @distributed_trace_async + async def begin_create_optimization_job( + self, + job: Union[_models.OptimizationJob, JSON, IO[bytes]], + *, + operation_id: Optional[str] = None, + **kwargs: Any, + ) -> AsyncAgentOptimizationLROPoller: + """Create an agent optimization job. + + :param job: The job to create. Required. + :type job: ~azure.ai.projects.models.OptimizationJob or JSON or IO[bytes] + :keyword operation_id: Client-generated unique ID for idempotent retries. When absent, the + server creates the job unconditionally. Default value is None. + :paramtype operation_id: str + :return: A poller that returns OptimizationJobResult and exposes the job ID in ``details``. + :rtype: ~azure.ai.projects.models.AsyncAgentOptimizationLROPoller + :raises ~azure.core.exceptions.HttpResponseError: + """ + headers = case_insensitive_dict(kwargs.pop("headers", {}) or {}) + params = kwargs.pop("params", {}) or {} + + content_type: Optional[str] = kwargs.pop("content_type", headers.pop("Content-Type", None)) + cls = kwargs.pop("cls", None) + polling: Union[bool, AsyncPollingMethod] = kwargs.pop("polling", True) + lro_delay = kwargs.pop("polling_interval", self._config.polling_interval) + continuation_token: Optional[str] = kwargs.pop("continuation_token", None) + raw_result = None + if continuation_token is None: + raw_result = await self._create_optimization_job_initial( + job=job, + operation_id=operation_id, + content_type=content_type, + cls=lambda x, y, z: x, + headers=headers, + params=params, + **kwargs, + ) + await raw_result.http_response.read() # type: ignore + kwargs.pop("error_map", None) + + def get_long_running_output(pipeline_response): + response_headers = {} + response = pipeline_response.http_response + response_headers["Operation-Location"] = self._deserialize( + "str", response.headers.get("Operation-Location") + ) + response_headers["Location"] = self._deserialize("str", response.headers.get("Location")) + + deserialized = _deserialize(_models.OptimizationJobResult, response.json().get("result", {})) + if cls: + return cls(pipeline_response, deserialized, response_headers) + return deserialized + + path_format_arguments = { + "endpoint": self._serialize.url("self._config.endpoint", self._config.endpoint, "str", skip_quote=True), + } + + if polling is True: + polling_method: AsyncPollingMethod = cast( + AsyncPollingMethod, + AsyncLROBasePolling(lro_delay, path_format_arguments=path_format_arguments, **kwargs), + ) + elif polling is False: + polling_method = cast(AsyncPollingMethod, AsyncNoPolling()) + else: + polling_method = polling + if continuation_token: + return AsyncAgentOptimizationLROPoller.from_continuation_token( + polling_method=polling_method, + continuation_token=continuation_token, + client=self._client, + deserialization_callback=get_long_running_output, + ) + assert raw_result is not None + return AsyncAgentOptimizationLROPoller( # type: ignore + self._client, raw_result, get_long_running_output, polling_method + ) diff --git a/sdk/ai/azure-ai-projects/azure/ai/projects/models/_patch.py b/sdk/ai/azure-ai-projects/azure/ai/projects/models/_patch.py index a6e3b0993605..33d30341ddc3 100644 --- a/sdk/ai/azure-ai-projects/azure/ai/projects/models/_patch.py +++ b/sdk/ai/azure-ai-projects/azure/ai/projects/models/_patch.py @@ -33,7 +33,13 @@ TracesPreviewEvalRunDataSource, ) from ._models import CustomCredential as CustomCredentialGenerated -from ..models import DataGenerationJobResult, EvaluatorVersion, MemoryStoreUpdateCompletedResult, MemoryStoreUpdateResult +from ..models import ( + DataGenerationJobResult, + EvaluatorVersion, + MemoryStoreUpdateCompletedResult, + MemoryStoreUpdateResult, + OptimizationJobResult, +) from ._enums import _FoundryFeaturesOptInKeys, _AgentDefinitionOptInKeys _FOUNDRY_FEATURES_HEADER_NAME: Final[str] = "Foundry-Features" @@ -533,7 +539,82 @@ def from_continuation_token( return cls(client, initial_response, deserialization_callback, polling_method) +class AgentOptimizationLROPoller(LROPoller[OptimizationJobResult]): + """Custom LROPoller for agent optimization job operations.""" + + def __init__(self, client: Any, initial_response: Any, deserialization_callback: Any, polling_method: Any) -> None: + self._job_id = DatasetGenerationLROPoller._get_job_id(initial_response) + super().__init__(client, initial_response, deserialization_callback, polling_method) + + @property + def details(self) -> Mapping[str, Any]: + """Returns metadata associated with the agent optimization job operation. + + :return: A mapping containing the created agent optimization job ID. + :rtype: Mapping[str, Any] + """ + return {"job_id": self._job_id} + + @classmethod + def from_continuation_token( + cls, polling_method: PollingMethod[OptimizationJobResult], continuation_token: str, **kwargs: Any + ) -> "AgentOptimizationLROPoller": + """Create a poller from a continuation token. + + :param polling_method: The polling strategy to adopt. + :type polling_method: ~azure.core.polling.PollingMethod + :param continuation_token: An opaque continuation token. + :type continuation_token: str + :return: An instance of AgentOptimizationLROPoller. + :rtype: AgentOptimizationLROPoller + """ + client, initial_response, deserialization_callback = polling_method.from_continuation_token( + continuation_token, **kwargs + ) + return cls(client, initial_response, deserialization_callback, polling_method) + + +class AsyncAgentOptimizationLROPoller(AsyncLROPoller[OptimizationJobResult]): + """Custom AsyncLROPoller for agent optimization job operations.""" + + def __init__(self, client: Any, initial_response: Any, deserialization_callback: Any, polling_method: Any) -> None: + super().__init__(client, initial_response, deserialization_callback, polling_method) + self._job_id = DatasetGenerationLROPoller._get_job_id(initial_response) + + @property + def details(self) -> Mapping[str, Any]: + """Returns metadata associated with the agent optimization job operation. + + :return: A mapping containing the created agent optimization job ID. + :rtype: Mapping[str, Any] + """ + return {"job_id": self._job_id} + + @classmethod + def from_continuation_token( + cls, + polling_method: AsyncPollingMethod[OptimizationJobResult], + continuation_token: str, + **kwargs: Any, + ) -> "AsyncAgentOptimizationLROPoller": + """Create a poller from a continuation token. + + :param polling_method: The polling strategy to adopt. + :type polling_method: ~azure.core.polling.AsyncPollingMethod + :param continuation_token: An opaque continuation token. + :type continuation_token: str + :return: An instance of AsyncAgentOptimizationLROPoller. + :rtype: AsyncAgentOptimizationLROPoller + """ + client, initial_response, deserialization_callback = polling_method.from_continuation_token( + continuation_token, **kwargs + ) + return cls(client, initial_response, deserialization_callback, polling_method) + + __all__: List[str] = [ + "AgentOptimizationLROPoller", + "AsyncAgentOptimizationLROPoller", "AsyncDatasetGenerationLROPoller", "AsyncEvaluatorGenerationLROPoller", "AsyncUpdateMemoriesLROPoller", diff --git a/sdk/ai/azure-ai-projects/azure/ai/projects/operations/_patch.py b/sdk/ai/azure-ai-projects/azure/ai/projects/operations/_patch.py index 8e88a3e7eafa..3970566daddf 100644 --- a/sdk/ai/azure-ai-projects/azure/ai/projects/operations/_patch.py +++ b/sdk/ai/azure-ai-projects/azure/ai/projects/operations/_patch.py @@ -12,7 +12,7 @@ import inspect from typing import Any, Callable, List from ..models._patch import _FOUNDRY_FEATURES_HEADER_NAME, _BETA_OPERATION_FEATURE_HEADERS, _has_header_case_insensitive -from ._patch_agents import AgentsOperations +from ._patch_agents import AgentsOperations, BetaAgentsOperations from ._patch_datasets import BetaDatasetsOperations, DatasetsOperations from ._patch_evaluators import BetaEvaluatorsOperations from ._patch_evaluation_rules import EvaluationRulesOperations @@ -21,7 +21,6 @@ from ._patch_memories import BetaMemoryStoresOperations from ._patch_models import BetaModelsOperations from ._operations import ( - BetaAgentsOperations, BetaEvaluationTaxonomiesOperations, BetaInsightsOperations, BetaOperations as GeneratedBetaOperations, @@ -122,7 +121,7 @@ def __init__(self, *args: Any, **kwargs: Any) -> None: super().__init__(*args, **kwargs) # Replace with patched class that returns EvaluatorGenerationLROPoller self.evaluators = BetaEvaluatorsOperations(self._client, self._config, self._serialize, self._deserialize) - # Replace with patched class that adds file-path overload to upload_session_file + # Replace with patched class that returns AgentOptimizationLROPoller self.agents = BetaAgentsOperations(self._client, self._config, self._serialize, self._deserialize) # Replace with patched class that includes begin_update_memories self.memory_stores = BetaMemoryStoresOperations(self._client, self._config, self._serialize, self._deserialize) diff --git a/sdk/ai/azure-ai-projects/azure/ai/projects/operations/_patch_agents.py b/sdk/ai/azure-ai-projects/azure/ai/projects/operations/_patch_agents.py index d72e81cf077d..d779e612ad75 100644 --- a/sdk/ai/azure-ai-projects/azure/ai/projects/operations/_patch_agents.py +++ b/sdk/ai/azure-ai-projects/azure/ai/projects/operations/_patch_agents.py @@ -10,11 +10,21 @@ import hashlib from io import IOBase -from typing import Union, Optional, Any, IO, overload +from typing import Union, Optional, Any, IO, cast, overload from azure.core.exceptions import HttpResponseError +from azure.core.polling import NoPolling, PollingMethod +from azure.core.polling.base_polling import LROBasePolling from azure.core.tracing.decorator import distributed_trace -from ._operations import AgentsOperations as GeneratedAgentsOperations, JSON, _Unset +from azure.core.utils import case_insensitive_dict +from ._operations import ( + AgentsOperations as GeneratedAgentsOperations, + BetaAgentsOperations as BetaAgentsOperationsGenerated, + JSON, + _Unset, +) from .. import models as _models +from .._utils.model_base import _deserialize +from ..models import AgentOptimizationLROPoller from ..models._patch import ( _FOUNDRY_FEATURES_HEADER_NAME, _has_header_case_insensitive, @@ -251,7 +261,6 @@ def create_version( new_exc.model = exc.model raise new_exc from exc raise - @distributed_trace def create_version_from_code( self, @@ -348,3 +357,113 @@ def create_version_from_code( new_exc.model = exc.model raise new_exc from exc raise + + +class BetaAgentsOperations(BetaAgentsOperationsGenerated): + """Custom operations for beta agent optimization jobs.""" + + @overload + def begin_create_optimization_job( + self, + job: _models.OptimizationJob, + *, + operation_id: Optional[str] = None, + content_type: str = "application/json", + **kwargs: Any, + ) -> AgentOptimizationLROPoller: ... + + @overload + def begin_create_optimization_job( + self, + job: JSON, + *, + operation_id: Optional[str] = None, + content_type: str = "application/json", + **kwargs: Any, + ) -> AgentOptimizationLROPoller: ... + + @overload + def begin_create_optimization_job( + self, + job: IO[bytes], + *, + operation_id: Optional[str] = None, + content_type: str = "application/json", + **kwargs: Any, + ) -> AgentOptimizationLROPoller: ... + + @distributed_trace + def begin_create_optimization_job( + self, + job: Union[_models.OptimizationJob, JSON, IO[bytes]], + *, + operation_id: Optional[str] = None, + **kwargs: Any, + ) -> AgentOptimizationLROPoller: + """Create an agent optimization job. + + :param job: The job to create. Required. + :type job: ~azure.ai.projects.models.OptimizationJob or JSON or IO[bytes] + :keyword operation_id: Client-generated unique ID for idempotent retries. When absent, the + server creates the job unconditionally. Default value is None. + :paramtype operation_id: str + :return: A poller that returns OptimizationJobResult and exposes the job ID in ``details``. + :rtype: ~azure.ai.projects.models.AgentOptimizationLROPoller + :raises ~azure.core.exceptions.HttpResponseError: + """ + headers = case_insensitive_dict(kwargs.pop("headers", {}) or {}) + params = kwargs.pop("params", {}) or {} + + content_type: Optional[str] = kwargs.pop("content_type", headers.pop("Content-Type", None)) + cls = kwargs.pop("cls", None) + polling: Union[bool, PollingMethod] = kwargs.pop("polling", True) + lro_delay = kwargs.pop("polling_interval", self._config.polling_interval) + continuation_token: Optional[str] = kwargs.pop("continuation_token", None) + raw_result = None + if continuation_token is None: + raw_result = self._create_optimization_job_initial( + job=job, + operation_id=operation_id, + content_type=content_type, + cls=lambda x, y, z: x, + headers=headers, + params=params, + **kwargs, + ) + raw_result.http_response.read() # type: ignore + kwargs.pop("error_map", None) + + def get_long_running_output(pipeline_response): + response_headers = {} + response = pipeline_response.http_response + response_headers["Operation-Location"] = self._deserialize( + "str", response.headers.get("Operation-Location") + ) + response_headers["Location"] = self._deserialize("str", response.headers.get("Location")) + + deserialized = _deserialize(_models.OptimizationJobResult, response.json().get("result", {})) + if cls: + return cls(pipeline_response, deserialized, response_headers) + return deserialized + + path_format_arguments = { + "endpoint": self._serialize.url("self._config.endpoint", self._config.endpoint, "str", skip_quote=True), + } + + if polling is True: + polling_method: PollingMethod = cast( + PollingMethod, LROBasePolling(lro_delay, path_format_arguments=path_format_arguments, **kwargs) + ) + elif polling is False: + polling_method = cast(PollingMethod, NoPolling()) + else: + polling_method = polling + if continuation_token: + return AgentOptimizationLROPoller.from_continuation_token( + polling_method=polling_method, + continuation_token=continuation_token, + client=self._client, + deserialization_callback=get_long_running_output, + ) + assert raw_result is not None + return AgentOptimizationLROPoller(self._client, raw_result, get_long_running_output, polling_method) # type: ignore diff --git a/sdk/ai/azure-ai-projects/samples/agents/optimization/sample_optimization_job_advanced_app_polling.py b/sdk/ai/azure-ai-projects/samples/agents/optimization/sample_optimization_job_advanced_app_polling.py index ab72b614aeb9..e47593f8a6fc 100644 --- a/sdk/ai/azure-ai-projects/samples/agents/optimization/sample_optimization_job_advanced_app_polling.py +++ b/sdk/ai/azure-ai-projects/samples/agents/optimization/sample_optimization_job_advanced_app_polling.py @@ -71,13 +71,6 @@ # 1. Create an optimization job without SDK polling. # ------------------------------------------------------------------ print("Creating optimization job...") - created_jobs: list[OptimizationJob] = [] - - def raw_response_hook(response): - # Since `polling=False` is set below, it is guaranteed that `raw_response_hook` will be - # invoked once on the initial "201 Created" response, and `response` is of type `OptimizationJob`. - response.http_response.read() - created_jobs.append(OptimizationJob(response.http_response.json())) job = OptimizationJob( inputs=OptimizationJobInputs( @@ -95,14 +88,16 @@ def raw_response_hook(response): ) ) - project_client.beta.agents.begin_create_optimization_job( + poller = project_client.beta.agents.begin_create_optimization_job( job=job, polling=False, - raw_response_hook=raw_response_hook, ) - if not created_jobs: - raise RuntimeError("The create operation did not return an optimization job.") - job = created_jobs[0] + job_id = poller.details["job_id"] + if not job_id: + raise RuntimeError( + "The create operation did not return an optimization job ID." + ) + job = project_client.beta.agents.get_optimization_job(job_id=job_id) print(f"Created job: id={job.id}, status={job.status}") # ------------------------------------------------------------------ diff --git a/sdk/ai/azure-ai-projects/samples/agents/optimization/sample_optimization_job_advanced_app_polling_async.py b/sdk/ai/azure-ai-projects/samples/agents/optimization/sample_optimization_job_advanced_app_polling_async.py index 7a8599ecb48b..59867af98793 100644 --- a/sdk/ai/azure-ai-projects/samples/agents/optimization/sample_optimization_job_advanced_app_polling_async.py +++ b/sdk/ai/azure-ai-projects/samples/agents/optimization/sample_optimization_job_advanced_app_polling_async.py @@ -73,13 +73,6 @@ async def main() -> None: # 1. Create an optimization job without SDK polling. # ------------------------------------------------------------------ print("Creating optimization job...") - pipeline_responses = [] - - def raw_response_hook(response): - # The raw_response_hook is called synchronously before the generated LRO method - # awaits read() on the initial response. Capture the pipeline response object here - # and parse the body afterwards, when read() has already been awaited. - pipeline_responses.append(response) job = OptimizationJob( inputs=OptimizationJobInputs( @@ -97,16 +90,16 @@ def raw_response_hook(response): ) ) - await project_client.beta.agents.begin_create_optimization_job( + poller = await project_client.beta.agents.begin_create_optimization_job( job=job, polling=False, - raw_response_hook=raw_response_hook, ) - # Alternatively, have the SDK handle polling by removing `polling=False`, assigning the awaited call - # to a poller, and then awaiting `poller.result()`. - if not pipeline_responses: - raise RuntimeError("The create operation did not return an optimization job.") - job = OptimizationJob(pipeline_responses[0].http_response.json()) + job_id = poller.details["job_id"] + if not job_id: + raise RuntimeError( + "The create operation did not return an optimization job ID." + ) + job = await project_client.beta.agents.get_optimization_job(job_id=job_id) print(f"Created job: id={job.id}, status={job.status}") # ------------------------------------------------------------------ @@ -134,7 +127,9 @@ def raw_response_hook(response): # 3. Inspect the results. # ------------------------------------------------------------------ if job.result is None: - raise RuntimeError(f"Optimization job `{job.id}` completed without a result.") + raise RuntimeError( + f"Optimization job `{job.id}` completed without a result." + ) result = job.result print(f"\nBaseline candidate: {result.baseline}") diff --git a/sdk/ai/azure-ai-projects/tests/agents/test_agent_optimization_poller.py b/sdk/ai/azure-ai-projects/tests/agents/test_agent_optimization_poller.py new file mode 100644 index 000000000000..21926017d044 --- /dev/null +++ b/sdk/ai/azure-ai-projects/tests/agents/test_agent_optimization_poller.py @@ -0,0 +1,35 @@ +# ------------------------------------ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. +# ------------------------------------ +"""Unit tests for sync agent optimization pollers.""" + +from unittest.mock import MagicMock + +from azure.ai.projects.models import AgentOptimizationLROPoller +from azure.ai.projects.operations._patch_agents import BetaAgentsOperations + + +def test_begin_create_optimization_job_exposes_job_id(): + """The sync create operation exposes its job ID without SDK polling.""" + operation = BetaAgentsOperations.__new__(BetaAgentsOperations) + operation._client = MagicMock() # pylint: disable=protected-access + operation._config = MagicMock( + polling_interval=0 + ) # pylint: disable=protected-access + operation._serialize = MagicMock() # pylint: disable=protected-access + operation._serialize.url.return_value = ( + "https://example.test" # pylint: disable=protected-access + ) + operation._deserialize = MagicMock() # pylint: disable=protected-access + + initial_response = MagicMock() + initial_response.http_response.json.return_value = {"id": "optimization-job-sync"} + operation._create_optimization_job_initial = MagicMock( + return_value=initial_response + ) # pylint: disable=protected-access + + poller = operation.begin_create_optimization_job(job={}, polling=False) + + assert isinstance(poller, AgentOptimizationLROPoller) + assert poller.details["job_id"] == "optimization-job-sync" diff --git a/sdk/ai/azure-ai-projects/tests/agents/test_agent_optimization_poller_async.py b/sdk/ai/azure-ai-projects/tests/agents/test_agent_optimization_poller_async.py new file mode 100644 index 000000000000..0df5d51d46d0 --- /dev/null +++ b/sdk/ai/azure-ai-projects/tests/agents/test_agent_optimization_poller_async.py @@ -0,0 +1,39 @@ +# ------------------------------------ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. +# ------------------------------------ +"""Unit tests for async agent optimization pollers.""" + +from unittest.mock import AsyncMock, MagicMock + +import pytest + +from azure.ai.projects.aio.operations._patch_agents_async import BetaAgentsOperations +from azure.ai.projects.models import AsyncAgentOptimizationLROPoller + + +@pytest.mark.asyncio +async def test_begin_create_optimization_job_exposes_job_id_async(): + """The async create operation exposes its job ID without SDK polling.""" + operation = BetaAgentsOperations.__new__(BetaAgentsOperations) + operation._client = MagicMock() # pylint: disable=protected-access + operation._config = MagicMock( + polling_interval=0 + ) # pylint: disable=protected-access + operation._serialize = MagicMock() # pylint: disable=protected-access + operation._serialize.url.return_value = ( + "https://example.test" # pylint: disable=protected-access + ) + operation._deserialize = MagicMock() # pylint: disable=protected-access + + initial_response = MagicMock() + initial_response.http_response.json.return_value = {"id": "optimization-job-async"} + initial_response.http_response.read = AsyncMock() + operation._create_optimization_job_initial = AsyncMock( + return_value=initial_response + ) # pylint: disable=protected-access + + poller = await operation.begin_create_optimization_job(job={}, polling=False) + + assert isinstance(poller, AsyncAgentOptimizationLROPoller) + assert poller.details["job_id"] == "optimization-job-async" From 4cc121f1fe04655702f2638c9037e5886f171043 Mon Sep 17 00:00:00 2001 From: Darren Cohen <39422044+dargilco@users.noreply.github.com> Date: Wed, 5 Aug 2026 17:37:55 -0700 Subject: [PATCH 6/8] update doc string to mention job_id --- .../azure/ai/projects/models/_patch.py | 24 ++++++++++++++----- 1 file changed, 18 insertions(+), 6 deletions(-) diff --git a/sdk/ai/azure-ai-projects/azure/ai/projects/models/_patch.py b/sdk/ai/azure-ai-projects/azure/ai/projects/models/_patch.py index 33d30341ddc3..3fcdf0571814 100644 --- a/sdk/ai/azure-ai-projects/azure/ai/projects/models/_patch.py +++ b/sdk/ai/azure-ai-projects/azure/ai/projects/models/_patch.py @@ -404,7 +404,9 @@ def _get_job_id(initial_response: Any) -> Optional[str]: def details(self) -> Mapping[str, Any]: """Returns metadata associated with the data generation job operation. - :return: A mapping containing the created data generation job ID. + The mapping contains a ``job_id`` key whose value is the created data generation job ID. + + :return: A mapping containing the ``job_id`` key. :rtype: Mapping[str, Any] """ return {"job_id": self._job_id} @@ -439,7 +441,9 @@ def __init__(self, client: Any, initial_response: Any, deserialization_callback: def details(self) -> Mapping[str, Any]: """Returns metadata associated with the data generation job operation. - :return: A mapping containing the created data generation job ID. + The mapping contains a ``job_id`` key whose value is the created data generation job ID. + + :return: A mapping containing the ``job_id`` key. :rtype: Mapping[str, Any] """ return {"job_id": self._job_id} @@ -477,7 +481,9 @@ def __init__(self, client: Any, initial_response: Any, deserialization_callback: def details(self) -> Mapping[str, Any]: """Returns metadata associated with the evaluator generation job operation. - :return: A mapping containing the created evaluator generation job ID. + The mapping contains a ``job_id`` key whose value is the created evaluator generation job ID. + + :return: A mapping containing the ``job_id`` key. :rtype: Mapping[str, Any] """ return {"job_id": self._job_id} @@ -512,7 +518,9 @@ def __init__(self, client: Any, initial_response: Any, deserialization_callback: def details(self) -> Mapping[str, Any]: """Returns metadata associated with the evaluator generation job operation. - :return: A mapping containing the created evaluator generation job ID. + The mapping contains a ``job_id`` key whose value is the created evaluator generation job ID. + + :return: A mapping containing the ``job_id`` key. :rtype: Mapping[str, Any] """ return {"job_id": self._job_id} @@ -550,7 +558,9 @@ def __init__(self, client: Any, initial_response: Any, deserialization_callback: def details(self) -> Mapping[str, Any]: """Returns metadata associated with the agent optimization job operation. - :return: A mapping containing the created agent optimization job ID. + The mapping contains a ``job_id`` key whose value is the created agent optimization job ID. + + :return: A mapping containing the ``job_id`` key. :rtype: Mapping[str, Any] """ return {"job_id": self._job_id} @@ -585,7 +595,9 @@ def __init__(self, client: Any, initial_response: Any, deserialization_callback: def details(self) -> Mapping[str, Any]: """Returns metadata associated with the agent optimization job operation. - :return: A mapping containing the created agent optimization job ID. + The mapping contains a ``job_id`` key whose value is the created agent optimization job ID. + + :return: A mapping containing the ``job_id`` key. :rtype: Mapping[str, Any] """ return {"job_id": self._job_id} From a6acb8aeef27aae4924c630ebafd509942af56b4 Mon Sep 17 00:00:00 2001 From: Darren Cohen <39422044+dargilco@users.noreply.github.com> Date: Wed, 5 Aug 2026 18:43:12 -0700 Subject: [PATCH 7/8] Fix tests --- sdk/ai/azure-ai-projects/tests/samples/test_samples.py | 1 + sdk/ai/azure-ai-projects/tests/samples/test_samples_async.py | 1 + 2 files changed, 2 insertions(+) diff --git a/sdk/ai/azure-ai-projects/tests/samples/test_samples.py b/sdk/ai/azure-ai-projects/tests/samples/test_samples.py index d7a9e0e886c6..d2428279b5ce 100644 --- a/sdk/ai/azure-ai-projects/tests/samples/test_samples.py +++ b/sdk/ai/azure-ai-projects/tests/samples/test_samples.py @@ -203,6 +203,7 @@ def test_models_samples(self, sample_path: str, **kwargs) -> None: "sample_dataset_generation_job_traces_for_evaluation.py", # PR #47067: recording not yet available "sample_dataset_generation_job_simpleqna_with_agent_source.py", # PR #47067: recording not yet available "sample_dataset_generation_job_simpleqna_with_file_source.py", # PR #47067: recording not yet available + "sample_dataset_generation_job_simpleqna_for_finetuning_with_app_polling.py", # Need test recordings ], ), ) diff --git a/sdk/ai/azure-ai-projects/tests/samples/test_samples_async.py b/sdk/ai/azure-ai-projects/tests/samples/test_samples_async.py index 1118a144c044..32174ad270db 100644 --- a/sdk/ai/azure-ai-projects/tests/samples/test_samples_async.py +++ b/sdk/ai/azure-ai-projects/tests/samples/test_samples_async.py @@ -168,6 +168,7 @@ async def test_models_samples(self, sample_path: str, **kwargs) -> None: samples_to_skip=[ "sample_datasets_async.py", # Skipped until re-enabled and recorded on Foundry endpoint that supports the new versioning schema "sample_dataset_generation_job_simpleqna_for_finetuning_async.py", # Need to add recordings + "sample_dataset_generation_job_simpleqna_for_finetuning_with_app_polling_async.py", # Need test recordings ], ), ) From d1f75eec23461c2b6b98b54a828d7bdacdb0f351 Mon Sep 17 00:00:00 2001 From: Darren Cohen <39422044+dargilco@users.noreply.github.com> Date: Wed, 5 Aug 2026 19:33:41 -0700 Subject: [PATCH 8/8] Update api.md --- sdk/ai/azure-ai-projects/api.md | 168 ++++++++++++++++++---- sdk/ai/azure-ai-projects/api.metadata.yml | 2 +- 2 files changed, 145 insertions(+), 25 deletions(-) diff --git a/sdk/ai/azure-ai-projects/api.md b/sdk/ai/azure-ai-projects/api.md index a683fb81cd13..16f8a58956a4 100644 --- a/sdk/ai/azure-ai-projects/api.md +++ b/sdk/ai/azure-ai-projects/api.md @@ -405,7 +405,7 @@ namespace azure.ai.projects.aio.operations ) -> SessionFileWriteResult: ... - class azure.ai.projects.aio.operations.BetaAgentsOperations: + class azure.ai.projects.aio.operations.BetaAgentsOperations(BetaAgentsOperationsGenerated): def __init__( self, @@ -421,7 +421,7 @@ namespace azure.ai.projects.aio.operations content_type: str = "application/json", operation_id: Optional[str] = ..., **kwargs: Any - ) -> AsyncLROPoller[OptimizationJobResult]: ... + ) -> AsyncAgentOptimizationLROPoller: ... @overload async def begin_create_optimization_job( @@ -431,7 +431,7 @@ namespace azure.ai.projects.aio.operations content_type: str = "application/json", operation_id: Optional[str] = ..., **kwargs: Any - ) -> AsyncLROPoller[OptimizationJobResult]: ... + ) -> AsyncAgentOptimizationLROPoller: ... @overload async def begin_create_optimization_job( @@ -441,7 +441,7 @@ namespace azure.ai.projects.aio.operations content_type: str = "application/json", operation_id: Optional[str] = ..., **kwargs: Any - ) -> AsyncLROPoller[OptimizationJobResult]: ... + ) -> AsyncAgentOptimizationLROPoller: ... @distributed_trace_async async def cancel_optimization_job( @@ -477,7 +477,7 @@ namespace azure.ai.projects.aio.operations ) -> AsyncItemPaged[OptimizationJobListItem]: ... - class azure.ai.projects.aio.operations.BetaDatasetsOperations: + class azure.ai.projects.aio.operations.BetaDatasetsOperations(BetaDatasetsOperationsGenerated): def __init__( self, @@ -493,7 +493,7 @@ namespace azure.ai.projects.aio.operations content_type: str = "application/json", operation_id: Optional[str] = ..., **kwargs: Any - ) -> AsyncLROPoller[DataGenerationJobResult]: ... + ) -> AsyncDatasetGenerationLROPoller: ... @overload async def begin_create_generation_job( @@ -503,7 +503,7 @@ namespace azure.ai.projects.aio.operations content_type: str = "application/json", operation_id: Optional[str] = ..., **kwargs: Any - ) -> AsyncLROPoller[DataGenerationJobResult]: ... + ) -> AsyncDatasetGenerationLROPoller: ... @overload async def begin_create_generation_job( @@ -513,7 +513,7 @@ namespace azure.ai.projects.aio.operations content_type: str = "application/json", operation_id: Optional[str] = ..., **kwargs: Any - ) -> AsyncLROPoller[DataGenerationJobResult]: ... + ) -> AsyncDatasetGenerationLROPoller: ... @distributed_trace_async async def cancel_generation_job( @@ -639,7 +639,7 @@ namespace azure.ai.projects.aio.operations ) -> EvaluationTaxonomy: ... - class azure.ai.projects.aio.operations.BetaEvaluatorsOperations: + class azure.ai.projects.aio.operations.BetaEvaluatorsOperations(BetaEvaluatorsOperationsGenerated): def __init__( self, @@ -655,7 +655,7 @@ namespace azure.ai.projects.aio.operations content_type: str = "application/json", operation_id: Optional[str] = ..., **kwargs: Any - ) -> AsyncLROPoller[EvaluatorVersion]: ... + ) -> AsyncEvaluatorGenerationLROPoller: ... @overload async def begin_create_generation_job( @@ -665,7 +665,7 @@ namespace azure.ai.projects.aio.operations content_type: str = "application/json", operation_id: Optional[str] = ..., **kwargs: Any - ) -> AsyncLROPoller[EvaluatorVersion]: ... + ) -> AsyncEvaluatorGenerationLROPoller: ... @overload async def begin_create_generation_job( @@ -675,7 +675,7 @@ namespace azure.ai.projects.aio.operations content_type: str = "application/json", operation_id: Optional[str] = ..., **kwargs: Any - ) -> AsyncLROPoller[EvaluatorVersion]: ... + ) -> AsyncEvaluatorGenerationLROPoller: ... @distributed_trace_async async def cancel_generation_job( @@ -2684,6 +2684,26 @@ namespace azure.ai.projects.models def __init__(self, mapping: Mapping[str, Any]) -> None: ... + class azure.ai.projects.models.AgentOptimizationLROPoller(LROPoller[OptimizationJobResult]): + property details: Mapping[str, Any] # Read-only + + def __init__( + self, + client: Any, + initial_response: Any, + deserialization_callback: Any, + polling_method: Any + ) -> None: ... + + @classmethod + def from_continuation_token( + cls, + polling_method: PollingMethod[OptimizationJobResult], + continuation_token: str, + **kwargs: Any + ) -> AgentOptimizationLROPoller: ... + + class azure.ai.projects.models.AgentSessionResource(_Model): agent_session_id: str created_at: datetime @@ -2890,6 +2910,66 @@ namespace azure.ai.projects.models def __init__(self, mapping: Mapping[str, Any]) -> None: ... + class azure.ai.projects.models.AsyncAgentOptimizationLROPoller(AsyncLROPoller[OptimizationJobResult]): + property details: Mapping[str, Any] # Read-only + + def __init__( + self, + client: Any, + initial_response: Any, + deserialization_callback: Any, + polling_method: Any + ) -> None: ... + + @classmethod + def from_continuation_token( + cls, + polling_method: AsyncPollingMethod[OptimizationJobResult], + continuation_token: str, + **kwargs: Any + ) -> AsyncAgentOptimizationLROPoller: ... + + + class azure.ai.projects.models.AsyncDatasetGenerationLROPoller(AsyncLROPoller[DataGenerationJobResult]): + property details: Mapping[str, Any] # Read-only + + def __init__( + self, + client: Any, + initial_response: Any, + deserialization_callback: Any, + polling_method: Any + ) -> None: ... + + @classmethod + def from_continuation_token( + cls, + polling_method: AsyncPollingMethod[DataGenerationJobResult], + continuation_token: str, + **kwargs: Any + ) -> AsyncDatasetGenerationLROPoller: ... + + + class azure.ai.projects.models.AsyncEvaluatorGenerationLROPoller(AsyncLROPoller[EvaluatorVersion]): + property details: Mapping[str, Any] # Read-only + + def __init__( + self, + client: Any, + initial_response: Any, + deserialization_callback: Any, + polling_method: Any + ) -> None: ... + + @classmethod + def from_continuation_token( + cls, + polling_method: AsyncPollingMethod[EvaluatorVersion], + continuation_token: str, + **kwargs: Any + ) -> AsyncEvaluatorGenerationLROPoller: ... + + class azure.ai.projects.models.AsyncUpdateMemoriesLROPoller(AsyncLROPoller[MemoryStoreUpdateCompletedResult]): property superseded_by: Optional[str] # Read-only property update_id: str # Read-only @@ -4310,6 +4390,26 @@ namespace azure.ai.projects.models def __init__(self, mapping: Mapping[str, Any]) -> None: ... + class azure.ai.projects.models.DatasetGenerationLROPoller(LROPoller[DataGenerationJobResult]): + property details: Mapping[str, Any] # Read-only + + def __init__( + self, + client: Any, + initial_response: Any, + deserialization_callback: Any, + polling_method: Any + ) -> None: ... + + @classmethod + def from_continuation_token( + cls, + polling_method: PollingMethod[DataGenerationJobResult], + continuation_token: str, + **kwargs: Any + ) -> DatasetGenerationLROPoller: ... + + class azure.ai.projects.models.DatasetReference(_Model): name: str version: str @@ -5058,6 +5158,26 @@ namespace azure.ai.projects.models TRACES = "traces" + class azure.ai.projects.models.EvaluatorGenerationLROPoller(LROPoller[EvaluatorVersion]): + property details: Mapping[str, Any] # Read-only + + def __init__( + self, + client: Any, + initial_response: Any, + deserialization_callback: Any, + polling_method: Any + ) -> None: ... + + @classmethod + def from_continuation_token( + cls, + polling_method: PollingMethod[EvaluatorVersion], + continuation_token: str, + **kwargs: Any + ) -> EvaluatorGenerationLROPoller: ... + + class azure.ai.projects.models.EvaluatorGenerationTokenUsage(_Model): input_tokens: int output_tokens: int @@ -9796,7 +9916,7 @@ namespace azure.ai.projects.operations ) -> SessionFileWriteResult: ... - class azure.ai.projects.operations.BetaAgentsOperations: + class azure.ai.projects.operations.BetaAgentsOperations(BetaAgentsOperationsGenerated): def __init__( self, @@ -9812,7 +9932,7 @@ namespace azure.ai.projects.operations content_type: str = "application/json", operation_id: Optional[str] = ..., **kwargs: Any - ) -> LROPoller[OptimizationJobResult]: ... + ) -> AgentOptimizationLROPoller: ... @overload def begin_create_optimization_job( @@ -9822,7 +9942,7 @@ namespace azure.ai.projects.operations content_type: str = "application/json", operation_id: Optional[str] = ..., **kwargs: Any - ) -> LROPoller[OptimizationJobResult]: ... + ) -> AgentOptimizationLROPoller: ... @overload def begin_create_optimization_job( @@ -9832,7 +9952,7 @@ namespace azure.ai.projects.operations content_type: str = "application/json", operation_id: Optional[str] = ..., **kwargs: Any - ) -> LROPoller[OptimizationJobResult]: ... + ) -> AgentOptimizationLROPoller: ... @distributed_trace def cancel_optimization_job( @@ -9868,7 +9988,7 @@ namespace azure.ai.projects.operations ) -> ItemPaged[OptimizationJobListItem]: ... - class azure.ai.projects.operations.BetaDatasetsOperations: + class azure.ai.projects.operations.BetaDatasetsOperations(BetaDatasetsOperationsGenerated): def __init__( self, @@ -9884,7 +10004,7 @@ namespace azure.ai.projects.operations content_type: str = "application/json", operation_id: Optional[str] = ..., **kwargs: Any - ) -> LROPoller[DataGenerationJobResult]: ... + ) -> DatasetGenerationLROPoller: ... @overload def begin_create_generation_job( @@ -9894,7 +10014,7 @@ namespace azure.ai.projects.operations content_type: str = "application/json", operation_id: Optional[str] = ..., **kwargs: Any - ) -> LROPoller[DataGenerationJobResult]: ... + ) -> DatasetGenerationLROPoller: ... @overload def begin_create_generation_job( @@ -9904,7 +10024,7 @@ namespace azure.ai.projects.operations content_type: str = "application/json", operation_id: Optional[str] = ..., **kwargs: Any - ) -> LROPoller[DataGenerationJobResult]: ... + ) -> DatasetGenerationLROPoller: ... @distributed_trace def cancel_generation_job( @@ -10030,7 +10150,7 @@ namespace azure.ai.projects.operations ) -> EvaluationTaxonomy: ... - class azure.ai.projects.operations.BetaEvaluatorsOperations: + class azure.ai.projects.operations.BetaEvaluatorsOperations(BetaEvaluatorsOperationsGenerated): def __init__( self, @@ -10046,7 +10166,7 @@ namespace azure.ai.projects.operations content_type: str = "application/json", operation_id: Optional[str] = ..., **kwargs: Any - ) -> LROPoller[EvaluatorVersion]: ... + ) -> EvaluatorGenerationLROPoller: ... @overload def begin_create_generation_job( @@ -10056,7 +10176,7 @@ namespace azure.ai.projects.operations content_type: str = "application/json", operation_id: Optional[str] = ..., **kwargs: Any - ) -> LROPoller[EvaluatorVersion]: ... + ) -> EvaluatorGenerationLROPoller: ... @overload def begin_create_generation_job( @@ -10066,7 +10186,7 @@ namespace azure.ai.projects.operations content_type: str = "application/json", operation_id: Optional[str] = ..., **kwargs: Any - ) -> LROPoller[EvaluatorVersion]: ... + ) -> EvaluatorGenerationLROPoller: ... @distributed_trace def cancel_generation_job( diff --git a/sdk/ai/azure-ai-projects/api.metadata.yml b/sdk/ai/azure-ai-projects/api.metadata.yml index 3d493abef420..77e7e6ec47f5 100644 --- a/sdk/ai/azure-ai-projects/api.metadata.yml +++ b/sdk/ai/azure-ai-projects/api.metadata.yml @@ -1,3 +1,3 @@ -apiMdSha256: 544c82773e2ee8b4aeb0ece5b64bb938f2d3703720950214d3e4c5c98e3e61fd +apiMdSha256: e2ea7472bfd266cb3abdd41b29c3014f7ab177eee9e9c2b0effe885a9ff36f4a parserVersion: 0.3.30 pythonVersion: 3.14.3