Add API-backed dataset cost estimates

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:

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.

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",
 ]

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

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):
         """

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())

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()