Improve types

great_ai/tracing/delete_ground_truth.py

@@ -1,5 +1,5 @@
 from datetime import datetime
-from typing import List, Optional, Union, cast
+from typing import List, Optional, Union
 
 from ..context import get_context
 
@@ -19,4 +19,4 @@ def delete_ground_truth(
         conjunctive_tags=tags, until=until, since=since, has_feedback=True
     )
 
-    db.delete_batch([cast(str, i.trace_id) for i in items])
+    db.delete_batch([i.trace_id for i in items])

great_ai/tracing/tracing_context.py

@@ -2,6 +2,7 @@ from contextvars import ContextVar
 from datetime import datetime
 from types import TracebackType
 from typing import Any, Dict, Generic, List, Literal, Optional, Type, TypeVar
+from uuid import uuid4
 
 from ..constants import DEVELOPMENT_TAG_NAME, ONLINE_TAG_NAME, PRODUCTION_TAG_NAME
 from ..context import get_context
@@ -25,11 +26,14 @@ class TracingContext(Generic[T]):
     def log_model(self, model: Model) -> None:
         self._models.append(model)
 
-    def finalise(self, output: T = None, exception: BaseException = None) -> Trace[T]:
+    def finalise(
+        self, output: Optional[T] = None, exception: Optional[BaseException] = None
+    ) -> Trace[T]:
         assert self._trace is None, "has been already finalised"
 
         delta_time = (datetime.utcnow() - self._start_time).microseconds / 1000
-        self._trace = Trace(
+        self._trace = Trace[T](
+            trace_id=str(uuid4()),
             created=self._start_time.isoformat(),
             original_execution_time_ms=delta_time,
             logged_values=self._values,