Bug fix and minor improvements

Author: pantonante

continuous_eval/eval/dataset.py

@@ -1,7 +1,9 @@
 import json
+import random
 import typing
 from dataclasses import dataclass
 from pathlib import Path
+from string import ascii_lowercase, digits
 
 import yaml
 
@@ -13,6 +15,15 @@
 _SAFE_DICT["ToolCall"] = ToolCall
 
 
+def _generate_uid():
+    return "".join(random.choices(ascii_lowercase + digits, k=8))
+
+
+@dataclass(frozen=True)
+class LambdaField:
+    func: typing.Callable
+
+
 @dataclass(frozen=True)
 class DatasetField:
     name: str
@@ -28,15 +39,15 @@ def to_dict(self):
         }
 
 
-@dataclass(frozen=True)
+@dataclass
 class DatasetManifest:
     name: str
     description: str
     format: str
     license: str
     fields: typing.Dict[str, DatasetField]
 
-    def to_yaml(self):
+    def to_dict(self):
         return {
             "name": self.name,
             "description": self.description,
@@ -45,6 +56,24 @@ def to_yaml(self):
             "fields": {field_name: field.to_dict() for field_name, field in self.fields.items()},
         }
 
+    @classmethod
+    def from_json(cls, data: typing.Dict):
+        return cls(
+            name=data.get("name", ""),
+            description=data.get("description", ""),
+            format=data.get("format", ""),
+            license=data.get("license", ""),
+            fields={
+                field_name: DatasetField(
+                    name=field_name,
+                    type=eval(field_info["type"], _SAFE_DICT),
+                    description=field_info.get("description", ""),
+                    is_ground_truth=field_info.get("ground_truth", False),
+                )
+                for field_name, field_info in data["fields"].items()
+            },
+        )
+
 
 class Dataset:
     def __init__(
@@ -68,14 +97,22 @@ def __init__(
         # load jsonl dataset
         with open(dataset_path, "r") as json_file:
             self._data = [json.loads(x) for x in json_file.readlines()]
+        for sample in self._data:
+            sample["uid"] = UID(sample["uid"]) if "uid" in sample else _generate_uid()
         self._manifest = self._load_or_infer_manifest(manifest_path)
         self._create_dynamic_properties()
 
     @classmethod
-    def from_data(cls, data: typing.List[typing.Dict[str, typing.Any]]):
+    def from_data(
+        cls,
+        data: typing.List[typing.Dict[str, typing.Any]],
+        manifest: typing.Optional[typing.Dict] = None,
+    ):
         dataset = cls.__new__(cls)
         dataset._data = data
-        dataset._manifest = dataset._infer_manifest()
+        for sample in dataset._data:
+            sample["uid"] = UID(sample["uid"]) if "uid" in sample else _generate_uid()
+        dataset._manifest = DatasetManifest.from_json(manifest) if manifest is not None else dataset._infer_manifest()
         dataset._create_dynamic_properties()
         return dataset
 
@@ -89,7 +126,7 @@ def save(self, file_path: typing.Union[str, Path], save_manifest: bool = False):
         if save_manifest:
             manifest_path = file_path.parent / "manifest.yaml"
             with open(manifest_path, "w") as manifest_file:
-                manifest_file.write(yaml.dump(self._manifest.to_yaml()))
+                manifest_file.write(yaml.dump(self._manifest.to_dict()))
 
     def _load_or_infer_manifest(self, manifest_path: typing.Optional[Path]) -> DatasetManifest:
         if manifest_path is None or not manifest_path.exists():
@@ -147,6 +184,10 @@ def _create_dynamic_properties(self):
     def filed_types(self, name: str) -> type:
         return getattr(self, name).type
 
+    @property
+    def manifest(self):
+        return self._manifest
+
     @property
     def data(self):
         return self._data
@@ -155,10 +196,18 @@ def data(self):
     def name(self):
         return self._manifest.name
 
+    @name.setter
+    def name(self, value):
+        self._manifest.name = value
+
     @property
     def description(self):
         return self._manifest.description
 
+    @description.setter
+    def description(self, value):
+        self._manifest.description = value
+
     @property
     def format(self):
         return self._manifest.format
@@ -167,13 +216,23 @@ def format(self):
     def license(self):
         return self._manifest.license
 
+    @license.setter
+    def license(self, value):
+        self._manifest.license = value
+
     @property
     def fields(self) -> typing.List[DatasetField]:
         return list(self._manifest.fields.values())
 
     def get_field(self, name: str) -> DatasetField:
         return self._manifest.fields[name]
 
+    def get_by_uid(self, uid: str) -> typing.Optional[typing.Dict]:
+        for sample in self._data:
+            if sample["uid"] == uid:
+                return sample
+        return None
+
     def __getitem__(self, key: str):
         return [x[key] for x in self._data]
 

continuous_eval/eval/modules.py

@@ -3,9 +3,25 @@
 
 from continuous_eval.eval.dataset import DatasetField
 from continuous_eval.eval.tests import Test
+from continuous_eval.eval.utils import type_hint_to_str
 from continuous_eval.metrics import Metric
 
 
+def _serialize_input_type(obj):
+    if isinstance(obj, DatasetField):
+        return {"__class__": obj.__class__.__name__, "name": obj.name}
+    elif isinstance(obj, Module):
+        return {"__class__": obj.__class__.__name__, "name": obj.name}
+    elif isinstance(obj, type):
+        return type_hint_to_str(obj)
+    elif isinstance(obj, (list, tuple)):
+        return [_serialize_input_type(x) for x in obj]
+    elif obj is None:
+        return "None"
+    else:
+        raise TypeError(f"Object of type {type(obj).__name__} is not serializable")
+
+
 @dataclass(frozen=True, eq=True)
 class Tool:
     name: str
@@ -33,6 +49,16 @@ def __post_init__(self):
             eval_names = {metric.name for metric in self.eval}
             assert len(eval_names) == len(self.eval), f"Each metric name must be unique"
 
+    def asdict(self):
+        return {
+            "name": self.name,
+            "input": _serialize_input_type(self.input),
+            "output": type_hint_to_str(self.output),
+            "description": self.description,
+            "eval": [metric.asdict() for metric in self.eval] if self.eval else None,
+            "tests": [test.asdict() for test in self.tests] if self.tests else None,
+        }
+
 
 @dataclass(frozen=True, eq=True)
 class AgentModule(Module):

continuous_eval/eval/pipeline.py

@@ -1,5 +1,5 @@
 from dataclasses import dataclass, field
-from typing import Any, Callable, List, Optional, Set, Tuple
+from typing import Any, Callable, List, Optional, Set, Tuple, Union
 
 from continuous_eval.eval.dataset import Dataset, DatasetField
 from continuous_eval.eval.modules import Module, SingleModule
@@ -11,7 +11,7 @@
 @dataclass
 class ModuleOutput:
     selector: Callable = field(default=lambda x: x)
-    module: Optional[Module] = None
+    module: Optional[Union[Module, str]] = None
 
     def __call__(self, *args: Any) -> Any:
         return self.selector(*args)
@@ -34,7 +34,7 @@ class Graph:
 
 
 class Pipeline:
-    def __init__(self, modules: List[Module], dataset: Dataset) -> None:
+    def __init__(self, modules: List[Module], dataset: Optional[Dataset] = None) -> None:
         self._modules = modules
         self._dataset = dataset
         self._graph = self._build_graph()
@@ -47,6 +47,10 @@ def modules(self):
     def dataset(self):
         return self._dataset
 
+    @dataset.setter
+    def dataset(self, dataset: Dataset):
+        self._dataset = dataset
+
     def module_by_name(self, name: str) -> Module:
         for module in self._modules:
             if module.name == name:
@@ -71,6 +75,8 @@ def _validate_modules(self):
             names.add(module.name)
 
     def _build_graph(self):
+        if self._dataset is None:
+            return None
         nodes = {m.name for m in self._modules}
         edges = set()
         dataset_edges = set()
@@ -96,6 +102,8 @@ def _build_graph(self):
         return Graph(nodes, edges, dataset_edges)
 
     def graph_repr(self, with_type_hints: bool = False):
+        if self._graph is None:
+            return None
         repr_str = "graph TD;\n"
         dataset_node_label = "Dataset"
         repr_str += f"    {dataset_node_label}(({dataset_node_label}));\n"
@@ -112,9 +120,14 @@ def graph_repr(self, with_type_hints: bool = False):
             repr_str += f'    {dataset_node_label} -. "{dataset_field_name}" .-> {end_node};\n'
         return repr_str
 
+    def asdict(self):
+        return {
+            "modules": [m.asdict() for m in self._modules],
+        }
+
 
 def SingleModulePipeline(
-    dataset: Dataset,
+    dataset: Optional[Dataset] = None,
     eval: Optional[List[Metric]] = None,
     tests: Optional[List[Test]] = None,
     name: str = "eval",

continuous_eval/eval/runner.py

@@ -1,7 +1,7 @@
 import logging
 from typing import Optional, Union
 
-from continuous_eval.eval.dataset import Dataset, DatasetField
+from continuous_eval.eval.dataset import Dataset, DatasetField, LambdaField
 from continuous_eval.eval.logger import PipelineLogger
 from continuous_eval.eval.modules import Module
 from continuous_eval.eval.pipeline import CalledTools, ModuleOutput, Pipeline
@@ -30,22 +30,49 @@ def dataset(self) -> Dataset:
         return self._pipeline.dataset
 
     # Evaluate
-
-    def _prepare(self, eval_results: PipelineResults, module: Module, metric: Metric):
+    @staticmethod
+    def prepare(dataset: Dataset, eval_results: PipelineResults, module: Module, metric: Metric):
         kwargs = dict()
         if metric.overloaded_params is not None:
             for key, val in metric.overloaded_params.items():
+                if key == "uid":
+                    continue
                 if isinstance(val, DatasetField):
-                    kwargs[key] = [x[val.name] for x in self.dataset.data]  # type: ignore
+                    kwargs[key] = [x[module.name][val.name] if module.name in x else x[val.name] for x in dataset.data]  # type: ignore
+                elif isinstance(val, LambdaField):
+                    kwargs[key] = list()
+                    for rx in eval_results.results:
+                        uid = rx["uid"]
+                        if module.name in rx:
+                            kwargs[key].append(val.func(rx[module.name]))
+                        else:
+                            for x in dataset.data:
+                                if x["uid"] == uid:
+                                    kwargs[key].append(val.func(x))
+                                    break
+                    # kwargs[key] = [
+                    #     val.func(x[module.name]) if module.name in x else val.func(x)
+                    #     for x in dataset.data
+                    # ]
                 elif isinstance(val, ModuleOutput):
-                    module_name = module.name if val.module is None else val.module.name
+                    module_name = module.name if val.module is None else val.module
+                    if isinstance(val, Module):
+                        module_name = val.name
                     kwargs[key] = [val(x[module_name]) for x in eval_results.results]
                 elif isinstance(val, CalledTools):
                     module_name = module.name if val.module is None else val.module.name
                     val_key = f"{TOOL_PREFIX}{module_name}"
                     kwargs[key] = [val(x[val_key]) for x in eval_results.results]
                 else:
                     raise ValueError(f"Invalid promised parameter {key}={val}")
+            return kwargs
+        else:
+            for item in eval_results.results:
+                itr = item[module.name] if module.name in item else item
+                for key, value in itr.items():
+                    if key not in kwargs:
+                        kwargs[key] = []
+                    kwargs[key].append(value)
         return kwargs
 
     @telemetry_event("eval_manager")
@@ -67,7 +94,8 @@ def evaluate(
         metrics_results = MetricsResults(self.pipeline)
         metrics_results.samples = {
             module.name: {
-                metric.name: metric.batch(**self._prepare(eval_results, module, metric)) for metric in module.eval
+                metric.name: metric.batch(**self.prepare(self.dataset, eval_results, module, metric))
+                for metric in module.eval
             }
             for module in self._pipeline.modules
             if module.eval is not None

continuous_eval/eval/tests.py

@@ -20,6 +20,12 @@ def run(self, metrics_per_sample) -> bool:
         """
         raise NotImplementedError
 
+    def asdict(self):
+        return {
+            "__class__": self.__class__.__name__,
+            "name": self.name,
+        }
+
 
 # Some common tests
 class GreaterOrEqualThan(Test):

continuous_eval/metrics/base.py

@@ -58,6 +58,12 @@ def aggregate(self, results: List[Any]) -> Any:
     def name(self):
         return self.__class__.__name__
 
+    def asdict(self):
+        return {
+            "__class__": self.__class__.__name__,
+            "name": self.name,
+        }
+
 
 class LLMBasedMetric(Metric):
     """