Add auto-parameter support & refactor

great_ai/setup.cfg

@@ -38,7 +38,7 @@ install_requires =
     pyaml >= 21.0.0
 
 [options.package_data]
-* = *.json, *.yaml, *.yml
+* = *.json, *.yaml, *.yml, *.css
 
 [options.packages.find]
 where = src

great_ai/src/great_ai/great_ai/__init__.py

@@ -1,4 +1,5 @@
 from .context import configure
-from .deploy import create_service, process_batch, process_single
-from .metrics import log_argument, log_metric
+from .deploy import GreatAI
+from .exceptions import ArgumentValidationError, MissingArgumentError
 from .models import save_model, use_model
+from .parameters import log_metric, parameter

great_ai/src/great_ai/great_ai/dashboard/__init__.py

@@ -0,0 +1 @@
+from .create_dash_app import create_dash_app

great_ai/src/great_ai/great_ai/dashboard/assets/index.css

@@ -0,0 +1,112 @@
+:root {
+  --accent-color: #47c2d0;
+  --error-color: #a30808;
+  --background-color: #edf5f6;
+  --small-padding: 10px;
+  --medium-padding: 20px;
+  --large-padding: 40px;
+  --border-radius: 10px;
+  --shadow: 0 4px 6px -1px rgb(0 0 0 / 10%), 0 2px 4px -1px rgb(0 0 0 / 6%);
+}
+
+* {
+  margin: 0;
+  box-sizing: border-box;
+}
+
+body {
+  background-color: var(--background-color);
+  font-family: Arial, Helvetica, sans-serif;
+}
+
+h1,
+h2,
+h3,
+h4,
+h5,
+h6 {
+  margin: var(--medium-padding) 0 var(--small-padding) 0;
+}
+
+.glance,
+.table-container,
+.parallel-coords {
+  padding: var(--large-padding);
+  margin: var(--large-padding);
+  border-radius: var(--border-radius);
+  box-shadow: var(--shadow);
+  background-color: white;
+  overflow: hidden;
+}
+
+.glance {
+  display: flex;
+}
+
+.glance .description {
+  width: 350px;
+  display: flex;
+  flex-direction: column;
+  justify-content: center;
+}
+
+.glance .dash-graph {
+  flex: 1;
+}
+
+.table-container {
+  padding: 0;
+}
+
+.table-container h2 {
+  padding: var(--small-padding);
+  margin: 0;
+}
+
+.table-container h2,
+footer.watermark {
+  opacity: 0.35;
+}
+
+.dash-spreadsheet {
+  overflow: auto;
+}
+
+th,
+td {
+  max-width: 350px;
+}
+
+.parallel-coords {
+  padding: 0;
+}
+
+footer.watermark {
+  display: flex;
+  justify-content: space-between;
+  align-items: center;
+  padding: var(--large-padding);
+  background-color: #ddd;
+  position: relative;
+}
+
+footer.watermark div {
+  display: flex;
+  flex-direction: column;
+}
+
+h6 {
+  font-size: 3rem;
+  margin-top: 0;
+}
+
+a img {
+  width: 64px;
+  height: 64px;
+  cursor: pointer;
+  transition: transform 300ms;
+}
+
+a img:hover {
+  transform: scale(1.1);
+}

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

