Fix typing and minor issues

great_ai/__init__.py

@@ -3,7 +3,7 @@ __version__ = "0.0.12"
 
 
 from .context import configure
-from .deploy import GreatAI
+from .deploy import GreatAI, RouteConfig
 from .exceptions import (
     ArgumentValidationError,
     MissingArgumentError,

great_ai/__main__.py

@@ -8,7 +8,6 @@ from threading import Event
 from typing import Optional
 
 import uvicorn
-from parse_arguments import parse_arguments
 from uvicorn._subprocess import get_subprocess
 from uvicorn.config import LOGGING_CONFIG, Config
 from uvicorn.supervisors.basereload import BaseReload
@@ -21,6 +20,8 @@ from great_ai.deploy import GreatAI
 from great_ai.exceptions import ArgumentValidationError, MissingArgumentError
 from great_ai.utilities import get_logger
 
+from .parse_arguments import parse_arguments
+
 logger = get_logger(SERVER_NAME)
 
 
@@ -28,11 +29,11 @@ GREAT_AI_LOGGING_CONFIG = {
     **LOGGING_CONFIG,
     "formatters": {
         "default": {
-            "()": "great_ai.logger.CustomFormatter",
+            "()": "great_ai.utilities.logger.CustomFormatter",
             "fmt": "%(asctime)s | %(levelname)8s | %(message)s",
         },
         "access": {
-            "()": "great_ai.logger.CustomFormatter",
+            "()": "great_ai.utilities.logger.CustomFormatter",
             "fmt": "%(asctime)s | %(levelname)8s | %(message)s",  # noqa: E501
         },
     },

great_ai/context.py

@@ -56,12 +56,12 @@ def configure(
     *,
     log_level: int = DEBUG,
     seed: int = 42,
-    tracing_database: Optional[Type[TracingDatabaseDriver]] = None,
+    tracing_database_factory: Optional[Type[TracingDatabaseDriver]] = None,
     large_file_implementation: Optional[Type[LargeFileBase]] = None,
     should_log_exception_stack: Optional[bool] = None,
     prediction_cache_size: int = 512,
     disable_se4ml_banner: bool = False,
-    dashboard_table_size: int = 50,
+    dashboard_table_size: int = 20,
 ) -> None:
     global _context
     logger = get_logger("great_ai", level=log_level)
@@ -76,17 +76,17 @@ def configure(
 
     _set_seed(seed)
 
-    tracing_database = _initialize_tracing_database(tracing_database, logger=logger)()
+    tracing_database_factory = _initialize_tracing_database(
+        tracing_database_factory, logger=logger
+    )
+    tracing_database = tracing_database_factory()
 
     if not tracing_database.is_production_ready:
+        message = f"The selected tracing database ({tracing_database_factory.__name__}) is not recommended for production"
         if is_production:
-            logger.error(
+            logger.error(message)
-                f"The selected tracing database ({type(tracing_database).__name__}) is not recommended for production"
-            )
         else:
-            logger.warning(
+            logger.warning(message)
-                f"The selected tracing database ({type(tracing_database).__name__}) is not recommended for production"
-            )
 
     _context = Context(
         tracing_database=tracing_database,

great_ai/deploy/__init__.py

@@ -1 +1,2 @@
 from .great_ai import GreatAI
+from .routes import RouteConfig

great_ai/deploy/great_ai.py

@@ -1,266 +1,216 @@
-import inspect
 from functools import lru_cache, partial, wraps
+from textwrap import dedent
 from typing import (
     Any,
+    Awaitable,
     Callable,
     Generic,
-    Iterable,
     List,
     Optional,
-    Type,
+    Sequence,
     TypeVar,
+    Union,
     cast,
     overload,
 )
 
-from fastapi import APIRouter, FastAPI, status
+from async_lru import alru_cache
-from pydantic import BaseModel, create_model
+from fastapi import FastAPI
 
 from ..constants import DASHBOARD_PATH
 from ..context import get_context
-from ..helper import (
+from ..helper import freeze_arguments, get_function_metadata_store, snake_case_to_text
-    freeze_arguments,
-    get_function_metadata_store,
-    snake_case_to_text,
-    use_http_exceptions,
-)
 from ..models import model_versions
 from ..parameters import automatically_decorate_parameters
 from ..tracing.tracing_context import TracingContext
 from ..utilities import parallel_map
-from ..views import ApiMetadata, CacheStatistics, HealthCheckResponse, Trace
+from ..views import ApiMetadata, Trace
-from .routes import (
-    bootstrap_docs_endpoints,
-    bootstrap_feedback_endpoints,
-    bootstrap_trace_endpoints,
-)
 from .routes.bootstrap_dashboard import bootstrap_dashboard
+from .routes.bootstrap_docs_endpoints import bootstrap_docs_endpoints
+from .routes.bootstrap_feedback_endpoints import bootstrap_feedback_endpoints
+from .routes.bootstrap_meta_endpoints import bootstrap_meta_endpoints
+from .routes.bootstrap_prediction_endpoint import bootstrap_prediction_endpoint
+from .routes.bootstrap_trace_endpoints import bootstrap_trace_endpoints
+from .routes.route_config import RouteConfig
 
-T = TypeVar("T")
+T = TypeVar("T", bound=Union[Trace, Awaitable[Trace]])
+V = TypeVar("V")
 
 
-class GreatAI(Generic[T]):
+class GreatAI(Generic[T, V]):
-    def __init__(self, func: Callable[..., Any], version: str, return_raw_result: bool):
+    __name__: str
-        is_asynchronous = inspect.iscoroutinefunction(func)
+    __doc__: str
+
+    def __init__(
+        self,
+        func: Callable[..., Union[V, Awaitable[V]]],
+        version: str,
+        route_config: RouteConfig,
+    ):
         func = automatically_decorate_parameters(func)
         get_function_metadata_store(func).is_finalised = True
 
-        self._func = func
+        self._cached_func = self._get_cached_traced_function(func)
-
+        self._wrapped_func = wraps(func)(freeze_arguments(self._cached_func))
-        def func_in_tracing_context_sync(
-            *args: Any, do_not_persist_traces: bool = False, **kwargs: Any
-        ) -> Trace[T]:
-            with TracingContext[T](
-                func.__name__, do_not_persist_traces=do_not_persist_traces
-            ) as t:
-                result = func(*args, **kwargs)
-                output = t.finalise(output=result)
-            return result if return_raw_result else output
-
-        async def func_in_tracing_context_async(
-            *args: Any, do_not_persist_traces: bool = False, **kwargs: Any
-        ) -> Trace[T]:
-            with TracingContext[T](
-                func.__name__, do_not_persist_traces=do_not_persist_traces
-            ) as t:
-                result = await func(*args, **kwargs)
-                output = t.finalise(output=result)
-            return result if return_raw_result else output
-
-        func_in_tracing_context = (
-            func_in_tracing_context_async
-            if is_asynchronous
-            else func_in_tracing_context_sync
-        )
-
-        self._cached_func = lru_cache(get_context().prediction_cache_size)(
-            func_in_tracing_context
-        )  # cannot put decorator on method, because it require the context to be setup
 
         wraps(func)(self)
+        self.__doc__ = f"GreatAI wrapper for interacting with the `{self.__name__}` function.\n\n{dedent(self.__doc__ or '')}"
 
-        self._version = version
+        self.version = version
+        flat_model_versions = ".".join(f"{k}-v{v}" for k, v in model_versions)
+        if flat_model_versions:
+            self.version += f"+{flat_model_versions}"
 
         self.app = FastAPI(
-            title=self.name,
+            title=snake_case_to_text(self.__name__),
             version=self.version,
-            description=self.documentation
+            description=self.__doc__
             + f"\n\nFind out more in the [dashboard]({DASHBOARD_PATH}).",
             docs_url=None,
             redoc_url=None,
         )
 
+        self._bootstrap_rest_api(route_config)
+
     @overload
     @staticmethod
-    def create(
+    def create(  # type: ignore
-        func: Optional[Callable[..., T]] = None,
+        # "Overloaded function signatures 1 and 2 overlap with incompatible return types"
-    ) -> "GreatAI[T]":
+        # https://github.com/python/mypy/issues/12759
+        func: Callable[..., Awaitable[V]],
+    ) -> "GreatAI[Awaitable[Trace[V]], V]":
         ...
 
     @overload
     @staticmethod
     def create(
-        version: str,
+        func: Callable[..., V],
-        return_raw_result: bool,
+    ) -> "GreatAI[Trace[V], V]":
-        disable_rest_api: bool,
+        ...
-        disable_docs: bool,
+
-        disable_dashboard: bool,
+    @overload
-    ) -> Callable[[Callable[..., T]], "GreatAI[T]"]:
+    @staticmethod
+    def create(
+        func: Optional[Union[Callable[..., V], Callable[..., Awaitable[V]]]] = ...,
+        *,
+        version: str = ...,
+        route_config: RouteConfig = ...,
+    ) -> Callable:
         ...
 
     @staticmethod
     def create(
-        func: Optional[Callable[..., T]] = None,
+        func: Optional[Callable] = None,
         *,
         version: str = "0.0.1",
-        return_raw_result: bool = False,
+        route_config: RouteConfig = RouteConfig(),
-        disable_rest_api: bool = False,
+    ) -> Union[Callable, "GreatAI"]:
-        disable_docs: bool = False,
-        disable_dashboard: bool = False,
-    ):
         if func is None:
-            return cast(
-                Callable[[Callable[..., T]], GreatAI[T]],
-                partial(
-                    GreatAI.create,
-                    version=version,
-                    return_raw_result=return_raw_result,
-                    disable_rest_api=disable_rest_api,
-                    disable_docs=disable_docs,
-                    disable_dashboard=disable_dashboard,
-                ),
-            )
 
-        instance = GreatAI[T](
+            @overload
-            func, version=version, return_raw_result=return_raw_result
+            def inner(func: Awaitable[V]) -> GreatAI[Awaitable[Trace[V]], V]:
+                ...
+
+            @overload
+            def inner(func: Callable[..., V]) -> GreatAI[Trace[V], V]:
+                ...
+
+            def inner(func):  # type: ignore
+                return GreatAI.create(
+                    func,
+                    version=version,
+                    route_config=route_config,
+                )
+
+            return inner
+
+        return GreatAI[T, V](
+            func,
+            version=version,
+            route_config=route_config,
         )
 
-        if not disable_rest_api:
+    def __call__(self, *args: Any, **kwargs: Any) -> T:
-            instance._bootstrap_rest_api(
+        return self._wrapped_func(*args, **kwargs)
-                disable_docs=disable_docs, disable_dashboard=disable_dashboard
-            )
-
-        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],
+        batch: Sequence,
         concurrency: Optional[int] = None,
-        do_not_persist_traces: bool = False,
+        do_not_persist_traces: Optional[bool] = False,
-    ) -> List[Trace[T]]:
+    ) -> List[Trace[V]]:
         return list(
             parallel_map(
-                freeze_arguments(
+                partial(
-                    partial(
+                    self._wrapped_func, do_not_persist_traces=do_not_persist_traces
-                        self._cached_func, do_not_persist_traces=do_not_persist_traces
-                    )
                 ),
                 batch,
                 concurrency=concurrency,
             )
         )
 
-    @property
+    @staticmethod
-    def name(self) -> str:
+    def _get_cached_traced_function(
-        return snake_case_to_text(self._func.__name__)
+        func: Callable[..., Union[V, Awaitable[V]]]
+    ) -> Callable[..., T]:
+        @lru_cache(maxsize=get_context().prediction_cache_size)
+        def func_in_tracing_context_sync(
+            *args: Any,
+            do_not_persist_traces: bool = False,
+            **kwargs: Any,
+        ) -> Trace[V]:
+            with TracingContext[V](
+                func.__name__, do_not_persist_traces=do_not_persist_traces
+            ) as t:
+                result = func(*args, **kwargs)
+                return t.finalise(output=result)
 
-    @property
+        @alru_cache(maxsize=get_context().prediction_cache_size)
-    def version(self) -> str:
+        async def func_in_tracing_context_async(
-        flat_model_versions = ".".join(f"{k}-v{v}" for k, v in model_versions)
+            *args: Any,
-        if flat_model_versions:
+            do_not_persist_traces: bool = False,
-            flat_model_versions = f"+{flat_model_versions}"
+            **kwargs: Any,
+        ) -> Trace[V]:
+            with TracingContext[V](
+                func.__name__, do_not_persist_traces=do_not_persist_traces
+            ) as t:
+                result = await cast(Callable[..., Awaitable], func)(*args, **kwargs)
+                return t.finalise(output=result)
 
-        return f"{self._version}{flat_model_versions}"
-
-    @property
-    def documentation(self) -> str:
         return (
-            f"GreatAI wrapper for interacting with the `{self._func.__name__}` function.\n\n"
+            func_in_tracing_context_async
-            + (
+            if get_function_metadata_store(func).is_asynchronous
-                "\n".join(
+            else func_in_tracing_context_sync
-                    line.strip()
-                    for line in (self._func.__doc__ or "").split("\n")
-                    if line.strip()
-                )
-            )
         )
 
-    def _bootstrap_rest_api(self, disable_docs: bool, disable_dashboard: bool) -> None:
+    def _bootstrap_rest_api(self, route_config: RouteConfig) -> None:
-        self._bootstrap_prediction_endpoint()
+        if route_config.prediction_endpoint_enabled:
+            bootstrap_prediction_endpoint(self.app, self._wrapped_func)
 
-        if not disable_docs:
+        if route_config.docs_endpoints_enabled:
             bootstrap_docs_endpoints(self.app)
 
-        if not disable_dashboard:
+        if route_config.dashboard_enabled:
             bootstrap_dashboard(
                 self.app,
-                function_name=self._func.__name__,
+                function_name=self.__name__,
-                documentation=self.documentation,
+                documentation=self.__doc__,
             )
+
+        if route_config.trace_endpoints_enabled:
             bootstrap_trace_endpoints(self.app)
 
-        bootstrap_feedback_endpoints(self.app)
+        if route_config.feedback_endpoints_enabled:
-        self._bootstrap_meta_endpoints()
+            bootstrap_feedback_endpoints(self.app)
 
-    def _bootstrap_prediction_endpoint(self) -> None:
+        if route_config.meta_endpoints_enabled:
-        router = APIRouter(
+            bootstrap_meta_endpoints(
-            tags=["predictions"],
+                self.app,
-        )
+                self._cached_func,
-
+                ApiMetadata(
-        schema = self._get_schema()
+                    name=self.__name__,
-
+                    version=self.version,
-        @router.post(
+                    documentation=self.__doc__,
-            "/predict", status_code=status.HTTP_200_OK, response_model=Trace[T]
+                    configuration=get_context().to_flat_dict(),
-        )
+                ),
-        @use_http_exceptions
-        def predict(input_value: schema) -> Trace[T]:  # type: ignore
-            return self(**cast(BaseModel, input_value).dict())
-
-        self.app.include_router(router)
-
-    def _get_schema(self) -> Type[BaseModel]:
-        signature = inspect.signature(self._func)
-        parameters = {
-            p.name: (
-                p.annotation if p.annotation != inspect._empty else Any,
-                p.default if p.default != inspect._empty else ...,
             )
-            for p in signature.parameters.values()
-            if p.name in get_function_metadata_store(self._func).input_parameter_names
-        }
-
-        schema: Type[BaseModel] = create_model("InputModel", **parameters)  # type: ignore
-        return schema
-
-    def _bootstrap_meta_endpoints(self) -> None:
-        router = APIRouter(
-            tags=["meta"],
-        )
-
-        @router.get("/health", status_code=status.HTTP_200_OK)
-        def check_health() -> HealthCheckResponse:
-            hits, misses, maxsize, cache_size = self._cached_func.cache_info()
-            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,
-                configuration=get_context().to_flat_dict(),
-            )
-
-        self.app.include_router(router)

great_ai/deploy/routes/__init__.py

@@ -1,4 +1,7 @@
 from .bootstrap_dashboard import bootstrap_dashboard
 from .bootstrap_docs_endpoints import bootstrap_docs_endpoints
 from .bootstrap_feedback_endpoints import bootstrap_feedback_endpoints
+from .bootstrap_meta_endpoints import bootstrap_meta_endpoints
+from .bootstrap_prediction_endpoint import bootstrap_prediction_endpoint
 from .bootstrap_trace_endpoints import bootstrap_trace_endpoints
+from .route_config import RouteConfig

great_ai/deploy/routes/bootstrap_dashboard.py

@@ -6,7 +6,7 @@ from fastapi.responses import RedirectResponse
 from fastapi.staticfiles import StaticFiles
 
 from ...constants import DASHBOARD_PATH
-from .dashboard import create_dash_app
+from .dashboard.create_dash_app import create_dash_app
 
 PATH = Path(__file__).parent.resolve()
 

great_ai/deploy/routes/bootstrap_meta_endpoints.py

@@ -0,0 +1,26 @@
+from typing import Any
+
+from fastapi import APIRouter, FastAPI, status
+
+from ...views import ApiMetadata, CacheStatistics, HealthCheckResponse
+
+
+def bootstrap_meta_endpoints(app: FastAPI, func: Any, metadata: ApiMetadata) -> None:
+    router = APIRouter(
+        tags=["meta"],
+    )
+
+    @router.get("/health", status_code=status.HTTP_200_OK)
+    def check_health() -> HealthCheckResponse:
+        hits, misses, maxsize, cache_size = func.cache_info()
+        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 metadata
+
+    app.include_router(router)

great_ai/deploy/routes/bootstrap_prediction_endpoint.py

@@ -0,0 +1,51 @@
+import inspect
+from typing import Any, Awaitable, Callable, Type, Union, cast
+
+from fastapi import APIRouter, FastAPI, HTTPException, status
+from pydantic import BaseModel, create_model
+
+from ...helper import get_function_metadata_store
+from ...views import Trace
+
+
+def bootstrap_prediction_endpoint(
+    app: FastAPI, func: Callable[..., Union[Trace, Awaitable[Trace]]]
+) -> None:
+    router = APIRouter(
+        tags=["predictions"],
+    )
+
+    schema = _get_schema(func)
+
+    @router.post("/predict", status_code=status.HTTP_200_OK, response_model=Trace)
+    async def predict(input_value: schema) -> Trace:  # type: ignore
+        try:
+            if inspect.iscoroutinefunction(func):
+                return await cast(Callable[..., Awaitable[Trace]], func)(
+                    **cast(BaseModel, input_value).dict()
+                )
+            return cast(Callable[..., Trace], func)(
+                **cast(BaseModel, input_value).dict()
+            )
+        except Exception as e:
+            raise HTTPException(
+                status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+                detail=f"The following exception has occurred: {type(e).__name__}: {e}",
+            )
+
+    app.include_router(router)
+
+
+def _get_schema(func: Callable) -> Type[BaseModel]:
+    signature = inspect.signature(func)
+    parameters = {
+        p.name: (
+            p.annotation if p.annotation != inspect._empty else Any,
+            p.default if p.default != inspect._empty else ...,
+        )
+        for p in signature.parameters.values()
+        if p.name in get_function_metadata_store(func).input_parameter_names
+    }
+
+    schema: Type[BaseModel] = create_model("InputModel", **parameters)  # type: ignore
+    return schema

great_ai/deploy/routes/dashboard/assets/index.css

@@ -129,13 +129,17 @@ main > header > div > h1 {
 
 .version-tag {
   border-radius: var(--border-radius);
-  background: #ddd;
   display: inline-block;
   font-size: 1rem;
-  padding: 3px 6px;
+  padding: 3px 8px;
   margin-left: var(--small-padding)
 }
 
+main > header .version-tag {
+  background: var(--background-color);
+  vertical-align: 4px;
+}
+
 main > header > *:nth-child(2) {
   min-width: 250px;
   max-width: 550px;

great_ai/deploy/routes/dashboard/create_dash_app.py

@@ -21,12 +21,13 @@ from .get_traces_table import get_traces_table
 
 def create_dash_app(function_name: str, version: str, function_docs: str) -> Flask:
     accent_color = text_to_hex_color(function_name)
+    function_name = snake_case_to_text(function_name)
 
     app = Dash(
         function_name,
         requests_pathname_prefix=DASHBOARD_PATH + "/",
         server=Flask(__name__),
-        title=snake_case_to_text(function_name),
+        title=function_name,
         update_title=None,
         external_stylesheets=[
             "/assets/index.css",

great_ai/deploy/routes/dashboard/get_footer.py

@@ -1,5 +1,7 @@
 from dash import html
 
+from great_ai import __version__
+
 from ....constants import GITHUB_LINK
 
 
@@ -8,7 +10,9 @@ def get_footer() -> html.Footer:
         [
             html.Div(
                 [
-                    html.H6("GreatAI"),
+                    html.H6(
+                        ["GreatAI", html.Span(__version__, className="version-tag")]
+                    ),
                     html.P(
                         "A human-friendly framework for robust end-to-end AI deployments."
                     ),

great_ai/deploy/routes/route_config.py

@@ -0,0 +1,10 @@
+from pydantic import BaseModel
+
+
+class RouteConfig(BaseModel):
+    prediction_endpoint_enabled: bool = True
+    docs_endpoints_enabled: bool = True
+    dashboard_enabled: bool = True
+    feedback_endpoints_enabled: bool = True
+    trace_endpoints_enabled: bool = True
+    meta_endpoints_enabled: bool = True

great_ai/helper/__init__.py

@@ -5,4 +5,3 @@ 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
-from .use_http_exceptions import use_http_exceptions

great_ai/helper/assert_function_is_not_finalised.py

@@ -1,10 +1,10 @@
-from typing import Any, Callable
+from typing import Callable
 
 from ..exceptions import WrongDecoratorOrderError
 from .get_function_metadata_store import get_function_metadata_store
 
 
-def assert_function_is_not_finalised(func: Callable[..., Any]) -> None:
+def assert_function_is_not_finalised(func: Callable) -> None:
     error_message = (
         "The outer-most (first) decorator has to be `@GreatAI.deploy`. "
         + f"In the case of `{func.__name__}`, it is not: fix this by moving `@GreatAI.deploy` to the top."

great_ai/helper/freeze_arguments.py

@@ -1,41 +1,50 @@
+import inspect
 from functools import wraps
-from typing import Any, Callable, Dict, List, Set, Union
+from typing import Any, Callable, Mapping, Sequence, Set, TypeVar, Union, cast
 
 from pydantic import BaseModel
 
+F = TypeVar("F", bound=Callable)
+
 
 class FrozenDict(dict):
-    def __hash__(self) -> int:
+    def __hash__(self) -> int:  # type: ignore
         return hash(frozenset((k, freeze(v)) for k, v in self.items()))
 
 
 class FrozenList(list):
-    def __hash__(self) -> int:
+    def __hash__(self) -> int:  # type: ignore
         return hash(tuple(freeze(i) for i in self))
 
 
 class FrozenSet(set):
-    def __hash__(self) -> int:
+    def __hash__(self) -> int:  # type: ignore
         return hash(frozenset(freeze(i) for i in self))
 
 
-def freeze_arguments(func: Callable[..., Any]) -> Callable[..., Any]:
+def freeze_arguments(func: F) -> F:
-    """Transform mutable dictionary
-    Into immutable
-    Useful to be compatible with cache
-    source: https://stackoverflow.com/questions/6358481/using-functools-lru-cache-with-dictionary-arguments
-    """
-
     @wraps(func)
-    def wrapper(*args: List[Any], **kwargs: Dict[str, Any]) -> Any:
+    def wrapper(*args: Any, **kwargs: Any) -> Any:
         args = tuple(freeze(arg) for arg in args)
         kwargs = {k: freeze(v) for k, v in kwargs.items()}
         return func(*args, **kwargs)
 
-    return wrapper
+    @wraps(func)
+    async def async_wrapper(*args: Any, **kwargs: Any) -> Any:
+        return await wrapper(*args, **kwargs)
+
+    return cast(F, async_wrapper if inspect.iscoroutinefunction(func) else wrapper)
 
 
-def freeze(value: Union[List[Any], Dict[str, Any], Set[Any]]) -> Any:
+def freeze(value: Union[Sequence[Any], Mapping[str, Any], Set[Any], BaseModel]) -> Any:
+    """
+    >>> class MyClass(BaseModel):
+    ...     a: int
+    >>> my_object = MyClass(a=3)
+    >>> my_other_object = MyClass(a=4)
+    >>> freeze(my_object) == freeze(my_other_object), freeze(my_object) == freeze(my_object)
+    (False, True)
+    """
     if isinstance(value, dict):
         return FrozenDict(value)
 
@@ -47,7 +56,7 @@ def freeze(value: Union[List[Any], Dict[str, Any], Set[Any]]) -> Any:
 
     if isinstance(value, BaseModel):
 
-        class HashableValue(type(value)):
+        class HashableValue(type(value)):  # type: ignore
             def __hash__(self) -> int:
                 return hash(frozenset((k, freeze(v)) for k, v in self.dict().items()))
 

great_ai/helper/get_arguments.py

@@ -3,7 +3,7 @@ from typing import Any, Callable, Dict, Mapping, Sequence
 
 
 def get_arguments(
-    func: Callable[..., Any], args: Sequence[Any], kwargs: Mapping[str, Any]
+    func: Callable, args: Sequence[Any], kwargs: Mapping[str, Any]
 ) -> Dict[str, Any]:
     """Return mapping from parameter names to actual argument values"""
 

great_ai/helper/get_function_metadata_store.py

@@ -1,12 +1,14 @@
+import inspect
 from typing import Any, Callable, cast
 
 from ..views.function_metadata import FunctionMetadata
 
 
-def get_function_metadata_store(func: Callable[..., Any]) -> FunctionMetadata:
+def get_function_metadata_store(func: Callable) -> FunctionMetadata:
     any_func = cast(Any, func)
 
     if not hasattr(any_func, "_great_ai_metadata"):
-        any_func._great_ai_metadata = FunctionMetadata()
+        is_asynchronous = inspect.iscoroutinefunction(func)
+        any_func._great_ai_metadata = FunctionMetadata(is_asynchronous=is_asynchronous)
 
     return any_func._great_ai_metadata

great_ai/helper/text_to_hex_color.py

@@ -4,10 +4,14 @@ from hashlib import md5
 
 def text_to_hex_color(text: str) -> str:
     ascii_bytes = text.encode("ascii")
+
     digest = md5(
         ascii_bytes
     ).hexdigest()  # the built-in hash function is salted differently in each process
+
     integer = int(digest, 16)
     hue = integer % 6311 / 6311.0
-    rgb = colorsys.hsv_to_rgb(hue, 0.75, 0.6)
+
+    rgb = colorsys.hsv_to_rgb(hue, 0.8, 0.6)
+
     return "#" + "".join("%02X" % round(i * 255) for i in rgb)

great_ai/helper/use_http_exceptions.py

@@ -1,20 +0,0 @@
-from functools import wraps
-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: F) -> F:
-    @wraps(func)
-    def wrapper(*args: List[Any], **kwargs: Dict[str, Any]) -> Any:
-        try:
-            return func(*args, **kwargs)
-        except Exception as e:
-            raise HTTPException(
-                status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
-                detail=f"The following exception has occurred: {type(e).__name__}: {e}",
-            )
-
-    return cast(F, wrapper)

great_ai/large_file/large_file/large_file_base.py

@@ -89,9 +89,7 @@ class LargeFileBase(ABC):
         cls.configure_credentials(**ConfigFile(secrets_path))
 
     @classmethod
-    def configure_credentials(
+    def configure_credentials(cls, **kwargs: str) -> None:
-        cls,
-    ) -> None:
         cls.initialized = True
 
     def __enter__(self) -> IO:

great_ai/large_file/large_file/large_file_mongo.py

@@ -1,7 +1,7 @@
 import re
 from functools import cached_property
 from pathlib import Path
-from typing import Any, List, Mapping
+from typing import Any, List
 
 from gridfs import DEFAULT_CHUNK_SIZE, Database, GridFSBucket
 from pymongo import MongoClient
@@ -27,7 +27,7 @@ class LargeFileMongo(LargeFileBase):
         *,
         mongo_connection_string: str,
         mongo_database: str,
-        **_: Mapping[str, Any],
+        **_: Any,
     ) -> None:
         cls.mongo_connection_string = mongo_connection_string
         cls.mongo_database = mongo_database

great_ai/large_file/large_file/large_file_s3.py

@@ -1,6 +1,6 @@
 from functools import cached_property
 from pathlib import Path
-from typing import Any, List, Mapping, Optional
+from typing import Any, List, Optional
 
 import boto3
 
@@ -58,7 +58,7 @@ class LargeFileS3(LargeFileBase):
         aws_secret_access_key: str,
         large_files_bucket_name: str,
         aws_endpoint_url: Optional[str] = None,
-        **_: Mapping[str, Any],
+        **_: Any,
     ) -> None:
         cls.region_name = aws_region_name
         cls.access_key_id = aws_access_key_id

great_ai/models/use_model.py

@@ -21,7 +21,7 @@ from ..helper.assert_function_is_not_finalised import assert_function_is_not_fin
 from ..tracing.tracing_context import TracingContext
 from ..views import Model
 
-F = TypeVar("F", bound=Callable[..., Any])
+F = TypeVar("F", bound=Callable)
 
 
 def use_model(
@@ -70,4 +70,6 @@ def _load_model(key: str, version: Optional[int] = None) -> Tuple[Any, int]:
         return path, file.version
 
     with file as f:
-        return load(f), file.version
+        loaded = load(f)
+
+    return loaded, file.version

great_ai/parameters/automatically_decorate_parameters.py

@@ -1,10 +1,10 @@
 import inspect
-from typing import Any, Callable, TypeVar
+from typing import Callable, TypeVar
 
 from ..helper.get_function_metadata_store import get_function_metadata_store
 from .parameter import parameter
 
-F = TypeVar("F", bound=Callable[..., Any])
+F = TypeVar("F", bound=Callable)
 
 
 def automatically_decorate_parameters(func: F) -> F:

great_ai/parameters/parameter.py

@@ -8,14 +8,13 @@ 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
 
-T = TypeVar("T")
+F = TypeVar("F", bound=Callable)
-F = TypeVar("F", bound=Callable[..., Any])
 
 
 def parameter(
     parameter_name: str,
     *,
-    validator: Callable[[T], bool] = lambda _: True,
+    validator: Callable[[Any], bool] = lambda _: True,
     disable_logging: bool = False,
 ) -> Callable[[F], F]:
     def decorator(func: F) -> F:

great_ai/persistence/mongodb_driver.py

@@ -1,5 +1,5 @@
 from datetime import datetime
-from typing import Any, List, Mapping, Optional, Sequence, Tuple
+from typing import Any, Dict, List, Optional, Sequence, Tuple
 
 from pymongo import MongoClient
 
@@ -20,6 +20,9 @@ operator_mapping = {
 class MongodbDriver(TracingDatabaseDriver):
     is_production_ready = True
 
+    mongo_connection_string: str
+    mongo_database: str
+
     def __init__(self) -> None:
         super().__init__()
         if self.mongo_connection_string is None or self.mongo_database is None:
@@ -33,7 +36,7 @@ class MongodbDriver(TracingDatabaseDriver):
         *,
         mongo_connection_string: str,
         mongo_database: str,
-        **_: Mapping[str, Any],
+        **_: Any,
     ) -> None:
         cls.mongo_connection_string = mongo_connection_string
         cls.mongo_database = mongo_database
@@ -43,21 +46,23 @@ class MongodbDriver(TracingDatabaseDriver):
         serialized = trace.to_flat_dict()
         serialized["_id"] = trace.trace_id
 
-        with MongoClient(self.mongo_connection_string) as client:
+        with MongoClient[Any](self.mongo_connection_string) as client:
-            return client[self.mongo_database].traces.insert_one(serialized)
+            return client[self.mongo_database].traces.insert_one(serialized).inserted_id
 
     def save_batch(self, documents: List[Trace]) -> List[str]:
         serialized = [d.to_flat_dict() for d in documents]
         for s in serialized:
             s["_id"] = s["trace_id"]
 
-        with MongoClient(self.mongo_connection_string) as client:
+        with MongoClient[Any](self.mongo_connection_string) as client:
-            return client[self.mongo_database].traces.insert_many(
+            return (
-                serialized, ordered=False
+                client[self.mongo_database]
+                .traces.insert_many(serialized, ordered=False)
+                .inserted_ids
             )
 
     def get(self, id: str) -> Optional[Trace]:
-        with MongoClient(self.mongo_connection_string) as client:
+        with MongoClient[Any](self.mongo_connection_string) as client:
             value = client[self.mongo_database].traces.find_one(id)
 
         if value:
@@ -83,15 +88,8 @@ class MongodbDriver(TracingDatabaseDriver):
         sort_by: Sequence[SortBy] = [],
     ) -> Tuple[List[Trace], int]:
 
-        query = {
+        query: Dict[str, Any] = {
-            "filter": {
+            "filter": {},
-                "$and": [{"tags": tag} for tag in conjunctive_tags]
-                + [
-                    {f.property: {self._get_operator(f): f.value}}
-                    for f in conjunctive_filters
-                ]
-                + [{}]
-            },
             "sort": [
                 (col.column_id, 1 if col.direction == "asc" else -1) for col in sort_by
             ],
@@ -103,35 +101,43 @@ class MongodbDriver(TracingDatabaseDriver):
         if take:
             query["limit"] = take
 
+        and_query: List[Dict[str, Any]] = [{}]
+        and_query.extend({"tags": tag} for tag in conjunctive_tags)
+        and_query.extend(
+            {f.property: {self._get_operator(f): f.value}} for f in conjunctive_filters
+        )
+
         if since:
-            query["filter"]["$and"].append({"created": {"$gte": since}})
+            and_query.append({"created": {"$gte": since}})
 
         if until:
-            query["filter"]["$and"].append({"created": {"$lte": until}})
+            and_query.append({"created": {"$lte": until}})
 
         if has_feedback is not None:
-            query["filter"]["$and"].append(
+            and_query.append(
                 {"feedback": {"$ne": None}} if has_feedback else {"feedback": None}
             )
+        query["filter"]["$and"] = and_query
 
-        with MongoClient(self.mongo_connection_string) as client:
+        with MongoClient[Any](self.mongo_connection_string) as client:
             values = client[self.mongo_database].traces.find(**query)
-            documents = [Trace.parse_obj(t) for t in values]
+            documents = [Trace[Any].parse_obj(t) for t in values]
 
         return documents, len(documents)
 
     def update(self, id: str, new_version: Trace) -> None:
         serialized = new_version.to_flat_dict()
         serialized["_id"] = new_version.trace_id
-        with MongoClient(self.mongo_connection_string) as client:
+
-            client[self.mongo_database].traces.update_one(id, new_version)
+        with MongoClient[Any](self.mongo_connection_string) as client:
+            client[self.mongo_database].traces.update_one({"_id": id}, serialized)
 
     def delete(self, id: str) -> None:
-        with MongoClient(self.mongo_connection_string) as client:
+        with MongoClient[Any](self.mongo_connection_string) as client:
-            client[self.mongo_database].traces.delete_one(id)
+            client[self.mongo_database].traces.delete_one({"_id": id})
 
-    def delete_batch(self, ids: List[str]) -> List[str]:
+    def delete_batch(self, ids: List[str]) -> None:
         delete_filter = {"_id": {"$in": ids}}
 
-        with MongoClient(self.mongo_connection_string) as client:
+        with MongoClient[Any](self.mongo_connection_string) as client:
-            return client[self.mongo_database].traces.delete_many(delete_filter)
+            client[self.mongo_database].traces.delete_many(delete_filter)

great_ai/persistence/parallel_tinydb_driver.py

@@ -1,7 +1,7 @@
 from datetime import datetime
 from multiprocessing import Lock
 from pathlib import Path
-from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple, cast
+from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple
 
 import pandas as pd
 from tinydb import TinyDB
@@ -48,14 +48,8 @@ class ParallelTinyDbDriver(TracingDatabaseDriver):
         def does_match(d: Dict[str, Any]) -> bool:
             return (
                 not set(conjunctive_tags) - set(d["tags"])
-                and (
+                and (since is None or datetime.fromisoformat(d["created"]) >= since)
-                    since is None
+                and (until is None or datetime.fromisoformat(d["created"]) <= until)
-                    or cast(datetime, datetime.fromisoformat(d["created"])) >= since
-                )
-                and (
-                    until is None
-                    or cast(datetime, datetime.fromisoformat(d["created"])) <= until
-                )
                 and (
                     has_feedback is None or has_feedback == (d["feedback"] is not None)
                 )
@@ -102,7 +96,7 @@ class ParallelTinyDbDriver(TracingDatabaseDriver):
     def delete(self, id: str) -> None:
         self._safe_execute(lambda db: db.remove(lambda d: d["trace_id"] == id))
 
-    def delete_batch(self, ids: List[str]) -> List[str]:
+    def delete_batch(self, ids: List[str]) -> None:
         for i in ids:
             self.delete(i)
 

great_ai/tracing/add_ground_truth.py

@@ -51,7 +51,7 @@ def add_ground_truth(
 
     created = datetime.utcnow().isoformat()
     traces = [
-        Trace(
+        Trace[T](
             created=created,
             original_execution_time_ms=0,
             logged_values=X if isinstance(X, dict) else {"input": X},

great_ai/tracing/delete_ground_truth.py

@@ -1,5 +1,5 @@
 from datetime import datetime
-from typing import List, Optional, Union
+from typing import List, Optional, Union, cast
 
 from ..context import get_context
 
@@ -19,4 +19,4 @@ def delete_ground_truth(
         conjunctive_tags=tags, until=until, since=since, has_feedback=True
     )
 
-    db.delete_batch([i.trace_id for i in items])
+    db.delete_batch([cast(str, i.trace_id) for i in items])

great_ai/utilities/chunk.py

@@ -3,7 +3,7 @@ from typing import Iterable, List, TypeVar
 T = TypeVar("T")
 
 
-def chunk(values: Iterable[T], chunk_size: int) -> Iterable[T]:
+def chunk(values: Iterable[T], chunk_size: int) -> Iterable[List[T]]:
     assert chunk_size >= 1
 
     result: List[T] = []

great_ai/utilities/config_file/config_file.py

@@ -1,6 +1,6 @@
 import os
 from pathlib import Path
-from typing import Dict, Iterable, Tuple, Union
+from typing import Dict, ItemsView, Iterator, KeysView, Mapping, Union, ValuesView
 
 from ..logger import get_logger
 from .parse_error import ParseError
@@ -11,7 +11,7 @@ ENVIRONMENT_VARIABLE_KEY_PREFIX = "ENV"
 logger = get_logger("ConfigFile")
 
 
-class ConfigFile:
+class ConfigFile(Mapping[str, str]):
     def __init__(self, path: Union[Path, str]) -> None:
         if not isinstance(path, Path):
             path = Path(path)
@@ -24,7 +24,7 @@ class ConfigFile:
 
         self._parse()
 
-    def _parse(self):
+    def _parse(self) -> None:
         with open(self._path, encoding="utf-8") as f:
             lines: str = f.read()
 
@@ -72,20 +72,20 @@ class ConfigFile:
 
     __getitem__ = __getattr__
 
-    def __iter__(self) -> Iterable[Tuple[str, str]]:
+    def __iter__(self) -> Iterator[str]:
         return iter(self._key_values)
 
     def __len__(self) -> int:
         return len(self._key_values)
 
-    def keys(self):
+    def keys(self) -> KeysView[str]:
         return self._key_values.keys()
 
-    def values(self):
+    def values(self) -> ValuesView[str]:
         return self._key_values.values()
 
-    def items(self):
+    def items(self) -> ItemsView[str, str]:
         return self._key_values.items()
 
-    def __repr__(self):
+    def __repr__(self) -> str:
         return f"{type(self).__name__}(path={self._path}) {self._key_values}"

great_ai/utilities/evaluate_ranking/evaluate_ranking.py

@@ -1,5 +1,5 @@
 from pathlib import Path
-from typing import Dict, List, Optional, Union
+from typing import Dict, List, Optional, TypeVar
 
 import matplotlib
 import matplotlib.pyplot as plt
@@ -9,9 +9,11 @@ from sklearn.metrics import average_precision_score, precision_recall_curve
 from ..unique import unique
 from .draw_f1_iso_lines import draw_f1_iso_lines
 
+T = TypeVar("T", str, float)
+
 
 def evaluate_ranking(
-    expected: List[Union[str, float]],
+    expected: List[T],
     actual_scores: List[float],
     target_recall: float,
     title: Optional[str] = "",
@@ -20,7 +22,7 @@ def evaluate_ranking(
     output_svg: Optional[Path] = None,
     reverse_order: bool = False,
     plot: bool = True,
-) -> Dict[Union[str, float], float]:
+) -> Dict[T, float]:
     assert 0 <= target_recall <= 1
 
     if plot and axes is None:
@@ -39,7 +41,7 @@ def evaluate_ranking(
 
             draw_f1_iso_lines(axes=ax)
 
-        results: Dict[Union[str, float], float] = {}
+        results: Dict[T, float] = {}
         for i in range(len(classes) - 1):
             binarized_expected = [
                 (v < classes[i]) if reverse_order else (v > classes[i])

great_ai/utilities/parallel_map/get_config.py

@@ -17,14 +17,15 @@ def get_config(
 ) -> ParallelMapConfiguration:
 
     is_input_sequence = hasattr(input_values, "__len__")
+    input_length = len(input_values) if is_input_sequence else None  # type: ignore
 
     if concurrency is None:
         concurrency = os.cpu_count() or 1
     assert concurrency >= 1, "At least one mapper process has to be created"
 
     if chunk_size is None:
-        if is_input_sequence:
+        if input_length is not None:
-            chunk_size = max(1, ceil(len(input_values) / concurrency / 10))
+            chunk_size = max(1, ceil(input_length / concurrency / 10))
         else:
             raise ValueError(
                 "The argument for `values` does not implement `__len__`, therefore, you must provide a `chunk_size`"
@@ -32,19 +33,14 @@ def get_config(
     assert chunk_size >= 1, "Chunks have to contain at least one element"
 
     chunk_count: Optional[int] = None
-    if is_input_sequence:
+    if input_length is not None:
-        chunk_count = ceil(len(input_values) / chunk_size)
+        chunk_count = ceil(input_length / chunk_size)
         if chunk_count < concurrency:
             logger.warning(
                 f"Limiting concurrency to {chunk_count} because there are only {chunk_count} chunks"
             )
             concurrency = chunk_count
 
-    if concurrency == 1:
-        logger.warning("Running in series, there is no reason for parallelism")
-
-    input_length = len(input_values) if is_input_sequence else None
-
     config = ParallelMapConfiguration(
         concurrency=concurrency,
         chunk_count=chunk_count,

great_ai/utilities/parallel_map/manage_communication.py

@@ -1,10 +1,11 @@
 import multiprocessing as mp
 import queue
 import traceback
-from typing import Dict, Iterable, TypeVar, Union
+from typing import Dict, Iterable, List, TypeVar, Union
 
 from ..chunk import chunk
 from ..logger import get_logger
+from .map_result import MapResult
 from .worker_exception import WorkerException
 
 logger = get_logger("parallel_map")
@@ -27,6 +28,7 @@ def manage_communication(
     next_output_index = 0
     read_input_length = 0
     is_iteration_over = False
+
     while not is_iteration_over or next_output_index < read_input_length:
         if not is_iteration_over:
             try:
@@ -36,30 +38,30 @@ def manage_communication(
             except StopIteration:
                 is_iteration_over = True
             except Exception as e:
-                if not ignore_exceptions:
+                if ignore_exceptions:
-                    raise
-                else:
                     logger.error(
                         f"Exception {e} encountered in input, traceback:\n{traceback.format_exc()}"
                     )
+                else:
+                    raise
 
         try:
-            result_chunk = output_queue.get_nowait()
+            result_chunk: List[MapResult] = output_queue.get_nowait()
-
+            for r in result_chunk:
-            for index, value, exception in result_chunk:
+                if r.exception is not None:
-                if exception is not None:
+                    if ignore_exceptions:
-                    e, tb = exception
-                    if not ignore_exceptions:
-                        raise WorkerException from e
-                    else:
                         logger.error(
-                            f"Exception {e} encountered in worker, traceback:\n{tb}"
+                            f"Exception {r.exception} encountered in worker, traceback:\n{r.worker_traceback}"
                         )
+                    else:
+                        raise WorkerException from r.exception
+
                 if unordered:
-                    yield value
+
+                    yield r.value
                     next_output_index += 1
                 else:
-                    indexed_results[index] = value
+                    indexed_results[r.order] = r.value
 
             if not unordered:
                 while next_output_index in indexed_results:

great_ai/utilities/parallel_map/manage_serial.py

@@ -1,36 +0,0 @@
-import traceback
-from typing import Callable, Iterable, TypeVar
-
-from ..logger import get_logger
-from .worker_exception import WorkerException
-
-logger = get_logger("parallel_map")
-
-T = TypeVar("T")
-V = TypeVar("V")
-
-
-def manage_serial(
-    *,
-    function: Callable[[T], V],
-    input_values: Iterable[T],
-    ignore_exceptions: bool,
-) -> Iterable[V]:
-    try:
-        for v in input_values:
-            try:
-                yield function(v)
-            except Exception as e:
-                if not ignore_exceptions:
-                    raise WorkerException from e
-                else:
-                    logger.error(
-                        f"Exception {e} encountered in worker, traceback:\n{traceback.format_exc()}"
-                    )
-    except Exception as e:
-        if not ignore_exceptions:
-            raise
-        else:
-            logger.error(
-                f"Exception {e} encountered in input, traceback:\n{traceback.format_exc()}"
-            )

great_ai/utilities/parallel_map/map_result.py

@@ -0,0 +1,8 @@
+from typing import Any, NamedTuple, Optional
+
+
+class MapResult(NamedTuple):
+    order: int
+    value: Any
+    exception: Optional[Exception] = None
+    worker_traceback: Optional[str] = None

great_ai/utilities/parallel_map/mapper_function.py

@@ -1,42 +1,71 @@
+import asyncio
+import inspect
 import multiprocessing as mp
 import queue
 import threading
 import traceback
-from typing import Callable, Union
+from multiprocessing.synchronize import Event
+from typing import Any, Awaitable, Callable, List, TypeVar, Union, cast
 
 import dill
 
+from .map_result import MapResult
+
+T = TypeVar("T")
+V = TypeVar("V")
+
 
 def mapper_function(
     input_queue: Union[mp.Queue, queue.Queue],
     output_queue: Union[mp.Queue, queue.Queue],
-    should_stop: Union[mp.Event, threading.Event],
+    should_stop: Union[Event, threading.Event],
-    map_function: Union[bytes, Callable],
+    func: Union[bytes, Callable[[T], V], Callable[[T], Awaitable[V]]],
 ) -> None:
     try:
-        if isinstance(map_function, bytes):
+        if isinstance(func, bytes):
-            map_function = dill.loads(map_function)
+            func = cast(Callable[[T], V], dill.loads(func))
 
-        last_chunk = None
+        is_asynchronous = inspect.iscoroutinefunction(func)
+
+        last_chunk: List[MapResult] = []
         while not should_stop.wait(0.1):
-            if last_chunk is None:
+            if not last_chunk:
                 try:
                     input_chunk = input_queue.get_nowait()
-                    last_chunk = []
+                    if is_asynchronous:
-                    for i, v in input_chunk:
+
-                        result, exception = None, None
+                        async def safe(i: int, value: T) -> Any:
-                        try:
+                            try:
-                            result = map_function(v)
+                                return MapResult(
-                        except Exception as e:
+                                    i,
-                            exception = e, traceback.format_exc()
+                                    await cast(Callable[[T], Awaitable[V]], func)(
-                        last_chunk.append((i, result, exception))
+                                        value
+                                    ),
+                                )
+                            except Exception as e:
+                                return MapResult(i, None, e, traceback.format_exc())
+
+                        async def main() -> List[MapResult]:
+                            return await asyncio.gather(
+                                *[safe(i, v) for i, v in input_chunk]
+                            )
+
+                        last_chunk = asyncio.run(main())
+                    else:
+                        for i, value in input_chunk:
+                            try:
+                                last_chunk.append(MapResult(i, func(value)))
+                            except Exception as e:
+                                last_chunk.append(
+                                    MapResult(i, None, e, traceback.format_exc())
+                                )
                 except queue.Empty:
                     pass
 
-            if last_chunk is not None:
+            if last_chunk:
                 try:
                     output_queue.put_nowait(last_chunk)
-                    last_chunk = None
+                    last_chunk = []
                 except queue.Full:
                     pass
     except (KeyboardInterrupt, BrokenPipeError):

great_ai/utilities/parallel_map/parallel_map.py

@@ -1,11 +1,20 @@
 import multiprocessing as mp
-from typing import Callable, Iterable, Literal, Optional, Sequence, TypeVar, overload
+from typing import (
+    Awaitable,
+    Callable,
+    Iterable,
+    Literal,
+    Optional,
+    Sequence,
+    TypeVar,
+    Union,
+    overload,
+)
 
 import dill
 
 from .get_config import get_config
 from .manage_communication import manage_communication
-from .manage_serial import manage_serial
 from .mapper_function import mapper_function
 from .worker_exception import WorkerException
 
@@ -15,79 +24,72 @@ V = TypeVar("V")
 
 @overload
 def parallel_map(
-    function: Callable[[T], V],
+    func: Callable[[T], Union[V, Awaitable[V]]],
     input_values: Sequence[T],
     *,
-    chunk_size: Optional[int],
+    ignore_exceptions: Literal[True],
-    concurrency: Optional[int],
+    chunk_size: Optional[int] = ...,
-    unordered: Optional[bool],
+    concurrency: Optional[int] = ...,
-    ignore_exceptions: Optional[Literal[False]],
+    unordered: bool = ...,
-) -> Iterable[V]:
-    ...
-
-
-@overload
-def parallel_map(
-    function: Callable[[T], V],
-    input_values: Iterable[T],
-    *,
-    chunk_size: int,
-    concurrency: Optional[int],
-    unordered: Optional[bool],
-    ignore_exceptions: Optional[Literal[False]],
-) -> Iterable[V]:
-    ...
-
-
-@overload
-def parallel_map(
-    function: Callable[[T], V],
-    input_values: Sequence[T],
-    *,
-    chunk_size: Optional[int],
-    concurrency: Optional[int],
-    unordered: Optional[bool],
-    ignore_exceptions: True,
 ) -> Iterable[Optional[V]]:
     ...
 
 
 @overload
 def parallel_map(
-    function: Callable[[T], V],
+    func: Callable[[T], Union[V, Awaitable[V]]],
-    input_values: Iterable[T],
+    input_values: Union[Iterable[T], Sequence[T]],
     *,
     chunk_size: int,
-    concurrency: Optional[int],
+    ignore_exceptions: Literal[True],
-    unordered: Optional[bool],
+    concurrency: Optional[int] = ...,
-    ignore_exceptions: True,
+    unordered: bool = ...,
 ) -> Iterable[Optional[V]]:
     ...
 
 
+@overload
 def parallel_map(
-    function,
+    func: Callable[[T], Union[V, Awaitable[V]]],
-    input_values,
+    input_values: Sequence[T],
     *,
-    chunk_size=None,
+    chunk_size: Optional[int] = ...,
-    concurrency=None,
+    ignore_exceptions: Literal[False] = ...,
-    unordered=False,
+    concurrency: Optional[int] = ...,
-    ignore_exceptions=False,
+    unordered: bool = ...,
-):
+) -> Iterable[V]:
+    ...
+
+
+@overload
+def parallel_map(
+    func: Callable[[T], Union[V, Awaitable[V]]],
+    input_values: Union[Iterable[T], Sequence[T]],
+    *,
+    chunk_size: int,
+    ignore_exceptions: Literal[False] = ...,
+    concurrency: Optional[int] = ...,
+    unordered: bool = ...,
+) -> Iterable[V]:
+    ...
+
+
+def parallel_map(
+    func: Callable[[T], Union[V, Awaitable[V]]],
+    input_values: Union[Iterable[T], Sequence[T]],
+    *,
+    chunk_size: Optional[int] = None,
+    ignore_exceptions: bool = False,
+    concurrency: Optional[int] = None,
+    unordered: bool = False,
+) -> Iterable[Optional[V]]:
     config = get_config(
-        function=function,
+        function=func,
         input_values=input_values,
         chunk_size=chunk_size,
         concurrency=concurrency,
     )
 
-    if config.concurrency == 1:
-        yield from manage_serial(
-            function=function,
-            input_values=input_values,
-            ignore_exceptions=ignore_exceptions,
-        )
-
     ctx = (
         mp.get_context("fork")
         if "fork" in mp.get_all_start_methods()
@@ -99,7 +101,7 @@ def parallel_map(
     output_queue = manager.Queue(config.concurrency * 2)
 
     should_stop = ctx.Event()
-    serialized_map_function = dill.dumps(function, byref=True, recurse=False)
+    serialized_map_function = dill.dumps(func, byref=True, recurse=False)
 
     processes = [
         ctx.Process(
@@ -110,7 +112,7 @@ def parallel_map(
                 input_queue=input_queue,
                 output_queue=output_queue,
                 should_stop=should_stop,
-                map_function=serialized_map_function,
+                func=serialized_map_function,
             ),
         )
         for i in range(config.concurrency)

great_ai/utilities/parallel_map/threaded_parallel_map.py

@@ -1,10 +1,19 @@
 import queue
 import threading
-from typing import Callable, Iterable, Literal, Optional, Sequence, TypeVar, overload
+from typing import (
+    Awaitable,
+    Callable,
+    Iterable,
+    Literal,
+    Optional,
+    Sequence,
+    TypeVar,
+    Union,
+    overload,
+)
 
 from .get_config import get_config
 from .manage_communication import manage_communication
-from .manage_serial import manage_serial
 from .mapper_function import mapper_function
 
 T = TypeVar("T")
@@ -13,81 +22,74 @@ V = TypeVar("V")
 
 @overload
 def threaded_parallel_map(
-    function: Callable[[T], V],
+    func: Callable[[T], Union[V, Awaitable[V]]],
     input_values: Sequence[T],
     *,
-    chunk_size: Optional[int],
+    ignore_exceptions: Literal[True],
-    concurrency: Optional[int],
+    chunk_size: Optional[int] = ...,
-    unordered: Optional[bool],
+    concurrency: Optional[int] = ...,
-    ignore_exceptions: Optional[Literal[False]],
+    unordered: bool = ...,
-) -> Iterable[V]:
-    ...
-
-
-@overload
-def threaded_parallel_map(
-    function: Callable[[T], V],
-    input_values: Iterable[T],
-    *,
-    chunk_size: int,
-    concurrency: Optional[int],
-    unordered: Optional[bool],
-    ignore_exceptions: Optional[Literal[False]],
-) -> Iterable[V]:
-    ...
-
-
-@overload
-def threaded_parallel_map(
-    function: Callable[[T], V],
-    input_values: Sequence[T],
-    *,
-    chunk_size: Optional[int],
-    concurrency: Optional[int],
-    unordered: Optional[bool],
-    ignore_exceptions: True,
 ) -> Iterable[Optional[V]]:
     ...
 
 
 @overload
 def threaded_parallel_map(
-    function: Callable[[T], V],
+    func: Callable[[T], Union[V, Awaitable[V]]],
-    input_values: Iterable[T],
+    input_values: Union[Iterable[T], Sequence[T]],
     *,
     chunk_size: int,
-    concurrency: Optional[int],
+    ignore_exceptions: Literal[True],
-    unordered: Optional[bool],
+    concurrency: Optional[int] = ...,
-    ignore_exceptions: True,
+    unordered: bool = ...,
 ) -> Iterable[Optional[V]]:
     ...
 
 
+@overload
 def threaded_parallel_map(
-    function,
+    func: Callable[[T], Union[V, Awaitable[V]]],
-    input_values,
+    input_values: Sequence[T],
     *,
-    chunk_size=None,
+    chunk_size: Optional[int] = ...,
-    concurrency=None,
+    ignore_exceptions: Literal[False] = ...,
-    unordered=False,
+    concurrency: Optional[int] = ...,
-    ignore_exceptions=False,
+    unordered: bool = ...,
-):
+) -> Iterable[V]:
+    ...
+
+
+@overload
+def threaded_parallel_map(
+    func: Callable[[T], Union[V, Awaitable[V]]],
+    input_values: Union[Iterable[T], Sequence[T]],
+    *,
+    chunk_size: int,
+    ignore_exceptions: Literal[False] = ...,
+    concurrency: Optional[int] = ...,
+    unordered: bool = ...,
+) -> Iterable[V]:
+    ...
+
+
+def threaded_parallel_map(
+    func: Callable[[T], Union[V, Awaitable[V]]],
+    input_values: Union[Iterable[T], Sequence[T]],
+    *,
+    chunk_size: Optional[int] = None,
+    ignore_exceptions: bool = False,
+    concurrency: Optional[int] = None,
+    unordered: bool = False,
+) -> Iterable[Optional[V]]:
     config = get_config(
-        function=function,
+        function=func,
         input_values=input_values,
         chunk_size=chunk_size,
         concurrency=concurrency,
     )
 
-    if config.concurrency == 1:
+    input_queue: queue.Queue = queue.Queue(config.concurrency * 2)
-        yield from manage_serial(
+    output_queue: queue.Queue = queue.Queue(config.concurrency * 2)
-            function=function,
-            input_values=input_values,
-            ignore_exceptions=ignore_exceptions,
-        )
-
-    input_queue = queue.Queue(config.concurrency * 2)
-    output_queue = queue.Queue(config.concurrency * 2)
     should_stop = threading.Event()
 
     threads = [
@@ -99,7 +101,7 @@ def threaded_parallel_map(
                 input_queue=input_queue,
                 output_queue=output_queue,
                 should_stop=should_stop,
-                map_function=function,
+                func=func,
             ),
         )
         for i in range(config.concurrency)

great_ai/views/function_metadata.py

@@ -4,6 +4,7 @@ from pydantic import BaseModel
 
 
 class FunctionMetadata(BaseModel):
+    is_asynchronous: bool
     input_parameter_names: List[str] = []
     model_parameter_names: List[str] = []
     is_finalised: bool = False

great_ai/views/trace.py

@@ -25,10 +25,10 @@ class Trace(Generic[T], HashableBaseModel):
         extra = Extra.ignore
 
     @validator("trace_id", always=True)
-    def generate_id(cls, v: Optional[str], values: Dict[str, Any]) -> Optional[str]:
+    def generate_id(cls, v: Optional[str], values: Dict[str, Any]) -> str:
-        if not v:
+        if v:
-            return str(uuid4())
+            return v
-        return v
+        return str(uuid4())
 
     @property
     def input(self) -> Any:

pyproject.toml

@@ -36,6 +36,7 @@ dependencies = [
     "typeguard >= 2.10.0",
     "pymongo >= 3.0.0",
     "dill >= 0.3.5.0",
+    "async_lru >= 1.0.0",
     "aiohttp[speedups] >= 3.8.0",
 ]
 
@@ -45,6 +46,7 @@ dev = [
     "mkdocs",
     "mkdocstrings[python]",
     "mkdocs-material",
+    "mkdocs-jupyter",
     "autoflake",
     "isort",
     "black[jupyter]",
@@ -53,7 +55,7 @@ dev = [
     "pytest",
     "pytest-cov",
     "pytest-subtests",
-    "pytest-asyncio"
+    "pytest-asyncio",
 ]
 
 [project.urls]

scripts/format-python.sh

@@ -26,5 +26,7 @@ for dir in "$@"; do
         python3 -m mypy --namespace-packages --ignore-missing-imports --install-types --non-interactive --disallow-untyped-defs --disallow-incomplete-defs --follow-imports=silent --exclude=external/ --exclude=/build/ --pretty . || true
     fi
 
+    python3 -m flake8 . --count --show-source --statistics --exclude=__init__.py,.env,external --ignore=E501,E402,F821,W503,E722,E203
+
     cd -
 done

tests/great_ai/test_great_ai.py

@@ -1,9 +0,0 @@
-from great_ai import GreatAI
-
-
-def test_process_batch() -> None:
-    @GreatAI.create(return_raw_result=True)
-    def f(x):
-        return x + 2
-
-    assert f.process_batch([3, 9, 34]) == [5, 11, 36]

tests/test_async_starters.py

@@ -1,13 +1,7 @@
 from asyncio import sleep
 
 import pytest
-
+from great_ai import GreatAI, WrongDecoratorOrderError, parameter
-from great_ai import (
-    ArgumentValidationError,
-    GreatAI,
-    WrongDecoratorOrderError,
-    parameter,
-)
 
 
 @pytest.mark.asyncio
@@ -17,28 +11,28 @@ async def test_create_trivial_cases() -> None:
         await sleep(0.5)
         return f"Hello {name}!"
 
-    assert (await hello_world_1("andras").output) == "Hello andras!"
+    assert (await hello_world_1("andras")).output == "Hello andras!"
 
     @GreatAI.create
     async def hello_world_2(name: str) -> str:
         await sleep(0.5)
         return f"Hello {name}!"
 
-    assert (await hello_world_2("andras").output) == "Hello andras!"
+    assert (await hello_world_2("andras")).output == "Hello andras!"
 
     @GreatAI.create()
     async def hello_world_3(name: str) -> str:
         await sleep(0.5)
         return f"Hello {name}!"
 
-    assert (await hello_world_3("andras").output) == "Hello andras!"
+    assert (await hello_world_3("andras")).output == "Hello andras!"
 
     @GreatAI.create()
-    async def hello_world_4(name):
+    async def hello_world_4(name):  # type: ignore
         await sleep(0.5)
         return f"Hello {name}!"
 
-    assert (await hello_world_4("andras").output) == "Hello andras!"
+    assert (await hello_world_4("andras")).output == "Hello andras!"
 
 
 @pytest.mark.asyncio
@@ -49,14 +43,23 @@ async def test_with_parameter() -> None:
         await sleep(0.5)
         return f"Hello {name}!"
 
-    assert (await hello_world("andras").output) == "Hello andras!"
 
-    with pytest.raises(ArgumentValidationError):
+@pytest.mark.asyncio
-        await hello_world("short")
+async def test_with_parameters() -> None:
+    @GreatAI.create
+    @parameter("name", validator=lambda v: len(v) > 5)
+    @parameter("unused", disable_logging=True)
+    async def hello_world(name: str, unused) -> str:  # type: ignore
+        await sleep(0.5)
+        return f"Hello {name}!"
 
+    assert (await hello_world("andras", "fr")).output == "Hello andras!"
+
+
+def test_wrong_order() -> None:
     with pytest.raises(WrongDecoratorOrderError):
 
         @parameter("name", validator=lambda v: len(v) > 5)
         @GreatAI.create
-        def hello_world(name: str) -> str:
+        async def hello_world(name: str) -> str:
             return f"Hello {name}!"

tests/test_basic_starters.py

@@ -1,7 +1,6 @@
 from functools import lru_cache
 
 import pytest
-
 from great_ai import (
     ArgumentValidationError,
     GreatAI,
@@ -18,7 +17,7 @@ def test_create_trivial_cases() -> None:
     assert hello_world_1("andras").output == "Hello andras!"
 
     @GreatAI.create
-    def hello_world_2(name):
+    def hello_world_2(name):  # type: ignore
         return f"Hello {name}!"
 
     assert hello_world_2("andras").output == "Hello andras!"
@@ -30,7 +29,7 @@ def test_create_trivial_cases() -> None:
     assert hello_world_3("andras").output == "Hello andras!"
 
     @GreatAI.create()
-    def hello_world_4(name):
+    def hello_world_4(name):  # type: ignore
         return f"Hello {name}!"
 
     assert hello_world_4("andras").output == "Hello andras!"
@@ -39,17 +38,17 @@ def test_create_trivial_cases() -> None:
 def test_create_with_other_decorator() -> None:
     @GreatAI.create
     @lru_cache
-    def hello_world(name: str) -> str:
+    def hello_world_1(name: str) -> str:
         return f"Hello {name}!"
 
-    assert hello_world("andras").output == "Hello andras!"
+    assert hello_world_1("andras").output == "Hello andras!"
 
     @lru_cache
     @GreatAI.create()
-    def hello_world(name: str) -> str:
+    def hello_world_2(name: str) -> str:
         return f"Hello {name}!"
 
-    assert hello_world("andras").output == "Hello andras!"
+    assert hello_world_2("andras").output == "Hello andras!"
 
 
 def test_with_parameter() -> None:
@@ -63,6 +62,8 @@ def test_with_parameter() -> None:
     with pytest.raises(ArgumentValidationError):
         hello_world("short")
 
+
+def test_wrong_order() -> None:
     with pytest.raises(WrongDecoratorOrderError):
 
         @parameter("name", validator=lambda v: len(v) > 5)

tests/test_great_ai.py

@@ -0,0 +1,22 @@
+from asyncio import sleep
+
+import pytest
+from great_ai import GreatAI
+
+
+def test_process_batch() -> None:
+    @GreatAI.create
+    def f(x: int) -> int:
+        return x + 2
+
+    assert [v.output for v in f.process_batch([3, 9, 34])] == [5, 11, 36]
+
+
+@pytest.mark.asyncio
+async def test_process_batch_async() -> None:
+    @GreatAI.create
+    async def f(x: int) -> int:
+        await sleep(0.2)
+        return x + 2
+
+    assert [v.output for v in f.process_batch([3, 9, 34])] == [5, 11, 36]

tests/utilities/test_chunk.py

@@ -1,5 +1,4 @@
 import pytest
-
 from great_ai.utilities import chunk
 
 
@@ -23,7 +22,7 @@ def test_bad_argument() -> None:
 
 
 def test_generator() -> None:
-    def my_generator():
+    def my_generator():  # type: ignore
         for i in range(1, 11):
             yield i
 

tests/utilities/test_config_file.py

@@ -2,7 +2,6 @@ import os
 from pathlib import Path
 
 import pytest
-
 from great_ai.utilities import ConfigFile
 
 DATA_PATH = Path(__file__).parent.resolve() / "data"

tests/utilities/test_language.py

@@ -20,11 +20,11 @@ def test_is_english() -> None:
     assert not is_english(None)
 
 
-def english_name_of_language() -> None:
+def test_english_name_of_language() -> None:
     assert english_name_of_language("en") == "English"
     assert english_name_of_language("hu") == "Hungarian"
     assert english_name_of_language("zh") == "Chinese"
-    assert english_name_of_language("zh-TW") == "Chinese"
+    assert english_name_of_language("zh-TW") == "Chinese (Taiwan)"
     assert english_name_of_language("und") == "Unknown language"
     assert english_name_of_language("") == "Unknown language"
     assert english_name_of_language(None) == "Unknown language"

tests/utilities/test_parallel_map.py

@@ -1,11 +1,13 @@
-import pytest
+from typing import Any, Iterable
 
+import pytest
 from great_ai.utilities import WorkerException, parallel_map
+from typing_extensions import Never
 
 COUNT = int(1e5) + 3
 
 
-def test_simple_case_with_progress_bar() -> None:
+def test_simple_case() -> None:
     assert list(parallel_map(lambda v: v**2, range(COUNT), concurrency=4)) == [
         v**2 for v in range(COUNT)
     ]
@@ -14,7 +16,7 @@ def test_simple_case_with_progress_bar() -> None:
 def test_with_iterable() -> None:
     from time import sleep
 
-    def my_generator():
+    def my_generator() -> Iterable[int]:
         for i in range(10):
             yield i
             sleep(0.1)
@@ -26,12 +28,6 @@ def test_with_iterable() -> None:
     )
 
 
-def test_simple_case_without_progress_bar() -> None:
-    assert list(parallel_map(lambda v: v**2, range(COUNT), concurrency=2)) == [
-        v**2 for v in range(COUNT)
-    ]
-
-
 def test_simple_case_invalid_values() -> None:
     with pytest.raises(AssertionError):
         list(parallel_map(lambda v: v**2, range(COUNT), concurrency=0))
@@ -41,7 +37,7 @@ def test_simple_case_invalid_values() -> None:
 
 
 def test_this_process_exception() -> None:
-    def my_generator():
+    def my_generator() -> Iterable[int]:
         yield 1
         yield 2
         yield 3
@@ -59,7 +55,7 @@ def test_this_process_exception() -> None:
 
 
 def test_ignore_this_process_exception() -> None:
-    def my_generator():
+    def my_generator() -> Iterable[float]:
         yield 1
         yield 2
         yield 3
@@ -74,19 +70,10 @@ def test_ignore_this_process_exception() -> None:
             ignore_exceptions=True,
         )
     ) == [1, 4]
-    assert list(
-        parallel_map(
-            lambda v: v**2,
-            my_generator(),
-            concurrency=1,
-            chunk_size=2,
-            ignore_exceptions=True,
-        )
-    ) == [1, 4, 9]
 
 
 def test_worker_process_exception() -> None:
-    def oh_no(_):
+    def oh_no(_: Any) -> Never:
         raise ValueError("hi")
 
     with pytest.raises(WorkerException):
@@ -97,7 +84,7 @@ def test_worker_process_exception() -> None:
 
 
 def test_ignore_worker_process_exception() -> None:
-    def oh_no(_):
+    def oh_no(_: Any) -> Never:
         raise ValueError("hi")
 
     assert (

tests/utilities/test_threaded_parallel_map.py

@@ -1,11 +1,13 @@
-import pytest
+from typing import Any, Iterable
 
+import pytest
 from great_ai.utilities import WorkerException, threaded_parallel_map
+from typing_extensions import Never
 
 COUNT = int(1e5) + 3
 
 
-def test_simple_case_with_progress_bar() -> None:
+def test_simple_case() -> None:
     assert list(
         threaded_parallel_map(lambda v: v**2, range(COUNT), concurrency=4)
     ) == [v**2 for v in range(COUNT)]
@@ -14,7 +16,7 @@ def test_simple_case_with_progress_bar() -> None:
 def test_with_iterable() -> None:
     from time import sleep
 
-    def my_generator():
+    def my_generator() -> Iterable[int]:
         for i in range(10):
             yield i
             sleep(0.1)
@@ -27,12 +29,6 @@ def test_with_iterable() -> None:
     )
 
 
-def test_simple_case_without_progress_bar() -> None:
-    assert list(
-        threaded_parallel_map(lambda v: v**2, range(COUNT), concurrency=2)
-    ) == [v**2 for v in range(COUNT)]
-
-
 def test_simple_case_invalid_values() -> None:
     with pytest.raises(AssertionError):
         list(threaded_parallel_map(lambda v: v**2, range(COUNT), concurrency=0))
@@ -42,7 +38,7 @@ def test_simple_case_invalid_values() -> None:
 
 
 def test_this_worker_exception() -> None:
-    def my_generator():
+    def my_generator() -> Iterable[int]:
         yield 1
         yield 2
         yield 3
@@ -64,7 +60,7 @@ def test_this_worker_exception() -> None:
 
 
 def test_ignore_this_worker_exception() -> None:
-    def my_generator():
+    def my_generator() -> Iterable[float]:
         yield 1
         yield 2
         yield 3
@@ -78,20 +74,14 @@ def test_ignore_this_worker_exception() -> None:
             chunk_size=2,
             ignore_exceptions=True,
         )
-    ) == [1, 4]
+    ) == [
-    assert list(
+        1,
-        threaded_parallel_map(
+        4,
-            lambda v: v**2,
+    ]  # the second chunk is ruined because of the error
-            my_generator(),
-            concurrency=1,
-            chunk_size=2,
-            ignore_exceptions=True,
-        )
-    ) == [1, 4, 9]
 
 
 def test_worker_worker_exception() -> None:
-    def oh_no(_):
+    def oh_no(_: Any) -> Never:
         raise ValueError("hi")
 
     with pytest.raises(WorkerException):
@@ -102,7 +92,7 @@ def test_worker_worker_exception() -> None:
 
 
 def test_ignore_worker_worker_exception() -> None:
-    def oh_no(_):
+    def oh_no(_: Any) -> Never:
         raise ValueError("hi")
 
     assert (