feat(ai): add data classes for time series forecasting and imputation by miladh · Pull Request #2873 · cognitedata/cognite-sdk-python

miladh · 2026-10-05T09:18:07Z

Description

First PR in a stack that adds SDK support for the ai/timeseries/forecast and ai/timeseries/impute endpoints. The API is in beta.

The stack builds the feature in two layers, both under client.ai.time_series:

client.ai.time_series.forecast / .impute: a low-level client that mirrors the API. PRs 1–4.
client.ai.time_series.forecast_with_retrieval / .impute_with_retrieval: convenience methods that work on CDF time series by instance id. PRs 5–8. They read history through the existing client.time_series.data.retrieve.

This PR only adds data classes, with no API calls yet:

InputDatapoint and InputTimeSeries, the request objects. They're dataclasses, like DatapointsQuery.
QuantileDatapoint.
ForecastResult(List) and ImputeResult(List). The lists carry quantile_levels and cohorts from the response metadata, the way NodeList carries typing/debug. They support .get(label=...) and to_pandas().

Results are CogniteResource subclasses, unlike the existing document-AI classes in data_classes/ai.py. That gives them dump/load, notebook rendering and list lookups for free.

Where the stack is going

The interface once all 8 PRs are in:

from cognite.client import CogniteClient
from cognite.client.data_classes.ai import InputTimeSeries
from cognite.client.data_classes.data_modeling import NodeId

client = CogniteClient()

# Low-level: your own series, close to the API (PRs 1-4)
res = client.ai.time_series.forecast(
    InputTimeSeries("21-PT-1019", [(1790812800000, 42.1), (1790812860000, 42.4)])
)
res.get(label="21-PT-1019").to_pandas()   # DatetimeIndex, one column per quantile level

# High-level: CDF time series by instance id (PRs 5-8)
node = NodeId("my_space", "21-PT-1019")
fc = client.ai.time_series.forecast_with_retrieval(
    instance_id=node, start="30d-ago", granularity="1h", aggregate="average"
)
imp = client.ai.time_series.impute_with_retrieval(
    instance_id=node, start="7d-ago", granularity="10m", aggregate="average", fill_gaps=True
)

What this PR adds, on its own:

from cognite.client.data_classes.ai import ForecastResultList, InputDatapoint, InputTimeSeries

series = InputTimeSeries(
    "23-PT-1101",
    [
        InputDatapoint(1790812800000, 42.1),
        InputDatapoint(1790812860000, None, missing=True),   # a gap: hidden from the model
        (1790812920000, 42.9),                                # tuples work too
    ],
    cohort="compression-train-a",
)
series.dump()   # the API request shape: {"label": ..., "cohort": ..., "datapoints": [...]}

res = ForecastResultList._load_response(api_response)   # done by the API in PR 3
res.quantile_levels        # [0.05, 0.1, ..., 0.95]
res.cohorts                # ["compression-train-a"]
res.get(label="23-PT-1101").forecast[0].quantiles   # {0.05: 41.0, ..., 0.95: 43.0}
res.to_pandas()            # DatetimeIndex, (label, quantile) columns

Checklist:

Tests added/updated: test_data_classes/test_ai.py. The pandas tests are marked dsl.
Documentation updated: docstrings. Rendered docs come in PR 3.
The PR title follows the Conventional Commit spec.

🤖 Generated with Claude Code

niekbeckers

For my own understanding, we cannot reuse the Datapoint and Timeseries data classes, or have composites that rely on a combination of primitives (e.g., Sequence) with those datamodels?

