arize-phoenix 8.32.1__py3-none-any.whl → 9.0.1__py3-none-any.whl

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

@@ -1,3 +1,4 @@
+import warnings
 from asyncio import get_running_loop
 from collections.abc import AsyncIterator
 from datetime import datetime, timezone
@@ -20,6 +21,7 @@ from phoenix.db.helpers import SupportedSQLDialect
 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.bearer_auth import PhoenixUser
 from phoenix.server.dml_event import SpanAnnotationInsertEvent
 from phoenix.trace.dsl import SpanQuery as SpanQuery_
 from phoenix.utilities.json import encode_df_as_json_string
@@ -171,10 +173,10 @@ class SpanAnnotationResult(V1RoutesBaseModel):
     )
 
 
-class SpanAnnotation(V1RoutesBaseModel):
+class SpanAnnotationData(V1RoutesBaseModel):
     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(
+    annotator_kind: Literal["LLM", "CODE", "HUMAN"] = Field(
         description="The kind of annotator used for the annotation"
     )
     result: Optional[SpanAnnotationResult] = Field(
@@ -183,8 +185,15 @@ class SpanAnnotation(V1RoutesBaseModel):
     metadata: Optional[dict[str, Any]] = Field(
         default=None, description="Metadata for the annotation"
     )
+    identifier: str = Field(
+        default="",
+        description=(
+            "The identifier of the annotation. "
+            "If provided, the annotation will be updated if it already exists."
+        ),
+    )
 
-    def as_precursor(self) -> Precursors.SpanAnnotation:
+    def as_precursor(self, *, user_id: Optional[int] = None) -> Precursors.SpanAnnotation:
         return Precursors.SpanAnnotation(
             self.span_id,
             models.SpanAnnotation(
@@ -194,12 +203,15 @@ class SpanAnnotation(V1RoutesBaseModel):
                 label=self.result.label if self.result else None,
                 explanation=self.result.explanation if self.result else None,
                 metadata_=self.metadata or {},
+                identifier=self.identifier,
+                source="API",
+                user_id=user_id,
             ),
         )
 
 
-class AnnotateSpansRequestBody(RequestBody[list[SpanAnnotation]]):
-    data: list[SpanAnnotation]
+class AnnotateSpansRequestBody(RequestBody[list[SpanAnnotationData]]):
+    data: list[SpanAnnotationData]
 
 
 class InsertedSpanAnnotation(V1RoutesBaseModel):
@@ -213,7 +225,7 @@ class AnnotateSpansResponseBody(ResponseBody[list[InsertedSpanAnnotation]]):
 @router.post(
     "/span_annotations",
     operation_id="annotateSpans",
-    summary="Create or update span annotations",
+    summary="Create span annotations",
     responses=add_errors_to_responses(
         [{"status_code": HTTP_404_NOT_FOUND, "description": "Span not found"}]
     ),
@@ -227,7 +239,22 @@ async def annotate_spans(
 ) -> AnnotateSpansResponseBody:
     if not request_body.data:
         return AnnotateSpansResponseBody(data=[])
-    precursors = [d.as_precursor() for d in request_body.data]
+
+    user_id: Optional[int] = None
+    if request.app.state.authentication_enabled and isinstance(request.user, PhoenixUser):
+        user_id = int(request.user.identity)
+
+    span_annotations = request_body.data
+    filtered_span_annotations = list(filter(lambda d: d.name != "note", span_annotations))
+    if len(filtered_span_annotations) != len(span_annotations):
+        warnings.warn(
+            (
+                "Span annotations with the name 'note' are not supported in this endpoint. "
+                "They will be ignored."
+            ),
+            UserWarning,
+        )
+    precursors = [d.as_precursor(user_id=user_id) for d in filtered_span_annotations]
     if not sync:
         await request.state.enqueue(*precursors)
         return AnnotateSpansResponseBody(data=[])
@@ -256,7 +283,7 @@ async def annotate_spans(
                     values,
                     dialect=dialect,
                     table=models.SpanAnnotation,
-                    unique_by=("name", "span_rowid"),
+                    unique_by=("name", "span_rowid", "identifier"),
                 ).returning(models.SpanAnnotation.id)
             )
             inserted_ids.append(span_annotation_id)

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

@@ -10,7 +10,7 @@ from opentelemetry.proto.collector.trace.v1.trace_service_pb2 import (
     ExportTraceServiceResponse,
 )
 from pydantic import Field
-from sqlalchemy import select
+from sqlalchemy import insert, select
 from starlette.concurrency import run_in_threadpool
 from starlette.datastructures import State
 from starlette.requests import Request
@@ -23,9 +23,9 @@ from starlette.status import (
 from strawberry.relay import GlobalID
 
 from phoenix.db import models
-from phoenix.db.helpers import SupportedSQLDialect
-from phoenix.db.insertion.helpers import as_kv, insert_on_conflict
+from phoenix.db.insertion.helpers import as_kv
 from phoenix.db.insertion.types import Precursors
+from phoenix.server.bearer_auth import PhoenixUser
 from phoenix.server.dml_event import TraceAnnotationInsertEvent
 from phoenix.trace.otel import decode_otlp_span
 from phoenix.utilities.project import get_project_name
@@ -114,8 +114,15 @@ class TraceAnnotation(V1RoutesBaseModel):
     metadata: Optional[dict[str, Any]] = Field(
         default=None, description="Metadata for the annotation"
     )
+    identifier: str = Field(
+        default="",
+        description=(
+            "The identifier of the annotation. "
+            "If provided, the annotation will be updated if it already exists."
+        ),
+    )
 
-    def as_precursor(self) -> Precursors.TraceAnnotation:
+    def as_precursor(self, *, user_id: Optional[int] = None) -> Precursors.TraceAnnotation:
         return Precursors.TraceAnnotation(
             self.trace_id,
             models.TraceAnnotation(
@@ -125,6 +132,9 @@ class TraceAnnotation(V1RoutesBaseModel):
                 label=self.result.label if self.result else None,
                 explanation=self.result.explanation if self.result else None,
                 metadata_=self.metadata or {},
+                identifier=self.identifier,
+                source="APP",
+                user_id=user_id,
             ),
         )
 
@@ -144,7 +154,7 @@ class AnnotateTracesResponseBody(ResponseBody[list[InsertedTraceAnnotation]]):
 @router.post(
     "/trace_annotations",
     operation_id="annotateTraces",
-    summary="Create or update trace annotations",
+    summary="Create trace annotations",
     responses=add_errors_to_responses(
         [{"status_code": HTTP_404_NOT_FOUND, "description": "Trace not found"}]
     ),
@@ -157,7 +167,12 @@ async def annotate_traces(
 ) -> AnnotateTracesResponseBody:
     if not request_body.data:
         return AnnotateTracesResponseBody(data=[])
-    precursors = [d.as_precursor() for d in request_body.data]
+
+    user_id: Optional[int] = None
+    if request.app.state.authentication_enabled and isinstance(request.user, PhoenixUser):
+        user_id = int(request.user.identity)
+
+    precursors = [d.as_precursor(user_id=user_id) for d in request_body.data]
     if not sync:
         await request.state.enqueue(*precursors)
         return AnnotateTracesResponseBody(data=[])
@@ -178,16 +193,10 @@ async def annotate_traces(
                 status_code=HTTP_404_NOT_FOUND,
             )
         inserted_ids = []
-        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,
-                    dialect=dialect,
-                    table=models.TraceAnnotation,
-                    unique_by=("name", "trace_rowid"),
-                ).returning(models.TraceAnnotation.id)
+                insert(models.TraceAnnotation).values(**values).returning(models.TraceAnnotation.id)
             )
             inserted_ids.append(trace_annotation_id)
     request.state.event_queue.put(TraceAnnotationInsertEvent(tuple(inserted_ids)))

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

@@ -1,7 +1,19 @@
 from typing import Any, Generic, Optional, TypedDict, TypeVar, Union
 
+from fastapi import HTTPException
+from sqlalchemy import select
+from sqlalchemy.ext.asyncio import AsyncSession
+from starlette.status import (
+    HTTP_404_NOT_FOUND,
+    HTTP_422_UNPROCESSABLE_ENTITY,
+)
+from strawberry.relay import GlobalID
 from typing_extensions import TypeAlias, assert_never
 
+from phoenix.db import models
+from phoenix.server.api.types.node import from_global_id_with_expected_type
+from phoenix.server.api.types.Project import Project as ProjectNodeType
+
 from .models import V1RoutesBaseModel
 
 StatusCode: TypeAlias = int
@@ -93,3 +105,51 @@ def add_text_csv_content_to_responses(
         "text/csv": {"schema": {"type": "string", "contentMediaType": "text/csv"}}
     }
     return output_responses
+
+
+async def _get_project_by_identifier(
+    session: AsyncSession,
+    project_identifier: str,
+) -> models.Project:
+    """
+    Get a project by its ID or name.
+
+    Args:
+        session: The database session.
+        project_identifier: The project ID or name.
+
+    Returns:
+        The project object.
+
+    Raises:
+        HTTPException: If the identifier format is invalid or the project is not found.
+    """
+    # Try to parse as a GlobalID first
+    try:
+        id_ = from_global_id_with_expected_type(
+            GlobalID.from_id(project_identifier),
+            ProjectNodeType.__name__,
+        )
+    except Exception:
+        try:
+            name = project_identifier
+        except HTTPException:
+            raise HTTPException(
+                status_code=HTTP_422_UNPROCESSABLE_ENTITY,
+                detail=f"Invalid project identifier format: {project_identifier}",
+            )
+        stmt = select(models.Project).filter_by(name=name)
+        project = await session.scalar(stmt)
+        if project is None:
+            raise HTTPException(
+                status_code=HTTP_404_NOT_FOUND,
+                detail=f"Project with name {name} not found",
+            )
+    else:
+        project = await session.get(models.Project, id_)
+        if project is None:
+            raise HTTPException(
+                status_code=HTTP_404_NOT_FOUND,
+                detail=f"Project with ID {project_identifier} not found",
+            )
+    return project

phoenix/server/api/types/Annotation.py

@@ -1,3 +1,4 @@
+from datetime import datetime
 from typing import Optional
 
 import strawberry
@@ -22,3 +23,9 @@ class Annotation:
         description="The annotator's explanation for the annotation result (i.e. "
         "score or label, or both) given to the subject."
     )
+    created_at: datetime = strawberry.field(
+        description="The date and time when the annotation was created."
+    )
+    updated_at: datetime = strawberry.field(
+        description="The date and time when the annotation was last updated."
+    )

phoenix/server/api/types/AnnotationConfig.py

@@ -0,0 +1,124 @@
+from typing import Annotated, Optional, Union
+
+import strawberry
+from strawberry.relay import Node, NodeID
+from typing_extensions import TypeAlias, assert_never
+
+from phoenix.db import models
+from phoenix.db.types.annotation_configs import (
+    AnnotationType,
+    OptimizationDirection,
+)
+from phoenix.db.types.annotation_configs import (
+    CategoricalAnnotationConfig as CategoricalAnnotationConfigModel,
+)
+from phoenix.db.types.annotation_configs import (
+    ContinuousAnnotationConfig as ContinuousAnnotationConfigModel,
+)
+from phoenix.db.types.annotation_configs import (
+    FreeformAnnotationConfig as FreeformAnnotationConfigModel,
+)
+
+
+@strawberry.interface
+class AnnotationConfigBase:
+    name: str
+    description: Optional[str]
+    annotation_type: AnnotationType
+
+
+@strawberry.type
+class CategoricalAnnotationValue:
+    label: str
+    score: Optional[float]
+
+
+@strawberry.type
+class CategoricalAnnotationConfig(Node, AnnotationConfigBase):
+    id_attr: NodeID[int]
+    optimization_direction: OptimizationDirection
+    values: list[CategoricalAnnotationValue]
+
+
+@strawberry.type
+class ContinuousAnnotationConfig(Node, AnnotationConfigBase):
+    id_attr: NodeID[int]
+    optimization_direction: OptimizationDirection
+    lower_bound: Optional[float]
+    upper_bound: Optional[float]
+
+
+@strawberry.type
+class FreeformAnnotationConfig(Node, AnnotationConfigBase):
+    id_attr: NodeID[int]
+
+
+AnnotationConfig: TypeAlias = Annotated[
+    Union[CategoricalAnnotationConfig, ContinuousAnnotationConfig, FreeformAnnotationConfig],
+    strawberry.union("AnnotationConfig"),
+]
+
+
+def _to_gql_categorical_annotation_config(
+    annotation_config: models.AnnotationConfig,
+) -> CategoricalAnnotationConfig:
+    config = annotation_config.config
+    assert isinstance(config, CategoricalAnnotationConfigModel)
+    values = [
+        CategoricalAnnotationValue(
+            label=val.label,
+            score=val.score,
+        )
+        for val in config.values
+    ]
+    return CategoricalAnnotationConfig(
+        id_attr=annotation_config.id,
+        name=annotation_config.name,
+        annotation_type=config.type,
+        optimization_direction=config.optimization_direction,
+        description=config.description,
+        values=values,
+    )
+
+
+def _to_gql_continuous_annotation_config(
+    annotation_config: models.AnnotationConfig,
+) -> ContinuousAnnotationConfig:
+    config = annotation_config.config
+    assert isinstance(config, ContinuousAnnotationConfigModel)
+    return ContinuousAnnotationConfig(
+        id_attr=annotation_config.id,
+        name=annotation_config.name,
+        annotation_type=config.type,
+        optimization_direction=config.optimization_direction,
+        description=config.description,
+        lower_bound=config.lower_bound,
+        upper_bound=config.upper_bound,
+    )
+
+
+def _to_gql_freeform_annotation_config(
+    annotation_config: models.AnnotationConfig,
+) -> FreeformAnnotationConfig:
+    config = annotation_config.config
+    assert isinstance(config, FreeformAnnotationConfigModel)
+    return FreeformAnnotationConfig(
+        id_attr=annotation_config.id,
+        name=annotation_config.name,
+        annotation_type=config.type,
+        description=config.description,
+    )
+
+
+def to_gql_annotation_config(annotation_config: models.AnnotationConfig) -> AnnotationConfig:
+    """
+    Convert an SQLAlchemy AnnotationConfig instance to one of the GraphQL types.
+    """
+    config = annotation_config.config
+    if isinstance(config, ContinuousAnnotationConfigModel):
+        return _to_gql_continuous_annotation_config(annotation_config)
+    if isinstance(config, CategoricalAnnotationConfigModel):
+        return _to_gql_categorical_annotation_config(annotation_config)
+    if isinstance(config, FreeformAnnotationConfigModel):
+        return _to_gql_freeform_annotation_config(annotation_config)
+    assert_never(annotation_config)

phoenix/server/api/types/AnnotationSource.py

@@ -0,0 +1,9 @@
+from enum import Enum
+
+import strawberry
+
+
+@strawberry.enum
+class AnnotationSource(Enum):
+    API = "API"
+    APP = "APP"

phoenix/server/api/types/AnnotationSummary.py

@@ -12,10 +12,9 @@ AnnotationType = Union[models.SpanAnnotation, models.TraceAnnotation]
 
 @strawberry.type
 class AnnotationSummary:
+    name: str
     df: Private[pd.DataFrame]
-
-    def __init__(self, dataframe: pd.DataFrame) -> None:
-        self.df = dataframe
+    simple_avg: Private[bool] = False
 
     @strawberry.field
     def count(self) -> int:
@@ -23,28 +22,43 @@ class AnnotationSummary:
 
     @strawberry.field
     def labels(self) -> list[str]:
-        return self.df.label.dropna().tolist()
+        unique_labels = self.df["label"].dropna().unique()
+        return [str(label) for label in unique_labels]
 
     @strawberry.field
     def label_fractions(self) -> list[LabelFraction]:
-        if not (n := self.df.label_count.sum()):
-            return []
+        if self.simple_avg:
+            if not (n := self.df.label_count.sum()):
+                return []
+            return [
+                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()
+            ]
         return [
             LabelFraction(
-                label=cast(str, row.label),
-                fraction=row.label_count / n,
+                label=row.label,
+                fraction=float(row.avg_label_fraction),
             )
-            for row in self.df.loc[
-                self.df.label.notna(),
-                ["label", "label_count"],
-            ].itertuples()
+            for row in self.df.itertuples()
+            if row.label is not None
         ]
 
     @strawberry.field
     def mean_score(self) -> Optional[float]:
-        if not (n := self.df.score_count.sum()):
+        if self.simple_avg:
+            if not (n := self.df.score_count.sum()):
+                return None
+            return cast(float, self.df.score_sum.sum() / n)
+        avg_scores = self.df["avg_score"].dropna()
+        if avg_scores.empty:
             return None
-        return cast(float, self.df.score_sum.sum() / n)
+        return float(avg_scores.mean())  # all avg_scores should be the same
 
     @strawberry.field
     def score_count(self) -> int:

phoenix/server/api/types/AnnotatorKind.py

@@ -14,3 +14,4 @@ class ExperimentRunAnnotatorKind(Enum):
 class AnnotatorKind(Enum):
     LLM = "LLM"
     HUMAN = "HUMAN"
+    CODE = "CODE"

phoenix/server/api/types/CronExpression.py

@@ -0,0 +1,15 @@
+from typing import NewType
+
+import strawberry
+
+from phoenix.db.types.trace_retention import TraceRetentionCronExpression
+
+
+def parse_value(value: str) -> str:
+    return TraceRetentionCronExpression.model_validate(value).root
+
+
+CronExpression = strawberry.scalar(
+    NewType("CronExpression", str),
+    parse_value=parse_value,
+)

phoenix/server/api/types/Evaluation.py

@@ -1,6 +1,5 @@
 import strawberry
 
-import phoenix.trace.v1 as pb
 from phoenix.db.models import DocumentAnnotation, TraceAnnotation
 
 from .Annotation import Annotation
@@ -8,19 +7,6 @@ from .Annotation import Annotation
 
 @strawberry.type
 class TraceEvaluation(Annotation):
-    @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
-        return TraceEvaluation(
-            name=evaluation.name,
-            score=score,
-            label=label,
-            explanation=explanation,
-        )
-
     @staticmethod
     def from_sql_trace_annotation(annotation: TraceAnnotation) -> "TraceEvaluation":
         return TraceEvaluation(
@@ -28,6 +14,8 @@ class TraceEvaluation(Annotation):
             score=annotation.score,
             label=annotation.label,
             explanation=annotation.explanation,
+            created_at=annotation.created_at,
+            updated_at=annotation.updated_at,
         )
 
 
@@ -38,22 +26,6 @@ class DocumentEvaluation(Annotation):
         "is collected as a list (even when ordering is not inherently meaningful)."
     )
 