@@ -0,0 +1,192 @@
+from typing import Any, Dict, List
+
+import pandas as pd
+import plotly.express as px
+import plotly.graph_objects as go
+from dash import Dash, dash_table, dcc, html
+from dash.dependencies import Input, Output
+from flask import Flask
+
+from great_ai.utilities.unique import unique
+
+from ..context import get_context
+from ..helper import snake_case_to_text, text_to_hex_color
+from ..views import SortBy
+from .get_description import get_description
+from .get_filter_from_datatable import get_filter_from_datatable
+from .get_footer import get_footer
+
+
+def create_dash_app(function_name: str, function_docs: str) -> Flask:
+    accent_color = text_to_hex_color(function_name)
+
+    flask_app = Flask(__name__)
+    app = Dash(
+        function_name,
+        requests_pathname_prefix=get_context().metrics_path + "/",
+        server=flask_app,
+        title=snake_case_to_text(function_name),
+        update_title=None,
+        external_stylesheets=[
+            "/assets/index.css",
+        ],
+    )
+
+    documents = get_context().persistence.get_documents()
+    df = pd.DataFrame(documents)
+
+    execution_time_histogram = dcc.Graph(config={"displaylogo": False})
+    table = dash_table.DataTable(
+        columns=[{"name": i, "id": i} for i in df.columns],
+        page_current=0,
+        page_size=20,
+        page_action="custom",
+        filter_action="custom",
+        filter_query="",
+        sort_action="custom",
+        sort_mode="multi",
+        sort_by=[
+            {"column_id": "created", "direction": "desc"},
+        ],
+    )
+
+    app.layout = html.Div(
+        [
+            html.Div(
+                [
+                    get_description(
+                        function_name=function_name,
+                        function_docs=function_docs,
+                        accent_color=accent_color,
+                    ),
+                    execution_time_histogram,
+                ],
+                className="glance",
+            ),
+            html.Div([html.H2("Latest traces"), table], className="table-container"),
+            parallel_coords := dcc.Graph(
+                className="parallel-coords", config={"displaylogo": False}
+            ),
+            get_footer(),
+            interval := dcc.Interval(
+                interval=4 * 1000,  # in milliseconds
+            ),
+        ]
+    )
+
+    @app.callback(
+        Output(table, "data"),
+        Input(table, "page_current"),
+        Input(table, "page_size"),
+        Input(table, "sort_by"),
+        Input(table, "filter_query"),
+        Input(interval, "n_intervals"),
+    )
+    def update_table(
+        page_current: int,
+        page_size: int,
+        sort_by: List[SortBy],
+        filter: str,
+        n_intervals: int,
+    ) -> List[Dict[str, Any]]:
+        conjunctive_filters = [
+            get_filter_from_datatable(f) for f in filter.split(" && ")
+        ]
+        non_null_conjunctive_filters = [f for f in conjunctive_filters if f is not None]
+
+        return get_context().persistence.query(
+            conjunctive_filters=non_null_conjunctive_filters,
+            sort_by=sort_by,
+            skip=page_current * page_size,
+            take=page_size,
+        )
+
+    @app.callback(
+        Output(execution_time_histogram, "figure"),
+        Input(table, "filter_query"),
+        Input(interval, "n_intervals"),
+    )
+    def update_execution_times(filter: str, _n_intervals: int) -> go.Figure:
+        conjunctive_filters = [
+            get_filter_from_datatable(f) for f in filter.split(" && ")
+        ]
+        non_null_conjunctive_filters = [f for f in conjunctive_filters if f is not None]
+
+        rows = get_context().persistence.query(
+            conjunctive_filters=non_null_conjunctive_filters
+        )
+
+        if not rows:
+            return go.Figure()
+
+        df = pd.DataFrame(rows)
+
+        fig = px.histogram(
+            df,
+            x="execution_time_ms",
+            labels={"execution_time_ms": "Execution time (ms)"},
+            nbins=20,
+            height=400,
+            log_y=True,
+            color_discrete_sequence=[accent_color],
+        )
+
+        fig.update_layout(
+            autosize=True,
+            margin=dict(l=0, r=0, b=0, t=0, pad=0),
+        )
+
+        return fig
+
+    @app.callback(
+        Output(parallel_coords, "figure"),
+        Input(table, "filter_query"),
+        Input(interval, "n_intervals"),
+    )
+    def update_parallel_coords(filter: str, _n_intervals: int) -> go.Figure:
+        conjunctive_filters = [
+            get_filter_from_datatable(f) for f in filter.split(" && ")
+        ]
+        non_null_conjunctive_filters = [f for f in conjunctive_filters if f is not None]
+
+        rows = get_context().persistence.query(
+            conjunctive_filters=non_null_conjunctive_filters
+        )
+
+        if not rows:
+            return go.Figure()
+
+        df = pd.DataFrame(rows)
+        return go.Figure(
+            go.Parcoords(
+                dimensions=[
+                    get_dimension_descriptor(df, c)
+                    for c in df.columns
+                    if c not in {"id", "created", "output"}
+                ],
+                line_color=accent_color,
+            )
+        )
+
+    return flask_app
+
+
+def get_dimension_descriptor(df: pd.DataFrame, column: str) -> Dict[str, Any]:
+    dimension: Dict[str, Any] = {
+        "label": snake_case_to_text(column),
+    }
+
+    values = df[column]
+
+    try:
+        dimension["values"] = [float(v) for v in values]
+    except (TypeError, ValueError):
+        MAX_LENGTH = 40
+        unique_values = unique(values)
+        value_mapping = {str(v)[-MAX_LENGTH:]: i for i, v in enumerate(unique_values)}
+
+        dimension["values"] = [value_mapping[str(v)[-MAX_LENGTH:]] for v in values]
+        dimension["tickvals"] = list(value_mapping.values())
+        dimension["ticktext"] = [k[-MAX_LENGTH:] for k in value_mapping.keys()]
+
+    return dimension

