Add exception logging

examples/complex/helpers.py

@@ -0,0 +1,16 @@
+import re
+
+from great_ai.utilities.clean import clean
+from great_ai.utilities.lemmatize_text import lemmatize_text
+
+
+def preprocess(text: str) -> str:
+    text = clean(text, convert_to_ascii=True)
+    text = re.sub(r"[^a-zA-Z0-9]", " ", text)
+    return text
+
+
+def lemmatize(text: str) -> str:
+    lemmatized = lemmatize_text(text)
+    clean_lemmas = [re.sub(r"\d[\d.,]*", "NUM", lemma) for lemma in lemmatized]
+    return " ".join(clean_lemmas)

examples/complex/main_service.py

@@ -1,7 +0,0 @@
-from great_ai import configure, create_service
-
-configure(development_mode_override=True)
-
-from predict_domain import predict_domain
-
-app = create_service(predict_domain)

examples/complex/predict_domain.py

@@ -1,39 +1,26 @@
-import re
 from typing import Dict, Iterable, List
-
+from great_ai import GreatAI, use_model, ClassificationOutput
-from great_ai import log_argument, log_metric, use_model
-from great_ai.utilities.clean import clean
-from pydantic import BaseModel
 from sklearn.pipeline import Pipeline
 
-from preprocess import preprocess
+from helpers import lemmatize, preprocess
-
-
-class DomainPrediction(BaseModel):
-    domain: str
-    probability: float
-    explanation: List[str]
 
 
+@GreatAI.deploy
 @use_model("small-domain-prediction-v2", version="latest")
-@log_argument("text", validator=lambda t: len(t) > 0)
+def predict_domain(text: str, model: Pipeline, target_confidence: float = 20) -> List[ClassificationOutput]:
-def predict_domain(
-    text: str, model: Pipeline, cut_off_probability: float = 0.2
-) -> List[DomainPrediction]:
-    assert 0 <= cut_off_probability <= 1
-
     """
     Predict the scientific domain of the input text.
-    Return labels until their sum likelihood is larger than cut_off_probability.
+    Return labels until their sum likelihood is larger than target_confidence.
     """
-    log_metric("text_length", len(text))
+    assert 0 <= target_confidence <= 100, "invalid argument"
 
-    cleaned = clean(text, convert_to_ascii=True)
+    text = preprocess(text)
-    text = re.sub(r"[^a-zA-Z0-9]", " ", cleaned)
 