Request and result types for the upcoming ai/timeseries/{forecast,impute} endpoints: - InputDatapoint / InputTimeSeries: request objects (dataclasses, like DatapointsQuery). Datapoints can be InputDatapoint, (timestamp, value) tuples or API-shaped dicts; timestamps accept ms, datetime or time-shift strings. - QuantileDatapoint: one step with quantile level -> value. Non-numeric keys are ignored on load. - ForecastResult(List) / ImputeResult(List): one result per input label, with quantile_levels and cohorts from the response metadata (like NodeList's typing/debug), lookup by .get(label=...), and to_pandas() with a DatetimeIndex. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

gemini-code-assist

Code Review

This pull request introduces new data classes and helper functions in cognite/client/data_classes/ai.py to support forecasting and imputation features, including InputDatapoint, InputTimeSeries, QuantileDatapoint, ForecastResult, ForecastResultList, ImputeResult, and ImputeResultList, along with corresponding unit tests. Feedback on the changes includes handling empty inputs in _concat_columns to prevent pandas errors, supporting lists in addition to tuples for input datapoints, and safely handling potential null metadata values in API responses to avoid attribute errors.

pd.concat raised "No objects to concatenate" for a ForecastResultList or ImputeResultList with no results (e.g. an empty slice). Other SDK lists return an empty DataFrame, so do the same. Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>

miladh · 2026-11-02T11:58:11Z

Re @niekbeckers' question on reusing Datapoint and TimeSeries: the existing classes don't fit this API. The input needs a nullable value and a missing flag, results are quantiles per step rather than a single value or aggregate, and series are identified by a caller-chosen label, not a CDF id. The CDF-identified variants (TimeSeriesForecast, TimeSeriesImpute, keyed by instance_id) come later in the stack, in PR 06.

— Written by Claude

Direct pd.to_datetime calls are only allowed in _pandas_helpers.py (enforced by test_pandas_timestamp_constructors_only_used_in_pandas_helpers), so build the empty DatetimeIndex with to_pandas_datetime_index instead. Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>

codecov · 2026-11-02T13:35:06Z

Codecov Report

❌ Patch coverage is 98.51485% with 3 lines in your changes missing coverage. Please review.
✅ Project coverage is 93.73%. Comparing base (1ab3df0) to head (cb8eba8).
⚠️ Report is 28 commits behind head on master.

Files with missing lines	Patch %	Lines
cognite/client/data_classes/ai.py	97.36%	3 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #2873      +/-   ##
==========================================
+ Coverage   93.49%   93.73%   +0.24%     
==========================================
  Files         527      541      +14     
  Lines       55489    57404    +1915     
==========================================
+ Hits        51879    53808    +1929     
+ Misses       3610     3596      -14

Files with missing lines	Coverage Δ
tests/tests_unit/test_base.py	`98.94% <100.00%> (+<0.01%)`	⬆️
tests/tests_unit/test_data_classes/test_ai.py	`100.00% <100.00%> (ø)`
cognite/client/data_classes/ai.py	`97.35% <97.36%> (-0.02%)`	⬇️

... and 112 files with indirect coverage changes

🚀 New features to boost your workflow:

❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

haakonvt · 2026-11-02T22:25:21Z

+
+
+@dataclass
+class InputDatapoint:


I think we need a more specific name that's tied to its use, e.g. ForecastInputDatapoint.

I guess this is also used for /impute 🤔 Consider just

ImputeInputDatapoint = ForecastInputDatapoint

Renamed to ForecastInputDatapoint / ForecastInputTimeSeries in f71110f. I skipped the Impute alias: impute takes the same request body, and two names for one class would make the docs and isinstance checks confusing. The docstrings say the classes are used for both.

haakonvt · 2026-11-04T10:29:40Z

+    Args:
+        timestamp (int | float | str | datetime): Milliseconds since epoch, a datetime, or a time-shift string like "2d-ago".
+        value (float | None): The observed value, or None when there is no reading.
+        missing (bool): Hide this point from the model. For impute, the missing points are the ones reconstructed.


Value being None and/or missing being True: this needs more text/explanations imo

Documented in b675bee.

haakonvt · 2026-11-04T10:31:22Z

+
+    Args:
+        label (str): Caller-chosen identifier, echoed back in the result.
+        datapoints (Sequence[InputDatapoint | tuple[int | float | str | datetime, float | None] | dict[str, Any]]):


Let's go with InputDatapoint only. InputDatapoint can then have a load method that is very accepting of different formats.

Done in d04aa67: datapoints takes ForecastInputDatapoint only.

haakonvt · 2026-11-04T10:31:48Z

+    def dump(self, camel_case: bool = True) -> dict[str, Any]:
+        output: dict[str, Any] = {
+            "label": self.label,
+            "datapoints": [_to_input_datapoint(dp).dump(camel_case) for dp in self.datapoints],


_to_input_datapoint should be a load method on the class

Added ForecastInputDatapoint.load in d04aa67. It accepts tuples, lists and API-style dicts.

haakonvt · 2026-11-04T10:32:33Z

+        cohort (str | None): Series sharing a cohort are forecast jointly and must have identical timestamps.
+            Omit it to forecast the series independently.


This needs to explain how the user uses the setting

Expanded in 5856c41. forecast() in #2875 has an example with a cohort.

haakonvt · 2026-11-04T10:33:54Z

+        # Only numeric keys are quantile levels; ignore anything else, like unknown fields elsewhere in the SDK.
+        quantiles = {float(level): value for level, value in resource["quantiles"].items() if _is_number(level)}


Are these fields really unknown? Is there a pattern for example, like key starting with quantile or a number or something? We can not decide this based on "is number" imo

No: the API keys quantiles by level only ("0.50", or "0.025" when more decimals are needed). Every key is now parsed, 6216571.

haakonvt · 2026-11-04T10:36:36Z

+        metadata = response.get("metadata", {})
+        return cls(
+            [ForecastResult._load(item) for item in response["timeSeries"]],
+            quantile_levels=[float(q) for q in metadata.get("quantileLevels", [])],


Why do you need to call float? Are you passing number as strings in the payload?

No need, the levels are numbers. Removed in 2d5e38a.

haakonvt · 2026-11-04T10:38:15Z

+    def get(
+        self,
+        id: int | None = None,
+        external_id: str | None = None,
+        instance_id: InstanceId | tuple[str, str] | None = None,
+        *,
+        label: str | None = None,
+    ) -> ForecastResult | None:


ForecastResult has no identifier, only label. Thus I think you should

Suggested change

def get(

self,

id: int | None = None,

external_id: str | None = None,

instance_id: InstanceId | tuple[str, str] | None = None,

*,

label: str | None = None,

) -> ForecastResult | None:

def get(

self,

label: str | None = None,

) -> ForecastResult | None: # with a type ignore-override-something

Done in 6a5bbcd.

haakonvt · 2026-11-04T10:38:55Z

+        Returns:
+            pd.DataFrame: The forecasts.
+        """
+        return _concat_columns({r.label: _points_to_pandas(r.forecast, self.quantile_levels) for r in self.data})


This does not handle repeat labels, either verify or handle duplicates like the datapoints retrieve methods do

Verified: forecast() and impute() reject duplicate labels (#2875). Noted in the get docstring, 6a5bbcd.

haakonvt · 2026-11-04T10:39:14Z

+            cohorts=metadata.get("cohorts", []),
+        )
+
+    def get(


Same comment here, lets go label only

Done in 6a5bbcd.

The forecast and impute endpoints return their series under "items", like other CDF list endpoints, rather than "timeSeries". Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

…ForecastInputTimeSeries InputDatapoint and InputTimeSeries were too generic for data_classes.ai, which also holds the document AI classes. The forecast and impute endpoints take the same request body, so both use these classes. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

… a load method ForecastInputTimeSeries.datapoints accepted datapoints, tuples and dicts, which widened its type for little gain. It now takes ForecastInputDatapoint only, and ForecastInputDatapoint.load creates one from a tuple, a list or an API-style dict. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

The API keys quantiles by their level only ("0.50", or more decimals when needed, like "0.025"), so there are no other keys to skip. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

metadata, with quantileLevels and cohorts, is always in the response, and the levels are numbers already, so there is nothing to default or convert. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

Results have no id, external id or instance id, so get takes just the label. Labels are unique, since duplicate labels in a request are rejected. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

test_handle_unknown_arguments_when_loading adds a made-up key to every dict in a resource. Every key in a quantiles dict must be a quantile level, like "0.50", so loading fails on the made-up key. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

miladh · 2026-11-04T20:07:44Z

Thanks @haakonvt for the useful feedback! It's addressed in 8772cc0e..cb8eba8b, with a reply on each thread. One question is still open, on the _is_number thread.

The range also includes a recent API change: the array of series is now called items instead of timeSeries (a3ec603). This PR only covers the response side; the request side follows in #2875.

miladh mentioned this pull request Oct 5, 2026

feat(ai): add request splitting for time series forecasting #2874

Open

3 tasks done

niekbeckers reviewed Oct 5, 2026

View reviewed changes

miladh force-pushed the ai-ts-forecasting/01-data-classes branch from ea1105b to 12e6a69 Compare November 2, 2026 11:15

miladh marked this pull request as ready for review November 2, 2026 11:20

miladh requested review from a team as code owners November 2, 2026 11:20

gemini-code-assist Bot reviewed Nov 2, 2026

View reviewed changes

Comment thread cognite/client/data_classes/ai.py Resolved

Comment thread cognite/client/data_classes/ai.py Resolved

Comment thread cognite/client/data_classes/ai.py Resolved

Comment thread cognite/client/data_classes/ai.py Resolved

haakonvt reviewed Nov 4, 2026

View reviewed changes

miladh and others added 9 commits November 4, 2026 19:12

fix(ai): read forecast and impute results from the items array

a3ec603

The forecast and impute endpoints return their series under "items", like other CDF list endpoints, rather than "timeSeries". Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

docs(ai): explain how value and missing affect what the model sees

b675bee

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

docs(ai): explain how to use cohort

5856c41

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

refactor(ai): parse every quantile key as a level

6216571

The API keys quantiles by their level only ("0.50", or more decimals when needed, like "0.025"), so there are no other keys to skip. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

refactor(ai): read the response metadata as the API returns it

2d5e38a

metadata, with quantileLevels and cohorts, is always in the response, and the levels are numbers already, so there is nothing to default or convert. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

refactor(ai): look up forecast and impute results by label only

6a5bbcd

Results have no id, external id or instance id, so get takes just the label. Labels are unique, since duplicate labels in a request are rejected. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

		cohort (str \| None): Series sharing a cohort are forecast jointly and must have identical timestamps.
		Omit it to forecast the series independently.

		# Only numeric keys are quantile levels; ignore anything else, like unknown fields elsewhere in the SDK.
		quantiles = {float(level): value for level, value in resource["quantiles"].items() if _is_number(level)}

Conversation

miladh commented Oct 5, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Where the stack is going

Checklist:

Uh oh!

niekbeckers left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

miladh commented Nov 2, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

codecov Bot commented Nov 2, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

miladh commented Nov 4, 2026

Uh oh!

This branch has not been deployed

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

miladh commented Oct 5, 2026 •

edited

Loading

niekbeckers left a comment •

edited

Loading

miladh commented Nov 2, 2026 •

edited

Loading

codecov Bot commented Nov 2, 2026 •

edited

Loading