updated logging and added convenience functions

src/database/utils.py

@@ -1,8 +1,10 @@
 from pymongo.collection import Collection
+import logging
 
 from .classes import (
     VisualCommunication,
-    NoDocumentFoundException
+    NoDocumentFoundException,
+    ModelOutputs
 )
 
 
@@ -25,8 +27,7 @@ def total_annotated(
 
 def get_visual_communication(
         collection: Collection,
-        with_annotation: bool = False,
+        with_annotation: bool = False
-        with_prediction: bool = False
     ) -> VisualCommunication:
     """Get a random visual communication from the database."""
     query = {}
@@ -34,15 +35,56 @@ def get_visual_communication(
         query["annotation"] = {"$ne": None}
     else:
         query["annotation"] = None
-    if with_prediction:
-        query["prediction"] = {"$ne": None}
-    else:
-        query["prediction"] = None
     data = collection.aggregate([
         { "$match": query },   # find using filters
         { "$sample": { "size": 1 } }    # get one random
     ])
     data = list(data)   # read data from cursor object
     if len(data) == 0:
-        raise NoDocumentFoundException
+        logging.error("failed getting visual communication")
-    return VisualCommunication.model_validate(data[0])
+        raise NoDocumentFoundException()
+    data = data[0]
+    logging.info("finished")
+    return VisualCommunication.model_validate(data)
+
+
+def upsert_predictions(
+    collection: Collection,
+    vis_com_name: str,
+    predictions: ModelOutputs,
+) -> None:
+    """Upsert prediction data in the database."""
+    query = {
+        "name": vis_com_name
+    }
+    update = {
+        "$set": { "prediction": predictions.model_dump() }
+    }
+    res = collection.update_one(
+        filter=query,
+        update=update,
+        upsert=True,
+    )
+    logging.debug("upserted document: %s", res)
+    logging.info("finished")
+
+
+def upsert_annotations(
+    collection: Collection,
+    vis_com_name: str,
+    annotations: ModelOutputs,
+) -> None:
+    """Upserts annotation data in the database."""
+    query = {
+        "name": vis_com_name
+    }
+    update = {
+        "annotation": annotations.model_dump()
+    }
+    res = collection.update_one(
+        filter=query,
+        update=update,
+        upsert=True,
+    )
+    logging.info("upserted document: %s", res)
+    logging.info("finished")