arize 8.0.0a14__py3-none-any.whl → 8.0.0a16__py3-none-any.whl

arize/experiments/tracing.py

@@ -0,0 +1,276 @@
+from __future__ import annotations
+
+import inspect
+import json
+from contextlib import contextmanager
+from contextvars import ContextVar
+from threading import Lock
+from typing import (
+    Any,
+    Callable,
+    Iterable,
+    Iterator,
+    List,
+    Mapping,
+    Sequence,
+    cast,
+)
+
+import numpy as np
+from openinference.semconv import trace
+from openinference.semconv.trace import DocumentAttributes, SpanAttributes
+from opentelemetry.sdk.resources import Resource
+from opentelemetry.sdk.trace import ReadableSpan
+from opentelemetry.trace import INVALID_TRACE_ID
+from typing_extensions import assert_never
+from wrapt import apply_patch, resolve_path, wrap_function_wrapper
+
+
+class SpanModifier:
+    """
+    A class that modifies spans with the specified resource attributes.
+    """
+
+    __slots__ = ("_resource",)
+
+    def __init__(self, resource: Resource) -> None:
+        self._resource = resource
+
+    def modify_resource(self, span: ReadableSpan) -> None:
+        """
+        Takes a span and merges in the resource attributes specified in the constructor.
+
+        Args:
+          span: ReadableSpan: the span to modify
+
+        """
+        if (ctx := span._context) is None or ctx.span_id == INVALID_TRACE_ID:
+            return
+        span._resource = span._resource.merge(self._resource)
+
+
+_ACTIVE_MODIFIER: ContextVar[SpanModifier | None] = ContextVar(
+    "active_modifier"
+)
+
+
+def override_span(
+    init: Callable[..., None], span: ReadableSpan, args: Any, kwargs: Any
+) -> None:
+    init(*args, **kwargs)
+    if isinstance(span_modifier := _ACTIVE_MODIFIER.get(None), SpanModifier):
+        span_modifier.modify_resource(span)
+
+
+_SPAN_INIT_MONKEY_PATCH_LOCK = Lock()
+_SPAN_INIT_MONKEY_PATCH_COUNT = 0
+_SPAN_INIT_MODULE = ReadableSpan.__init__.__module__
+_SPAN_INIT_NAME = ReadableSpan.__init__.__qualname__
+_SPAN_INIT_PARENT, _SPAN_INIT_ATTR, _SPAN_INIT_ORIGINAL = resolve_path(
+    _SPAN_INIT_MODULE, _SPAN_INIT_NAME
+)
+
+
+@contextmanager
+def _monkey_patch_span_init() -> Iterator[None]:
+    global _SPAN_INIT_MONKEY_PATCH_COUNT
+    with _SPAN_INIT_MONKEY_PATCH_LOCK:
+        _SPAN_INIT_MONKEY_PATCH_COUNT += 1
+        if _SPAN_INIT_MONKEY_PATCH_COUNT == 1:
+            wrap_function_wrapper(
+                module=_SPAN_INIT_MODULE,
+                name=_SPAN_INIT_NAME,
+                wrapper=override_span,
+            )
+    yield
+    with _SPAN_INIT_MONKEY_PATCH_LOCK:
+        _SPAN_INIT_MONKEY_PATCH_COUNT -= 1
+        if _SPAN_INIT_MONKEY_PATCH_COUNT == 0:
+            apply_patch(_SPAN_INIT_PARENT, _SPAN_INIT_ATTR, _SPAN_INIT_ORIGINAL)
+
+
+@contextmanager
+def capture_spans(resource: Resource) -> Iterator[SpanModifier]:
+    """
+    A context manager that captures spans and modifies them with the specified resources.
+
+    Args:
+      resource: Resource: The resource to merge into the spans created within the context.
+
+    Returns:
+        modifier: Iterator[SpanModifier]: The span modifier that is active within the context.
+
+    """
+    modifier = SpanModifier(resource)
+    with _monkey_patch_span_init():
+        token = _ACTIVE_MODIFIER.set(modifier)
+        yield modifier
+        _ACTIVE_MODIFIER.reset(token)
+
+
+"""
+Span attribute keys have a special relationship with the `.` separator. When
+a span attribute is ingested from protobuf, it's in the form of a key value
+pair such as `("llm.token_count.completion", 123)`. What we need to do is to split
+the key by the `.` separator and turn it into part of a nested dictionary such
+as {"llm": {"token_count": {"completion": 123}}}. We also need to reverse this
+process, which is to flatten the nested dictionary into a list of key value
+pairs. This module provides functions to do both of these operations.
+
+Note that digit keys are treated as indices of a nested array. For example,
+the digits inside `("retrieval.documents.0.document.content", 'A')` and
+`("retrieval.documents.1.document.content": 'B')` turn the sub-keys following
+them into a nested list of dictionaries i.e.
+{`retrieval: {"documents": [{"document": {"content": "A"}}, {"document":
+{"content": "B"}}]}`.
+"""
+
+
+DOCUMENT_METADATA = DocumentAttributes.DOCUMENT_METADATA
+LLM_PROMPT_TEMPLATE_VARIABLES = SpanAttributes.LLM_PROMPT_TEMPLATE_VARIABLES
+METADATA = SpanAttributes.METADATA
+TOOL_PARAMETERS = SpanAttributes.TOOL_PARAMETERS
+
+# attributes interpreted as JSON strings during ingestion
+JSON_STRING_ATTRIBUTES = (
+    DOCUMENT_METADATA,
+    LLM_PROMPT_TEMPLATE_VARIABLES,
+    METADATA,
+    TOOL_PARAMETERS,
+)
+
+SEMANTIC_CONVENTIONS: List[str] = sorted(
+    # e.g. "input.value", "llm.token_count.total", etc.
+    (
+        cast(str, getattr(klass, attr))
+        for name in dir(trace)
+        if name.endswith("Attributes")
+        and inspect.isclass(klass := getattr(trace, name))
+        for attr in dir(klass)
+        if attr.isupper()
+    ),
+    key=len,
+    reverse=True,
+)  # sorted so the longer strings go first
+
+
+def flatten(
+    obj: Mapping[str, Any] | Iterable[Any],
+    *,
+    prefix: str = "",
+    separator: str = ".",
+    recurse_on_sequence: bool = False,
+    json_string_attributes: Sequence[str] | None = None,
+) -> Iterator[tuple[str, Any]]:
+    """
+    Flatten a nested dictionary or a sequence of dictionaries into a list of
+    key value pairs. If `recurse_on_sequence` is True, then the function will
+    also recursively flatten nested sequences of dictionaries. If
+    `json_string_attributes` is provided, then the function will interpret the
+    attributes in the list as JSON strings and convert them into dictionaries.
+    The `prefix` argument is used to prefix the keys in the output list, but
+    it's mostly used internally to facilitate recursion.
+    """
+    if isinstance(obj, Mapping):
+        yield from _flatten_mapping(
+            obj,
+            prefix=prefix,
+            recurse_on_sequence=recurse_on_sequence,
+            json_string_attributes=json_string_attributes,
+            separator=separator,
+        )
+    elif isinstance(obj, Iterable):
+        yield from _flatten_sequence(
+            obj,
+            prefix=prefix,
+            recurse_on_sequence=recurse_on_sequence,
+            json_string_attributes=json_string_attributes,
+            separator=separator,
+        )
+    else:
+        assert_never(obj)
+
+
+def has_mapping(sequence: Iterable[Any]) -> bool:
+    """
+    Check if a sequence contains a dictionary. We don't flatten sequences that
+    only contain primitive types, such as strings, integers, etc. Conversely,
+    we'll only un-flatten digit sub-keys if it can be interpreted the index of
+    an array of dictionaries.
+    """
+    return any(isinstance(item, Mapping) for item in sequence)
+
+
+def _flatten_mapping(
+    mapping: Mapping[str, Any],
+    *,
+    prefix: str = "",
+    recurse_on_sequence: bool = False,
+    json_string_attributes: Sequence[str] | None = None,
+    separator: str = ".",
+) -> Iterator[tuple[str, Any]]:
+    """
+    Flatten a nested dictionary into a list of key value pairs. If `recurse_on_sequence`
+    is True, then the function will also recursively flatten nested sequences of dictionaries.
+    If `json_string_attributes` is provided, then the function will interpret the attributes
+    in the list as JSON strings and convert them into dictionaries. The `prefix` argument is
+    used to prefix the keys in the output list, but it's mostly used internally to facilitate
+    recursion.
+    """
+    for key, value in mapping.items():
+        prefixed_key = f"{prefix}{separator}{key}" if prefix else key
+        if isinstance(value, Mapping):
+            if json_string_attributes and prefixed_key.endswith(
+                JSON_STRING_ATTRIBUTES
+            ):
+                yield prefixed_key, json.dumps(value)
+            else:
+                yield from _flatten_mapping(
+                    value,
+                    prefix=prefixed_key,
+                    recurse_on_sequence=recurse_on_sequence,
+                    json_string_attributes=json_string_attributes,
+                    separator=separator,
+                )
+        elif (
+            isinstance(value, (Sequence, np.ndarray))
+        ) and recurse_on_sequence:
+            yield from _flatten_sequence(
+                value,
+                prefix=prefixed_key,
+                recurse_on_sequence=recurse_on_sequence,
+                json_string_attributes=json_string_attributes,
+                separator=separator,
+            )
+        elif value is not None:
+            yield prefixed_key, value
+
+
+def _flatten_sequence(
+    sequence: Iterable[Any],
+    *,
+    prefix: str = "",
+    recurse_on_sequence: bool = False,
+    json_string_attributes: Sequence[str] | None = None,
+    separator: str = ".",
+) -> Iterator[tuple[str, Any]]:
+    """
+    Flatten a sequence of dictionaries into a list of key value pairs. If `recurse_on_sequence`
+    is True, then the function will also recursively flatten nested sequences of dictionaries.
+    If `json_string_attributes` is provided, then the function will interpret the attributes
+    in the list as JSON strings and convert them into dictionaries. The `prefix` argument is
+    used to prefix the keys in the output list, but it's mostly used internally to facilitate
+    recursion.
+    """
+    if isinstance(sequence, str) or not has_mapping(sequence):
+        yield prefix, sequence
+    for idx, obj in enumerate(sequence):
+        if not isinstance(obj, Mapping):
+            continue
+        yield from _flatten_mapping(
+            obj,
+            prefix=f"{prefix}{separator}{idx}" if prefix else f"{idx}",
+            recurse_on_sequence=recurse_on_sequence,
+            json_string_attributes=json_string_attributes,
+            separator=separator,
+        )

