Add ground_truth API and refactor

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

@@ -1,35 +1,67 @@
-from typing import Any, Dict, Optional
+from typing import Any, Dict, Generic, List, Optional, TypeVar
+from uuid import uuid4
 
 import yaml
 from pydantic import validator
 
-from .trace_view import TraceView
+from ..helper import HashableBaseModel
+from .model import Model
+
+T = TypeVar("T")
 
 
-class Trace(TraceView):
-    models_flat: Optional[str]
-    output_flat: Optional[str]
-    feedback_flat: Optional[str]
+class Trace(HashableBaseModel, Generic[T]):
+    trace_id: Optional[str]
+    created: str
+    original_execution_time_ms: float
+    logged_values: Dict[str, Any]
+    models: List[Model]
+    exception: Optional[str]
+    output: T
+    feedback: Any = None
+    tags: List[str]
 
-    @validator("models_flat", always=True)
-    def flatten_models(cls, v: Optional[str], values: Dict[str, Any]) -> str:
-        return ", ".join(f"{m.key}:{m.version}" for m in values["models"])
+    @validator("trace_id", always=True)
+    def generate_id(cls, v: Optional[str], values: Dict[str, Any]) -> Optional[str]:
+        if not v:
+            return str(uuid4())
+        return v
 
-    @validator("output_flat", always=True)
-    def flatten_output(cls, v: Optional[str], values: Dict[str, Any]) -> str:
-        return yaml.dump(values["output"], default_flow_style=False, indent=2)
+    @property
+    def input(self) -> Any:
+        return (
+            self.logged_values["input"]
+            if list(self.logged_values.keys()) == ["input"]
+            else self.logged_values
+        )
+
+    @property
+    def models_flat(self) -> str:
+        return ", ".join(f"{m.key}:{m.version}" for m in self.models)
+
+    @property
+    def output_flat(self) -> str:
+        return yaml.dump(self.output, stream=None)
+
+    @property
+    def feedback_flat(self) -> str:
+        return (
+            "null" if self.feedback is None else yaml.dump(self.feedback, stream=None)
+        )
+
+    @property
+    def tags_flat(self) -> str:
+        return ",\n".join(self.tags)
 
     def to_flat_dict(self) -> Dict[str, Any]:
         return {
             "trace_id": self.trace_id,
             "created": self.created,
             "original_execution_time_ms": self.original_execution_time_ms,
-            "models_flat": self.models_flat,
-            "output_flat": self.output_flat,
-            "exception": self.exception or "null",
-            "feedback_flat": self.feedback_flat or "null",
             **self.logged_values,
+            "models_flat": self.models_flat,
+            "exception": "null" if self.exception is None else self.exception,
+            "output_flat": self.output_flat,
+            "feedback_flat": self.feedback_flat,
+            "tags_flat": self.tags_flat,
         }
-
-    def __hash__(self) -> int:
-        return hash((type(self),) + tuple(self.__dict__.values()))