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

phoenix/trace/fixtures.py

@@ -1,9 +1,8 @@
 from dataclasses import dataclass
-from typing import List, Optional
+from typing import List, Optional, cast
 from urllib import request
 
-import pandas as pd
-
+from phoenix.trace.trace_dataset import TraceDataset
 from phoenix.trace.utils import json_lines_to_df
 
 
@@ -14,13 +13,29 @@ class TracesFixture:
     file_name: str
 
 
+llama_index_rag_fixture = TracesFixture(
+    name="llama_index_rag",
+    description="Traces from running the llama_index on a RAG use case.",
+    file_name="llama_index_rag_v5.jsonl",
+)
+
+langchain_rag_stuff_document_chain_fixture = TracesFixture(
+    name="langchain_rag_stuff_document_chain",
+    description="LangChain RAG data",
+    file_name="langchain_rag.jsonl",
+)
+
 random_fixture = TracesFixture(
     name="random",
     description="Randomly generated traces",
     file_name="random.jsonl",
 )
 
-TRACES_FIXTURES: List[TracesFixture] = [random_fixture]
+TRACES_FIXTURES: List[TracesFixture] = [
+    llama_index_rag_fixture,
+    langchain_rag_stuff_document_chain_fixture,
+    random_fixture,
+]
 
 NAME_TO_TRACES_FIXTURE = {fixture.name: fixture for fixture in TRACES_FIXTURES}
 
@@ -45,20 +60,20 @@ def _download_traces_fixture(
     host: Optional[str] = "https://storage.googleapis.com/",
     bucket: Optional[str] = "arize-assets",
     prefix: Optional[str] = "phoenix/traces/",
-) -> pd.DataFrame:
+) -> List[str]:
     """
     Downloads the traces fixture from the phoenix bucket.
     """
     url = f"{host}{bucket}/{prefix}{fixture.file_name}"
     with request.urlopen(url) as f:
-        return json_lines_to_df(f.readlines())
+        return cast(List[str], f.readlines())
 
 
-def load_example_traces(use_case: str) -> pd.DataFrame:
+def load_example_traces(use_case: str) -> TraceDataset:
     """
     Loads a trace dataframe by name.
 
     NB: this functionality is under active construction.
     """
     fixture = _get_trace_fixture_by_name(use_case)
-    return _download_traces_fixture(fixture)
+    return TraceDataset(json_lines_to_df(_download_traces_fixture(fixture)))

phoenix/trace/schemas.py

@@ -4,14 +4,29 @@ from enum import Enum
 from typing import Any, Dict, List, Optional, Union
 from uuid import UUID
 
+from phoenix.trace.semantic_conventions import (
+    EXCEPTION_ESCAPED,
+    EXCEPTION_MESSAGE,
+    EXCEPTION_STACKTRACE,
+    EXCEPTION_TYPE,
+)
+
 
 class SpanStatusCode(Enum):
     UNSET = "UNSET"
     OK = "OK"
     ERROR = "ERROR"
 
+    def __str__(self) -> str:
+        return self.value
+
+    @classmethod
+    def _missing_(cls, v: Any) -> Optional["SpanStatusCode"]:
+        if v and isinstance(v, str) and not v.isupper():
+            return cls(v.upper())
+        return None if v else cls.UNSET
+
 
