arize-phoenix 0.0.32__py3-none-any.whl → 0.0.33__py3-none-any.whl

phoenix/core/traces.py

@@ -0,0 +1,285 @@
+import weakref
+from collections import defaultdict
+from datetime import datetime, timezone
+from queue import SimpleQueue
+from threading import RLock, Thread
+from types import MethodType
+from typing import (
+    Any,
+    DefaultDict,
+    Dict,
+    Iterable,
+    Iterator,
+    List,
+    Optional,
+    SupportsFloat,
+    Tuple,
+    Union,
+    cast,
+)
+from uuid import UUID
+
+from sortedcontainers import SortedKeyList
+from typing_extensions import TypeAlias
+from wrapt import ObjectProxy
+
+import phoenix.trace.v1.trace_pb2 as pb
+from phoenix.datetime_utils import right_open_time_range
+from phoenix.trace import semantic_conventions
+from phoenix.trace.schemas import (
+    ATTRIBUTE_PREFIX,
+    COMPUTED_PREFIX,
+    CONTEXT_PREFIX,
+    Span,
+    SpanAttributes,
+    SpanID,
+    TraceID,
+)
+from phoenix.trace.v1 import decode, encode
+
+NAME = "name"
+STATUS_CODE = "status_code"
+SPAN_KIND = "span_kind"
+TRACE_ID = CONTEXT_PREFIX + "trace_id"
+SPAN_ID = CONTEXT_PREFIX + "span_id"
+PARENT_ID = "parent_id"
+START_TIME = "start_time"
+END_TIME = "end_time"
+LLM_TOKEN_COUNT_TOTAL = ATTRIBUTE_PREFIX + semantic_conventions.LLM_TOKEN_COUNT_TOTAL
+LLM_TOKEN_COUNT_PROMPT = ATTRIBUTE_PREFIX + semantic_conventions.LLM_TOKEN_COUNT_PROMPT
+LLM_TOKEN_COUNT_COMPLETION = ATTRIBUTE_PREFIX + semantic_conventions.LLM_TOKEN_COUNT_COMPLETION
+LATENCY_MS = COMPUTED_PREFIX + "latency_ms"  # The latency (or duration) of the span in milliseconds
+CUMULATIVE_LLM_TOKEN_COUNT_TOTAL = COMPUTED_PREFIX + "cumulative_token_count_total"
+CUMULATIVE_LLM_TOKEN_COUNT_PROMPT = COMPUTED_PREFIX + "cumulative_token_count_prompt"
+CUMULATIVE_LLM_TOKEN_COUNT_COMPLETION = COMPUTED_PREFIX + "cumulative_token_count_completion"
+
+
+class ReadableSpan(ObjectProxy):  # type: ignore
+    """
+    A wrapped a protobuf Span, with access methods and ability to decode to
+    a python span. It's meant to be interface layer separating use from
+    implementation. It can also provide computed values that are not intrinsic
+    to the span, e.g. the latency rank percent which can change as more spans
+    are ingested, and would need to be re-computed on the fly.
+    """
+
+    __wrapped__: pb.Span
+
+    def __init__(self, span: pb.Span) -> None:
+        super().__init__(span)
+        self._self_computed_values: Dict[str, SupportsFloat] = {}
+
+    @property
+    def span(self) -> Span:
+        span = decode(self.__wrapped__)
+        span.attributes.update(cast(SpanAttributes, self._self_computed_values))
+        # TODO: compute latency rank percent (which can change depending on how
+        # many spans already ingested).
+        return span
+
+    def __getitem__(self, key: str) -> Any:
+        if key.startswith(COMPUTED_PREFIX):
+            return self._self_computed_values.get(key)
+        if key.startswith(CONTEXT_PREFIX):
+            suffix_key = key[len(CONTEXT_PREFIX) :]
+            return getattr(self.__wrapped__.context, suffix_key, None)
+        if key.startswith(ATTRIBUTE_PREFIX):
+            suffix_key = key[len(ATTRIBUTE_PREFIX) :]
+            if suffix_key not in self.__wrapped__.attributes:
+                return None
+            return self.__wrapped__.attributes[suffix_key]
+        return getattr(self.__wrapped__, key, None)
+
+    def __setitem__(self, key: str, value: Any) -> None:
+        if not key.startswith(COMPUTED_PREFIX):
+            raise KeyError(f"{key} is not a computed value")
+        self._self_computed_values[key] = value
+
+
+ParentSpanID: TypeAlias = SpanID
+ChildSpanID: TypeAlias = SpanID
+
+
+class Traces:
+    def __init__(self, spans: Optional[Iterable[Span]] = None) -> None:
+        self._queue: "SimpleQueue[Optional[pb.Span]]" = SimpleQueue()
+        # Putting `None` as the sentinel value for queue termination.
+        weakref.finalize(self, self._queue.put, None)
+        for span in spans or ():
+            self.put(span)
+        self._lock = RLock()
+        self._spans: Dict[SpanID, ReadableSpan] = {}
+        self._parent_span_ids: Dict[SpanID, ParentSpanID] = {}
+        self._traces: Dict[TraceID, List[SpanID]] = defaultdict(list)
+        self._child_span_ids: DefaultDict[SpanID, List[ChildSpanID]] = defaultdict(list)
+        self._orphan_spans: DefaultDict[ParentSpanID, List[pb.Span]] = defaultdict(list)
+        self._start_time_sorted_span_ids: SortedKeyList[SpanID] = SortedKeyList(
+            key=lambda span_id: self._spans[span_id].start_time.ToDatetime(timezone.utc),
+        )
+        self._start_time_sorted_root_span_ids: SortedKeyList[SpanID] = SortedKeyList(
+            key=lambda span_id: self._spans[span_id].start_time.ToDatetime(timezone.utc),
+        )
+        self._latency_sorted_root_span_ids: SortedKeyList[SpanID] = SortedKeyList(
+            key=lambda span_id: self._spans[span_id][LATENCY_MS],
+        )
+        self._min_start_time: Optional[datetime] = None
+        self._max_start_time: Optional[datetime] = None
+        self._start_consumer()
+
+    def put(self, span: Optional[Union[Span, pb.Span]] = None) -> None:
+        self._queue.put(encode(span) if isinstance(span, Span) else span)
+
+    def get_trace(self, trace_id: TraceID) -> Iterator[Span]:
+        for span_id in self._traces[trace_id]:
+            if span := self[span_id]:
+                yield span
+
+    def get_spans(
+        self,
+        start_time: Optional[datetime] = None,
+        stop_time: Optional[datetime] = None,
+        root_spans_only: Optional[bool] = False,
+    ) -> Iterator[Span]:
+        if not self._spans:
+            return
+        min_start_time, max_stop_time = cast(
+            Tuple[datetime, datetime],
+            self.right_open_time_range,
+        )
+        start_time = start_time or min_start_time
+        stop_time = stop_time or max_stop_time
+        sorted_span_ids = (
+            self._start_time_sorted_root_span_ids
+            if root_spans_only
+            else self._start_time_sorted_span_ids
+        )
+        for span_id in sorted_span_ids.irange_key(
+            start_time.astimezone(timezone.utc),
+            stop_time.astimezone(timezone.utc),
+            inclusive=(True, False),
+            reverse=True,  # most recent spans first
+        ):
+            if span := self[span_id]:
+                yield span
+
+    def latency_rank_percent(self, latency_ms: float) -> Optional[float]:
+        """
+        Returns a value between 0 and 100 approximating the rank of the
+        latency value as percent of the total count of root spans. E.g., for
+        a latency value at the 75th percentile, the result is roughly 75.
+        """
+        root_span_ids = self._latency_sorted_root_span_ids
+        if not (n := len(root_span_ids)):
+            return None
+        rank = cast(int, root_span_ids.bisect_key_left(latency_ms))
+        return rank / n * 100
+
+    def get_descendant_span_ids(self, span_id: SpanID) -> Iterator[SpanID]:
+        for child_span_id in self._child_span_ids.get(span_id) or ():
+            yield child_span_id
+            yield from self.get_descendant_span_ids(child_span_id)
+
+    @property
+    def span_count(self) -> int:
+        """Total number of spans (excluding orphan spans if any)"""
+        return len(self._spans)
+
+    @property
+    def right_open_time_range(self) -> Tuple[Optional[datetime], Optional[datetime]]:
+        return right_open_time_range(self._min_start_time, self._max_start_time)
+
+    def __getitem__(self, span_id: SpanID) -> Optional[Span]:
+        with self._lock:
+            if span := self._spans.get(span_id):
+                return span.span
+        return None
+
+    def _start_consumer(self) -> None:
+        Thread(
+            target=MethodType(
+                self.__class__._consume_spans,
+                weakref.proxy(self),
+            ),
+            daemon=True,
+        ).start()
+
+    def _consume_spans(self) -> None:
+        while True:
+            if not (span := self._queue.get()):
+                return
+            with self._lock:
+                self._process_span(span)
+
+    def _process_span(self, span: pb.Span) -> None:
+        span_id = UUID(bytes=span.context.span_id)
+        existing_span = self._spans.get(span_id)
+        if existing_span and existing_span.end_time:
+            # Reject updates if span has ended.
+            return
+        is_root_span = not span.HasField("parent_span_id")
+        if not is_root_span:
+            parent_span_id = UUID(bytes=span.parent_span_id.value)
+            if parent_span_id not in self._spans:
+                # Span can't be processed before its parent.
+                self._orphan_spans[parent_span_id].append(span)
+                return
+            self._child_span_ids[parent_span_id].append(span_id)
+            self._parent_span_ids[span_id] = parent_span_id
+        new_span = ReadableSpan(span)
+        start_time = span.start_time.ToDatetime(timezone.utc)
+        end_time = span.end_time.ToDatetime(timezone.utc) if span.HasField("end_time") else None
+        if end_time:
+            new_span[LATENCY_MS] = (end_time - start_time).total_seconds() * 1000
+        self._spans[span_id] = new_span
+        if is_root_span and end_time:
+            self._latency_sorted_root_span_ids.add(span_id)
+        if not existing_span:
+            trace_id = UUID(bytes=span.context.trace_id)
+            self._traces[trace_id].append(span_id)
+            if is_root_span:
+                self._start_time_sorted_root_span_ids.add(span_id)
+            self._start_time_sorted_span_ids.add(span_id)
+            self._min_start_time = (
+                start_time
+                if self._min_start_time is None
+                else min(self._min_start_time, start_time)
+            )
+            self._max_start_time = (
+                start_time
+                if self._max_start_time is None
+                else max(self._max_start_time, start_time)
+            )
+        # Update cumulative values for span's ancestors.
+        for attribute_name, cumulative_attribute_name in (
+            (LLM_TOKEN_COUNT_TOTAL, CUMULATIVE_LLM_TOKEN_COUNT_TOTAL),
+            (LLM_TOKEN_COUNT_PROMPT, CUMULATIVE_LLM_TOKEN_COUNT_PROMPT),
+            (LLM_TOKEN_COUNT_COMPLETION, CUMULATIVE_LLM_TOKEN_COUNT_COMPLETION),
+        ):
+            existing_value = (existing_span[attribute_name] or 0) if existing_span else 0
+            new_value = new_span[attribute_name] or 0
+            if not (difference := new_value - existing_value):
+                continue
+            existing_cumulative_value = (
+                (existing_span[cumulative_attribute_name] or 0) if existing_span else 0
+            )
+            new_span[cumulative_attribute_name] = difference + existing_cumulative_value
+            self._add_value_to_span_ancestors(
+                span_id,
+                cumulative_attribute_name,
+                difference,
+            )
+        # Process previously orphaned spans, if any.
+        for orphan_span in self._orphan_spans[span_id]:
+            self._process_span(orphan_span)
+
+    def _add_value_to_span_ancestors(
+        self,
+        span_id: SpanID,
+        attribute_name: str,
+        value: float,
+    ) -> None:
+        while parent_span_id := self._parent_span_ids.get(span_id):
+            parent_span = self._spans[parent_span_id]
+            cumulative_value = parent_span[attribute_name] or 0
+            parent_span[attribute_name] = cumulative_value + value
+            span_id = parent_span_id

