[miladh]

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:

• forecast / impute: a low-level client that mirrors the API. PRs 1–4.
• forecast_with_retrieval / impute_with_retrieval: convenience methods for CDF time series by instance id. PRs 5–8.

This PR only adds data classes, shaped like the API's request and response, with no API calls yet:

• ForecastInputDatapoint and ForecastInputTimeSeries: the request objects, used by both endpoints. They're dataclasses, like DatapointsQuery. ForecastInputDatapoint.load creates a datapoint from a tuple, a list or an API-style dict.
• QuantileDatapoint: one predicted point, quantile level -> value.
• ForecastResult(List) and ImputeResult(List): one result per input series. The lists hold the response's metadata (quantile levels and cohorts) as ForecastMetadata, following SDK precedent: _load_raw_api_response and dump_raw like NodeList, and a non-optional field with an empty default for lists built from results alone, like SyncRecordList.
• to_pandas() gives one column per series with the median; include_all_quantiles=True gives (label, quantile level) columns for every level.

Usage (this PR on its own)

from cognite.client.data_classes.ai import ForecastInputDatapoint, ForecastInputTimeSeries, ForecastResultList

series = ForecastInputTimeSeries(
    "23-PT-1101",
    [
        ForecastInputDatapoint(1790812800000, 42.1),
        ForecastInputDatapoint(1790812860000, None),                 # no reading: hidden from the model
        ForecastInputDatapoint(1790812920000, 42.9, missing=True),   # hidden; impute returns a prediction for it
        ForecastInputDatapoint.load((1790812980000, 43.1)),          # from a tuple
    ],
    cohort="compression-train-a",
)
series.dump()   # the API request shape: {"label": ..., "cohort": ..., "datapoints": [...]}

res = ForecastResultList._load_raw_api_response([api_response])   # done by the API in PR 3
res.metadata.quantile_levels                   # as the API returns them
res.metadata.cohorts                           # ["compression-train-a"]
res.get("23-PT-1101").forecast[0].quantiles    # {0.05: 41.0, ..., 0.95: 43.0}
res.to_pandas()                                # DatetimeIndex, one median column per series
res.to_pandas(include_all_quantiles=True)      # (label, quantile level) columns
res.dump_raw()                                 # {"items": [...], "metadata": {...}}, like the API response

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?

[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.

[miladh]

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

[codecov]

Codecov Report

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

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

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #2873      +/-   ##
==========================================
+ Coverage   93.49%   93.72%   +0.22%     
==========================================
  Files         527      541      +14     
  Lines       55489    57533    +2044     
==========================================
+ Hits        51879    53920    +2041     
- Misses       3610     3613       +3     

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.65% <97.82%> (+0.28%)	⬆️

... and 109 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.

[miladh]

Thanks @haakonvt for the useful feedback! It's addressed in 8772cc0e..9ab0e18f, 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.

It also adds a quantile_levels argument to to_pandas, which now includes only the median (0.5) by default, so DataFrames and the notebook view stay narrow (9ab0e18).

[haakonvt]

Some more comments 🥳

[haakonvt]

LGTM, a few more nits, but nothing blocking as we are "a long stack away" from making these endpoints public 😅

[haakonvt]

🦄 one nit you can put in a later pr

[haakonvt]

You can use from cognite.client.utils._auxiliary import remove_duplicates_keep_order

[miladh]

Added it to #2875 in d55bb65.