-@dataclass(frozen=True)
 class SpanKind(Enum):
     """
     SpanKind is loosely inspired by OpenTelemetry's SpanKind
@@ -25,8 +40,20 @@ class SpanKind(Enum):
     LLM = "LLM"
     RETRIEVER = "RETRIEVER"
     EMBEDDING = "EMBEDDING"
+    AGENT = "AGENT"
+    UNKNOWN = "UNKNOWN"
+
+    def __str__(self) -> str:
+        return self.value
 
+    @classmethod
+    def _missing_(cls, v: Any) -> Optional["SpanKind"]:
+        if v and isinstance(v, str) and not v.isupper():
+            return cls(v.upper())
+        return None if v else cls.UNKNOWN
 
+
+TraceID = UUID
 SpanID = UUID
 AttributePrimitiveValue = Union[str, bool, float, int]
 AttributeValue = Union[AttributePrimitiveValue, List[AttributePrimitiveValue]]
@@ -37,7 +64,7 @@ SpanAttributes = Dict[str, AttributeValue]
 class SpanContext:
     """Context propagation for a span"""
 
-    trace_id: UUID
+    trace_id: TraceID
     span_id: SpanID
 
 
@@ -58,10 +85,11 @@ class SpanEvent(Dict[str, Any]):
     """
 
     name: str
-    message: str
     timestamp: datetime
+    attributes: SpanAttributes
 
 
+@dataclass(frozen=True)
 class SpanException(SpanEvent):
     """
     A Span Exception is a special type of Span Event that denotes an error
@@ -73,8 +101,28 @@ class SpanException(SpanEvent):
     https://github.com/open-telemetry/opentelemetry-specification/blob/main/specification/trace/semantic_conventions/exceptions.md
     """
 
-    def __init__(self, timestamp: datetime, message: str):
-        super().__init__(name="exception", message=message, timestamp=timestamp)
+    def __init__(
+        self,
+        timestamp: datetime,
+        message: str,
+        exception_type: Optional[str] = None,
+        exception_escaped: Optional[bool] = None,
+        exception_stacktrace: Optional[str] = None,
+    ):
+        super().__init__(
+            name="exception",
+            timestamp=timestamp,
+            attributes={
+                k: v
+                for k, v in {
+                    EXCEPTION_TYPE: exception_type,
+                    EXCEPTION_MESSAGE: message,
+                    EXCEPTION_ESCAPED: exception_escaped,
+                    EXCEPTION_STACKTRACE: exception_stacktrace,
+                }.items()
+                if v is not None
+            },
+        )
 
 
 @dataclass(frozen=True)
@@ -95,7 +143,7 @@ class Span:
     "If the parent_id is None, this is the root span"
     parent_id: Optional[SpanID]
     start_time: datetime
-    end_time: datetime
+    end_time: Optional[datetime]
     status_code: SpanStatusCode
     status_message: str
     """
@@ -131,3 +179,8 @@ class Span:
     conversation_id
     """
     conversation: Optional[SpanConversationAttributes]
+
+
+ATTRIBUTE_PREFIX = "attributes."
+CONTEXT_PREFIX = "context."
+COMPUTED_PREFIX = "__computed__."

phoenix/trace/semantic_conventions.py

