Skip to content

Poller changes #48482

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
Howie Leung (howieleung) wants to merge 1 commit into
base: dargilco/custom-lro-poller
Choose a base branch
from
+95 −12
Open

Poller changes #48482

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
83 changes: 82 additions & 1 deletion sdk/ai/azure-ai-projects/azure/ai/projects/models/_patch.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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."""
Expand Down Expand Up @@ -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


Expand Down
24 changes: 13 additions & 11 deletions sdk/ai/azure-ai-projects/azure/ai/projects/operations/_patch_agents.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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(
Expand All @@ -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(
Expand All @@ -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(
Expand All @@ -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.
Expand All @@ -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 {})
Expand Down Expand Up @@ -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