Add exception logging

great_ai/src/great_ai/__main__.py

@@ -1,4 +1,5 @@
 #!/usr/bin/env python3
+
 import re
 from importlib import import_module
 
@@ -33,9 +34,7 @@ def main() -> None:
         if function_name:
             logger.warning(f"Found `{function_name}` as the value of function_name")
         else:
-            raise MissingArgumentError(
-                "Argument function_name not provided and could not be guessed"
-            )
+            raise MissingArgumentError("Argument function_name could not be guessed")
 
     app_name = f"{file_name}:{function_name}"
     logger.info(f"Starting uvicorn server with app={app_name}")

great_ai/src/great_ai/great_ai/__init__.py

@@ -2,4 +2,5 @@ from .context import configure
 from .deploy import GreatAI
 from .exceptions import ArgumentValidationError, MissingArgumentError
 from .models import save_model, use_model
+from .output_models import ClassificationOutput, RegressionOutput
 from .parameters import log_metric, parameter

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

@@ -1,6 +1,18 @@
 import inspect
+from functools import partial
 from pathlib import Path
-from typing import Any, Callable, Dict, Iterable, List, Optional, Sequence, Type, cast
+from typing import (
+    Any,
+    Callable,
+    Dict,
+    Iterable,
+    List,
+    Optional,
+    Sequence,
+    Type,
+    Union,
+    cast,
+)
 
 from fastapi import FastAPI, HTTPException, status
 from fastapi.middleware.wsgi import WSGIMiddleware
@@ -10,6 +22,7 @@ from fastapi.staticfiles import StaticFiles
 from pydantic import BaseModel, create_model
 from starlette.responses import HTMLResponse
 
+from great_ai.great_ai.helper.use_http_exceptions import use_http_exceptions
 from great_ai.utilities.parallel_map import parallel_map
 
 from ..context import get_context
@@ -26,22 +39,12 @@ 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
+        schema = self._get_schema()
 
         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)
+                output = t.finalise(output=result)
             return output
 
         self.process_single = process_single
@@ -58,14 +61,24 @@ class GreatAI(FastAPI):
 
     @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
+        func: Optional[Callable[..., Any]] = None,
+        *,
+        disable_docs: bool = False,
+        disable_metrics: bool = False,
+    ) -> Union[Callable[[Callable[..., Any]], "GreatAI"], "GreatAI"]:
+        if func is None:
+            return cast(
+                Callable[..., Any],
+                partial(
+                    GreatAI.deploy,
+                    disable_docs=disable_docs,
+                    disable_metrics=disable_metrics,
+                ),
             )
 
-        return decorator
+        return GreatAI(func)._bootstrap_rest_api(
+            disable_docs=disable_docs, disable_metrics=disable_metrics
+        )
 
     def process_batch(
         self,
@@ -85,11 +98,25 @@ class GreatAI(FastAPI):
             + (self._func.__doc__ or "")
         )
 
+    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_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
+            use_http_exceptions(self.process_single)
         )
 
         @self.get("/evaluations/:evaluation_id", status_code=status.HTTP_200_OK)

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

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

great_ai/src/great_ai/great_ai/helper/use_http_exceptions.py

@@ -0,0 +1,18 @@
+from functools import wraps
+from typing import Any, Callable, Dict, List
+
+from fastapi import HTTPException, status
+
+
+def use_http_exceptions(func: Callable[..., Any]) -> Callable[..., Any]:
+    @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 wrapper

great_ai/src/great_ai/great_ai/output_models/__init__.py

@@ -0,0 +1,2 @@
+from .classification_output import ClassificationOutput
+from .regression_output import RegressionOutput

great_ai/src/great_ai/great_ai/output_models/classification_output.py

@@ -0,0 +1,12 @@
+from typing import Any, Optional, Union
+
+from pydantic import BaseModel
+
+
+class ClassificationOutput(BaseModel):
+    label: Union[str, int]
+    confidence: float
+    explanation: Optional[Any]
+
+    def __hash__(self) -> int:
+        return hash((type(self),) + tuple(self.__dict__.values()))