great_ai/src/great_ai/great_ai/dashboard/get_description.py

@@ -0,0 +1,36 @@
+from dash import dcc, html
+
+from ..helper import snake_case_to_text, strip_lines
+
+
+def get_description(
+    function_name: str, function_docs: str, accent_color: str
+) -> html.Div:
+    return html.Div(
+        [
+            html.H1(
+                f"{snake_case_to_text(function_name)} - metrics",
+                style={"color": accent_color},
+            ),
+            dcc.Markdown(
+                strip_lines(
+                    f"""
+                > View the live data of your deployments here.
+
+                ## Using the API
+
+                You can find the available endpoints at [/docs](/docs).
+
+                ### Details
+
+                {function_docs}
+
+                ## Metrics
+
+                Recent traces and aggregated metrics are presented below. Try filtering the table.
+                """
+                ),
+                className="description",
+            ),
+        ]
+    )

great_ai/src/great_ai/great_ai/dashboard/get_filter_from_datatable.py

@@ -0,0 +1,23 @@
+from typing import Optional, Union
+
+from ..views import Filter, operators
+
+
+def get_filter_from_datatable(description: str) -> Optional[Filter]:
+    for operator in operators:
+        if operator in description:
+            name_part, value_part = description.split(operator, 1)
+            value_part = value_part.strip()
+            name_part = name_part[name_part.find("{") + 1 : name_part.rfind("}")]
+
+            v0 = value_part[0]
+            if v0 == value_part[-1] and v0 in ("'", '"', "`"):
+                value: Union[str, float] = value_part[1:-1].replace("\\" + v0, v0)
+            else:
+                try:
+                    value = float(value_part)
+                except ValueError:
+                    value = value_part
+            return Filter(property=name_part, operator=operator, value=value)
+
+    return None

great_ai/src/great_ai/great_ai/dashboard/get_footer.py

@@ -0,0 +1,22 @@
+from dash import html
+
+
+def get_footer() -> html.Footer:
+    return html.Footer(
+        [
+            html.Div(
+                [
+                    html.H6("GreatAI"),
+                    html.P(
+                        "A human-friendly framework for robust end-to-end AI deployments."
+                    ),
+                ]
+            ),
+            html.A(
+                html.Img(src="/assets/github.png"),
+                href="https://github.com/ScoutinScience/great-ai",
+                target="_blank",
+            ),
+        ],
+        className="watermark",
+    )

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

@@ -1,3 +1 @@
-from .create_service import create_service
-from .process_batch import process_batch
-from .process_single import process_single
+from .great_ai import GreatAI

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

