pymongopredict.predictions_api

The predictions API: a $predict aggregation stage backed by prediction models stored in the database.

You define a model once on a collection, as you would a vector search index, then call it by name as the first stage of an aggregation pipeline; later stages read its output through $meta. No server changes are needed.

$vectorSearch $predict
defined with create_search_index create_prediction_model
managed with list_ / update_ / drop_search_index list_ / update_ / drop_prediction_model
stage options index, path, queryVector, filter, limit model, filter, limit, batchSize
position first stage only first stage only
output {"$meta": "vectorSearchScore"} {"$meta": "prediction"}, "predictScore", "predictProbabilities"
runs inside Atlas in your Python program, then on the server

What you use:

Without enable(), call aggregate(coll, pipeline) and the module-level model functions with the collection as their first argument.

Model definition

The definition given to PredictionModel takes:

  • target: the dotted path of the field to predict (required).
  • columns, embedFields, excludeFields, embeddingModel, embeddingKwargs, maskTarget, maxTokens, maxJsonChars: the settings of the pymongopredict.embeddings_api.DocumentFeaturizer; featurizerOptions passes any other of its arguments by its Python name. Without embedFields the whole document is embedded.
  • embeddingModel: "voyage-4-nano" (the default, runs locally), any Hugging Face model id (runs locally), or a Voyage API model such as "voyage-4-lite", "voyage-4" or "voyage-4-large" (needs VOYAGE_API_KEY); embeddingKwargs passes its settings, e.g. {"output_dimension": 512} for an API model (but not api_key, which would be stored with the model: set VOYAGE_API_KEY instead). See the embeddings API documentation for every option.
  • task: "auto", "classification" or "regression". "auto" predicts a number when every value of the target is a number (numbers stored as text count; empty strings and text such as "NA", "none" or "-" are missing values, while in a field of categories they are labels) and there are at least 8 distinct values, and a category otherwise, so a 0/1 flag or a 1 to 5 star rating stored as numbers is a category.
  • estimator: "auto" (the default: XGBoost with its default settings, for categories and numbers, on at most 16 threads unless OMP_NUM_THREADS is set), "xgboost" (the same), "logistic", "ridge" or "hgb" (histogram gradient boosting), or any scikit-learn estimator passed to create_prediction_model().
  • filter: which documents to learn from (documents without the target are always left out); numTrainingDocuments: a cap on their number.
  • testSize: the fraction, or the number, of training documents held out to measure the stored metrics (default 0.2; 0 skips them). The stored model is then refitted on all training documents.

{"$meta": "predictScore"} is the probability of the predicted category and {"$meta": "predictProbabilities"} maps each category to its probability; both exist only when the estimator has predict_proba (not for numbers or "ridge").

How a pipeline runs

  1. create_prediction_model() fits the featurizer and the estimator on the documents that have the target, records the held-out metrics, refits on all of them and stores the fitted model (a pickle of at most 15 MB) with its status and version in the __mongopredict_models collection of the same database. As with a search index, the previous version keeps answering while a new one trains.
  2. aggregate() fetches the documents selected by the stage's filter and limit (only the fields the model reads), predicts them in batches in your program, and uploads only {_id, prediction, score, probabilities} rows to a temporary collection.
  3. The rest of the pipeline runs on the server: the temporary collection is joined back to the full documents ($lookup on _id), the $meta references are rewritten to the joined fields, and the helper field is removed before the output (or before a final $merge / $out). The temporary collection is then dropped.

On the 14 MongoDB sample tasks, $predict returns the same predictions as calling the model directly, at a median 1.39 times the time of that code (reports/01_predictions_api.html).

