arize-phoenix 4.15.0__py3-none-any.whl → 4.16.0__py3-none-any.whl

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

@@ -1,6 +1,6 @@
 import gzip
 from itertools import chain
-from typing import Iterator, Optional, Tuple
+from typing import Any, Callable, Iterator, Optional, Tuple, Union, cast
 
 import pandas as pd
 import pyarrow as pa
@@ -24,10 +24,10 @@ from typing_extensions import TypeAlias
 import phoenix.trace.v1 as pb
 from phoenix.config import DEFAULT_PROJECT_NAME
 from phoenix.db import models
+from phoenix.db.insertion.types import Precursors
 from phoenix.exceptions import PhoenixEvaluationNameIsMissing
 from phoenix.server.api.routers.utils import table_to_bytes
 from phoenix.server.types import DbSessionFactory
-from phoenix.session.evaluation import encode_evaluations
 from phoenix.trace.span_evaluations import (
     DocumentEvaluations,
     Evaluations,
@@ -194,8 +194,94 @@ async def _process_pyarrow(request: Request) -> Response:
 
 
 async def _add_evaluations(state: State, evaluations: Evaluations) -> None:
-    for evaluation in encode_evaluations(evaluations):
-        await state.queue_evaluation_for_bulk_insert(evaluation)
+    dataframe = evaluations.dataframe
+    eval_name = evaluations.eval_name
+    names = dataframe.index.names
+    if (
+        len(names) == 2
+        and "document_position" in names
+        and ("context.span_id" in names or "span_id" in names)
+    ):
+        cls = _document_annotation_factory(
+            names.index("span_id") if "span_id" in names else names.index("context.span_id"),
+            names.index("document_position"),
+        )
+        for index, row in dataframe.iterrows():
+            score, label, explanation = _get_annotation_result(row)
+            document_annotation = cls(cast(Union[Tuple[str, int], Tuple[int, str]], index))(
+                name=eval_name,
+                annotator_kind="LLM",
+                score=score,
+                label=label,
+                explanation=explanation,
+                metadata_={},
+            )
+            await state.enqueue(document_annotation)
+    elif len(names) == 1 and names[0] in ("context.span_id", "span_id"):
+        for index, row in dataframe.iterrows():
+            score, label, explanation = _get_annotation_result(row)
+            span_annotation = _span_annotation_factory(cast(str, index))(
+                name=eval_name,
+                annotator_kind="LLM",
+                score=score,
+                label=label,
+                explanation=explanation,
+                metadata_={},
+            )
+            await state.enqueue(span_annotation)
+    elif len(names) == 1 and names[0] in ("context.trace_id", "trace_id"):
+        for index, row in dataframe.iterrows():
+            score, label, explanation = _get_annotation_result(row)
+            trace_annotation = _trace_annotation_factory(cast(str, index))(
+                name=eval_name,
+                annotator_kind="LLM",
+                score=score,
+                label=label,
+                explanation=explanation,
+                metadata_={},
+            )
+            await state.enqueue(trace_annotation)
+
+
+def _get_annotation_result(
+    row: "pd.Series[Any]",
+) -> Tuple[Optional[float], Optional[str], Optional[str]]:
+    return (
+        cast(Optional[float], row.get("score")),
+        cast(Optional[str], row.get("label")),
+        cast(Optional[str], row.get("explanation")),
+    )
+
+
+def _document_annotation_factory(
+    span_id_idx: int,
+    document_position_idx: int,
+) -> Callable[
+    [Union[Tuple[str, int], Tuple[int, str]]],
+    Callable[..., Precursors.DocumentAnnotation],
+]:
+    return lambda index: lambda **kwargs: Precursors.DocumentAnnotation(
+        span_id=str(index[span_id_idx]),
+        document_position=int(index[document_position_idx]),
+        obj=models.DocumentAnnotation(
+            document_position=int(index[document_position_idx]),
+            **kwargs,
+        ),
+    )
+
+
+def _span_annotation_factory(span_id: str) -> Callable[..., Precursors.SpanAnnotation]:
+    return lambda **kwargs: Precursors.SpanAnnotation(
+        span_id=str(span_id),
+        obj=models.SpanAnnotation(**kwargs),
+    )
+
+
+def _trace_annotation_factory(trace_id: str) -> Callable[..., Precursors.TraceAnnotation]:
+    return lambda **kwargs: Precursors.TraceAnnotation(
+        trace_id=str(trace_id),
+        obj=models.TraceAnnotation(**kwargs),
+    )
 
 
 def _read_sql_trace_evaluations_into_dataframe(

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

@@ -13,9 +13,9 @@ from phoenix.config import DEFAULT_PROJECT_NAME
 from phoenix.datetime_utils import normalize_datetime
 from phoenix.db import models
 from phoenix.db.helpers import SupportedSQLDialect
-from phoenix.db.insertion.helpers import insert_on_conflict
+from phoenix.db.insertion.helpers import as_kv, insert_on_conflict
+from phoenix.db.insertion.types import Precursors
 from phoenix.server.api.routers.utils import df_to_bytes
-from phoenix.server.api.types.node import from_global_id_with_expected_type
 from phoenix.trace.dsl import SpanQuery as SpanQuery_
 
 from .pydantic_compat import V1RoutesBaseModel
@@ -143,7 +143,7 @@ class SpanAnnotationResult(V1RoutesBaseModel):
 
 
 class SpanAnnotation(V1RoutesBaseModel):
-    span_id: str = Field(description="The ID of the span being annotated")
+    span_id: str = Field(description="OpenTelemetry Span ID (hex format w/o 0x prefix)")
     name: str = Field(description="The name of the annotation")
     annotator_kind: Literal["LLM", "HUMAN"] = Field(
         description="The kind of annotator used for the annotation"
@@ -155,6 +155,19 @@ class SpanAnnotation(V1RoutesBaseModel):
         default=None, description="Metadata for the annotation"
     )
 
+    def as_precursor(self) -> Precursors.SpanAnnotation:
+        return Precursors.SpanAnnotation(
+            self.span_id,
+            models.SpanAnnotation(
+                name=self.name,
+                annotator_kind=self.annotator_kind,
+                score=self.result.score if self.result else None,
+                label=self.result.label if self.result else None,
+                explanation=self.result.explanation if self.result else None,
+                metadata_=self.metadata or {},
+            ),
+        )
+
 
 class AnnotateSpansRequestBody(RequestBody[List[SpanAnnotation]]):
     data: List[SpanAnnotation]
@@ -178,59 +191,36 @@ class AnnotateSpansResponseBody(ResponseBody[List[InsertedSpanAnnotation]]):
     response_description="Span annotations inserted successfully",
 )
 async def annotate_spans(
-    request: Request, request_body: AnnotateSpansRequestBody
+    request: Request,
+    request_body: AnnotateSpansRequestBody,
+    sync: bool = Query(default=True, description="If true, fulfill request synchronously."),
 ) -> AnnotateSpansResponseBody:
-    span_annotations = request_body.data
-    span_gids = [GlobalID.from_id(annotation.span_id) for annotation in span_annotations]
-
-    resolved_span_ids = []
-    for span_gid in span_gids:
-        try:
-            resolved_span_ids.append(from_global_id_with_expected_type(span_gid, "Span"))
-        except ValueError:
-            raise HTTPException(
-                detail="Span with ID {span_gid} does not exist",
-                status_code=HTTP_404_NOT_FOUND,
-            )
+    precursors = [d.as_precursor() for d in request_body.data]
+    if not sync:
+        await request.state.enqueue(*precursors)
+        return AnnotateSpansResponseBody(data=[])
 
+    span_ids = {p.span_id for p in precursors}
     async with request.app.state.db() as session:
-        spans = await session.execute(
-            select(models.Span).filter(models.Span.id.in_(resolved_span_ids))
-        )
-        existing_span_ids = {span.id for span in spans.scalars()}
+        existing_spans = {
+            span.span_id: span.id
+            async for span in await session.stream_scalars(
+                select(models.Span).filter(models.Span.span_id.in_(span_ids))
+            )
+        }
 
-        missing_span_ids = set(resolved_span_ids) - existing_span_ids
+        missing_span_ids = span_ids - set(existing_spans.keys())
         if missing_span_ids:
-            missing_span_gids = [
-                str(GlobalID("Span", str(span_gid))) for span_gid in missing_span_ids
-            ]
             raise HTTPException(
-                detail=f"Spans with IDs {', '.join(missing_span_gids)} do not exist.",
+                detail=f"Spans with IDs {', '.join(missing_span_ids)} do not exist.",
                 status_code=HTTP_404_NOT_FOUND,
             )
 
         inserted_annotations = []
-        for annotation in span_annotations:
-            span_gid = GlobalID.from_id(annotation.span_id)
-            span_id = from_global_id_with_expected_type(span_gid, "Span")
-            name = annotation.name
-            annotator_kind = annotation.annotator_kind
-            result = annotation.result
-            label = result.label if result else None
-            score = result.score if result else None
-            explanation = result.explanation if result else None
-            metadata = annotation.metadata or {}
-
-            values = dict(
-                span_rowid=span_id,
-                name=name,
-                label=label,
-                score=score,
-                explanation=explanation,
-                annotator_kind=annotator_kind,
-                metadata_=metadata,
-            )
-            dialect = SupportedSQLDialect(session.bind.dialect.name)
+
+        dialect = SupportedSQLDialect(session.bind.dialect.name)
+        for p in precursors:
+            values = dict(as_kv(p.as_insertable(existing_spans[p.span_id]).row))
             span_annotation_id = await session.scalar(
                 insert_on_conflict(
                     values,

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

@@ -2,7 +2,7 @@ import gzip
 import zlib
 from typing import Any, Dict, List, Literal, Optional
 
-from fastapi import APIRouter, BackgroundTasks, Header, HTTPException
+from fastapi import APIRouter, BackgroundTasks, Header, HTTPException, Query
 from google.protobuf.message import DecodeError
 from opentelemetry.proto.collector.trace.v1.trace_service_pb2 import (
     ExportTraceServiceRequest,
@@ -22,8 +22,8 @@ from strawberry.relay import GlobalID
 
 from phoenix.db import models
 from phoenix.db.helpers import SupportedSQLDialect
-from phoenix.db.insertion.helpers import insert_on_conflict
-from phoenix.server.api.types.node import from_global_id_with_expected_type
+from phoenix.db.insertion.helpers import as_kv, insert_on_conflict
+from phoenix.db.insertion.types import Precursors
 from phoenix.trace.otel import decode_otlp_span
 from phoenix.utilities.project import get_project_name
 
@@ -100,7 +100,7 @@ class TraceAnnotationResult(V1RoutesBaseModel):
 
 
 class TraceAnnotation(V1RoutesBaseModel):
-    trace_id: str = Field(description="The ID of the trace being annotated")
+    trace_id: str = Field(description="OpenTelemetry Trace ID (hex format w/o 0x prefix)")
     name: str = Field(description="The name of the annotation")
     annotator_kind: Literal["LLM", "HUMAN"] = Field(
         description="The kind of annotator used for the annotation"
@@ -112,6 +112,19 @@ class TraceAnnotation(V1RoutesBaseModel):
         default=None, description="Metadata for the annotation"
     )
 
+    def as_precursor(self) -> Precursors.TraceAnnotation:
+        return Precursors.TraceAnnotation(
+            self.trace_id,
+            models.TraceAnnotation(
+                name=self.name,
+                annotator_kind=self.annotator_kind,
+                score=self.result.score if self.result else None,
+                label=self.result.label if self.result else None,
+                explanation=self.result.explanation if self.result else None,
+                metadata_=self.metadata or {},
+            ),
+        )
+
 
 class AnnotateTracesRequestBody(RequestBody[List[TraceAnnotation]]):
     data: List[TraceAnnotation] = Field(description="The trace annotations to be upserted")
@@ -134,61 +147,36 @@ class AnnotateTracesResponseBody(ResponseBody[List[InsertedTraceAnnotation]]):
     ),
 )
 async def annotate_traces(
-    request: Request, request_body: AnnotateTracesRequestBody
+    request: Request,
+    request_body: AnnotateTracesRequestBody,
+    sync: bool = Query(default=True, description="If true, fulfill request synchronously."),
 ) -> AnnotateTracesResponseBody:
-    trace_annotations = request_body.data
-    trace_gids = [GlobalID.from_id(annotation.trace_id) for annotation in trace_annotations]
-
-    resolved_trace_ids = []
-    for trace_gid in trace_gids:
-        try:
-            resolved_trace_ids.append(from_global_id_with_expected_type(trace_gid, "Trace"))
-        except ValueError:
-            raise HTTPException(
-                detail="Trace with ID {trace_gid} does not exist",
-                status_code=HTTP_404_NOT_FOUND,
-            )
+    precursors = [d.as_precursor() for d in request_body.data]
+    if not sync:
+        await request.state.enqueue(*precursors)
+        return AnnotateTracesResponseBody(data=[])
 
+    trace_ids = {p.trace_id for p in precursors}
     async with request.app.state.db() as session:
-        traces = await session.execute(
-            select(models.Trace).filter(models.Trace.id.in_(resolved_trace_ids))
-        )
-        existing_trace_ids = {trace.id for trace in traces.scalars()}
+        existing_traces = {
+            trace.trace_id: trace.id
+            async for trace in await session.stream_scalars(
+                select(models.Trace).filter(models.Trace.trace_id.in_(trace_ids))
+            )
+        }
 
-        missing_trace_ids = set(resolved_trace_ids) - existing_trace_ids
+        missing_trace_ids = trace_ids - set(existing_traces.keys())
         if missing_trace_ids:
-            missing_trace_gids = [
-                str(GlobalID("Trace", str(trace_gid))) for trace_gid in missing_trace_ids
-            ]
             raise HTTPException(
-                detail=f"Traces with IDs {', '.join(missing_trace_gids)} do not exist.",
+                detail=f"Traces with IDs {', '.join(missing_trace_ids)} do not exist.",
                 status_code=HTTP_404_NOT_FOUND,
             )
 
         inserted_annotations = []
 
-        for annotation in trace_annotations:
-            trace_gid = GlobalID.from_id(annotation.trace_id)
-            trace_id = from_global_id_with_expected_type(trace_gid, "Trace")
-
-            name = annotation.name
-            annotator_kind = annotation.annotator_kind
-            result = annotation.result
-            label = result.label if result else None
-            score = result.score if result else None
-            explanation = result.explanation if result else None
-            metadata = annotation.metadata or {}
-
-            values = dict(
-                trace_rowid=trace_id,
-                name=name,
-                label=label,
-                score=score,
-                explanation=explanation,
-                annotator_kind=annotator_kind,
-                metadata_=metadata,
-            )
-            dialect = SupportedSQLDialect(session.bind.dialect.name)
+        dialect = SupportedSQLDialect(session.bind.dialect.name)
+        for p in precursors:
+            values = dict(as_kv(p.as_insertable(existing_traces[p.trace_id]).row))
             trace_annotation_id = await session.scalar(
                 insert_on_conflict(
                     values,

phoenix/server/api/types/Span.py

@@ -7,6 +7,7 @@ from typing import TYPE_CHECKING, Any, List, Mapping, Optional, Sized, cast
 import numpy as np
 import strawberry
 from openinference.semconv.trace import EmbeddingAttributes, SpanAttributes
+from sqlalchemy import select
 from strawberry import ID, UNSET
 from strawberry.relay import Node, NodeID
 from strawberry.types import Info
@@ -19,6 +20,9 @@ from phoenix.server.api.helpers.dataset_helpers import (
     get_dataset_example_input,
     get_dataset_example_output,
 )
+from phoenix.server.api.input_types.SpanAnnotationSort import SpanAnnotationSort
+from phoenix.server.api.types.SortDir import SortDir
+from phoenix.server.api.types.SpanAnnotation import to_gql_span_annotation
 from phoenix.trace.attributes import get_attribute_value
 
 from .DocumentRetrievalMetrics import DocumentRetrievalMetrics
@@ -177,12 +181,27 @@ class Span(Node):
 
     @strawberry.field(
         description=(
-            "Annotations of the span's parent span. This encompasses both "
+            "Annotations associated with the span. This encompasses both "
             "LLM and human annotations."
         )
     )  # type: ignore
-    async def span_annotations(self, info: Info[Context, None]) -> List[SpanAnnotation]:
-        return await info.context.data_loaders.span_annotations.load(self.id_attr)
+    async def span_annotations(
+        self,
+        info: Info[Context, None],
+        sort: Optional[SpanAnnotationSort] = UNSET,
+    ) -> List[SpanAnnotation]:
+        async with info.context.db() as session:
+            stmt = select(models.SpanAnnotation).filter_by(span_rowid=self.id_attr)
+            if sort:
+                sort_col = getattr(models.SpanAnnotation, sort.col.value)
+                if sort.dir is SortDir.desc:
+                    stmt = stmt.order_by(sort_col.desc(), models.SpanAnnotation.id.desc())
+                else:
+                    stmt = stmt.order_by(sort_col.asc(), models.SpanAnnotation.id.asc())
+            else:
+                stmt = stmt.order_by(models.SpanAnnotation.created_at.desc())
+            annotations = await session.scalars(stmt)
+        return [to_gql_span_annotation(annotation) for annotation in annotations]
 
     @strawberry.field(
         description="Evaluations of the documents associated with the span, e.g. "

phoenix/server/api/types/Trace.py

@@ -11,13 +11,15 @@ from strawberry.types import Info
 
 from phoenix.db import models
 from phoenix.server.api.context import Context
-from phoenix.server.api.types.Evaluation import TraceEvaluation
+from phoenix.server.api.input_types.TraceAnnotationSort import TraceAnnotationSort
 from phoenix.server.api.types.pagination import (
     ConnectionArgs,
     CursorString,
     connection_from_list,
 )
+from phoenix.server.api.types.SortDir import SortDir
 from phoenix.server.api.types.Span import Span, to_gql_span
+from phoenix.server.api.types.TraceAnnotation import TraceAnnotation, to_gql_trace_annotation
 
 
 @strawberry.type
@@ -62,6 +64,21 @@ class Trace(Node):
             data = [to_gql_span(span) async for span in spans]
         return connection_from_list(data=data, args=args)
 
-    @strawberry.field(description="Evaluations associated with the trace")  # type: ignore
-    async def trace_evaluations(self, info: Info[Context, None]) -> List[TraceEvaluation]:
-        return await info.context.data_loaders.trace_evaluations.load(self.id_attr)
+    @strawberry.field(description="Annotations associated with the trace.")  # type: ignore
+    async def span_annotations(
+        self,
+        info: Info[Context, None],
+        sort: Optional[TraceAnnotationSort] = None,
+    ) -> List[TraceAnnotation]:
+        async with info.context.db() as session:
+            stmt = select(models.TraceAnnotation).filter_by(span_rowid=self.id_attr)
+            if sort:
+                sort_col = getattr(models.TraceAnnotation, sort.col.value)
+                if sort.dir is SortDir.desc:
+                    stmt = stmt.order_by(sort_col.desc(), models.TraceAnnotation.id.desc())
+                else:
+                    stmt = stmt.order_by(sort_col.asc(), models.TraceAnnotation.id.asc())
+            else:
+                stmt = stmt.order_by(models.TraceAnnotation.created_at.desc())
+            annotations = await session.scalars(stmt)
+        return [to_gql_trace_annotation(annotation) for annotation in annotations]

phoenix/server/app.py

@@ -229,6 +229,7 @@ def _lifespan(
         global DB_MUTEX
         DB_MUTEX = asyncio.Lock() if dialect is SupportedSQLDialect.SQLITE else None
         async with bulk_inserter as (
+            enqueue,
             queue_span,
             queue_evaluation,
             enqueue_operation,
@@ -239,6 +240,7 @@ def _lifespan(
             enable_prometheus=enable_prometheus,
         ):
             yield {
+                "enqueue": enqueue,
                 "queue_span_for_bulk_insert": queue_span,
                 "queue_evaluation_for_bulk_insert": queue_evaluation,
                 "enqueue_operation": enqueue_operation,

phoenix/server/static/.vite/manifest.json

@@ -1,22 +1,22 @@
 {
-  "_components-kGgeFkHp.js": {
-    "file": "assets/components-kGgeFkHp.js",
+  "_components-Ci5kMOk5.js": {
+    "file": "assets/components-Ci5kMOk5.js",
     "name": "components",
     "imports": [
       "_vendor-CP0b0YG0.js",
-      "_vendor-arizeai-B5Hti8OB.js",
-      "_pages-DabDCmVd.js",
+      "_vendor-arizeai-DTbiPGp6.js",
+      "_pages-BrevprVW.js",
       "_vendor-three-DwGkEfCM.js",
       "_vendor-codemirror-DtdPDzrv.js"
     ]
   },
-  "_pages-DabDCmVd.js": {
-    "file": "assets/pages-DabDCmVd.js",
+  "_pages-BrevprVW.js": {
+    "file": "assets/pages-BrevprVW.js",
     "name": "pages",
     "imports": [
       "_vendor-CP0b0YG0.js",
-      "_components-kGgeFkHp.js",
-      "_vendor-arizeai-B5Hti8OB.js",
+      "_components-Ci5kMOk5.js",
+      "_vendor-arizeai-DTbiPGp6.js",
       "_vendor-recharts-A0DA1O99.js",
       "_vendor-codemirror-DtdPDzrv.js"
     ]
@@ -35,8 +35,8 @@
       "assets/vendor-DxkFTwjz.css"
     ]
   },
-  "_vendor-arizeai-B5Hti8OB.js": {
-    "file": "assets/vendor-arizeai-B5Hti8OB.js",
+  "_vendor-arizeai-DTbiPGp6.js": {
+    "file": "assets/vendor-arizeai-DTbiPGp6.js",
     "name": "vendor-arizeai",
     "imports": [
       "_vendor-CP0b0YG0.js"
@@ -61,15 +61,15 @@
     "name": "vendor-three"
   },
   "index.tsx": {
-    "file": "assets/index-BctFO6S7.js",
+    "file": "assets/index-BQG5WVX7.js",
     "name": "index",
     "src": "index.tsx",
     "isEntry": true,
     "imports": [
       "_vendor-CP0b0YG0.js",
-      "_vendor-arizeai-B5Hti8OB.js",
-      "_pages-DabDCmVd.js",
-      "_components-kGgeFkHp.js",
+      "_vendor-arizeai-DTbiPGp6.js",
+      "_pages-BrevprVW.js",
+      "_components-Ci5kMOk5.js",
       "_vendor-three-DwGkEfCM.js",
       "_vendor-recharts-A0DA1O99.js",
       "_vendor-codemirror-DtdPDzrv.js"