diff --git a/.vscode/settings.json b/.vscode/settings.json index a6c7dace8..752653713 100644 --- a/.vscode/settings.json +++ b/.vscode/settings.json @@ -6,8 +6,8 @@ "-p", "test_*.py" ], - "python.testing.unittestEnabled": true, - "python.testing.pytestEnabled": false, + "python.testing.unittestEnabled": false, + "python.testing.pytestEnabled": true, "python.autoComplete.extraPaths": [ "./src" ], @@ -19,5 +19,8 @@ "cursorpyright.analysis.extraPaths": [ "./src" ], - "cursorpyright.analysis.typeCheckingMode": "basic" + "cursorpyright.analysis.typeCheckingMode": "basic", + "python.testing.pytestArgs": [ + "tests" + ] } diff --git a/docs/audiences.md b/docs/audiences.md index 21c967292..3723afb1c 100644 --- a/docs/audiences.md +++ b/docs/audiences.md @@ -99,6 +99,19 @@ results = job.get_results() print(results) ``` +### Cost warnings and jobs under review + +`assign_job` never blocks on funds: the job is always created. If its estimated cost +exceeds your account balance, `assign_job` logs a warning with the estimate, your +balance, and the expected shortfall — the job still runs, but may pause partway until +you top up. + +Some jobs don't go straight to running. A job can enter manual review +(`ManualApproval`) or, once out of funds mid-run, become spend-limited +(`SpendLimited`). Neither state completes on its own, so `get_results()` raises an +informative error naming the state (and the review reason, when available) instead of +blocking indefinitely — top up or wait for a reviewer, then call it again. + ## Complete Example Here's the full workflow — creating a custom audience, adding qualification examples, and running a labeling job: diff --git a/docs/quickstart.md b/docs/quickstart.md index f5c2a65fc..3ceba845b 100644 --- a/docs/quickstart.md +++ b/docs/quickstart.md @@ -105,7 +105,7 @@ results = job.get_results() # (3)! 1. Assigns the job definition to the audience and starts collecting responses. 2. Opens your browser on the running job, where you can watch responses come in and monitor progress. -3. Blocks until the job is complete and returns the results. You can also monitor progress on the [Rapidata Dashboard](https://app.rapidata.ai/dashboard). +3. Blocks until the job is complete and returns the results. If the job needs manual review or runs out of funds, this raises an informative error instead of blocking — see [Cost warnings and jobs under review](audiences.md#cost-warnings-and-jobs-under-review). You can also monitor progress on the [Rapidata Dashboard](https://app.rapidata.ai/dashboard). To understand the results format, see the [Understanding the Results](understanding_the_results.md) guide. diff --git a/pyproject.toml b/pyproject.toml index 85327dec3..43c33c7dc 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -34,6 +34,7 @@ dev = [ "ipykernel>=7.1.0,<8", "jupyter>=1.1.1,<2", "pyright>=1.1.407", + "pytest>=8.0.0,<9", ] docs = [ "mkdocs-material>=9.5.34,<10", @@ -49,6 +50,9 @@ docs = [ "mkdocs-llmstxt-md>=0.1.0", ] +[tool.pytest.ini_options] +testpaths = ["tests"] + [tool.hatch.build.targets.sdist] include = ["src/rapidata"] diff --git a/src/rapidata/rapidata_client/audience/_audience_base.py b/src/rapidata/rapidata_client/audience/_audience_base.py index 87bc2a755..1ae0735be 100644 --- a/src/rapidata/rapidata_client/audience/_audience_base.py +++ b/src/rapidata/rapidata_client/audience/_audience_base.py @@ -10,6 +10,9 @@ RapidataJobDefinition, ) from rapidata.rapidata_client.job.rapidata_job import RapidataJob + from rapidata.api_client.models.create_job_endpoint_cost_warning_model import ( + CreateJobEndpointCostWarningModel, + ) class RapidataAudienceBase: @@ -89,8 +92,31 @@ def assign_job(self, job_definition: RapidataJobDefinition) -> RapidataJob: managed_print( f"Job '{job.name}' is now viewable under: {job.job_details_page}" ) + self._warn_if_cost_exceeds_balance(job, response.cost_warning) return job + @staticmethod + def _warn_if_cost_exceeds_balance( + job: RapidataJob, cost_warning: CreateJobEndpointCostWarningModel | None + ) -> None: + """Surface the create response's optional cost warning via the SDK logger. + + The job is created and runs regardless — this is an advisory estimate that it + may pause for funds before finishing, not an error. + """ + if cost_warning is None: + return + logger.warning( + "Job '%s' has an estimated cost of %.2f, but the account balance is %.2f — " + "it will likely pause about %.2f short of finishing. The job was created and " + "runs as far as the balance allows; top up the account to let it complete. " + "This is an estimate.", + job.name, + cost_warning.estimated_cost, + cost_warning.available_balance, + cost_warning.shortfall, + ) + def find_jobs( self, name: str = "", amount: int = 10, page: int = 1 ) -> list[RapidataJob]: diff --git a/src/rapidata/rapidata_client/job/rapidata_job.py b/src/rapidata/rapidata_client/job/rapidata_job.py index 129454f9b..b8ef61525 100644 --- a/src/rapidata/rapidata_client/job/rapidata_job.py +++ b/src/rapidata/rapidata_client/job/rapidata_job.py @@ -9,6 +9,7 @@ from colorama import Fore from tqdm.auto import tqdm +from rapidata.api_client.models.audience_job_state import AudienceJobState from rapidata.service.openapi_service import OpenAPIService from rapidata.rapidata_client.config import ( logger, @@ -27,6 +28,9 @@ ) if TYPE_CHECKING: + from rapidata.api_client.models.get_job_by_id_endpoint_output import ( + GetJobByIdEndpointOutput, + ) from rapidata.rapidata_client.results.rapidata_results import RapidataResults T = TypeVar("T") @@ -70,15 +74,47 @@ def __init__( self.job_details_page = f"https://app.{self._openapi_service.environment}/audiences/{self.audience_id}/job/{self.id}" logger.debug("RapidataJob initialized") + # States a job can settle into that will never progress to Completed/Failed on + # their own: ManualApproval needs a Rapidata reviewer to act, SpendLimited needs + # an account top-up. Waiting on either (e.g. from get_results) would hang the + # caller forever, so we surface them as informative errors instead. + _BLOCKING_STATUSES = ( + AudienceJobState.MANUALAPPROVAL, + AudienceJobState.SPENDLIMITED, + ) + + def _fetch_job(self) -> GetJobByIdEndpointOutput: + """Fetches the job's GET output (state, failure message, review reason).""" + return self._openapi_service.order.job_api.job_job_id_get(self.id) + def _get_job_failure_message(self) -> str | None: """Retrieves the failure message from the job if available.""" try: - job = self._openapi_service.order.job_api.job_job_id_get(self.id) - return job.failure_message + return self._fetch_job().failure_message except Exception: logger.debug("Failed to get job failure message", self, exc_info=True) return None + def _raise_for_blocking_status(self, job: GetJobByIdEndpointOutput) -> None: + """Raises an informative error for a job state that can't reach completion + on its own, instead of letting the caller block on it indefinitely.""" + if job.state == AudienceJobState.SPENDLIMITED: + raise Exception( + f"Job '{self}' is spend-limited: the account ran out of funds while " + f"running, so it stopped collecting responses. Partial results remain " + f"available; top up the account to resume and let the job finish." + ) + + # ManualApproval — reviewReason is optional; a job can legitimately be under + # review with no customer-facing reason recorded yet. + reason = job.review_reason + reason_detail = f" ({reason.value})" if reason else "" + raise Exception( + f"Job '{self}' is being reviewed{reason_detail}: a Rapidata reviewer " + f"needs to approve it before it runs. This can take a couple of hours; " + f"once approved, call this again." + ) + def _retry_operation( self, operation: Callable[[], T], @@ -133,14 +169,23 @@ def _wait_for_status( Returns: The final status reached + + Raises: + Exception: If the job enters a state it can't progress out of on its own + (``ManualApproval`` or ``SpendLimited``) while a different status is + being awaited — with the review reason when the API provides one. """ - while (current_status := self.get_status()) not in target_statuses: + while True: + job = self._fetch_job() + current_status = job.state.value + if current_status in target_statuses: + return current_status + if current_status in self._BLOCKING_STATUSES: + self._raise_for_blocking_status(job) if status_message: logger.debug(status_message, self, current_status) sleep(check_interval) - return current_status - @property def completed_at(self) -> datetime | None: """Returns the completion date of the job, or None if not completed.""" @@ -190,9 +235,7 @@ def get_status(self) -> str: The current status of the job as a string. """ with tracer.start_as_current_span("RapidataJob.get_status"): - return self._openapi_service.order.job_api.job_job_id_get( - self.id - ).state.value + return self._fetch_job().state.value def _regenerate_results(self) -> None: """Triggers regeneration of a job whose results have gone stale. @@ -221,7 +264,9 @@ def get_results(self) -> RapidataResults: RapidataResults: The results of the job. Raises: - Exception: If failed to get job results. + Exception: If failed to get job results, or if the job is in manual + review (``ManualApproval``) or spend-limited (``SpendLimited``) and + therefore cannot complete without intervention. """ with tracer.start_as_current_span("RapidataJob.get_results"): from rapidata.api_client.exceptions import ApiException @@ -260,18 +305,22 @@ def display_progress_bar(self, refresh_rate: int = 5) -> None: Raises: ValueError: If refresh_rate is less than 1. + Exception: If the job has failed, or is in a state it can't progress out + of on its own (``ManualApproval`` or ``SpendLimited``). """ if refresh_rate < 1: raise ValueError("refresh_rate must be at least 1") - current_status = self.get_status() - if current_status == "Completed": + job = self._fetch_job() + if job.state == AudienceJobState.COMPLETED: managed_print(f"Job '{self}' is already completed.") return - if current_status == "Failed": - failure_message = self._get_job_failure_message() - raise Exception(f"Job '{self}' has failed: {failure_message}") + if job.state == AudienceJobState.FAILED: + raise Exception(f"Job '{self}' has failed: {job.failure_message}") + + if job.state in self._BLOCKING_STATUSES: + self._raise_for_blocking_status(job) # Get progress from pipeline if available with tqdm( @@ -283,15 +332,17 @@ def display_progress_bar(self, refresh_rate: int = 5) -> None: ) as pbar: last_percentage = 0 while True: - current_status = self.get_status() + job = self._fetch_job() - if current_status == "Completed": + if job.state == AudienceJobState.COMPLETED: pbar.update(100 - last_percentage) break - if current_status == "Failed": - failure_message = self._get_job_failure_message() - raise Exception(f"Job '{self}' has failed: {failure_message}") + if job.state == AudienceJobState.FAILED: + raise Exception(f"Job '{self}' has failed: {job.failure_message}") + + if job.state in self._BLOCKING_STATUSES: + self._raise_for_blocking_status(job) # Try to get progress from workflow try: diff --git a/tests/rapidata_client/__init__.py b/tests/rapidata_client/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/rapidata_client/audience/__init__.py b/tests/rapidata_client/audience/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/rapidata_client/audience/test_audience_base.py b/tests/rapidata_client/audience/test_audience_base.py new file mode 100644 index 000000000..1256418a1 --- /dev/null +++ b/tests/rapidata_client/audience/test_audience_base.py @@ -0,0 +1,64 @@ +"""Tests for the create-time cost warning surfaced by assign_job. + +When job creation returns a costWarning (estimate exceeds balance), assign_job +logs an advisory warning. The job is created regardless; behaviour is otherwise +unchanged. +""" + +from __future__ import annotations + +from unittest.mock import MagicMock + +import rapidata.rapidata_client.audience._audience_base as base_module +from rapidata.api_client.models.create_job_endpoint_cost_warning_model import ( + CreateJobEndpointCostWarningModel, +) +from rapidata.rapidata_client.audience._audience_base import RapidataAudienceBase + + +def _make_audience() -> tuple[RapidataAudienceBase, MagicMock]: + openapi_service = MagicMock() + openapi_service.environment = "rapidata.ai" + audience = RapidataAudienceBase( + id="aud-1", name="Aud", filters=[], openapi_service=openapi_service + ) + return audience, openapi_service + + +def test_assign_job_warns_on_cost_warning(monkeypatch): + audience, openapi_service = _make_audience() + response = MagicMock() + response.job_id = "job-1" + response.cost_warning = CreateJobEndpointCostWarningModel( + estimatedCost=120.0, availableBalance=50.0, shortfall=70.0 + ) + openapi_service.order.job_api.job_post.return_value = response + + warn = MagicMock() + monkeypatch.setattr(base_module.logger, "warning", warn) + + audience.assign_job(MagicMock(id="def-1", name="My Job")) + + warn.assert_called_once() + args = warn.call_args.args + assert "estimate" in args[0].lower() + # Values are passed as lazy %-args, not pre-formatted into the message. + assert 120.0 in args + assert 50.0 in args + assert 70.0 in args + + +def test_assign_job_no_warning_when_cost_within_balance(monkeypatch): + audience, openapi_service = _make_audience() + response = MagicMock() + response.job_id = "job-1" + response.cost_warning = None + openapi_service.order.job_api.job_post.return_value = response + + warn = MagicMock() + monkeypatch.setattr(base_module.logger, "warning", warn) + + job = audience.assign_job(MagicMock(id="def-1", name="My Job")) + + warn.assert_not_called() + assert job.id == "job-1" diff --git a/tests/rapidata_client/job/__init__.py b/tests/rapidata_client/job/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/rapidata_client/job/test_rapidata_job.py b/tests/rapidata_client/job/test_rapidata_job.py new file mode 100644 index 000000000..6512e4ff8 --- /dev/null +++ b/tests/rapidata_client/job/test_rapidata_job.py @@ -0,0 +1,101 @@ +"""Tests for RapidataJob surfacing states it can't progress out of on its own. + +A job that enters ManualApproval or SpendLimited never reaches Completed/Failed, +so waiting on it (e.g. via get_results) must raise an informative error instead +of blocking the caller forever. +""" + +from __future__ import annotations + +import json +from unittest.mock import MagicMock + +import pytest + +from rapidata.api_client.models.audience_job_state import AudienceJobState +from rapidata.api_client.models.review_reason_model import ReviewReasonModel +from rapidata.rapidata_client.job.rapidata_job import RapidataJob + + +def _job_get(state: str, review_reason: ReviewReasonModel | None = None) -> MagicMock: + """A stand-in for the job GET response with just the fields the code reads.""" + job = MagicMock() + job.state = AudienceJobState(state) + job.review_reason = review_reason + job.failure_message = None + return job + + +def _make_job(job_get: MagicMock) -> tuple[RapidataJob, MagicMock]: + openapi_service = MagicMock() + openapi_service.environment = "rapidata.ai" + openapi_service.order.job_api.job_job_id_get.return_value = job_get + job = RapidataJob( + job_id="job-1", + name="My Job", + audience_id="aud-1", + created_at=MagicMock(), + definition_id="def-1", + openapi_service=openapi_service, + ) + return job, openapi_service + + +def test_get_results_raises_on_manual_approval_with_reason(): + job, _ = _make_job( + _job_get("ManualApproval", review_reason=ReviewReasonModel.CONTENTFLAGGED) + ) + + with pytest.raises(Exception) as excinfo: + job.get_results() + + message = str(excinfo.value) + assert "being reviewed" in message + assert "ContentFlagged" in message + + +def test_get_results_raises_on_manual_approval_without_reason(): + job, _ = _make_job(_job_get("ManualApproval", review_reason=None)) + + with pytest.raises(Exception) as excinfo: + job.get_results() + + message = str(excinfo.value) + assert "being reviewed" in message + # A null reason must not leak into the message as empty parentheses. + assert "()" not in message + + +def test_get_results_raises_on_spend_limited(): + job, _ = _make_job(_job_get("SpendLimited")) + + with pytest.raises(Exception) as excinfo: + job.get_results() + + message = str(excinfo.value).lower() + assert "spend-limited" in message + assert "partial results" in message + assert "top up" in message + + +def test_display_progress_bar_raises_on_spend_limited(): + job, _ = _make_job(_job_get("SpendLimited")) + + with pytest.raises(Exception) as excinfo: + job.display_progress_bar() + + assert "spend-limited" in str(excinfo.value).lower() + + +def test_get_results_happy_path_returns_results(): + job, openapi_service = _make_job(_job_get("Completed")) + openapi_service.order.job_api.job_job_id_download_results_get.return_value = ( + json.dumps({"info": {}, "results": []}) + ) + + results = job.get_results() + + assert results == {"info": {}, "results": []} + openapi_service.order.job_api.job_job_id_download_results_get.assert_called_once_with( + job_id="job-1" + ) diff --git a/uv.lock b/uv.lock index 0f823f568..b04fdaae5 100644 --- a/uv.lock +++ b/uv.lock @@ -799,6 +799,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/a4/ed/1f1afb2e9e7f38a545d628f864d562a5ae64fe6f7a10e28ffb9b185b4e89/importlib_resources-6.5.2-py3-none-any.whl", hash = "sha256:789cfdc3ed28c78b67a06acb8126751ced69a3d5f79c095a98298cd8a760ccec", size = 37461, upload-time = "2025-01-03T18:51:54.306Z" }, ] +[[package]] +name = "iniconfig" +version = "2.3.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/72/34/14ca021ce8e5dfedc35312d08ba8bf51fdd999c576889fc2c24cb97f4f10/iniconfig-2.3.0.tar.gz", hash = "sha256:c76315c77db068650d49c5b56314774a7804df16fee4402c1f19d6d15d8c4730", size = 20503, upload-time = "2025-10-18T21:55:43.219Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/cb/b1/3846dd7f199d53cb17f49cba7e651e9ce294d8497c8c150530ed11865bb8/iniconfig-2.3.0-py3-none-any.whl", hash = "sha256:f631c04d2c48c52b84d0d0549c99ff3859c98df65b3101406327ecc7d53fbf12", size = 7484, upload-time = "2025-10-18T21:55:41.639Z" }, +] + [[package]] name = "ipykernel" version = "7.1.0" @@ -2357,6 +2366,24 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/dc/93/b69052907d032b00c40cb656d21438ec00b3a471733de137a3f65a49a0a0/pyright-1.1.407-py3-none-any.whl", hash = "sha256:6dd419f54fcc13f03b52285796d65e639786373f433e243f8b94cf93a7444d21", size = 5997008, upload-time = "2025-10-24T23:17:13.159Z" }, ] +[[package]] +name = "pytest" +version = "8.4.2" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "colorama", marker = "sys_platform == 'win32'" }, + { name = "exceptiongroup", marker = "python_full_version < '3.11'" }, + { name = "iniconfig" }, + { name = "packaging" }, + { name = "pluggy" }, + { name = "pygments" }, + { name = "tomli", marker = "python_full_version < '3.11'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/a3/5c/00a0e072241553e1a7496d638deababa67c5058571567b92a7eaa258397c/pytest-8.4.2.tar.gz", hash = "sha256:86c0d0b93306b961d58d62a4db4879f27fe25513d4b969df351abdddb3c30e01", size = 1519618, upload-time = "2025-09-04T14:34:22.711Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/a8/a4/20da314d277121d6534b3a980b29035dcd51e6744bd79075a6ce8fa4eb8d/pytest-8.4.2-py3-none-any.whl", hash = "sha256:872f880de3fc3a5bdc88a11b39c9710c3497a547cfa9320bc3c5e62fbf272e79", size = 365750, upload-time = "2025-09-04T14:34:20.226Z" }, +] + [[package]] name = "python-dateutil" version = "2.9.0.post0" @@ -2583,7 +2610,7 @@ wheels = [ [[package]] name = "rapidata" -version = "3.15.6" +version = "3.16.4" source = { editable = "." } dependencies = [ { name = "authlib" }, @@ -2612,6 +2639,7 @@ dev = [ { name = "ipykernel" }, { name = "jupyter" }, { name = "pyright" }, + { name = "pytest" }, { name = "python-dotenv" }, ] docs = [ @@ -2656,6 +2684,7 @@ dev = [ { name = "ipykernel", specifier = ">=7.1.0,<8" }, { name = "jupyter", specifier = ">=1.1.1,<2" }, { name = "pyright", specifier = ">=1.1.407" }, + { name = "pytest", specifier = ">=8.0.0,<9" }, { name = "python-dotenv", specifier = ">=1.0.1,<2" }, ] docs = [