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

phoenix/server/api/routers/v1/traces.py

@@ -0,0 +1,82 @@
+import gzip
+import zlib
+
+from google.protobuf.message import DecodeError
+from opentelemetry.proto.collector.trace.v1.trace_service_pb2 import (
+    ExportTraceServiceRequest,
+)
+from starlette.background import BackgroundTask
+from starlette.concurrency import run_in_threadpool
+from starlette.datastructures import State
+from starlette.requests import Request
+from starlette.responses import Response
+from starlette.status import (
+    HTTP_403_FORBIDDEN,
+    HTTP_415_UNSUPPORTED_MEDIA_TYPE,
+    HTTP_422_UNPROCESSABLE_ENTITY,
+)
+
+from phoenix.trace.otel import decode_otlp_span
+from phoenix.utilities.project import get_project_name
+
+
+async def post_traces(request: Request) -> Response:
+    """
+    summary: Send traces to Phoenix
+    operationId: addTraces
+    tags:
+      - traces
+    requestBody:
+      required: true
+      content:
+        application/x-protobuf:
+          schema:
+            type: string
+            format: binary
+    responses:
+      200:
+        description: Success
+      403:
+        description: Forbidden
+      415:
+        description: Unsupported content type, only gzipped protobuf
+      422:
+        description: Request body is invalid
+    """
+    if request.app.state.read_only:
+        return Response(status_code=HTTP_403_FORBIDDEN)
+    content_type = request.headers.get("content-type")
+    if content_type != "application/x-protobuf":
+        return Response(
+            content=f"Unsupported content type: {content_type}",
+            status_code=HTTP_415_UNSUPPORTED_MEDIA_TYPE,
+        )
+    content_encoding = request.headers.get("content-encoding")
+    if content_encoding and content_encoding not in ("gzip", "deflate"):
+        return Response(
+            content=f"Unsupported content encoding: {content_encoding}",
+            status_code=HTTP_415_UNSUPPORTED_MEDIA_TYPE,
+        )
+    body = await request.body()
+    if content_encoding == "gzip":
+        body = await run_in_threadpool(gzip.decompress, body)
+    elif content_encoding == "deflate":
+        body = await run_in_threadpool(zlib.decompress, body)
+    req = ExportTraceServiceRequest()
+    try:
+        await run_in_threadpool(req.ParseFromString, body)
+    except DecodeError:
+        return Response(
+            content="Request body is invalid ExportTraceServiceRequest",
+            status_code=HTTP_422_UNPROCESSABLE_ENTITY,
+        )
+    return Response(background=BackgroundTask(_add_spans, req, request.state))
+
+
+async def _add_spans(req: ExportTraceServiceRequest, state: State) -> None:
+    for resource_spans in req.resource_spans:
+        project_name = get_project_name(resource_spans.resource.attributes)
+        for scope_span in resource_spans.scope_spans:
+            for otlp_span in scope_span.spans:
+                span = await run_in_threadpool(decode_otlp_span, otlp_span)
+                await state.queue_span_for_bulk_insert(span, project_name)

phoenix/server/api/schema.py

@@ -4,11 +4,17 @@ from typing import Dict, List, Optional, Set, Union
 import numpy as np
 import numpy.typing as npt
 import strawberry
+from sqlalchemy import delete, select
+from sqlalchemy.orm import contains_eager, load_only
 from strawberry import ID, UNSET
 from strawberry.types import Info
 from typing_extensions import Annotated
 
+from phoenix.config import DEFAULT_PROJECT_NAME
+from phoenix.db import models
+from phoenix.db.insertion.span import ClearProjectSpansEvent
 from phoenix.pointcloud.clustering import Hdbscan