Details and limits

  • Requirements: MongoDB 5.0 or later (5.1 or later when the source collection is sharded); fitting a model or running $predict on an older server raises RuntimeError. The tests run on MongoDB 5.0, 6.0, 7.0 and 8.0, and on Atlas (8.0). On 5.0 a time-series collection is a view, so collation= cannot be passed for it there (6.0 and later accept it). $predict needs find on the source collection and createCollection, insert and dropCollection in its database (for the temporary __mongopredict_tmp_* collection); training needs find, insert and update on __mongopredict_models (remove to drop a model). Model records and temporary collections are read and written on the primary, even with a secondary read preference.
  • A killed process can leave a temporary collection behind: cleanup_temporary_collections() drops those older than 24 hours (older_than). A call whose temporary collection is dropped while it runs raises RuntimeError instead of returning partial results. A record still BUILDING after 24 hours is flagged stale: True by list_prediction_models().
  • Documents are predicted in batches of at most batchSize documents and 64 MB of BSON. Scores and probabilities are computed only when the pipeline reads predictScore or predictProbabilities. An embedding cache directory that cannot be used gives a warning, and the embeddings are then computed without a cache.
  • The source can be a collection, a view or a time-series collection; views and time-series collections are joined from the source side, through the temporary collection's _id index. _id values must be unique (a view with an $unwind stage is refused).
  • A class is returned as first stored in the training documents (an ObjectId, a date, an int), so predictions can be compared with the labels or joined to them with $lookup. The predictProbabilities keys and the record's classes are str() of those values, so str(prediction) is always a key.
  • Model names may not contain . or $. limit and batchSize must be positive integers; without limit every matching document is predicted.
  • hint=, collation=, let= and session= passed to the aggregate apply to the query on the source collection, and the temporary collection gets the source collection's default collation, so string matching and sorting behave the same.
  • $meta works in $sort, sortBy ($setWindowFields, $top, $bottom), expressions, $group, $facet, $redact, and after an inclusion $project or a $replaceRoot. It is not available inside $lookup / $unionWith sub-pipelines (add it as a field first and pass it with let), nor in $merge's let or whenMatched. __mongopredict is reserved as a field name.
  • $predict cannot run inside a multi-document transaction, because it creates and drops a temporary collection.
  • The documents travel to your program and the model runs there; the result is read in full before it is returned; $predict works from Python only (not from mongosh or other drivers).
  • The stored model is a pickle, so only users trusted to run code in your program should be able to write __mongopredict_models. Set MONGOPREDICT_MODEL_KEY (the same secret where models are trained and where they are used) to sign stored models: a model without a valid signature is then never unpickled.
  • The record holds the library versions the model was stored with (versions). Loading it under other scikit-learn, numpy or pymongopredict versions gives a warning; a model that cannot be loaded raises a ValueError that says to rebuild it with update_prediction_model().
  • $predict caches embeddings only when MONGOPREDICT_PREDICT_CACHE_DIR names a directory (the training cache_dir is not used), so a long-running process does not grow a cache of every document it predicts.
  • Training on a collection streams the documents: only _id and the target of each are held in memory (if documents are deleted meanwhile, it trains again from a single read). Views and time-series collections are read once into memory, and training on them also needs unique _id values.
Examples

Usage:

from pymongopredict import PredictionModel, predictions_api

predictions_api.enable()                       # coll.aggregate() understands $predict
coll.create_prediction_model(PredictionModel({"target": "room_type"}, name="room_type"))
coll.list_prediction_models()                  # status, definition, holdout metrics
coll.aggregate([{"$predict": {"model": "room_type", "limit": 1000}},
                {"$project": {"p": {"$meta": "prediction"}}}])
class PredictionModel:

A prediction model definition, the counterpart of pymongo.operations.SearchIndexModel.

Only the arguments are checked here; the values are checked when the model is trained.

