From 44ff83ca16de531d272a75917e0461e1499daba0 Mon Sep 17 00:00:00 2001 From: Howie Leung Date: Thu, 6 Aug 2026 11:44:06 -0700 Subject: [PATCH] Poller changes --- .../azure/ai/projects/models/_patch.py | 83 ++++++++++++++++++- .../ai/projects/operations/_patch_agents.py | 24 +++--- 2 files changed, 95 insertions(+), 12 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 3fcdf0571814..b868126d5344 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 @@ -8,7 +8,8 @@ Follow our quickstart for examples: https://aka.ms/azsdk/python/dpcodegen/python/customize """ -from typing import Final, FrozenSet, List, Dict, Mapping, Optional, Any, Tuple +from dataclasses import MISSING, dataclass, fields +from typing import Final, FrozenSet, Generic, List, Dict, Mapping, Optional, Any, Tuple, Type, TypeVar, cast from azure.core.polling import LROPoller, AsyncLROPoller, PollingMethod, AsyncPollingMethod from azure.core.polling.base_polling import ( LROBasePolling, @@ -385,6 +386,85 @@ def from_continuation_token( return cls(client, initial_response, deserialization_callback, polling_method) +TResult = TypeVar("TResult") +TJob = TypeVar("TJob") + + +@dataclass +class DatasetGenerationJob: + id: str + + +@dataclass +class EvaluatorGenerationJob: + id: str + + +@dataclass +class OptimizationJob: + id: str + + +class AdvanceLROPoller(LROPoller[TResult], Generic[TResult, TJob]): + _job_type: Type[TJob] + + def __init__(self, client: Any, initial_response: Any, deserialization_callback: Any, polling_method: Any) -> None: + self._job = self._get_job(initial_response) + super().__init__(client, initial_response, deserialization_callback, polling_method) + + def _get_job(self, initial_response: Any) -> TJob: + try: + payload = initial_response.http_response.json() + except (AttributeError, TypeError, ValueError) as exc: + raise ValueError("Failed to read job details from initial response.") from exc + + if not isinstance(payload, Mapping): + raise ValueError("Failed to read job details from initial response.") + + job_kwargs = {} + missing_fields = [] + for field in fields(cast(Any, self._job_type)): + if field.name in payload: + job_kwargs[field.name] = payload[field.name] + elif field.default is MISSING and field.default_factory is MISSING: + missing_fields.append(field.name) + + if missing_fields: + missing_field_list = ", ".join(missing_fields) + raise ValueError(f"Failed to extract required job fields from initial response: {missing_field_list}") + + return self._job_type(**job_kwargs) + + @property + def details(self) -> TJob: + """Returns metadata associated with the long-running operation. + + The returned job model is populated from matching fields in the initial HTTP response. + + :return: A job model populated from the initial response payload. + :rtype: TJob + """ + return self._job + + @classmethod + def from_continuation_token( + cls, polling_method: PollingMethod[TResult], continuation_token: str, **kwargs: Any + ) -> "AdvanceLROPoller[TResult, TJob]": + """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 AdvanceLROPoller. + :rtype: AdvanceLROPoller + """ + client, initial_response, deserialization_callback = polling_method.from_continuation_token( + continuation_token, **kwargs + ) + return cls(client, initial_response, deserialization_callback, polling_method) + + class DatasetGenerationLROPoller(LROPoller[DataGenerationJobResult]): """Custom LROPoller for data generation job operations.""" @@ -648,6 +728,7 @@ def from_continuation_token( "ToolDescriptionParam", "TracesPreviewEvalRunDataSource", "UpdateMemoriesLROPoller", + "AdvanceLROPoller" ] # Add all objects you want publicly available to users at this package level 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 d779e612ad75..6d0672834634 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 @@ -24,9 +24,11 @@ ) from .. import models as _models from .._utils.model_base import _deserialize -from ..models import AgentOptimizationLROPoller +from ..models import AdvanceLROPoller from ..models._patch import ( + OptimizationJobResult, _FOUNDRY_FEATURES_HEADER_NAME, + OptimizationJob, _has_header_case_insensitive, _AGENT_OPERATION_FEATURE_HEADERS, _PREVIEW_FEATURE_REQUIRED_CODE, @@ -370,7 +372,7 @@ def begin_create_optimization_job( operation_id: Optional[str] = None, content_type: str = "application/json", **kwargs: Any, - ) -> AgentOptimizationLROPoller: ... + ) -> AdvanceLROPoller[OptimizationJobResult, OptimizationJob]: ... @overload def begin_create_optimization_job( @@ -380,7 +382,7 @@ def begin_create_optimization_job( operation_id: Optional[str] = None, content_type: str = "application/json", **kwargs: Any, - ) -> AgentOptimizationLROPoller: ... + ) -> AdvanceLROPoller[OptimizationJobResult, OptimizationJob]: ... @overload def begin_create_optimization_job( @@ -390,7 +392,7 @@ def begin_create_optimization_job( operation_id: Optional[str] = None, content_type: str = "application/json", **kwargs: Any, - ) -> AgentOptimizationLROPoller: ... + ) -> AdvanceLROPoller[OptimizationJobResult, OptimizationJob]: ... @distributed_trace def begin_create_optimization_job( @@ -399,7 +401,7 @@ def begin_create_optimization_job( *, operation_id: Optional[str] = None, **kwargs: Any, - ) -> AgentOptimizationLROPoller: + ) -> AdvanceLROPoller[OptimizationJobResult, OptimizationJob]: """Create an agent optimization job. :param job: The job to create. Required. @@ -408,7 +410,7 @@ def begin_create_optimization_job( 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 + :rtype: ~azure.ai.projects.models.AdvanceLROPoller[OptimizationJobResult, OptimizationJob] :raises ~azure.core.exceptions.HttpResponseError: """ headers = case_insensitive_dict(kwargs.pop("headers", {}) or {}) @@ -451,19 +453,19 @@ def get_long_running_output(pipeline_response): } if polling is True: - polling_method: PollingMethod = cast( - PollingMethod, LROBasePolling(lro_delay, path_format_arguments=path_format_arguments, **kwargs) + polling_method: PollingMethod[OptimizationJobResult] = cast( + PollingMethod[OptimizationJobResult], LROBasePolling(lro_delay, path_format_arguments=path_format_arguments, **kwargs) ) elif polling is False: - polling_method = cast(PollingMethod, NoPolling()) + polling_method = cast(PollingMethod[OptimizationJobResult], NoPolling()) else: polling_method = polling if continuation_token: - return AgentOptimizationLROPoller.from_continuation_token( + return AdvanceLROPoller.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 + return AdvanceLROPoller(self._client, raw_result, get_long_running_output, polling_method) # type: ignore