Fix whylogs compat functions for refactored metric definition

whylabs · Mar 26, 2024 · e90f691 · e90f691
1 parent 394c033
commit e90f691
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 5 deletions.
diff --git a/langkit/metrics/whylogs_compat.py b/langkit/metrics/whylogs_compat.py
@@ -2,7 +2,7 @@
 
 from dataclasses import dataclass
 from functools import reduce
-from typing import Any, Dict, List, Optional, Union
+from typing import Any, Dict, List, Optional, Type, Union
 
 import pandas as pd
 
@@ -63,8 +63,11 @@ def udf(text: Union[pd.DataFrame, Dict[str, List[Any]]]) -> Any:
         types = {"prompt": str, "response": str}
         column_names = ["prompt", "response"]
     else:
-        types = {metric.input_names: str}
-        column_names = [metric.input_names]
+        types: Dict[str, Type[str]] = {}
+        for name in metric.input_names:
+            types[name] = str
+
+        column_names = [*metric.input_names]
 
     schema = UdfSchemaArgs(
         types=types,
@@ -89,7 +92,7 @@ def udf(text: Union[pd.DataFrame, Dict[str, List[Any]]]) -> Any:
     return UdfSchemaArgs(
         resolvers=[],
         types={k: str for k in metric.names},
-        udf_specs=[UdfSpec(column_names=[metric.input_names], udf=udf, prefix="")],
+        udf_specs=[UdfSpec(column_names=metric.input_names, udf=udf, prefix="")],
     )
 
 

diff --git a/tests/langkit/metrics/test_text_statistics.py b/tests/langkit/metrics/test_text_statistics.py
@@ -414,7 +414,7 @@ def udf(text: pd.DataFrame) -> MultiMetricResult:
 
         return MultiMetric(
             names=[f"{column_name}.custom_textstat1", f"{column_name}.custom_textstat2"],
-            input_names=column_name,
+            input_names=[column_name],
             evaluate=udf,
         )