diff --git a/great_ai/persistence/parallel_tinydb_driver.py b/great_ai/persistence/parallel_tinydb_driver.py index e648477..c278549 100644 --- a/great_ai/persistence/parallel_tinydb_driver.py +++ b/great_ai/persistence/parallel_tinydb_driver.py @@ -67,7 +67,7 @@ class ParallelTinyDbDriver(TracingDatabaseDriver): if not documents: return [], 0 - df = pd.DataFrame(documents) + df = pd.DataFrame([Trace.parse_obj(d).to_flat_dict() for d in documents]) for f in conjunctive_filters: operator = f.operator.lower() @@ -76,7 +76,12 @@ class ParallelTinyDbDriver(TracingDatabaseDriver): getattr(df[f.property], operator_mapping[f.operator])(f.value) ] elif operator == "contains": - df = df.loc[df[f.property].str.contains(f.value, case=False)] + df = df.loc[ + df[f.property].str.contains( + str(int(f.value)) if isinstance(f.value, float) else f.value, + case=False, + ) + ] if sort_by: df.sort_values( diff --git a/great_ai/views/trace.py b/great_ai/views/trace.py index fa71d81..76f03f5 100644 --- a/great_ai/views/trace.py +++ b/great_ai/views/trace.py @@ -88,7 +88,9 @@ class Trace(Generic[T], HashableBaseModel): } ), **{ - k: 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,