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

@@ -0,0 +1,151 @@
+import os
+import random
+from logging import DEBUG, Logger
+from pathlib import Path
+from typing import Any, Dict, Optional, Type, cast
+
+from pydantic import BaseModel
+
+from great_ai.large_file import LargeFile, LargeFileLocal
+from great_ai.utilities.logger import get_logger
+
+from .constants import (
+    DEFAULT_LARGE_FILE_CONFIG_PATHS,
+    DEFAULT_TRACING_DB_FILENAME,
+    ENV_VAR_KEY,
+    PRODUCTION_KEY,
+)
+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: 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 is not None:
+        logger.warn(
+            "Configuration has been already initialised, overwriting.\n"
+            + "Make sure to call `configure()` before importing your application code."
+        )
+
+    is_production = _is_in_production_mode(logger=logger)
+    _initialize_large_file(large_file_implementation, logger=logger)
+    _set_seed(seed)
+
+    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(
+        tracing_database=tracing_database,
+        large_file_implementation=large_file_implementation,
+        is_production=is_production,
+        logger=logger,
+        should_log_exception_stack=not is_production
+        if should_log_exception_stack is None
+        else should_log_exception_stack,
+        prediction_cache_size=prediction_cache_size,
+    )
+
+    logger.info("Options: configured ✅")
+
+
+def _is_in_production_mode(logger: Optional[Logger]) -> bool:
+    environment = os.environ.get(ENV_VAR_KEY)
+
+    if environment is None:
+        if logger:
+            logger.warning(
+                f"Environment variable {ENV_VAR_KEY} is not set, defaulting to development mode ‼️"
+            )
+        is_production = False
+    else:
+        is_production = environment.lower() == PRODUCTION_KEY
+        if logger:
+            if not is_production:
+                logger.info(
+                    f"Value of {ENV_VAR_KEY} is `{environment}` which is not equal to `{PRODUCTION_KEY}`"
+                    + "defaulting to development mode ‼️"
+                )
+            else:
+                logger.info("Running in production mode ✅")
+
+    return is_production
+
+
+def _initialize_large_file(large_file: Type[LargeFile], logger: Logger) -> None:
+    path = DEFAULT_LARGE_FILE_CONFIG_PATHS[large_file]
+    if path is None:
+        return
+
+    if large_file.initialized:
+        logger.warning(
+            f"{large_file.__name__} has been already configured: skipping initialisation"
+        )
+        return
+
+    if path.exists():
+        large_file.configure_credentials_from_file(path)
+        logger.info(f"{large_file.__name__} initialised with config ({path.resolve()})")
+    else:
+        logger.warning(
+            f"Default {large_file.__name__} config ({path.resolve()}) not found, skipping {large_file.__name__} initialisation"
+        )
+
+
+def _set_seed(seed: int) -> None:
+    random.seed(seed)
+
+    try:
+        import numpy
+
+        numpy.random.seed(seed + 1)
+    except ImportError:
+        pass

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] = []