@@ -1,94 +0,0 @@
-from pathlib import Path
-from typing import Any, Callable, Dict, List
-
-from fastapi import FastAPI, HTTPException, status
-from fastapi.middleware.wsgi import WSGIMiddleware
-from fastapi.openapi.docs import get_swagger_ui_html
-from fastapi.responses import RedirectResponse
-from fastapi.staticfiles import StaticFiles
-from starlette.responses import HTMLResponse
-
-from ..context import get_context
-from ..helper import snake_case_to_text
-from ..metrics import create_dash_app
-from ..tracing import TracingContext
-from ..views import EvaluationFeedbackRequest, HealthCheckResponse, Query, Trace
-
-PATH = Path(__file__).parent.resolve()
-
-
-def create_service(
-    func: Callable[..., Any], disable_docs: bool = False, disable_metrics: bool = False
-) -> FastAPI:
-    function_name = func.__name__
-    function_docs = func.__doc__
-
-    documentation = (
-        f"REST API wrapper for interacting with the '{function_name}' function.\n"
-    )
-    if function_docs:
-        documentation += function_docs
-
-    app = FastAPI(
-        title=snake_case_to_text(function_name),
-        description=documentation,
-        docs_url=None,
-        redoc_url=None,
-    )
-
-    @app.post("/evaluations", status_code=status.HTTP_200_OK, response_model=Trace)
-    def score(input: Any) -> Trace:
-        with TracingContext() as t:
-            result = func(input)
-            output = t.log_output(result)
-        return output
-
-    @app.get("/evaluations/:evaluation_id", status_code=status.HTTP_200_OK)
-    def get_evaluation(evaluation_id: str) -> Trace:
-        result = get_context().persistence.get_trace(evaluation_id)
-        if result is None:
-            raise HTTPException(status_code=status.HTTP_404_NOT_FOUND)
-        return result
-
-    @app.post(
-        "/evaluations/:evaluation_id/feedback", status_code=status.HTTP_202_ACCEPTED
-    )
-    def give_feedback(evaluation_id: str, input: EvaluationFeedbackRequest) -> None:
-        get_context().persistence.add_evaluation(evaluation_id, input.evaluation)
-
-    if not disable_docs:
-
-        @app.get("/docs", include_in_schema=False)
-        def custom_swagger_ui_html() -> HTMLResponse:
-            return get_swagger_ui_html(openapi_url="openapi.json", title=app.title)
-
-        @app.get("/docs/index.html", include_in_schema=False)
-        def redirect_to_entrypoint() -> RedirectResponse:
-            return RedirectResponse("/docs")
-
-    if not disable_metrics:
-        dash_app = create_dash_app(function_name, documentation)
-        app.mount(get_context().metrics_path, WSGIMiddleware(dash_app))
-
-        @app.get("/", include_in_schema=False)
-        def redirect_to_entrypoint() -> RedirectResponse:
-            return RedirectResponse("/metrics")
-
-        app.mount(
-            "/assets", StaticFiles(directory=PATH / "../metrics/assets"), name="static"
-        )
-
-        @app.post("/query", status_code=status.HTTP_200_OK)
-        def query_metrics(query: Query) -> List[Dict[str, Any]]:
-            return get_context().persistence.query(
-                conjunctive_filters=query.filter,
-                sort_by=query.sort,
-                skip=query.skip,
-                take=query.take,
-            )
-
-    @app.get("/health", status_code=status.HTTP_200_OK)
-    def check_health() -> HealthCheckResponse:
-        return HealthCheckResponse(is_healthy=True)
-
-    return app

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

