Integrate MongoDB

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

@@ -1,8 +1,8 @@
+from pprint import pformat
 from typing import Any, Dict, Generic, List, Optional, TypeVar
 from uuid import uuid4
 
-import yaml
-from pydantic import validator
+from pydantic import Extra, validator
 
 from ..helper import HashableBaseModel
 from .model import Model
@@ -10,7 +10,7 @@ from .model import Model
 T = TypeVar("T")
 
 
-class Trace(HashableBaseModel, Generic[T]):
+class Trace(Generic[T], HashableBaseModel):
     trace_id: Optional[str]
     created: str
     original_execution_time_ms: float
@@ -21,6 +21,9 @@ class Trace(HashableBaseModel, Generic[T]):
     feedback: Any = None
     tags: List[str]
 
+    class Config:
+        extra = Extra.ignore
+
     @validator("trace_id", always=True)
     def generate_id(cls, v: Optional[str], values: Dict[str, Any]) -> Optional[str]:
         if not v:
@@ -41,26 +44,42 @@ class Trace(HashableBaseModel, Generic[T]):
 
     @property
     def output_flat(self) -> str:
-        return yaml.dump(self.output, stream=None)
+        return pformat(self.output, indent=2, compact=True)
+
+    @property
+    def exception_flat(self) -> str:
+        return (
+            "null"
+            if self.exception is None
+            else pformat(self.exception, indent=2, compact=True)
+        )
 
     @property
     def feedback_flat(self) -> str:
         return (
-            "null" if self.feedback is None else yaml.dump(self.feedback, stream=None)
+            "null"
+            if self.feedback is None
+            else pformat(self.feedback, indent=2, compact=True)
         )
 
     @property
     def tags_flat(self) -> str:
         return ",\n".join(self.tags)
 
-    def to_flat_dict(self) -> Dict[str, Any]:
+    def to_flat_dict(self, include_original: bool = True) -> Dict[str, Any]:
         return {
-            "trace_id": self.trace_id,
-            "created": self.created,
-            "original_execution_time_ms": self.original_execution_time_ms,
+            **(
+                self.dict()
+                if include_original
+                else {
+                    "trace_id": self.trace_id,
+                    "created": self.created,
+                    "original_execution_time_ms": self.original_execution_time_ms,
+                }
+            ),
             **self.logged_values,
             "models_flat": self.models_flat,
-            "exception": "null" if self.exception is None else self.exception,
+            "exception_flat": self.exception_flat,
             "output_flat": self.output_flat,
             "feedback_flat": self.feedback_flat,
             "tags_flat": self.tags_flat,