@@ -5,7 +5,8 @@ Inspiration from OpenTelemetry:
 https://opentelemetry.io/docs/specs/otel/trace/semantic_conventions/span-general/
 """
 from dataclasses import dataclass
-from typing import Any, Dict
+from enum import Enum
+from typing import Any, Dict, Optional
 
 
 @dataclass(frozen=True)
@@ -35,3 +36,142 @@ class DeploymentAttributes(AttributeGroup):
             type="string",
         ),
     }
+
+
+EXCEPTION_TYPE = "exception.type"
+EXCEPTION_MESSAGE = "exception.message"
+EXCEPTION_ESCAPED = "exception.escaped"
+EXCEPTION_STACKTRACE = "exception.stacktrace"
+
+
+OUTPUT_VALUE = "output.value"
+OUTPUT_MIME_TYPE = "output.mime_type"
+"""
+The type of output.value. If unspecified, the type is plain text by default.
+If type is JSON, the value is a string representing a JSON object.
+"""
+INPUT_VALUE = "input.value"
+INPUT_MIME_TYPE = "input.mime_type"
+"""
+The type of input.value. If unspecified, the type is plain text by default.
+If type is JSON, the value is a string representing a JSON object.
+"""
+
+
+class MimeType(Enum):
+    TEXT = "text/plain"
+    JSON = "application/json"
+
+    @classmethod
+    def _missing_(cls, v: Any) -> Optional["MimeType"]:
+        return None if v else cls.TEXT
+
+
+EMBEDDING_EMBEDDINGS = "embedding.embeddings"
+"""
+A list of objects containing embedding data, including the vector and represented piece of text.
+"""
+EMBEDDING_MODEL_NAME = "embedding.model_name"
+"""
+The name of the embedding model.
+"""
+EMBEDDING_TEXT = "embedding.text"
+"""
+The text represented by the embedding.
+"""
+EMBEDDING_VECTOR = "embedding.vector"
+"""
+The embedding vector.
+"""
+
+MESSAGE_ROLE = "message.role"
+"""
+The role of the message, such as "user", "agent", "function".
+"""
+MESSAGE_NAME = "message.name"
+"""
+The name of the message, often used to identify the function
+that was used to generate the message.
+"""
+MESSAGE_FUNCTION_CALL_NAME = "message.function_call_name"
+"""
+The function name that is a part of the message list.
+This is populated for role 'function' or 'agent' as a mechanism to identify
+the function that was called during the execution of a tool
+"""
+MESSAGE_FUNCTION_CALL_ARGUMENTS_JSON = "message.function_call_arguments_json"
+"""
+The JSON string representing the arguments passed to the function
+during a function call
+"""
+MESSAGE_CONTENT = "message.content"
+"""
+The content of the message to the llm
+"""
+LLM_FUNCTION_CALL = "llm.function_call"
+"""
+For models and APIs that support function calling. Records attributes such as the function name and
+arguments to the called function.
+"""
+LLM_INVOCATION_PARAMETERS = "llm.invocation_parameters"
+"""
+Invocation parameters passed to the LLM or API, such as the model name, temperature, etc.
+"""
+LLM_MESSAGES = "llm.messages"
+"""
+Messages provided to a chat API.
+"""
+LLM_MODEL_NAME = "llm.model_name"
+"""
+The name of the model being used.
+"""
+LLM_PROMPT = "llm.prompt"
+"""
+Messages provided to a completions API.
+"""
+LLM_PROMPT_TEMPLATE = "llm.prompt_template.template"
+"""
+The prompt template as a Python f-string.
+"""
+LLM_PROMPT_TEMPLATE_VARIABLES = "llm.prompt_template.variables"
+"""
+A list of input variables to the prompt template.
+"""
+LLM_PROMPT_TEMPLATE_VERSION = "llm.prompt_template.version"
+"""
+The version of the prompt template being used.
+"""
+LLM_TOKEN_COUNT_PROMPT = "llm.token_count.prompt"
+"""
+Number of tokens in the prompt.
+"""
+LLM_TOKEN_COUNT_COMPLETION = "llm.token_count.completion"
+"""
+Number of tokens in the completion.
+"""
+LLM_TOKEN_COUNT_TOTAL = "llm.token_count.total"
+"""
+Total number of tokens, including both prompt and completion.
+"""
+
+TOOL_NAME = "tool.name"
+"""
+Name of the tool being used.
+"""
+TOOL_DESCRIPTION = "tool.description"
+"""
+Description of the tool's purpose, typically used to select the tool.
+"""
+TOOL_PARAMETERS = "tool.parameters"
+"""
+Parameters of the tool, e.g. see https://platform.openai.com/docs/guides/gpt/function-calling
+"""
+
+RETRIEVAL_DOCUMENTS = "retrieval.documents"
+DOCUMENT_ID = "document.id"
+DOCUMENT_SCORE = "document.score"
+DOCUMENT_CONTENT = "document.content"
+DOCUMENT_METADATA = "document.metadata"
+"""
+Document metadata as a string representing a JSON object
+"""

phoenix/trace/span_json_decoder.py

@@ -1,22 +1,52 @@
 import json
 from datetime import datetime
-from typing import Any, Dict
+from typing import Any, Dict, Optional
+from uuid import UUID
 
 from phoenix.trace.schemas import (
     Span,
     SpanContext,
     SpanConversationAttributes,
     SpanEvent,
+    SpanException,
     SpanKind,
     SpanStatusCode,
 )
+from phoenix.trace.semantic_conventions import (
+    DOCUMENT_METADATA,
+    EXCEPTION_MESSAGE,
+    INPUT_MIME_TYPE,
+    OUTPUT_MIME_TYPE,
+    RETRIEVAL_DOCUMENTS,
+    MimeType,
+)
+
+
+def json_to_document(obj: Optional[Dict[str, Any]]) -> Dict[str, Any]:
+    if obj is None:
+        return {}
+    if document_metadata := obj.get(DOCUMENT_METADATA):
+        obj[DOCUMENT_METADATA] = json.loads(document_metadata)
+    return obj
+
+
+def json_to_attributes(obj: Optional[Dict[str, Any]]) -> Dict[str, Any]:
+    if obj is None:
+        return {}
+    if not isinstance(obj, dict):
+        raise ValueError(f"attributes should be dict, but attributes={obj}")
+    if mime_type := obj.get(INPUT_MIME_TYPE):
+        obj[INPUT_MIME_TYPE] = MimeType(mime_type)
+    if mime_type := obj.get(OUTPUT_MIME_TYPE):
+        obj[OUTPUT_MIME_TYPE] = MimeType(mime_type)
+    if documents := obj.get(RETRIEVAL_DOCUMENTS):
+        obj[RETRIEVAL_DOCUMENTS] = [json_to_document(document) for document in documents]
+    return obj
 
 
 def json_to_span(data: Dict[str, Any]) -> Any:
     """
     A hook for json.loads to convert a dict to a Span object.
-
-    NB: this function is mainly used for testing purposes. Consider swapping this out for pydantic.
     """
     # Check if the dict can be interpreted as a Span
     if set(data.keys()) == {
@@ -32,12 +62,36 @@ def json_to_span(data: Dict[str, Any]) -> Any:
         "events",
         "conversation",
     }:
-        data["context"] = SpanContext(**data["context"])
+        context = data["context"]
+        if not isinstance(context, dict):
+            raise ValueError(f"context should be dict, but context={context}")
+        data["context"] = SpanContext(
+            trace_id=UUID(context["trace_id"]),
+            span_id=UUID(context["span_id"]),
+        )
+        parent_id = data.get("parent_id")
+        data["parent_id"] = UUID(parent_id) if parent_id else None
+        attributes = data.get("attributes")
+        data["attributes"] = json_to_attributes(attributes)
         data["start_time"] = datetime.fromisoformat(data["start_time"])
-        data["end_time"] = datetime.fromisoformat(data["end_time"])
+        data["end_time"] = (
+            datetime.fromisoformat(end_time) if (end_time := data.get("end_time")) else None
+        )
         data["span_kind"] = SpanKind(data["span_kind"])
         data["status_code"] = SpanStatusCode(data["status_code"])
-        data["events"] = [SpanEvent(**event) for event in data["events"]]  # Build SpanEvent objects
+        data["events"] = [
+            SpanException(
+                message=(event.get("attributes") or {}).get(EXCEPTION_MESSAGE) or "",
+                timestamp=datetime.fromisoformat(event["timestamp"]),
+            )
+            if event["name"] == "exception"
+            else SpanEvent(
+                name=event["name"],
+                attributes=event.get("attributes") or {},
+                timestamp=datetime.fromisoformat(event["timestamp"]),
+            )
+            for event in data["events"]
+        ]
         data["conversation"] = (
             SpanConversationAttributes(**data["conversation"])
             if data["conversation"] is not None

phoenix/trace/span_json_encoder.py

@@ -10,7 +10,6 @@ from .schemas import (
     SpanContext,
     SpanConversationAttributes,
     SpanEvent,
-    SpanException,
 )
 
 
@@ -27,14 +26,7 @@ class SpanJSONEncoder(json.JSONEncoder):
         elif isinstance(obj, SpanEvent):
             return {
                 "name": obj.name,
-                "message": obj.message,
-                "timestamp": obj.timestamp.isoformat(),
-            }
-        elif isinstance(obj, SpanException):
-            # TODO: add stacktrace etc.
-            return {
-                "name": obj.name,
-                "message": obj.message,
+                "attributes": obj.attributes,
                 "timestamp": obj.timestamp.isoformat(),
             }
         elif isinstance(obj, Span):

phoenix/trace/trace_dataset.py

@@ -1,10 +1,21 @@
+import json
+import uuid
+from datetime import datetime
+from typing import Iterator, List, Optional, cast
+
 import pandas as pd
-from pandas import DataFrame
+from pandas import DataFrame, read_parquet
+
+from phoenix.datetime_utils import normalize_timestamps
+
+from ..config import DATASET_DIR, GENERATED_DATASET_NAME_PREFIX
+from .schemas import ATTRIBUTE_PREFIX, CONTEXT_PREFIX, Span
+from .span_json_decoder import json_to_span
+from .span_json_encoder import span_to_json
 
 # A set of columns that is required
 REQUIRED_COLUMNS = [
     "name",
-    "message",
     "span_kind",
     "parent_id",
     "start_time",
@@ -16,6 +27,15 @@ REQUIRED_COLUMNS = [
 ]
 
 
+def normalize_dataframe(dataframe: DataFrame) -> "DataFrame":
+    """Makes the dataframe have appropriate data types"""
+
+    # Convert the start and end times to datetime
+    dataframe["start_time"] = normalize_timestamps(dataframe["start_time"])
+    dataframe["end_time"] = normalize_timestamps(dataframe["end_time"])
+    return dataframe
+
+
 class TraceDataset:
     """
     A TraceDataset is a wrapper around a dataframe which is a flattened representation
