sientia-dataops-laborious-legacy/laborious/activities/prediction_metadata.py

from temporalio import activity, workflow

with workflow.unsafe.imports_passed_through():
    from typing import Any

    from bson import ObjectId
    from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
    from sientia_do.observability.logger import Logger
    from sientia_do.observability.metrics_controller import MetricsController
    from sientia_do.observability.sientia_monitoring import SientiaMonitoring
    from sientia_do.repository.mongodb_repository_sync import MongoDBRepository


class PredictionMetadata(SientiaMonitoring):
    """
    Centralized, best-effort writer for prediction gate metadata (triggered filter, policy,
    counter) to MongoDB.

    Gates only compute and return a metadata document; they never write it themselves. Keeping
    the write here, instead of in each gate, puts the one real write point in the subworkflow
    that's about to persist a `predictions` row — the only place that reliably knows whether one
    is being written at all — and lets other metadata sources reuse the same activity instead of
    duplicating write logic wherever they originate.
    """

    mongo_repository: MongoDBRepository | None = None
    _PREDICTION_METADATA_COLLECTION = 'prediction_metadata'

    def __init__(
        self,
        mongo_repository: MongoDBRepository | None = None,
        logger: Logger | None = None,
        notification_handler: NotificationHandler | None = None,
        metrics_controller: MetricsController | None = None,
    ):
        """
        Initialize the prediction metadata writer.

        Args:
            mongo_repository: MongoDB, where gate metadata (triggered filter, policy, counter) is
                written best-effort. `None` is a valid state (see `write_prediction_metadata`).

        Raises:
            Exception: If BaseActivity initialization fails
        """
        self.mongo_repository = mongo_repository
        SientiaMonitoring.__init__(
            self,
            logger=logger,
            notification_handler=notification_handler,
            metrics_controller=metrics_controller,
        )

    def close(self) -> None:
        """
        Close the prediction metadata writer and clean up resources.
        """
        if self.mongo_repository is not None:
            try:
                self.mongo_repository.close()
            finally:
                self.mongo_repository = None
        SientiaMonitoring.shutdown(self)

    def __del__(self):
        self.close()

    @activity.defn(name='write_prediction_metadata')
    def write_prediction_metadata(self, input_data: dict[str, Any]) -> str | None:
        """
        Best-effort write of one or more gate metadata documents to Mongo, combined into a single
        `prediction_metadata` document.

        Implements both business rules for this feature in one place:
        - The prediction write must never be blocked or failed by a metadata write. A failed
          `insert(...)` is swallowed, never re-raised, logged here with call-site context (which
          `mongo_repository.insert(...)` itself can't know) — matching how `storage.py`'s
          `cleanup_minio_objects_expired` also logs again at the call site even though the
          repository already logs internally and emits its own Prometheus metric
          (`MONGODB_WRITE_ERROR_COUNT`); the generic library log is for the metric/trace, this
          one is for whoever is reading logs for this specific call.
        - Nothing is written when there's nothing relevant to report. An empty/missing
          `documents` list (or one full of falsy/empty entries) skips Mongo entirely, so callers
          can call this unconditionally and get `None` back for "no metadata this time" instead
          of deciding upfront whether there's something worth saving.

        Takes a *list* of documents, not a single one: each caller contributes its own document
        (shaped like `{'filters': [...]}`), and this flattens every `filters` entry across all of
        them into one `prediction_metadata` document before the single insert. Today only one
        gate ever contributes a document per prediction cycle (the gate chain short-circuits on
        the first non-`None` `path_flag`), but the list shape means this activity won't need to
        change again once more than one source contributes metadata for the same cycle.

        `insert(...)` doesn't hand back the inserted document's `_id`, so it is generated here,
        client-side, before the insert, as a plain string (not a BSON ObjectId) — matching this
        codebase's convention of string `_id`s (e.g. the `models` collection) and the 24-hex-char
        size of the `predictions.metadata_id` Postgres column that will reference it.

        Args:
            input_data: Dict with:
                - documents (list[dict[str, Any]]): Metadata payloads to combine, each shaped
                  like `{'filters': [...]}`. Missing/empty, or entries with no `filters`, skip
                  the write entirely.
                - metadata (dict): Workflow metadata for contextual logging.

        Return:
            The generated id, or `None` when there's nothing to write, Mongo isn't configured, or
            the insert failed.
        """
        metadata = input_data['metadata']
        documents = input_data.get('documents') or []

        filters = [
            entry for document in documents if document for entry in document.get('filters', [])
        ]

        if not filters or self.mongo_repository is None:
            return None

        generated_id = str(ObjectId())

        try:
            self.mongo_repository.insert(
                self._PREDICTION_METADATA_COLLECTION,
                {'filters': filters, '_id': generated_id},
                metadata,
            )
        except Exception as e:
            self.error(f'Failed to write prediction metadata: {e}', metadata)
            return None

        return generated_id