Add ground_truth API and refactor

.vscode/settings.json

@@ -15,6 +15,7 @@
     "initialised",
     "inplace",
     "ipynb",
+    "joblib",
     "lemmatize",
     "levelname",
     "levelno",

great_ai/src/great_ai/__main__.py

@@ -14,13 +14,14 @@ from uvicorn.supervisors.basereload import BaseReload
 from watchdog.events import FileSystemEvent, PatternMatchingEventHandler
 from watchdog.observers import Observer
 
-from .great_ai.context.configure import _is_in_production_mode
+from .great_ai.constants import SERVER_NAME
+from .great_ai.context import _is_in_production_mode
 from .great_ai.deploy import GreatAI
 from .great_ai.exceptions import ArgumentValidationError, MissingArgumentError
 from .parse_arguments import parse_arguments
 from .utilities.logger import get_logger
 
-logger = get_logger("GreatAI-Server")
+logger = get_logger(SERVER_NAME)
 
 
 GREAT_AI_LOGGING_CONFIG = {

great_ai/src/great_ai/great_ai/__init__.py

@@ -1,5 +1,11 @@
 from .context import configure
 from .deploy import GreatAI
 from .models import save_model, use_model
-from .output_models import ClassificationOutput, RegressionOutput
+from .output_models import (
+    ClassificationOutput,
+    MultiLabelClassificationOutput,
+    RegressionOutput,
+)
 from .parameters import log_metric, parameter
+from .persistence import MongoDbDriver, ParallelTinyDbDriver, TracingDatabaseDriver
+from .tracing import add_ground_truth, query_ground_truth

great_ai/src/great_ai/great_ai/constants.py

@@ -14,3 +14,13 @@ DEFAULT_LARGE_FILE_CONFIG_PATHS = {
 }
 
 GITHUB_LINK = "https://github.com/ScoutinScience/great-ai"
+
+TRAIN_SPLIT_TAG_NAME = "train"
+TEST_SPLIT_TAG_NAME = "test"
+VALIDATION_SPLIT_TAG_NAME = "validation"
+GROUND_TRUTH_TAG_NAME = "ground_truth"
+PRODUCTION_TAG_NAME = "production"
+DEVELOPMENT_TAG_NAME = "development"
+ONLINE_TAG_NAME = "online"
+
+SERVER_NAME = "GreatAI-Server"

great_ai/src/great_ai/great_ai/context.py

@@ -2,34 +2,68 @@ import os
 import random
 from logging import DEBUG, Logger
 from pathlib import Path
-from typing import Optional, Type
+from typing import Any, Dict, Optional, Type, cast
+
+from pydantic import BaseModel
 
-import great_ai.great_ai.context.context as context
 from great_ai.large_file import LargeFile, LargeFileLocal
 from great_ai.utilities.logger import get_logger
 
-from ..constants import (
+from .constants import (
     DEFAULT_LARGE_FILE_CONFIG_PATHS,
     DEFAULT_TRACING_DB_FILENAME,
     ENV_VAR_KEY,
     PRODUCTION_KEY,
 )
-from ..tracing.parallel_tinydb_driver import ParallelTinyDbDriver, TracingDatabase
+from .persistence import ParallelTinyDbDriver, TracingDatabaseDriver
+
+
+class Context(BaseModel):
+    tracing_database: TracingDatabaseDriver
+    large_file_implementation: Type[LargeFile]
+    is_production: bool
+    logger: Logger
+    should_log_exception_stack: bool
+    prediction_cache_size: int
+
+    class Config:
+        arbitrary_types_allowed = True
+
+    def to_flat_dict(self) -> Dict[str, Any]:
+        return {
+            "tracing_database": type(self.tracing_database).__name__,
+            "large_file_implementation": self.large_file_implementation.__name__,
+            "is_production": self.is_production,
+            "should_log_exception_stack": self.should_log_exception_stack,
+            "prediction_cache_size": self.prediction_cache_size,
+        }
+
+
+_context: Optional[Context] = None
+
+
+def get_context() -> Context:
+    if _context is None:
+        configure()
+
+    return cast(Context, _context)
 
 
 def configure(
+    *,
     log_level: int = DEBUG,
     seed: int = 42,
-    tracing_database: TracingDatabase = ParallelTinyDbDriver(
+    tracing_database: TracingDatabaseDriver = ParallelTinyDbDriver(
         Path(DEFAULT_TRACING_DB_FILENAME)
     ),
     large_file_implementation: Type[LargeFile] = LargeFileLocal,
     should_log_exception_stack: Optional[bool] = None,
     prediction_cache_size: int = 512,
 ) -> None:
+    global _context
     logger = get_logger("great_ai", level=log_level)
 
-    if context._context is not None:
+    if _context is not None:
         logger.warn(
             "Configuration has been already initialised, overwriting.\n"
             + "Make sure to call `configure()` before importing your application code."
@@ -39,12 +73,17 @@ def configure(
     _initialize_large_file(large_file_implementation, logger=logger)
     _set_seed(seed)
 
-    if not tracing_database.is_threadsafe:
-        logger.warning(
-            f"The selected persistence driver ({type(tracing_database).__name__}) is not threadsafe"
-        )
+    if not tracing_database.is_production_ready:
+        if is_production:
+            logger.error(
+                f"The selected tracing database ({type(tracing_database).__name__}) is not recommended for production"
+            )
+        else:
+            logger.warning(
+                f"The selected tracing database ({type(tracing_database).__name__}) is not recommended for production"
+            )
 
-    context._context = context.Context(
+    _context = Context(
         tracing_database=tracing_database,
         large_file_implementation=large_file_implementation,
         is_production=is_production,

great_ai/src/great_ai/great_ai/context/__init__.py

@@ -1,3 +0,0 @@
-from .configure import configure
-from .context import Context
-from .get_context import get_context

great_ai/src/great_ai/great_ai/context/context.py

@@ -1,32 +0,0 @@
-from logging import Logger
-from typing import Any, Dict, Optional, Type
-
-from pydantic import BaseModel
-
-from great_ai.large_file.large_file.large_file import LargeFile
-
-from ..tracing.tracing_database import TracingDatabase
-
-
-class Context(BaseModel):
-    tracing_database: TracingDatabase
-    large_file_implementation: Type[LargeFile]
-    is_production: bool
-    logger: Logger
-    should_log_exception_stack: bool
-    prediction_cache_size: int
-
-    class Config:
-        arbitrary_types_allowed = True
-
-    def to_flat_dict(self) -> Dict[str, Any]:
-        return {
-            "tracing_database": type(self.tracing_database).__name__,
-            "large_file_implementation": self.large_file_implementation.__name__,
-            "is_production": self.is_production,
-            "should_log_exception_stack": self.should_log_exception_stack,
-            "prediction_cache_size": self.prediction_cache_size,
-        }
-
-
-_context: Optional[Context] = None

great_ai/src/great_ai/great_ai/context/get_context.py

@@ -1,12 +0,0 @@
-from typing import cast
-
-import great_ai.great_ai.context.context as context
-
-from .configure import configure
-
-
-def get_context() -> context.Context:
-    if context._context is None:
-        configure()
-
-    return cast(context.Context, context._context)

great_ai/src/great_ai/great_ai/deploy/great_ai.py

@@ -1,7 +1,20 @@
 import inspect
+from asyncio.log import logger
 from functools import lru_cache, partial, wraps
-from typing import Any, Callable, Iterable, Optional, Sequence, Type, Union, cast
+from typing import (
+    Any,
+    Callable,
+    Generic,
+    Iterable,
+    List,
+    Optional,
+    Type,
+    TypeVar,
+    Union,
+    cast,
+)
 
+import yaml
 from fastapi import APIRouter, FastAPI, status
 from pydantic import BaseModel, create_model
 
@@ -26,14 +39,24 @@ from .routes import (
     bootstrap_trace_endpoints,
 )
 
+T = TypeVar("T")
 
-class GreatAI:
+
+class GreatAI(Generic[T]):
     def __init__(self, func: Callable[..., Any], version: str):
-        self._func = automatically_decorate_parameters(func)
-        get_function_metadata_store(self._func).is_finalised = True
+        func = automatically_decorate_parameters(func)
+        get_function_metadata_store(func).is_finalised = True
+
+        self._func = func
+
+        def func_in_tracing_context(*args: Any, **kwargs: Any) -> Trace[T]:
+            with TracingContext[T](func.__name__) as t:
+                result = func(*args, **kwargs)
+                output = t.finalise(output=result)
+            return output
 
         self._cached_func = lru_cache(get_context().prediction_cache_size)(
-            self._func
+            func_in_tracing_context
         )  # cannot put decorator on method, because it require the context to be setup
 
         wraps(func)(self)
@@ -49,25 +72,22 @@ class GreatAI:
             redoc_url=None,
         )
 
-    @freeze_arguments
-    def __call__(self, *args: Any, **kwargs: Any) -> Trace:
-        with TracingContext() as t:
-            result = self._cached_func(*args, **kwargs)
-            output = t.finalise(output=result)
-        return output
+        logger.info(
+            f"Current configuration: {yaml.dump(get_context().to_flat_dict(), stream=None)}"
+        )
 
     @staticmethod
     def deploy(
-        func: Optional[Callable[..., Any]] = None,
+        func: Optional[Callable[..., T]] = None,
         *,
         version: str = "0.0.1",
         disable_rest_api: bool = False,
         disable_docs: bool = False,
         disable_dashboard: bool = False,
-    ) -> Union[Callable[[Callable[..., Any]], "GreatAI"], "GreatAI"]:
+    ) -> Union[Callable[[Callable[..., T]], "GreatAI[T]"], "GreatAI[T]"]:
         if func is None:
             return cast(
-                Callable[..., Any],
+                Callable[[Callable[..., T]], GreatAI[T]],
                 partial(
                     GreatAI.deploy,
                     disable_http=disable_rest_api,
@@ -76,7 +96,7 @@ class GreatAI:
                 ),
             )
 
-        instance = GreatAI(func, version=version)
+        instance = GreatAI[T](func, version=version)
 
         if not disable_rest_api:
             instance._bootstrap_rest_api(
@@ -85,16 +105,18 @@ class GreatAI:
 
         return instance
 
+    @freeze_arguments
+    def __call__(self, *args: Any, **kwargs: Any) -> Trace[T]:
+        return self._cached_func(*args, **kwargs)
+
     def process_batch(
         self,
         batch: Iterable[Any],
         concurrency: Optional[int] = None,
-    ) -> Sequence[Trace]:
-        if not get_context().tracing_database.is_threadsafe:
-            concurrency = 1
-            get_context().logger.warning("Concurrency is ignored")
-
-        return parallel_map(self, batch, concurrency=concurrency)
+    ) -> List[Trace[T]]:
+        return parallel_map(
+            freeze_arguments(self._cached_func), batch, concurrency=concurrency
+        )
 
     @property
     def name(self) -> str:
@@ -144,9 +166,9 @@ class GreatAI:
 
         schema = self._get_schema()
 
-        @router.post("/", status_code=status.HTTP_200_OK, response_model=Trace)
+        @router.post("/", status_code=status.HTTP_200_OK, response_model=Trace[T])
         @use_http_exceptions
-        def predict(input_value: schema) -> Trace:  # type: ignore
+        def predict(input_value: schema) -> Trace[T]:  # type: ignore
             return self(**cast(BaseModel, input_value).dict())
 
         self.app.include_router(router)

great_ai/src/great_ai/great_ai/deploy/routes/bootstrap_feedback_endpoints.py

@@ -1,6 +1,5 @@
 from typing import Any
 
-import yaml
 from fastapi import APIRouter, FastAPI, HTTPException, Response, status
 
 from ...context import get_context
@@ -20,9 +19,6 @@ def bootstrap_feedback_endpoints(app: FastAPI) -> None:
             raise HTTPException(status_code=status.HTTP_404_NOT_FOUND)
 
         trace.feedback = input.feedback
-        trace.feedback_flat = yaml.dump(
-            input.feedback, default_flow_style=False, indent=2
-        )
 
         get_context().tracing_database.update(trace_id, trace)
         return Response(status_code=status.HTTP_202_ACCEPTED)
@@ -41,7 +37,6 @@ def bootstrap_feedback_endpoints(app: FastAPI) -> None:
             raise HTTPException(status_code=status.HTTP_404_NOT_FOUND)
 
         trace.feedback = None
-        trace.feedback_flat = None
 
         get_context().tracing_database.update(trace_id, trace)
         return Response(status_code=status.HTTP_204_NO_CONTENT)

great_ai/src/great_ai/great_ai/deploy/routes/bootstrap_trace_endpoints.py

@@ -3,7 +3,7 @@ from typing import List
 from fastapi import APIRouter, FastAPI, HTTPException, Response, status
 
 from ...context import get_context
-from ...views import Query, Trace, TraceView
+from ...views import Query, Trace
 
 
 def bootstrap_trace_endpoints(app: FastAPI) -> None:
@@ -12,7 +12,7 @@ def bootstrap_trace_endpoints(app: FastAPI) -> None:
         tags=["traces"],
     )
 
-    @router.post("/", status_code=status.HTTP_200_OK, response_model=List[TraceView])
+    @router.post("/", status_code=status.HTTP_200_OK, response_model=List[Trace])
     def query_traces(
         query: Query,
         skip: int = 0,
@@ -25,7 +25,7 @@ def bootstrap_trace_endpoints(app: FastAPI) -> None:
             take=take,
         )[0]
 
-    @router.get("/{trace_id}", status_code=status.HTTP_200_OK, response_model=TraceView)
+    @router.get("/{trace_id}", status_code=status.HTTP_200_OK, response_model=Trace)
     def get_trace(trace_id: str) -> Trace:
         result = get_context().tracing_database.get(trace_id)
         if result is None:

great_ai/src/great_ai/great_ai/deploy/routes/dashboard/create_dash_app.py

@@ -1,5 +1,5 @@
 from math import ceil
-from typing import Any, Dict, List, Tuple
+from typing import Any, Dict, List, Sequence, Tuple
 
 import pandas as pd
 import plotly.express as px
@@ -10,7 +10,7 @@ from flask import Flask
 
 from great_ai.utilities.unique import unique
 
-from ....constants import DASHBOARD_PATH
+from ....constants import DASHBOARD_PATH, ONLINE_TAG_NAME
 from ....context import get_context
 from ....helper import snake_case_to_text, text_to_hex_color
 from ....views import SortBy
@@ -23,11 +23,10 @@ from .get_traces_table import get_traces_table
 def create_dash_app(function_name: str, function_docs: str) -> Flask:
     accent_color = text_to_hex_color(function_name)
 
-    flask_app = Flask(__name__)
     app = Dash(
         function_name,
         requests_pathname_prefix=DASHBOARD_PATH + "/",
-        server=flask_app,
+        server=Flask(__name__),
         title=snake_case_to_text(function_name),
         update_title=None,
         external_stylesheets=[
@@ -129,6 +128,7 @@ def create_dash_app(function_name: str, function_docs: str) -> Flask:
         elements, count = get_context().tracing_database.query(
             skip=page_current * page_size,
             take=page_size,
+            conjunctive_tags=[ONLINE_TAG_NAME],
             conjunctive_filters=non_null_conjunctive_filters,
             sort_by=sort_by,
         )
@@ -145,8 +145,10 @@ def create_dash_app(function_name: str, function_docs: str) -> Flask:
     )
     def update_layout(
         n_intervals: int,
-    ) -> Tuple[List[Dict[str, str]], Dict[str, Any]]:
-        elements, count = get_context().tracing_database.query(take=1)
+    ) -> Tuple[List[Dict[str, Sequence[str]]], Dict[str, Any]]:
+        elements, count = get_context().tracing_database.query(
+            take=1, conjunctive_tags=[ONLINE_TAG_NAME]
+        )
 
         if elements:
             keys = list(elements[0].to_flat_dict().keys())
@@ -185,11 +187,10 @@ def create_dash_app(function_name: str, function_docs: str) -> Flask:
         non_null_conjunctive_filters = [f for f in conjunctive_filters if f is not None]
 
         elements, count = get_context().tracing_database.query(
-            conjunctive_filters=non_null_conjunctive_filters
+            conjunctive_tags=[ONLINE_TAG_NAME],
+            conjunctive_filters=non_null_conjunctive_filters,
         )
 
-        elements = [e.to_flat_dict() for e in elements]
-
         if not elements:
             return (
                 html.Span(
@@ -200,8 +201,10 @@ def create_dash_app(function_name: str, function_docs: str) -> Flask:
                 {"display": "none"},
             )
 
+        flat_elements = [e.to_flat_dict() for e in elements]
+
         execution_time_histogram = dcc.Graph(config={"displaylogo": False})
-        df = pd.DataFrame(elements)
+        df = pd.DataFrame(flat_elements)
         fig = px.histogram(
             df,
             x="original_execution_time_ms",
@@ -230,7 +233,7 @@ def create_dash_app(function_name: str, function_docs: str) -> Flask:
         )
         return execution_time_histogram, parallel_coords_fig, {}
 
-    return flask_app
+    return app.server
 
 
 def get_dimension_descriptor(df: pd.DataFrame, column: str) -> Dict[str, Any]:

great_ai/src/great_ai/great_ai/helper/__init__.py

@@ -1,7 +1,7 @@
-from .assert_function_is_not_finalised import assert_function_is_not_finalised
 from .freeze_arguments import freeze_arguments
 from .get_arguments import get_arguments
 from .get_function_metadata_store import get_function_metadata_store
+from .hashable_base_model import HashableBaseModel
 from .snake_case_to_text import snake_case_to_text
 from .strip_lines import strip_lines
 from .text_to_hex_color import text_to_hex_color

great_ai/src/great_ai/great_ai/helper/hashable_base_model.py

@@ -0,0 +1,6 @@
+from pydantic import BaseModel
+
+
+class HashableBaseModel(BaseModel):
+    def __hash__(self) -> int:
+        return hash((type(self),) + tuple(self.__dict__.values()))

great_ai/src/great_ai/great_ai/helper/use_http_exceptions.py

@@ -1,10 +1,12 @@
 from functools import wraps
-from typing import Any, Callable, Dict, List
+from typing import Any, Callable, Dict, List, TypeVar, cast
 
 from fastapi import HTTPException, status
 
+F = TypeVar("F", bound=Callable[..., Any])
 
-def use_http_exceptions(func: Callable[..., Any]) -> Callable[..., Any]:
+
+def use_http_exceptions(func: F) -> F:
     @wraps(func)
     def wrapper(*args: List[Any], **kwargs: Dict[str, Any]) -> Any:
         try:
@@ -15,4 +17,4 @@ def use_http_exceptions(func: Callable[..., Any]) -> Callable[..., Any]:
                 detail=f"The following exception has occurred: {type(e).__name__}: {e}",
             )
 
-    return wrapper
+    return cast(F, wrapper)

great_ai/src/great_ai/great_ai/models/save_model.py

@@ -7,7 +7,7 @@ from ..context import get_context
 
 
 def save_model(
-    model: Union[Path, str, object], key: str, keep_last_n: Optional[int] = None
+    model: Union[Path, str, object], key: str, *, keep_last_n: Optional[int] = None
 ) -> str:
     file = get_context().large_file_implementation(
         name=key, mode="wb", keep_last_n=keep_last_n

great_ai/src/great_ai/great_ai/models/use_model.py

@@ -1,11 +1,14 @@
 from functools import wraps
-from typing import Any, Callable, Dict, List, Literal, Union
+from typing import Any, Callable, Dict, List, Literal, TypeVar, Union, cast
 
-from ..helper import assert_function_is_not_finalised, get_function_metadata_store
+from ..helper import get_function_metadata_store
+from ..helper.assert_function_is_not_finalised import assert_function_is_not_finalised
 from ..tracing.tracing_context import TracingContext
 from ..views import Model
 from .load_model import load_model
 
+F = TypeVar("F", bound=Callable[..., Any])
+
 
 def use_model(
     key: str,
@@ -13,7 +16,7 @@ def use_model(
     version: Union[int, Literal["latest"]],
     return_path: bool = False,
     model_kwarg_name: str = "model",
-) -> Callable[..., Any]:
+) -> Callable[[F], F]:
     assert (
         isinstance(version, int) or version == "latest"
     ), "Only integers or the string literal `latest` is allowed as version"
@@ -24,7 +27,7 @@ def use_model(
         return_path=return_path,
     )
 
-    def decorator(func: Callable[..., Any]) -> Callable[..., Any]:
+    def decorator(func: F) -> F:
         assert_function_is_not_finalised(func)
 
         store = get_function_metadata_store(func)
@@ -35,11 +38,11 @@ def use_model(
 
         @wraps(func)
         def wrapper(*args: List[Any], **kwargs: Dict[str, Any]) -> Any:
-            tracing_context = TracingContext.get_current_context()
+            tracing_context = TracingContext.get_current_tracing_context()
             if tracing_context:
                 tracing_context.log_model(Model(key=key, version=actual_version))
             return func(*args, **kwargs, **{model_kwarg_name: model})
 
-        return wrapper
+        return cast(F, wrapper)
 
     return decorator

great_ai/src/great_ai/great_ai/output_models/__init__.py

@@ -1,2 +1,3 @@
 from .classification_output import ClassificationOutput
+from .multi_label_classification_output import MultiLabelClassificationOutput
 from .regression_output import RegressionOutput

great_ai/src/great_ai/great_ai/output_models/classification_output.py

@@ -1,12 +1,9 @@
 from typing import Any, Optional, Union
 
-from pydantic import BaseModel
+from ..helper import HashableBaseModel
 
 
-class ClassificationOutput(BaseModel):
+class ClassificationOutput(HashableBaseModel):
     label: Union[str, int]
     confidence: float
     explanation: Optional[Any]
-
-    def __hash__(self) -> int:
-        return hash((type(self),) + tuple(self.__dict__.values()))

great_ai/src/great_ai/great_ai/output_models/multi_label_classification_output.py

@@ -0,0 +1,8 @@
+from typing import List
+
+from ..helper import HashableBaseModel
+from .classification_output import ClassificationOutput
+
+
+class MultiLabelClassificationOutput(HashableBaseModel):
+    labels: List[ClassificationOutput] = []

great_ai/src/great_ai/great_ai/output_models/regression_output.py

@@ -1,11 +1,8 @@
 from typing import Any, Optional, Union
 
-from pydantic import BaseModel
+from ..helper import HashableBaseModel
 
 
-class RegressionOutput(BaseModel):
+class RegressionOutput(HashableBaseModel):
     value: Union[int, float]
     explanation: Optional[Any]
-
-    def __hash__(self) -> int:
-        return hash((type(self),) + tuple(self.__dict__.values()))

great_ai/src/great_ai/great_ai/parameters/automatically_decorate_parameters.py

@@ -1,5 +1,5 @@
 import inspect
-from typing import Any, Callable
+from typing import Any, Callable, TypeVar
 
 from great_ai.great_ai.helper.get_function_metadata_store import (
     get_function_metadata_store,
@@ -7,8 +7,10 @@ from great_ai.great_ai.helper.get_function_metadata_store import (
 
 from .parameter import parameter
 
+F = TypeVar("F", bound=Callable[..., Any])
 
-def automatically_decorate_parameters(func: Callable[..., Any]) -> Callable[..., Any]:
+
+def automatically_decorate_parameters(func: F) -> F:
     signature = inspect.signature(func)
     parameter_names = [
         param.name

great_ai/src/great_ai/great_ai/parameters/log_metric.py

@@ -1,13 +1,12 @@
 import inspect
 from typing import Any
 
-from great_ai.great_ai.context.get_context import get_context
-
-from ..tracing.tracing_context import TracingContext
+from ..context import get_context
+from ..tracing import TracingContext
 
 
 def log_metric(argument_name: str, value: Any) -> None:
-    tracing_context = TracingContext.get_current_context()
+    tracing_context = TracingContext.get_current_tracing_context()
     caller = inspect.stack()[1].function
     actual_name = f"metric:{caller}:{argument_name}"
     if tracing_context:

great_ai/src/great_ai/great_ai/parameters/parameter.py

@@ -1,22 +1,21 @@
 from functools import wraps
-from typing import Any, Callable, Dict
+from typing import Any, Callable, Dict, TypeVar, cast
 
 from ..exceptions import ArgumentValidationError
-from ..helper import (
-    assert_function_is_not_finalised,
-    get_arguments,
-    get_function_metadata_store,
-)
+from ..helper import get_arguments, get_function_metadata_store
+from ..helper.assert_function_is_not_finalised import assert_function_is_not_finalised
 from ..tracing.tracing_context import TracingContext
 
+F = TypeVar("F", bound=Callable[..., Any])
+
 
 def parameter(
     parameter_name: str,
     *,
     validator: Callable[[Any], bool] = lambda _: True,
     disable_logging: bool = False,
-) -> Callable[..., Any]:
-    def decorator(func: Callable[..., Any]) -> Callable[..., Any]:
+) -> Callable[[F], F]:
+    def decorator(func: F) -> F:
         get_function_metadata_store(func).input_parameter_names.append(parameter_name)
         assert_function_is_not_finalised(func)
 
@@ -39,7 +38,7 @@ def parameter(
                     f"Argument {parameter_name} in {func.__name__} did not pass validation"
                 )
 
-            context = TracingContext.get_current_context()
+            context = TracingContext.get_current_tracing_context()
             if context and not disable_logging:
                 context.log_value(name=f"{actual_name}:value", value=argument)
                 if isinstance(argument, str):
@@ -47,6 +46,6 @@ def parameter(
 
             return func(*args, **kwargs)
 
-        return wrapper
+        return cast(F, wrapper)
 
     return decorator

great_ai/src/great_ai/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/src/great_ai/great_ai/persistence/mongodb_driver.py

@@ -0,0 +1,5 @@
+from .tracing_database_driver import TracingDatabaseDriver
+
+
+class MongoDbDriver(TracingDatabaseDriver):
+    is_production_ready = True

great_ai/src/great_ai/great_ai/persistence/parallel_tinydb_driver.py

@@ -1,12 +1,13 @@
+from datetime import datetime
 from multiprocessing import Lock
 from pathlib import Path
-from typing import Any, Callable, Optional, Sequence, Tuple
+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 import TracingDatabase
+from .tracing_database_driver import TracingDatabaseDriver
 
 lock = Lock()
 
@@ -14,8 +15,8 @@ lock = Lock()
 operator_mapping = {"=": "eq", "!=": "ne", "<": "lt", "<=": "le", ">": "gt", ">=": "ge"}
 
 
-class ParallelTinyDbDriver(TracingDatabase):
-    is_threadsafe = True
+class ParallelTinyDbDriver(TracingDatabaseDriver):
+    is_production_ready = False
 
     def __init__(self, path_to_db: Path) -> None:
         super().__init__()
@@ -24,6 +25,10 @@ class ParallelTinyDbDriver(TracingDatabase):
     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:
@@ -32,13 +37,30 @@ class ParallelTinyDbDriver(TracingDatabase):
 
     def query(
         self,
+        *,
         skip: int = 0,
         take: Optional[int] = None,
         conjunctive_filters: Sequence[Filter] = [],
+        conjunctive_tags: Sequence[str] = [],
+        since: Optional[datetime] = None,
         sort_by: Sequence[SortBy] = [],
-    ) -> Tuple[Sequence[Trace], int]:
-        documents = [
-            Trace.parse_obj(t) for t in self._safe_execute(lambda db: db.all())
+        has_feedback: Optional[bool] = None
+    ) -> 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 (
+                    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:

great_ai/src/great_ai/great_ai/persistence/tracing_database_driver.py

@@ -1,16 +1,24 @@
 from abc import ABC, abstractmethod
-from typing import Optional, Sequence, Tuple
+from datetime import datetime
+from typing import List, Optional, Sequence, Tuple
 
 from ..views import Filter, SortBy, Trace
 
 
-class TracingDatabase(ABC):
-    is_threadsafe: bool
+class TracingDatabaseDriver(ABC):
+    is_production_ready: bool
 
     @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
@@ -18,11 +26,15 @@ class TracingDatabase(ABC):
     @abstractmethod
     def query(
         self,
+        *,
         skip: int = 0,
         take: Optional[int] = None,
         conjunctive_filters: Sequence[Filter] = [],
+        conjunctive_tags: Sequence[str] = [],
+        since: Optional[datetime] = None,
         sort_by: Sequence[SortBy] = [],
-    ) -> Tuple[Sequence[Trace], int]:
+        has_feedback: Optional[bool] = None
+    ) -> Tuple[List[Trace], int]:
         pass
 
     @abstractmethod

great_ai/src/great_ai/great_ai/tracing/__init__.py

@@ -1,3 +1,3 @@
-from .mongodb_driver import MongoDbDriver
-from .parallel_tinydb_driver import ParallelTinyDbDriver
-from .tracing_database import TracingDatabase
+from .add_ground_truth import add_ground_truth
+from .query_ground_truth import query_ground_truth
+from .tracing_context import TracingContext

great_ai/src/great_ai/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,
+    test_split_ratio: float,
+    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[T](
+            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/src/great_ai/great_ai/tracing/mongodb_driver.py

@@ -1,5 +0,0 @@
-from .tracing_database import TracingDatabase
-
-
-class MongoDbDriver(TracingDatabase):
-    pass

great_ai/src/great_ai/great_ai/tracing/query_ground_truth.py

@@ -0,0 +1,24 @@
+from datetime import datetime
+from typing import List, Optional, TypeVar, Union
+
+from ..context import get_context
+from ..views.trace import Trace
+
+T = TypeVar("T")
+
+
+def query_ground_truth(
+    conjunctive_tags: Union[List[str], str] = [],
+    *,
+    since: Optional[datetime] = None,
+    return_max_count: Optional[int] = None
+) -> List[Trace[T]]:
+    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/src/great_ai/great_ai/tracing/tracing_context.py

@@ -2,20 +2,34 @@ import threading
 from collections import defaultdict
 from datetime import datetime
 from types import TracebackType
-from typing import Any, DefaultDict, Dict, List, Literal, Optional, Type
+from typing import (
+    Any,
+    DefaultDict,
+    Dict,
+    Generic,
+    List,
+    Literal,
+    Optional,
+    Type,
+    TypeVar,
+)
 
-from ..context.get_context import get_context
+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:
+
+class TracingContext(Generic[T]):
     _contexts: DefaultDict[int, List["TracingContext"]] = defaultdict(lambda: [])
 
-    def __init__(self) -> None:
+    def __init__(self, function_name: str) -> None:
         self._models: List[Model] = []
         self._values: Dict[str, Any] = {}
-        self._trace: Optional[Trace] = None
+        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
@@ -23,7 +37,7 @@ class TracingContext:
     def log_model(self, model: Model) -> None:
         self._models.append(model)
 
-    def finalise(self, output: Any = None, exception: BaseException = None) -> Trace:
+    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
@@ -36,12 +50,19 @@ class TracingContext:
             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
 
     @classmethod
-    def get_current_context(cls) -> Optional["TracingContext"]:
+    def get_current_tracing_context(cls) -> Optional["TracingContext"]:
         if cls._contexts[threading.get_ident()]:
             return cls._contexts[threading.get_ident()][-1]
         return None

great_ai/src/great_ai/great_ai/views/__init__.py

@@ -9,4 +9,3 @@ from .operators import operators
 from .query import Query
 from .sort_by import SortBy
 from .trace import Trace
-from .trace_view import TraceView

great_ai/src/great_ai/great_ai/views/trace.py

@@ -1,35 +1,67 @@
-from typing import Any, Dict, Optional
+from typing import Any, Dict, Generic, List, Optional, TypeVar
+from uuid import uuid4
 
 import yaml
 from pydantic import validator
 
-from .trace_view import TraceView
+from ..helper import HashableBaseModel
+from .model import Model
+
+T = TypeVar("T")
 
 
-class Trace(TraceView):
-    models_flat: Optional[str]
-    output_flat: Optional[str]
-    feedback_flat: Optional[str]
+class Trace(HashableBaseModel, Generic[T]):
+    trace_id: Optional[str]
+    created: str
+    original_execution_time_ms: float
+    logged_values: Dict[str, Any]
+    models: List[Model]
+    exception: Optional[str]
+    output: T
+    feedback: Any = None
+    tags: List[str]
 
-    @validator("models_flat", always=True)
-    def flatten_models(cls, v: Optional[str], values: Dict[str, Any]) -> str:
-        return ", ".join(f"{m.key}:{m.version}" for m in values["models"])
+    @validator("trace_id", always=True)
+    def generate_id(cls, v: Optional[str], values: Dict[str, Any]) -> Optional[str]:
+        if not v:
+            return str(uuid4())
+        return v
 
-    @validator("output_flat", always=True)
-    def flatten_output(cls, v: Optional[str], values: Dict[str, Any]) -> str:
-        return yaml.dump(values["output"], default_flow_style=False, indent=2)
+    @property
+    def input(self) -> Any:
+        return (
+            self.logged_values["input"]
+            if list(self.logged_values.keys()) == ["input"]
+            else self.logged_values
+        )
+
+    @property
+    def models_flat(self) -> str:
+        return ", ".join(f"{m.key}:{m.version}" for m in self.models)
+
+    @property
+    def output_flat(self) -> str:
+        return yaml.dump(self.output, stream=None)
+
+    @property
+    def feedback_flat(self) -> str:
+        return (
+            "null" if self.feedback is None else yaml.dump(self.feedback, stream=None)
+        )
+
+    @property
+    def tags_flat(self) -> str:
+        return ",\n".join(self.tags)
 
     def to_flat_dict(self) -> Dict[str, Any]:
         return {
             "trace_id": self.trace_id,
             "created": self.created,
             "original_execution_time_ms": self.original_execution_time_ms,
-            "models_flat": self.models_flat,
-            "output_flat": self.output_flat,
-            "exception": self.exception or "null",
-            "feedback_flat": self.feedback_flat or "null",
             **self.logged_values,
+            "models_flat": self.models_flat,
+            "exception": "null" if self.exception is None else self.exception,
+            "output_flat": self.output_flat,
+            "feedback_flat": self.feedback_flat,
+            "tags_flat": self.tags_flat,
         }
-
-    def __hash__(self) -> int:
-        return hash((type(self),) + tuple(self.__dict__.values()))

great_ai/src/great_ai/great_ai/views/trace_view.py

@@ -1,26 +0,0 @@
-from typing import Any, Dict, List, Optional
-from uuid import uuid4
-
-from pydantic import BaseModel, validator
-
-from .model import Model
-
-
-class TraceView(BaseModel):
-    trace_id: Optional[str]
-    created: str
-    original_execution_time_ms: float
-    logged_values: Dict[str, Any]
-    models: List[Model]
-    exception: Optional[str]
-    output: Any
-    feedback: Any = None
-
-    @validator("trace_id", always=True)
-    def generate_id(cls, v: Optional[str], values: Dict[str, Any]) -> Optional[str]:
-        if not v:
-            return str(uuid4())
-        return v
-
-    def __hash__(self) -> int:
-        return hash((type(self),) + tuple(self.__dict__.values()))