@@ -27,12 +47,84 @@ class TraceDataset:
         the pandas dataframe containing the tracing data. Each row represents a span.
     """
 
+    name: str
     dataframe: pd.DataFrame
+    _data_file_name: str = "data.parquet"
 
-    def __init__(self, dataframe: DataFrame):
+    def __init__(self, dataframe: DataFrame, name: Optional[str] = None):
         # Validate the the dataframe has required fields
-        columns = dataframe.columns.values
-        missing_columns = [column for column in REQUIRED_COLUMNS if column not in columns]
-        if missing_columns:
-            raise ValueError(f"The dataframe is missing some required columns: {missing_columns}")
-        self.dataframe = dataframe
+        if missing_columns := set(REQUIRED_COLUMNS) - set(dataframe.columns):
+            raise ValueError(
+                f"The dataframe is missing some required columns: {', '.join(missing_columns)}"
+            )
+        self.dataframe = normalize_dataframe(dataframe)
+        self.name = name or f"{GENERATED_DATASET_NAME_PREFIX}{str(uuid.uuid4())}"
+
+    @classmethod
+    def from_spans(cls, spans: List[Span]) -> "TraceDataset":
+        """Creates a TraceDataset from a list of spans.
+
+        Args:
+            spans (List[Span]): A list of spans.
+
+        Returns:
+            TraceDataset: A TraceDataset containing the spans.
+        """
+        return cls(pd.json_normalize(map(json.loads, map(span_to_json, spans))))  # type: ignore
+
+    def to_spans(self) -> Iterator[Span]:
+        for _, row in self.dataframe.iterrows():
+            is_attribute = row.index.str.startswith(ATTRIBUTE_PREFIX)
+            attribute_keys = row.index[is_attribute]
+            attributes = (
+                row.loc[is_attribute]
+                .rename(
+                    {key: key[len(ATTRIBUTE_PREFIX) :] for key in attribute_keys},
+                )
+                .dropna()
+                .to_dict()
+            )
+            is_context = row.index.str.startswith(CONTEXT_PREFIX)
+            context_keys = row.index[is_context]
+            context = (
+                row.loc[is_context]
+                .rename(
+                    {key: key[len(CONTEXT_PREFIX) :] for key in context_keys},
+                )
+                .to_dict()
+            )
+            end_time: Optional[datetime] = cast(datetime, row.get("end_time"))
+            if end_time is pd.NaT:
+                end_time = None
+            yield json_to_span(
+                {
+                    "name": row["name"],
+                    "context": context,
+                    "span_kind": row["span_kind"],
+                    "parent_id": row.get("parent_id"),
+                    "start_time": cast(datetime, row["start_time"]).isoformat(),
+                    "end_time": end_time.isoformat() if end_time else None,
+                    "status_code": row["status_code"],
+                    "status_message": row.get("status_message") or "",
+                    "attributes": attributes,
+                    "events": row.get("events") or [],
+                    "conversation": row.get("conversation"),
+                }
+            )
+
+    @classmethod
+    def from_name(cls, name: str) -> "TraceDataset":
+        """Retrieves a dataset by name from the file system"""
+        directory = DATASET_DIR / name
+        df = read_parquet(directory / cls._data_file_name)
+        return cls(df, name)
+
+    def to_disc(self) -> None:
+        """writes the data to disc"""
+        directory = DATASET_DIR / self.name
+        directory.mkdir(parents=True, exist_ok=True)
+        self.dataframe.to_parquet(
+            directory / self._data_file_name,
+            allow_truncated_timestamps=True,
+            coerce_timestamps="ms",
+        )

phoenix/trace/tracer.py

@@ -1,5 +1,6 @@
+import logging
 from datetime import datetime
-from typing import Callable, List, Optional
+from typing import Any, Callable, Iterator, List, Optional, Protocol
 from uuid import UUID, uuid4
 
 from .schemas import (
@@ -13,6 +14,14 @@ from .schemas import (
     SpanStatusCode,
 )
 
+logger = logging.getLogger(__name__)
+logger.addHandler(logging.NullHandler())
+
+
+class SpanExporter(Protocol):
+    def export(self, span: Span) -> None:
+        ...
+
 
 class Tracer:
     """