phoenix/datasets/dataset.py

@@ -9,17 +9,14 @@ from typing import Any, Dict, List, Optional, Set, Tuple, Union
 
 import numpy as np
 import pandas as pd
-import pytz
-from pandas import DataFrame, Series, Timestamp, read_parquet, to_datetime
+from pandas import DataFrame, Series, Timestamp, read_parquet
 from pandas.api.types import (
-    is_datetime64_any_dtype,
-    is_datetime64tz_dtype,
     is_numeric_dtype,
-    is_object_dtype,
 )
 from typing_extensions import TypeAlias
 
 from phoenix.config import DATASET_DIR, GENERATED_DATASET_NAME_PREFIX
+from phoenix.datetime_utils import normalize_timestamps
 
 from . import errors as err
 from .schema import (
@@ -70,6 +67,7 @@ class Dataset:
     _data_file_name: str = "data.parquet"
     _schema_file_name: str = "schema.json"
     _is_persisted: bool = False
+    _is_empty: bool = False
 
     def __init__(
         self,
@@ -85,8 +83,6 @@ class Dataset:
             schema=schema,
         )
         if errors:
-            for e in errors:
-                logger.error(e)
             raise err.DatasetError(errors)
         dataframe, schema = _parse_dataframe_and_schema(dataframe, schema)
         dataframe, schema = _normalize_timestamps(
@@ -98,6 +94,7 @@ class Dataset:
         self.__name: str = (
             name if name is not None else f"{GENERATED_DATASET_NAME_PREFIX}{str(uuid.uuid4())}"
         )
+        self._is_empty = self.dataframe.empty
         logger.info(f"""Dataset: {self.__name} initialized""")
 
     def __repr__(self) -> str:
@@ -531,7 +528,7 @@ def _create_and_normalize_dataframe_and_schema(
         if column_name_to_include.get(str(column_name), False):
             included_column_names.append(str(column_name))
     parsed_dataframe = dataframe[included_column_names].copy()
-    parsed_schema = replace(schema, excluded_column_names=None, **schema_patch)
+    parsed_schema = replace(schema, excluded_column_names=None, **schema_patch)  # type: ignore
     pred_id_col_name = parsed_schema.prediction_id_column_name
     if pred_id_col_name is None:
         parsed_schema = replace(parsed_schema, prediction_id_column_name="prediction_id")
@@ -594,8 +591,9 @@ def _normalize_timestamps(
     """
     Ensures that the dataframe has a timestamp column and the schema has a timestamp field. If the
     input dataframe contains a Unix or datetime timestamp or ISO8601 timestamp strings column, it
-    is converted to UTC timestamps. If the input dataframe and schema do not contain timestamps,
-    the default timestamp is used.
+    is converted to UTC timezone-aware timestamp. If the input dataframe and schema do not contain
+    timestamps, the default timestamp is used. If a timestamp is timezone-naive, it is localized
+    as per local timezone and then converted to UTC
     """
     timestamp_column: Series[Timestamp]
     if (timestamp_column_name := schema.timestamp_column_name) is None:
@@ -606,18 +604,9 @@ def _normalize_timestamps(
             if len(dataframe)
             else Series([default_timestamp]).iloc[:0].set_axis(dataframe.index, axis=0)
         )
-    elif is_numeric_dtype(timestamp_column_dtype := dataframe[timestamp_column_name].dtype):
-        timestamp_column = to_datetime(dataframe[timestamp_column_name], unit="s", utc=True)
-    elif is_datetime64tz_dtype(timestamp_column_dtype):
-        timestamp_column = dataframe[timestamp_column_name].dt.tz_convert(pytz.utc)
-    elif is_datetime64_any_dtype(timestamp_column_dtype):
-        timestamp_column = dataframe[timestamp_column_name].dt.tz_localize(pytz.utc)
-    elif is_object_dtype(timestamp_column_dtype):
-        timestamp_column = to_datetime(dataframe[timestamp_column_name], utc=True)
     else:
-        raise ValueError(
-            "When provided, input timestamp column must have numeric or datetime dtype, "
-            f"but found {timestamp_column_dtype} instead."
+        timestamp_column = normalize_timestamps(
+            dataframe[timestamp_column_name],
         )
     dataframe[timestamp_column_name] = timestamp_column
     return dataframe, schema
@@ -731,3 +720,7 @@ def _parse_open_inference_column_name(column_name: str) -> _OpenInferenceColumnN
             name=extract.get("name", ""),
         )
     raise ValueError(f"Invalid format for column name: {column_name}")
+
+
+# A dataset with no data. Useful for stubs
+EMPTY_DATASET = Dataset(pd.DataFrame(), schema=Schema())

phoenix/datasets/errors.py

@@ -58,7 +58,10 @@ class DatasetError(Exception):
     """An error raised when the dataset is invalid or incomplete"""
 
     def __init__(self, errors: Union[ValidationError, List[ValidationError]]):
-        self.errors = errors
+        self.errors: List[ValidationError] = errors if isinstance(errors, list) else [errors]
+
+    def __str__(self) -> str:
+        return "\n".join(map(str, self.errors))
 
 
 class InvalidColumnType(ValidationError):

phoenix/datasets/schema.py

@@ -105,7 +105,7 @@ class Schema:
             object.__setattr__(self, "prediction_id_column_name", self.id_column_name)
             object.__setattr__(self, "id_column_name", None)
 
-    def replace(self, **changes: str) -> "Schema":
+    def replace(self, **changes: Any) -> "Schema":
         return replace(self, **changes)
 
     def asdict(self) -> Dict[str, str]:

phoenix/datetime_utils.py

@@ -0,0 +1,87 @@
+from datetime import datetime, timedelta, timezone
+from typing import Any, Optional, Tuple, cast
+
+import pandas as pd
+import pytz
+from pandas import Timestamp, to_datetime
+from pandas.core.dtypes.common import (
+    is_datetime64_any_dtype,
+    is_datetime64tz_dtype,
+    is_numeric_dtype,
+    is_object_dtype,
+)
+
+
+def normalize_timestamps(
+    timestamps: "pd.Series[Any]",
+) -> "pd.Series[Timestamp]":
+    """
+    If the input timestamps contains a Unix or datetime timestamp or ISO8601
+    timestamp strings column, it is converted to UTC timezone-aware timestamp.
+    If a timestamp is timezone-naive, it is localized as per local timezone
+    and then converted to UTC.
+    """
+    if is_numeric_dtype(timestamps):
+        return to_datetime(timestamps, unit="s", utc=True)
+    if is_datetime64tz_dtype(timestamps):
+        return timestamps.dt.tz_convert(pytz.utc)
+    if is_datetime64_any_dtype(timestamps):
+        return timestamps.dt.tz_localize(
+            datetime.now().astimezone().tzinfo,
+        ).dt.tz_convert(
+            timezone.utc,
+        )
+    if is_object_dtype(timestamps):
+        timestamps = to_datetime(timestamps)
+        if timestamps.dt.tz is None:
+            timestamps = timestamps.dt.tz_localize(
+                datetime.now().astimezone().tzinfo,
+            )
+        return timestamps.dt.tz_convert(
+            timezone.utc,
+        )
+    raise ValueError(
+        "When provided, input timestamp column must have numeric or datetime dtype, "
+        f"but found {timestamps.dtype} instead."
+    )
+
+
+def floor_to_minute(dt: datetime) -> datetime:
+    """Floor datetime to the minute by taking a round-trip through string
+    format because there isn't always an available function to strip the
+    nanoseconds if present."""
+    try:
+        dt_as_string = dt.astimezone(
+            timezone.utc,
+        ).strftime(
+            MINUTE_DATETIME_FORMAT,
+        )
+    except ValueError:
+        # NOTE: as of Python 3.8.16, pandas 1.5.3:
+        # >>> isinstance(pd.NaT, datetime.datetime)
+        # True
+        return cast(datetime, pd.NaT)
+    return datetime.strptime(
+        dt_as_string,
+        MINUTE_DATETIME_FORMAT,
+    ).astimezone(
+        timezone.utc,
+    )
+
+
+MINUTE_DATETIME_FORMAT = "%Y-%m-%dT%H:%M:00%z"
+
+
+def right_open_time_range(
+    min_time: Optional[datetime],
+    max_time: Optional[datetime],
+) -> Tuple[Optional[datetime], Optional[datetime]]:
+    """
+    First adds one minute to `max_time`, because time intervals are right
+    open and one minute is the smallest interval allowed, then rounds down
+    the times to the nearest minute.
+    """
+    return (
+        floor_to_minute(min_time) if min_time else None,
+        floor_to_minute(max_time + timedelta(minutes=1)) if max_time else None,
+    )