Document and reformat

great_ai/persistence/__init__.py

@@ -1,3 +0,0 @@
-from .mongodb_driver import MongodbDriver
-from .parallel_tinydb_driver import ParallelTinyDbDriver
-from .tracing_database_driver import TracingDatabaseDriver

great_ai/persistence/mongodb_driver.py

@@ -18,7 +18,16 @@ operator_mapping = {
 }
 
 
-class MongodbDriver(TracingDatabaseDriver):
+class MongoDbDriver(TracingDatabaseDriver):
+    """TracingDatabaseDriver implementation using MongoDB as a backend.
+
+    A production-ready database driver suitable for efficiently handling semi-structured
+    data.
+
+    Checkout [MongoDB Atlas](https://www.mongodb.com/cloud/atlas/register) for a hosted
+    MongoDB solution.
+    """
+
     is_production_ready = True
 
     mongo_connection_string: str
@@ -28,7 +37,8 @@ class MongodbDriver(TracingDatabaseDriver):
         super().__init__()
         if self.mongo_connection_string is None or self.mongo_database is None:
             raise ValueError(
-                "Please configure the MongoDB access options by calling MongodbDriver.configure_credentials"
+                "Please configure the MongoDB access options by calling "
+                "MongoDbDriver.configure_credentials"
             )
 
         with MongoClient[Any](self.mongo_connection_string) as client:
@@ -44,6 +54,14 @@ class MongodbDriver(TracingDatabaseDriver):
         mongo_database: str,
         **_: Any,
     ) -> None:
+        """Configure the connection details for MongoDB.
+
+        Args:
+            mongo_connection_string: For example:
+                'mongodb://my_user:my_pass@localhost:27017'
+            mongo_database: Name of the database to use. If doesn't exist, it is
+                created and initialised.
+        """
         cls.mongo_connection_string = mongo_connection_string
         cls.mongo_database = mongo_database
         super().configure_credentials()

great_ai/persistence/parallel_tinydb_driver.py

@@ -17,6 +17,14 @@ operator_mapping = {"=": "eq", "!=": "ne", "<": "lt", "<=": "le", ">": "gt", ">=
 
 
 class ParallelTinyDbDriver(TracingDatabaseDriver):
+    """TracingDatabaseDriver with TinyDB as a backend.
+
+    Saves the database as a JSON into a single file. Highly inefficient on inserting,
+    not advised for production use.
+
+    A multiprocessing lock protects the database file to avoid parallelisation issues.
+    """
+
     is_production_ready = False
     path_to_db = Path(DEFAULT_TRACING_DB_FILENAME)
 

great_ai/persistence/tracing_database_driver.py

@@ -8,6 +8,8 @@ from ..views import Filter, SortBy, Trace
 
 
 class TracingDatabaseDriver(ABC):
+    """Interface expected from a database to be used for storing traces."""
+
     is_production_ready: bool
     initialized: bool = False