Improve REST API

.vscode/settings.json

@@ -3,6 +3,7 @@
     "boto",
     "botocore",
     "fastapi",
+    "gridfs",
     "iloc",
     "inplace",
     "ipynb",
@@ -10,8 +11,10 @@
     "matplotlib",
     "nbconvert",
     "plotly",
+    "proba",
     "psutil",
     "pydantic",
+    "pymongo",
     "pyplot",
     "redoc",
     "sklearn",

examples/simple/__deploy__.py

@@ -1,82 +0,0 @@
-#!/usr/bin/env python
-# coding: utf-8
-
-# # Train a domain classifier on the [semantic scholar dataset](https://api.semanticscholar.org/corpus)
-#
-# ## Part 3: Create production inference function
-#
-# In the [previous notebook](train.ipynb), we trained our AI model. Now, it's time to create **G**eneral **R**obust **E**nd-to-end **A**utomated **T**rustworthy deployment from it using the `GreatAI` Python package.
-
-# In[1]:
-
-
-import re
-from typing import List
-
-from great_ai import ClassificationOutput, GreatAI, use_model
-from great_ai.utilities.clean import clean
-from sklearn.pipeline import Pipeline
-
-# In[2]:
-
-
-@GreatAI.deploy
-@use_model("small-domain-prediction-v2", version="latest")
-def predict_domain(
-    text: str, model: Pipeline, target_confidence: int = 20
-) -> List[ClassificationOutput]:
-    """
-    Predict the scientific domain of the input text.
-    Return labels until their sum likelihood is larger than target_confidence.
-    """
-    assert 0 <= target_confidence <= 100, "invalid argument"
-
-    preprocessed = re.sub(r"[^a-zA-Z ]", "", clean(text, convert_to_ascii=True))
-    features = model.named_steps["vectorizer"].transform([preprocessed])
-    prediction = model.named_steps["classifier"].predict_proba(features)[0]
-
-    best_classes = sorted(enumerate(prediction), key=lambda v: v[1], reverse=True)
-
-    results: List[ClassificationOutput] = []
-    for class_index, probability in best_classes:
-        results.append(
-            ClassificationOutput(
-                label=model.named_steps["classifier"].classes_[class_index],
-                confidence=round(probability * 100),
-                explanation=[
-                    word
-                    for _, word in sorted(
-                        (
-                            (weight, word)
-                            for weight, word, count in zip(
-                                model.named_steps["classifier"].feature_log_prob_[
-                                    class_index
-                                ],
-                                model.named_steps["vectorizer"].get_feature_names_out(),
-                                features.A[0],
-                            )
-                            if count > 0
-                        ),
-                        reverse=True,
-                    )
-                ][:5],
-            )
-        )
-
-        if sum(r.confidence for r in results) >= target_confidence:
-            break
-
-    return results
-
-
-# In[3]:
-
-
-result = predict_domain(
-    """
-    State-of-the-art methods for zero-shot visual recognition formulate learning as a joint embedding problem of images and side information. In these formulations the current best complement to visual features are attributes: manually encoded vectors describing shared characteristics among categories. Despite good performance, attributes have limitations: (1) finer-grained recognition requires commensurately more, and (2) attributes do not provide a natural language interface. We propose to overcome these limitations by training neural language models from scratch; i.e. without pre-training and only consuming words and characters. Our proposed models train end-to-end to align with the fine-grained and category-specific content of images. Natural language provides a flexible and compact way of encoding only the salient visual aspects for distinguishing categories. By training on raw text, our model can do inference on raw text as well, providing humans a familiar mode both for annotation and retrieval. Our model achieves strong performance on zero-shot text-based image retrieval and significantly outperforms the attribute-based state-of-the-art for zero-shot classification on the CaltechUCSD Birds 200-2011 dataset. """
-)
-
-from pprint import pprint
-
-pprint(result.dict(), width=120)

examples/simple/deploy.ipynb

