diff --git a/README.md b/README.md index 4c6b3b1..f201d44 100644 --- a/README.md +++ b/README.md @@ -332,6 +332,31 @@ We're building the future of tabular machine learning and would love your involv Each API request consumes usage credits; the cost grows with the number of rows and columns in your dataset. You can check your current usage at [platform.priorlabs.ai/account/usage](https://platform.priorlabs.ai/account/usage). +### Estimate an operation before running it + +```python +from tabpfn_client import estimate_cost + +quote = estimate_cost(X_train, X_test, model_version="v3", n_estimators=8) +print(quote.estimated_cost, quote.pricing_version) +print(quote.inputs) # dimensions and resolved server defaults + +fit_quote = estimate_cost(X_train, operation="thinking_fit", thinking_effort="medium") +``` + +This authenticated call sends only dataset dimensions and settings. It does not +upload feature values or consume quota. Quota v3 must be enabled on the server. +Omitted model version and estimator count use the server's defaults. Quotes use +raw feature counts before preprocessing; `quota_v3` costs are tokens and +`legacy_v2` costs are cell-prediction credits. + +Operations are `predict`, `thinking_fit`, `thinking_predict`, and `cache_predict`. +Thinking fit takes no test dataset and defaults to medium effort. For Thinking +prediction, supply the fitted model's version and actual per-base-estimator +count. Cache quotes assume a hit; a fallback or different fitted estimator count +can change the final charge. A quote does not guarantee dataset eligibility or +model access. There is no local fallback formula when the server cannot quote. + ### Monitoring Usage Track your API usage through response headers: diff --git a/changelog/382.added.md b/changelog/382.added.md new file mode 100644 index 0000000..ca4653a --- /dev/null +++ b/changelog/382.added.md @@ -0,0 +1 @@ +Add `estimate_cost()` to check prediction, Thinking, and cached prediction costs from dataset dimensions without uploading data or consuming quota. diff --git a/src/tabpfn_client/__init__.py b/src/tabpfn_client/__init__.py index f80a4eb..79bd350 100644 --- a/src/tabpfn_client/__init__.py +++ b/src/tabpfn_client/__init__.py @@ -10,6 +10,7 @@ set_access_token, get_api_usage, ) +from tabpfn_client.cost import estimate_cost from tabpfn_client.estimator import TabPFNClassifier, TabPFNRegressor from tabpfn_client.interactive_auth import InteractiveLoginError, interactive_login from tabpfn_client.service_wrapper import UserDataClient @@ -29,6 +30,7 @@ "get_access_token", "set_access_token", "get_api_usage", + "estimate_cost", "interactive_login", "InteractiveLoginError", ] diff --git a/src/tabpfn_client/api_models.py b/src/tabpfn_client/api_models.py index 1b0a7f2..e1f7f4d 100644 --- a/src/tabpfn_client/api_models.py +++ b/src/tabpfn_client/api_models.py @@ -336,3 +336,26 @@ class SubmitFitJobRequest(BaseModel): class SubmitFitJobResponse(BaseModel): fitted_train_set_id: UUID + + +class QuotaOperation(str, Enum): + PREDICT = "predict" + THINKING_PREDICT = "thinking_predict" + THINKING_FIT = "thinking_fit" + CACHE_PREDICT = "cache_predict" + + +class EstimateCostRequest(BaseModel): + train_rows: int + raw_columns: int + test_rows: int | None = None + model_version: Annotated[ModelVersion | UnknownEnum, Field(union_mode="left_to_right")] | None = None + operation: Annotated[QuotaOperation | UnknownEnum, Field(union_mode="left_to_right")] | None = None + n_estimators: int | None = None + thinking_effort: Annotated[ThinkingEffort | str, Field(union_mode="left_to_right")] | None = None + + +class EstimateCostResponse(BaseModel): + estimated_cost: int + pricing_version: str + inputs: EstimateCostRequest diff --git a/src/tabpfn_client/client.py b/src/tabpfn_client/client.py index 45557b6..da87313 100644 --- a/src/tabpfn_client/client.py +++ b/src/tabpfn_client/client.py @@ -40,6 +40,8 @@ from tabpfn_common_utils.utils import Singleton from tabpfn_client.api_models import ( GetSettingsResponse, + EstimateCostRequest, + EstimateCostResponse, TabPFNSystem, PrepareTrainSetUploadRequest, PrepareTrainSetUploadResponse, @@ -1262,6 +1264,18 @@ def delete_user_account(cls) -> None: cls._raise_on_error(response, "delete_user_account") + @classmethod + def estimate_cost( + cls, req: EstimateCostRequest, *, access_token: str + ) -> EstimateCostResponse: + """Quote one operation using dimensions only; never upload or reserve quota.""" + response = cls.httpx_client.post( + "/tabpfn/estimate_cost", + json=req.model_dump(mode="json", exclude_none=True), + headers={"Authorization": f"Bearer {access_token}"}, + ) + return cls._validate_response(response, "estimate_cost", EstimateCostResponse) + @classmethod def get_api_usage(cls, access_token: str): """ diff --git a/src/tabpfn_client/cost.py b/src/tabpfn_client/cost.py new file mode 100644 index 0000000..ada4c96 --- /dev/null +++ b/src/tabpfn_client/cost.py @@ -0,0 +1,71 @@ +"""Estimate quota cost on the server using only dataset dimensions.""" + +from typing import Literal + +import numpy as np +import pandas as pd + +from tabpfn_client.api_models import EstimateCostRequest, EstimateCostResponse +from tabpfn_client.client import ServiceClient +from tabpfn_client.config import get_access_token + + +def _shape(X: np.ndarray | pd.DataFrame, name: str) -> tuple[int, int]: + shape = getattr(X, "shape", None) + if shape is None or len(shape) != 2: + raise ValueError(f"{name} must be a two-dimensional array or DataFrame") + rows, columns = int(shape[0]), int(shape[1]) + if rows < 0 or columns <= 0 or (name == "X_train" and rows == 0): + raise ValueError(f"{name} must have valid row and feature counts") + return rows, columns + + +def estimate_cost( + X_train: np.ndarray | pd.DataFrame, + X_test: np.ndarray | pd.DataFrame | None = None, + *, + model_version: str | None = None, + operation: Literal[ + "predict", "thinking_fit", "thinking_predict", "cache_predict" + ] = "predict", + n_estimators: int | None = None, + thinking_effort: Literal["medium", "high"] | None = None, +) -> EstimateCostResponse: + """Estimate one operation without uploading data or consuming quota. + + Only raw row/column counts and the supplied configuration are sent. The + server resolves omitted model/ensemble defaults and returns them in + ``result.inputs`` alongside ``estimated_cost`` and ``pricing_version``. + Quota v3 must be enabled on the server. Costs are tokens for ``quota_v3`` + and legacy cell-prediction credits for ``legacy_v2``. + + ``thinking_fit`` takes no X_test; its default effort is medium. For + ``thinking_predict``, supply the fitted model's version and actual + per-base-estimator count. ``cache_predict`` assumes a cache hit; fallback + or different fitted estimator counts can change the final charge. + A quote does not guarantee model access or dataset eligibility. + """ + train_rows, raw_columns = _shape(X_train, "X_train") + test_rows = 0 + if X_test is not None: + if operation == "thinking_fit": + raise ValueError("thinking_fit does not use X_test") + test_rows, test_columns = _shape(X_test, "X_test") + if test_columns != raw_columns: + raise ValueError("X_train and X_test must have the same number of features") + if n_estimators is not None and ( + type(n_estimators) is not int or n_estimators <= 0 + ): + raise ValueError("n_estimators must be a positive integer") + req = EstimateCostRequest.model_validate( + { + "train_rows": train_rows, + "test_rows": test_rows, + "raw_columns": raw_columns, + "model_version": model_version, + "operation": operation, + "n_estimators": n_estimators, + "thinking_effort": thinking_effort, + } + ) + return ServiceClient.estimate_cost(req, access_token=get_access_token()) diff --git a/tests/unit/test_estimate_cost.py b/tests/unit/test_estimate_cost.py new file mode 100644 index 0000000..e88c535 --- /dev/null +++ b/tests/unit/test_estimate_cost.py @@ -0,0 +1,132 @@ +import json +from types import SimpleNamespace +from typing import Literal, cast +from unittest.mock import Mock + +import httpx +import numpy as np +import pandas as pd +import pytest + +from tabpfn_client import estimate_cost +from tabpfn_client.config import Config +from tabpfn_client.api_models import EstimateCostResponse +from tabpfn_client.client import ServiceClient +from tabpfn_client.errors import RetryableServerError + + +@pytest.fixture +def transport(monkeypatch): + handler = Mock() + monkeypatch.setattr(Config, "is_initialized", True) + monkeypatch.setattr(ServiceClient, "_access_token", "estimate-test-token") + with httpx.Client( + base_url="https://estimate.invalid", transport=httpx.MockTransport(handler) + ) as client: + monkeypatch.setattr(ServiceClient, "httpx_client", client) + yield handler + + +def quote(request): + body = json.loads(request.content) + return httpx.Response( + 200, + json={ + "estimated_cost": 20284, + "pricing_version": "quota_v3", + "inputs": { + "model_version": "v3", + "n_estimators": 8, + "thinking_effort": None, + **body, + }, + }, + ) + + +@pytest.mark.parametrize("as_frame", [False, True]) +def test_estimate_posts_only_shape_and_options(transport, as_frame): + transport.side_effect = quote + X_train = np.zeros((3, 2)) + X_test = np.zeros((4, 2)) + if as_frame: + X_train = pd.DataFrame(X_train) + X_test = pd.DataFrame(X_test) + result = estimate_cost(X_train, X_test) + assert isinstance(result, EstimateCostResponse) + assert result.estimated_cost == 20284 + assert result.inputs.n_estimators == 8 + transport.assert_called_once() + request = transport.call_args.args[0] + assert request.method == "POST" + assert request.url.path == "/tabpfn/estimate_cost" + assert request.headers["Authorization"] == "Bearer estimate-test-token" + assert json.loads(request.content) == { + "train_rows": 3, + "test_rows": 4, + "raw_columns": 2, + "operation": "predict", + } + + +@pytest.mark.parametrize( + "operation", ["thinking_fit", "thinking_predict", "cache_predict"] +) +def test_options_and_shape_without_reading_array_values( + transport, + operation: Literal["thinking_fit", "thinking_predict", "cache_predict"], +): + transport.side_effect = quote + # No contents, array conversion, or upload methods exist on this object. + X = cast(np.ndarray, SimpleNamespace(shape=(100000, 100))) + effort: Literal["high"] | None = "high" if operation == "thinking_fit" else None + estimate_cost( + X, + model_version="v3.5", + operation=operation, + n_estimators=4, + thinking_effort=effort, + ) + body = json.loads(transport.call_args.args[0].content) + assert body == { + "train_rows": 100000, + "test_rows": 0, + "raw_columns": 100, + "operation": operation, + "model_version": "v3.5", + "n_estimators": 4, + **({"thinking_effort": "high"} if operation == "thinking_fit" else {}), + } + + +@pytest.mark.parametrize( + "X_train,X_test,options", + [ + (np.zeros(3), None, {}), + (np.zeros((0, 2)), None, {}), + (np.zeros((2, 0)), None, {}), + (np.zeros((2, 3)), np.zeros((2, 4)), {}), + (np.zeros((2, 3)), np.zeros(2), {}), + (np.zeros((2, 3)), np.zeros((2, 3)), {"operation": "thinking_fit"}), + (np.zeros((2, 3)), None, {"n_estimators": True}), + (np.zeros((2, 3)), None, {"n_estimators": 1.5}), + (np.zeros((2, 3)), None, {"n_estimators": 0}), + ], +) +def test_invalid_dimensions_and_counts_do_not_make_http_requests( + transport, X_train, X_test, options +): + with pytest.raises(ValueError): + estimate_cost(X_train, X_test, **options) + transport.assert_not_called() + + +@pytest.mark.parametrize("status", [401, 422, 503]) +def test_server_errors_surface_without_local_price_fallback(transport, status): + transport.return_value = httpx.Response( + status, json={"detail": "estimate unavailable"} + ) + error = RetryableServerError if status == 503 else RuntimeError + with pytest.raises(error, match="estimate unavailable"): + estimate_cost(np.zeros((2, 3))) + transport.assert_called_once()