arize/experiments/types.py

@@ -0,0 +1,394 @@
+from __future__ import annotations
+
+import json
+import textwrap
+from copy import copy, deepcopy
+from dataclasses import dataclass, field
+from datetime import datetime, timezone
+from importlib.metadata import version
+from random import getrandbits
+from typing import (
+    Any,
+    Awaitable,
+    Callable,
+    Iterable,
+    Mapping,
+    cast,
+)
+
+import pandas as pd
+from wrapt import ObjectProxy
+
+from arize.experiments.evaluators.types import (
+    EvaluationResult,
+    JSONSerializable,
+)
+
+ExperimentId = str
+# DatasetId= str
+# DatasetVersionId= str
+ExampleId = str
+RepetitionNumber = int
+ExperimentRunId = str
+TraceId = str
+
+
+@dataclass(frozen=True)
+class Example:
+    """
+    Represents an example in an experiment dataset.
+    Args:
+        id: The unique identifier for the example.
+        updated_at: The timestamp when the example was last updated.
+        input: The input data for the example.
+        output: The output data for the example.
+        metadata: Additional metadata for the example.
+        dataset_row: The original dataset row containing the example data.
+    """
+
+    id: ExampleId = field(default_factory=str)
+    updated_at: datetime = field(default_factory=datetime.now)
+    input: Mapping[str, JSONSerializable] = field(default_factory=dict)
+    output: Mapping[str, JSONSerializable] = field(default_factory=dict)
+    metadata: Mapping[str, JSONSerializable] = field(default_factory=dict)
+    dataset_row: Mapping[str, JSONSerializable] = field(default_factory=dict)
+
+    def __post_init__(self) -> None:
+        if self.dataset_row is not None:
+            object.__setattr__(
+                self, "dataset_row", _make_read_only(self.dataset_row)
+            )
+            if "attributes.input.value" in self.dataset_row:
+                object.__setattr__(
+                    self,
+                    "input",
+                    _make_read_only(self.dataset_row["attributes.input.value"]),
+                )
+            if "attributes.output.value" in self.dataset_row:
+                object.__setattr__(
+                    self,
+                    "output",
+                    _make_read_only(
+                        self.dataset_row["attributes.output.value"]
+                    ),
+                )
+            if "attributes.metadata" in self.dataset_row:
+                object.__setattr__(
+                    self,
+                    "metadata",
+                    _make_read_only(self.dataset_row["attributes.metadata"]),
+                )
+            if "id" in self.dataset_row:
+                object.__setattr__(self, "id", self.dataset_row["id"])
+            if "updated_at" in self.dataset_row:
+                object.__setattr__(
+                    self, "updated_at", self.dataset_row["updated_at"]
+                )
+        else:
+            object.__setattr__(self, "input", self.input)
+            object.__setattr__(self, "output", self.output)
+            object.__setattr__(self, "metadata", self.metadata)
+
+    @classmethod
+    def from_dict(cls, obj: Mapping[str, Any]) -> Example:
+        return cls(
+            id=obj["id"],
+            input=obj["input"],
+            output=obj["output"],
+            metadata=obj.get("metadata") or {},
+            updated_at=obj["updated_at"],
+        )
+
+    def __repr__(self) -> str:
+        spaces = " " * 4
+        name = self.__class__.__name__
+        identifiers = [f'{spaces}id="{self.id}",']
+        contents = []
+        for key in ("input", "output", "metadata", "dataset_row"):
+            value = getattr(self, key, None)
+            if value:
+                contents.append(
+                    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")
+                    + ","
+                )
+        return "\n".join([f"{name}(", *identifiers, *contents, ")"])
+
+
+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__)
+
+
+def _blue(text: str) -> str:
+    return f"\033[1m\033[94m{text}\033[0m"
+
+
+@dataclass(frozen=True)
+class TestCase:
+    example: Example
+    repetition_number: RepetitionNumber
+
+
+EXP_ID: ExperimentId = "EXP_ID"
+
+
+def _exp_id() -> str:
+    suffix = getrandbits(24).to_bytes(3, "big").hex()
+    return f"{EXP_ID}_{suffix}"
+
+
+@dataclass(frozen=True)
+class ExperimentRun:
+    """
+    Represents a single run of an experiment.
+    Args:
+        start_time: The start time of the experiment run.
+        end_time: The end time of the experiment run.
+        experiment_id: The unique identifier for the experiment.
+        dataset_example_id: The unique identifier for the dataset example.
+        repetition_number: The repetition number of the experiment run.
+        output: The output of the experiment run.
+        error: The error message if the experiment run failed.
+        id: The unique identifier for the experiment run.
+        trace_id: The trace identifier for the experiment run.
+    """
+
+    start_time: datetime
+    end_time: datetime
+    experiment_id: ExperimentId
+    dataset_example_id: ExampleId
+    repetition_number: RepetitionNumber
+    output: JSONSerializable
+    error: str | None = None
+    id: ExperimentRunId = field(default_factory=_exp_id)
+    trace_id: TraceId | None = 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) == (self.error is None):
+            raise ValueError(
+                "Must specify exactly one of experiment_run_output or error"
+            )
+
+
+@dataclass(frozen=True)
+class ExperimentEvaluationRun:
+    """
+    Represents a single evaluation run of an experiment.
+    Args:
+        experiment_run_id: The unique identifier for the experiment run.
+        start_time: The start time of the evaluation run.
+        end_time: The end time of the evaluation run.
+        name: The name of the evaluation run.
+        annotator_kind: The kind of annotator used in the evaluation run.
+        error: The error message if the evaluation run failed.
+        result (Optional[EvaluationResult]): The result of the evaluation run.
+        id (str): The unique identifier for the evaluation run.
+        trace_id (Optional[TraceId]): The trace identifier for the evaluation run.
+    """
+
+    experiment_run_id: ExperimentRunId
+    start_time: datetime
+    end_time: datetime
+    name: str
+    annotator_kind: str
+    error: str | None = None
+    result: EvaluationResult | None = None
+    id: str = field(default_factory=_exp_id)
+    trace_id: TraceId | None = 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 bool(self.result) == bool(self.error):
+            raise ValueError("Must specify either result or error")
+
+
+_LOCAL_TIMEZONE = datetime.now(timezone.utc).astimezone().tzinfo
+
+
+def local_now() -> datetime:
+    return datetime.now(timezone.utc).astimezone(tz=_LOCAL_TIMEZONE)
+
+
+@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  # type: ignore
+
+
+@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, n_examples: int, task_runs: Iterable[ExperimentRun | None]
+    ) -> _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": n_examples,
+            "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, length: int = 100) -> str | None:
+    if (cnt := s.dropna().str.slice(0, length).value_counts()).empty:
+        return None
+    return cast(str, cnt.sort_values(ascending=False).index[0])
+
+
+@dataclass
+class ExperimentTaskResultFieldNames:
+    """Column names for mapping experiment task results in a DataFrame.
+
+    Args:
+        example_id: Name of column containing example IDs.
+            The ID values must match the id of the dataset rows.
+        result: Name of column containing task results
+    """
+
+    example_id: str
+    result: str
+
+
+TaskOutput = JSONSerializable
+ExampleOutput = Mapping[str, JSONSerializable]
+ExampleMetadata = Mapping[str, JSONSerializable]
+ExampleInput = Mapping[str, JSONSerializable]
+ExperimentTask = (
+    Callable[[Example], TaskOutput] | Callable[[Example], Awaitable[TaskOutput]]
+)

