arize-phoenix 3.16.1__py3-none-any.whl → 7.7.0__py3-none-any.whl

phoenix/experiments/types.py

@@ -0,0 +1,726 @@
+from __future__ import annotations
+
+import json
+import textwrap
+from collections import Counter
+from collections.abc import (
+    Awaitable,
+    Callable,
+    Iterable,
+    Iterator,
+    Mapping,
+)
+from copy import copy, deepcopy
+from dataclasses import dataclass, field, fields
+from datetime import datetime
+from enum import Enum
+from functools import cached_property
+from importlib.metadata import version
+from random import getrandbits
+from typing import Any, Optional, TypeVar, Union, cast, overload
+
+import pandas as pd
+from typing_extensions import TypeAlias
+from wrapt import ObjectProxy
+
+from phoenix.datetime_utils import local_now
+from phoenix.experiments.utils import get_experiment_url
+
+
+class AnnotatorKind(Enum):
+    CODE = "CODE"
+    LLM = "LLM"
+
+
+JSONSerializable: TypeAlias = Optional[Union[dict[str, Any], list[Any], str, int, float, bool]]
+ExperimentId: TypeAlias = str
+DatasetId: TypeAlias = str
+DatasetVersionId: TypeAlias = str
+ExampleId: TypeAlias = str
+RepetitionNumber: TypeAlias = int
+ExperimentRunId: TypeAlias = str
+TraceId: TypeAlias = str
+
+TaskOutput: TypeAlias = JSONSerializable
+
+ExampleOutput: TypeAlias = Mapping[str, JSONSerializable]
+ExampleMetadata: TypeAlias = Mapping[str, JSONSerializable]
+ExampleInput: TypeAlias = Mapping[str, JSONSerializable]
+
+Score: TypeAlias = Optional[Union[bool, int, float]]
+Label: TypeAlias = Optional[str]
+Explanation: TypeAlias = Optional[str]
+
+EvaluatorName: TypeAlias = str
+EvaluatorKind: TypeAlias = str
+EvaluatorOutput: TypeAlias = Union[
+    "EvaluationResult", bool, int, float, str, tuple[Score, Label, Explanation]
+]
+
+DRY_RUN: ExperimentId = "DRY_RUN"
+
+
+def _dry_run_id() -> str:
+    suffix = getrandbits(24).to_bytes(3, "big").hex()
+    return f"{DRY_RUN}_{suffix}"
+
+
+@dataclass(frozen=True)
+class Example:
+    id: ExampleId
+    updated_at: datetime
+    input: Mapping[str, JSONSerializable] = field(default_factory=dict)
+    output: Mapping[str, JSONSerializable] = field(default_factory=dict)
+    metadata: Mapping[str, JSONSerializable] = field(default_factory=dict)
+
+    def __post_init__(self) -> None:
+        object.__setattr__(self, "input", _make_read_only(self.input))
+        object.__setattr__(self, "output", _make_read_only(self.output))
+        object.__setattr__(self, "metadata", _make_read_only(self.metadata))
+
+    @classmethod
+    def from_dict(cls, obj: Mapping[str, Any]) -> Example:
+        return cls(
+            input=obj["input"],
+            output=obj["output"],
+            metadata=obj.get("metadata") or {},
+            id=obj["id"],
+            updated_at=obj["updated_at"],
+        )
+
+    def __repr__(self) -> str:
+        spaces = " " * 4
+        name = self.__class__.__name__
+        identifiers = [f'{spaces}id="{self.id}",']
+        contents = [
+            spaces
+            + f"{_blue(key)}="
+            + json.dumps(
+                _shorten(value),
+                ensure_ascii=False,
+                sort_keys=True,
+                indent=len(spaces),
+            )
+            .replace("\n", f"\n{spaces}")
+            .replace(' "..."\n', " ...\n")
+            + ","
+            for key in ("input", "output", "metadata")
+            if (value := getattr(self, key, None))
+        ]
+        return "\n".join([f"{name}(", *identifiers, *contents, ")"])
+
+
+@dataclass(frozen=True)
+class Dataset:
+    id: DatasetId
+    version_id: DatasetVersionId
+    examples: Mapping[ExampleId, Example] = field(repr=False, default_factory=dict)
+
+    def __post_init__(self) -> None:
+        object.__setattr__(self, "examples", _ReadOnly(self.examples))
+
+    def __len__(self) -> int:
+        return len(self.examples)
+
+    def __iter__(self) -> Iterator[Example]:
+        return iter(self.examples.values())
+
+    @cached_property
+    def _keys(self) -> tuple[str, ...]:
+        return tuple(self.examples.keys())
+
+    @overload
+    def __getitem__(self, key: int) -> Example: ...
+    @overload
+    def __getitem__(self, key: slice) -> list[Example]: ...
+    def __getitem__(self, key: Union[int, slice]) -> Union[Example, list[Example]]:
+        if isinstance(key, int):
+            return self.examples[self._keys[key]]
+        return [self.examples[k] for k in self._keys[key]]
+
+    def as_dataframe(self, drop_empty_columns: bool = True) -> pd.DataFrame:
+        df = pd.DataFrame.from_records(
+            [
+                {
+                    "example_id": example.id,
+                    "input": deepcopy(example.input),
+                    "output": deepcopy(example.output),
+                    "metadata": deepcopy(example.metadata),
+                }
+                for example in self.examples.values()
+            ]
+        ).set_index("example_id")
+        if drop_empty_columns:
+            return df.reindex([k for k, v in df.items() if v.astype(bool).any()], axis=1)
+        return df
+
+    @classmethod
+    def from_dict(cls, obj: Mapping[str, Any]) -> Dataset:
+        examples = tuple(map(Example.from_dict, obj.get("examples") or ()))
+        return cls(
+            id=obj["dataset_id"],
+            version_id=obj["version_id"],
+            examples={ex.id: ex for ex in examples},
+        )
+
+
+@dataclass(frozen=True)
+class TestCase:
+    example: Example
+    repetition_number: RepetitionNumber
+
+
+@dataclass(frozen=True)
+class Experiment:
+    id: ExperimentId
+    dataset_id: DatasetId
+    dataset_version_id: DatasetVersionId
+    repetitions: int
+    project_name: str = field(repr=False)
+
+    @classmethod
+    def from_dict(cls, obj: Mapping[str, Any]) -> Experiment:
+        return cls(
+            id=obj["id"],
+            dataset_id=obj["dataset_id"],
+            dataset_version_id=obj["dataset_version_id"],
+            repetitions=obj.get("repetitions") or 1,
+            project_name=obj.get("project_name") or "",
+        )
+
+
+@dataclass(frozen=True)
+class ExperimentRun:
+    start_time: datetime
+    end_time: datetime
+    experiment_id: ExperimentId
+    dataset_example_id: ExampleId
+    repetition_number: RepetitionNumber
+    output: JSONSerializable
+    error: Optional[str] = None
+    id: ExperimentRunId = field(default_factory=_dry_run_id)
+    trace_id: Optional[TraceId] = None
+
+    @classmethod
+    def from_dict(cls, obj: Mapping[str, Any]) -> ExperimentRun:
+        return cls(
+            start_time=obj["start_time"],
+            end_time=obj["end_time"],
+            experiment_id=obj["experiment_id"],
+            dataset_example_id=obj["dataset_example_id"],
+            repetition_number=obj.get("repetition_number") or 1,
+            output=_make_read_only(obj.get("output")),
+            error=obj.get("error"),
+            id=obj["id"],
+            trace_id=obj.get("trace_id"),
+        )
+
+    def __post_init__(self) -> None:
+        if self.output is None and self.error is None:
+            raise ValueError("Must specify exactly one of experiment_run_output or error")
+
+
+@dataclass(frozen=True)
+class EvaluationResult:
+    score: Optional[float] = None
+    label: Optional[str] = None
+    explanation: Optional[str] = None
+    metadata: Mapping[str, JSONSerializable] = field(default_factory=dict)
+
+    @classmethod
+    def from_dict(cls, obj: Optional[Mapping[str, Any]]) -> Optional[EvaluationResult]:
+        if not obj:
+            return None
+        return cls(
+            score=obj.get("score"),
+            label=obj.get("label"),
+            explanation=obj.get("explanation"),
+            metadata=obj.get("metadata") or {},
+        )
+
+    def __post_init__(self) -> None:
+        if self.score is None and not self.label:
+            raise ValueError("Must specify score or label, or both")
+        if self.score is None and not self.label:
+            object.__setattr__(self, "score", 0)
+        for k in ("label", "explanation"):
+            if (v := getattr(self, k, None)) is not None:
+                object.__setattr__(self, k, str(v) or None)
+
+
+@dataclass(frozen=True)
+class ExperimentEvaluationRun:
+    experiment_run_id: ExperimentRunId
+    start_time: datetime
+    end_time: datetime
+    name: str
+    annotator_kind: str
+    error: Optional[str] = None
+    result: Optional[EvaluationResult] = None
+    id: str = field(default_factory=_dry_run_id)
+    trace_id: Optional[TraceId] = None
+
+    @classmethod
+    def from_dict(cls, obj: Mapping[str, Any]) -> ExperimentEvaluationRun:
+        return cls(
+            experiment_run_id=obj["experiment_run_id"],
+            start_time=obj["start_time"],
+            end_time=obj["end_time"],
+            name=obj["name"],
+            annotator_kind=obj["annotator_kind"],
+            error=obj.get("error"),
+            result=EvaluationResult.from_dict(obj.get("result")),
+            id=obj["id"],
+            trace_id=obj.get("trace_id"),
+        )
+
+    def __post_init__(self) -> None:
+        if self.result is None and self.error is None:
+            raise ValueError("Must specify either result or error")
+
+
+ExperimentTask: TypeAlias = Union[
+    Callable[[Example], TaskOutput],
+    Callable[[Example], Awaitable[TaskOutput]],
+]
+
+
+@dataclass(frozen=True)
+class ExperimentParameters:
+    n_examples: int
+    n_repetitions: int = 1
+
+    @property
+    def count(self) -> int:
+        return self.n_examples * self.n_repetitions
+
+
+@dataclass(frozen=True)
+class EvaluationParameters:
+    eval_names: frozenset[str]
+    exp_params: ExperimentParameters
+
+
+@dataclass(frozen=True)
+class _HasStats:
+    _title: str = field(repr=False, default="")
+    _timestamp: datetime = field(repr=False, default_factory=local_now)
+    stats: pd.DataFrame = field(repr=False, default_factory=pd.DataFrame)
+
+    @property
+    def title(self) -> str:
+        return f"{self._title} ({self._timestamp:%x %I:%M %p %z})"
+
+    def __str__(self) -> str:
+        try:
+            assert int(version("pandas").split(".")[0]) >= 1
+            # `tabulate` is used by pandas >= 1.0 in DataFrame.to_markdown()
+            import tabulate  # noqa: F401
+        except (AssertionError, ImportError):
+            text = self.stats.__str__()
+        else:
+            text = self.stats.to_markdown(index=False)
+        return f"{self.title}\n{'-'*len(self.title)}\n" + text
+
+
+@dataclass(frozen=True)
+class EvaluationSummary(_HasStats):
+    """
+    Summary statistics of experiment evaluations.
+
+    Users should not instantiate this directly.
+    """
+
+    _title: str = "Experiment Summary"
+
+    @classmethod
+    def from_eval_runs(
+        cls,
+        params: EvaluationParameters,
+        *eval_runs: Optional[ExperimentEvaluationRun],
+    ) -> EvaluationSummary:
+        df = pd.DataFrame.from_records(
+            [
+                {
+                    "evaluator": run.name,
+                    "error": run.error,
+                    "score": run.result.score if run.result else None,
+                    "label": run.result.label if run.result else None,
+                }
+                for run in eval_runs
+                if run is not None
+            ]
+        )
+        if df.empty:
+            df = pd.DataFrame.from_records(
+                [
+                    {"evaluator": name, "error": None, "score": None, "label": None}
+                    for name in params.eval_names
+                ]
+            )
+        has_error = bool(df.loc[:, "error"].astype(bool).sum())
+        has_score = bool(df.loc[:, "score"].dropna().count())
+        has_label = bool(df.loc[:, "label"].astype(bool).sum())
+        agg = {
+            **(
+                dict(n_errors=("error", "count"), top_error=("error", _top_string))
+                if has_error
+                else {}
+            ),
+            **(dict(n_scores=("score", "count"), avg_score=("score", "mean")) if has_score else {}),
+            **(
+                dict(
+                    n_labels=("label", "count"),
+                    top_2_labels=(
+                        "label",
+                        lambda s: (dict(Counter(s.dropna()).most_common(2)) or None),
+                    ),
+                )
+                if has_label
+                else {}
+            ),
+        }
+        stats = (
+            df.groupby("evaluator").agg(**agg)  # type: ignore[call-overload]
+            if agg
+            else pd.DataFrame()
+        )
+        sorted_eval_names = sorted(params.eval_names)
+        eval_names = pd.DataFrame(
+            {
+                "evaluator": sorted_eval_names,
+                "n": [params.exp_params.count] * len(sorted_eval_names),
+            }
+        ).set_index("evaluator")
+        stats = pd.concat([eval_names, stats], axis=1).reset_index()
+        summary: EvaluationSummary = object.__new__(cls)
+        summary.__init__(stats=stats)  # type: ignore[misc]
+        return summary
+
+    @classmethod
+    def __new__(cls, *args: Any, **kwargs: Any) -> Any:
+        # Direct instantiation by users is discouraged.
+        raise NotImplementedError
+
+    @classmethod
+    def __init_subclass__(cls, **kwargs: Any) -> None:
+        # Direct sub-classing by users is discouraged.
+        raise NotImplementedError
+
+
+@dataclass(frozen=True)
+class TaskSummary(_HasStats):
+    """
+    Summary statistics of experiment task executions.
+
+    **Users should not instantiate this object directly.**
+    """
+
+    _title: str = "Tasks Summary"
+
+    @classmethod
+    def from_task_runs(
+        cls,
+        params: ExperimentParameters,
+        task_runs: Iterable[Optional[ExperimentRun]],
+    ) -> "TaskSummary":
+        df = pd.DataFrame.from_records(
+            [
+                {
+                    "example_id": run.dataset_example_id,
+                    "error": run.error,
+                }
+                for run in task_runs
+                if run is not None
+            ]
+        )
+        n_runs = len(df)
+        n_errors = 0 if df.empty else df.loc[:, "error"].astype(bool).sum()
+        record = {
+            "n_examples": params.count,
+            "n_runs": n_runs,
+            "n_errors": n_errors,
+            **(dict(top_error=_top_string(df.loc[:, "error"])) if n_errors else {}),
+        }
+        stats = pd.DataFrame.from_records([record])
+        summary: TaskSummary = object.__new__(cls)
+        summary.__init__(stats=stats)  # type: ignore[misc]
+        return summary
+
+    @classmethod
+    def __new__(cls, *args: Any, **kwargs: Any) -> Any:
+        # Direct instantiation by users is discouraged.
+        raise NotImplementedError
+
+    @classmethod
+    def __init_subclass__(cls, **kwargs: Any) -> None:
+        # Direct sub-classing by users is discouraged.
+        raise NotImplementedError
+
+
+def _top_string(s: "pd.Series[Any]", length: int = 100) -> Optional[str]:
+    if (cnt := s.dropna().str.slice(0, length).value_counts()).empty:
+        return None
+    return cast(str, cnt.sort_values(ascending=False).index[0])
+
+
+@dataclass(frozen=True)
+class RanExperiment(Experiment):
+    """
+    An experiment that has been run.
+
+    **Users should not instantiate this object directly.**
+    """
+
+    params: ExperimentParameters = field(repr=False)
+    dataset: Dataset = field(repr=False)
+    runs: Mapping[ExperimentRunId, ExperimentRun] = field(repr=False)
+    task_summary: TaskSummary = field(repr=False)
+    eval_runs: tuple[ExperimentEvaluationRun, ...] = field(repr=False, default=())
+    eval_summaries: tuple[EvaluationSummary, ...] = field(repr=False, default=())
+
+    @property
+    def url(self) -> str:
+        return get_experiment_url(dataset_id=self.dataset.id, experiment_id=self.id)
+
+    @property
+    def info(self) -> str:
+        return f"🔗 View this experiment: {self.url}"
+
+    def __post_init__(self) -> None:
+        runs = {
+            id_: (
+                _ExperimentRunWithExample(run, example)
+                if (example := self.dataset.examples.get(run.dataset_example_id))
+                else run
+            )
+            for id_, run in self.runs.items()
+        }
+        object.__setattr__(self, "runs", runs)
+
+    def __len__(self) -> int:
+        return len(self.runs)
+
+    def __iter__(self) -> Iterator[ExperimentRun]:
+        return iter(self.runs.values())
+
+    @cached_property
+    def _keys(self) -> tuple[str, ...]:
+        return tuple(self.runs.keys())
+
+    @overload
+    def __getitem__(self, key: int) -> ExperimentRun: ...
+    @overload
+    def __getitem__(self, key: slice) -> list[ExperimentRun]: ...
+    def __getitem__(self, key: Union[int, slice]) -> Union[ExperimentRun, list[ExperimentRun]]:
+        if isinstance(key, int):
+            return self.runs[self._keys[key]]
+        return [self.runs[k] for k in self._keys[key]]
+
+    def get_evaluations(
+        self,
+        drop_empty_columns: bool = True,
+    ) -> pd.DataFrame:
+        df = pd.DataFrame.from_records(
+            [
+                {
+                    "run_id": run.experiment_run_id,
+                    "name": run.name,
+                    "error": run.error,
+                    "score": run.result.score if run.result else None,
+                    "label": run.result.label if run.result else None,
+                    "explanation": run.result.explanation if run.result else None,
+                }
+                for run in self.eval_runs
+            ]
+        ).set_index("run_id")
+        if drop_empty_columns:
+            df = df.reindex([k for k, v in df.items() if v.astype(bool).any()], axis=1)
+        return df.join(self.as_dataframe())
+
+    def as_dataframe(self, drop_empty_columns: bool = True) -> pd.DataFrame:
+        df = pd.DataFrame.from_records(
+            [
+                {
+                    "run_id": run.id,
+                    "error": run.error,
+                    "output": deepcopy(run.output),
+                    "input": deepcopy((ex := self.dataset.examples[run.dataset_example_id]).input),
+                    "expected": deepcopy(ex.output),
+                    "metadata": deepcopy(ex.metadata),
+                    "example_id": run.dataset_example_id,
+                }
+                for run in self.runs.values()
+            ]
+        ).set_index("run_id")
+        if drop_empty_columns:
+            return df.reindex([k for k, v in df.items() if v.astype(bool).any()], axis=1)
+        return df
+
+    def add(
+        self,
+        eval_summary: EvaluationSummary,
+        *eval_runs: Optional[ExperimentEvaluationRun],
+    ) -> "RanExperiment":
+        return _replace(
+            self,
+            eval_runs=(*self.eval_runs, *filter(bool, eval_runs)),
+            eval_summaries=(*self.eval_summaries, eval_summary),
+        )
+
+    def __str__(self) -> str:
+        summaries = (*reversed(self.eval_summaries), self.task_summary)
+        return (
+            "\n"
+            + ("" if self.id.startswith(DRY_RUN) else f"{self.info}\n\n")
+            + "\n\n".join(map(str, summaries))
+        )
+
+    @classmethod
+    def __new__(cls, *args: Any, **kwargs: Any) -> Any:
+        # Direct instantiation by users is discouraged.
+        raise NotImplementedError
+
+    @classmethod
+    def __init_subclass__(cls, **kwargs: Any) -> None:
+        # Direct sub-classing by users is discouraged.
+        raise NotImplementedError
+
+
+def _asdict(dc: Any) -> dict[str, Any]:
+    # non-recursive version of `dataclasses.asdict()`
+    return {field.name: getattr(dc, field.name) for field in fields(dc)}
+
+
+T = TypeVar("T")
+
+
+def _replace(obj: T, **kwargs: Any) -> T:
+    new_obj = object.__new__(obj.__class__)
+    new_obj.__init__(**{**_asdict(obj), **kwargs})  # type: ignore[misc]
+    return new_obj
+
+
+def _shorten(obj: Any, width: int = 50) -> Any:
+    if isinstance(obj, str):
+        return textwrap.shorten(obj, width=width, placeholder="...")
+    if isinstance(obj, dict):
+        return {k: _shorten(v) for k, v in obj.items()}
+    if isinstance(obj, list):
+        if len(obj) > 2:
+            return [_shorten(v) for v in obj[:2]] + ["..."]
+        return [_shorten(v) for v in obj]
+    return obj
+
+
+def _make_read_only(obj: Any) -> Any:
+    if isinstance(obj, dict):
+        return _ReadOnly({k: _make_read_only(v) for k, v in obj.items()})
+    if isinstance(obj, str):
+        return obj
+    if isinstance(obj, list):
+        return _ReadOnly(list(map(_make_read_only, obj)))
+    return obj
+
+
+class _ReadOnly(ObjectProxy):  # type: ignore[misc]
+    def __setitem__(self, *args: Any, **kwargs: Any) -> Any:
+        raise NotImplementedError
+
+    def __delitem__(self, *args: Any, **kwargs: Any) -> Any:
+        raise NotImplementedError
+
+    def __iadd__(self, *args: Any, **kwargs: Any) -> Any:
+        raise NotImplementedError
+
+    def pop(self, *args: Any, **kwargs: Any) -> Any:
+        raise NotImplementedError
+
+    def append(self, *args: Any, **kwargs: Any) -> Any:
+        raise NotImplementedError
+
+    def __copy__(self, *args: Any, **kwargs: Any) -> Any:
+        return copy(self.__wrapped__)
+
+    def __deepcopy__(self, *args: Any, **kwargs: Any) -> Any:
+        return deepcopy(self.__wrapped__)
+
+    def __repr__(self) -> str:
+        return repr(self.__wrapped__)
+
+    def __str__(self) -> str:
+        return str(self.__wrapped__)
+
+
+class _ExperimentRunWithExample(ObjectProxy):  # type: ignore[misc]
+    def __init__(self, wrapped: ExperimentRun, example: Example) -> None:
+        super().__init__(wrapped)
+        self._self_example = example
+
+    @property
+    def expected(self) -> ExampleOutput:
+        return deepcopy(self._self_example.output)
+
+    @property
+    def reference(self) -> ExampleOutput:
+        return deepcopy(self._self_example.output)
+
+    @property
+    def input(self) -> ExampleInput:
+        return deepcopy(self._self_example.input)
+
+    @property
+    def metadata(self) -> ExampleMetadata:
+        return deepcopy(self._self_example.metadata)
+
+    def __repr__(self) -> str:
+        spaces = " " * 4
+        name = self.__class__.__name__
+        identifiers = [
+            f'{spaces}id="{self.id}",',
+            f'{spaces}example_id="{self.dataset_example_id}",',
+        ]
+        outputs = [
+            *([f'{spaces}error="{self.error}",'] if self.error else []),
+            *(
+                [
+                    f"{spaces}{_blue('output')}="
+                    + json.dumps(
+                        _shorten(self.output),
+                        ensure_ascii=False,
+                        sort_keys=True,
+                        indent=len(spaces),
+                    )
+                    .replace("\n", f"\n{spaces}")
+                    .replace(' "..."\n', " ...\n")
+                ]
+                if not self.error
+                else []
+            ),
+        ]
+        dicts = [
+            spaces
+            + f"{_blue(alias)}={{"
+            + (f" # {comment}" if comment else "")
+            + json.dumps(
+                _shorten(value),
+                ensure_ascii=False,
+                sort_keys=True,
+                indent=len(spaces),
+            )[1:]
+            .replace("\n", f"\n{spaces}")
+            .replace(' "..."\n', " ...\n")
+            + ","
+            for alias, value, comment in (
+                ("expected", self.expected, f"alias for the example.{_blue('output')} dict"),
+                ("reference", self.reference, f"alias for the example.{_blue('output')} dict"),
+                ("input", self.input, f"alias for the example.{_blue('input')} dict"),
+                ("metadata", self.metadata, f"alias for the example.{_blue('metadata')} dict"),
+            )
+            if value
+        ]
+        return "\n".join([f"{name}(", *identifiers, *outputs, *dicts, ")"])
+
+
+def _blue(text: str) -> str:
+    return f"\033[1m\033[94m{text}\033[0m"