@@ -0,0 +1,145 @@
+import inspect
+from pathlib import Path
+from typing import Any, Callable, Dict, Iterable, List, Optional, Sequence, Type, cast
+
+from fastapi import FastAPI, HTTPException, status
+from fastapi.middleware.wsgi import WSGIMiddleware
+from fastapi.openapi.docs import get_swagger_ui_html
+from fastapi.responses import RedirectResponse
+from fastapi.staticfiles import StaticFiles
+from pydantic import BaseModel, create_model
+from starlette.responses import HTMLResponse
+
+from great_ai.utilities.parallel_map import parallel_map
+
+from ..context import get_context
+from ..dashboard import create_dash_app
+from ..helper import get_function_metadata_store, snake_case_to_text
+from ..parameters import automatically_decorate_parameters
+from ..tracing import TracingContext
+from ..views import EvaluationFeedbackRequest, HealthCheckResponse, Query, Trace
+
+PATH = Path(__file__).parent.resolve()
+
+
+class GreatAI(FastAPI):
+    def __init__(self, func: Callable[..., Any], *args: Any, **kwargs: Any):
+        self._func = automatically_decorate_parameters(func)
+
+        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
+
+        def process_single(input_value: schema) -> Trace:  # type: ignore
+            with TracingContext() as t:
+                result = self._func(**cast(BaseModel, input_value).dict())
+                output = t.log_output(result)
+            return output
+
+        self.process_single = process_single
+
+        super().__init__(
+            *args,
+            title=snake_case_to_text(func.__name__),
+            description=self.documentation,
+            docs_url=None,
+            version=get_function_metadata_store(func).version,
+            redoc_url=None,
+            **kwargs,
+        )
+
+    @staticmethod
+    def deploy(
+        disable_docs: bool = False, disable_metrics: bool = False
+    ) -> Callable[[Callable[..., Any]], "GreatAI"]:
+        def decorator(func: Callable[..., Any]) -> GreatAI:
+            return GreatAI(func)._bootstrap_rest_api(
+                disable_docs=disable_docs, disable_metrics=disable_metrics
+            )
+
+        return decorator
+
+    def process_batch(
+        self,
+        batch: Iterable[Any],
+        concurrency: Optional[int] = None,
+    ) -> Sequence[Trace]:
+        if not get_context().persistence.is_threadsafe:
+            concurrency = 1
+            get_context().logger.warning("Concurrency is ignored")
+
+        return parallel_map(self.process_single, batch, concurrency=concurrency)
+
+    @property
+    def documentation(self) -> str:
+        return (
+            f"GreatAI wrapper for interacting with the '{self._func.__name__}' function.\n"
+            + (self._func.__doc__ or "")
+        )
+
+    def _bootstrap_rest_api(
+        self, disable_docs: bool, disable_metrics: bool
+    ) -> "GreatAI":
+        self.post("/evaluations", status_code=status.HTTP_200_OK, response_model=Trace)(
+            self.process_single
+        )
+
+        @self.get("/evaluations/:evaluation_id", status_code=status.HTTP_200_OK)
+        def get_evaluation(evaluation_id: str) -> Trace:
+            result = get_context().persistence.get_trace(evaluation_id)
+            if result is None:
+                raise HTTPException(status_code=status.HTTP_404_NOT_FOUND)
+            return result
+
+        @self.post(
+            "/evaluations/:evaluation_id/feedback", status_code=status.HTTP_202_ACCEPTED
+        )
+        def give_feedback(evaluation_id: str, input: EvaluationFeedbackRequest) -> None:
+            get_context().persistence.add_evaluation(evaluation_id, input.evaluation)
+
+        if not disable_docs:
+
+            @self.get("/docs", include_in_schema=False)
+            def custom_swagger_ui_html() -> HTMLResponse:
+                return get_swagger_ui_html(openapi_url="openapi.json", title=self.title)
+
+            @self.get("/docs/index.html", include_in_schema=False)
+            def redirect_to_entrypoint() -> RedirectResponse:
+                return RedirectResponse("/docs")
+
+        if not disable_metrics:
+            dash_app = create_dash_app(self._func.__name__, self.documentation)
+            self.mount(get_context().metrics_path, WSGIMiddleware(dash_app))
+
+            @self.get("/", include_in_schema=False)
+            def redirect_to_entrypoint() -> RedirectResponse:
+                return RedirectResponse("/metrics")
+
+            self.mount(
+                "/assets",
+                StaticFiles(directory=PATH / "../dashboard/assets"),
+                name="static",
+            )
+
+            @self.post("/query", status_code=status.HTTP_200_OK)
+            def query_metrics(query: Query) -> List[Dict[str, Any]]:
+                return get_context().persistence.query(
+                    conjunctive_filters=query.filter,
+                    sort_by=query.sort,
+                    skip=query.skip,
+                    take=query.take,
+                )
+
+        @self.get("/health", status_code=status.HTTP_200_OK)
+        def check_health() -> HealthCheckResponse:
+            return HealthCheckResponse(is_healthy=True)
+
+        return self

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

