From a34826652d3bb50dd44e518626b621ad9b00569c Mon Sep 17 00:00:00 2001 From: Stefan Hackmann Date: Mon, 14 Sep 2026 14:06:24 +0000 Subject: [PATCH 1/3] docs: Update AGENTS.md --- AGENTS.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/AGENTS.md b/AGENTS.md index e5c2631..823357d 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -1 +1 @@ -Agents must ground their thoughts in facts, not assumptions: before planning, claiming, or editing, read the relevant material — source, configs, data, docs, test output — and only act on beliefs backed by something just read or run. Agents must run ad-hoc Python and CLI commands through `task run -- ` (which wraps `uv run` with the project's canonical parameters), never bare `python` or hand-written `uv run` invocations — this keeps every invocation from accidentally re-resolving or mutating the `.venv` and `uv.lock`. Agents must explain what each command will do and why it is being run before running it. Before asking the user a question, agents must first explain the corresponding context and terminology — what the question concerns, why it arises, and what any project-specific terms mean — so the user can answer without digging through the code themselves. Agents must run `task pre-commit` and fix all reported issues before reporting success to the user. After `task pre-commit` succeeds, show the diff and explain why each change is necessary before reporting success to the user. Agents must never perform destructive or state-changing git operations unless the user explicitly instructs the agent to run that specific operation — no `git push --force`, no `git reset --hard`, no `git stash` (which hides uncommitted work), no branch/tag deletion, no history rewrites (`rebase`, `commit --amend` on published commits, `filter-branch`), no `git clean -fdx`, no `--no-verify` to bypass hooks, and no discarding of uncommitted work. Read-only inspection commands (`git status`, `git diff`, `git log`, `git show`) are always allowed. If a task seems to require a state-changing git operation and the user has not explicitly asked for it, stop and ask the user to run it. Agents must never mention temporary planning identifiers — phase names or numbers such as "Phase 3g", milestone, sprint, or ticket codes — in docstrings, comments, help or description strings, error messages, or test docstrings; those labels are deleted when the plan is retired and leave readers with a dangling reference nobody can decode, so describe the concept by its lasting behavior or stable configuration key instead (phase names belong only in roadmap and planning docs, and linking to such a doc by its actual filename is fine). Agents must not mention or recommend key rotation (rotating API keys, tokens, or other credentials) — the user takes care of key rotation themselves. Agents must keep comments in configuration files value-independent: a comment attached to a configuration parameter — in `config.yaml`/`config.sample.yaml`, a `Taskfile`, `pyproject.toml`, or any other settings file — may only state what remains true whatever the value is (what the parameter controls, what its options mean, which invariant or trade-off it participates in), and must never restate the current value, derive from it, or do arithmetic on it (`# 64 of 576 rows`, `# 2x the 288-step horizon`, `# keeps 90% of draws at full context`); configuration is retuned constantly, so such a comment turns false the moment somebody edits the value on the line below it and nothing catches the drift, so express a needed derivation as a rule over the parameter or record concrete numbers in a dated notes doc where they read as a historical observation. ripgrep (`rg`) is available (installed by `task setup` via `task install:ripgrep`); prefer it for fast code and text search. +Agents must ground their thoughts in facts, not assumptions: before planning, claiming, or editing, read the relevant material — source, configs, data, docs, test output — and only act on beliefs backed by something just read or run. Agents must run ad-hoc Python and CLI commands through `task run -- ` (which wraps `uv run` with the project's canonical parameters), never bare `python` or hand-written `uv run` invocations — this keeps every invocation from accidentally re-resolving or mutating the `.venv` and `uv.lock`. Agents must explain what each command will do and why it is being run before running it. Before asking the user a question, agents must first explain the corresponding context and terminology — what the question concerns, why it arises, and what any project-specific terms mean — so the user can answer without digging through the code themselves. Agents must run `task pre-commit` and fix all reported issues before reporting success to the user. After `task pre-commit` succeeds, show the diff and explain why each change is necessary before reporting success to the user. When what an agent reports is a set of items measured or classified along the same few dimensions — files and their status, config keys and their values, metrics before and after, options and their trade-offs — it must be presented as a Markdown table instead of prose or a bullet list, with short cells and only the columns that carry information; prose stays the right form for a single item, a narrative explanation, or an ordered procedure, because a one-row table, or one whose cells are sentences, is harder to read than the paragraph it replaced. Agents must never perform destructive or state-changing git operations unless the user explicitly instructs the agent to run that specific operation — no `git push --force`, no `git reset --hard`, no `git stash` (which hides uncommitted work), no branch/tag deletion, no history rewrites (`rebase`, `commit --amend` on published commits, `filter-branch`), no `git clean -fdx`, no `--no-verify` to bypass hooks, and no discarding of uncommitted work. Read-only inspection commands (`git status`, `git diff`, `git log`, `git show`) are always allowed. If a task seems to require a state-changing git operation and the user has not explicitly asked for it, stop and ask the user to run it. Agents must never mention temporary planning identifiers — phase names or numbers such as "Phase 3g", milestone, sprint, or ticket codes — in docstrings, comments, help or description strings, error messages, or test docstrings; those labels are deleted when the plan is retired and leave readers with a dangling reference nobody can decode, so describe the concept by its lasting behavior or stable configuration key instead (phase names belong only in roadmap and planning docs, and linking to such a doc by its actual filename is fine). Agents must not mention or recommend key rotation (rotating API keys, tokens, or other credentials) — the user takes care of key rotation themselves. Agents must keep comments in configuration files value-independent: a comment attached to a configuration parameter — in `config.yaml`/`config.sample.yaml`, a `Taskfile`, `pyproject.toml`, or any other settings file — may only state what remains true whatever the value is (what the parameter controls, what its options mean, which invariant or trade-off it participates in), and must never restate the current value, derive from it, or do arithmetic on it (`# 64 of 576 rows`, `# 2x the 288-step horizon`, `# keeps 90% of draws at full context`); configuration is retuned constantly, so such a comment turns false the moment somebody edits the value on the line below it and nothing catches the drift, so express a needed derivation as a rule over the parameter or record concrete numbers in a dated notes doc where they read as a historical observation. ripgrep (`rg`) is available (installed by `task setup` via `task install:ripgrep`); prefer it for fast code and text search. From d32614df78c81cc97425dd98249d89244d3536f0 Mon Sep 17 00:00:00 2001 From: Stefan Hackmann Date: Thu, 17 Sep 2026 20:20:45 +0000 Subject: [PATCH 2/3] feat: schema v3 for conditioning --- config.sample.yaml | 2 +- src/jointfm_client/__init__.py | 16 + src/jointfm_client/adapters.py | 11 +- src/jointfm_client/client.py | 28 +- src/jointfm_client/contract.py | 413 +++++++++++++++++- .../fixtures/condition_interval_response.json | 33 ++ tests/fixtures/condition_mean_request.json | 41 ++ tests/fixtures/condition_mean_response.json | 30 ++ tests/fixtures/forecast_mean_request.json | 2 +- tests/fixtures/forecast_mean_response.json | 2 +- .../fixtures/forecast_quantiles_request.json | 2 +- .../fixtures/forecast_quantiles_response.json | 2 +- tests/fixtures/forecast_samples_request.json | 2 +- tests/fixtures/forecast_samples_response.json | 2 +- tests/fixtures/health_metadata.json | 3 +- .../input_size_exceeded_response.json | 2 +- .../model_version_mismatch_response.json | 2 +- .../schema_version_mismatch_response.json | 2 +- tests/fixtures/validation_error_response.json | 2 +- tests/test_cli.py | 7 +- tests/test_condition_mode.py | 353 +++++++++++++++ tests/test_configuration.py | 7 +- tests/test_contract.py | 31 +- tests/test_contract_models.py | 14 +- tests/test_feature_importance.py | 2 +- tests/test_pool.py | 19 +- tests/test_settings.py | 14 +- tests/test_surfaces.py | 4 +- tests/test_transport.py | 59 +-- 29 files changed, 1012 insertions(+), 95 deletions(-) create mode 100644 tests/fixtures/condition_interval_response.json create mode 100644 tests/fixtures/condition_mean_request.json create mode 100644 tests/fixtures/condition_mean_response.json create mode 100644 tests/test_condition_mode.py diff --git a/config.sample.yaml b/config.sample.yaml index 05064be..768f4c5 100644 --- a/config.sample.yaml +++ b/config.sample.yaml @@ -52,7 +52,7 @@ transport: - X-DataRobot-Execution-ID user_agent_header: User-Agent forecast: - schema_version: v2 + schema_version: v3 query_mode: forecast return_mode: mean time_index_mode: ordinal diff --git a/src/jointfm_client/__init__.py b/src/jointfm_client/__init__.py index 79de7bb..5972c76 100644 --- a/src/jointfm_client/__init__.py +++ b/src/jointfm_client/__init__.py @@ -67,6 +67,13 @@ PACKAGE_VERSION, PREDICT_REQUEST_TYPE, SCHEMA_VERSION, + ConditionBlock, + ConditionKind, + ConditionPlausibility, + EqualityCondition, + IntervalCondition, + IntervalEstimator, + require_condition_support, MeanForecastResult, QuantileForecast, QuantileForecastResult, @@ -75,6 +82,7 @@ SUPPORTED_COLUMN_ROLES, STRUCTURED_ERROR_CODES, SUPPORTED_DECODING_STRATEGIES, + SUPPORTED_CONDITION_KINDS, SUPPORTED_QUERY_MODES, SUPPORTED_REQUEST_TYPES, SUPPORTED_RETURN_MODES, @@ -212,11 +220,19 @@ "PREDICT_REQUEST_TYPE", "RetryConfig", "SCHEMA_VERSION", + "ConditionBlock", + "ConditionKind", + "ConditionPlausibility", + "EqualityCondition", + "IntervalCondition", + "IntervalEstimator", + "require_condition_support", "SampleForecastResult", "SUPPORTED_COLUMN_MODALITIES", "SUPPORTED_COLUMN_ROLES", "STRUCTURED_ERROR_CODES", "SUPPORTED_DECODING_STRATEGIES", + "SUPPORTED_CONDITION_KINDS", "SUPPORTED_QUERY_MODES", "SUPPORTED_RETURN_MODES", "SUPPORTED_TIME_INDEX_MODES", diff --git a/src/jointfm_client/adapters.py b/src/jointfm_client/adapters.py index d23b861..988cb22 100644 --- a/src/jointfm_client/adapters.py +++ b/src/jointfm_client/adapters.py @@ -24,6 +24,7 @@ from jointfm_client.configuration import DEFAULT_FORECAST_SCHEMA_VERSION from jointfm_client.contract import ( DEFAULT_CALENDAR_ID, + ConditionBlock, ColumnModality, ColumnRole, ColumnSpec, @@ -296,8 +297,14 @@ def build_forecast_payload_from_dataframe( time_value_columns: Sequence[str] | Mapping[str, TimeValueKind] | None = None, nullable_columns: Sequence[str] | None = None, bounds: ColumnBounds | None = None, + condition: ConditionBlock | None = None, ) -> dict[str, Any]: - """Build a validated forecast payload from a pandas ``DataFrame``.""" + """Build a validated forecast payload from a pandas ``DataFrame``. + + Passing ``condition`` makes this a conditioning request: the payload then + carries ``query_mode='condition'`` and the block, and the deployment + answers the conditional at the one future position the block names. + """ column_specs = ( infer_column_specs_from_dataframe( frame, @@ -346,6 +353,8 @@ def build_forecast_payload_from_dataframe( quantiles=quantiles, seed=seed, schema_version=schema_version, + query_mode="forecast" if condition is None else "condition", + condition=condition, ) diff --git a/src/jointfm_client/client.py b/src/jointfm_client/client.py index ec252da..f0d9b19 100644 --- a/src/jointfm_client/client.py +++ b/src/jointfm_client/client.py @@ -36,6 +36,7 @@ load_configuration, ) from jointfm_client.contract import ( + ConditionBlock, DEFAULT_CALENDAR_ID, HEALTH_REQUEST_TYPE, SCHEMA_VERSION, @@ -47,6 +48,7 @@ TimeIndexMode, TimeValueKind, build_forecast_payload, + require_condition_support, validate_service_metadata, ) from jointfm_client.exceptions import ( @@ -310,9 +312,18 @@ def forecast( nullable_columns: Sequence[str] | None = None, bounds: Mapping[str, tuple[float | int | None, float | int | None]] | None = None, + condition: ConditionBlock | None = None, ) -> ForecastResponse: - """Build and submit a forecast request from tabular history inputs.""" + """Build and submit a forecast request from tabular history inputs. + + Passing ``condition`` asks the deployment for the conditional at the one + future position the block names, instead of the unconditional forecast. + The deployment's advertised capability is checked first, so a deployment + that cannot condition is refused here rather than after a round trip. + """ self._require_predict_url("forecast") + if condition is not None: + require_condition_support(self.health(cache=True), condition) resolved_model_version = self._resolve_model_version( model_version=model_version, ) @@ -336,6 +347,7 @@ def forecast( use_local_normalized_time=use_local_normalized_time, calendar_id=calendar_id, timezone=timezone, + condition=condition, ) else: payload = build_forecast_payload_from_dataframe( @@ -369,6 +381,7 @@ def forecast( time_value_columns=time_value_columns, nullable_columns=nullable_columns, bounds=bounds, + condition=condition, ) sample_cap = self._resolve_sample_batch_cap(payload) if sample_cap is not None: @@ -636,6 +649,7 @@ def _forecast_payload_from_rows( use_local_normalized_time: bool, calendar_id: str, timezone: str | None, + condition: ConditionBlock | None = None, ) -> dict[str, Any]: if schema is None: if columns is None: @@ -662,6 +676,8 @@ def _forecast_payload_from_rows( quantiles=quantiles, seed=seed, schema_version=schema_version, + query_mode="forecast" if condition is None else "condition", + condition=condition, ) def _resolve_sample_batch_cap(self, payload: Mapping[str, Any]) -> int | None: @@ -1095,7 +1111,17 @@ def _merge_sample_forecast_results( history_rows=first_result.diagnostics.history_rows, horizon_count=first_result.diagnostics.horizon_count, seed=diagnostics_seed, + condition_draws=( + None + if first_result.diagnostics.condition_draws is None + else len(merged_samples) + ), + interval_estimator=first_result.diagnostics.interval_estimator, ), + # Every batch answers the same request, so what the model thinks of the + # conditions does not vary across them; the region probability reported + # here is the first batch's estimate of it. + plausibility=first_result.plausibility, errors=(), samples=merged_samples, ) diff --git a/src/jointfm_client/contract.py b/src/jointfm_client/contract.py index ba40b09..1675796 100644 --- a/src/jointfm_client/contract.py +++ b/src/jointfm_client/contract.py @@ -37,7 +37,7 @@ # pyproject.toml the way a hand-maintained literal here did. PACKAGE_VERSION: Final = importlib.metadata.version(DISTRIBUTION_NAME) -SCHEMA_VERSION: Final = "v2" +SCHEMA_VERSION: Final = "v3" DATAROBOT_UNSTRUCTURED_PREDICTION_ROUTE_TEMPLATE: Final = ( "deployments/{deployment_id}/predictionsUnstructured" ) @@ -52,7 +52,8 @@ HEALTH_REQUEST_TYPE, ) -QueryMode: TypeAlias = Literal["forecast"] +QueryMode: TypeAlias = Literal["forecast", "condition"] +ConditionKind: TypeAlias = Literal["equality", "interval"] ReturnMode: TypeAlias = Literal["mean", "samples", "quantiles", "log_prob"] TimeIndexMode: TypeAlias = Literal[ "ordinal", @@ -82,13 +83,16 @@ TimeValueKind: TypeAlias = Literal["continuous_float", "absolute_datetime"] StructuredErrorCode: TypeAlias = Literal[ "VALIDATION_ERROR", + "UNSUPPORTED_HEAD_QUERY_COMBINATION", + "UNSUPPORTED_RETURN_MODE", "SCHEMA_VERSION_MISMATCH", "MODEL_VERSION_MISMATCH", "INPUT_SIZE_EXCEEDED", "INTERNAL_ERROR", ] -SUPPORTED_QUERY_MODES: Final[tuple[QueryMode, ...]] = ("forecast",) +SUPPORTED_QUERY_MODES: Final[tuple[QueryMode, ...]] = ("forecast", "condition") +SUPPORTED_CONDITION_KINDS: Final[tuple[ConditionKind, ...]] = ("equality", "interval") SUPPORTED_RETURN_MODES: Final[tuple[ReturnMode, ...]] = ( "mean", "samples", @@ -126,6 +130,8 @@ ) STRUCTURED_ERROR_CODES: Final[tuple[StructuredErrorCode, ...]] = ( "VALIDATION_ERROR", + "UNSUPPORTED_HEAD_QUERY_COMBINATION", + "UNSUPPORTED_RETURN_MODE", "SCHEMA_VERSION_MISMATCH", "MODEL_VERSION_MISMATCH", "INPUT_SIZE_EXCEEDED", @@ -302,6 +308,150 @@ def to_payload(self) -> dict[str, Any]: return payload +@dataclass(frozen=True, slots=True) +class EqualityCondition: + """One column of the conditioned position pinned to a value. + + An equality condition is an event of probability zero under a continuous + head, so the deployment answers it analytically rather than by filtering + draws. It fixes the column's value, which is why that column leaves the + read-out projection. + """ + + column: str + value: float + + def __post_init__(self) -> None: + """Reject a pin the service would refuse anyway.""" + _require_string(self.column, field="condition.column") + if isinstance(self.value, bool) or not isinstance(self.value, (int, float)): + raise ValueError("condition value must be a number") + if not math.isfinite(float(self.value)): + raise ValueError("condition value must be finite") + + def to_payload(self) -> dict[str, Any]: + """Return this condition as JSON-compatible payload fields.""" + return {"column": self.column, "kind": "equality", "value": float(self.value)} + + +@dataclass(frozen=True, slots=True) +class IntervalCondition: + """One column of the conditioned position confined to a range. + + ``None`` on either side leaves it open. The region carries positive mass, + so the deployment answers it by composition on the reduced mixture, and the + column stays readable: what it returns is its distribution inside the + region. + """ + + column: str + lower: float | None = None + upper: float | None = None + + def __post_init__(self) -> None: + """Reject a range that bounds nothing or bounds it backwards.""" + _require_string(self.column, field="condition.column") + for bound, name in ((self.lower, "lower"), (self.upper, "upper")): + if bound is None: + continue + if isinstance(bound, bool) or not isinstance(bound, (int, float)): + raise ValueError(f"condition {name} bound must be a number or None") + if not math.isfinite(float(bound)): + raise ValueError(f"condition {name} bound must be finite or None") + if self.lower is None and self.upper is None: + raise ValueError( + "an interval open on both sides conditions nothing; give at least one bound" + ) + if ( + self.lower is not None + and self.upper is not None + and float(self.lower) >= float(self.upper) + ): + raise ValueError( + f"condition needs lower < upper, got [{self.lower}, {self.upper}]" + ) + + def to_payload(self) -> dict[str, Any]: + """Return this condition as JSON-compatible payload fields.""" + return { + "column": self.column, + "kind": "interval", + "lower": None if self.lower is None else float(self.lower), + "upper": None if self.upper is None else float(self.upper), + } + + +@dataclass(frozen=True, slots=True) +class ConditionBlock: + """Every condition of one request, against one future position. + + The position is named once, by its index into the request's + ``query_times``, because conditioning relates variables to each other + within a position and never across horizons. A statement spanning + horizons is out of reach of this mode and cannot be expressed here. + """ + + query_time_index: int + conditions: Sequence[EqualityCondition | IntervalCondition] + + def __post_init__(self) -> None: + """Validate the block independently of any schema or deployment.""" + if isinstance(self.query_time_index, bool) or not isinstance( + self.query_time_index, int + ): + raise ValueError("query_time_index must be an integer") + if self.query_time_index < 0: + raise ValueError("query_time_index must not be negative") + conditions = _require_sequence(self.conditions, field="conditions") + if not conditions: + raise ValueError("a condition block must carry at least one condition") + seen: set[str] = set() + for index, condition in enumerate(conditions): + if not isinstance(condition, (EqualityCondition, IntervalCondition)): + raise ValueError( + f"conditions[{index}] must be an EqualityCondition or an " + "IntervalCondition" + ) + if condition.column in seen: + raise ValueError( + f"column {condition.column!r} carries more than one condition" + ) + seen.add(condition.column) + + @property + def pinned_columns(self) -> tuple[str, ...]: + """Columns an equality condition fixes, which leave the read-out set.""" + return tuple( + condition.column + for condition in self.conditions + if isinstance(condition, EqualityCondition) + ) + + @property + def conditioned_columns(self) -> tuple[str, ...]: + """Every column this block conditions, of either kind.""" + return tuple(condition.column for condition in self.conditions) + + @property + def kinds(self) -> tuple[ConditionKind, ...]: + """The condition kinds this block uses, which a deployment must advertise.""" + kinds: list[ConditionKind] = [] + for condition in self.conditions: + kind: ConditionKind = ( + "equality" if isinstance(condition, EqualityCondition) else "interval" + ) + if kind not in kinds: + kinds.append(kind) + return tuple(kinds) + + def to_payload(self) -> dict[str, Any]: + """Return the block as JSON-compatible payload fields.""" + return { + "query_time_index": self.query_time_index, + "conditions": [condition.to_payload() for condition in self.conditions], + } + + @dataclass(frozen=True, slots=True) class ForecastRequestMetadata: """Version and mode metadata for one forecast request.""" @@ -352,6 +502,7 @@ class ForecastRequest: quantiles: Sequence[float | int] | None = None seed: int | None = None query_row_ids: Sequence[int] | None = None + condition: ConditionBlock | None = None def __post_init__(self) -> None: """Validate payload controls and JSON-facing request arrays.""" @@ -373,6 +524,14 @@ def __post_init__(self) -> None: _optional_positive_int(self.n_samples, field="n_samples") _optional_int(self.seed, field="seed") + if (self.metadata.query_mode == "condition") != (self.condition is not None): + raise ValueError( + "query_mode='condition' and a condition block go together: one " + "without the other means a different request than the caller wrote" + ) + if self.condition is not None: + self._validate_condition() + if self.metadata.return_mode == "quantiles": _require_quantiles(self.quantiles) elif self.quantiles is not None: @@ -380,6 +539,50 @@ def __post_init__(self) -> None: "quantiles may be provided only when return_mode='quantiles'" ) + def _validate_condition(self) -> None: + """Check the condition block against this request's schema and horizons. + + The deployment rejects all of this too, before executing anything; the + point of repeating it here is that a caller learns which column name it + got wrong without paying for a round trip. + """ + block = self.condition + assert block is not None + declared = {column.name for column in self.schema.columns} + unknown = [name for name in block.conditioned_columns if name not in declared] + if unknown: + raise ValueError(f"condition references undeclared columns {unknown}") + horizon_count = len(_require_sequence(self.query_times, field="query_times")) + if block.query_time_index >= horizon_count: + raise ValueError( + f"condition.query_time_index {block.query_time_index} is outside the " + f"{horizon_count} requested future positions" + ) + if len(set(block.conditioned_columns)) >= len(declared): + if len(block.pinned_columns) == len(set(block.conditioned_columns)): + raise ValueError( + "pinning every column leaves nothing to read out; the joint " + "density of those values is what return_mode='log_prob' answers" + ) + raise ValueError( + "a request must leave at least one column unconditioned: what it " + "reads out is the conditional distribution of the columns it does " + "not condition" + ) + requested = _resolve_requested_columns( + self.schema.columns, self.requested_columns + ) + if requested is None: + return + requested_names = {value for value in requested if isinstance(value, str)} + pinned = sorted(set(block.pinned_columns) & requested_names) + if pinned: + raise ValueError( + f"requested_columns lists pinned columns {pinned}; an equality " + "condition fixes the value, so reading it back returns only what " + "the request supplied" + ) + def to_payload(self) -> dict[str, Any]: """Return a JSON-compatible forecast request without mutating inputs.""" payload = self.metadata.to_payload() @@ -401,6 +604,8 @@ def to_payload(self) -> dict[str, Any]: payload["quantiles"] = _require_quantiles(self.quantiles) if self.seed is not None: payload["seed"] = self.seed + if self.condition is not None: + payload["condition"] = self.condition.to_payload() return payload @@ -486,6 +691,7 @@ class HealthMetadata: head: str decoding_strategy: DecodingStrategy supported_query_modes: tuple[str, ...] + supported_condition_kinds: tuple[str, ...] supported_return_modes: tuple[str, ...] supported_time_index_modes: tuple[str, ...] time_index_encoding: str @@ -534,6 +740,13 @@ def from_payload(cls, payload: Mapping[str, Any]) -> Self: payload.get("supported_query_modes"), field="supported_query_modes", ), + # A deployment that cannot condition advertises an empty list, + # which is an answer rather than a missing field. + supported_condition_kinds=_string_tuple( + payload.get("supported_condition_kinds"), + field="supported_condition_kinds", + allow_empty=True, + ), supported_return_modes=_string_tuple( payload.get("supported_return_modes"), field="supported_return_modes", @@ -573,6 +786,70 @@ def from_payload(cls, payload: Mapping[str, Any]) -> Self: ) +@dataclass(frozen=True, slots=True) +class IntervalEstimator: + """How a multi-column region's probability was estimated. + + Present only when the interval block spans more than one column, which is + the only case where a component's box probability has to be estimated; one + interval column is exact, and a number that never varies would say nothing. + """ + + points: int + effective_sample_size: float + + @classmethod + def from_payload(cls, payload: Mapping[str, Any]) -> Self: + """Parse the estimator accounting from a response payload.""" + return cls( + points=_require_positive_int( + payload.get("points"), + field="diagnostics.interval_estimator.points", + ), + effective_sample_size=_require_output_float( + payload.get("effective_sample_size"), + field="diagnostics.interval_estimator.effective_sample_size", + ), + ) + + +@dataclass(frozen=True, slots=True) +class ConditionPlausibility: + """What the model thinks of the request it was asked to condition on. + + Without it a caller cannot separate "the model is confident" from "you + conditioned on something the model finds absurd". Both numbers are in log + space, and each is ``None`` when the request carried no condition of that + kind. The service reports them and never refuses on them, so acting on them + is the caller's decision. + """ + + equality_log_density: float | None = None + region_log_probability: float | None = None + + @classmethod + def from_payload(cls, payload: Mapping[str, Any]) -> Self: + """Parse the plausibility block from a response payload.""" + equality = payload.get("equality_log_density") + region = payload.get("region_log_probability") + return cls( + equality_log_density=( + None + if equality is None + else _require_output_float( + equality, field="plausibility.equality_log_density" + ) + ), + region_log_probability=( + None + if region is None + else _require_output_float( + region, field="plausibility.region_log_probability" + ) + ), + ) + + @dataclass(frozen=True, slots=True) class ForecastDiagnostics: """Diagnostics block returned with one forecast response.""" @@ -580,11 +857,27 @@ class ForecastDiagnostics: history_rows: int horizon_count: int seed: int | None = None + condition_draws: int | None = None + interval_estimator: IntervalEstimator | None = None @classmethod def from_payload(cls, payload: Mapping[str, Any]) -> Self: """Parse response diagnostics from a service payload.""" + raw_estimator = payload.get("interval_estimator") return cls( + condition_draws=_optional_positive_int( + payload.get("condition_draws"), + field="diagnostics.condition_draws", + ), + interval_estimator=( + None + if raw_estimator is None + else IntervalEstimator.from_payload( + _require_mapping( + raw_estimator, field="diagnostics.interval_estimator" + ) + ) + ), history_rows=_require_positive_int( payload.get("history_rows"), field="diagnostics.history_rows", @@ -710,6 +1003,7 @@ class ForecastResponse: requested_columns: tuple[str, ...] diagnostics: ForecastDiagnostics errors: tuple[StructuredError, ...] + plausibility: ConditionPlausibility | None @property def outputs(self) -> ForecastOutputs: @@ -859,6 +1153,13 @@ def from_payload( "requested_columns": outputs.requested_columns, "diagnostics": diagnostics, "errors": _structured_error_tuple(payload.get("errors")), + "plausibility": ( + None + if payload.get("plausibility") is None + else ConditionPlausibility.from_payload( + _require_mapping(payload.get("plausibility"), field="plausibility") + ) + ), } if return_mode == "mean": @@ -1082,6 +1383,7 @@ def build_forecast_payload( seed: int | None = None, schema_version: str = SCHEMA_VERSION, query_mode: QueryMode = "forecast", + condition: ConditionBlock | None = None, ) -> dict[str, Any]: """Build a validated JSON-compatible forecast request payload.""" return ForecastRequest( @@ -1098,9 +1400,41 @@ def build_forecast_payload( n_samples=n_samples, quantiles=quantiles, seed=seed, + condition=condition, ).to_payload() +def require_condition_support( + metadata: HealthMetadata, + block: ConditionBlock, +) -> None: + """Refuse a conditioning request the mounted deployment has not advertised. + + The gate exists so a caller learns a deployment cannot condition *before* + paying for the request, rather than from the error on one it already sent. + It reads the advertisement alone, which the deployment derives from its own + head, so it can never claim a capability the checkpoint lacks. + + Raises: + UnsupportedServiceContractError: when the deployment serves no + ``condition`` mode, or none of the condition kinds this block uses. + """ + if "condition" not in metadata.supported_query_modes: + raise UnsupportedServiceContractError( + f"The mounted {metadata.head!r} deployment does not serve condition " + f"queries; it advertises {list(metadata.supported_query_modes)}" + ) + missing = [ + kind for kind in block.kinds if kind not in metadata.supported_condition_kinds + ] + if missing: + raise UnsupportedServiceContractError( + f"The mounted {metadata.head!r} deployment does not serve condition " + f"kinds {missing}; it advertises " + f"{list(metadata.supported_condition_kinds)}" + ) + + def validate_service_metadata( metadata: Mapping[str, Any], *, @@ -1126,10 +1460,21 @@ def validate_service_metadata( f"expected {expected_model_version!r}, got {model_version!r}" ) - _require_exact_values( + # Query modes are per deployment: the service derives them from the mounted + # head, so a head that cannot condition advertises fewer modes than this SDK + # knows. Requiring equality here would reject exactly the deployments the + # advertisement exists to describe. + _require_advertised_subset( metadata, field="supported_query_modes", supported_values=SUPPORTED_QUERY_MODES, + allow_empty=False, + ) + _require_advertised_subset( + metadata, + field="supported_condition_kinds", + supported_values=SUPPORTED_CONDITION_KINDS, + allow_empty=True, ) _require_exact_values( metadata, @@ -1169,6 +1514,48 @@ def _require_decoding_strategy(value: Any) -> DecodingStrategy: return cast(DecodingStrategy, value) +def _require_advertised_subset( + metadata: Mapping[str, Any], + *, + field: str, + supported_values: Sequence[str], + allow_empty: bool, +) -> None: + """Require one advertised capability list to be a subset this SDK understands. + + Used where the deployment legitimately serves fewer values than the SDK + knows, because the service derives the list from its own checkpoint. An + advertised value the SDK does not know is still a contract mismatch: it + means the deployment is newer than this client. + """ + advertised_values = metadata.get(field) + if not isinstance(advertised_values, Sequence) or isinstance( + advertised_values, str | bytes | bytearray + ): + raise UnsupportedServiceContractError( + f"JointFM health metadata field {field!r} must be a JSON array of strings" + ) + + parsed_values: list[str] = [] + for index, advertised_value in enumerate(advertised_values): + if not isinstance(advertised_value, str) or advertised_value == "": + raise UnsupportedServiceContractError( + f"JointFM health metadata field {field!r}[{index}] must be a non-empty string" + ) + parsed_values.append(advertised_value) + + if not parsed_values and not allow_empty: + raise UnsupportedServiceContractError( + f"JointFM health metadata field {field!r} advertises nothing" + ) + unknown = sorted(set(parsed_values) - set(supported_values)) + if unknown: + raise UnsupportedServiceContractError( + f"Unsupported JointFM {field}: this client does not know {unknown!r}; " + f"it supports {sorted(supported_values)!r}" + ) + + def _require_exact_values( metadata: Mapping[str, Any], *, @@ -1366,8 +1753,10 @@ def _structured_error_tuple(value: Any) -> tuple[StructuredError, ...]: return tuple(parsed_errors) -def _string_tuple(value: Any, *, field: str) -> tuple[str, ...]: - values = _require_sequence(value, field=field) +def _string_tuple( + value: Any, *, field: str, allow_empty: bool = False +) -> tuple[str, ...]: + values = _require_sequence(value, field=field, allow_empty=allow_empty) return tuple(_require_string(item, field=f"{field}[]") for item in values) @@ -1419,6 +1808,18 @@ def _forecast_response_expectations( ): requested_columns = tuple(requested_column_values) + condition_value = request_payload.get("condition") + if condition_value is not None and query_times is not None: + condition_block = _require_mapping( + condition_value, field="request_payload.condition" + ) + conditioned_index = condition_block.get("query_time_index") + if isinstance(conditioned_index, int) and not isinstance( + conditioned_index, bool + ): + if 0 <= conditioned_index < len(query_times): + query_times = (query_times[conditioned_index],) + query_mode_value = request_payload.get("query_mode") return_mode_value = request_payload.get("return_mode") quantiles_value = request_payload.get("quantiles") diff --git a/tests/fixtures/condition_interval_response.json b/tests/fixtures/condition_interval_response.json new file mode 100644 index 0000000..e9ab16e --- /dev/null +++ b/tests/fixtures/condition_interval_response.json @@ -0,0 +1,33 @@ +{ + "schema_version": "v3", + "image_version": "0.3.0", + "model_version": "jointfm-inference:0.3.0+ckpt.sdk-test", + "checkpoint_version": "sdk-test", + "head": "gmm", + "query_mode": "condition", + "return_mode": "mean", + "outputs": { + "query_times": [3], + "requested_columns": ["target"], + "mean": [ + [14.25] + ], + "samples": null, + "quantiles": null + }, + "plausibility": { + "equality_log_density": null, + "region_log_probability": -8.91 + }, + "diagnostics": { + "history_rows": 2, + "horizon_count": 1, + "seed": 7, + "condition_draws": 500, + "interval_estimator": { + "points": 16384, + "effective_sample_size": 11453.2 + } + }, + "errors": [] +} diff --git a/tests/fixtures/condition_mean_request.json b/tests/fixtures/condition_mean_request.json new file mode 100644 index 0000000..ba7d9f4 --- /dev/null +++ b/tests/fixtures/condition_mean_request.json @@ -0,0 +1,41 @@ +{ + "schema_version": "v3", + "model_version": "jointfm-inference:0.3.0+ckpt.sdk-test", + "query_mode": "condition", + "return_mode": "mean", + "time_index_mode": "ordinal", + "columns": [ + { + "name": "driver", + "modality": "numeric" + }, + { + "name": "target", + "modality": "numeric", + "role": "target" + } + ], + "history_rows": [ + { + "driver": 1.0, + "target": 10.0 + }, + { + "driver": 1.1, + "target": 11.0 + } + ], + "query_times": [2, 3], + "requested_columns": ["target"], + "condition": { + "query_time_index": 1, + "conditions": [ + { + "column": "driver", + "kind": "equality", + "value": 1.5 + } + ] + }, + "seed": 7 +} diff --git a/tests/fixtures/condition_mean_response.json b/tests/fixtures/condition_mean_response.json new file mode 100644 index 0000000..622745b --- /dev/null +++ b/tests/fixtures/condition_mean_response.json @@ -0,0 +1,30 @@ +{ + "schema_version": "v3", + "image_version": "0.3.0", + "model_version": "jointfm-inference:0.3.0+ckpt.sdk-test", + "checkpoint_version": "sdk-test", + "head": "gmm", + "query_mode": "condition", + "return_mode": "mean", + "outputs": { + "query_times": [3], + "requested_columns": ["target"], + "mean": [ + [13.5] + ], + "samples": null, + "quantiles": null + }, + "plausibility": { + "equality_log_density": -1.27, + "region_log_probability": null + }, + "diagnostics": { + "history_rows": 2, + "horizon_count": 1, + "seed": 7, + "condition_draws": 500, + "interval_estimator": null + }, + "errors": [] +} diff --git a/tests/fixtures/forecast_mean_request.json b/tests/fixtures/forecast_mean_request.json index a824605..c4b7ab7 100644 --- a/tests/fixtures/forecast_mean_request.json +++ b/tests/fixtures/forecast_mean_request.json @@ -1,5 +1,5 @@ { - "schema_version": "v2", + "schema_version": "v3", "model_version": "jointfm-inference:0.3.0+ckpt.sdk-test", "query_mode": "forecast", "return_mode": "mean", diff --git a/tests/fixtures/forecast_mean_response.json b/tests/fixtures/forecast_mean_response.json index 299aba9..d6aa094 100644 --- a/tests/fixtures/forecast_mean_response.json +++ b/tests/fixtures/forecast_mean_response.json @@ -1,5 +1,5 @@ { - "schema_version": "v2", + "schema_version": "v3", "image_version": "0.3.0", "model_version": "jointfm-inference:0.3.0+ckpt.sdk-test", "checkpoint_version": "sdk-test", diff --git a/tests/fixtures/forecast_quantiles_request.json b/tests/fixtures/forecast_quantiles_request.json index 9a2f4c5..602ddc4 100644 --- a/tests/fixtures/forecast_quantiles_request.json +++ b/tests/fixtures/forecast_quantiles_request.json @@ -1,5 +1,5 @@ { - "schema_version": "v2", + "schema_version": "v3", "model_version": "jointfm-inference:0.3.0+ckpt.sdk-test", "query_mode": "forecast", "return_mode": "quantiles", diff --git a/tests/fixtures/forecast_quantiles_response.json b/tests/fixtures/forecast_quantiles_response.json index ac6b06c..aa8e9bd 100644 --- a/tests/fixtures/forecast_quantiles_response.json +++ b/tests/fixtures/forecast_quantiles_response.json @@ -1,5 +1,5 @@ { - "schema_version": "v2", + "schema_version": "v3", "image_version": "0.3.0", "model_version": "jointfm-inference:0.3.0+ckpt.sdk-test", "checkpoint_version": "sdk-test", diff --git a/tests/fixtures/forecast_samples_request.json b/tests/fixtures/forecast_samples_request.json index 9fe3e84..0f161c2 100644 --- a/tests/fixtures/forecast_samples_request.json +++ b/tests/fixtures/forecast_samples_request.json @@ -1,5 +1,5 @@ { - "schema_version": "v2", + "schema_version": "v3", "model_version": "jointfm-inference:0.3.0+ckpt.sdk-test", "query_mode": "forecast", "return_mode": "samples", diff --git a/tests/fixtures/forecast_samples_response.json b/tests/fixtures/forecast_samples_response.json index d77b3df..60e60e1 100644 --- a/tests/fixtures/forecast_samples_response.json +++ b/tests/fixtures/forecast_samples_response.json @@ -1,5 +1,5 @@ { - "schema_version": "v2", + "schema_version": "v3", "image_version": "0.3.0", "model_version": "jointfm-inference:0.3.0+ckpt.sdk-test", "checkpoint_version": "sdk-test", diff --git a/tests/fixtures/health_metadata.json b/tests/fixtures/health_metadata.json index dbaf71d..08d9aeb 100644 --- a/tests/fixtures/health_metadata.json +++ b/tests/fixtures/health_metadata.json @@ -1,6 +1,6 @@ { "status": "ok", - "schema_version": "v2", + "schema_version": "v3", "image_version": "0.3.0", "model_version": "jointfm-inference:0.3.0+ckpt.sdk-test", "checkpoint_version": "sdk-test", @@ -9,6 +9,7 @@ "head": "studentt", "decoding_strategy": "parallel_dense", "supported_query_modes": ["forecast"], + "supported_condition_kinds": [], "supported_return_modes": ["mean", "samples", "quantiles", "log_prob"], "supported_time_index_modes": [ "ordinal", diff --git a/tests/fixtures/input_size_exceeded_response.json b/tests/fixtures/input_size_exceeded_response.json index 30a1d5a..e244f5c 100644 --- a/tests/fixtures/input_size_exceeded_response.json +++ b/tests/fixtures/input_size_exceeded_response.json @@ -1,5 +1,5 @@ { - "schema_version": "v2", + "schema_version": "v3", "errors": [ { "code": "INPUT_SIZE_EXCEEDED", diff --git a/tests/fixtures/model_version_mismatch_response.json b/tests/fixtures/model_version_mismatch_response.json index 0b293e5..b5fa39e 100644 --- a/tests/fixtures/model_version_mismatch_response.json +++ b/tests/fixtures/model_version_mismatch_response.json @@ -1,5 +1,5 @@ { - "schema_version": "v2", + "schema_version": "v3", "errors": [ { "code": "MODEL_VERSION_MISMATCH", diff --git a/tests/fixtures/schema_version_mismatch_response.json b/tests/fixtures/schema_version_mismatch_response.json index bf3f34b..ffb299b 100644 --- a/tests/fixtures/schema_version_mismatch_response.json +++ b/tests/fixtures/schema_version_mismatch_response.json @@ -3,7 +3,7 @@ "errors": [ { "code": "SCHEMA_VERSION_MISMATCH", - "message": "Unsupported schema_version: expected 'v2', got 'v3'", + "message": "Unsupported schema_version: expected 'v3', got 'v2'", "field": "schema_version", "detail": {} } diff --git a/tests/fixtures/validation_error_response.json b/tests/fixtures/validation_error_response.json index 126d537..7f12104 100644 --- a/tests/fixtures/validation_error_response.json +++ b/tests/fixtures/validation_error_response.json @@ -1,5 +1,5 @@ { - "schema_version": "v2", + "schema_version": "v3", "errors": [ { "code": "VALIDATION_ERROR", diff --git a/tests/test_cli.py b/tests/test_cli.py index 3e3be99..440447a 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -48,7 +48,7 @@ class FakeHealthClient: "deployment-id/predictionsUnstructured" ), deployment_selector="deployment_id", - schema_version="v2", + schema_version="v3", model_version="jointfm-inference:0.3.0+ckpt.sdk-test", deployment_id="deployment-id", ) @@ -58,7 +58,7 @@ def health(self, *, cache: bool = False, refresh: bool = False) -> HealthMetadat del cache, refresh return HealthMetadata( status="ok", - schema_version="v2", + schema_version="v3", image_version="0.3.0", model_version="jointfm-inference:0.3.0+ckpt.sdk-test", checkpoint_version="sdk-test", @@ -67,6 +67,7 @@ def health(self, *, cache: bool = False, refresh: bool = False) -> HealthMetadat head="studentt", decoding_strategy="parallel_dense", supported_query_modes=("forecast",), + supported_condition_kinds=(), supported_return_modes=("mean", "samples", "quantiles", "log_prob"), supported_time_index_modes=( "ordinal", @@ -171,7 +172,7 @@ def test_predict_command_writes_response_file(monkeypatch, tmp_path: Path) -> No request_file.write_text( json.dumps( { - "schema_version": "v2", + "schema_version": "v3", "model_version": "jointfm-inference:0.3.0+ckpt.sdk-test", } ), diff --git a/tests/test_condition_mode.py b/tests/test_condition_mode.py new file mode 100644 index 0000000..5995393 --- /dev/null +++ b/tests/test_condition_mode.py @@ -0,0 +1,353 @@ +# Copyright 2026 DataRobot, Inc. and its affiliates. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests for the ``condition`` query mode on the client side. + +Three surfaces are covered, and the split matters. The condition objects +validate what is wrong independently of any deployment, so a caller learns it +without a round trip. The capability gate refuses what *this* deployment has +not advertised, which is the only way to know before paying for a request. +The response parser reads back the two numbers the service reports and never +refuses on, so acting on them stays the caller's decision. +""" + +from __future__ import annotations + +from collections.abc import Callable +from typing import Any + +import pytest + +from jointfm_client import ( + ColumnSpec, + ConditionBlock, + ConditionPlausibility, + DataFrameSchema, + EqualityCondition, + ForecastRequest, + ForecastRequestMetadata, + ForecastResponse, + HealthMetadata, + IntervalCondition, + build_forecast_payload, + require_condition_support, +) +from jointfm_client.contract import QueryMode +from jointfm_client.exceptions import UnsupportedServiceContractError + +_MODEL_VERSION = "jointfm-inference:0.3.0+ckpt.sdk-test" + + +def _schema() -> DataFrameSchema: + """Build one three-column ordinal schema for request-level tests.""" + return DataFrameSchema( + columns=( + ColumnSpec(name="driver", modality="numeric"), + ColumnSpec(name="hedge", modality="numeric"), + ColumnSpec(name="target", modality="numeric", role="target"), + ), + time_index_mode="ordinal", + ) + + +def _request( + block: ConditionBlock | None, + *, + requested_columns: tuple[str, ...] | None = ("target",), + query_mode: QueryMode = "condition", +) -> ForecastRequest: + """Build one condition request against the two-column schema.""" + return ForecastRequest( + metadata=ForecastRequestMetadata( + model_version=_MODEL_VERSION, + query_mode=query_mode, + ), + schema=_schema(), + history_rows=( + {"driver": 1.0, "hedge": 0.5, "target": 10.0}, + {"driver": 1.1, "hedge": 0.6, "target": 11.0}, + ), + query_times=(2, 3), + requested_columns=requested_columns, + condition=block, + ) + + +def _health( + *, + query_modes: tuple[str, ...] = ("forecast", "condition"), + condition_kinds: tuple[str, ...] = ("equality", "interval"), +) -> HealthMetadata: + """Build one advertisement without going through a transport.""" + return HealthMetadata( + status="ok", + schema_version="v3", + image_version="0.3.0", + model_version=_MODEL_VERSION, + checkpoint_version="sdk-test", + checkpoint_path="/models/jointfm.pt", + device="cpu", + head="gmm", + decoding_strategy="parallel_dense", + supported_query_modes=query_modes, + supported_condition_kinds=condition_kinds, + supported_return_modes=("mean", "samples", "quantiles", "log_prob"), + supported_time_index_modes=("ordinal",), + time_index_encoding="ordinal", + max_sample_count=10000, + ) + + +def test_an_equality_condition_rejects_a_value_it_cannot_pin() -> None: + """A pin must be a finite number; anything else is not an event.""" + with pytest.raises(ValueError, match="finite"): + EqualityCondition(column="driver", value=float("inf")) + + +@pytest.mark.parametrize( + ("lower", "upper", "message"), + [ + (None, None, "conditions nothing"), + (2.0, 1.0, "lower < upper"), + (float("nan"), 1.0, "finite"), + ], +) +def test_an_interval_condition_rejects_a_range_that_bounds_nothing( + lower: float | None, upper: float | None, message: str +) -> None: + """An interval open on both sides, or backwards, describes no region.""" + with pytest.raises(ValueError, match=message): + IntervalCondition(column="driver", lower=lower, upper=upper) + + +def test_a_block_rejects_two_conditions_on_one_column() -> None: + """A column carries at most one condition, of either kind.""" + with pytest.raises(ValueError, match="more than one condition"): + ConditionBlock( + query_time_index=0, + conditions=( + EqualityCondition(column="driver", value=1.0), + IntervalCondition(column="driver", lower=0.0, upper=1.0), + ), + ) + + +def test_a_block_reports_the_kinds_a_deployment_must_advertise() -> None: + """The gate needs the kinds, and each kind appears once however many columns use it.""" + block = ConditionBlock( + query_time_index=0, + conditions=( + EqualityCondition(column="driver", value=1.0), + IntervalCondition(column="target", lower=0.0, upper=None), + ), + ) + + assert block.kinds == ("equality", "interval") + assert block.pinned_columns == ("driver",) + assert block.conditioned_columns == ("driver", "target") + + +def test_the_mode_and_the_block_must_agree() -> None: + """One without the other means a different request than the caller wrote.""" + block = ConditionBlock( + query_time_index=0, + conditions=(EqualityCondition(column="driver", value=1.0),), + ) + + with pytest.raises(ValueError, match="go together"): + _request(block, query_mode="forecast") + with pytest.raises(ValueError, match="go together"): + _request(None) + + +@pytest.mark.parametrize( + ("block", "requested_columns", "message"), + [ + ( + ConditionBlock( + query_time_index=0, + conditions=(EqualityCondition(column="absent", value=1.0),), + ), + ("target",), + "undeclared columns", + ), + ( + ConditionBlock( + query_time_index=9, + conditions=(EqualityCondition(column="driver", value=1.0),), + ), + ("target",), + "outside the 2 requested future positions", + ), + ( + ConditionBlock( + query_time_index=0, + conditions=(EqualityCondition(column="driver", value=1.0),), + ), + ("driver", "target"), + "pinned columns", + ), + ( + ConditionBlock( + query_time_index=0, + conditions=( + EqualityCondition(column="driver", value=1.0), + EqualityCondition(column="hedge", value=0.5), + EqualityCondition(column="target", value=2.0), + ), + ), + None, + "log_prob", + ), + ( + ConditionBlock( + query_time_index=0, + conditions=( + EqualityCondition(column="driver", value=1.0), + EqualityCondition(column="hedge", value=0.5), + IntervalCondition(column="target", lower=0.0, upper=None), + ), + ), + None, + "at least one column unconditioned", + ), + ], +) +def test_a_request_is_checked_against_its_own_schema_before_any_round_trip( + block: ConditionBlock, requested_columns: tuple[str, ...] | None, message: str +) -> None: + """The caller learns which column name it got wrong without paying for a request.""" + with pytest.raises(ValueError, match=message): + _request(block, requested_columns=requested_columns) + + +def test_an_interval_column_may_still_be_read_out() -> None: + """An interval fixes a region, so the column's distribution inside it is an answer.""" + block = ConditionBlock( + query_time_index=0, + conditions=(IntervalCondition(column="driver", lower=0.5, upper=1.5),), + ) + + request = _request(block, requested_columns=("driver", "target")) + + assert request.to_payload()["requested_columns"] == ["driver", "target"] + + +def test_the_payload_carries_the_block_the_service_parses() -> None: + """The wire form names its position once and each condition names its kind.""" + payload = build_forecast_payload( + model_version=_MODEL_VERSION, + schema=_schema(), + history_rows=({"driver": 1.0, "hedge": 0.5, "target": 10.0},), + query_times=(2, 3), + requested_columns=("target",), + query_mode="condition", + condition=ConditionBlock( + query_time_index=1, + conditions=( + EqualityCondition(column="driver", value=1.5), + IntervalCondition(column="target", lower=None, upper=12.0), + ), + ), + ) + + assert payload["query_mode"] == "condition" + assert payload["condition"] == { + "query_time_index": 1, + "conditions": [ + {"column": "driver", "kind": "equality", "value": 1.5}, + {"column": "target", "kind": "interval", "lower": None, "upper": 12.0}, + ], + } + + +def test_the_gate_refuses_a_deployment_that_advertises_no_condition_mode() -> None: + """Discovery before the request is the point: the error must not cost a round trip.""" + block = ConditionBlock( + query_time_index=0, + conditions=(EqualityCondition(column="driver", value=1.0),), + ) + + with pytest.raises( + UnsupportedServiceContractError, match="does not serve condition" + ): + require_condition_support( + _health(query_modes=("forecast",), condition_kinds=()), block + ) + + +def test_the_gate_refuses_a_kind_the_deployment_does_not_answer() -> None: + """A deployment may serve one kind before the other; the block says which it needs.""" + block = ConditionBlock( + query_time_index=0, + conditions=(IntervalCondition(column="driver", lower=0.0, upper=1.0),), + ) + + with pytest.raises(UnsupportedServiceContractError, match=r"kinds \['interval'\]"): + require_condition_support(_health(condition_kinds=("equality",)), block) + + require_condition_support(_health(), block) + + +def test_the_equality_response_reads_back_its_plausibility( + json_fixture_loader: Callable[[str], dict[str, Any]], +) -> None: + """The density of the pinned values is what separates confident from absurd.""" + request_payload = json_fixture_loader("condition_mean_request") + response = ForecastResponse.from_payload( + json_fixture_loader("condition_mean_response"), + request_payload=request_payload, + ) + + assert response.query_mode == "condition" + assert response.plausibility == ConditionPlausibility( + equality_log_density=-1.27, + region_log_probability=None, + ) + assert response.diagnostics.condition_draws == 500 + assert response.diagnostics.interval_estimator is None + + +def test_the_response_describes_the_conditioned_position_alone( + json_fixture_loader: Callable[[str], dict[str, Any]], +) -> None: + """The request asked for two future rows; a condition answers about one.""" + request_payload = json_fixture_loader("condition_mean_request") + assert request_payload["query_times"] == [2, 3] + + response = ForecastResponse.from_payload( + json_fixture_loader("condition_mean_response"), + request_payload=request_payload, + ) + + assert response.query_times == (3,) + assert response.diagnostics.horizon_count == 1 + + +def test_an_interval_response_reads_back_its_estimator_accuracy( + json_fixture_loader: Callable[[str], dict[str, Any]], +) -> None: + """A multi-column region is estimated, and the caller judges the estimate itself.""" + response = ForecastResponse.from_payload( + json_fixture_loader("condition_interval_response"), + request_payload=json_fixture_loader("condition_mean_request"), + ) + + assert response.plausibility is not None + assert response.plausibility.equality_log_density is None + assert response.plausibility.region_log_probability == pytest.approx(-8.91) + estimator = response.diagnostics.interval_estimator + assert estimator is not None + assert estimator.points == 16384 + assert estimator.effective_sample_size == pytest.approx(11453.2) diff --git a/tests/test_configuration.py b/tests/test_configuration.py index 3445f50..5ee8cab 100644 --- a/tests/test_configuration.py +++ b/tests/test_configuration.py @@ -47,7 +47,7 @@ class _HealthTransport: _METADATA: dict[str, object] = { "status": "ok", - "schema_version": "v2", + "schema_version": "v3", "image_version": "0.3.0", "model_version": "jointfm-inference:0.3.0+ckpt.yaml", "checkpoint_version": "yaml", @@ -56,6 +56,7 @@ class _HealthTransport: "head": "studentt", "decoding_strategy": "parallel_dense", "supported_query_modes": ["forecast"], + "supported_condition_kinds": [], "supported_return_modes": ["mean", "samples", "quantiles", "log_prob"], "supported_time_index_modes": [ "ordinal", @@ -146,7 +147,7 @@ def test_load_settings_layers_config_below_dotenv_and_environment( "datarobot_endpoint": "https://app.datarobot.com/api/v2", "datarobot_api_token": "yaml-token", "deployment_id": "yaml-deployment-id", - "schema_version": "v2", + "schema_version": "v3", "model_version": "jointfm-inference:0.3.0+ckpt.yaml", } }, @@ -190,7 +191,7 @@ def test_client_from_env_uses_transport_defaults_from_config( "datarobot_endpoint": "https://app.datarobot.com/api/v2", "datarobot_api_token": "yaml-token", "deployment_id": "yaml-deployment-id", - "schema_version": "v2", + "schema_version": "v3", "model_version": "jointfm-inference:0.3.0+ckpt.yaml", }, "transport": { diff --git a/tests/test_contract.py b/tests/test_contract.py index 410b0bd..90af77b 100644 --- a/tests/test_contract.py +++ b/tests/test_contract.py @@ -48,6 +48,7 @@ SUPPORTED_COLUMN_ROLES, STRUCTURED_ERROR_CODES, SUPPORTED_DECODING_STRATEGIES, + SUPPORTED_CONDITION_KINDS, SUPPORTED_QUERY_MODES, SUPPORTED_RETURN_MODES, SUPPORTED_TIME_INDEX_MODES, @@ -75,7 +76,7 @@ def _health_metadata() -> dict[str, object]: """Health metadata.""" return { "status": "ok", - "schema_version": "v2", + "schema_version": "v3", "image_version": "0.3.0", "model_version": "jointfm-inference:0.3.0+ckpt.smoke-1", "checkpoint_version": "smoke-1", @@ -84,6 +85,7 @@ def _health_metadata() -> dict[str, object]: "head": "dummy", "decoding_strategy": "parallel_dense", "supported_query_modes": ["forecast"], + "supported_condition_kinds": [], "supported_return_modes": ["mean", "quantiles", "samples", "log_prob"], "supported_time_index_modes": [ "absolute_datetime", @@ -109,7 +111,7 @@ def test_package_identity_contract() -> None: assert DISTRIBUTION_NAME == "jointfm-client" assert IMPORT_NAMESPACE == "jointfm_client" assert FIRST_SUPPORTED_PYTHON_VERSION == "3.11" - assert SCHEMA_VERSION == "v2" + assert SCHEMA_VERSION == "v3" def test_package_version_matches_installed_distribution() -> None: @@ -131,7 +133,8 @@ def test_service_route_contract() -> None: def test_mode_and_error_contract() -> None: """Mode and error contract.""" - assert SUPPORTED_QUERY_MODES == ("forecast",) + assert SUPPORTED_QUERY_MODES == ("forecast", "condition") + assert SUPPORTED_CONDITION_KINDS == ("equality", "interval") assert SUPPORTED_RETURN_MODES == ("mean", "samples", "quantiles", "log_prob") assert SUPPORTED_TIME_INDEX_MODES == ( "ordinal", @@ -264,7 +267,7 @@ def test_forecast_payload_matches_service_contract_without_mutating_inputs() -> ) assert payload == { - "schema_version": "v2", + "schema_version": "v3", "model_version": "jointfm-inference:0.3.0+ckpt.smoke-1", "query_mode": "forecast", "return_mode": "quantiles", @@ -359,7 +362,7 @@ def test_dataframe_payload_matches_service_forecast_request_shape() -> None: ) assert payload == { - "schema_version": "v2", + "schema_version": "v3", "model_version": "jointfm-inference:0.3.0+ckpt.smoke-1", "query_mode": "forecast", "return_mode": "mean", @@ -774,7 +777,7 @@ def test_health_and_response_models_parse_current_payloads() -> None: health = HealthMetadata.from_payload(_health_metadata()) response = ForecastResponse.from_payload( { - "schema_version": "v2", + "schema_version": "v3", "image_version": "0.3.0", "model_version": "jointfm-inference:0.3.0+ckpt.smoke-1", "checkpoint_version": "smoke-1", @@ -837,7 +840,7 @@ def test_forecast_result_conversion_helpers_cover_mean_samples_and_quantiles() - """Forecast result conversion helpers cover mean samples and quantiles.""" mean_result = ForecastResponse.from_payload( { - "schema_version": "v2", + "schema_version": "v3", "image_version": "0.3.0", "model_version": "jointfm-inference:0.3.0+ckpt.smoke-1", "checkpoint_version": "smoke-1", @@ -857,7 +860,7 @@ def test_forecast_result_conversion_helpers_cover_mean_samples_and_quantiles() - ) sample_result = ForecastResponse.from_payload( { - "schema_version": "v2", + "schema_version": "v3", "image_version": "0.3.0", "model_version": "jointfm-inference:0.3.0+ckpt.smoke-1", "checkpoint_version": "smoke-1", @@ -880,7 +883,7 @@ def test_forecast_result_conversion_helpers_cover_mean_samples_and_quantiles() - ) quantile_result = ForecastResponse.from_payload( { - "schema_version": "v2", + "schema_version": "v3", "image_version": "0.3.0", "model_version": "jointfm-inference:0.3.0+ckpt.smoke-1", "checkpoint_version": "smoke-1", @@ -954,7 +957,7 @@ def test_forecast_result_conversion_helpers_cover_mean_samples_and_quantiles() - def test_forecast_response_validates_request_scoped_shapes() -> None: """Forecast response validates request scoped shapes.""" request_payload = { - "schema_version": "v2", + "schema_version": "v3", "model_version": "jointfm-inference:0.3.0+ckpt.smoke-1", "query_mode": "forecast", "return_mode": "samples", @@ -963,7 +966,7 @@ def test_forecast_response_validates_request_scoped_shapes() -> None: "n_samples": 3, } response_payload = { - "schema_version": "v2", + "schema_version": "v3", "image_version": "0.3.0", "model_version": "jointfm-inference:0.3.0+ckpt.smoke-1", "checkpoint_version": "smoke-1", @@ -993,7 +996,7 @@ def test_forecast_response_raises_typed_error_for_success_payload_errors() -> No with pytest.raises(JointFMServiceError) as exc_info: ForecastResponse.from_payload( { - "schema_version": "v2", + "schema_version": "v3", "image_version": "0.3.0", "model_version": "jointfm-inference:0.3.0+ckpt.smoke-1", "checkpoint_version": "smoke-1", @@ -1030,7 +1033,7 @@ def test_validate_service_metadata_accepts_current_v1_contract() -> None: def test_validate_service_metadata_rejects_schema_mismatch() -> None: """Validate service metadata rejects schema mismatch.""" metadata = _health_metadata() - metadata["schema_version"] = "v3" + metadata["schema_version"] = "v2" with pytest.raises(UnsupportedSchemaVersionError, match="schema_version"): validate_service_metadata(metadata) @@ -1159,7 +1162,7 @@ def post_json(self, url: str, payload: Mapping[str, Any]) -> Mapping[str, Any]: def _forecast_response_payload(*, return_mode: str) -> dict[str, object]: """Forecast response payload.""" return { - "schema_version": "v2", + "schema_version": "v3", "image_version": "0.3.0", "model_version": "jointfm-inference:0.3.0+ckpt.smoke-1", "checkpoint_version": "smoke-1", diff --git a/tests/test_contract_models.py b/tests/test_contract_models.py index 5b40b29..9d5186b 100644 --- a/tests/test_contract_models.py +++ b/tests/test_contract_models.py @@ -76,7 +76,7 @@ def test_request_models_serialize_direct_payloads_without_mutating_inputs() -> N ) assert metadata.to_payload() == { - "schema_version": "v2", + "schema_version": "v3", "model_version": "jointfm-inference:0.3.0+ckpt.smoke-1", "query_mode": "forecast", "return_mode": "quantiles", @@ -91,7 +91,7 @@ def test_request_models_serialize_direct_payloads_without_mutating_inputs() -> N "timezone": "UTC", } assert request.to_payload() == { - "schema_version": "v2", + "schema_version": "v3", "model_version": "jointfm-inference:0.3.0+ckpt.smoke-1", "query_mode": "forecast", "return_mode": "quantiles", @@ -131,7 +131,7 @@ def test_request_models_reject_direct_validation_edges() -> None: with pytest.raises(ValueError, match="schema_version"): ForecastRequestMetadata( model_version="jointfm-inference:0.3.0+ckpt.smoke-1", - schema_version="v3", + schema_version="v2", ) with pytest.raises(ValueError, match="query_mode"): @@ -303,7 +303,7 @@ def test_response_models_reject_direct_validation_edges() -> None: def test_forecast_response_rejects_request_scoped_metadata_mismatches() -> None: """Forecast response rejects request scoped metadata mismatches.""" request_payload = { - "schema_version": "v2", + "schema_version": "v3", "model_version": "jointfm-inference:0.3.0+ckpt.smoke-1", "query_mode": "forecast", "return_mode": "mean", @@ -354,7 +354,7 @@ def test_forecast_response_rejects_request_scoped_metadata_mismatches() -> None: def test_forecast_response_rejects_sample_bound_violations() -> None: """Forecast response rejects sample bound violations.""" request_payload = { - "schema_version": "v2", + "schema_version": "v3", "model_version": "jointfm-inference:0.3.0+ckpt.smoke-1", "query_mode": "forecast", "return_mode": "samples", @@ -397,7 +397,7 @@ def test_forecast_response_rejects_sample_bound_violations() -> None: def _mean_response_payload() -> dict[str, Any]: """Mean response payload.""" return { - "schema_version": "v2", + "schema_version": "v3", "image_version": "0.3.0", "model_version": "jointfm-inference:0.3.0+ckpt.smoke-1", "checkpoint_version": "smoke-1", @@ -419,7 +419,7 @@ def _mean_response_payload() -> dict[str, Any]: def _sample_response_payload() -> dict[str, Any]: """Sample response payload.""" return { - "schema_version": "v2", + "schema_version": "v3", "image_version": "0.3.0", "model_version": "jointfm-inference:0.3.0+ckpt.smoke-1", "checkpoint_version": "smoke-1", diff --git a/tests/test_feature_importance.py b/tests/test_feature_importance.py index 8c3dcf9..3d6fe0a 100644 --- a/tests/test_feature_importance.py +++ b/tests/test_feature_importance.py @@ -47,7 +47,7 @@ def post_json(self, url: str, payload: Mapping[str, Any]) -> Mapping[str, Any]: self.payloads.append(dict(payload)) samples = self.sample_batches[len(self.payloads) - 1] return { - "schema_version": "v2", + "schema_version": "v3", "image_version": "0.3.0", "model_version": _MODEL_VERSION, "checkpoint_version": "sdk-test", diff --git a/tests/test_pool.py b/tests/test_pool.py index 73fc2fd..837f19e 100644 --- a/tests/test_pool.py +++ b/tests/test_pool.py @@ -56,7 +56,7 @@ def _health( """Health.""" return { "status": "ok", - "schema_version": "v2", + "schema_version": "v3", "image_version": "0.3.0", "model_version": model_version, "checkpoint_version": checkpoint_version, @@ -65,6 +65,7 @@ def _health( "head": "studentt", "decoding_strategy": "parallel_dense", "supported_query_modes": ["forecast"], + "supported_condition_kinds": [], "supported_return_modes": ["mean", "samples", "quantiles", "log_prob"], "supported_time_index_modes": [ "ordinal", @@ -133,7 +134,7 @@ def test_pool_retries_next_instance_on_470() -> None: pool = _pool(transport=_Transport(fail_ids=frozenset({"a"}))) assert pool.next_instance().deployment_id == "a" assert pool.next_instance().deployment_id == "b" - assert pool.post_json({"schema_version": "v2"}) == { + assert pool.post_json({"schema_version": "v3"}) == { "ok": True, "deployment_id": "b", } @@ -143,7 +144,7 @@ def test_pool_raises_when_all_instances_unavailable() -> None: """Pool raises when all instances unavailable.""" pool = _pool(transport=_Transport(fail_ids=frozenset({"a", "b"}))) with pytest.raises(JointFMHTTPStatusError, match="unavailable"): - pool.post_json({"schema_version": "v2"}) + pool.post_json({"schema_version": "v3"}) def test_pool_health_rejects_mismatch_and_aligns_sample_cap() -> None: @@ -274,7 +275,7 @@ def post_json(self, url: str, payload: Mapping[str, Any]) -> Mapping[str, Any]: executor.submit( pool.post_json_to, pool.instance_at(index), - {"schema_version": "v2"}, + {"schema_version": "v3"}, ) for index in range(2) ] @@ -298,7 +299,7 @@ def test_pool_health_routes_only_reachable_peers() -> None: assert pool.instance_at(0).deployment_id == "b" assert pool.instance_at(1).deployment_id == "b" assert pool.next_instance().deployment_id == "b" - assert pool.post_json({"schema_version": "v2"})["deployment_id"] == "b" + assert pool.post_json({"schema_version": "v3"})["deployment_id"] == "b" def test_pool_failover_retries_health_excluded_peer() -> None: @@ -312,7 +313,7 @@ def test_pool_failover_retries_health_excluded_peer() -> None: assert pool.instance_at(0).deployment_id == "b" transport.fail_ids = frozenset({"b"}) - assert pool.post_json({"schema_version": "v2"})["deployment_id"] == "a" + assert pool.post_json({"schema_version": "v3"})["deployment_id"] == "a" assert pool.instance_at(0).deployment_id == "a" @@ -332,7 +333,7 @@ def test_pool_health_skips_incompatible_peer_when_another_matches_pin() -> None: assert metadata.model_version == pinned assert pool.instance_at(0).deployment_id == "a" assert pool.instance_at(1).deployment_id == "a" - assert pool.post_json({"schema_version": "v2"})["deployment_id"] == "a" + assert pool.post_json({"schema_version": "v3"})["deployment_id"] == "a" def test_pool_cooldown_restores_peer_after_transient_failure( @@ -344,7 +345,7 @@ def test_pool_cooldown_restores_peer_after_transient_failure( transport = _Transport(fail_ids=frozenset({"a"})) pool = _pool(transport=transport, peer_cooldown_seconds=10.0) - assert pool.post_json({"schema_version": "v2"})["deployment_id"] == "b" + assert pool.post_json({"schema_version": "v3"})["deployment_id"] == "b" assert pool.instance_at(0).deployment_id == "b" assert pool.instance_at(1).deployment_id == "b" @@ -354,4 +355,4 @@ def test_pool_cooldown_restores_peer_after_transient_failure( clock["now"] = 110.0 assert {pool.instance_at(i).deployment_id for i in range(2)} == {"a", "b"} - assert pool.post_json({"schema_version": "v2"})["deployment_id"] in {"a", "b"} + assert pool.post_json({"schema_version": "v3"})["deployment_id"] in {"a", "b"} diff --git a/tests/test_settings.py b/tests/test_settings.py index fcf49e3..22c74f3 100644 --- a/tests/test_settings.py +++ b/tests/test_settings.py @@ -54,7 +54,7 @@ def _hosted_env(**overrides: str) -> dict[str, str]: DATAROBOT_ENDPOINT_ENV: "https://app.datarobot.com/api/v2/", DATAROBOT_API_TOKEN_ENV: "secret-token", JOINTFM_DEPLOYMENT_ID_ENV: "deployment-id", - JOINTFM_SCHEMA_VERSION_ENV: "v2", + JOINTFM_SCHEMA_VERSION_ENV: "v3", JOINTFM_MODEL_VERSION_ENV: "jointfm-inference:0.3.0+ckpt.sdk-test", } env.update(overrides) @@ -66,7 +66,7 @@ def test_load_settings_from_environment_with_deployment_id_builds_hosted_url() - settings = load_settings(env=_hosted_env(), dotenv_path=None) assert settings.datarobot_endpoint == "https://app.datarobot.com/api/v2" - assert settings.schema_version == "v2" + assert settings.schema_version == "v3" assert settings.model_version == "jointfm-inference:0.3.0+ckpt.sdk-test" assert settings.deployment_id == "deployment-id" assert settings.predict_url == ( @@ -82,7 +82,7 @@ def test_load_settings_with_local_service_base_url_builds_direct_urls() -> None: settings = load_settings( env={ JOINTFM_LOCAL_BASE_URL_ENV: "http://127.0.0.1:8080/", - JOINTFM_SCHEMA_VERSION_ENV: "v2", + JOINTFM_SCHEMA_VERSION_ENV: "v3", JOINTFM_MODEL_VERSION_ENV: "jointfm-inference:0.3.0+ckpt.local-test", }, dotenv_path=None, @@ -158,7 +158,7 @@ def test_load_settings_reads_dotenv_without_overriding_environment(tmp_path) -> "DATAROBOT_ENDPOINT=https://app.datarobot.com/api/v2", "DATAROBOT_API_TOKEN=file-token", "JOINTFM_DEPLOYMENT_ID=file-deployment-id", - "JOINTFM_SCHEMA_VERSION=v2", + "JOINTFM_SCHEMA_VERSION=v3", "JOINTFM_MODEL_VERSION=jointfm-inference:0.3.0+ckpt.sdk-test", ] ), @@ -197,7 +197,7 @@ def test_load_settings_rejects_unsupported_schema_version() -> None: """Load settings rejects unsupported schema version.""" with pytest.raises(JointFMConfigurationError, match=JOINTFM_SCHEMA_VERSION_ENV): load_settings( - env=_hosted_env(**{JOINTFM_SCHEMA_VERSION_ENV: "v3"}), + env=_hosted_env(**{JOINTFM_SCHEMA_VERSION_ENV: "v2"}), dotenv_path=None, ) @@ -209,7 +209,7 @@ def test_load_settings_rejects_missing_credentials_without_defaults() -> None: env={ DATAROBOT_API_TOKEN_ENV: "secret-token", JOINTFM_DEPLOYMENT_ID_ENV: "deployment-id", - JOINTFM_SCHEMA_VERSION_ENV: "v2", + JOINTFM_SCHEMA_VERSION_ENV: "v3", JOINTFM_MODEL_VERSION_ENV: "jointfm-inference:0.3.0+ckpt.sdk-test", }, dotenv_path=None, @@ -220,7 +220,7 @@ def test_load_settings_rejects_missing_credentials_without_defaults() -> None: env={ DATAROBOT_ENDPOINT_ENV: "https://app.datarobot.com/api/v2", JOINTFM_DEPLOYMENT_ID_ENV: "deployment-id", - JOINTFM_SCHEMA_VERSION_ENV: "v2", + JOINTFM_SCHEMA_VERSION_ENV: "v3", JOINTFM_MODEL_VERSION_ENV: "jointfm-inference:0.3.0+ckpt.sdk-test", }, dotenv_path=None, diff --git a/tests/test_surfaces.py b/tests/test_surfaces.py index 032d840..1d73b3d 100644 --- a/tests/test_surfaces.py +++ b/tests/test_surfaces.py @@ -70,7 +70,7 @@ def test_hosted_surface_uses_datarobot_routes_and_auth_headers( health_url=predict_url, predict_url=predict_url, deployment_selector="deployment_id", - schema_version="v2", + schema_version="v3", model_version=request_payload["model_version"], deployment_id="deployment-id", ) @@ -209,7 +209,7 @@ def test_hosted_surface_auto_discovers_model_version_when_settings_unpinned( health_url=predict_url, predict_url=predict_url, deployment_selector="deployment_id", - schema_version="v2", + schema_version="v3", deployment_id="deployment-id", ) assert settings.model_version is None diff --git a/tests/test_transport.py b/tests/test_transport.py index c05d64f..441d7aa 100644 --- a/tests/test_transport.py +++ b/tests/test_transport.py @@ -154,7 +154,7 @@ def _health_payload( """Health payload.""" return { "status": "ok", - "schema_version": "v2", + "schema_version": "v3", "image_version": "0.3.0", "model_version": model_version, "checkpoint_version": checkpoint_version, @@ -163,6 +163,7 @@ def _health_payload( "head": "studentt", "decoding_strategy": "parallel_dense", "supported_query_modes": ["forecast"], + "supported_condition_kinds": [], "supported_return_modes": ["mean", "samples", "quantiles", "log_prob"], "supported_time_index_modes": [ "ordinal", @@ -189,7 +190,7 @@ def _forecast_response_payload(*, return_mode: str = "mean") -> dict[str, object else None, } return { - "schema_version": "v2", + "schema_version": "v3", "image_version": "0.3.0", "model_version": "jointfm-inference:0.3.0+ckpt.sdk-test", "checkpoint_version": "sdk-test", @@ -233,7 +234,7 @@ def test_transport_posts_json_with_headers_timeout_and_user_agent() -> None: session.mount("https://", adapter) result = transport.post_json( - "https://example.com/predict", {"schema_version": "v2"} + "https://example.com/predict", {"schema_version": "v3"} ) assert result == {"ok": True} @@ -247,7 +248,7 @@ def test_transport_posts_json_with_headers_timeout_and_user_agent() -> None: assert adapter.kwargs[0]["timeout"] == (1.5, 2.5) request_body = request.body assert isinstance(request_body, bytes) - assert json.loads(request_body.decode("utf-8")) == {"schema_version": "v2"} + assert json.loads(request_body.decode("utf-8")) == {"schema_version": "v3"} def test_transport_from_settings_attaches_hosted_auth_headers_and_closes_session() -> ( @@ -264,7 +265,7 @@ def test_transport_from_settings_attaches_hosted_auth_headers_and_closes_session "deployment-id/predictionsUnstructured" ), deployment_selector="deployment_id", - schema_version="v2", + schema_version="v3", model_version="jointfm-inference:0.3.0+ckpt.sdk-test", deployment_id="deployment-id", ) @@ -292,7 +293,7 @@ def test_transport_from_local_settings_omits_hosted_auth_headers() -> None: health_url="http://127.0.0.1:8080/healthz", predict_url="http://127.0.0.1:8080/predict", deployment_selector="local_service", - schema_version="v2", + schema_version="v3", model_version="jointfm-inference:0.3.0+ckpt.local-test", local_base_url="http://127.0.0.1:8080", ) @@ -318,7 +319,7 @@ def test_transport_retries_retryable_server_responses() -> None: retry_config=JointFMRetryConfig(max_attempts=2) ) - result = transport.post_json(_server_url(server), {"schema_version": "v2"}) + result = transport.post_json(_server_url(server), {"schema_version": "v3"}) assert result == {"ok": True} assert handler.request_count == 2 @@ -350,7 +351,7 @@ def test_transport_retries_html_bodied_gateway_errors() -> None: ), ) - result = transport.post_json(_server_url(server), {"schema_version": "v2"}) + result = transport.post_json(_server_url(server), {"schema_version": "v3"}) assert result == {"ok": True} assert handler.request_count == 2 @@ -379,7 +380,7 @@ def test_transport_raises_status_error_when_gateway_html_persists() -> None: ) with pytest.raises(JointFMHTTPStatusError) as exc_info: - transport.post_json(_server_url(server), {"schema_version": "v2"}) + transport.post_json(_server_url(server), {"schema_version": "v3"}) assert exc_info.value.status_code == HTTPStatus.BAD_GATEWAY assert "502 Bad Gateway" in exc_info.value.response_body_excerpt @@ -443,7 +444,7 @@ def request(self, *args: Any, **kwargs: Any) -> requests.Response: ) result = transport.post_json( - "https://example.com/predict", {"schema_version": "v2"} + "https://example.com/predict", {"schema_version": "v3"} ) assert result == {"ok": True} @@ -462,7 +463,7 @@ def test_status_error_carries_parsed_retry_after_seconds() -> None: ) with pytest.raises(JointFMHTTPStatusError) as exc_info: - transport.post_json(_server_url(server), {"schema_version": "v2"}) + transport.post_json(_server_url(server), {"schema_version": "v3"}) assert exc_info.value.retry_after_seconds == 0.5 assert handler.request_count == 1 @@ -480,7 +481,7 @@ def test_transport_does_not_retry_validation_errors() -> None: ) with pytest.raises(JointFMHTTPStatusError) as exc_info: - transport.post_json(_server_url(server), {"schema_version": "v2"}) + transport.post_json(_server_url(server), {"schema_version": "v3"}) assert exc_info.value.status_code == HTTPStatus.BAD_REQUEST assert exc_info.value.datarobot_request_id == "request-id-1" @@ -506,7 +507,7 @@ def test_transport_rejects_non_json_serializable_payloads() -> None: with pytest.raises(JointFMRequestEncodingError, match="JSON-serializable"): transport.post_json( "https://example.com/predict", - {"schema_version": "v2", "bad": object()}, + {"schema_version": "v3", "bad": object()}, ) @@ -571,14 +572,14 @@ def test_client_predict_uses_configured_transport_and_settings() -> None: health_url="https://app.datarobot.com/api/v2/deployments/deployment-id/healthz", predict_url="https://app.datarobot.com/api/v2/deployments/deployment-id/predictionsUnstructured", deployment_selector="deployment_id", - schema_version="v2", + schema_version="v3", model_version="jointfm-inference:0.3.0+ckpt.sdk-test", deployment_id="deployment-id", ) transport = RecordingTransport() client = JointFMClient(settings=settings, transport=transport) payload = { - "schema_version": "v2", + "schema_version": "v3", "model_version": "jointfm-inference:0.3.0+ckpt.sdk-test", } @@ -597,7 +598,7 @@ def test_client_health_returns_typed_metadata_and_caches_only_when_requested() - health_url="https://app.datarobot.com/api/v2/deployments/deployment-id/predictionsUnstructured", predict_url="https://app.datarobot.com/api/v2/deployments/deployment-id/predictionsUnstructured", deployment_selector="deployment_id", - schema_version="v2", + schema_version="v3", model_version="jointfm-inference:0.3.0+ckpt.sdk-test", deployment_id="deployment-id", ) @@ -628,7 +629,7 @@ def test_client_health_instances_returns_one_entry_for_single_endpoint() -> None health_url="https://app.datarobot.com/api/v2/deployments/deployment-id/predictionsUnstructured", predict_url="https://app.datarobot.com/api/v2/deployments/deployment-id/predictionsUnstructured", deployment_selector="deployment_id", - schema_version="v2", + schema_version="v3", model_version="jointfm-inference:0.3.0+ckpt.sdk-test", deployment_id="deployment-id", ) @@ -690,7 +691,7 @@ def capture_transport( "DATAROBOT_ENDPOINT": "https://app.datarobot.com/api/v2", "DATAROBOT_API_TOKEN": "secret-token", "JOINTFM_DEPLOYMENT_ID": "deployment-id", - "JOINTFM_SCHEMA_VERSION": "v2", + "JOINTFM_SCHEMA_VERSION": "v3", "JOINTFM_MODEL_VERSION": "jointfm-inference:0.3.0+ckpt.sdk-test", }, dotenv_path=None, @@ -712,7 +713,7 @@ def test_client_health_rejects_cached_model_mismatch() -> None: health_url="https://app.datarobot.com/api/v2/deployments/deployment-id/predictionsUnstructured", predict_url="https://app.datarobot.com/api/v2/deployments/deployment-id/predictionsUnstructured", deployment_selector="deployment_id", - schema_version="v2", + schema_version="v3", model_version="jointfm-inference:0.3.0+ckpt.sdk-test", deployment_id="deployment-id", ) @@ -738,7 +739,7 @@ def test_client_hosted_health_posts_request_type_health_to_predict_url() -> None health_url=predict_url, predict_url=predict_url, deployment_selector="deployment_id", - schema_version="v2", + schema_version="v3", model_version="jointfm-inference:0.3.0+ckpt.sdk-test", deployment_id="deployment-id", ) @@ -762,7 +763,7 @@ def test_client_local_health_keeps_get_request_to_healthz_route() -> None: health_url="http://127.0.0.1:8080/healthz", predict_url="http://127.0.0.1:8080/predict", deployment_selector="local_service", - schema_version="v2", + schema_version="v3", model_version="jointfm-inference:0.3.0+ckpt.sdk-test", local_base_url="http://127.0.0.1:8080", ) @@ -804,7 +805,7 @@ def test_client_forecast_builds_payload_from_rows_and_returns_typed_response() - assert result.outputs.mean == ((12.0,),) assert transport.predict_url == "http://localhost:8080/predict" assert transport.payload == { - "schema_version": "v2", + "schema_version": "v3", "model_version": "jointfm-inference:0.3.0+ckpt.sdk-test", "query_mode": "forecast", "return_mode": "mean", @@ -979,7 +980,7 @@ def post_json(self, url: str, payload: Mapping[str, Any]) -> Mapping[str, Any]: health_url="http://127.0.0.1:8080/healthz", predict_url="http://127.0.0.1:8080/predict", deployment_selector="local_service", - schema_version="v2", + schema_version="v3", model_version="jointfm-inference:0.3.0+ckpt.sdk-test", local_base_url="http://127.0.0.1:8080", ) @@ -1031,13 +1032,13 @@ def test_client_predict_raises_typed_service_error_for_success_payload_errors() health_url="https://app.datarobot.com/api/v2/deployments/deployment-id/healthz", predict_url="https://app.datarobot.com/api/v2/deployments/deployment-id/predictionsUnstructured", deployment_selector="deployment_id", - schema_version="v2", + schema_version="v3", model_version="jointfm-inference:0.3.0+ckpt.sdk-test", deployment_id="deployment-id", ) transport = RecordingTransport() transport.predict_payload = { - "schema_version": "v2", + "schema_version": "v3", "errors": [ { "code": "VALIDATION_ERROR", @@ -1051,7 +1052,7 @@ def test_client_predict_raises_typed_service_error_for_success_payload_errors() with pytest.raises(JointFMServiceError) as exc_info: client.predict( { - "schema_version": "v2", + "schema_version": "v3", "model_version": settings.model_version, } ) @@ -1140,7 +1141,7 @@ def do_POST(self) -> None: payload = {"ok": True} else: payload = { - "schema_version": "v2", + "schema_version": "v3", "errors": [ { "code": "VALIDATION_ERROR", @@ -1224,7 +1225,7 @@ def _pool_settings(primary: str, backup: str) -> JointFMSettings: health_url=primary, predict_url=primary, deployment_selector="deployment_ids", - schema_version="v2", + schema_version="v3", instances=( JointFMInstanceSettings(deployment_id="primary-id", predict_url=primary), JointFMInstanceSettings(deployment_id="backup-id", predict_url=backup), @@ -1266,7 +1267,7 @@ def post_json(self, url: str, payload: Mapping[str, Any]) -> Mapping[str, Any]: transport = PoolTransport() client = JointFMClient(settings=settings, transport=transport) payload = { - "schema_version": "v2", + "schema_version": "v3", "model_version": "jointfm-inference:0.3.0+ckpt.sdk-test", } client.predict(payload) From aca06f8006f2f4d2471fe535c35dd498292983e3 Mon Sep 17 00:00:00 2001 From: Stefan Hackmann Date: Mon, 21 Sep 2026 15:08:25 +0000 Subject: [PATCH 3/3] feat: Thread condition blocks through the forecast helpers Accept a ConditionBlock on forecast_mean, forecast_samples, and forecast_quantiles so conditional queries reach every high-level surface, not only forecast(). Document conditional queries in the README and API reference, add the forecast_condition notebook, and move the sample env schema pin to v3. --- .env.sample | 2 +- README.md | 60 +++++-- docs/api-reference.md | 45 +++-- notebooks/forecast_condition.ipynb | 243 +++++++++++++++++++++++++++ src/jointfm_client/client.py | 24 ++- tests/test_condition_mode.py | 253 ++++++++++++++++++++++++++++- tests/test_notebooks.py | 1 + 7 files changed, 593 insertions(+), 35 deletions(-) create mode 100644 notebooks/forecast_condition.ipynb diff --git a/.env.sample b/.env.sample index 395cb55..e98b8c8 100644 --- a/.env.sample +++ b/.env.sample @@ -4,7 +4,7 @@ DATAROBOT_ENDPOINT=https://app.datarobot.com/api/v2 DATAROBOT_API_TOKEN= # JointFM compatibility metadata for the selected deployment. -JOINTFM_SCHEMA_VERSION=v1 +JOINTFM_SCHEMA_VERSION=v3 # Optional drift-detection pin; the SDK discovers the model version from /healthz when unset: # JOINTFM_MODEL_VERSION=jointfm-inference:0.2.0+ckpt.fin-2026-05-22 diff --git a/README.md b/README.md index 7d42277..827109b 100644 --- a/README.md +++ b/README.md @@ -10,7 +10,7 @@ The SDK targets the DataRobot-hosted unstructured prediction route and the same - Import namespace: `jointfm_client` - Supported Python: `>=3.11` - Current SDK package version: `0.7.0` -- Current JointFM service schema: `schema_version="v2"` +- Current JointFM service schema: `schema_version="v3"` The public API shape is a synchronous low-level `JointFMClient` with `health()`, `health_instances()`, and `predict(payload)` methods plus high-level `forecast(...)`, `forecast_mean(...)`, `forecast_samples(...)`, and `forecast_quantiles(...)` helpers. The SDK is not a proxy service; callers use it as a local Python library that talks to the hosted or local JointFM endpoint. @@ -50,7 +50,7 @@ Example deployment configuration: deployment: datarobot_endpoint: https://app.datarobot.com/api/v2 datarobot_api_token: - schema_version: v2 + schema_version: v3 deployment_id: # Optional model-version pin; the SDK discovers it from /healthz when unset: # model_version: jointfm-inference:0.3.0+ckpt.fin-2026-05-22 @@ -68,7 +68,7 @@ Equivalent `.env` deployment configuration: ```dotenv DATAROBOT_ENDPOINT=https://app.datarobot.com/api/v2 DATAROBOT_API_TOKEN= -JOINTFM_SCHEMA_VERSION=v2 +JOINTFM_SCHEMA_VERSION=v3 JOINTFM_DEPLOYMENT_ID= # Optional drift-detection pin; the SDK discovers the model version from /healthz when unset: # JOINTFM_MODEL_VERSION=jointfm-inference:0.3.0+ckpt.fin-2026-05-22 @@ -78,7 +78,7 @@ Equivalent local REST configuration for a service started from the `joint` repos ```dotenv JOINTFM_LOCAL_BASE_URL=http://127.0.0.1:8080 -JOINTFM_SCHEMA_VERSION=v2 +JOINTFM_SCHEMA_VERSION=v3 # Optional drift-detection pin; the SDK discovers the model version from /healthz when unset: # JOINTFM_MODEL_VERSION=jointfm-inference:0.3.0+ckpt.fin_i504_o63_f0_t10_h16l16_mam7_af_t3r1_cnn_k3l4_hpst_h16l2_studentt_m4cr2df8skew ``` @@ -163,19 +163,55 @@ bootstrap_notebook(add_src_root=True) Run `task setup` first so VS Code can select the registered `Python (joint-client-python)` notebook kernel backed by this repository's `.venv`. -The bootstrap helper resolves the nearest src-layout Python project root, switches the working directory there, and prepends that project's local `src` tree during development. The examples cover hosted health checks, low-level JSON prediction, mean forecasts, sample forecasts, quantile forecasts, pandas/NumPy result conversion, and CSV forecast workflows. They use `.env.sample` placeholders and checked-in fixture payloads; no real tokens or deployment IDs are stored in notebooks. +The bootstrap helper resolves the nearest src-layout Python project root, switches the working directory there, and prepends that project's local `src` tree during development. The examples cover hosted health checks, low-level JSON prediction, mean forecasts, sample forecasts, quantile forecasts, conditional forecasts, pandas/NumPy result conversion, and CSV forecast workflows. They use `.env.sample` placeholders and checked-in fixture payloads; no real tokens or deployment IDs are stored in notebooks. The current forecast request contract is: -- `schema_version`: exactly `"v2"`, configured as `JOINTFM_SCHEMA_VERSION` for `from_env()` clients +- `schema_version`: exactly `"v3"`, configured as `JOINTFM_SCHEMA_VERSION` for `from_env()` clients - `model_version`: exact model version advertised by `/healthz` or otherwise selected by the caller. Optional for `from_env()` clients: when `JOINTFM_MODEL_VERSION` is unset the SDK reads it from `/healthz` on first use; when set it acts as a drift-detection pin -- `query_mode`: `"forecast"` +- `query_mode`: `"forecast"` for the unconditional forecast, or `"condition"` for a conditional query at one future position; the high-level helpers set it from whether a `condition` block was passed - `return_mode`: one of `"mean"`, `"samples"`, or `"quantiles"` - `time_index_mode`: one of `"ordinal"`, `"continuous_float"`, or `"absolute_datetime"` - `time_column`: required for `"absolute_datetime"`, and used for ordered ordinal or continuous histories when supplied - `query_times`: non-empty future forecast times only - `requested_columns`: optional column names or integer column indices, with duplicates rejected - `n_samples`: positive sample count for sampled forecasts and quantile estimation. When `return_mode="samples"` exceeds the `max_sample_count` advertised by the deployment's health metadata, `forecast_samples(...)` splits the request into capped prediction batches up front and returns one merged `SampleForecastResult`. +- `condition`: required with `query_mode="condition"` and forbidden otherwise. A `ConditionBlock` naming one future position by its index into `query_times` and one condition per column, see below. + +### Conditional Queries + +The `condition` query mode asks for the model's joint distribution at one future position *given* something about some of its columns at that same position. Conditioning relates columns to each other within one position and never across horizons, so the block names the position once and the response describes that position alone: `outputs.query_times` carries exactly one entry however many `query_times` the request listed. + +A column carries at most one condition, of either kind: + +- `EqualityCondition(column, value)` pins the column to a finite value. The pinned column leaves the read-out set, so it must not appear in `requested_columns`. +- `IntervalCondition(column, lower=None, upper=None)` confines the column to a range; `None` leaves that side open, and at least one side must be bounded. The column stays readable, and what comes back is its distribution inside the range. + +Every column without a condition is a read-out column, and at least one must remain. Pass the block to `forecast(...)`, `forecast_mean(...)`, `forecast_samples(...)`, or `forecast_quantiles(...)`: + +```python +from jointfm_client import ConditionBlock, EqualityCondition, IntervalCondition + +block = ConditionBlock( + query_time_index=0, + conditions=[ + EqualityCondition(column="equity_index_level", value=4780.0), + IntervalCondition(column="treasury_10y_yield", lower=0.041, upper=0.045), + ], +) +result = client.forecast_mean( + history, + query_times=query_times, + requested_columns=["portfolio_nav", "realized_volatility"], + columns=plan.columns, + condition=block, +) +print(result.plausibility) +``` + +Whether a deployment can condition depends on the mounted checkpoint's head. `/healthz` advertises `condition` in `supported_query_modes` and the kinds it answers in `supported_condition_kinds` (empty when the mode is absent). The client checks that advertisement before sending, so a deployment that cannot condition is refused with `UnsupportedServiceContractError` rather than after a paid round trip; `require_condition_support(metadata, block)` exposes the same check. + +A condition response carries a `plausibility` block: `equality_log_density` is the log density the model assigns to the pinned values and `region_log_probability` the log probability it gives the interval region, each `None` when the request carried no condition of that kind. They separate a confident answer from one conditioned on something the model finds implausible; the service reports them and never refuses on them. `diagnostics.condition_draws` counts the draws behind sampled outputs, and `diagnostics.interval_estimator` reports the numerical accounting (`points`, `effective_sample_size`) when more than one column carries an interval and the region probability had to be estimated. Column descriptors support the server fields `name`, `modality`, `role`, `nullable`, `vocabulary_size`, `level_count`, `mapping`, `lower_bound`, `upper_bound`, `time_value_kind`, `time_value_scale_seconds`, `time_value_use_local_normalized_time`, `time_value_calendar_id`, and `time_value_timezone`. @@ -187,11 +223,11 @@ uv add "jointfm-client[notebooks]" Use `build_forecast_payload_from_dataframe(...)` when history is already in a pandas DataFrame. It can accept explicit `ColumnSpec` objects or infer basic numeric, categorical, ordinal, count, binary, and time-valued columns from the DataFrame plus role, mapping, nullable, and bounds hints. The helper emits `history_rows` in the same order as the service frame builder: `time_column` first when present, followed by the ordered modeled columns. `build_forecast_payload_from_arrays(...)` provides the same request path for two-dimensional NumPy-like arrays when callers already have array values and column metadata. `build_datetime_query_times(...)`, `build_ordinal_query_times(...)`, `build_continuous_query_times(...)`, and `validate_forecast_horizon(...)` perform local future-horizon validation before the SDK sends the request. -Successful forecast responses preserve `schema_version`, `image_version`, `model_version`, `checkpoint_version`, `head`, `query_mode`, `return_mode`, `outputs`, and `diagnostics`. Structured service errors use this shape: +Successful forecast responses preserve `schema_version`, `image_version`, `model_version`, `checkpoint_version`, `head`, `query_mode`, `return_mode`, `outputs`, `plausibility`, and `diagnostics`. Structured service errors use this shape: ```json { - "schema_version": "v2", + "schema_version": "v3", "errors": [ { "code": "VALIDATION_ERROR", @@ -202,11 +238,11 @@ Successful forecast responses preserve `schema_version`, `image_version`, `model } ``` -Known error codes are `VALIDATION_ERROR`, `SCHEMA_VERSION_MISMATCH`, `MODEL_VERSION_MISMATCH`, `INPUT_SIZE_EXCEEDED`, and `INTERNAL_ERROR`. +Known error codes are `VALIDATION_ERROR`, `UNSUPPORTED_HEAD_QUERY_COMBINATION`, `UNSUPPORTED_RETURN_MODE`, `SCHEMA_VERSION_MISMATCH`, `MODEL_VERSION_MISMATCH`, `INPUT_SIZE_EXCEEDED`, and `INTERNAL_ERROR`. ## Compatibility Policy -The SDK supports only `schema_version="v2"`. `validate_service_metadata()` checks `/healthz` metadata and raises typed compatibility errors before prediction if the service advertises a different schema, an unexpected model version, mode capabilities outside the recorded service contract, or an unsupported `decoding_strategy`. +The SDK supports only `schema_version="v3"`. `validate_service_metadata()` checks `/healthz` metadata and raises typed compatibility errors before prediction if the service advertises a different schema, an unexpected model version, mode capabilities outside the recorded service contract, or an unsupported `decoding_strategy`. Return modes and time-index modes must match the SDK's lists exactly. Query modes and condition kinds are derived by the service from the mounted head, so a deployment may advertise fewer of them than the SDK knows; it must advertise at least one query mode and nothing the SDK does not know. Callers should pass an expected `model_version` when they already know which deployment artifact they intend to use. A mismatch is treated as a hard compatibility error rather than silently downgrading, guessing, or retrying another model. @@ -248,7 +284,7 @@ Create `.env` from `.env.sample` or set the same values in your shell. A hosted ```dotenv DATAROBOT_ENDPOINT=https://app.datarobot.com/api/v2 DATAROBOT_API_TOKEN= -JOINTFM_SCHEMA_VERSION=v2 +JOINTFM_SCHEMA_VERSION=v3 JOINTFM_DEPLOYMENT_ID= # Or: JOINTFM_DEPLOYMENT_IDS=chevron-id,research-id # Optional drift-detection pin; the SDK discovers the model version from /healthz when unset: diff --git a/docs/api-reference.md b/docs/api-reference.md index 650b235..1f05497 100644 --- a/docs/api-reference.md +++ b/docs/api-reference.md @@ -6,7 +6,7 @@ This reference covers the supported public Python surface exported by `jointfm_c | Name | Purpose | | --- | --- | -| `JointFMClient` | Synchronous client for hosted or local JointFM endpoints. Use `from_env()` for `.env` and `config.yaml` backed hosted settings, `health()` for consensus typed service metadata, `health_instances()` for per-deployment probe results and pooled sample topology, `predict(payload)` for low-level JSON prediction, `forecast(...)` for validated tabular forecasts, and the `forecast_mean(...)`, `forecast_samples(...)`, and `forecast_quantiles(...)` convenience methods for typed forecast results. `health()` probes `GET /healthz` for local deployments and POSTs `{"request_type": "health"}` to `predict_url` for hosted DataRobot deployments because the DataRobot deployment gateway only proxies the unstructured prediction route. `feature_importance(...)` runs permutation feature importance: one baseline `forecast_samples` call plus one per shuffled feature column, returning a list of `{"feature", "mean", "distance"}` dicts, each holding that feature's absolute forecast-mean shift and centered squared 2-Wasserstein distance indexed by target and horizon. | +| `JointFMClient` | Synchronous client for hosted or local JointFM endpoints. Use `from_env()` for `.env` and `config.yaml` backed hosted settings, `health()` for consensus typed service metadata, `health_instances()` for per-deployment probe results and pooled sample topology, `predict(payload)` for low-level JSON prediction, `forecast(...)` for validated tabular forecasts, and the `forecast_mean(...)`, `forecast_samples(...)`, and `forecast_quantiles(...)` convenience methods for typed forecast results. Each forecast method accepts `condition=ConditionBlock(...)`, which switches the request to the `condition` query mode and, before anything is sent, checks that the deployment's health metadata advertises the mode and every condition kind the block uses. `health()` probes `GET /healthz` for local deployments and POSTs `{"request_type": "health"}` to `predict_url` for hosted DataRobot deployments because the DataRobot deployment gateway only proxies the unstructured prediction route. `feature_importance(...)` runs permutation feature importance: one baseline `forecast_samples` call plus one per shuffled feature column, returning a list of `{"feature", "mean", "distance"}` dicts, each holding that feature's absolute forecast-mean shift and centered squared 2-Wasserstein distance indexed by target and horizon. | `JointFMClient.from_env()` loads `config.yaml`, optional `.env` values, and process environment variables. `JointFMClient.health(cache=True)` caches health metadata only when requested. `JointFMClient.health_instances()` returns the same probe as a `HealthInstances` object: one `InstanceHealth` per configured deployment (including failures), `max_sample_count` as the sum of reachable caps (overall parallel capacity), and `topology` / `topology_label` grouping those caps (unavailable peers are listed but excluded from the sum and topology). Each endpoint's health payload describes only that endpoint; the client aggregates by calling each configured peer. `health()` still exposes the minimum reachable `max_sample_count`, which is the sample-batch cap used by forecast helpers. `JointFMClient.predict(payload)` requires `payload["model_version"]`; high-level forecast helpers resolve the configured model version when the caller does not pass one explicitly. When `forecast_samples(...)` requests an explicit `n_samples`, the client learns the deployment's `max_sample_count` from health metadata before the first prediction, splits oversized requests into capped prediction batches, and returns one merged `SampleForecastResult`. Clients configured without a reachable health route fall back to discovering the cap from the structured service error. @@ -16,18 +16,23 @@ This reference covers the supported public Python surface exported by `jointfm_c | --- | --- | | `ColumnSpec` | Describes one modeled request column. Fields are `name`, `modality`, `role`, `nullable`, `vocabulary_size`, `level_count`, `mapping`, `lower_bound`, `upper_bound`, `time_value_kind`, `time_value_scale_seconds`, `time_value_use_local_normalized_time`, `time_value_calendar_id`, and `time_value_timezone`. | | `DataFrameSchema` | Describes tabular history layout. Fields are `columns`, `time_index_mode`, `time_column`, `time_scale_seconds`, `use_local_normalized_time`, `calendar_id`, and `timezone`. | -| `ForecastRequestMetadata` | Holds `schema_version`, `model_version`, `query_mode`, and `return_mode` for one forecast request. | -| `ForecastRequest` | Validated request object that combines metadata, schema, history rows, query times, requested columns, sample or quantile controls, and `seed`, then emits a JSON-compatible payload with `to_payload()`. | -| `HealthMetadata` | Typed service-health payload with service status, schema and model versions, checkpoint metadata, device, head, `decoding_strategy`, advertised modes, time-index encoding, `max_sample_count`, and an optional `data_generation` block carrying advertised capacity limits. The container exposes it on `GET /healthz` for direct local access and as the response to `POST {"request_type": "health"}` on the unstructured prediction route for DataRobot-hosted deployments. Each endpoint reports only its own capabilities. | +| `ForecastRequestMetadata` | Holds `schema_version`, `model_version`, `query_mode` (`forecast` or `condition`), and `return_mode` for one forecast request. | +| `ForecastRequest` | Validated request object that combines metadata, schema, history rows, query times, requested columns, sample or quantile controls, `seed`, and an optional `condition` block, then emits a JSON-compatible payload with `to_payload()`. `query_mode="condition"` and a `condition` block must appear together, and the block is validated against the request's own schema and `query_times` before any round trip. | +| `EqualityCondition` | One column of the conditioned position pinned to a finite `value`. The pinned column leaves the read-out set, so it must not appear in `requested_columns`. | +| `IntervalCondition` | One column of the conditioned position confined to `[lower, upper]`; `None` leaves a side open, at least one side must be bounded, and `lower < upper`. The column stays readable and the response describes its distribution inside the range. | +| `ConditionBlock` | Every condition of one request: `query_time_index` into the request's `query_times` and a sequence of `EqualityCondition` or `IntervalCondition` values, at most one per column. `kinds`, `pinned_columns`, and `conditioned_columns` expose what the capability gate and the request validation need. | +| `ConditionPlausibility` | What the model thinks of the conditions it was given: `equality_log_density` of the pinned values and `region_log_probability` of the interval region, each `None` when the request carried no condition of that kind. Reported by the service and never refused on. | +| `IntervalEstimator` | Numerical accounting (`points`, `effective_sample_size`) behind a region probability that had to be estimated, which happens when more than one column carries an interval condition. | +| `HealthMetadata` | Typed service-health payload with service status, schema and model versions, checkpoint metadata, device, head, `decoding_strategy`, advertised query modes, `supported_condition_kinds` (empty when the deployment cannot condition), return modes, time-index modes, time-index encoding, `max_sample_count`, and an optional `data_generation` block carrying advertised capacity limits. The container exposes it on `GET /healthz` for direct local access and as the response to `POST {"request_type": "health"}` on the unstructured prediction route for DataRobot-hosted deployments. Each endpoint reports only its own capabilities. | | `InstanceHealth` | One configured deployment's probe outcome: `deployment_id`, optional `metadata` (`HealthMetadata` when reachable), and optional `error` when the peer was skipped. | | `HealthInstances` | Client aggregation of `health_instances()`: `instances` (one `InstanceHealth` per configured ID), `max_sample_count` (sum of reachable caps = overall parallel capacity), `topology` as `(count, cap)` pairs sorted by descending cap, and `topology_label` such as `2x5000` or `1x7000, 1x3000`. Unavailable peers stay in `instances` but are omitted from the sum and topology. | | `DataGenerationCapabilities` | Optional service-health block describing the deployed checkpoint's data-generation capacity. Fields are `sampler_type`, `min_series`, `max_series`, `t_input`, `t_output`, `n_input`, and `n_output`. `min_series` and `max_series` bound the total number of series in one request, features and targets counted together — there is no separate per-role cap. | | `ForecastPlan` | Validated forecast plan returned by `plan_forecast_columns`. Fields are `columns` (ordered `ColumnSpec` tuple), `feature_columns`, `target_columns` (both as the caller declared them), and `requested_columns` (the caller's target list). | | `StructuredError` | One structured JointFM service error with `code`, `message`, and optional `field`. | -| `ForecastDiagnostics` | Response diagnostics containing `history_rows`, `horizon_count`, and optional `seed`. | +| `ForecastDiagnostics` | Response diagnostics containing `history_rows`, `horizon_count`, optional `seed`, and, on condition responses, `condition_draws` and an optional `interval_estimator`. | | `QuantileForecast` | One quantile surface with `quantile` and `values`; `to_numpy()` returns axis order `(horizon, column)`. | | `ForecastOutputs` | Legacy nested view of parsed output arrays, including `query_times`, `requested_columns`, and exactly one of `mean`, `samples`, or `quantiles`. | -| `ForecastResponse` | Shared base for parsed forecast results. It preserves schema, image, model, checkpoint, head, mode, query-time, requested-column, diagnostic, and error metadata. | +| `ForecastResponse` | Shared base for parsed forecast results. It preserves schema, image, model, checkpoint, head, mode, query-time, requested-column, diagnostic, and error metadata, plus `plausibility` (`None` on forecast responses). | | `MeanForecastResult` | Parsed mean forecast. `to_numpy()` returns `(horizon, column)` and pandas helpers return tidy or wide frames. | | `SampleForecastResult` | Parsed sample forecasts. `to_numpy()` returns `(sample, horizon, column)` and pandas helpers return tidy or wide frames. | | `QuantileForecastResult` | Parsed quantile forecasts. `to_numpy()` returns `(quantile, horizon, column)`, `quantile_levels` exposes the ordered levels, and pandas helpers return tidy or wide frames. | @@ -73,9 +78,9 @@ All SDK-specific exceptions inherit from `JointFMError`. | `JointFMHTTPStatusError` | The service returns an HTTP error status. | | `JointFMServiceError` | A response body contains non-empty JointFM `errors`, including the case where HTTP status unexpectedly succeeded. | | `JointFMCompatibilityError` | Base class for fail-fast service compatibility failures. | -| `UnsupportedSchemaVersionError` | The service or response advertises a schema version other than `v2`. | +| `UnsupportedSchemaVersionError` | The service or response advertises a schema version other than `v3`. | | `UnsupportedModelVersionError` | The service or response model version differs from the configured or requested version. | -| `UnsupportedServiceContractError` | The service-health payload advertises mode capabilities or a `decoding_strategy` outside the recorded service contract. | +| `UnsupportedServiceContractError` | The service-health payload advertises mode capabilities or a `decoding_strategy` outside the recorded service contract, or a condition request targets a deployment that does not advertise the `condition` mode or one of the block's condition kinds. | ## Public Functions @@ -98,7 +103,8 @@ All SDK-specific exceptions inherit from `JointFMError`. | `build_local_predict_url(service_base_url)` | Build a direct local service `/predict` URL. | | `build_datarobot_prediction_headers(api_token)` | Build hosted prediction headers: bearer authorization, broad accept header, and JSON content type. | | `build_forecast_payload(...)` | Build a validated JSON-compatible forecast payload from explicit schema, history rows, query times, and return-mode controls. | -| `validate_service_metadata(metadata, expected_model_version=None)` | Validate the service-health metadata against the supported schema version, the expected model when supplied, advertised mode capabilities, and a supported `decoding_strategy`. | +| `validate_service_metadata(metadata, expected_model_version=None)` | Validate the service-health metadata against the supported schema version, the expected model when supplied, advertised mode capabilities, and a supported `decoding_strategy`. Return and time-index modes must match the SDK's lists exactly; `supported_query_modes` (non-empty) and `supported_condition_kinds` (may be empty) must be subsets of what the SDK knows, because the service derives them from the mounted head. | +| `require_condition_support(metadata, block)` | Raise `UnsupportedServiceContractError` when `HealthMetadata` does not advertise the `condition` query mode or one of the kinds the `ConditionBlock` uses. The forecast helpers call it before sending a condition request. | | `infer_column_specs_from_dataframe(frame, ...)` | Infer ordered `ColumnSpec` objects from a pandas `DataFrame` and explicit role, modality, mapping, nullability, time-value, and bounds hints. | | `dataframe_to_history_rows(frame, schema)` | Convert a pandas `DataFrame` into server-compatible `history_rows`. | | `arrays_to_history_rows(values, columns=..., ...)` | Convert a two-dimensional NumPy-like array plus column metadata into `history_rows`. | @@ -118,7 +124,7 @@ All SDK-specific exceptions inherit from `JointFMError`. | --- | --- | --- | | `DATAROBOT_ENDPOINT` | Hosted calls | HTTPS DataRobot API v2 endpoint, normalized without a trailing slash and required to end in `/api/v2`. | | `DATAROBOT_API_TOKEN` | Hosted calls | Non-empty, whitespace-free API token used in the hosted bearer authorization header. | -| `JOINTFM_SCHEMA_VERSION` | Hosted calls | Request schema pin. The SDK supports only `v2`. | +| `JOINTFM_SCHEMA_VERSION` | Hosted calls | Request schema pin. The SDK supports only `v3`. | | `JOINTFM_MODEL_VERSION` | Hosted calls | Exact JointFM deployment model version expected from the service-health payload and prediction responses. | | `JOINTFM_DEPLOYMENT_ID` | One selector | Deployment ID used to build hosted health and prediction URLs. | | `JOINTFM_DEPLOYMENT_IDS` | One selector | Comma-separated hosted deployment IDs for round-robin load balancing (at least two unique IDs). Mutually exclusive with other selectors. Peers must share `model_version` and `checkpoint_version`. `health()` uses the minimum reachable `max_sample_count` as the sample-batch cap; `health_instances()` sums reachable caps for overall parallel capacity and reports topology. | @@ -149,9 +155,9 @@ The string literals are exposed as `PREDICT_REQUEST_TYPE`, `HEALTH_REQUEST_TYPE` | Field | Required | Description | | --- | --- | --- | | `request_type` | Optional | One of `"predict"` (default) or `"health"`. Forecast requests omit this field or set it to `"predict"`. | -| `schema_version` | Yes | Must be `"v2"`. | +| `schema_version` | Yes | Must be `"v3"`. | | `model_version` | Yes | Exact deployed model version expected by the caller. | -| `query_mode` | Yes | Must be `"forecast"`. | +| `query_mode` | Yes | `"forecast"` for the unconditional forecast or `"condition"` for a conditional query at one future position. | | `return_mode` | Yes | One of `"mean"`, `"samples"`, `"quantiles"`, or `"log_prob"`. The high-level `forecast_mean`, `forecast_samples`, and `forecast_quantiles` helpers cover the first three; `"log_prob"` is reachable through the low-level `predict(payload)` path. | | `time_index_mode` | Yes | One of `"ordinal"`, `"continuous_float"`, or `"absolute_datetime"`. | | `columns` | Yes | Non-empty array of column descriptors for modeled columns. | @@ -162,6 +168,7 @@ The string literals are exposed as `PREDICT_REQUEST_TYPE`, `HEALTH_REQUEST_TYPE` | `n_samples` | Samples and quantiles controls | Positive sample count when sampling controls are needed. Oversized sample forecasts are batched automatically against the cap advertised in health metadata. | | `quantiles` | Quantiles mode | Quantile levels in `(0, 1)`, required for `return_mode="quantiles"`. | | `seed` | Optional | Integer random seed for reproducible stochastic outputs. | +| `condition` | With `query_mode="condition"` | Object with `query_time_index` (index into `query_times`) and `conditions`, a list of `{"column", "kind": "equality", "value"}` or `{"column", "kind": "interval", "lower", "upper"}` entries with `null` for an open bound. At most one condition per column, at least one column left unconditioned, and no pinned column in `requested_columns`. Forbidden with any other query mode. | | `time_scale_seconds` | Optional | Positive scale for continuous time indexes. | | `use_local_normalized_time` | Optional | Whether the service should use local normalized time features. | | `calendar_id` | Optional | Calendar identifier, defaulting to `pandas-default`. | @@ -190,14 +197,14 @@ The string literals are exposed as `PREDICT_REQUEST_TYPE`, `HEALTH_REQUEST_TYPE` | Field | Description | | --- | --- | -| `schema_version` | Response schema, expected to be `"v2"`. | +| `schema_version` | Response schema, expected to be `"v3"`. | | `image_version` | Service image version that produced the response. | | `model_version` | Model version that produced the response. | | `checkpoint_version` | Checkpoint version that produced the response. | | `head` | Forecast head used by the service. | -| `query_mode` | Response query mode, expected to be `"forecast"`. | +| `query_mode` | Response query mode, matching the request: `"forecast"` or `"condition"`. | | `return_mode` | Response return mode matching the request. | -| `outputs.query_times` | Forecast horizon values preserved from the request. | +| `outputs.query_times` | Forecast horizon values preserved from the request. A condition response carries only the conditioned position. | | `outputs.requested_columns` | Output columns in response order. | | `outputs.mean` | Mean values with axis order `(horizon, column)` when `return_mode="mean"`. | | `outputs.samples` | Sample values with axis order `(sample, horizon, column)` when `return_mode="samples"`. | @@ -205,6 +212,9 @@ The string literals are exposed as `PREDICT_REQUEST_TYPE`, `HEALTH_REQUEST_TYPE` | `diagnostics.history_rows` | Number of history rows processed. | | `diagnostics.horizon_count` | Number of forecast horizon steps returned. | | `diagnostics.seed` | Optional seed used by the service. | +| `diagnostics.condition_draws` | Condition responses only: number of draws behind sampled outputs. Batched sample requests report the merged count. | +| `diagnostics.interval_estimator` | Condition responses only, and only when more than one column carries an interval: `points` and `effective_sample_size` of the numerical region-probability estimate. | +| `plausibility` | `null` on forecast responses. On condition responses an object with `equality_log_density` (log density of the pinned values) and `region_log_probability` (log probability of the interval region), each `null` when the request carried no condition of that kind. | | `errors` | Structured service errors. Non-empty arrays raise typed SDK exceptions. | ### Health Metadata @@ -212,7 +222,7 @@ The string literals are exposed as `PREDICT_REQUEST_TYPE`, `HEALTH_REQUEST_TYPE` | Field | Description | | --- | --- | | `status` | Service status string. | -| `schema_version` | Advertised schema version. The SDK requires `v2`. | +| `schema_version` | Advertised schema version. The SDK requires `v3`. | | `image_version` | Running service image version. | | `model_version` | Running model version. | | `checkpoint_version` | Loaded checkpoint version. | @@ -220,7 +230,8 @@ The string literals are exposed as `PREDICT_REQUEST_TYPE`, `HEALTH_REQUEST_TYPE` | `device` | Device used by inference. | | `head` | Active forecast head. | | `decoding_strategy` | Horizon decoding mode advertised by the mounted model. Must be one of `SUPPORTED_DECODING_STRATEGIES`: `parallel_dense`, `parallel_scalable`, or `autoregressive`. Parallel strategies decode every horizon in one pass; `autoregressive` rolls horizons sequentially. | -| `supported_query_modes` | Must match the SDK's query modes. | +| `supported_query_modes` | Non-empty subset of the SDK's query modes (`forecast`, `condition`); the service derives it from the mounted head. | +| `supported_condition_kinds` | Subset of the SDK's condition kinds (`equality`, `interval`), empty when `condition` is not advertised. Condition requests are refused locally when the block uses a kind that is missing here. | | `supported_return_modes` | Must match the SDK's return modes (`mean`, `samples`, `quantiles`, `log_prob`). | | `supported_time_index_modes` | Must match the SDK's time-index modes. | | `time_index_encoding` | Time-index encoding advertised by the service. | diff --git a/notebooks/forecast_condition.ipynb b/notebooks/forecast_condition.ipynb new file mode 100644 index 0000000..521c81e --- /dev/null +++ b/notebooks/forecast_condition.ipynb @@ -0,0 +1,243 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "0", + "metadata": { + "id": "license-header", + "language": "markdown" + }, + "source": [ + "Copyright 2026 DataRobot, Inc. and its affiliates.\n", + "\n", + "Licensed under the Apache License, Version 2.0 (the \"License\");\n", + "you may not use this file except in compliance with the License.\n", + "You may obtain a copy of the License at\n", + "\n", + " http://www.apache.org/licenses/LICENSE-2.0\n", + "\n", + "Unless required by applicable law or agreed to in writing, software\n", + "distributed under the License is distributed on an \"AS IS\" BASIS,\n", + "WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n", + "See the License for the specific language governing permissions and\n", + "limitations under the License." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "1", + "metadata": { + "id": "condition-bootstrap", + "language": "python" + }, + "outputs": [], + "source": [ + "from jointfm_client import bootstrap_notebook\n", + "\n", + "bootstrap_notebook(add_src_root=True)" + ] + }, + { + "cell_type": "markdown", + "id": "2", + "metadata": { + "id": "condition-description", + "language": "markdown" + }, + "source": [ + "# Conditional Forecast\n", + "Forecast USD portfolio NAV and risk from 100 positive float daily observations: equity index level, 10-year Treasury yield, and one EUR/USD FX rate. Instead of the unconditional forecast, ask the model a *what-if* question about the first future step: what does it expect for portfolio NAV and realized volatility **given** something about the other columns at that same step.\n", + "\n", + "The `condition` query mode answers this in closed form from the model's joint distribution at one future position. A request names that position once, by its index into `query_times`, and attaches one condition per column it wants to fix:\n", + "\n", + "- An **equality condition** pins a column to a value (`EqualityCondition`). The pinned column leaves the read-out set, because reading it back would only repeat the request.\n", + "- An **interval condition** confines a column to a range whose bounds may be open on either side (`IntervalCondition`). The column stays readable: what comes back is its distribution inside the range.\n", + "\n", + "Every column without a condition is a read-out column, and the response describes the conditional distribution of those columns at the conditioned position only, so `outputs.query_times` has exactly one entry however many `query_times` the request carried.\n", + "\n", + "Whether a deployment can condition depends on the checkpoint's head, so `/healthz` advertises `condition` in `supported_query_modes` and the kinds it answers in `supported_condition_kinds`. The client checks that advertisement before sending, and this notebook reads it explicitly so the check is visible." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "3", + "metadata": { + "id": "condition-equality-example", + "language": "python" + }, + "outputs": [], + "source": [ + "from pathlib import Path\n", + "\n", + "import pandas as pd\n", + "\n", + "from jointfm_client import (\n", + " ConditionBlock,\n", + " EqualityCondition,\n", + " JointFMClient,\n", + " plan_forecast_columns,\n", + ")\n", + "\n", + "HISTORY_PATH = Path(\"notebooks/history.csv\")\n", + "FEATURE_COLUMNS = [\"equity_index_level\", \"treasury_10y_yield\", \"eur_usd_rate\"]\n", + "TARGET_COLUMNS = [\"portfolio_nav\", \"realized_volatility\"]\n", + "INPUT_STEPS = 100\n", + "OUTPUT_HORIZONS = 10\n", + "CONDITIONED_STEP = 0\n", + "EQUITY_RALLY = 1.02\n", + "EXPECTED_COLUMNS = FEATURE_COLUMNS + TARGET_COLUMNS\n", + "QUERY_TIMES = list(range(INPUT_STEPS, INPUT_STEPS + OUTPUT_HORIZONS))\n", + "\n", + "history = pd.read_csv(HISTORY_PATH, dtype=float)\n", + "if list(history.columns) != EXPECTED_COLUMNS:\n", + " raise ValueError(\n", + " f\"Expected columns {EXPECTED_COLUMNS!r}, got {list(history.columns)!r}\"\n", + " )\n", + "if len(history) != INPUT_STEPS:\n", + " raise ValueError(f\"Expected {INPUT_STEPS} history rows, got {len(history)}\")\n", + "\n", + "client = JointFMClient.from_env()\n", + "health = client.health(cache=True)\n", + "if \"condition\" not in health.supported_query_modes:\n", + " raise RuntimeError(\n", + " f\"This deployment serves {list(health.supported_query_modes)} only; \"\n", + " \"mount a checkpoint whose head can condition\"\n", + " )\n", + "print(\"condition kinds served:\", list(health.supported_condition_kinds))\n", + "\n", + "plan = plan_forecast_columns(\n", + " health=health,\n", + " feature_columns=FEATURE_COLUMNS,\n", + " target_columns=TARGET_COLUMNS,\n", + " history_length=len(history),\n", + " query_times_length=len(QUERY_TIMES),\n", + ")\n", + "\n", + "last_equity_level = float(history[\"equity_index_level\"].iloc[-1])\n", + "rally = ConditionBlock(\n", + " query_time_index=CONDITIONED_STEP,\n", + " conditions=[\n", + " EqualityCondition(\n", + " column=\"equity_index_level\", value=last_equity_level * EQUITY_RALLY\n", + " )\n", + " ],\n", + ")\n", + "result = client.forecast_mean(\n", + " history,\n", + " query_times=QUERY_TIMES,\n", + " requested_columns=plan.requested_columns,\n", + " columns=plan.columns,\n", + " seed=7,\n", + " condition=rally,\n", + ")\n", + "if result.query_times != (QUERY_TIMES[CONDITIONED_STEP],):\n", + " raise ValueError(\n", + " f\"Expected the conditioned position alone, got {result.query_times!r}\"\n", + " )\n", + "if result.plausibility is None:\n", + " raise ValueError(\"A condition response must carry its plausibility block\")\n", + "print(\"log density of the pinned value:\", result.plausibility.equality_log_density)\n", + "forecast = result.to_pandas_tidy()\n", + "expected_forecast_rows = len(plan.requested_columns)\n", + "if len(forecast) != expected_forecast_rows:\n", + " raise ValueError(\n", + " f\"Expected {expected_forecast_rows} forecast rows, got {len(forecast)}\"\n", + " )\n", + "forecast" + ] + }, + { + "cell_type": "markdown", + "id": "4", + "metadata": { + "id": "condition-plausibility", + "language": "markdown" + }, + "source": [ + "## Reading the plausibility\n", + "`result.plausibility.equality_log_density` is the log density the model assigns to the pinned value before conditioning. It separates *the model is confident about NAV given this rally* from *the model finds a rally of this size absurd and is extrapolating*. The service reports the number and never refuses on it; comparing it across candidate pins, or against the density of a pin at the model's own unconditional mean, is the caller's decision." + ] + }, + { + "cell_type": "markdown", + "id": "5", + "metadata": { + "id": "condition-interval-description", + "language": "markdown" + }, + "source": [ + "## Interval condition\n", + "Now confine the 10-year yield to a band around its last observed value instead of pinning it, and pin the equity index at the same time: a request may mix both kinds across the columns of one position. The response then carries `region_log_probability`, the log probability the model gives the yield band, and the yield column itself stays readable because its distribution inside the band is a genuine answer.\n", + "\n", + "With one interval column the region probability is exact. When several columns carry intervals the service estimates the probability of the box numerically and reports the accounting in `diagnostics.interval_estimator`, so the caller can judge the estimate." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "6", + "metadata": { + "id": "condition-interval-example", + "language": "python" + }, + "outputs": [], + "source": [ + "from jointfm_client import IntervalCondition\n", + "\n", + "YIELD_BAND_HALF_WIDTH = 0.002\n", + "\n", + "last_yield = float(history[\"treasury_10y_yield\"].iloc[-1])\n", + "rally_with_yield_band = ConditionBlock(\n", + " query_time_index=CONDITIONED_STEP,\n", + " conditions=[\n", + " EqualityCondition(\n", + " column=\"equity_index_level\", value=last_equity_level * EQUITY_RALLY\n", + " ),\n", + " IntervalCondition(\n", + " column=\"treasury_10y_yield\",\n", + " lower=last_yield - YIELD_BAND_HALF_WIDTH,\n", + " upper=last_yield + YIELD_BAND_HALF_WIDTH,\n", + " ),\n", + " ],\n", + ")\n", + "banded = client.forecast_mean(\n", + " history,\n", + " query_times=QUERY_TIMES,\n", + " requested_columns=[\"treasury_10y_yield\", *plan.requested_columns],\n", + " columns=plan.columns,\n", + " seed=7,\n", + " condition=rally_with_yield_band,\n", + ")\n", + "if banded.plausibility is None:\n", + " raise ValueError(\"A condition response must carry its plausibility block\")\n", + "print(\"log density of the pinned value:\", banded.plausibility.equality_log_density)\n", + "print(\"log probability of the yield band:\", banded.plausibility.region_log_probability)\n", + "print(\"interval estimator accounting:\", banded.diagnostics.interval_estimator)\n", + "banded.to_pandas_tidy()" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "joint-client-python (3.13.3)", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.13.3" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/src/jointfm_client/client.py b/src/jointfm_client/client.py index f0d9b19..2b969a8 100644 --- a/src/jointfm_client/client.py +++ b/src/jointfm_client/client.py @@ -410,8 +410,13 @@ def forecast_mean( requested_columns: Sequence[str | int] | None = None, model_version: str | None = None, seed: int | None = None, + condition: ConditionBlock | None = None, ) -> MeanForecastResult: - """Forecast mean values through the shared forecast validation path.""" + """Forecast mean values through the shared forecast validation path. + + With ``condition`` the mean is the conditional mean at the one future + position the block names; see :meth:`forecast`. + """ return cast( MeanForecastResult, self.forecast( @@ -425,6 +430,7 @@ def forecast_mean( return_mode="mean", model_version=model_version, seed=seed, + condition=condition, ), ) @@ -441,8 +447,13 @@ def forecast_samples( model_version: str | None = None, n_samples: int | None = None, seed: int | None = None, + condition: ConditionBlock | None = None, ) -> SampleForecastResult: - """Forecast sample paths through the shared forecast validation path.""" + """Forecast sample paths through the shared forecast validation path. + + With ``condition`` the draws come from the conditional at the one future + position the block names; see :meth:`forecast`. + """ return cast( SampleForecastResult, self.forecast( @@ -457,6 +468,7 @@ def forecast_samples( model_version=model_version, n_samples=n_samples, seed=seed, + condition=condition, ), ) @@ -474,8 +486,13 @@ def forecast_quantiles( n_samples: int | None = None, quantiles: Sequence[float | int] | None = None, seed: int | None = None, + condition: ConditionBlock | None = None, ) -> QuantileForecastResult: - """Forecast quantiles through the shared forecast validation path.""" + """Forecast quantiles through the shared forecast validation path. + + With ``condition`` the quantiles describe the conditional at the one + future position the block names; see :meth:`forecast`. + """ return cast( QuantileForecastResult, self.forecast( @@ -491,6 +508,7 @@ def forecast_quantiles( n_samples=n_samples, quantiles=quantiles, seed=seed, + condition=condition, ), ) diff --git a/tests/test_condition_mode.py b/tests/test_condition_mode.py index 5995393..31d7deb 100644 --- a/tests/test_condition_mode.py +++ b/tests/test_condition_mode.py @@ -24,8 +24,8 @@ from __future__ import annotations -from collections.abc import Callable -from typing import Any +from collections.abc import Callable, Mapping +from typing import Any, cast import pytest @@ -40,8 +40,13 @@ ForecastResponse, HealthMetadata, IntervalCondition, + JointFMClient, + MeanForecastResult, + SampleForecastResult, build_forecast_payload, + build_forecast_payload_from_dataframe, require_condition_support, + validate_service_metadata, ) from jointfm_client.contract import QueryMode from jointfm_client.exceptions import UnsupportedServiceContractError @@ -351,3 +356,247 @@ def test_an_interval_response_reads_back_its_estimator_accuracy( assert estimator is not None assert estimator.points == 16384 assert estimator.effective_sample_size == pytest.approx(11453.2) + + +class _ConditionTransport: + """Fake JSON transport for a deployment that advertises the condition mode. + + It records every predict payload and answers each sample request with as + many draws as it asked for, so the batching path can be checked end to end + without a service. + """ + + def __init__( + self, + *, + health_payload: dict[str, Any], + predict_payload: dict[str, Any] | None = None, + ) -> None: + """Remember the advertisement and the canned prediction to serve.""" + self.health_payload = health_payload + self.predict_payload = predict_payload + self.payloads: list[dict[str, Any]] = [] + self.health_count = 0 + + def get_json(self, url: str) -> Mapping[str, Any]: + """Serve the health advertisement on the local health route.""" + assert url == "http://127.0.0.1:8080/healthz" + self.health_count += 1 + return self.health_payload + + def post_json(self, url: str, payload: Mapping[str, Any]) -> Mapping[str, Any]: + """Record the predict payload and answer it.""" + assert url == "http://127.0.0.1:8080/predict" + self.payloads.append(dict(payload)) + if self.predict_payload is not None: + return self.predict_payload + sample_count = payload["n_samples"] + assert isinstance(sample_count, int) + start = sum(cast(int, earlier["n_samples"]) for earlier in self.payloads[:-1]) + return { + "schema_version": "v3", + "image_version": "0.3.0", + "model_version": _MODEL_VERSION, + "checkpoint_version": "sdk-test", + "head": "gmm", + "query_mode": "condition", + "return_mode": "samples", + "outputs": { + "query_times": [3], + "requested_columns": ["target"], + "mean": None, + "samples": [[[float(start + index)]] for index in range(sample_count)], + "quantiles": None, + }, + "plausibility": { + "equality_log_density": -1.27, + "region_log_probability": None, + }, + "diagnostics": { + "history_rows": 2, + "horizon_count": 1, + "seed": payload.get("seed"), + "condition_draws": sample_count, + "interval_estimator": None, + }, + "errors": [], + } + + +def _health_payload( + *, + query_modes: list[str], + condition_kinds: list[str], + max_sample_count: int = 4096, +) -> dict[str, Any]: + """Build one health advertisement as the service serializes it.""" + return { + "status": "ok", + "schema_version": "v3", + "image_version": "0.3.0", + "model_version": _MODEL_VERSION, + "checkpoint_version": "sdk-test", + "checkpoint_path": "/models/jointfm.pt", + "device": "cpu", + "head": "gmm", + "decoding_strategy": "parallel_dense", + "supported_query_modes": query_modes, + "supported_condition_kinds": condition_kinds, + "supported_return_modes": ["mean", "samples", "quantiles", "log_prob"], + "supported_time_index_modes": [ + "ordinal", + "continuous_float", + "absolute_datetime", + ], + "time_index_encoding": "legacy_discrete_grid", + "max_sample_count": max_sample_count, + } + + +def _client(transport: _ConditionTransport) -> JointFMClient: + """Build a local-service client over the fake transport.""" + return JointFMClient( + health_url="http://127.0.0.1:8080/healthz", + predict_url="http://127.0.0.1:8080/predict", + transport=transport, + ) + + +_HISTORY_ROWS = ( + {"driver": 1.0, "hedge": 0.5, "target": 10.0}, + {"driver": 1.1, "hedge": 0.6, "target": 11.0}, +) +_PIN_DRIVER = ConditionBlock( + query_time_index=1, + conditions=(EqualityCondition(column="driver", value=1.5),), +) + + +def test_the_client_sends_the_condition_and_reads_the_answer_back( + json_fixture_loader: Callable[[str], dict[str, Any]], +) -> None: + """The typed helper carries the block onto the wire and types what comes back.""" + transport = _ConditionTransport( + health_payload=_health_payload( + query_modes=["forecast", "condition"], + condition_kinds=["equality", "interval"], + ), + predict_payload=json_fixture_loader("condition_mean_response"), + ) + + result = _client(transport).forecast_mean( + list(_HISTORY_ROWS), + schema=_schema(), + query_times=[2, 3], + requested_columns=["target"], + model_version=_MODEL_VERSION, + seed=7, + condition=_PIN_DRIVER, + ) + + assert isinstance(result, MeanForecastResult) + assert result.query_mode == "condition" + assert result.query_times == (3,) + assert result.mean == ((13.5,),) + assert result.plausibility == ConditionPlausibility(equality_log_density=-1.27) + assert len(transport.payloads) == 1 + sent = transport.payloads[0] + assert sent["query_mode"] == "condition" + assert sent["condition"] == _PIN_DRIVER.to_payload() + + +def test_the_client_refuses_before_posting_when_the_deployment_cannot_condition() -> ( + None +): + """The gate runs on the advertisement, so the predict route is never touched.""" + transport = _ConditionTransport( + health_payload=_health_payload(query_modes=["forecast"], condition_kinds=[]), + ) + + with pytest.raises(UnsupportedServiceContractError, match="does not serve"): + _client(transport).forecast_mean( + list(_HISTORY_ROWS), + schema=_schema(), + query_times=[2, 3], + requested_columns=["target"], + model_version=_MODEL_VERSION, + condition=_PIN_DRIVER, + ) + + assert transport.health_count == 1 + assert transport.payloads == [] + + +def test_batched_condition_samples_merge_into_one_conditional_answer() -> None: + """Every batch repeats the same block; the merge recounts the draws and keeps the plausibility.""" + transport = _ConditionTransport( + health_payload=_health_payload( + query_modes=["forecast", "condition"], + condition_kinds=["equality", "interval"], + max_sample_count=2, + ), + ) + + result = _client(transport).forecast_samples( + list(_HISTORY_ROWS), + schema=_schema(), + query_times=[2, 3], + requested_columns=["target"], + model_version=_MODEL_VERSION, + n_samples=3, + seed=7, + condition=_PIN_DRIVER, + ) + + assert isinstance(result, SampleForecastResult) + assert result.samples == (((0.0,),), ((1.0,),), ((2.0,),)) + assert result.query_times == (3,) + assert result.diagnostics.condition_draws == 3 + assert result.plausibility == ConditionPlausibility(equality_log_density=-1.27) + assert [payload["n_samples"] for payload in transport.payloads] == [2, 1] + assert all( + payload["query_mode"] == "condition" + and payload["condition"] == _PIN_DRIVER.to_payload() + for payload in transport.payloads + ) + + +def test_the_dataframe_adapter_builds_a_condition_request() -> None: + """A pandas caller passes the block and gets the condition envelope.""" + pandas = pytest.importorskip("pandas") + frame = pandas.DataFrame(list(_HISTORY_ROWS)) + + payload = build_forecast_payload_from_dataframe( + frame, + model_version=_MODEL_VERSION, + time_index_mode="ordinal", + query_times=[2, 3], + target_columns=["target"], + requested_columns=["target"], + condition=_PIN_DRIVER, + ) + + assert payload["query_mode"] == "condition" + assert payload["condition"] == _PIN_DRIVER.to_payload() + assert payload["requested_columns"] == ["target"] + + +@pytest.mark.parametrize( + ("field", "advertised", "message"), + [ + ("supported_query_modes", [], "advertises nothing"), + ("supported_condition_kinds", ["parametric"], "does not know"), + ], +) +def test_metadata_validation_rejects_an_advertisement_this_client_cannot_serve( + field: str, advertised: list[str], message: str +) -> None: + """Fewer capabilities than the SDK knows are fine; none, or unknown ones, are not.""" + payload = _health_payload( + query_modes=["forecast", "condition"], + condition_kinds=["equality", "interval"], + ) + payload[field] = advertised + + with pytest.raises(UnsupportedServiceContractError, match=message): + validate_service_metadata(payload, expected_model_version=_MODEL_VERSION) diff --git a/tests/test_notebooks.py b/tests/test_notebooks.py index 76b7bda..a4e2428 100644 --- a/tests/test_notebooks.py +++ b/tests/test_notebooks.py @@ -128,6 +128,7 @@ def test_example_notebooks_start_with_bootstrap_cell() -> None: ) assert [path.name for path in notebook_paths] == [ + "forecast_condition.ipynb", "forecast_csv.ipynb", "forecast_mean.ipynb", "forecast_quantiles.ipynb",