Parameters

  • definition (dict): How to train the model. All keys are optional except target:

    • target: dotted path of the field to predict.
    • columns, embedFields, excludeFields, embeddingModel, embeddingKwargs, maskTarget, maxTokens, maxJsonChars: passed to pymongopredict.DocumentFeaturizer as columns, embed_fields, exclude, embedding_model, embedding_kwargs, mask_target, max_tokens and max_json_chars (embedFields omitted means embed the whole document). featurizerOptions is a dict of any other DocumentFeaturizer argument by its Python name; it overrides the keys above.
    • task: "classification", "regression" or "auto" (default). "auto" picks regression when every target is a number (or numeric text; empty strings are dropped) and there are at least MIN_REGRESSION_VALUES (8) distinct values; otherwise classification, so a 0/1 flag, 1-5 stars, booleans and text are categories.
    • estimator: "auto" (default) or "xgboost": XGBoost with its default settings (_XGBLabelClassifier for classification, xgboost.XGBRegressor for regression); "logistic" (classification only), "ridge" (RidgeClassifierCV / RidgeCV), "hgb" (histogram gradient boosting), or an sklearn estimator object (usually passed as the estimator argument of create_prediction_model(); it is deep-copied before fitting).
    • filter: query selecting the training documents. Documents without the target are always left out, as are documents whose target is null (and, for regression, not a number).
    • numTrainingDocuments: positive integer cap on training documents (default: all).
    • testSize: holdout fraction in [0, 1), or integer count of documents, used for the stored metrics (default 0.2; 0 skips them). The split is stratified for classification when every class has at least 2 documents. The stored model is refitted on all training documents.
  • name (str, default "default"): Model name, unique per collection. Must be a non-empty string without . or $.
  • type (str, default "prediction"): Model type; only "prediction" is supported.
Raises
  • ValueError: If name is invalid, type is not "prediction", target is missing, or definition has a key outside DEFINITION_KEYS.
Examples

Usage:

PredictionModel({"target": "imdb.rating", "embedFields": ["plot", "genres"]},
                name="imdb_rating")
PredictionModel(definition: dict, name: str = 'default', type: str = 'prediction')
document: dict

The model as a {"name", "type", "definition"} dict (SearchIndexModel.document).

create_prediction_model() accepts this dict in place of the object.

def create_prediction_model( collection, model: PredictionModel | dict, estimator=None, progress: bool | None = None) -> str:

Train and store a prediction model; the counterpart of Collection.create_search_index.

Training runs synchronously in this process. If a model with this name already exists on the collection, it is rebuilt (as update_prediction_model would): as with a search index rebuild, the previous version stays stored and queryable until the new one is written.

The record in MODELS_COLLECTION (_id "<collection>.<name>") holds:

  • collection, name, type and latestDefinition (the definition of the latest build, with values BSON cannot store, such as estimator objects, as their repr).
  • status: "BUILDING" while a build runs, then "READY", or "FAILED" with message set to the error (an interrupted build, e.g. KeyboardInterrupt, is marked FAILED too). Only a killed process (SIGKILL, out of memory, a lost node) leaves a record BUILDING; buildStartedAt records when the build started and list_prediction_models() flags such builds as stale after STALE_BUILD.
  • queryable: False while the first build runs; True once a version is stored. A failed rebuild keeps the previous version and its queryable value.
  • version: increments with each build started (taken atomically, so concurrent builds get distinct versions); buildId: unique id of the stored build.
  • After a successful build: task, classes (empty for regression), numTrainingDocuments, metrics (r2 / mae, or accuracy / auroc, plus numTestDocuments; empty when testSize is 0), modelBytes, updatedAt (UTC) and model (the pickled featurizer and estimator).
  • Internal bookkeeping: nextVersion and buildingId (the build allowed to write).

Only the latest build started may write its result. If a newer build started meanwhile, or the model was dropped, this build's result is discarded with a warning.

Parameters
  • collection (pymongo.collection.Collection): Collection holding the training documents; the model is registered on it.
  • model (PredictionModel or dict): The model, or its {"name", "type", "definition"} document (name defaults to "default", type to "prediction"). See PredictionModel for the definition keys.
  • estimator (sklearn estimator, optional): Overrides the definition's estimator with any sklearn estimator object.
  • progress (bool or None, default None): Show a progress bar for each training step: embedding the training part, fitting on it, embedding and scoring the holdout part, embedding all documents (mostly from the cache) and fitting the final model (only these last two when testSize is 0). Embedding steps count documents. None shows the bars on a terminal or in a notebook and hides them when stderr is redirected; True / False force them on or off.