@@ -33,18 +33,18 @@
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "\u001b[38;5;226m2022-05-28 15:02:20,852 |  WARNING | Environment variable ENVIRONMENT is not set, defaulting to development mode ‼️\u001b[0m\n",
+      "\u001b[38;5;226m2022-05-28 18:17:13,344 |  WARNING | Environment variable ENVIRONMENT is not set, defaulting to development mode ‼️\u001b[0m\n",
-      "\u001b[38;5;39m2022-05-28 15:02:20,853 |     INFO | Options: configured ✅\u001b[0m\n",
+      "\u001b[38;5;39m2022-05-28 18:17:13,346 |     INFO | Options: configured ✅\u001b[0m\n",
-      "\u001b[38;5;39m2022-05-28 15:02:21,168 |     INFO | Latest version of small-domain-prediction-v2 is 8 (from versions: 3, 4, 5, 6, 7, 8)\u001b[0m\n",
+      "\u001b[38;5;39m2022-05-28 18:17:13,694 |     INFO | Latest version of small-domain-prediction is 10 (from versions: 9, 10)\u001b[0m\n",
-      "\u001b[38;5;39m2022-05-28 15:02:21,169 |     INFO | File small-domain-prediction-v2-8 found in cache\u001b[0m\n"
+      "\u001b[38;5;39m2022-05-28 18:17:13,694 |     INFO | File small-domain-prediction-10 found in cache\u001b[0m\n"
      ]
     }
    ],
    "source": [
     "@GreatAI.deploy\n",
-    "@use_model(\"small-domain-prediction-v2\", version=\"latest\")\n",
+    "@use_model(\"small-domain-prediction\", version=\"latest\")\n",
     "def predict_domain(\n",
-    "    text: str, model: Pipeline, target_confidence: int = 20\n",
+    "    text: str, model: Pipeline, target_confidence: int = 50\n",
     ") -> List[ClassificationOutput]:\n",
     "    \"\"\"\n",
     "    Predict the scientific domain of the input text.\n",
@@ -52,7 +52,7 @@
     "    \"\"\"\n",
     "    assert 0 <= target_confidence <= 100, \"invalid argument\"\n",
     "\n",
-    "    preprocessed = re.sub(r\"[^a-zA-Z ]\", \"\", clean(text, convert_to_ascii=True))\n",
+    "    preprocessed = re.sub(r\"[^a-zA-Z\\s]\", \"\", clean(text, convert_to_ascii=True))\n",
     "    features = model.named_steps[\"vectorizer\"].transform([preprocessed])\n",
     "    prediction = model.named_steps[\"classifier\"].predict_proba(features)[0]\n",
     "\n",
@@ -99,12 +99,11 @@
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "{'created': '2022-05-28T13:02:22.035886',\n",
+      "{'created': '2022-05-28T16:17:14.581693',\n",
-      " 'evaluation': None,\n",
-      " 'evaluation_id': 'c035d78e-dea9-415f-8f93-0dc4cd8dd7b5',\n",
       " 'exception': None,\n",
-      " 'execution_time_ms': 94.814,\n",
+      " 'execution_time_ms': 93.638,\n",
-      " 'logged_values': {'arg:predict_domain:target_confidence': 20,\n",
+      " 'feedback': None,\n",
+      " 'logged_values': {'arg:predict_domain:target_confidence': 50,\n",
       "                   'arg:predict_domain:text': '\\n'\n",
       "                                              '    State-of-the-art methods for zero-shot visual recognition formulate '\n",
       "                                              'learning as a joint embedding problem of images and side information. '\n",
@@ -125,10 +124,11 @@
       "                                              'outperforms the attribute-based state-of-the-art for zero-shot '\n",
       "                                              'classification on the CaltechUCSD Birds 200-2011 dataset. ',\n",
       "                   'arg:predict_domain:text:length': 1236},\n",
-      " 'models': [{'key': 'small-domain-prediction-v2', 'version': 8}],\n",
+      " 'models': [{'key': 'small-domain-prediction', 'version': 10}],\n",
       " 'output': [{'confidence': 99.0,\n",
       "             'explanation': ['information', 'model', 'learning', 'proposed', 'image'],\n",
-      "             'label': 'Computer Science'}]}\n"
+      "             'label': 'Computer Science'}],\n",
+      " 'trace_id': 'd35fcb96-0a95-45f8-93e4-8967160b17dd'}\n"
      ]
     }
    ],
@@ -140,7 +140,7 @@
     "\n",
     "from pprint import pprint\n",
     "\n",
-    "pprint(result.dict(), width=120)"
+    "# pprint(result.dict(), width=120)"
    ]
   }
  ],

great_ai/src/great_ai/great_ai/dashboard/create_dash_app.py

@@ -33,7 +33,7 @@ def create_dash_app(function_name: str, function_docs: str) -> Flask:
         ],
     )
 
-    documents = get_context().persistence.get_documents()
+    documents = get_context().tracing_database.query()
     df = pd.DataFrame(documents)
 
     execution_time_histogram = dcc.Graph(config={"displaylogo": False})
