virtual record should use supplied perf, cost (#1266)

* virtual record should use supplied perf, cost * fmt
truera · Jun 27, 2024 · 26848c6 · 26848c6
1 parent 65828bd
commit 26848c6
Showing 1 changed file with 4 additions and 6 deletions.
diff --git a/trulens_eval/trulens_eval/tru_virtual.py b/trulens_eval/trulens_eval/tru_virtual.py
@@ -407,12 +407,10 @@ def __init__(
         if (end_time - start_time).total_seconds() == 0.0:
             end_time += datetime.timedelta(microseconds=1)
 
-        if "cost" not in kwargs:
-            kwargs['cost'] = mod_base_schema.Cost()
-        if "perf" not in kwargs:
-            kwargs['perf'] = mod_base_schema.Perf(
-                start_time=start_time, end_time=end_time
-            )
+        kwargs['cost'] = cost or mod_base_schema.Cost()
+        kwargs['perf'] = perf or mod_base_schema.Perf(
+            start_time=start_time, end_time=end_time
+        )
 
         if "main_input" not in kwargs:
             kwargs['main_input'] = "No main_input provided."