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]],
+            @overload
-                partial(
+            def inner(func: Awaitable[V]) -> GreatAI[Awaitable[Trace[V]], V]:
-                    GreatAI.create,
+                ...
+
+            @overload
+            def inner(func: Callable[..., V]) -> GreatAI[Trace[V], V]:
+                ...
+
+            def inner(func):  # type: ignore
+                return GreatAI.create(
+                    func,
                     version=version,
-                    return_raw_result=return_raw_result,
+                    route_config=route_config,
-                    disable_rest_api=disable_rest_api,
-                    disable_docs=disable_docs,
-                    disable_dashboard=disable_dashboard,
-                ),
                 )
 
-        instance = GreatAI[T](
+            return inner
-            func, version=version, return_raw_result=return_raw_result
+
+        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(
-                        self._cached_func, do_not_persist_traces=do_not_persist_traces
+                    self._wrapped_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)
 
+        if route_config.feedback_endpoints_enabled:
             bootstrap_feedback_endpoints(self.app)
-        self._bootstrap_meta_endpoints()
 
-    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__,
-
-        @router.post(
-            "/predict", status_code=status.HTTP_200_OK, response_model=Trace[T]
-        )
-        @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,
+                    documentation=self.__doc__,
                     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)