@@ -28,7 +37,10 @@ class Tracer:
 
     def __init__(
         self,
+        exporter: Optional[SpanExporter] = None,
         on_append: Optional[Callable[[List[Span]], None]] = None,
+        *args: Any,
+        **kwargs: Any,
     ):
         """
         Create a new Tracer. A Tracer's main purpose is to create spans.
@@ -42,14 +54,16 @@ class Tracer:
         """
         self.span_buffer = []
         self.on_append = on_append
+        self._exporter: Optional[SpanExporter] = exporter
+        super().__init__(*args, **kwargs)
 
     def create_span(
         self,
         name: str,
         span_kind: SpanKind,
         start_time: datetime,
-        end_time: datetime,
-        status_code: SpanStatusCode,
+        end_time: Optional[datetime] = None,
+        status_code: SpanStatusCode = SpanStatusCode.UNSET,
         status_message: Optional[str] = "",
         parent_id: Optional[SpanID] = None,
         trace_id: Optional[UUID] = None,
@@ -86,8 +100,17 @@ class Tracer:
             conversation=conversation,
         )
 
+        if self._exporter:
+            self._exporter.export(span)
         self.span_buffer.append(span)
 
         if self.on_append is not None:
             self.on_append(self.span_buffer)
         return span
+
+    def get_spans(self) -> Iterator[Span]:
+        """
+        Returns the spans stored in the tracer. This is useful if you are running
+        in a notebook environment and you want to inspect the spans.
+        """
+        yield from self.span_buffer