great_ai/src/great_ai/great_ai/output_models/regression_output.py

@@ -0,0 +1,11 @@
+from typing import Any, Optional, Union
+
+from pydantic import BaseModel
+
+
+class RegressionOutput(BaseModel):
+    value: Union[int, float]
+    explanation: Optional[Any]
+
+    def __hash__(self) -> int:
+        return hash((type(self),) + tuple(self.__dict__.values()))

great_ai/src/great_ai/great_ai/output_models/sequence_labeling_output.py

@@ -0,0 +1 @@
+# todo

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

@@ -10,6 +10,7 @@ def parameter(
     parameter_name: str,
     *,
     validator: Callable[[Any], bool] = lambda _: True,
+    disable_logging: bool = False,
 ) -> Callable[..., Any]:
     def decorator(func: Callable[..., Any]) -> Callable[..., Any]:
         get_function_metadata_store(func).input_parameter_names.append(parameter_name)
@@ -34,7 +35,7 @@ def parameter(
                 )
 
             context = TracingContext.get_current_context()
-            if context:
+            if context and not disable_logging:
                 context.log_value(name=actual_name, value=argument)
                 if isinstance(argument, str):
                     context.log_value(name=f"{actual_name}:length", value=len(argument))

great_ai/src/great_ai/great_ai/tracing/tracing_context.py

@@ -14,7 +14,6 @@ class TracingContext:
     def __init__(self) -> None:
         self._models: List[Model] = []
         self._values: Dict[str, Any] = {}
-        self._output: Any = None
         self._trace: Optional[Trace] = None
         self._start_time = datetime.utcnow()
 
@@ -24,18 +23,21 @@ class TracingContext:
     def log_model(self, model: Model) -> None:
         self._models.append(model)
 
-    def log_output(self, output: Any, evaluation_id: Optional[str] = None) -> Trace:
-        self._output = output
+    def finalise(self, output: Any = None, exception: BaseException = None) -> Trace:
+        assert self._trace is None, "has been already finalised"
 
         delta_time = (datetime.utcnow() - self._start_time).microseconds / 1000
         self._trace = Trace(
-            evaluation_id=evaluation_id,
             created=self._start_time.isoformat(),
             execution_time_ms=delta_time,
             logged_values=self._values,
             models=self._models,
-            output=self._output,
+            output=output,
+            exception=None
+            if exception is None
+            else f"{type(exception).__name__}: {exception}",
         )
+
         return self._trace
 
     @classmethod
@@ -57,10 +59,16 @@ class TracingContext:
         assert self._contexts[threading.get_ident()][-1] == self
         self._contexts[threading.get_ident()].remove(self)
 
-        if type is None:
-            assert self._trace is not None
-            get_context().persistence.save_trace(self._trace)
-        else:
-            get_context().logger.exception(f"Could not finish operation: {exception}")
+        if exception is not None and type is not None:
+            self.finalise(exception=exception)
+            if get_context().is_production:
+                get_context().logger.error(
+                    f"Could not finish operation because of {type.__name__}: {exception}"
+                )
+            else:
+                get_context().logger.exception("Could not finish operation")
+
+        assert self._trace is not None
+        get_context().persistence.save_trace(self._trace)
 
         return False

great_ai/src/great_ai/great_ai/views/trace.py

@@ -13,6 +13,7 @@ class Trace(BaseModel):
     execution_time_ms: float
     logged_values: Dict[str, Any]
     models: List[Model]
+    exception: Optional[str]
     output: Any
     evaluation: Any = None
 
@@ -31,6 +32,7 @@ class Trace(BaseModel):
             "execution_time_ms": self.execution_time_ms,
             "models": ", ".join(f"{m.key}:{m.version}" for m in self.models),
             "output": dumps(self.output),
+            "exception": self.exception or "null",
             "evaluation": self.evaluation,
             **self.logged_values,
         }