added database convenience functions

src/database/utils.py

@@ -0,0 +1,48 @@
+from pymongo.collection import Collection
+
+from .classes import (
+    VisualCommunication,
+    NoDocumentFoundException
+)
+
+
+def total_documents(
+    collection: Collection
+) -> int:
+    """Get total number of documents in database."""
+    return collection.count_documents(filter={})
+
+
+def total_annotated(
+    collection: Collection
+) -> int:
+    """Get total number of annotated documents in database."""
+    query = {
+        "annotation": { "$ne": None }
+    }
+    return collection.count_documents(filter=query)
+
+
+def get_visual_communication(
+        collection: Collection,
+        with_annotation: bool = False,
+        with_prediction: bool = False
+    ) -> VisualCommunication:
+    """Get a random visual communication from the database."""
+    query = {}
+    if with_annotation:
+        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
+    return VisualCommunication.model_validate(data[0])