-    @staticmethod
-    def from_pb_evaluation(evaluation: pb.Evaluation) -> "DocumentEvaluation":
-        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
-        document_retrieval_id = evaluation.subject_id.document_retrieval_id
-        document_position = document_retrieval_id.document_position
-        return DocumentEvaluation(
-            name=evaluation.name,
-            score=score,
-            label=label,
-            explanation=explanation,
-            document_position=document_position,
-        )
-
     @staticmethod
     def from_sql_document_annotation(annotation: DocumentAnnotation) -> "DocumentEvaluation":
         return DocumentEvaluation(
@@ -62,4 +34,6 @@ class DocumentEvaluation(Annotation):
             label=annotation.label,
             explanation=annotation.explanation,
             document_position=annotation.document_position,
+            created_at=annotation.created_at,
+            updated_at=annotation.updated_at,
         )

phoenix/server/api/types/Project.py

@@ -1,6 +1,8 @@
+from __future__ import annotations
+
 import operator
 from datetime import datetime
-from typing import Any, ClassVar, Optional
+from typing import TYPE_CHECKING, Annotated, Any, ClassVar, Optional
 
 import strawberry
 from aioitertools.itertools import islice
@@ -8,7 +10,7 @@ from openinference.semconv.trace import SpanAttributes
 from sqlalchemy import desc, distinct, func, or_, select
 from sqlalchemy.sql.elements import ColumnElement
 from sqlalchemy.sql.expression import tuple_