arize/models/client.py

@@ -584,7 +584,7 @@ class MLModelsClient:
             # pyarrow will err if a mixed type column exist in the dataset even if
             # the column is not specified in schema. Caveat: There may be other
             # error conditions that we're currently not aware of.
-            pa_table = pa.Table.from_pandas(dataframe)
+            pa_table = pa.Table.from_pandas(dataframe, preserve_index=False)
         except pa.ArrowInvalid as e:
             logger.error(f"{INVALID_ARROW_CONVERSION_MSG}: {str(e)}")
             raise pa.ArrowInvalid(
@@ -660,6 +660,7 @@ class MLModelsClient:
             headers=headers,
             timeout=timeout,
             verify=self._sdk_config.request_verify,
+            max_chunksize=self._sdk_config.pyarrow_max_chunksize,
             tmp_dir=tmp_dir,
         )
 
@@ -688,6 +689,7 @@ class MLModelsClient:
             port=self._sdk_config.flight_server_port,
             scheme=self._sdk_config.flight_scheme,
             request_verify=self._sdk_config.request_verify,
+            max_chunksize=self._sdk_config.pyarrow_max_chunksize,
         ) as flight_client:
             exporter = ArizeExportClient(
                 flight_client=flight_client,
@@ -732,6 +734,7 @@ class MLModelsClient:
             port=self._sdk_config.flight_server_port,
             scheme=self._sdk_config.flight_scheme,
             request_verify=self._sdk_config.request_verify,
+            max_chunksize=self._sdk_config.pyarrow_max_chunksize,
         ) as flight_client:
             exporter = ArizeExportClient(
                 flight_client=flight_client,