Modernise

great_ai/views/outputs/classification_output.py

@@ -6,4 +6,4 @@ from ..hashable_base_model import HashableBaseModel
 class ClassificationOutput(HashableBaseModel):
     label: Union[str, int]
     confidence: float
-    explanation: Optional[Any]
+    explanation: Optional[Any] = None

great_ai/views/outputs/regression_output.py

@@ -5,4 +5,4 @@ from ..hashable_base_model import HashableBaseModel
 
 class RegressionOutput(HashableBaseModel):
     value: Union[int, float]
-    explanation: Optional[Any]
+    explanation: Optional[Any] = None

great_ai/views/outputs/sequence_labeling_output.py

@@ -9,9 +9,9 @@ class LabeledToken(HashableBaseModel):
     token: str
     tag: Literal["B", "I", "O", "E", "S"]
     confidence: float
-    explanation: Optional[Any]
+    explanation: Optional[Any] = None
 
 
 class SequenceLabelingOutput(HashableBaseModel):
     labeled_tokens: List[LabeledToken]
-    explanation: Optional[Any]
+    explanation: Optional[Any] = None

great_ai/views/query.py

@@ -1,7 +1,7 @@
 from datetime import datetime
 from typing import List, Optional, Sequence
 
-from pydantic import BaseModel
+from pydantic import BaseModel, ConfigDict
 
 from .filter import Filter
 from .sort_by import SortBy
@@ -15,8 +15,8 @@ class Query(BaseModel):
     until: Optional[datetime] = None
     has_feedback: Optional[bool] = None
 
-    class Config:
-        schema_extra = {
+    model_config = ConfigDict(
+        json_schema_extra={
             "example": {
                 "filter": [
                     {
@@ -33,3 +33,4 @@ class Query(BaseModel):
                 "has_feedback": False,
             }
         }
+    )

great_ai/views/trace.py

@@ -1,7 +1,7 @@
 from pprint import pformat
 from typing import Any, Dict, Generic, List, Optional, TypeVar
 
-from pydantic import Extra
+from pydantic import ConfigDict
 
 from .hashable_base_model import HashableBaseModel
 from .model import Model
@@ -32,13 +32,12 @@ class Trace(Generic[T], HashableBaseModel):
     original_execution_time_ms: float
     logged_values: Dict[str, Any]
     models: List[Model]
-    exception: Optional[str]
-    output: Optional[T]
+    exception: Optional[str] = None
+    output: Optional[T] = None
     feedback: Any = None
     tags: List[str]
 
-    class Config:
-        extra = Extra.ignore
+    model_config = ConfigDict(extra="ignore")
 
     @property
     def input(self) -> Any:
@@ -79,7 +78,7 @@ class Trace(Generic[T], HashableBaseModel):
     def to_flat_dict(self, include_original: bool = True) -> Dict[str, Any]:
         return {
             **(
-                self.dict()
+                self.model_dump()
                 if include_original
                 else {
                     "trace_id": self.trace_id,
@@ -88,9 +87,11 @@ class Trace(Generic[T], HashableBaseModel):
                 }
             ),
             **{
-                k: v
-                if (isinstance(v, float) or isinstance(v, int))
-                else pformat(v, indent=2, compact=True)
+                k: (
+                    v
+                    if (isinstance(v, float) or isinstance(v, int))
+                    else pformat(v, indent=2, compact=True)
+                )
                 for k, v in self.logged_values.items()
             },
             "models_flat": self.models_flat,
@@ -101,6 +102,7 @@ class Trace(Generic[T], HashableBaseModel):
         }
 
     def __repr__(self) -> str:
-        return f"""Trace[{type(self.output).__name__}]({
-            pformat(self.dict(), indent=2, compact=True).replace('{ ', '{', 1)
-        })"""
+        formatted = pformat(self.model_dump(), indent=2, compact=True).replace(
+            "{ ", "{", 1
+        )
+        return f"Trace[{type(self.output).__name__}]({formatted})"