-from strawberry import ID, UNSET, Private
+from strawberry import ID, UNSET, Private, lazy
 from strawberry.relay import Connection, Node, NodeID
 from strawberry.types import Info
 from typing_extensions import assert_never
@@ -22,13 +24,16 @@ from phoenix.server.api.input_types.ProjectSessionSort import (
 )
 from phoenix.server.api.input_types.SpanSort import SpanSort, SpanSortConfig
 from phoenix.server.api.input_types.TimeRange import TimeRange
+from phoenix.server.api.types.AnnotationConfig import AnnotationConfig, to_gql_annotation_config
 from phoenix.server.api.types.AnnotationSummary import AnnotationSummary
 from phoenix.server.api.types.DocumentEvaluationSummary import DocumentEvaluationSummary
 from phoenix.server.api.types.pagination import (
+    ConnectionArgs,
     Cursor,
     CursorSortColumn,
     CursorString,
     connection_from_cursors_and_nodes,
+    connection_from_list,
 )
 from phoenix.server.api.types.ProjectSession import ProjectSession, to_gql_project_session
 from phoenix.server.api.types.SortDir import SortDir
@@ -38,6 +43,8 @@ from phoenix.server.api.types.ValidationResult import ValidationResult
 from phoenix.trace.dsl import SpanFilter
 
 DEFAULT_PAGE_SIZE = 30
