arize-phoenix 3.25.0__py3-none-any.whl → 4.0.0__py3-none-any.whl

phoenix/trace/fixtures.py

@@ -1,11 +1,15 @@
-from dataclasses import dataclass, field
-from typing import Iterable, Iterator, List, NamedTuple, Optional, Tuple, cast
+from binascii import hexlify
+from dataclasses import dataclass, field, replace
+from datetime import datetime, timezone
+from random import getrandbits
+from typing import Dict, Iterable, Iterator, List, NamedTuple, Optional, Tuple, cast
 from urllib import request
 
 import pandas as pd
 from google.protobuf.wrappers_pb2 import DoubleValue, StringValue
 
 import phoenix.trace.v1 as pb
+from phoenix.trace.schemas import Span
 from phoenix.trace.trace_dataset import TraceDataset
 from phoenix.trace.utils import json_lines_to_df
 
@@ -105,7 +109,7 @@ TRACES_FIXTURES: List[TracesFixture] = [
 NAME_TO_TRACES_FIXTURE = {fixture.name: fixture for fixture in TRACES_FIXTURES}
 
 
-def _get_trace_fixture_by_name(fixture_name: str) -> TracesFixture:
+def get_trace_fixture_by_name(fixture_name: str) -> TracesFixture:
     """
     Returns the fixture whose name matches the input name.
 
@@ -120,7 +124,7 @@ def _get_trace_fixture_by_name(fixture_name: str) -> TracesFixture:
     return NAME_TO_TRACES_FIXTURE[fixture_name]
 
 
-def _download_traces_fixture(
+def download_traces_fixture(
     fixture: TracesFixture,
     host: Optional[str] = "https://storage.googleapis.com/",
     bucket: Optional[str] = "arize-assets",
@@ -138,12 +142,12 @@ def load_example_traces(use_case: str) -> TraceDataset:
     """
     Loads a trace dataframe by name.
     """
-    fixture = _get_trace_fixture_by_name(use_case)
-    return TraceDataset(json_lines_to_df(_download_traces_fixture(fixture)))
+    fixture = get_trace_fixture_by_name(use_case)
+    return TraceDataset(json_lines_to_df(download_traces_fixture(fixture)))
 
 
 def get_evals_from_fixture(use_case: str) -> Iterator[pb.Evaluation]:
-    fixture = _get_trace_fixture_by_name(use_case)
+    fixture = get_trace_fixture_by_name(use_case)
     for eval_fixture in fixture.evaluation_fixtures:
         yield from _read_eval_fixture(eval_fixture)
 
@@ -195,3 +199,61 @@ def _url(
     prefix: Optional[str] = "phoenix/traces/",
 ) -> str:
     return f"{host}{bucket}/{prefix}{file_name}"
+
+
+def reset_fixture_span_ids_and_timestamps(
+    spans: Iterable[Span],
+    evals: Iterable[pb.Evaluation] = (),
+) -> Tuple[List[Span], List[pb.Evaluation]]:
+    old_spans, old_evals = list(spans), list(evals)
+    new_trace_ids: Dict[str, str] = {}
+    new_span_ids: Dict[str, str] = {}
+    for old_span in old_spans:
+        new_trace_ids[old_span.context.trace_id] = _new_trace_id()
+        new_span_ids[old_span.context.span_id] = _new_span_id()
+        if old_span.parent_id:
+            new_span_ids[old_span.parent_id] = _new_span_id()
+    for old_eval in old_evals:
+        subject_id = old_eval.subject_id
+        if trace_id := subject_id.trace_id:
+            new_trace_ids[trace_id] = _new_trace_id()
+        elif span_id := subject_id.span_id:
+            new_span_ids[span_id] = _new_span_id()
+        elif span_id := subject_id.document_retrieval_id.span_id:
+            new_span_ids[span_id] = _new_span_id()
+    max_end_time = max(old_span.end_time for old_span in old_spans)
+    time_diff = datetime.now(timezone.utc) - max_end_time
+    new_spans: List[Span] = []
+    new_evals: List[pb.Evaluation] = []
+    for old_span in old_spans:
+        new_trace_id = new_trace_ids[old_span.context.trace_id]
+        new_span_id = new_span_ids[old_span.context.span_id]
+        new_parent_id = new_span_ids[old_span.parent_id] if old_span.parent_id else None
+        new_span = replace(
+            old_span,
+            context=replace(old_span.context, trace_id=new_trace_id, span_id=new_span_id),
+            parent_id=new_parent_id,
+            start_time=old_span.start_time + time_diff,
+            end_time=old_span.end_time + time_diff,
+        )
+        new_spans.append(new_span)
+    for old_eval in old_evals:
+        new_eval = pb.Evaluation()
+        new_eval.CopyFrom(old_eval)
+        subject_id = new_eval.subject_id
+        if trace_id := subject_id.trace_id:
+            subject_id.trace_id = new_trace_ids[trace_id]
+        elif span_id := subject_id.span_id:
+            subject_id.span_id = new_span_ids[span_id]
+        elif span_id := subject_id.document_retrieval_id.span_id:
+            subject_id.document_retrieval_id.span_id = new_span_ids[span_id]
+        new_evals.append(new_eval)
+    return new_spans, new_evals
+
+
+def _new_trace_id() -> str:
+    return hexlify(getrandbits(128).to_bytes(16, "big")).decode()
+
+
+def _new_span_id() -> str:
+    return hexlify(getrandbits(64).to_bytes(8, "big")).decode()

phoenix/trace/otel.py

@@ -1,39 +1,40 @@
-import inspect
 import json
 from binascii import hexlify, unhexlify
 from datetime import datetime, timezone
 from types import MappingProxyType
 from typing import (
     Any,
-    DefaultDict,
     Dict,
     Iterable,
     Iterator,
-    List,
     Mapping,
     Optional,
     Sequence,
-    Set,
     SupportsFloat,
     Tuple,
-    Union,
     cast,
 )
 
 import numpy as np
 import opentelemetry.proto.trace.v1.trace_pb2 as otlp
-from openinference.semconv import trace
 from openinference.semconv.trace import DocumentAttributes, SpanAttributes
 from opentelemetry.proto.common.v1.common_pb2 import AnyValue, ArrayValue, KeyValue
 from opentelemetry.util.types import Attributes, AttributeValue
 from typing_extensions import TypeAlias, assert_never
 
+from phoenix.trace.attributes import (
+    JSON_STRING_ATTRIBUTES,
+    flatten,
+    get_attribute_value,
+    has_mapping,
+    load_json_strings,
+    unflatten,
+)
 from phoenix.trace.schemas import (
     EXCEPTION_ESCAPED,
     EXCEPTION_MESSAGE,
     EXCEPTION_STACKTRACE,
     EXCEPTION_TYPE,
-    MimeType,
     Span,
     SpanContext,
     SpanEvent,
@@ -61,16 +62,10 @@ def decode_otlp_span(otlp_span: otlp.Span) -> Span:
     parent_id = _decode_identifier(otlp_span.parent_span_id)
 
     start_time = _decode_unix_nano(otlp_span.start_time_unix_nano)
-    end_time = (
-        _decode_unix_nano(otlp_span.end_time_unix_nano) if otlp_span.end_time_unix_nano else None
-    )
+    end_time = _decode_unix_nano(otlp_span.end_time_unix_nano)
 
-    attributes = dict(_unflatten(_load_json_strings(_decode_key_values(otlp_span.attributes))))
-    span_kind = SpanKind(attributes.pop(OPENINFERENCE_SPAN_KIND, None))
-
-    for mime_type in (INPUT_MIME_TYPE, OUTPUT_MIME_TYPE):
-        if mime_type in attributes:
-            attributes[mime_type] = MimeType(attributes[mime_type])
+    attributes = unflatten(load_json_strings(_decode_key_values(otlp_span.attributes)))
+    span_kind = SpanKind(get_attribute_value(attributes, OPENINFERENCE_SPAN_KIND))
 
     status_code, status_message = _decode_status(otlp_span.status)
     events = [_decode_event(event) for event in otlp_span.events]
@@ -152,28 +147,6 @@ def _decode_value(any_value: AnyValue) -> Any:
     assert_never(which)
 
 
-_JSON_STRING_ATTRIBUTES = (
-    DOCUMENT_METADATA,
-    LLM_PROMPT_TEMPLATE_VARIABLES,
-    METADATA,
-    TOOL_PARAMETERS,
-)
-
-
-def _load_json_strings(key_values: Iterable[Tuple[str, Any]]) -> Iterator[Tuple[str, Any]]:
-    for key, value in key_values:
-        if key.endswith(_JSON_STRING_ATTRIBUTES):
-            try:
-                dict_value = json.loads(value)
-            except Exception:
-                yield key, value
-            else:
-                if dict_value:
-                    yield key, dict_value
-        else:
-            yield key, value
-
-
 StatusMessage: TypeAlias = str
 
 _STATUS_DECODING = MappingProxyType(
@@ -190,120 +163,6 @@ def _decode_status(otlp_status: otlp.Status) -> Tuple[SpanStatusCode, StatusMess
     return status_code, otlp_status.message
 
 
-_SEMANTIC_CONVENTIONS: List[str] = sorted(
-    (
-        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()
-    ),
-    reverse=True,
-)  # sorted so the longer strings go first
-
-
-def _semantic_convention_prefix_partition(key: str, separator: str = ".") -> Tuple[str, str, str]:
-    """Return the longest prefix of `key` that is a semantic convention, and the remaining suffix
-    separated by `.`. For example, if `key` is "retrieval.documents.2.document.score", return
-    ("retrieval.documents", ".", "2.document.score"). The return signature is based on Python's
-    `.partition` method for strings.
-    """
-    for prefix in _SEMANTIC_CONVENTIONS:
-        if key == prefix:
-            return key, "", ""
-        if key.startswith(prefix) and key[len(prefix) :].startswith(separator):
-            return prefix, separator, key[len(prefix) + len(separator) :]
-    return "", "", ""
-
-
-class _Trie(DefaultDict[Union[str, int], "_Trie"]):
-    """Prefix Tree with special handling for indices (i.e. all-digit keys)."""
-
-    def __init__(self) -> None:
-        super().__init__(_Trie)
-        self.value: Any = None
-        self.indices: Set[int] = set()
-        self.branches: Set[Union[str, int]] = set()
-
-    def set_value(self, value: Any) -> None:
-        self.value = value
-        # value and indices must not coexist
-        self.branches.update(self.indices)
-        self.indices.clear()
-
-    def add_index(self, index: int) -> "_Trie":
-        if self.value is not None:
-            self.branches.add(index)
-        elif index not in self.branches:
-            self.indices.add(index)
-        return self[index]
-
-    def add_branch(self, branch: Union[str, int]) -> "_Trie":
-        if branch in self.indices:
-            self.indices.discard(cast(int, branch))
-        self.branches.add(branch)
-        return self[branch]
-
-
-# FIXME: Ideally we should not need something so complicated as a Trie, but it's useful here
-# for backward compatibility reasons regarding some deeply nested objects such as TOOL_PARAMETERS.
-# In the future, we should `json_dumps` them and not let things get too deeply nested.
-def _build_trie(
-    key_value_pairs: Iterable[Tuple[str, Any]],
-    separator: str = ".",
-) -> _Trie:
-    """Build a Trie (a.k.a. prefix tree) from `key_value_pairs`, by partitioning the keys by
-    separator. Each partition is a branch in the Trie. Special handling is done for partitions
-    that are all digits, e.g. "0", "12", etc., which are converted to integers and collected
-    as indices.
-    """
-    trie = _Trie()
-    for key, value in key_value_pairs:
-        if value is None:
-            continue
-        t = trie
-        while True:
-            prefix, _, suffix = _semantic_convention_prefix_partition(key, separator)
-            if prefix:
-                t = t.add_branch(prefix)
-            else:
-                prefix, _, suffix = key.partition(separator)
-                if prefix.isdigit():
-                    index = int(prefix)
-                    t = t.add_index(index) if suffix else t.add_branch(index)
-                else:
-                    t = t.add_branch(prefix)
-            if not suffix:
-                break
-            key = suffix
-        t.set_value(value)
-    return trie
-
-
-def _walk(trie: _Trie, prefix: str = "") -> Iterator[Tuple[str, Any]]:
-    if trie.value is not None:
-        yield prefix, trie.value
-    elif prefix and trie.indices:
-        yield prefix, [dict(_walk(trie[index])) for index in sorted(trie.indices)]
-    elif trie.indices:
-        for index in trie.indices:
-            yield from _walk(trie[index], prefix=f"{index}")
-    elif prefix:
-        yield prefix, dict(_walk(trie))
-        return
-    for branch in trie.branches:
-        new_prefix = f"{prefix}.{branch}" if prefix else f"{branch}"
-        yield from _walk(trie[branch], new_prefix)
-
-
-def _unflatten(
-    key_value_pairs: Iterable[Tuple[str, Any]],
-    separator: str = ".",
-) -> Iterator[Tuple[str, Any]]:
-    trie = _build_trie(key_value_pairs, separator)
-    yield from _walk(trie)
-
-
 _BILLION = 1_000_000_000  # for converting seconds to nanoseconds
 
 
@@ -316,11 +175,7 @@ def encode_span_to_otlp(span: Span) -> otlp.Span:
     start_time_unix_nano: int = int(span.start_time.timestamp() * _BILLION)
     end_time_unix_nano: int = int(span.end_time.timestamp() * _BILLION) if span.end_time else 0
 
-    attributes: Dict[str, Any] = span.attributes.copy()
-
-    for mime_type in (INPUT_MIME_TYPE, OUTPUT_MIME_TYPE):
-        if mime_type in attributes:
-            attributes[mime_type] = attributes[mime_type].value
+    attributes: Dict[str, Any] = dict(span.attributes)
 
     for key, value in span.attributes.items():
         if value is None:
@@ -328,19 +183,34 @@ def encode_span_to_otlp(span: Span) -> otlp.Span:
             attributes.pop(key, None)
         elif isinstance(value, Mapping):
             attributes.pop(key, None)
-            if key.endswith(_JSON_STRING_ATTRIBUTES):
+            if key.endswith(JSON_STRING_ATTRIBUTES):
                 attributes[key] = json.dumps(value)
             else:
-                attributes.update(_flatten_mapping(value, key))
+                attributes.update(
+                    flatten(
+                        value,
+                        prefix=key,
+                        recurse_on_sequence=True,
+                        json_string_attributes=JSON_STRING_ATTRIBUTES,
+                    )
+                )
         elif (
             not isinstance(value, str)
             and (isinstance(value, Sequence) or isinstance(value, np.ndarray))
-            and _has_mapping(value)
+            and has_mapping(value)
         ):
             attributes.pop(key, None)
-            attributes.update(_flatten_sequence(value, key))
-
-    attributes[OPENINFERENCE_SPAN_KIND] = span.span_kind.value
+            attributes.update(
+                flatten(
+                    value,
+                    prefix=key,
+                    recurse_on_sequence=True,
+                    json_string_attributes=JSON_STRING_ATTRIBUTES,
+                )
+            )
+
+    if OPENINFERENCE_SPAN_KIND not in attributes:
+        attributes[OPENINFERENCE_SPAN_KIND] = span.span_kind.value
 
     status = _encode_status(span.status_code, span.status_message)
     events = map(_encode_event, span.events)
@@ -381,42 +251,6 @@ def _encode_identifier(identifier: Optional[str]) -> bytes:
     return unhexlify(identifier)
 
 
-def _has_mapping(sequence: Sequence[Any]) -> bool:
-    for item in sequence:
-        if isinstance(item, Mapping):
-            return True
-    return False
-
-
-def _flatten_mapping(
-    mapping: Mapping[str, Any],
-    prefix: str,
-) -> Iterator[Tuple[str, Any]]:
-    for key, value in mapping.items():
-        prefixed_key = f"{prefix}.{key}"
-        if isinstance(value, Mapping):
-            if key.endswith(_JSON_STRING_ATTRIBUTES):
-                yield prefixed_key, json.dumps(value)
-            else:
-                yield from _flatten_mapping(value, prefixed_key)
-        elif isinstance(value, Sequence):
-            yield from _flatten_sequence(value, prefixed_key)
-        elif value is not None:
-            yield prefixed_key, value
-
-
-def _flatten_sequence(
-    sequence: Sequence[Any],
-    prefix: str,
-) -> Iterator[Tuple[str, Any]]:
-    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, f"{prefix}.{idx}")
-
-
 def _encode_event(event: SpanEvent) -> otlp.Span.Event:
     return otlp.Span.Event(
         name=event.name,

phoenix/trace/schemas.py

@@ -1,7 +1,7 @@
 from dataclasses import dataclass
 from datetime import datetime
 from enum import Enum
-from typing import Any, Dict, List, Optional, Union
+from typing import Any, List, Mapping, NamedTuple, Optional
 from uuid import UUID
 
 EXCEPTION_TYPE = "exception.type"
@@ -47,16 +47,14 @@ class SpanKind(Enum):
 
     @classmethod
     def _missing_(cls, v: Any) -> Optional["SpanKind"]:
-        if v and isinstance(v, str) and not v.isupper():
+        if v and isinstance(v, str) and v.isascii() and not v.isupper():
             return cls(v.upper())
-        return None if v else cls.UNKNOWN
+        return cls.UNKNOWN
 
 
 TraceID = str
 SpanID = str
-AttributePrimitiveValue = Union[str, bool, float, int]
-AttributeValue = Union[AttributePrimitiveValue, List[AttributePrimitiveValue]]
-SpanAttributes = Dict[str, AttributeValue]
+SpanAttributes = Mapping[str, Any]
 
 
 @dataclass(frozen=True)
@@ -73,7 +71,7 @@ class SpanConversationAttributes:
 
 
 @dataclass(frozen=True)
-class SpanEvent(Dict[str, Any]):
+class SpanEvent:
     """
     A Span Event can be thought of as a structured log message (or annotation)
     on a Span, typically used to denote a meaningful, singular point in time
@@ -142,7 +140,7 @@ class Span:
     "If the parent_id is None, this is the root span"
     parent_id: Optional[SpanID]
     start_time: datetime
-    end_time: Optional[datetime]
+    end_time: datetime
     status_code: SpanStatusCode
     status_message: str
     """
@@ -202,3 +200,11 @@ class ComputedAttributes(Enum):
     CUMULATIVE_LLM_TOKEN_COUNT_COMPLETION = "cumulative_token_count.completion"
     ERROR_COUNT = "error_count"
     CUMULATIVE_ERROR_COUNT = "cumulative_error_count"
+
+
+class ComputedValues(NamedTuple):
+    latency_ms: float
+    cumulative_error_count: int
+    cumulative_llm_token_count_prompt: int
+    cumulative_llm_token_count_completion: int
+    cumulative_llm_token_count_total: int

phoenix/trace/span_evaluations.py

@@ -12,6 +12,7 @@ from pandas.api.types import is_integer_dtype, is_numeric_dtype, is_string_dtype
 from pyarrow import RecordBatchStreamReader, Schema, Table, parquet
 
 from phoenix.config import TRACE_DATASET_DIR
+from phoenix.exceptions import PhoenixEvaluationNameIsMissing
 from phoenix.trace.errors import InvalidParquetMetadataError
 
 EVAL_NAME_COLUMN_PREFIX = "eval."
@@ -335,8 +336,10 @@ def _parse_schema_metadata(schema: Schema) -> Tuple[UUID, str, Type[Evaluations]
         arize_metadata = json.loads(metadata[b"arize"])
         eval_classes = {subclass.__name__: subclass for subclass in Evaluations.__subclasses__()}
         eval_id = UUID(arize_metadata["eval_id"])
-        if not isinstance((eval_name := arize_metadata["eval_name"]), str):
-            raise ValueError('Arize metadata must contain a string value for key "eval_name"')
+        if not isinstance((eval_name := arize_metadata["eval_name"]), str) or not eval_name.strip():
+            raise PhoenixEvaluationNameIsMissing(
+                'Arize metadata must contain a non-empty string value for key "eval_name"'
+            )
         evaluations_cls = eval_classes[arize_metadata["eval_type"]]
         return eval_id, eval_name, evaluations_cls
     except Exception as err:

phoenix/utilities/__init__.py

@@ -1,26 +0,0 @@
-from datetime import datetime
-from typing import List, Optional
-
-import pandas as pd
-
-from phoenix.core.project import Project
-from phoenix.trace.dsl import SpanQuery
-
-
-def query_spans(
-    project: Optional[Project],
-    *queries: SpanQuery,
-    start_time: Optional[datetime] = None,
-    stop_time: Optional[datetime] = None,
-    root_spans_only: Optional[bool] = None,
-) -> List[pd.DataFrame]:
-    if not queries or not project:
-        return []
-    spans = tuple(
-        project.get_spans(
-            start_time=start_time,
-            stop_time=stop_time,
-            root_spans_only=root_spans_only,
-        )
-    )
-    return [query(spans) for query in queries]

phoenix/utilities/span_store.py

@@ -1,23 +0,0 @@
-from typing import Optional
-
-from phoenix.config import get_env_span_storage_type, get_storage_dir
-from phoenix.core.traces import Traces
-from phoenix.storage.span_store import SPAN_STORE_FACTORIES, SpanStore
-from phoenix.trace.otel import decode_otlp_span
-from phoenix.utilities.project import get_project_name
-
-
-def get_span_store() -> Optional[SpanStore]:
-    if span_store_type := get_env_span_storage_type():
-        span_store_factory = SPAN_STORE_FACTORIES[span_store_type]
-        return span_store_factory(get_storage_dir())
-    return None
-
-
-def load_traces_data_from_store(traces: Traces, span_store: SpanStore) -> None:
-    for traces_data in span_store.load():
-        for resource_spans in traces_data.resource_spans:
-            project_name = get_project_name(resource_spans.resource.attributes)
-            for scope_span in resource_spans.scope_spans:
-                for span in scope_span.spans:
-                    traces.put(decode_otlp_span(span), project_name=project_name)

phoenix/version.py

@@ -1 +1 @@
-__version__ = "3.25.0"
+__version__ = "4.0.0"