Move files

great_ai/tracing/__init__.py

@@ -0,0 +1,4 @@
+from .add_ground_truth import add_ground_truth
+from .delete_ground_truth import delete_ground_truth
+from .query_ground_truth import query_ground_truth
+from .tracing_context import TracingContext

great_ai/tracing/add_ground_truth.py

@@ -0,0 +1,67 @@
+from datetime import datetime
+from math import ceil
+from random import shuffle
+from typing import Any, Iterable, List, TypeVar
+
+from ..constants import (
+    GROUND_TRUTH_TAG_NAME,
+    TEST_SPLIT_TAG_NAME,
+    TRAIN_SPLIT_TAG_NAME,
+    VALIDATION_SPLIT_TAG_NAME,
+)
+from ..context import get_context
+from ..views import Trace
+
+T = TypeVar("T")
+
+
+def add_ground_truth(
+    inputs: Iterable[Any],
+    expected_outputs: Iterable[T],
+    *,
+    tags: List[str] = [],
+    train_split_ratio: float = 1,
+    test_split_ratio: float = 0,
+    validation_split_ratio: float = 0
+) -> None:
+    get_context()  # this resets the seed
+
+    inputs = list(inputs)
+    expected_outputs = list(expected_outputs)
+    assert len(inputs) == len(
+        expected_outputs
+    ), "The length of the inputs and expected_outputs must be equal"
+
+    sum_ratio = train_split_ratio + test_split_ratio + validation_split_ratio
+    assert sum_ratio > 0, "The sum of the split ratios must be a positive number"
+
+    train_split_ratio /= sum_ratio
+    test_split_ratio /= sum_ratio
+    validation_split_ratio /= sum_ratio
+
+    values = list(zip(inputs, expected_outputs))
+    shuffle(values)
+
+    split_tags = (
+        [TRAIN_SPLIT_TAG_NAME] * ceil(train_split_ratio * len(inputs))
+        + [TEST_SPLIT_TAG_NAME] * ceil(test_split_ratio * len(inputs))
+        + [VALIDATION_SPLIT_TAG_NAME] * ceil(validation_split_ratio * len(inputs))
+    )
+    shuffle(split_tags)
+
+    created = datetime.utcnow().isoformat()
+    traces = [
+        Trace(
+            created=created,
+            original_execution_time_ms=0,
+            logged_values=X if isinstance(X, dict) else {"input": X},
+            models=[],
+            output=y,
+            feedback=y,
+            exception=None,
+            tags=[GROUND_TRUTH_TAG_NAME, split_tag, *tags],
+        )
+        for ((X, y), split_tag) in zip(values, split_tags)
+    ]
+
+    get_context().tracing_database.save_batch(traces)

great_ai/tracing/delete_ground_truth.py

@@ -0,0 +1,22 @@
+from datetime import datetime
+from typing import List, Optional, Union
+
+from ..context import get_context
+
+
+def delete_ground_truth(
+    conjunctive_tags: Union[List[str], str] = [],
+    *,
+    until: Optional[datetime] = None,
+    since: Optional[datetime] = None,
+) -> None:
+    tags = (
+        conjunctive_tags if isinstance(conjunctive_tags, list) else [conjunctive_tags]
+    )
+    db = get_context().tracing_database
+
+    items, length = db.query(
+        conjunctive_tags=tags, until=until, since=since, has_feedback=True
+    )
+
+    db.delete_batch([i.trace_id for i in items])

great_ai/tracing/query_ground_truth.py

@@ -0,0 +1,22 @@
+from datetime import datetime
+from typing import List, Optional, Union
+
+from ..context import get_context
+from ..views import Trace
+
+
+def query_ground_truth(
+    conjunctive_tags: Union[List[str], str] = [],
+    *,
+    since: Optional[datetime] = None,
+    return_max_count: Optional[int] = None
+) -> List[Trace]:
+    tags = (
+        conjunctive_tags if isinstance(conjunctive_tags, list) else [conjunctive_tags]
+    )
+    db = get_context().tracing_database
+
+    items, length = db.query(
+        conjunctive_tags=tags, since=since, take=return_max_count, has_feedback=True
+    )
+    return items

great_ai/tracing/tracing_context.py

@@ -0,0 +1,86 @@
+from contextvars import ContextVar
+from datetime import datetime
+from types import TracebackType
+from typing import Any, Dict, Generic, List, Literal, Optional, Type, TypeVar
+
+from ..constants import DEVELOPMENT_TAG_NAME, ONLINE_TAG_NAME, PRODUCTION_TAG_NAME
+from ..context import get_context
+from ..views import Model, Trace
+
+T = TypeVar("T")
+
+
+class TracingContext(Generic[T]):
+    def __init__(self, function_name: str, do_not_persist_traces: bool) -> None:
+        self._do_not_persist_traces = do_not_persist_traces
+        self._models: List[Model] = []
+        self._values: Dict[str, Any] = {}
+        self._trace: Optional[Trace[T]] = None
+        self._start_time = datetime.utcnow()
+        self._name = function_name
+
+    def log_value(self, name: str, value: Any) -> None:
+        self._values[name] = value
+
+    def log_model(self, model: Model) -> None:
+        self._models.append(model)
+
+    def finalise(self, output: T = None, exception: BaseException = None) -> Trace[T]:
+        assert self._trace is None, "has been already finalised"
+
+        delta_time = (datetime.utcnow() - self._start_time).microseconds / 1000
+        self._trace = Trace(
+            created=self._start_time.isoformat(),
+            original_execution_time_ms=delta_time,
+            logged_values=self._values,
+            models=self._models,
+            output=output,
+            exception=None
+            if exception is None
+            else f"{type(exception).__name__}: {exception}",
+            tags=[
+                self._name,
+                ONLINE_TAG_NAME,
+                PRODUCTION_TAG_NAME
+                if get_context().is_production
+                else DEVELOPMENT_TAG_NAME,
+            ],
+        )
+
+        return self._trace
+
+    @staticmethod
+    def get_current_tracing_context() -> Optional["TracingContext"]:
+        return _current_tracing_context.get()
+
+    def __enter__(self) -> "TracingContext":
+        _current_tracing_context.set(self)
+        return self
+
+    def __exit__(
+        self,
+        type: Optional[Type[BaseException]],
+        exception: Optional[BaseException],
+        traceback: Optional[TracebackType],
+    ) -> Literal[False]:
+        _current_tracing_context.set(None)
+
+        if exception is not None and type is not None:
+            self.finalise(exception=exception)
+            if get_context().should_log_exception_stack:
+                get_context().logger.exception("Could not finish operation")
+            else:
+                get_context().logger.error(
+                    f"Could not finish operation because of {type.__name__}: {exception}"
+                )
+
+        assert self._trace is not None
+        if not self._do_not_persist_traces:
+            get_context().tracing_database.save(self._trace)
+
+        return False
+
+
+_current_tracing_context: ContextVar[Optional[TracingContext]] = ContextVar(
+    "_current_tracing_context"
+)