+if TYPE_CHECKING:
+    from phoenix.server.api.types.ProjectTraceRetentionPolicy import ProjectTraceRetentionPolicy
 
 
 @strawberry.type
@@ -536,6 +543,47 @@ class Project(Node):
                 error_message=e.msg,
             )
 
+    @strawberry.field
+    async def annotation_configs(
+        self,
+        info: Info[Context, None],
+        first: Optional[int] = 50,
+        last: Optional[int] = None,
+        after: Optional[str] = None,
+        before: Optional[str] = None,
+    ) -> Connection[AnnotationConfig]:
+        args = ConnectionArgs(
+            first=first,
+            after=after if isinstance(after, CursorString) else None,
+            last=last,
+            before=before if isinstance(before, CursorString) else None,
+        )
+        async with info.context.db() as session:
+            annotation_configs = await session.stream_scalars(
+                select(models.AnnotationConfig)
+                .join(
+                    models.ProjectAnnotationConfig,
+                    models.AnnotationConfig.id
+                    == models.ProjectAnnotationConfig.annotation_config_id,
+                )
+                .where(models.ProjectAnnotationConfig.project_id == self.project_rowid)
+                .order_by(models.AnnotationConfig.name)
+            )
+            data = [to_gql_annotation_config(config) async for config in annotation_configs]
+        return connection_from_list(data=data, args=args)
+
+    @strawberry.field
+    async def trace_retention_policy(
+        self,
+        info: Info[Context, None],
+    ) -> Annotated[ProjectTraceRetentionPolicy, lazy(".ProjectTraceRetentionPolicy")]:
+        from .ProjectTraceRetentionPolicy import ProjectTraceRetentionPolicy
+
+        id_ = await info.context.data_loaders.trace_retention_policy_id_by_project_id.load(
+            self.project_rowid
+        )
+        return ProjectTraceRetentionPolicy(id=id_)
+
 
 INPUT_VALUE = SpanAttributes.INPUT_VALUE.split(".")
 OUTPUT_VALUE = SpanAttributes.OUTPUT_VALUE.split(".")