Returns
  • str: The model name.
Raises
  • ValueError: If the definition is invalid (see PredictionModel), there are fewer than 2 training documents, the estimator name does not fit the task, the pickled model is over 15 MB, or the record is over MongoDB's 16 MB document limit. The record is marked FAILED before any error from training is re-raised.
Warns
  • UserWarning: If the build was superseded by a newer one (or the model dropped) and was not stored.
Examples

Usage:

create_prediction_model(coll, PredictionModel({"target": "room_type"}, "room_type"))
create_prediction_model(coll, {"name": "price", "definition": {"target": "price"}},
                        estimator=HistGradientBoostingRegressor())
def update_prediction_model( collection, name: str, definition: dict, estimator=None, progress: bool | None = None) -> str:

Retrain an existing model with a new definition (Collection.update_search_index).

The previous version stays queryable while the new one trains, as described in create_prediction_model().

Parameters
  • collection (pymongo.collection.Collection): Collection the model is registered on.
  • name (str): Name of the existing model.
  • definition (dict): The new definition; it replaces the old one entirely (see PredictionModel).
  • estimator (sklearn estimator, optional): Overrides the definition's estimator.
  • progress (bool or None, default None): Progress bars for the training steps, as in create_prediction_model().
Returns
  • str: The model name.
Raises
def list_prediction_models(collection, name: str | None = None) -> list[dict]:

Status, definition and metrics of the models on a collection (list_search_indexes).

Parameters
  • collection (pymongo.collection.Collection): Collection the models are registered on.
  • name (str, optional): Return only this model (an empty list if it does not exist).
Returns
  • list of dict: One record per model, as described in create_prediction_model(), without the pickled model and the internal nextVersion / buildingId fields. A record still BUILDING after STALE_BUILD gets stale: True: its build was most likely killed (the process cannot record that); create the model again.
def drop_prediction_model(collection, name: str) -> None:

Delete a stored model and its cached copy (Collection.drop_search_index).

Dropping a model that does not exist is a no-op. A build still running for it is not stored when it finishes.

Parameters
  • collection (pymongo.collection.Collection): Collection the model is registered on.
  • name (str): Model name.
def aggregate(collection, pipeline: list[dict], _aggregate=None, **kwargs):

collection.aggregate that understands a leading $predict stage.

The stage is {"$predict": {"model": <name>, "filter": <query>, "limit": <n>, "batchSize": <n>}}:

  • model (required): name of a stored prediction model of collection with a queryable version (see pymongopredict.predictions_api.load_prediction_model()).
  • filter: any query selecting the documents to predict (default: all).
  • limit: positive integer, the most documents to predict (default: every match).
  • batchSize: positive integer, how many documents are featurized and uploaded at a time (default 512).

Later stages read the predictions with {"$meta": "prediction"}, {"$meta": "predictScore"} and {"$meta": "predictProbabilities"}; see the module docstring for how the pipeline runs. Pipelines without $predict go straight to the server unchanged.

Parameters
  • collection (pymongo.collection.Collection): Source collection.
  • pipeline (list of dict): Aggregation pipeline; $predict, if present, must be its first stage and nowhere else.
  • _aggregate (callable, optional): Internal: (collection, pipeline, **kwargs) -> cursor that runs a plain pipeline; enable() passes the unpatched Collection.aggregate here.

  • **kwargs: aggregate options, handled as on the server where possible:

    • session: used for every operation (the query, creating, filling and dropping the temporary collection, the server pipeline). A session in a transaction, or a snapshot session, is refused.
    • hint: applies only to the query on the source collection.
    • collation: applies to the source query and to the user's stages. The _id join then runs separately with the source's default collation (into a second temporary collection), so a case-insensitive collation cannot join "a" to "A". Without it, the temporary collection gets the source's default collation, so string matching and sorting behave the same as on the source.
    • comment, maxTimeMS: passed to the source query and to the server pipeline (comment is dropped from the query on backends that do not support it).
    • let: passed to the server pipeline, and to the source query when the stage has a filter (for $$ variables in its $expr).
    • anything else: passed to the server pipeline.
