Move files

great_ai/persistence/__init__.py

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

great_ai/persistence/mongodb_driver.py

@@ -0,0 +1,137 @@
+from datetime import datetime
+from typing import Any, List, Mapping, Optional, Sequence, Tuple
+
+from pymongo import MongoClient
+
+from ..views import Filter, SortBy, Trace
+from .tracing_database_driver import TracingDatabaseDriver
+
+operator_mapping = {
+    "=": "$eq",
+    "!=": "$ne",
+    "<": "$lt",
+    "<=": "$lte",
+    ">": "$gt",
+    ">=": "$gte",
+    "contains": "$regex",
+}
+
+
+class MongodbDriver(TracingDatabaseDriver):
+    is_production_ready = True
+
+    def __init__(self) -> None:
+        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"
+            )
+
+    @classmethod
+    def configure_credentials(  # type: ignore
+        cls,
+        *,
+        mongo_connection_string: str,
+        mongo_database: str,
+        **_: Mapping[str, Any],
+    ) -> None:
+        cls.mongo_connection_string = mongo_connection_string
+        cls.mongo_database = mongo_database
+        super().configure_credentials()
+
+    def save(self, trace: Trace) -> str:
+        serialized = trace.to_flat_dict()
+        serialized["_id"] = trace.trace_id
+
+        with MongoClient(self.mongo_connection_string) as client:
+            return client[self.mongo_database].traces.insert_one(serialized)
+
+    def save_batch(self, documents: List[Trace]) -> List[str]:
+        serialized = [d.to_flat_dict() for d in documents]
+        for s in serialized:
+            s["_id"] = s["trace_id"]
+
+        with MongoClient(self.mongo_connection_string) as client:
+            return client[self.mongo_database].traces.insert_many(
+                serialized, ordered=False
+            )
+
+    def get(self, id: str) -> Optional[Trace]:
+        with MongoClient(self.mongo_connection_string) as client:
+            value = client[self.mongo_database].traces.find_one(id)
+
+        if value:
+            value = Trace.parse_obj(value)
+
+        return value
+
+    def _get_operator(self, filter: Filter) -> str:
+        if filter.operator == "contains" and not isinstance(filter.value, str):
+            return operator_mapping["="]
+        return operator_mapping[filter.operator]
+
+    def query(
+        self,
+        *,
+        skip: int = 0,
+        take: Optional[int] = None,
+        conjunctive_filters: Sequence[Filter] = [],
+        conjunctive_tags: Sequence[str] = [],
+        since: Optional[datetime] = None,
+        until: Optional[datetime] = None,
+        has_feedback: Optional[bool] = None,
+        sort_by: Sequence[SortBy] = [],
+    ) -> Tuple[List[Trace], int]:
+
+        query = {
+            "filter": {
+                "$and": [{"tags": tag} for tag in conjunctive_tags]
+                + [
+                    {f.property: {self._get_operator(f): f.value}}
+                    for f in conjunctive_filters
+                ]
+                + [{}]
+            },
+            "sort": [
+                (col.column_id, 1 if col.direction == "asc" else -1) for col in sort_by
+            ],
+        }
+
+        if skip:
+            query["skip"] = skip
+
+        if take:
+            query["limit"] = take
+
+        if since:
+            query["filter"]["$and"].append({"created": {"$gte": since}})
+
+        if until:
+            query["filter"]["$and"].append({"created": {"$lte": until}})
+
+        if has_feedback is not None:
+            query["filter"]["$and"].append(
+                {"feedback": {"$ne": None}} if has_feedback else {"feedback": None}
+            )
+
+        with MongoClient(self.mongo_connection_string) as client:
+            values = client[self.mongo_database].traces.find(**query)
+            documents = [Trace.parse_obj(t) for t in values]
+
+        return documents, len(documents)
+
+    def update(self, id: str, new_version: Trace) -> None:
+        serialized = new_version.to_flat_dict()
+        serialized["_id"] = new_version.trace_id
+        with MongoClient(self.mongo_connection_string) as client:
+            client[self.mongo_database].traces.update_one(id, new_version)
+
+    def delete(self, id: str) -> None:
+        with MongoClient(self.mongo_connection_string) as client:
+            client[self.mongo_database].traces.delete_one(id)
+
+    def delete_batch(self, ids: List[str]) -> List[str]:
+        delete_filter = {"_id": {"$in": ids}}
+
+        with MongoClient(self.mongo_connection_string) as client:
+            return client[self.mongo_database].traces.delete_many(delete_filter)

great_ai/persistence/parallel_tinydb_driver.py