@@ -1,25 +0,0 @@
-from typing import Any, Callable, Iterable, Optional, Sequence
-
-from great_ai.utilities.parallel_map import parallel_map
-
-from ..context import get_context
-from ..tracing import TracingContext
-from ..views import Trace
-
-
-def process_batch(
-    function: Callable[..., Any],
-    batch: Iterable[Any],
-    concurrency: Optional[int] = None,
-) -> Sequence[Trace]:
-    def inner(input: Any) -> Trace:
-        with TracingContext() as t:
-            result = function(input)
-            output = t.log_output(result)
-        return output
-
-    if not get_context().persistence.is_threadsafe:
-        concurrency = 1
-        get_context().logger.warning("Concurrency is ignored")
-
-    return parallel_map(inner, batch, concurrency=concurrency)

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

@@ -1,11 +0,0 @@
-from typing import Any, Callable
-
-from ..tracing import TracingContext
-from ..views import Trace
-
-
-def process_single(function: Callable[..., Any], input_value: Any) -> Trace:
-    with TracingContext() as t:
-        result = function(input_value)
-        output = t.log_output(result)
-    return output

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

@@ -1,4 +1,5 @@
 from .get_args import get_args
+from .get_function_metadata_store import get_function_metadata_store
 from .snake_case_to_text import snake_case_to_text
 from .strip_lines import strip_lines
 from .text_to_hex_color import text_to_hex_color

great_ai/src/great_ai/great_ai/helper/get_args.py

@@ -5,6 +5,7 @@ from typing import Any, Callable, Dict, Mapping, Sequence
 def get_args(
     func: Callable[..., Any], args: Sequence[Any], kwargs: Mapping[str, Any]
 ) -> Dict[str, Any]:
+    """Return mapping from parameter names to actual argument values"""
     signature = inspect.signature(func)
     filter_keys = [
         param.name

great_ai/src/great_ai/great_ai/helper/get_function_metadata_store.py

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

great_ai/src/great_ai/great_ai/metrics/__init__.py

@@ -1,3 +0,0 @@
-from .create_dash_app import create_dash_app
-from .log_argument import log_argument
-from .log_metric import log_metric

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

@@ -1,6 +1,7 @@
 from functools import wraps
 from typing import Any, Callable, Dict, List, Literal, Union
 
+from ..helper import get_function_metadata_store
 from ..tracing import TracingContext
 from ..views import Model
 from .load_model import load_model
@@ -11,7 +12,7 @@ def use_model(
     *,
     version: Union[int, Literal["latest"]],
     return_path: bool = False,
-    model_kwarg_name: str = "model"
+    model_kwarg_name: str = "model",
 ) -> Callable[..., Any]:
     assert isinstance(version, int) or version == "latest"
 
@@ -22,6 +23,12 @@ def use_model(
     )
 
     def decorator(func: Callable[..., Any]) -> Callable[..., Any]:
+        store = get_function_metadata_store(func)
+        store.model_parameter_names.append(model_kwarg_name)
+        if store.version:
+            store.version += "|"
+        store.version += f"{key}:{actual_version}"
+
         @wraps(func)
         def wrapper(*args: List[Any], **kwargs: Dict[str, Any]) -> Any:
             context = TracingContext.get_current_context()

great_ai/src/great_ai/great_ai/parameters/__init__.py

@@ -0,0 +1,3 @@
+from .automatically_decorate_parameters import automatically_decorate_parameters
+from .log_metric import log_metric
+from .parameter import parameter