+from phoenix.server.api.context import Context
 from phoenix.server.api.helpers import ensure_list
 from phoenix.server.api.input_types.ClusterInput import ClusterInput
 from phoenix.server.api.input_types.Coordinates import (
@@ -16,56 +22,70 @@ from phoenix.server.api.input_types.Coordinates import (
     InputCoordinate3D,
 )
 from phoenix.server.api.types.Cluster import Cluster, to_gql_clusters
-from phoenix.server.api.types.Project import Project
-
-from .context import Context
-from .types.DatasetRole import AncillaryDatasetRole, DatasetRole
-from .types.Dimension import to_gql_dimension
-from .types.EmbeddingDimension import (
+from phoenix.server.api.types.DatasetRole import AncillaryDatasetRole, DatasetRole
+from phoenix.server.api.types.Dimension import to_gql_dimension
+from phoenix.server.api.types.EmbeddingDimension import (
     DEFAULT_CLUSTER_SELECTION_EPSILON,
     DEFAULT_MIN_CLUSTER_SIZE,
     DEFAULT_MIN_SAMPLES,
     to_gql_embedding_dimension,
 )
-from .types.Event import create_event_id, unpack_event_id
-from .types.ExportEventsMutation import ExportEventsMutation
-from .types.Functionality import Functionality
-from .types.Model import Model
-from .types.node import GlobalID, Node, from_global_id, from_global_id_with_expected_type
-from .types.pagination import Connection, ConnectionArgs, Cursor, connection_from_list
+from phoenix.server.api.types.Event import create_event_id, unpack_event_id
+from phoenix.server.api.types.ExportEventsMutation import ExportEventsMutation
+from phoenix.server.api.types.Functionality import Functionality
+from phoenix.server.api.types.Model import Model
+from phoenix.server.api.types.node import (
+    GlobalID,
+    Node,
+    from_global_id,
+    from_global_id_with_expected_type,
+)
+from phoenix.server.api.types.pagination import (
+    Connection,
+    ConnectionArgs,
+    CursorString,
+    connection_from_list,
+)
+from phoenix.server.api.types.Project import Project
+from phoenix.server.api.types.Span import to_gql_span
+from phoenix.server.api.types.Trace import Trace
 
 
 @strawberry.type
 class Query:
     @strawberry.field
-    def projects(
+    async def projects(
         self,
         info: Info[Context, None],
         first: Optional[int] = 50,
         last: Optional[int] = UNSET,
-        after: Optional[Cursor] = UNSET,
-        before: Optional[Cursor] = UNSET,
+        after: Optional[CursorString] = UNSET,
+        before: Optional[CursorString] = UNSET,
     ) -> Connection[Project]:
         args = ConnectionArgs(
             first=first,
-            after=after if isinstance(after, Cursor) else None,
+            after=after if isinstance(after, CursorString) else None,
             last=last,
-            before=before if isinstance(before, Cursor) else None,
-        )
-        data = (
-            []
-            if (traces := info.context.traces) is None
-            else [
-                Project(id_attr=project_id, name=project_name, project=project)
-                for project_id, project_name, project in traces.get_projects()
-            ]
+            before=before if isinstance(before, CursorString) else None,
         )
+        async with info.context.db() as session:
+            projects = await session.scalars(select(models.Project))
+        data = [
+            Project(
+                id_attr=project.id,
+                name=project.name,
+                gradient_start_color=project.gradient_start_color,
+                gradient_end_color=project.gradient_end_color,
+            )
+            for project in projects
+        ]
         return connection_from_list(data=data, args=args)
 
     @strawberry.field
-    def functionality(self, info: Info[Context, None]) -> "Functionality":
+    async def functionality(self, info: Info[Context, None]) -> "Functionality":
         has_model_inferences = not info.context.model.is_empty
-        has_traces = info.context.traces is not None
+        async with info.context.db() as session:
+            has_traces = (await session.scalar(select(models.Trace).limit(1))) is not None
         return Functionality(
             model_inferences=has_model_inferences,
             tracing=has_traces,
@@ -76,7 +96,7 @@ class Query:
         return Model()
 
     @strawberry.field
-    def node(self, id: GlobalID, info: Info[Context, None]) -> Node:
+    async def node(self, id: GlobalID, info: Info[Context, None]) -> Node:
         type_name, node_id = from_global_id(str(id))
         if type_name == "Dimension":
             dimension = info.context.model.scalar_dimensions[node_id]
@@ -85,17 +105,42 @@ class Query:
             embedding_dimension = info.context.model.embedding_dimensions[node_id]
             return to_gql_embedding_dimension(node_id, embedding_dimension)
         elif type_name == "Project":
-            if (traces := info.context.traces) is not None:
-                projects = {
-                    project_id: (project_name, project)
-                    for project_id, project_name, project in traces.get_projects()
-                }
-                if node_id in projects:
-                    name, project = projects[node_id]
-                    return Project(id_attr=node_id, name=name, project=project)
-            raise Exception(f"Unknown project: {id}")
-
-        raise Exception(f"Unknown node type: {type}")
+            project_stmt = select(
+                models.Project.id,
+                models.Project.name,
+                models.Project.gradient_start_color,
+                models.Project.gradient_end_color,
+            ).where(models.Project.id == node_id)
+            async with info.context.db() as session:
+                project = (await session.execute(project_stmt)).first()
+            if project is None:
+                raise ValueError(f"Unknown project: {id}")
+            return Project(
+                id_attr=project.id,
+                name=project.name,
+                gradient_start_color=project.gradient_start_color,
+                gradient_end_color=project.gradient_end_color,
+            )
+        elif type_name == "Trace":
+            trace_stmt = select(models.Trace.id).where(models.Trace.id == node_id)
+            async with info.context.db() as session:
+                id_attr = await session.scalar(trace_stmt)
+            if id_attr is None:
+                raise ValueError(f"Unknown trace: {id}")
+            return Trace(id_attr=id_attr)
+        elif type_name == "Span":
+            span_stmt = (
+                select(models.Span)
+                .join(models.Trace)
+                .options(contains_eager(models.Span.trace))
+                .where(models.Span.id == node_id)
+            )
+            async with info.context.db() as session:
+                span = await session.scalar(span_stmt)
+            if span is None:
+                raise ValueError(f"Unknown span: {id}")
+            return to_gql_span(span)
+        raise Exception(f"Unknown node type: {type_name}")
 
     @strawberry.field
     def clusters(
@@ -229,18 +274,37 @@ class Query:
 @strawberry.type
 class Mutation(ExportEventsMutation):
     @strawberry.mutation
-    def delete_project(self, info: Info[Context, None], id: GlobalID) -> Query:
-        if (traces := info.context.traces) is not None:
-            node_id = from_global_id_with_expected_type(str(id), "Project")
-            traces.archive_project(node_id)
+    async def delete_project(self, info: Info[Context, None], id: GlobalID) -> Query:
+        node_id = from_global_id_with_expected_type(str(id), "Project")
+        async with info.context.db() as session:
+            project = await session.scalar(
+                select(models.Project)
+                .where(models.Project.id == node_id)
+                .options(load_only(models.Project.name))
+            )
+            if project is None:
+                raise ValueError(f"Unknown project: {id}")
+            if project.name == DEFAULT_PROJECT_NAME:
+                raise ValueError(f"Cannot delete the {DEFAULT_PROJECT_NAME} project")
+            await session.delete(project)
         return Query()
 
     @strawberry.mutation
-    def archive_project(self, info: Info[Context, None], id: GlobalID) -> Query:
-        if (traces := info.context.traces) is not None:
-            node_id = from_global_id_with_expected_type(str(id), "Project")
-            traces.archive_project(node_id)
+    async def clear_project(self, info: Info[Context, None], id: GlobalID) -> Query:
+        project_id = from_global_id_with_expected_type(str(id), "Project")
+        delete_statement = delete(models.Trace).where(models.Trace.project_rowid == project_id)
+        async with info.context.db() as session:
+            await session.execute(delete_statement)
+            if cache := info.context.cache_for_dataloaders:
+                cache.invalidate(ClearProjectSpansEvent(project_rowid=project_id))
         return Query()
 
 
-schema = strawberry.Schema(query=Query, mutation=Mutation)
+# This is the schema for generating `schema.graphql`.
+# See https://strawberry.rocks/docs/guides/schema-export
+# It should be kept in sync with the server's runtime-initialized
+# instance. To do so, search for the usage of `strawberry.Schema(...)`.
+schema = strawberry.Schema(
+    query=Query,
+    mutation=Mutation,
+)

phoenix/server/api/types/DocumentEvaluationSummary.py

@@ -80,7 +80,7 @@ class DocumentEvaluationSummary:
             return result
         values = self.metrics_collection.apply(lambda m: m.precision(k))
         result = (values.mean(), values.count())
-        self._cached_average_ndcg_results[k] = result
+        self._cached_average_precision_results[k] = result
         return result
 
     @cached_property

phoenix/server/api/types/Evaluation.py

@@ -3,7 +3,7 @@ from typing import Optional
 import strawberry
 
 import phoenix.trace.v1 as pb
-from phoenix.trace.schemas import SpanID, TraceID
+from phoenix.db.models import DocumentAnnotation, SpanAnnotation, TraceAnnotation
 
 
 @strawberry.interface
@@ -26,47 +26,56 @@ class Evaluation:
 
 @strawberry.type
 class TraceEvaluation(Evaluation):
-    trace_id: strawberry.Private[TraceID]
-
     @staticmethod
     def from_pb_evaluation(evaluation: pb.Evaluation) -> "TraceEvaluation":
         result = evaluation.result
         score = result.score.value if result.HasField("score") else None
         label = result.label.value if result.HasField("label") else None
         explanation = result.explanation.value if result.HasField("explanation") else None
-        trace_id = TraceID(evaluation.subject_id.trace_id)
         return TraceEvaluation(
             name=evaluation.name,
             score=score,
             label=label,
             explanation=explanation,
-            trace_id=trace_id,
+        )
+
+    @staticmethod
+    def from_sql_trace_annotation(annotation: TraceAnnotation) -> "TraceEvaluation":
+        return TraceEvaluation(
+            name=annotation.name,
+            score=annotation.score,
+            label=annotation.label,
+            explanation=annotation.explanation,
         )
 
 
 @strawberry.type
 class SpanEvaluation(Evaluation):
-    span_id: strawberry.Private[SpanID]
-
     @staticmethod
     def from_pb_evaluation(evaluation: pb.Evaluation) -> "SpanEvaluation":
         result = evaluation.result
         score = result.score.value if result.HasField("score") else None
         label = result.label.value if result.HasField("label") else None
         explanation = result.explanation.value if result.HasField("explanation") else None
-        span_id = SpanID(evaluation.subject_id.span_id)
         return SpanEvaluation(
             name=evaluation.name,
             score=score,
             label=label,
             explanation=explanation,
-            span_id=span_id,
+        )
+
+    @staticmethod
+    def from_sql_span_annotation(annotation: SpanAnnotation) -> "SpanEvaluation":
+        return SpanEvaluation(
+            name=annotation.name,
+            score=annotation.score,
+            label=annotation.label,
+            explanation=annotation.explanation,
         )
 
 
 @strawberry.type
 class DocumentEvaluation(Evaluation):
-    span_id: strawberry.Private[SpanID]
     document_position: int = strawberry.field(
         description="The zero-based index among retrieved documents, which "
         "is collected as a list (even when ordering is not inherently meaningful)."
@@ -80,12 +89,20 @@ class DocumentEvaluation(Evaluation):
         explanation = result.explanation.value if result.HasField("explanation") else None
         document_retrieval_id = evaluation.subject_id.document_retrieval_id
         document_position = document_retrieval_id.document_position
-        span_id = SpanID(document_retrieval_id.span_id)
         return DocumentEvaluation(
             name=evaluation.name,
             score=score,
             label=label,
             explanation=explanation,
             document_position=document_position,
-            span_id=span_id,
+        )
+
+    @staticmethod
+    def from_sql_document_annotation(annotation: DocumentAnnotation) -> "DocumentEvaluation":
+        return DocumentEvaluation(
+            name=annotation.name,
+            score=annotation.score,
+            label=annotation.label,
+            explanation=annotation.explanation,
+            document_position=annotation.document_position,
         )

phoenix/server/api/types/EvaluationSummary.py

@@ -1,13 +1,12 @@
-import math
-from functools import cached_property
-from typing import List, Optional, Tuple, cast
+from typing import List, Optional, Union, cast
 
 import pandas as pd
 import strawberry
-from pandas.api.types import CategoricalDtype
 from strawberry import Private
 
-import phoenix.trace.v1 as pb
+from phoenix.db import models
+
+AnnotationType = Union[models.SpanAnnotation, models.TraceAnnotation]
 
 
 @strawberry.type
@@ -18,58 +17,44 @@ class LabelFraction:
 
 @strawberry.type
 class EvaluationSummary:
-    count: int
-    labels: Tuple[str, ...]
-    evaluations: Private[Tuple[pb.Evaluation, ...]]
+    df: Private[pd.DataFrame]
+
+    def __init__(self, dataframe: pd.DataFrame) -> None:
+        self.df = dataframe
 
-    def __init__(
-        self,
-        evaluations: Tuple[pb.Evaluation, ...],
-        labels: Tuple[str, ...],
-    ) -> None:
-        self.evaluations = evaluations
-        self.labels = labels
-        self.count = len(evaluations)
+    @strawberry.field
+    def count(self) -> int:
+        return cast(int, self.df.record_count.sum())
+
+    @strawberry.field
+    def labels(self) -> List[str]:
+        return self.df.label.dropna().tolist()
 
     @strawberry.field
     def label_fractions(self) -> List[LabelFraction]:
-        if not self.labels or not (n := len(self._eval_labels)):
+        if not (n := self.df.label_count.sum()):
             return []
-        counts = self._eval_labels.value_counts(dropna=True)
         return [
-            LabelFraction(label=cast(str, label), fraction=count / n)
-            for label, count in counts.items()
+            LabelFraction(
+                label=cast(str, row.label),
+                fraction=row.label_count / n,
+            )
+            for row in self.df.loc[
+                self.df.label.notna(),
+                ["label", "label_count"],
+            ].itertuples()
         ]
 
     @strawberry.field
     def mean_score(self) -> Optional[float]:
-        value = self._eval_scores.mean()
-        return None if math.isnan(value) else value
+        if not (n := self.df.score_count.sum()):
+            return None
+        return cast(float, self.df.score_sum.sum() / n)
 
     @strawberry.field
     def score_count(self) -> int:
-        return self._eval_scores.count()
+        return cast(int, self.df.score_count.sum())
 
     @strawberry.field
     def label_count(self) -> int:
-        return self._eval_labels.count()
-
-    @cached_property
-    def _eval_scores(self) -> "pd.Series[float]":
-        return pd.Series(
-            (
-                evaluation.result.score.value if evaluation.result.HasField("score") else None
-                for evaluation in self.evaluations
-            ),
-            dtype=float,
-        )
-
-    @cached_property
-    def _eval_labels(self) -> "pd.Series[CategoricalDtype]":
-        return pd.Series(
-            (
-                evaluation.result.label.value if evaluation.result.HasField("label") else None
-                for evaluation in self.evaluations
-            ),
-            dtype=CategoricalDtype(categories=self.labels),  # type: ignore
-        )
+        return cast(int, self.df.label_count.sum())

phoenix/server/api/types/MimeType.py

@@ -8,8 +8,8 @@ import phoenix.trace.schemas as trace_schemas
 
 @strawberry.enum
 class MimeType(Enum):
-    text = trace_schemas.MimeType.TEXT
-    json = trace_schemas.MimeType.JSON
+    text = trace_schemas.MimeType.TEXT.value
+    json = trace_schemas.MimeType.JSON.value
 
     @classmethod
     def _missing_(cls, v: Any) -> Optional["MimeType"]:

phoenix/server/api/types/Model.py

@@ -19,7 +19,7 @@ from .DatasetRole import AncillaryDatasetRole, DatasetRole
 from .Dimension import Dimension, to_gql_dimension
 from .EmbeddingDimension import EmbeddingDimension, to_gql_embedding_dimension
 from .ExportedFile import ExportedFile
-from .pagination import Connection, ConnectionArgs, Cursor, connection_from_list
+from .pagination import Connection, ConnectionArgs, CursorString, connection_from_list
 from .TimeSeries import (
     PerformanceTimeSeries,
     ensure_timeseries_parameters,
@@ -35,8 +35,8 @@ class Model:
         info: Info[Context, None],
         first: Optional[int] = 50,
         last: Optional[int] = UNSET,
-        after: Optional[Cursor] = UNSET,
-        before: Optional[Cursor] = UNSET,
+        after: Optional[CursorString] = UNSET,
+        before: Optional[CursorString] = UNSET,
         include: Optional[DimensionFilter] = UNSET,
         exclude: Optional[DimensionFilter] = UNSET,
     ) -> Connection[Dimension]:
@@ -50,9 +50,9 @@ class Model:
             ],
             args=ConnectionArgs(
                 first=first,
-                after=after if isinstance(after, Cursor) else None,
+                after=after if isinstance(after, CursorString) else None,
                 last=last,
-                before=before if isinstance(before, Cursor) else None,
+                before=before if isinstance(before, CursorString) else None,
             ),
         )
 
@@ -105,8 +105,8 @@ class Model:
         info: Info[Context, None],
         first: Optional[int] = 50,
         last: Optional[int] = UNSET,
-        after: Optional[Cursor] = UNSET,
-        before: Optional[Cursor] = UNSET,
+        after: Optional[CursorString] = UNSET,
+        before: Optional[CursorString] = UNSET,
     ) -> Connection[EmbeddingDimension]:
         """
         A non-trivial implementation should efficiently fetch only
@@ -123,9 +123,9 @@ class Model:
             ],
             args=ConnectionArgs(
                 first=first,
-                after=after if isinstance(after, Cursor) else None,
+                after=after if isinstance(after, CursorString) else None,
                 last=last,
-                before=before if isinstance(before, Cursor) else None,
+                before=before if isinstance(before, CursorString) else None,
             ),
         )