From 551fcfc1dc5b8b67c0c6179d49c3eef2c0a38676 Mon Sep 17 00:00:00 2001 From: evan-runloopai <155475557+evan-runloopai@users.noreply.github.com> Date: Fri, 31 Jan 2025 22:29:56 -0800 Subject: [PATCH] feat: add helper methods for scenarios --- .../resources/scenarios/runs.py | 286 ++++++++++++++++++ 1 file changed, 286 insertions(+) diff --git a/src/runloop_api_client/resources/scenarios/runs.py b/src/runloop_api_client/resources/scenarios/runs.py index 94b450a36..351fb8f3d 100644 --- a/src/runloop_api_client/resources/scenarios/runs.py +++ b/src/runloop_api_client/resources/scenarios/runs.py @@ -18,6 +18,9 @@ from ..._base_client import AsyncPaginator, make_request_options from ...types.scenarios import run_list_params from ...types.scenario_run_view import ScenarioRunView +from ..._exceptions import RunloopError +from ...lib.polling import PollingConfig, poll_until +from ...lib.polling_async import async_poll_until __all__ = ["RunsResource", "AsyncRunsResource"] @@ -206,6 +209,147 @@ def score( cast_to=ScenarioRunView, ) + def await_scored( + self, + id: str, + *, + polling_config: PollingConfig | None = None, + # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs. + # The extra values given here take precedence over values defined on the client or passed to this method. + extra_headers: Headers | None = None, + extra_query: Query | None = None, + extra_body: Body | None = None, + timeout: float | httpx.Timeout | None | NotGiven = NOT_GIVEN, + ) -> ScenarioRunView: + """Wait for a scenario run to be scored. + + Args: + id: The ID of the scenario run to wait for + polling_config: Optional polling configuration + extra_headers: Send extra headers + extra_query: Add additional query parameters to the request + extra_body: Add additional JSON properties to the request + timeout: Override the client-level default timeout for this request, in seconds + + Returns: + The scored scenario run + + Raises: + PollingTimeout: If polling times out before scenario run is scored + RunloopError: If scenario run enters a non-scored terminal state + """ + def retrieve_run() -> ScenarioRunView: + return self.retrieve( + id, + extra_headers=extra_headers, + extra_query=extra_query, + extra_body=extra_body, + timeout=timeout + ) + + def is_done_scoring(run: ScenarioRunView) -> bool: + return run.state not in ["scoring"] + + run = poll_until(retrieve_run, is_done_scoring, polling_config) + + if run.state != "scored": + raise RunloopError( + f"Scenario run entered non-scored state unexpectedly: {run.state}" + ) + + return run + + def score_and_await( + self, + id: str, + *, + polling_config: PollingConfig | None = None, + # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs. + # The extra values given here take precedence over values defined on the client or passed to this method. + extra_headers: Headers | None = None, + extra_query: Query | None = None, + extra_body: Body | None = None, + timeout: float | httpx.Timeout | None | NotGiven = NOT_GIVEN, + ) -> ScenarioRunView: + """Score a scenario run and wait for it to be scored. + + Args: + id: The ID of the scenario run to score and wait for + polling_config: Optional polling configuration + extra_headers: Send extra headers + extra_query: Add additional query parameters to the request + extra_body: Add additional JSON properties to the request + timeout: Override the client-level default timeout for this request, in seconds + + Returns: + The scored scenario run + + Raises: + PollingTimeout: If polling times out before scenario run is scored + RunloopError: If scenario run enters a non-scored terminal state + """ + self.score( + id, + extra_headers=extra_headers, + extra_query=extra_query, + extra_body=extra_body, + timeout=timeout, + ) + + return self.await_scored( + id, + polling_config=polling_config, + extra_headers=extra_headers, + extra_query=extra_query, + extra_body=extra_body, + timeout=timeout, + ) + + def score_and_complete( + self, + id: str, + *, + polling_config: PollingConfig | None = None, + # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs. + # The extra values given here take precedence over values defined on the client or passed to this method. + extra_headers: Headers | None = None, + extra_query: Query | None = None, + extra_body: Body | None = None, + timeout: float | httpx.Timeout | None | NotGiven = NOT_GIVEN, + ) -> ScenarioRunView: + """Score a scenario run, wait for it to be scored, then complete it. + + Args: + id: The ID of the scenario run to score, wait for, and complete + polling_config: Optional polling configuration + extra_headers: Send extra headers + extra_query: Add additional query parameters to the request + extra_body: Add additional JSON properties to the request + timeout: Override the client-level default timeout for this request, in seconds + + Returns: + The completed scenario run + + Raises: + PollingTimeout: If polling times out before scenario run is scored + RunloopError: If scenario run enters a non-scored terminal state + """ + self.score_and_await( + id, + polling_config=polling_config, + extra_headers=extra_headers, + extra_query=extra_query, + extra_body=extra_body, + timeout=timeout, + ) + + return self.complete( + id, + extra_headers=extra_headers, + extra_query=extra_query, + extra_body=extra_body, + timeout=timeout, + ) class AsyncRunsResource(AsyncAPIResource): @cached_property @@ -391,6 +535,148 @@ async def score( cast_to=ScenarioRunView, ) + async def await_scored( + self, + id: str, + *, + polling_config: PollingConfig | None = None, + # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs. + # The extra values given here take precedence over values defined on the client or passed to this method. + extra_headers: Headers | None = None, + extra_query: Query | None = None, + extra_body: Body | None = None, + timeout: float | httpx.Timeout | None | NotGiven = NOT_GIVEN, + ) -> ScenarioRunView: + """Wait for a scenario run to be scored. + + Args: + id: The ID of the scenario run to wait for + polling_config: Optional polling configuration + extra_headers: Send extra headers + extra_query: Add additional query parameters to the request + extra_body: Add additional JSON properties to the request + timeout: Override the client-level default timeout for this request, in seconds + + Returns: + The scored scenario run + + Raises: + PollingTimeout: If polling times out before scenario run is scored + RunloopError: If scenario run enters a non-scored terminal state + """ + async def retrieve_run() -> ScenarioRunView: + return await self.retrieve( + id, + extra_headers=extra_headers, + extra_query=extra_query, + extra_body=extra_body, + timeout=timeout + ) + + def is_done_scoring(run: ScenarioRunView) -> bool: + return run.state not in ["scoring"] + + run = await async_poll_until(retrieve_run, is_done_scoring, polling_config) + + if run.state != "scored": + raise RunloopError( + f"Scenario run entered non-scored state unexpectedly: {run.state}" + ) + + return run + + async def score_and_await( + self, + id: str, + *, + polling_config: PollingConfig | None = None, + # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs. + # The extra values given here take precedence over values defined on the client or passed to this method. + extra_headers: Headers | None = None, + extra_query: Query | None = None, + extra_body: Body | None = None, + timeout: float | httpx.Timeout | None | NotGiven = NOT_GIVEN, + ) -> ScenarioRunView: + """Score a scenario run and wait for it to be scored. + + Args: + id: The ID of the scenario run to score and wait for + polling_config: Optional polling configuration + extra_headers: Send extra headers + extra_query: Add additional query parameters to the request + extra_body: Add additional JSON properties to the request + timeout: Override the client-level default timeout for this request, in seconds + + Returns: + The scored scenario run + + Raises: + PollingTimeout: If polling times out before scenario run is scored + RunloopError: If scenario run enters a non-scored terminal state + """ + await self.score( + id, + extra_headers=extra_headers, + extra_query=extra_query, + extra_body=extra_body, + timeout=timeout, + ) + + return await self.await_scored( + id, + polling_config=polling_config, + extra_headers=extra_headers, + extra_query=extra_query, + extra_body=extra_body, + timeout=timeout, + ) + + async def score_and_complete( + self, + id: str, + *, + polling_config: PollingConfig | None = None, + # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs. + # The extra values given here take precedence over values defined on the client or passed to this method. + extra_headers: Headers | None = None, + extra_query: Query | None = None, + extra_body: Body | None = None, + timeout: float | httpx.Timeout | None | NotGiven = NOT_GIVEN, + ) -> ScenarioRunView: + """Score a scenario run, wait for it to be scored, then complete it. + + Args: + id: The ID of the scenario run to score, wait for, and complete + polling_config: Optional polling configuration + extra_headers: Send extra headers + extra_query: Add additional query parameters to the request + extra_body: Add additional JSON properties to the request + timeout: Override the client-level default timeout for this request, in seconds + + Returns: + The completed scenario run + + Raises: + PollingTimeout: If polling times out before scenario run is scored + RunloopError: If scenario run enters a non-scored terminal state + """ + await self.score_and_await( + id, + polling_config=polling_config, + extra_headers=extra_headers, + extra_query=extra_query, + extra_body=extra_body, + timeout=timeout, + ) + + return await self.complete( + id, + extra_headers=extra_headers, + extra_query=extra_query, + extra_body=extra_body, + timeout=timeout, + ) + class RunsResourceWithRawResponse: def __init__(self, runs: RunsResource) -> None: