Add exception logging

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