arize-phoenix 1.6.0__py3-none-any.whl → 1.7.1__py3-none-any.whl

phoenix/trace/openai/instrumentor.py

@@ -6,6 +6,7 @@ from inspect import BoundArguments, signature
 from types import TracebackType
 from typing import (
     Any,
+    AsyncIterator,
     Callable,
     ContextManager,
     Coroutine,
@@ -22,7 +23,7 @@ from typing import (
 )
 
 import openai
-from openai import Stream
+from openai import AsyncStream, Stream
 from openai.types.chat import ChatCompletion, ChatCompletionChunk
 from typing_extensions import ParamSpec
 from wrapt import ObjectProxy
@@ -128,13 +129,14 @@ class ChatCompletionContext(ContextManager["ChatCompletionContext"]):
         """
         self.tracer = tracer
         self.start_time: Optional[datetime] = None
-        self.end_time: Optional[datetime] = None
         self.status_code = SpanStatusCode.UNSET
         self.status_message = ""
         self.events: List[SpanEvent] = []
         self.attributes: SpanAttributes = dict()
         parameters = _parameters(bound_arguments)
         self.num_choices = parameters.get("n", 1)
+        self.chat_completion_chunks: List[ChatCompletionChunk] = []
+        self.stream_complete = False
         self._process_parameters(parameters)
         self._span_created = False
 
@@ -150,14 +152,13 @@ class ChatCompletionContext(ContextManager["ChatCompletionContext"]):
     ) -> None:
         if exc_value is None:
             return
-        self.end_time = datetime.now(tz=timezone.utc)
         self.status_code = SpanStatusCode.ERROR
         status_message = str(exc_value)
         self.status_message = status_message
         self.events.append(
             SpanException(
                 message=status_message,
-                timestamp=self.end_time,
+                timestamp=datetime.now(tz=timezone.utc),
                 exception_type=type(exc_value).__name__,
                 exception_stacktrace=get_stacktrace(exc_value),
             )
@@ -171,17 +172,17 @@ class ChatCompletionContext(ContextManager["ChatCompletionContext"]):
         Args:
             response (ChatCompletion): The chat completion object.
         """
-        self.end_time = datetime.now(tz=timezone.utc)
-        self.status_code = SpanStatusCode.OK
         if isinstance(response, ChatCompletion):
             self._process_chat_completion(response)
         elif isinstance(response, Stream):
-            self.end_time = None  # set end time to None to indicate that the stream is still open
             return StreamWrapper(stream=response, context=self)
+        elif isinstance(response, AsyncStream):
+            return AsyncStreamWrapper(stream=response, context=self)
         elif hasattr(response, "parse") and callable(
             response.parse
         ):  # handle raw response by converting them to chat completions
             self._process_chat_completion(response.parse())
+        self.status_code = SpanStatusCode.OK
         self.create_span()
         return response
 
@@ -195,7 +196,7 @@ class ChatCompletionContext(ContextManager["ChatCompletionContext"]):
             name="OpenAI Chat Completion",
             span_kind=SpanKind.LLM,
             start_time=cast(datetime, self.start_time),
-            end_time=self.end_time,
+            end_time=datetime.now(tz=timezone.utc),
             status_code=self.status_code,
             status_message=self.status_message,
             attributes=self.attributes,
@@ -237,9 +238,8 @@ class ChatCompletionContext(ContextManager["ChatCompletionContext"]):
 
 class StreamWrapper(ObjectProxy):  # type: ignore
     """
-    A wrapper for streams of chat completion chunks that records each span
-    stream event and updates the span upon completion of the stream or upon an
-    exception.
+    A wrapper for streams of chat completion chunks that records events and
+    creates the span upon completion of the stream or upon an exception.
     """
 
     def __init__(self, stream: Stream[ChatCompletionChunk], context: ChatCompletionContext) -> None:
@@ -252,8 +252,7 @@ class StreamWrapper(ObjectProxy):  # type: ignore
             fields and attributes.
         """
         super().__init__(stream)
-        self._self_context = context
-        self._self_chunks: List[ChatCompletionChunk] = []
+        self._self_context = ChatCompletionStreamEventContext(context)
 
     def __next__(self) -> ChatCompletionChunk:
         """
@@ -263,48 +262,10 @@ class StreamWrapper(ObjectProxy):  # type: ignore
         Returns:
             ChatCompletionChunk: The forwarded chat completion chunk.
         """
-        finished_streaming = False
-        try:
+        with self._self_context as context:
             chat_completion_chunk = next(self.__wrapped__)
-            if not self._self_chunks:
-                self._self_context.events.append(
-                    SpanEvent(
-                        name="First Token Stream Event",
-                        timestamp=datetime.now(tz=timezone.utc),
-                        attributes={},
-                    )
-                )
-            self._self_chunks.append(chat_completion_chunk)
+            context.process_chat_completion_chunk(chat_completion_chunk)
             return cast(ChatCompletionChunk, chat_completion_chunk)
-        except StopIteration:
-            finished_streaming = True
-            raise
-        except Exception as error:
-            finished_streaming = True
-            status_message = str(error)
-            self._self_context.status_code = SpanStatusCode.ERROR
-            self._self_context.status_message = status_message
-            self._self_context.events.append(
-                SpanException(
-                    message=status_message,
-                    timestamp=datetime.now(tz=timezone.utc),
-                    exception_type=type(error).__name__,
-                    exception_stacktrace=get_stacktrace(error),
-                )
-            )
-            raise
-        finally:
-            if finished_streaming:
-                self._self_context.end_time = datetime.now(tz=timezone.utc)
-                self._self_context.attributes = {
-                    **self._self_context.attributes,
-                    LLM_OUTPUT_MESSAGES: _accumulate_messages(
-                        chunks=self._self_chunks, num_choices=self._self_context.num_choices
-                    ),  # type: ignore
-                    OUTPUT_VALUE: json.dumps([chunk.dict() for chunk in self._self_chunks]),
-                    OUTPUT_MIME_TYPE: MimeType.JSON,  # type: ignore
-                }
-                self._self_context.create_span()
 
     def __iter__(self) -> Iterator[ChatCompletionChunk]:
         """
@@ -314,6 +275,120 @@ class StreamWrapper(ObjectProxy):  # type: ignore
         return self
 
 
+class AsyncStreamWrapper(ObjectProxy):  # type: ignore
+    """
+    A wrapper for asynchronous streams of chat completion chunks that records
+    events and creates the span upon completion of the stream or upon an
+    exception.
+    """
+
+    def __init__(
+        self, stream: AsyncStream[ChatCompletionChunk], context: ChatCompletionContext
+    ) -> None:
+        """Initializes the stream wrapper.
+
+        Args:
+            stream (AsyncStream[ChatCompletionChunk]): The stream to wrap.
+
+            context (ChatCompletionContext): The context used to store span
+            fields and attributes.
+        """
+        super().__init__(stream)
+        self._self_context = ChatCompletionStreamEventContext(context)
+
+    async def __anext__(self) -> ChatCompletionChunk:
+        """
+        A wrapped __anext__ method that records span stream events and updates
+        the span upon completion of the stream or upon exception.
+
+        Returns:
+            ChatCompletionChunk: The forwarded chat completion chunk.
+        """
+        with self._self_context as context:
+            chat_completion_chunk = await self.__wrapped__.__anext__()
+            context.process_chat_completion_chunk(chat_completion_chunk)
+            return cast(ChatCompletionChunk, chat_completion_chunk)
+
+    def __aiter__(self) -> AsyncIterator[ChatCompletionChunk]:
+        """
+        An __aiter__ method that bypasses the wrapped class' __aiter__ method so
+        that __aiter__ is automatically instrumented using __anext__.
+        """
+        return self
+
+
+class ChatCompletionStreamEventContext(ContextManager["ChatCompletionStreamEventContext"]):
+    """
+    A context manager that surrounds stream events in a stream of chat
+    completions and processes each chat completion chunk.
+    """
+
+    def __init__(self, chat_completion_context: ChatCompletionContext) -> None:
+        """Initializes the context manager.
+
+        Args:
+            chat_completion_context (ChatCompletionContext): The chat completion
+            context storing span fields and attributes.
+        """
+        self._context = chat_completion_context
+
+    def __exit__(
+        self,
+        exc_type: Optional[Type[BaseException]],
+        exc_value: Optional[BaseException],
+        traceback: Optional[TracebackType],
+    ) -> None:
+        if isinstance(exc_value, StopIteration) or isinstance(exc_value, StopAsyncIteration):
+            self._context.stream_complete = True
+            self._context.status_code = SpanStatusCode.OK
+        elif exc_value is not None:
+            self._context.stream_complete = True
+            self._context.status_code = SpanStatusCode.ERROR
+            status_message = str(exc_value)
+            self._context.status_message = status_message
+            self._context.events.append(
+                SpanException(
+                    message=status_message,
+                    timestamp=datetime.now(tz=timezone.utc),
+                    exception_type=type(exc_value).__name__,
+                    exception_stacktrace=get_stacktrace(exc_value),
+                )
+            )
+        if not self._context.stream_complete:
+            return
+        self._context.attributes = {
+            **self._context.attributes,
+            LLM_OUTPUT_MESSAGES: _accumulate_messages(
+                chunks=self._context.chat_completion_chunks,
+                num_choices=self._context.num_choices,
+            ),  # type: ignore
+            OUTPUT_VALUE: json.dumps(
+                [chunk.dict() for chunk in self._context.chat_completion_chunks]
+            ),
+            OUTPUT_MIME_TYPE: MimeType.JSON,  # type: ignore
+        }
+        self._context.create_span()
+
+    def process_chat_completion_chunk(self, chat_completion_chunk: ChatCompletionChunk) -> None:
+        """
+        Processes a chat completion chunk and adds relevant information to the
+        context.
+
+        Args:
+            chat_completion_chunk (ChatCompletionChunk): The chat completion
+            chunk to be processed.
+        """
+        if not self._context.chat_completion_chunks:
+            self._context.events.append(
+                SpanEvent(
+                    name="First Token Stream Event",
+                    timestamp=datetime.now(tz=timezone.utc),
+                    attributes={},
+                )
+            )
+        self._context.chat_completion_chunks.append(chat_completion_chunk)
+
+
 def _wrapped_openai_sync_client_request_function(
     request_fn: Callable[ParameterSpec, GenericType], tracer: Tracer
 ) -> Callable[ParameterSpec, GenericType]:
@@ -363,22 +438,11 @@ def _wrapped_openai_async_client_request_function(
 
     async def wrapped(*args: Any, **kwargs: Any) -> Any:
         bound_arguments = call_signature.bind(*args, **kwargs)
-        if (
-            _is_streaming_request(bound_arguments)
-            or _request_type(bound_arguments) is not RequestType.CHAT_COMPLETION
-        ):
+        if _request_type(bound_arguments) is not RequestType.CHAT_COMPLETION:
             return await request_fn(*args, **kwargs)
         with ChatCompletionContext(bound_arguments, tracer) as context:
             response = await request_fn(*args, **kwargs)
-            context.process_response(
-                cast(
-                    ChatCompletion,
-                    response.parse()
-                    if hasattr(response, "parse") and callable(response.parse)
-                    else response,
-                )
-            )
-            return response
+            return context.process_response(response)
 
     return wrapped
 
@@ -521,19 +585,6 @@ _CHAT_COMPLETION_ATTRIBUTE_FUNCTIONS: Dict[str, Callable[[ChatCompletion], Any]]
 }
 
 
-def _is_streaming_request(bound_arguments: BoundArguments) -> bool:
-    """
-    Determines whether the request is a streaming request.
-
-    Args:
-        bound_arguments (BoundArguments): The bound arguments to the request function.
-
-    Returns:
-        bool: True if the request is a streaming request, False otherwise.
-    """
-    return cast(bool, bound_arguments.arguments["stream"])
-
-
 def _parameters(bound_arguments: BoundArguments) -> Parameters:
     """
     The parameters for the LLM call, e.g., temperature.

phoenix/trace/span_evaluations.py

@@ -1,12 +1,163 @@
-import pandas as pd
+from abc import ABC
+from dataclasses import dataclass, field
+from itertools import product
+from types import MappingProxyType
+from typing import Any, Callable, List, Mapping, Optional, Sequence, Set, Tuple
 
-EVALUATIONS_INDEX_NAME = "context.span_id"
-RESULTS_COLUMN_NAMES = ["score", "label", "explanation"]
+import pandas as pd
+from pandas.api.types import is_integer_dtype, is_numeric_dtype, is_string_dtype
 
 EVAL_NAME_COLUMN_PREFIX = "eval."
 
 
-class SpanEvaluations:
+class NeedsNamedIndex(ABC):
+    index_names: Mapping[Tuple[str, ...], Callable[[Any], bool]]
+    all_valid_index_name_sorted_combos: Set[Tuple[str, ...]]
+
+    @classmethod
+    def preferred_names(cls) -> List[str]:
+        return [choices[0] for choices in cls.index_names.keys()]
+
+    @classmethod
+    def aliases(cls) -> Mapping[str, str]:
+        return {alias: choices[0] for choices in cls.index_names.keys() for alias in choices[1:]}
+
+    @classmethod
+    def unalias(cls, name: str) -> str:
+        return cls.aliases().get(name, name)
+
+    @classmethod
+    def is_valid_index_names(cls, names: Sequence[str]) -> bool:
+        return (
+            len(names) == len(cls.index_names)
+            and tuple(sorted(names)) in cls.all_valid_index_name_sorted_combos
+        )
+
+    @classmethod
+    def find_valid_index_names(cls, dtypes: "pd.Series[Any]") -> Optional[List[str]]:
+        valid_names = []
+        for names, check_type in cls.index_names.items():
+            for name in names:
+                if name in dtypes.index and check_type(dtypes[name]):
+                    valid_names.append(name)
+                    break
+            else:
+                return None
+        return valid_names
+
+
+class NeedsResultColumns(ABC):
+    result_column_names: Mapping[str, Callable[[Any], bool]] = MappingProxyType(
+        {
+            "score": is_numeric_dtype,
+            "label": is_string_dtype,
+            "explanation": is_string_dtype,
+        }
+    )
+
+    @classmethod
+    def is_valid_result_columns(cls, dtypes: "pd.Series[Any]") -> bool:
+        names = cls.result_column_names.keys()
+        intersection = dtypes.index.intersection(names)  # type: ignore
+        if not len(intersection):
+            return False
+        for name in intersection:
+            check_type = cls.result_column_names[name]
+            if not check_type(dtypes[name]):
+                return False
+        return True
+
+
+@dataclass(frozen=True)
+class Evaluations(NeedsNamedIndex, NeedsResultColumns, ABC):
+    eval_name: str  # The name for the evaluation, e.g. 'toxicity'
+    dataframe: pd.DataFrame = field(repr=False)
+
+    def __len__(self) -> int:
+        return len(self.dataframe)
+
+    def __repr__(self) -> str:
+        return (
+            f"{self.__class__.__name__}(eval_name={self.eval_name!r}, "
+            f"dataframe=<rows: {len(self.dataframe)!r}>)"
+        )
+
+    def __dir__(self) -> List[str]:
+        return ["get_dataframe"]
+
+    def get_dataframe(self, prefix_columns_with_name: bool = True) -> pd.DataFrame:
+        """
+        Returns a copy of the dataframe with the evaluation annotations
+
+        Parameters
+        __________
+        prefix_columns_with_name: bool
+            if True, the columns will be prefixed with the eval_name, e.g. 'eval.toxicity.value'
+        """
+        if prefix_columns_with_name:
+            prefix = f"{EVAL_NAME_COLUMN_PREFIX}{self.eval_name}."
+            return self.dataframe.add_prefix(prefix)
+        return self.dataframe.copy(deep=False)
+
+    def __bool__(self) -> bool:
+        return not self.dataframe.empty
+
+    def __post_init__(self) -> None:
+        dataframe = (
+            pd.DataFrame() if self.dataframe.empty else self._clean_dataframe(self.dataframe)
+        )
+        object.__setattr__(self, "dataframe", dataframe)
+
+    def _clean_dataframe(self, dataframe: pd.DataFrame) -> pd.DataFrame:
+        # Ensure column names are strings.
+        column_names = dataframe.columns.astype(str)
+        dataframe = dataframe.set_axis(column_names, axis=1)
+
+        # If the dataframe contains the index columns, set the index to those columns
+        if not self.is_valid_index_names(dataframe.index.names) and (
+            index_names := self.find_valid_index_names(dataframe.dtypes)
+        ):
+            dataframe = dataframe.set_index(index_names)
+
+        # Validate that the dataframe is indexed correctly.
+        if not self.is_valid_index_names(dataframe.index.names):
+            raise ValueError(
+                f"The dataframe index must be {self.preferred_names()} but was "
+                f"'{dataframe.index.name or dataframe.index.names}'"
+            )
+
+        # Validate that the dataframe contains result columns of appropriate types.
+        if not self.is_valid_result_columns(dataframe.dtypes):
+            raise ValueError(
+                f"The dataframe must contain one of these columns with appropriate "
+                f"value types: {self.result_column_names.keys()} "
+            )
+
+        # Un-alias to the preferred names.
+        preferred_names = [self.unalias(name) for name in dataframe.index.names]
+        dataframe = dataframe.rename_axis(preferred_names, axis=0)
+
+        # Drop the unnecessary columns.
+        result_column_names = dataframe.columns.intersection(self.result_column_names.keys())  # type: ignore
+        return dataframe.loc[:, result_column_names]  # type: ignore
+
+    def __init_subclass__(
+        cls,
+        index_names: Mapping[Tuple[str, ...], Callable[[Any], bool]],
+        **kwargs: Any,
+    ) -> None:
+        super().__init_subclass__(**kwargs)
+        cls.index_names = index_names
+        cls.all_valid_index_name_sorted_combos = set(
+            tuple(sorted(prod)) for prod in product(*cls.index_names.keys())
+        )
+
+
+@dataclass(frozen=True)
+class SpanEvaluations(
+    Evaluations,
+    index_names=MappingProxyType({("context.span_id", "span_id"): is_string_dtype}),
+):
     """
     SpanEvaluations is a set of evaluation annotations for a set of spans.
     SpanEvaluations encompasses the evaluation annotations for a single evaluation task
@@ -18,7 +169,7 @@ class SpanEvaluations:
     Parameters
     __________
     eval_name: str
-        the name of the evaluation, e.x. 'toxicity'
+        the name of the evaluation, e.g. 'toxicity'
     dataframe: pandas.DataFrame
         the pandas dataframe containing the evaluation annotations Each row
         represents the evaluations on a span.
@@ -35,38 +186,52 @@ class SpanEvaluations:
     | span_3  | 1                  | toxic              | discrimination     |
     """
 
-    dataframe: pd.DataFrame
 
-    eval_name: str  # The name for the evaluation, e.x. 'toxicity'
+@dataclass(frozen=True)
+class DocumentEvaluations(
+    Evaluations,
+    index_names=MappingProxyType(
+        {
+            ("context.span_id", "span_id"): is_string_dtype,
+            ("document_position", "position"): is_integer_dtype,
+        }
+    ),
+):
+    """
+    DocumentEvaluations is a set of evaluation annotations for a set of documents.
+    DocumentEvaluations encompasses the evaluation annotations for a single evaluation task
+    such as relevance.
 
-    def __init__(self, eval_name: str, dataframe: pd.DataFrame):
-        self.eval_name = eval_name
+    Parameters
+    __________
+    eval_name: str
+        the name of the evaluation, e.g. 'relevance'
+    dataframe: pandas.DataFrame
+        the pandas dataframe containing the evaluation annotations. Each row
+        represents the evaluations on a document.
 
-        # If the dataframe contains the index column, set the index to that column
-        if EVALUATIONS_INDEX_NAME in dataframe.columns:
-            dataframe = dataframe.set_index(EVALUATIONS_INDEX_NAME)
+    Example
+    _______
 
-        # validate that the dataframe is indexed by context.span_id
-        if dataframe.index.name != EVALUATIONS_INDEX_NAME:
-            raise ValueError(
-                f"The dataframe index must be '{EVALUATIONS_INDEX_NAME}' but was "
-                f"'{dataframe.index.name}'"
-            )
+    DataFrame of document evaluations for relevance may look like:
 
-        # Drop the unnecessary columns
-        extra_column_names = dataframe.columns.difference(RESULTS_COLUMN_NAMES)
-        self.dataframe = dataframe.drop(extra_column_names, axis=1)
+    | span_id | position | score | label      | explanation  |
+    |---------|----------|-------|------------|--------------|
+    | span_1  | 0        | 1     | relevant   | it's apropos |
+    | span_1  | 1        | 1     | relevant   | it's germane |
+    | span_2  | 0        | 0     | irrelevant | it's rubbish |
+    """
 
-    def get_dataframe(self, prefix_columns_with_name: bool = True) -> pd.DataFrame:
-        """
-        Returns a copy of the dataframe with the evaluation annotations
+    def _clean_dataframe(self, dataframe: pd.DataFrame) -> pd.DataFrame:
+        dataframe = super()._clean_dataframe(dataframe)
+        if dataframe.index.names != self.preferred_names():
+            return dataframe.swaplevel()
+        return dataframe
 
-        Parameters
-        __________
-        prefix_columns_with_name: bool
-            if True, the columns will be prefixed with the eval_name, e.x. 'eval.toxicity.value'
-        """
-        if prefix_columns_with_name:
-            prefix = f"{EVAL_NAME_COLUMN_PREFIX}{self.eval_name}."
-            return self.dataframe.add_prefix(prefix)
-        return self.dataframe.copy()
+
+@dataclass(frozen=True)
+class TraceEvaluations(
+    Evaluations,
+    index_names=MappingProxyType({("context.trace_id", "trace_id"): is_string_dtype}),
+):
+    ...

phoenix/trace/trace_dataset.py

@@ -16,7 +16,7 @@ from .semantic_conventions import (
     RERANKER_OUTPUT_DOCUMENTS,
     RETRIEVAL_DOCUMENTS,
 )
-from .span_evaluations import EVALUATIONS_INDEX_NAME, SpanEvaluations
+from .span_evaluations import Evaluations, SpanEvaluations
 from .span_json_decoder import json_to_span
 from .span_json_encoder import span_to_json
 
@@ -93,14 +93,14 @@ class TraceDataset:
 
     name: str
     dataframe: pd.DataFrame
-    evaluations: List[SpanEvaluations] = []
+    evaluations: List[Evaluations] = []
     _data_file_name: str = "data.parquet"
 
     def __init__(
         self,
         dataframe: DataFrame,
         name: Optional[str] = None,
-        evaluations: Iterable[SpanEvaluations] = (),
+        evaluations: Iterable[Evaluations] = (),
     ):
         """
         Constructs a TraceDataset from a dataframe of spans. Optionally takes in
@@ -199,7 +199,7 @@ class TraceDataset:
             coerce_timestamps="ms",
         )
 