@@ -0,0 +1,112 @@
+from datetime import datetime
+from multiprocessing import Lock
+from pathlib import Path
+from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple, cast
+
+import pandas as pd
+from tinydb import TinyDB
+
+from ..views import Filter, SortBy, Trace
+from .tracing_database_driver import TracingDatabaseDriver
+
+DEFAULT_TRACING_DB_FILENAME = "tracing_database.json"
+lock = Lock()
+
+
+operator_mapping = {"=": "eq", "!=": "ne", "<": "lt", "<=": "le", ">": "gt", ">=": "ge"}
+
+
+class ParallelTinyDbDriver(TracingDatabaseDriver):
+    is_production_ready = False
+    path_to_db = Path(DEFAULT_TRACING_DB_FILENAME)
+
+    def save(self, trace: Trace) -> str:
+        return self._safe_execute(lambda db: db.insert(trace.dict()))
+
+    def save_batch(self, documents: List[Trace]) -> List[str]:
+        traces = [d.dict() for d in documents]
+        return self._safe_execute(lambda db: db.insert_multiple(traces))
+
+    def get(self, id: str) -> Optional[Trace]:
+        value = self._safe_execute(lambda db: db.get(lambda d: d["trace_id"] == id))
+        if value:
+            value = Trace.parse_obj(value)
+        return value
+
+    def query(
+        self,
+        *,
+        skip: int = 0,
+        take: Optional[int] = None,
+        conjunctive_filters: Sequence[Filter] = [],
+        conjunctive_tags: Sequence[str] = [],
+        since: Optional[datetime] = None,
+        until: Optional[datetime] = None,
+        has_feedback: Optional[bool] = None,
+        sort_by: Sequence[SortBy] = []
+    ) -> Tuple[List[Trace], int]:
+        def does_match(d: Dict[str, Any]) -> bool:
+            return (
+                not set(conjunctive_tags) - set(d["tags"])
+                and (
+                    since is None
+                    or cast(datetime, datetime.fromisoformat(d["created"])) >= since
+                )
+                and (
+                    until is None
+                    or cast(datetime, datetime.fromisoformat(d["created"])) <= until
+                )
+                and (
+                    has_feedback is None or has_feedback == (d["feedback"] is not None)
+                )
+            )
+
+        documents: List[Trace] = [
+            Trace.parse_obj(t)
+            for t in self._safe_execute(lambda db: db.search(does_match))
+        ]
+
+        if not documents:
+            return [], 0
+
+        df = pd.DataFrame([d.to_flat_dict() for d in documents])
+
+        for f in conjunctive_filters:
+            operator = f.operator.lower()
+            if operator in operator_mapping:
+                df = df.loc[
+                    getattr(df[f.property], operator_mapping[f.operator])(f.value)
+                ]
+            elif operator == "contains":
+                df = df.loc[df[f.property].str.contains(f.value, case=False)]
+
+        if sort_by:
+            df.sort_values(
+                [col.column_id for col in sort_by],
+                ascending=[col.direction == "asc" for col in sort_by],
+                inplace=True,
+            )
+
+        count = len(df)
+        result = df.iloc[skip:] if take is None else df.iloc[skip : skip + take]
+        return [
+            next(d for d in documents if d.trace_id == trace_id)
+            for trace_id in result["trace_id"]
+        ], count
+
+    def update(self, id: str, new_version: Trace) -> None:
+        self._safe_execute(
+            lambda db: db.update(new_version.dict(), lambda d: d["trace_id"] == id)
+        )
+
+    def delete(self, id: str) -> None:
+        self._safe_execute(lambda db: db.remove(lambda d: d["trace_id"] == id))
+
+    def delete_batch(self, ids: List[str]) -> List[str]:
+        for i in ids:
+            self.delete(i)
+
+    def _safe_execute(self, func: Callable[[TinyDB], Any]) -> Any:
+        with lock:
+            with TinyDB(self.path_to_db) as db:
+                return func(db)

great_ai/persistence/tracing_database_driver.py

@@ -0,0 +1,70 @@
+from abc import ABC, abstractmethod
+from datetime import datetime
+from pathlib import Path
+from typing import List, Optional, Sequence, Tuple, Union
+
+from ..utilities import ConfigFile
+from ..views import Filter, SortBy, Trace
+
+
+class TracingDatabaseDriver(ABC):
+    is_production_ready: bool
+    initialized: bool = False
+
+    @classmethod
+    def configure_credentials_from_file(
+        cls,
+        secrets_path: Union[Path, str],
+    ) -> None:
+        cls.configure_credentials(**ConfigFile(secrets_path))
+
+    @classmethod
+    def configure_credentials(
+        cls,
+    ) -> None:
+        cls.initialized = True
+
+    @abstractmethod
+    def save(self, document: Trace) -> str:
+        pass
+
+    @abstractmethod
+    def save_batch(
+        self,
+        documents: List[Trace],
+    ) -> List[str]:
+        pass
+
+    @abstractmethod
+    def get(self, id: str) -> Optional[Trace]:
+        pass
+
+    @abstractmethod
+    def query(
+        self,
+        *,
+        skip: int = 0,
+        take: Optional[int] = None,
+        conjunctive_filters: Sequence[Filter] = [],
+        conjunctive_tags: Sequence[str] = [],
+        until: Optional[datetime] = None,
+        since: Optional[datetime] = None,
+        has_feedback: Optional[bool] = None,
+        sort_by: Sequence[SortBy] = []
+    ) -> Tuple[List[Trace], int]:
+        pass
+
+    @abstractmethod
+    def update(self, id: str, new_version: Trace) -> None:
+        pass
+
+    @abstractmethod
+    def delete(self, id: str) -> None:
+        pass
+
+    @abstractmethod
+    def delete_batch(
+        self,
+        ids: List[str],
+    ) -> None:
+        pass