@@ -95,11 +95,11 @@ 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]
 
-        return get_context().persistence.query(
+        return get_context().tracing_database.query(
-            conjunctive_filters=non_null_conjunctive_filters,
-            sort_by=sort_by,
             skip=page_current * page_size,
             take=page_size,
+            conjunctive_filters=non_null_conjunctive_filters,
+            sort_by=sort_by,
         )
 
     @app.callback(
@@ -114,7 +114,7 @@ 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]
 
-        rows = get_context().persistence.query(
+        rows = get_context().tracing_database.query(
             conjunctive_filters=non_null_conjunctive_filters
         )
 

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

@@ -1,27 +1,11 @@
 import inspect
 from functools import lru_cache, partial, wraps
-from pathlib import Path
+from typing import Any, Callable, Iterable, Optional, Sequence, Type, Union, cast
-from typing import (
-    Any,
-    Callable,
-    Dict,
-    Iterable,
-    List,
-    Optional,
-    Sequence,
-    Type,
-    Union,
-    cast,
-)
 
-from fastapi import APIRouter, FastAPI, HTTPException, status
+from fastapi import APIRouter, FastAPI, status
-from fastapi.middleware.wsgi import WSGIMiddleware
-from fastapi.openapi.docs import get_swagger_ui_html
-from fastapi.responses import RedirectResponse
-from fastapi.staticfiles import StaticFiles
 from pydantic import BaseModel, create_model
-from starlette.responses import HTMLResponse
 
+from great_ai.great_ai.views.cache_statistics import CacheStatistics
 from great_ai.utilities.parallel_map import parallel_map
 
 from ..constants import METRICS_PATH
@@ -34,36 +18,32 @@ from ..helper import (
     use_http_exceptions,
 )
 from ..parameters import automatically_decorate_parameters
-from ..tracing import TracingContext
+from ..tracing.tracing_context import TracingContext
-from ..views import (
+from ..views import ApiMetadata, HealthCheckResponse, Trace
-    ApiMetadata,
+from .routes import (
-    EvaluationFeedbackRequest,
+    bootstrap_docs_endpoints,
-    HealthCheckResponse,
+    bootstrap_feedback_endpoints,
-    Query,
+    bootstrap_trace_endpoints,
-    Trace,
 )
 
-PATH = Path(__file__).parent.resolve()
-
 
 class GreatAI:
     def __init__(self, func: Callable[..., Any]):
         self._func = automatically_decorate_parameters(func)
-        self._func = freeze_arguments(
-            lru_cache(get_context().prediction_cache_size)(self._func)
-        )
-
         get_function_metadata_store(self._func).is_finalised = True
         wraps(func)(self)
 
         self.app = FastAPI(
             title=self.name,
             version=self.version,
-            description=self.documentation,
+            description=self.documentation
+            + f"\n\n Find out more on the [metrics page]({METRICS_PATH}).",
             docs_url=None,
             redoc_url=None,
         )
 
+    @freeze_arguments
+    @lru_cache(get_context().prediction_cache_size)
     def __call__(self, *args: Any, **kwargs: Any) -> Trace:
         with TracingContext() as t:
             result = self._func(*args, **kwargs)
@@ -103,7 +83,7 @@ class GreatAI:
         batch: Iterable[Any],
         concurrency: Optional[int] = None,
     ) -> Sequence[Trace]:
-        if not get_context().persistence.is_threadsafe:
+        if not get_context().tracing_database.is_threadsafe:
             concurrency = 1
             get_context().logger.warning("Concurrency is ignored")
 
@@ -132,14 +112,17 @@ class GreatAI:
 
     def _bootstrap_rest_api(self, disable_docs: bool, disable_metrics: bool) -> None:
         self._bootstrap_prediction_endpoints()
-        self._bootstrap_feedback_endpoints()
-        self._bootstrap_meta_endpoints()
 
         if not disable_docs:
-            self._bootstrap_docs_endpoints()
+            bootstrap_docs_endpoints(self.app)
 
         if not disable_metrics:
-            self._bootstrap_metrics_endpoints()
+            dash_app = create_dash_app(self._func.__name__, self.documentation)
+            bootstrap_trace_endpoints(self.app, dash_app)
+
+        bootstrap_feedback_endpoints(self.app)
+
+        self._bootstrap_meta_endpoints()
 
     def _bootstrap_prediction_endpoints(self) -> None:
         router = APIRouter(
@@ -154,19 +137,6 @@ class GreatAI:
         def predict(input_value: schema) -> Trace:  # type: ignore
             return self(**cast(BaseModel, input_value).dict())
 
-        @router.get(
-            "/:prediction_id", response_model=Trace, status_code=status.HTTP_200_OK
-        )
-        def get_prediction(prediction_id: str) -> Trace:
-            result = get_context().persistence.get_trace(prediction_id)
-            if result is None:
-                raise HTTPException(status_code=status.HTTP_404_NOT_FOUND)
-            return result
-
-        @router.delete("/:prediction_id", status_code=status.HTTP_204_NO_CONTENT)
-        def delete_prediction(prediction_id: str) -> None:
-            get_context().persistence.delete_trace(prediction_id)
-
         self.app.include_router(router)
 
     def _get_schema(self) -> Type[BaseModel]:
@@ -183,26 +153,6 @@ class GreatAI:
         schema: Type[BaseModel] = create_model("InputModel", **parameters)  # type: ignore
         return schema
 
-    def _bootstrap_feedback_endpoints(self) -> None:
-        router = APIRouter(
-            prefix="/predictions/:prediction_id/feedback",
-            tags=["feedback"],
-        )
-
-        @router.put("/", status_code=status.HTTP_202_ACCEPTED)
-        def set_feedback(prediction_id: str, input: EvaluationFeedbackRequest) -> None:
-            get_context().persistence.add_feedback(prediction_id, input.evaluation)
-
-        @router.get("/", status_code=status.HTTP_200_OK)
-        def get_feedback(prediction_id: str) -> Any:
-            return get_context().persistence.get_feedback(prediction_id)
-
-        @router.delete("/", status_code=status.HTTP_200_OK)
-        def delete_feedback(prediction_id: str) -> Any:
-            get_context().persistence.delete_feedback(prediction_id)
-
-        self.app.include_router(router)
-
     def _bootstrap_meta_endpoints(self) -> None:
         router = APIRouter(
             tags=["meta"],
@@ -210,53 +160,24 @@ class GreatAI:
 
         @router.get("/health", status_code=status.HTTP_200_OK)
         def check_health() -> HealthCheckResponse:
-            return HealthCheckResponse(is_healthy=True)
+            hits, misses, maxsize, cache_size = self.__call__.cache_info()  # type: ignore
+            cache_statistics = CacheStatistics(
+                hits=hits, misses=misses, size=cache_size, max_size=maxsize
+            )
+
+            return HealthCheckResponse(
+                is_healthy=True, cache_statistics=cache_statistics
+            )
 
         @router.get(
             "/version", response_model=ApiMetadata, status_code=status.HTTP_200_OK
         )
         def get_version() -> ApiMetadata:
             return ApiMetadata(
-                name=self.name, version=self.version, documentation=self.documentation
+                name=self.name,
-            )
+                version=self.version,
-
+                documentation=self.documentation,
-        self.app.include_router(router)
+                configuration=get_context().to_flat_dict(),
-
-    def _bootstrap_docs_endpoints(self) -> None:
-        @self.app.get("/docs", include_in_schema=False)
-        def custom_swagger_ui_html() -> HTMLResponse:
-            return get_swagger_ui_html(openapi_url="openapi.json", title=self.app.title)
-
-        @self.app.get("/docs/index.html", include_in_schema=False)
-        def redirect_to_docs() -> RedirectResponse:
-            return RedirectResponse("/docs")
-
-    def _bootstrap_metrics_endpoints(self) -> None:
-        dash_app = create_dash_app(self._func.__name__, self.documentation)
-        self.app.mount(METRICS_PATH, WSGIMiddleware(dash_app))
-
-        @self.app.get("/", include_in_schema=False)
-        def redirect_to_entrypoint() -> RedirectResponse:
-            return RedirectResponse("/metrics")
-
-        self.app.mount(
-            "/assets",
-            StaticFiles(directory=PATH / "../dashboard/assets"),
-            name="static",
-        )
-
-        router = APIRouter(
-            prefix="/metrics",
-            tags=["metrics"],
-        )
-
-        @router.post("/query", status_code=status.HTTP_200_OK)
-        def query_metrics(query: Query) -> List[Dict[str, Any]]:
-            return get_context().persistence.query(
-                conjunctive_filters=query.filter,
-                sort_by=query.sort,
-                skip=query.skip,
-                take=query.take,
             )
 
         self.app.include_router(router)

great_ai/src/great_ai/great_ai/deploy/routes/__init__.py

@@ -0,0 +1,3 @@
+from .bootstrap_docs_endpoints import bootstrap_docs_endpoints
+from .bootstrap_feedback_endpoints import bootstrap_feedback_endpoints
+from .bootstrap_trace_endpoints import bootstrap_trace_endpoints

great_ai/src/great_ai/great_ai/deploy/routes/bootstrap_docs_endpoints.py

@@ -0,0 +1,14 @@
+from fastapi import FastAPI
+from fastapi.openapi.docs import get_swagger_ui_html
+from fastapi.responses import RedirectResponse
+from starlette.responses import HTMLResponse
+
+
+def bootstrap_docs_endpoints(app: FastAPI) -> None:
+    @app.get("/docs", include_in_schema=False)
+    def custom_swagger_ui_html() -> HTMLResponse:
+        return get_swagger_ui_html(openapi_url="openapi.json", title=app.title)
+
+    @app.get("/docs/index.html", include_in_schema=False)
+    def redirect_to_docs() -> RedirectResponse:
+        return RedirectResponse("/docs")

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

@@ -0,0 +1,40 @@
+from typing import Any
+
+from fastapi import APIRouter, FastAPI, HTTPException, Response, status
+
+from ...context import get_context
+from ...views import EvaluationFeedbackRequest
+
+
+def bootstrap_feedback_endpoints(app: FastAPI) -> None:
+    router = APIRouter(
+        prefix="/traces/{trace_id}/feedback",
+        tags=["feedback"],
+    )
+
+    @router.put("/", status_code=status.HTTP_202_ACCEPTED)
+    def set_feedback(trace_id: str, input: EvaluationFeedbackRequest) -> Response:
+        trace = get_context().tracing_database.get(trace_id)
+        if trace is None:
+            raise HTTPException(status_code=status.HTTP_404_NOT_FOUND)
+        trace.feedback = input.feedback
+        get_context().tracing_database.update(trace_id, trace)
+        return Response(status_code=status.HTTP_202_ACCEPTED)
+
+    @router.get("/", status_code=status.HTTP_200_OK)
+    def get_feedback(trace_id: str) -> Any:
+        trace = get_context().tracing_database.get(trace_id)
+        if trace is None:
+            raise HTTPException(status_code=status.HTTP_404_NOT_FOUND)
+        return trace.feedback
+
+    @router.delete("/", status_code=status.HTTP_204_NO_CONTENT)
+    def delete_feedback(trace_id: str) -> Any:
+        trace = get_context().tracing_database.get(trace_id)
+        if trace is None:
+            raise HTTPException(status_code=status.HTTP_404_NOT_FOUND)
+        trace.feedback = None
+        get_context().tracing_database.update(trace_id, trace)
+        return Response(status_code=status.HTTP_204_NO_CONTENT)
+
+    app.include_router(router)

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

@@ -0,0 +1,60 @@
+from pathlib import Path
+from typing import List
+
+from fastapi import APIRouter, FastAPI, HTTPException, Response, status
+from fastapi.middleware.wsgi import WSGIMiddleware
+from fastapi.responses import RedirectResponse
+from fastapi.staticfiles import StaticFiles
+from flask import Flask
+
+from ...constants import METRICS_PATH
+from ...context import get_context
+from ...views import Query, Trace
+
+PATH = Path(__file__).parent.resolve()
+
+
+def bootstrap_trace_endpoints(app: FastAPI, dash_app: Flask) -> None:
+    app.mount(METRICS_PATH, WSGIMiddleware(dash_app))
+
+    @app.get("/", include_in_schema=False)
+    def redirect_to_entrypoint() -> RedirectResponse:
+        return RedirectResponse("/metrics")
+
+    app.mount(
+        "/assets",
+        StaticFiles(directory=PATH / "../../dashboard/assets"),
+        name="static",
+    )
+
+    router = APIRouter(
+        prefix="/traces",
+        tags=["traces"],
+    )
+
+    @router.post("/", status_code=status.HTTP_200_OK, response_model=List[Trace])
+    def query_traces(
+        query: Query,
+        skip: int = 0,
+        take: int = 100,
+    ) -> List[Trace]:
+        return get_context().tracing_database.query(
+            conjunctive_filters=query.filter,
+            sort_by=query.sort,
+            skip=skip,
+            take=take,
+        )
+
+    @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:
+            raise HTTPException(status_code=status.HTTP_404_NOT_FOUND)
+        return result
+
+    @router.delete("/{trace_id}", status_code=status.HTTP_204_NO_CONTENT)
+    def delete_trace(trace_id: str) -> Response:
+        get_context().tracing_database.delete(trace_id)
+        return Response(status_code=status.HTTP_204_NO_CONTENT)
+
+    app.include_router(router)

great_ai/src/great_ai/great_ai/helper/freeze_arguments.py

@@ -22,4 +22,5 @@ def freeze_arguments(func: Callable[..., Any]) -> Callable[..., Any]:
         }
         return func(*args, **kwargs)
 
+    wrapper.cache_info = func.cache_info  # type: ignore
     return wrapper

great_ai/src/great_ai/great_ai/models/load_model.py

@@ -2,17 +2,13 @@ from typing import Any, Optional, Tuple
 
 from joblib import load
 
-from great_ai.open_s3 import LargeFile
-
 from ..context import get_context
 
 
 def load_model(
     key: str, version: Optional[int] = None, return_path: bool = False
 ) -> Tuple[Any, int]:
-    get_context()  # will setup LargeFile if there was no config set
+    file = get_context().large_file_implementation(name=key, mode="rb", version=version)
-
-    file = LargeFile(name=key, mode="rb", version=version)
 
     if return_path:
         return file.get(), file.version

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

@@ -3,17 +3,15 @@ from typing import Optional, Union
 
 from joblib import dump
 
-from great_ai.open_s3 import LargeFile
-
 from ..context import get_context
 
 
 def save_model(
     model: Union[Path, str, object], key: str, keep_last_n: Optional[int] = None
 ) -> str:
-    get_context()  # will setup LargeFile if there was no config set
+    file = get_context().large_file_implementation(
-
+        name=key, mode="wb", keep_last_n=keep_last_n
-    file = LargeFile(name=key, mode="wb", keep_last_n=keep_last_n)
+    )
 
     if isinstance(model, Path) or isinstance(model, str):
         file.push(model)

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

@@ -2,7 +2,7 @@ from functools import wraps
 from typing import Any, Callable, Dict, List, Literal, Union
 
 from ..helper import assert_function_is_not_finalised, get_function_metadata_store
-from ..tracing import TracingContext
+from ..tracing.tracing_context import TracingContext
 from ..views import Model
 from .load_model import load_model
 

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

@@ -3,7 +3,7 @@ from typing import Any
 
 from great_ai.great_ai.context.get_context import get_context
 
-from ..tracing import TracingContext
+from ..tracing.tracing_context import TracingContext
 
 
 def log_metric(argument_name: str, value: Any) -> None:

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

@@ -7,7 +7,7 @@ from ..helper import (
     get_arguments,
     get_function_metadata_store,
 )
-from ..tracing import TracingContext
+from ..tracing.tracing_context import TracingContext
 
 
 def parameter(

great_ai/src/great_ai/great_ai/persistence/__init__.py

@@ -1,3 +0,0 @@
-from .mongodb_driver import MongoDbDriver
-from .parallel_tinydb_driver import ParallelTinyDbDriver
-from .persistence_driver import PersistenceDriver

great_ai/src/great_ai/great_ai/persistence/mongodb_driver.py

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

great_ai/src/great_ai/great_ai/persistence/persistence_driver.py

@@ -1,38 +0,0 @@
-from abc import ABC, abstractmethod
-from typing import Any, Dict, List, Optional
-
-from ..views import Filter, SortBy, Trace
-
-
-class PersistenceDriver(ABC):
-    is_threadsafe: bool
-
-    @abstractmethod
-    def save_trace(self, document: Trace) -> str:
-        pass
-
-    @abstractmethod
-    def add_feedback(self, id: str, evaluation: Any) -> None:
-        pass
-
-    @abstractmethod
-    def get_trace(self, id: str) -> Optional[Trace]:
-        pass
-
-    @abstractmethod
-    def get_traces(self) -> List[Trace]:
-        pass
-
-    @abstractmethod
-    def get_documents(self) -> List[Dict[str, Any]]:
-        pass
-
-    @abstractmethod
-    def query(
-        self,
-        conjunctive_filters: List[Filter],
-        sort_by: List[SortBy] = [],
-        skip: int = 0,
-        take: Optional[int] = None,
-    ) -> List[Dict[str, Any]]:
-        pass

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

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