-    def append_evaluations(self, evaluations: SpanEvaluations) -> None:
+    def append_evaluations(self, evaluations: Evaluations) -> None:
         """adds an evaluation to the traces"""
         # Append the evaluations to the list of evaluations
         self.evaluations.append(evaluations)
@@ -209,7 +209,11 @@ class TraceDataset:
         Creates a flat dataframe of all the evaluations for the dataset.
         """
         return pd.concat(
-            [evals.get_dataframe(prefix_columns_with_name=True) for evals in self.evaluations],
+            [
+                evals.get_dataframe(prefix_columns_with_name=True)
+                for evals in self.evaluations
+                if isinstance(evals, SpanEvaluations)
+            ],
             axis=1,
         )
 
@@ -227,5 +231,5 @@ class TraceDataset:
             return self.dataframe.copy()
         evals_df = self.get_evals_dataframe()
         # Make sure the index is set to the span_id
-        df = self.dataframe.set_index(EVALUATIONS_INDEX_NAME, drop=False)
+        df = self.dataframe.set_index("context.span_id", drop=False)
         return pd.concat([df, evals_df], axis=1)

phoenix/trace/tracer.py

@@ -72,6 +72,7 @@ class Tracer:
         attributes: Optional[SpanAttributes] = None,
         events: Optional[List[SpanEvent]] = None,
         conversation: Optional[SpanConversationAttributes] = None,
+        span_id: Optional[UUID] = None,
     ) -> Span:
         """
         create_span creates a new span with the given name and options.
@@ -90,7 +91,7 @@ class Tracer:
 
         span = Span(
             name=name,
-            context=SpanContext(trace_id=trace_id, span_id=uuid4()),
+            context=SpanContext(trace_id=trace_id, span_id=span_id or uuid4()),
             span_kind=span_kind,
             parent_id=parent_id,
             start_time=start_time,