Returns
  • PredictCursor or cursor: A PredictCursor for a $predict pipeline; otherwise whatever collection.aggregate returns.
Raises
  • ValueError: If $predict is not only the first stage, the stage has other fields beside it, model is missing, the spec has unknown keys, limit / batchSize are not positive integers, the model does not exist or is not queryable, the session is in a transaction or is a snapshot session, a document contains the reserved META_FIELD, a view or time-series source returns an _id twice, or $meta is used where it is not available ($lookup / $unionWith sub-pipelines, $merge's let / whenMatched).
  • RuntimeError: If the temporary collection lost rows while the call ran (dropped by cleanup_temporary_collections() with a short older_than).
Examples

Usage:

aggregate(coll, [
    {"$predict": {"model": "imdb_rating", "filter": {"year": {"$gte": 2010}}}},
    {"$sort": {"p": {"$meta": "prediction"}}},   # a $meta sort key sorts descending
    {"$limit": 10},
], maxTimeMS=60_000).to_list()
def enable(*classes) -> None:

Let Collection.aggregate run $predict and add the model helpers as methods.

After this, coll.aggregate(pipeline, ...) calls aggregate() (pymongo's positional session, let and comment arguments are still accepted), and coll.create_prediction_model, update_prediction_model, list_prediction_models and drop_prediction_model call the functions of the same name with coll first. Calling it again for a class already patched does nothing.

Parameters
  • *classes (type): Collection classes to patch (e.g. mongomock.Collection). Default: pymongo.collection.Collection.
Examples

Usage:

predictions_api.enable()
coll.create_prediction_model(PredictionModel({"target": "room_type"}, "room_type"))
coll.aggregate([{"$predict": {"model": "room_type"}}]).to_list()
predictions_api.disable()
def disable() -> None:

Undo enable() for every patched class, restoring their original attributes.

Attributes a class inherited before enable() are deleted from it again, so the inherited ones apply. Does nothing if nothing is patched.

class PredictCursor:

The materialised result of a $predict pipeline.

Behaves like a pymongo CommandCursor over a list: iterate it, or call next(), try_next(), to_list(), close(), check alive, or use it as a context manager.

Parameters
  • docs (list of dict): The pipeline output.
  • stats (dict, optional): Timing statistics (default: empty).
Attributes
  • stats (dict): nDocuments: number of documents predicted (after removing duplicates); fetchSeconds: seconds spent fetching them; predictSeconds: featurizing and predicting on the client; uploadSeconds: uploading the predictions; serverSeconds: running the rest of the pipeline on the server.
PredictCursor(docs: list[dict], stats: dict | None = None)
def next(self):

The next document; raises StopIteration when none are left.

def try_next(self) -> dict | None:

The next document, or None when none are left.

alive: bool

Whether documents are left (like pymongo's cursors).

def to_list(self, length: int | None = None) -> list[dict]:

The remaining documents as a list (like pymongo's Cursor.to_list).

Parameters
  • length (int, optional): Return at most this many documents (the rest stay in the cursor). Default: all.
Raises
  • ValueError: If length is less than 1.
def close(self) -> None:

Discard the remaining documents.

def cleanup_temporary_collections( database, older_than: datetime.timedelta = datetime.timedelta(days=1)):

Drop the temporary collections that $predict calls left behind (a killed process, or a user without the dropCollection privilege). Returns the dropped names.

Collections created less than older_than ago may belong to a running call and are kept; names without a creation time (older versions) count as old. A call whose collection is dropped while it runs raises RuntimeError rather than returning partial results, so keep older_than above the longest $predict call.