-    feature_names = model.named_steps["vectorizer"].get_feature_names_out()
+    token_mapping = {lemmatize(original): original for original in text.split(" ")}
-
+    feature_names = [
-    token_mapping = {preprocess(original): original for original in text.split(" ")}
+        token_mapping.get(name)
+        for name in model.named_steps["vectorizer"].get_feature_names_out()
+    ]
 
     features = model.named_steps["vectorizer"].transform(
         [" ".join(token_mapping.keys())]
@@ -41,42 +28,38 @@ def predict_domain(
     prediction = model.named_steps["classifier"].predict_proba(features)[0]
     best_classes = sorted(enumerate(prediction), key=lambda v: v[1], reverse=True)
 
-    results: List[DomainPrediction] = []
+    results: List[ClassificationOutput] = []
     for class_index, probability in best_classes:
         weights = model.named_steps["classifier"].feature_log_prob_[class_index]
         domain = model.named_steps["classifier"].classes_[class_index]
 
         results.append(
-            DomainPrediction(
+            ClassificationOutput(
-                domain=domain,
+                label=domain,
-                probability=round(probability * 100),
+                confidence=round(probability * 100),
                 explanation=_get_explanation(
-                    feature_names=feature_names,
-                    features=features.A[0],
                     weights=weights,
-                    token_mapping=token_mapping,
+                    counts=features.A[0],
+                    words=feature_names,
                 ),
             )
         )
 
-        if sum(r.probability for r in results) >= cut_off_probability * 100:
+        if sum(r.confidence for r in results) >= target_confidence:
             break
 
     return results
 
 
 def _get_explanation(
-    feature_names: Iterable[str],
-    features: Iterable[float],
     weights: Iterable[float],
-    token_mapping: Dict[str, str],
+    counts: Iterable[float],
+    words: Iterable[str],
 ) -> List[str]:
-    influential = [
+    most_influential = sorted((
-        (weight, name)
+        (weight, word)
-        for weight, value, name in zip(weights, features, feature_names)
+        for weight, count, word in zip(weights, counts, words)
-        if value
+        if count > 0
-    ]
+    ), reverse=True)[:5]
 
-    most_influential = sorted(influential, reverse=True)[:5]
+    return [word for _, word in most_influential]
-
-    return [token_mapping[name] for _, name in most_influential]

examples/complex/preprocess.py

@@ -1,8 +0,0 @@
-import re
-
-from great_ai.utilities.lemmatize_text import lemmatize_text
-
-
-def preprocess(text: str) -> str:
-    lemmas = [re.sub(r"\d[\d.,]*", "NUM", lemma) for lemma in lemmatize_text(text)]
-    return " ".join(lemmas)

examples/complex/train.ipynb

@@ -0,0 +1,269 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Train domain classifier on the [semantic scholar dataset](https://api.semanticscholar.org/corpus)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import json\n",
+    "import matplotlib.pyplot as plt\n",
+    "from pathlib import Path\n",
+    "import pandas as pd\n",
+    "from sklearn import metrics, set_config\n",
+    "from sklearn.model_selection import train_test_split\n",
+    "from sklearn.naive_bayes import MultinomialNB\n",
+    "from sklearn.pipeline import Pipeline\n",
+    "from sklearn.feature_extraction.text import TfidfVectorizer\n",
+    "from sklearn.model_selection import GridSearchCV\n",
+    "\n",
+    "from great_ai.utilities.parallel_map import parallel_map\n",
+    "from great_ai.utilities.language import is_english, predict_language\n",
+    "from great_ai import save_model, configure, LargeFile\n",
+    "\n",
+    "from helpers import preprocess, lemmatize"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Configuration"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "PREFIX = \"domain-\"\n",
+    "DATASET_KEY = \"data\"\n",
+    "MAX_FILE_COUNT = 5\n",
+    "MODEL_KEY = \"small-domain-prediction-v2\""
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "configure()\n",
+    "corpus_path = LargeFile(DATASET_KEY).get()\n",
+    "\n",
+    "set_config(display=\"diagram\")\n",
+    "plt.rcParams[\"figure.figsize\"] = (30, 15)\n",
+    "plt.rcParams[\"figure.facecolor\"] = \"white\"\n",
+    "plt.rcParams[\"font.size\"] = 12\n",
+    "plt.rcParams[\"axes.xmargin\"] = 0"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Preprocessing"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def clean_file(p: Path) -> None:\n",
+    "    try:\n",
+    "        processed_path = p.with_name(f\"{PREFIX}{p.stem}{p.suffix}\")\n",
+    "\n",
+    "        if processed_path.exists():\n",
+    "            return\n",
+    "\n",
+    "        with open(p) as f:\n",
+    "            content = json.load(f)\n",
+    "\n",
+    "        result = {\n",
+    "            lemmatize(preprocess(f'{c[\"title\"]} {c[\"abstract\"]}')): c[\"domain\"]\n",
+    "            for c in content\n",
+    "            if (\n",
+    "                c[\"domain\"]\n",
+    "                and c[\"abstract\"]\n",
+    "                and is_english(predict_language(c[\"abstract\"]))\n",
+    "            )\n",
+    "        }\n",
+    "\n",
+    "        with open(processed_path, \"w\") as f:\n",
+    "            json.dump(result, f)\n",
+    "    except Exception as e:\n",
+    "        print(f\"Error ({e}) processing {p}\")\n",
+    "\n",
+    "\n",
+    "parallel_map(\n",
+    "    clean_file,\n",
+    "    list(corpus_path.glob(\"s2-corpus-*.json\"))[:MAX_FILE_COUNT],\n",
+    "    chunk_size=1,\n",
+    ")\n",
+    "None"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "corpora = list(corpus_path.glob(f\"{PREFIX}*.json\"))[:MAX_FILE_COUNT]\n",
+    "print(f\"Found {len(corpora)} files\")\n",
+    "\n",
+    "data = []\n",
+    "for p in corpora:\n",
+    "    with open(p) as f:\n",
+    "        data.extend(json.load(f).items())\n",
+    "\n",
+    "print(f\"Found {len(data)} documents\")\n",
+    "\n",
+    "X_train, X_test, y_train, y_test = train_test_split(\n",
+    "    [d[0] for d in data], [d[1] for d in data], test_size=0.1, random_state=1\n",
+    ")\n",
+    "\n",
+    "X_train = [x for x, y in zip(X_train, y_train) for domain in y]\n",
+    "y_train = [domain for x, y in zip(X_train, y_train) for domain in y]"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Optimise and train Multinomial Naive Bayes classifier"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def create_pipeline() -> Pipeline:\n",
+    "    return Pipeline(\n",
+    "        steps=[\n",
+    "            (\"vectorizer\", TfidfVectorizer()),\n",
+    "            (\"classifier\", MultinomialNB()),\n",
+    "        ]\n",
+    "    )"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "optimisation_pipeline = GridSearchCV(\n",
+    "    create_pipeline(),\n",
+    "    {\n",
+    "        \"vectorizer__max_df\": [0.05, 0.1],\n",
+    "        \"vectorizer__min_df\": [5, 20],\n",
+    "        \"vectorizer__sublinear_tf\": [True, False],\n",
+    "        \"classifier__alpha\": [0.5, 1],\n",
+    "        \"classifier__fit_prior\": [True, False],\n",
+    "    },\n",
+    "    scoring=\"f1_macro\",\n",
+    "    cv=3,\n",
+    "    n_jobs=8,\n",
+    "    verbose=1,\n",
+    ")\n",
+    "optimisation_pipeline.fit(X_train, y_train)\n",
+    "\n",
+    "results = pd.DataFrame(optimisation_pipeline.cv_results_)\n",
+    "results.sort_values(\"rank_test_score\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "classifier = create_pipeline()\n",
+    "classifier.set_params(**optimisation_pipeline.best_params_)\n",
+    "classifier.fit(X_train, y_train)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Check accuracy on the test split"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "predicted = classifier.predict(X_test)\n",
+    "\n",
+    "y_test_aligned = [p if p in y else y[0] for p, y in zip(predicted, y_test)]\n",
+    "\n",
+    "print(metrics.classification_report(y_test_aligned, predicted))\n",
+    "metrics.ConfusionMatrixDisplay.from_predictions(\n",
+    "    y_true=y_test_aligned,\n",
+    "    y_pred=predicted,\n",
+    "    xticks_rotation=\"vertical\",\n",
+    "    normalize=\"pred\",\n",
+    "    values_format=\".2f\",\n",
+    ")\n",
+    "None"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Export the model using GreatAI"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "save_model(classifier, key=MODEL_KEY, keep_last_n=5)"
+   ]
+  }
+ ],
+ "metadata": {
+  "interpreter": {
+   "hash": "acc70e949538f42041ccc57bc4df2261507e3fd7d6b9ce5dcc28e3bcf9d48274"
+  },
+  "kernelspec": {
+   "display_name": "Python 3.8.5 ('.env': venv)",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.8.5"
+  },
+  "orig_nbformat": 4
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

examples/simple/README.md

@@ -1,12 +0,0 @@
-# Train Domain classifier on the [semantic scholar dataset](https://api.semanticscholar.org/corpus)
-
-## Upload the dataset (or a part of it) to shared infrastructure
-
-```sh
-mkdir ss-data && cd ss-data
-wget https://s3-us-west-2.amazonaws.com/ai2-s2-research-public/open-corpus/2022-02-01/manifest.txt
-wget -B https://s3-us-west-2.amazonaws.com/ai2-s2-research-public/open-corpus/2022-02-01/ -i manifest.txt
-cd -
-python3 -m great_ai.open_s3 --secrets s3.ini --push ss-data
-rm -rf ss-data
-```

examples/simple/predict_domain.py

@@ -1,77 +0,0 @@
-from typing import Dict, Iterable, List
-
-from great_ai import GreatAI, use_model
-from pydantic import BaseModel
-from sklearn.pipeline import Pipeline
-
-from helpers import lemmatize, preprocess
-
-
-class DomainPrediction(BaseModel):
-    domain: str
-    probability: float
-    explanation: List[str]
-
-
-@GreatAI.deploy()
-@use_model("small-domain-prediction-v2", version="latest")
-def predict_domain(
-    text: str, model: Pipeline, cut_off_probability: float = 0.2
-) -> List[DomainPrediction]:
-    """
-    Predict the scientific domain of the input text.
-    Return labels until their sum likelihood is larger than cut_off_probability.
-    """
-    assert 0 <= cut_off_probability <= 1
-
-    text = preprocess(text)
-
-    feature_names = model.named_steps["vectorizer"].get_feature_names_out()
-
-    token_mapping = {lemmatize(original): original for original in text.split(" ")}
-
-    features = model.named_steps["vectorizer"].transform(
-        [" ".join(token_mapping.keys())]
-    )
-    prediction = model.named_steps["classifier"].predict_proba(features)[0]
-    best_classes = sorted(enumerate(prediction), key=lambda v: v[1], reverse=True)
-
-    results: List[DomainPrediction] = []
-    for class_index, probability in best_classes:
-        weights = model.named_steps["classifier"].feature_log_prob_[class_index]
-        domain = model.named_steps["classifier"].classes_[class_index]
-
-        results.append(
-            DomainPrediction(
-                domain=domain,
-                probability=round(probability * 100),
-                explanation=_get_explanation(
-                    feature_names=feature_names,
-                    features=features.A[0],
-                    weights=weights,
-                    token_mapping=token_mapping,
-                ),
-            )
-        )
-
-        if sum(r.probability for r in results) >= cut_off_probability * 100:
-            break
-
-    return results
-
-
-def _get_explanation(
-    feature_names: Iterable[str],
-    features: Iterable[float],
-    weights: Iterable[float],
-    token_mapping: Dict[str, str],
-) -> List[str]:
-    influential = [
-        (weight, name)
-        for weight, value, name in zip(weights, features, feature_names)
-        if value
-    ]
-
-    most_influential = sorted(influential, reverse=True)[:5]
-
-    return [token_mapping[name] for _, name in most_influential]

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(
+            raise MissingArgumentError("Argument function_name could not be guessed")
-                "Argument function_name not provided and 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)
+        schema = self._get_schema()
-        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)
+                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
+        func: Optional[Callable[..., Any]] = None,
-    ) -> Callable[[Callable[..., Any]], "GreatAI"]:
+        *,
-        def decorator(func: Callable[..., Any]) -> GreatAI:
+        disable_docs: bool = False,
-            return GreatAI(func)._bootstrap_rest_api(
+        disable_metrics: bool = False,
-                disable_docs=disable_docs, disable_metrics=disable_metrics
+    ) -> 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:
+    def finalise(self, output: Any = None, exception: BaseException = None) -> Trace:
-        self._output = output
+        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:
+        if exception is not None and type is not None:
-            assert self._trace is not None
+            self.finalise(exception=exception)
-            get_context().persistence.save_trace(self._trace)
+            if get_context().is_production:
-        else:
+                get_context().logger.error(
-            get_context().logger.exception(f"Could not finish operation: {exception}")
+                    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,
         }