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/create_dash_app.py

@@ -162,7 +162,7 @@ def create_dash_app(function_name: str, function_docs: str) -> Flask:
                 dimensions=[
                     get_dimension_descriptor(df, c)
                     for c in df.columns
-                    if not c.startswith("arg:") and c not in {"id", "created", "output"}
+                    if c not in {"id", "created", "output"}
                 ],
                 line_color=accent_color,
             )
@@ -181,11 +181,12 @@ def get_dimension_descriptor(df: pd.DataFrame, column: str) -> Dict[str, Any]:
     try:
         dimension["values"] = [float(v) for v in values]
     except (TypeError, ValueError):
+        MAX_LENGTH = 40
         unique_values = unique(values)
-        value_mapping = {str(v): i for i, v in enumerate(unique_values)}
+        value_mapping = {str(v)[-MAX_LENGTH:]: i for i, v in enumerate(unique_values)}
 
-        dimension["values"] = [value_mapping[str(v)] for v in values]
+        dimension["values"] = [value_mapping[str(v)[-MAX_LENGTH:]] for v in values]
         dimension["tickvals"] = list(value_mapping.values())
-        dimension["ticktext"] = list(value_mapping.keys())
+        dimension["ticktext"] = [k[-MAX_LENGTH:] for k in value_mapping.keys()]
 
     return dimension

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

great_ai/src/great_ai/great_ai/parameters/automatically_decorate_parameters.py

@@ -0,0 +1,28 @@
+import inspect
+from typing import Any, Callable
+
+from great_ai.great_ai.helper.get_function_metadata_store import (
+    get_function_metadata_store,
+)
+
+from .parameter import parameter
+
+
+def automatically_decorate_parameters(func: Callable[..., Any]) -> Callable[..., Any]:
+    signature = inspect.signature(func)
+    parameter_names = [
+        param.name
+        for param in signature.parameters.values()
+        if param.kind == param.POSITIONAL_OR_KEYWORD
+    ]
+
+    metadata = get_function_metadata_store(func)
+
+    for name in parameter_names:
+        if (
+            name not in metadata.model_parameter_names
+            and name not in metadata.input_parameter_names
+        ):
+            func = parameter(name)(func)
+
+    return func

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

@@ -2,38 +2,43 @@ from functools import wraps
 from typing import Any, Callable, Dict
 
 from ..exceptions import ArgumentValidationError
-from ..helper import get_args
+from ..helper import get_args, get_function_metadata_store
 from ..tracing import TracingContext
 
 
-def log_argument(
-    argument_name: str,
+def parameter(
+    parameter_name: str,
     *,
     validator: Callable[[Any], bool] = lambda _: True,
 ) -> Callable[..., Any]:
     def decorator(func: Callable[..., Any]) -> Callable[..., Any]:
-        actual_name = f"arg:{func.__name__}:{argument_name}"
+        get_function_metadata_store(func).input_parameter_names.append(parameter_name)
+
+        actual_name = f"arg:{func.__name__}:{parameter_name}"
 
         @wraps(func)
         def wrapper(*args: Any, **kwargs: Dict[str, Any]) -> Any:
             arguments = get_args(func, args, kwargs)
-            argument = arguments[argument_name]
+            argument = arguments[parameter_name]
 
-            expected_type = func.__annotations__.get(argument_name)
+            expected_type = func.__annotations__.get(parameter_name)
 
             if expected_type is not None and not isinstance(argument, expected_type):
                 raise ArgumentValidationError(
-                    f"Argument {argument_name} in {func.__name__} has the wrong type, expected: {expected_type.__name__}, got: {type(argument).__name__}"
+                    f"Argument {parameter_name} in {func.__name__} has the wrong type, expected: {expected_type.__name__}, got: {type(argument).__name__}"
                 )
 
             if not validator(argument):
                 raise ArgumentValidationError(
-                    f"Argument {argument_name} in {func.__name__} did not pass validation"
+                    f"Argument {parameter_name} in {func.__name__} did not pass validation"
                 )
 
             context = TracingContext.get_current_context()
             if context:
                 context.log_value(name=actual_name, value=argument)
+                if isinstance(argument, str):
+                    context.log_value(name=f"{actual_name}:length", value=len(argument))
+
             return func(*args, **kwargs)
 
         return wrapper

great_ai/src/great_ai/great_ai/views/__init__.py

@@ -1,5 +1,6 @@
 from .evaluation_feedback_request import EvaluationFeedbackRequest
 from .filter import Filter
+from .function_metadata import FunctionMetadata
 from .health_check_response import HealthCheckResponse
 from .model import Model
 from .operators import operators

great_ai/src/great_ai/great_ai/views/function_metadata.py

@@ -0,0 +1,9 @@
+from typing import List
+
+from pydantic import BaseModel
+
+
+class FunctionMetadata(BaseModel):
+    input_parameter_names: List[str] = []
+    model_parameter_names: List[str] = []
+    version: str = ""