arize-phoenix 11.23.1__py3-none-any.whl → 12.28.1__py3-none-any.whl

phoenix/server/api/types/ProjectSession.py

@@ -1,14 +1,19 @@
+from collections import defaultdict
+from dataclasses import asdict, dataclass
 from datetime import datetime
-from typing import TYPE_CHECKING, Annotated, ClassVar, Optional, Type
+from typing import TYPE_CHECKING, Annotated, Optional
 
+import pandas as pd
 import strawberry
 from openinference.semconv.trace import SpanAttributes
 from sqlalchemy import select
-from strawberry import UNSET, Info, Private, lazy
-from strawberry.relay import Connection, GlobalID, Node, NodeID
+from strawberry import UNSET, Info, lazy
+from strawberry.relay import Connection, Node, NodeID
 
 from phoenix.db import models
 from phoenix.server.api.context import Context
+from phoenix.server.api.input_types.AnnotationFilter import AnnotationFilter, satisfies_filter
+from phoenix.server.api.types.AnnotationSummary import AnnotationSummary
 from phoenix.server.api.types.CostBreakdown import CostBreakdown
 from phoenix.server.api.types.MimeType import MimeType
 from phoenix.server.api.types.pagination import ConnectionArgs, CursorString, connection_from_list
@@ -18,44 +23,94 @@ from phoenix.server.api.types.SpanIOValue import SpanIOValue
 from phoenix.server.api.types.TokenUsage import TokenUsage
 
 if TYPE_CHECKING:
+    from phoenix.server.api.types.Project import Project
+    from phoenix.server.api.types.ProjectSessionAnnotation import ProjectSessionAnnotation
     from phoenix.server.api.types.Trace import Trace
 
 
 @strawberry.type
 class ProjectSession(Node):
-    _table: ClassVar[Type[models.ProjectSession]] = models.ProjectSession
-    id_attr: NodeID[int]
-    project_rowid: Private[int]
-    session_id: str
-    start_time: datetime
-    end_time: datetime
+    id: NodeID[int]
+    db_record: strawberry.Private[Optional[models.ProjectSession]] = None
+
+    def __post_init__(self) -> None:
+        if self.db_record and self.id != self.db_record.id:
+            raise ValueError("ProjectSession ID mismatch")
 
     @strawberry.field
-    async def project_id(self) -> GlobalID:
+    async def session_id(
+        self,
+        info: Info[Context, None],
+    ) -> str:
+        if self.db_record:
+            val = self.db_record.session_id
+        else:
+            val = await info.context.data_loaders.project_session_fields.load(
+                (self.id, models.ProjectSession.session_id),
+            )
+        return val
+
+    @strawberry.field
+    async def start_time(
+        self,
+        info: Info[Context, None],
+    ) -> datetime:
+        if self.db_record:
+            val = self.db_record.start_time
+        else:
+            val = await info.context.data_loaders.project_session_fields.load(
+                (self.id, models.ProjectSession.start_time),
+            )
+        return val
+
+    @strawberry.field
+    async def end_time(
+        self,
+        info: Info[Context, None],
+    ) -> datetime:
+        if self.db_record:
+            val = self.db_record.end_time
+        else:
+            val = await info.context.data_loaders.project_session_fields.load(
+                (self.id, models.ProjectSession.end_time),
+            )
+        return val
+
+    @strawberry.field
+    async def project(
+        self,
+        info: Info[Context, None],
+    ) -> Annotated["Project", lazy(".Project")]:
         from phoenix.server.api.types.Project import Project
 
-        return GlobalID(type_name=Project.__name__, node_id=str(self.project_rowid))
+        if self.db_record:
+            project_rowid = self.db_record.project_id
+        else:
+            project_rowid = await info.context.data_loaders.project_session_fields.load(
+                (self.id, models.ProjectSession.project_id),
+            )
+        return Project(id=project_rowid)
 
     @strawberry.field
     async def num_traces(
         self,
         info: Info[Context, None],
     ) -> int:
-        return await info.context.data_loaders.session_num_traces.load(self.id_attr)
+        return await info.context.data_loaders.session_num_traces.load(self.id)
 
     @strawberry.field
     async def num_traces_with_error(
         self,
         info: Info[Context, None],
     ) -> int:
-        return await info.context.data_loaders.session_num_traces_with_error.load(self.id_attr)
+        return await info.context.data_loaders.session_num_traces_with_error.load(self.id)
 
     @strawberry.field
     async def first_input(
         self,
         info: Info[Context, None],
     ) -> Optional[SpanIOValue]:
-        record = await info.context.data_loaders.session_first_inputs.load(self.id_attr)
+        record = await info.context.data_loaders.session_first_inputs.load(self.id)
         if record is None:
             return None
         return SpanIOValue(
@@ -68,7 +123,7 @@ class ProjectSession(Node):
         self,
         info: Info[Context, None],
     ) -> Optional[SpanIOValue]:
-        record = await info.context.data_loaders.session_last_outputs.load(self.id_attr)
+        record = await info.context.data_loaders.session_last_outputs.load(self.id)
         if record is None:
             return None
         return SpanIOValue(
@@ -81,7 +136,7 @@ class ProjectSession(Node):
         self,
         info: Info[Context, None],
     ) -> TokenUsage:
-        usage = await info.context.data_loaders.session_token_usages.load(self.id_attr)
+        usage = await info.context.data_loaders.session_token_usages.load(self.id)
         return TokenUsage(
             prompt=usage.prompt,
             completion=usage.completion,
@@ -106,12 +161,12 @@ class ProjectSession(Node):
         )
         stmt = (
             select(models.Trace)
-            .filter_by(project_session_rowid=self.id_attr)
+            .filter_by(project_session_rowid=self.id)
             .order_by(models.Trace.start_time)
         )
         async with info.context.db() as session:
             traces = await session.stream_scalars(stmt)
-            data = [Trace(trace_rowid=trace.id, db_trace=trace) async for trace in traces]
+            data = [Trace(id=trace.id, db_record=trace) async for trace in traces]
         return connection_from_list(data=data, args=args)
 
     @strawberry.field
@@ -121,7 +176,7 @@ class ProjectSession(Node):
         probability: float,
     ) -> Optional[float]:
         return await info.context.data_loaders.session_trace_latency_ms_quantile.load(
-            (self.id_attr, probability)
+            (self.id, probability)
         )
 
     @strawberry.field
@@ -130,7 +185,7 @@ class ProjectSession(Node):
         info: Info[Context, None],
     ) -> SpanCostSummary:
         loader = info.context.data_loaders.span_cost_summary_by_project_session
-        summary = await loader.load(self.id_attr)
+        summary = await loader.load(self.id)
         return SpanCostSummary(
             prompt=CostBreakdown(
                 tokens=summary.prompt.tokens,
@@ -152,7 +207,7 @@ class ProjectSession(Node):
         info: Info[Context, None],
     ) -> list[SpanCostDetailSummaryEntry]:
         loader = info.context.data_loaders.span_cost_detail_summary_entries_by_project_session
-        summary = await loader.load(self.id_attr)
+        summary = await loader.load(self.id)
         return [
             SpanCostDetailSummaryEntry(
                 token_type=entry.token_type,
@@ -165,15 +220,77 @@ class ProjectSession(Node):
             for entry in summary
         ]
 
+    @strawberry.field
+    async def session_annotations(
+        self,
+        info: Info[Context, None],
+    ) -> list[Annotated["ProjectSessionAnnotation", lazy(".ProjectSessionAnnotation")]]:
+        """Get all annotations for this session."""
+        from .ProjectSessionAnnotation import ProjectSessionAnnotation
+
+        stmt = select(models.ProjectSessionAnnotation).filter_by(project_session_id=self.id)
+        async with info.context.db() as session:
+            annotations = await session.stream_scalars(stmt)
+            return [
+                ProjectSessionAnnotation(id=annotation.id, db_record=annotation)
+                async for annotation in annotations
+            ]
+
+    @strawberry.field(
+        description="Summarizes each annotation (by name) associated with the session"
+    )  # type: ignore
+    async def session_annotation_summaries(
+        self,
+        info: Info[Context, None],
+        filter: Optional[AnnotationFilter] = None,
+    ) -> list[AnnotationSummary]:
+        """
+        Retrieves and summarizes annotations associated with this span.
+
+        This method aggregates annotation data by name and label, calculating metrics
+        such as count of occurrences and sum of scores. The results are organized
+        into a structured format that can be easily converted to a DataFrame.
+
+        Args:
+            info: GraphQL context information
+            filter: Optional filter to apply to annotations before processing
+
+        Returns:
+            A list of AnnotationSummary objects, each containing:
+            - name: The name of the annotation
+            - data: A list of dictionaries with label statistics
+        """
+        # Load all annotations for this span from the data loader
+        annotations = await info.context.data_loaders.session_annotations_by_session.load(self.id)
+
+        # Apply filter if provided to narrow down the annotations
+        if filter:
+            annotations = [
+                annotation for annotation in annotations if satisfies_filter(annotation, filter)
+            ]
+
+        @dataclass
+        class Metrics:
+            record_count: int = 0
+            label_count: int = 0
+            score_sum: float = 0
+            score_count: int = 0
+
+        summaries: defaultdict[str, defaultdict[Optional[str], Metrics]] = defaultdict(
+            lambda: defaultdict(Metrics)
+        )
+        for annotation in annotations:
+            metrics = summaries[annotation.name][annotation.label]
+            metrics.record_count += 1
+            metrics.label_count += int(annotation.label is not None)
+            metrics.score_sum += annotation.score or 0
+            metrics.score_count += int(annotation.score is not None)
 
-def to_gql_project_session(project_session: models.ProjectSession) -> ProjectSession:
-    return ProjectSession(
-        id_attr=project_session.id,
-        session_id=project_session.session_id,
-        start_time=project_session.start_time,
-        project_rowid=project_session.project_id,
-        end_time=project_session.end_time,
-    )
+        result: list[AnnotationSummary] = []
+        for name, label_metrics in summaries.items():
+            rows = [{"label": label, **asdict(metrics)} for label, metrics in label_metrics.items()]
+            result.append(AnnotationSummary(name=name, df=pd.DataFrame(rows), simple_avg=True))
+        return result
 
 
 INPUT_VALUE = SpanAttributes.INPUT_VALUE.split(".")

phoenix/server/api/types/ProjectSessionAnnotation.py

@@ -0,0 +1,187 @@
+from math import isfinite
+from typing import TYPE_CHECKING, Annotated, Optional
+
+import strawberry
+from strawberry.relay import GlobalID, Node, NodeID
+from strawberry.scalars import JSON
+from strawberry.types import Info
+
+from phoenix.db import models
+from phoenix.server.api.context import Context
+from phoenix.server.api.types.AnnotatorKind import AnnotatorKind
+
+from .Annotation import Annotation
+from .AnnotationSource import AnnotationSource
+
+if TYPE_CHECKING:
+    from .ProjectSession import ProjectSession
+    from .User import User
+
+
+@strawberry.type
+class ProjectSessionAnnotation(Node, Annotation):
+    id: NodeID[int]
+    db_record: strawberry.Private[Optional[models.ProjectSessionAnnotation]] = None
+
+    def __post_init__(self) -> None:
+        if self.db_record and self.id != self.db_record.id:
+            raise ValueError("ProjectSessionAnnotation ID mismatch")
+
+    @strawberry.field(description="Name of the annotation, e.g. 'helpfulness' or 'relevance'.")  # type: ignore
+    async def name(
+        self,
+        info: Info[Context, None],
+    ) -> str:
+        if self.db_record:
+            val = self.db_record.name
+        else:
+            val = await info.context.data_loaders.project_session_annotation_fields.load(
+                (self.id, models.ProjectSessionAnnotation.name),
+            )
+        return val
+
+    @strawberry.field(description="The kind of annotator that produced the annotation.")  # type: ignore
+    async def annotator_kind(
+        self,
+        info: Info[Context, None],
+    ) -> AnnotatorKind:
+        if self.db_record:
+            val = self.db_record.annotator_kind
+        else:
+            val = await info.context.data_loaders.project_session_annotation_fields.load(
+                (self.id, models.ProjectSessionAnnotation.annotator_kind),
+            )
+        return AnnotatorKind(val)
+
+    @strawberry.field(
+        description="Value of the annotation in the form of a string, e.g. 'helpful' or 'not helpful'. Note that the label is not necessarily binary."  # noqa: E501
+    )  # type: ignore
+    async def label(
+        self,
+        info: Info[Context, None],
+    ) -> Optional[str]:
+        if self.db_record:
+            val = self.db_record.label
+        else:
+            val = await info.context.data_loaders.project_session_annotation_fields.load(
+                (self.id, models.ProjectSessionAnnotation.label),
+            )
+        return val
+
+    @strawberry.field(description="Value of the annotation in the form of a numeric score.")  # type: ignore
+    async def score(
+        self,
+        info: Info[Context, None],
+    ) -> Optional[float]:
+        if self.db_record:
+            val = self.db_record.score
+        else:
+            val = await info.context.data_loaders.project_session_annotation_fields.load(
+                (self.id, models.ProjectSessionAnnotation.score),
+            )
+        return val if val is not None and isfinite(val) else None
+
+    @strawberry.field(
+        description="The annotator's explanation for the annotation result (i.e. score or label, or both) given to the subject."  # noqa: E501
+    )  # type: ignore
+    async def explanation(
+        self,
+        info: Info[Context, None],
+    ) -> Optional[str]:
+        if self.db_record:
+            val = self.db_record.explanation
+        else:
+            val = await info.context.data_loaders.project_session_annotation_fields.load(
+                (self.id, models.ProjectSessionAnnotation.explanation),
+            )
+        return val
+
+    @strawberry.field(description="Metadata about the annotation.")  # type: ignore
+    async def metadata(
+        self,
+        info: Info[Context, None],
+    ) -> JSON:
+        if self.db_record:
+            val = self.db_record.metadata_
+        else:
+            val = await info.context.data_loaders.project_session_annotation_fields.load(
+                (self.id, models.ProjectSessionAnnotation.metadata_),
+            )
+        return val
+
+    @strawberry.field(description="The identifier of the annotation.")  # type: ignore
+    async def identifier(
+        self,
+        info: Info[Context, None],
+    ) -> str:
+        if self.db_record:
+            val = self.db_record.identifier
+        else:
+            val = await info.context.data_loaders.project_session_annotation_fields.load(
+                (self.id, models.ProjectSessionAnnotation.identifier),
+            )
+        return val
+
+    @strawberry.field(description="The source of the annotation.")  # type: ignore
+    async def source(
+        self,
+        info: Info[Context, None],
+    ) -> AnnotationSource:
+        if self.db_record:
+            val = self.db_record.source
+        else:
+            val = await info.context.data_loaders.project_session_annotation_fields.load(
+                (self.id, models.ProjectSessionAnnotation.source),
+            )
+        return AnnotationSource(val)
+
+    @strawberry.field(description="The project session associated with the annotation.")  # type: ignore
+    async def project_session_id(
+        self,
+        info: Info[Context, None],
+    ) -> GlobalID:
+        from phoenix.server.api.types.ProjectSession import ProjectSession
+
+        if self.db_record:
+            project_session_id = self.db_record.project_session_id
+        else:
+            project_session_id = (
+                await info.context.data_loaders.project_session_annotation_fields.load(
+                    (self.id, models.ProjectSessionAnnotation.project_session_id),
+                )
+            )
+        return GlobalID(type_name=ProjectSession.__name__, node_id=str(project_session_id))
+
+    @strawberry.field(description="The project session associated with the annotation.")  # type: ignore
+    async def project_session(
+        self,
+        info: Info[Context, None],
+    ) -> Annotated["ProjectSession", strawberry.lazy(".ProjectSession")]:
+        if self.db_record:
+            project_session_id = self.db_record.project_session_id
+        else:
+            project_session_id = (
+                await info.context.data_loaders.project_session_annotation_fields.load(
+                    (self.id, models.ProjectSessionAnnotation.project_session_id),
+                )
+            )
+        from .ProjectSession import ProjectSession
+
+        return ProjectSession(id=project_session_id)
+
+    @strawberry.field(description="The user that produced the annotation.")  # type: ignore
+    async def user(
+        self,
+        info: Info[Context, None],
+    ) -> Optional[Annotated["User", strawberry.lazy(".User")]]:
+        if self.db_record:
+            user_id = self.db_record.user_id
+        else:
+            user_id = await info.context.data_loaders.project_session_annotation_fields.load(
+                (self.id, models.ProjectSessionAnnotation.user_id),
+            )
+        if user_id is None:
+            return None
+        from .User import User
+
+        return User(id=user_id)

phoenix/server/api/types/ProjectTraceRetentionPolicy.py

@@ -106,5 +106,5 @@ class ProjectTraceRetentionPolicy(Node):
         project_rowids = await info.context.data_loaders.projects_by_trace_retention_policy_id.load(
             self.id
         )
-        data = [Project(project_rowid=project_rowid) for project_rowid in project_rowids]
+        data = [Project(id=project_rowid) for project_rowid in project_rowids]
         return connection_from_list(data=data, args=args)

phoenix/server/api/types/Prompt.py

@@ -6,9 +6,11 @@ import strawberry
 from sqlalchemy import func, select
 from strawberry import UNSET
 from strawberry.relay import Connection, GlobalID, Node, NodeID
+from strawberry.scalars import JSON
 from strawberry.types import Info
 
 from phoenix.db import models
+from phoenix.db.types.identifier import Identifier as IdentifierModel
 from phoenix.server.api.context import Context
 from phoenix.server.api.exceptions import NotFound
 from phoenix.server.api.types.Identifier import Identifier
@@ -19,24 +21,96 @@ from phoenix.server.api.types.pagination import (
     connection_from_list,
 )
 
+from .PromptLabel import PromptLabel
 from .PromptVersion import (
     PromptVersion,
     to_gql_prompt_version,
 )
-from .PromptVersionTag import PromptVersionTag, to_gql_prompt_version_tag
+from .PromptVersionTag import PromptVersionTag
 
 
 @strawberry.type
 class Prompt(Node):
-    id_attr: NodeID[int]
-    source_prompt_id: Optional[GlobalID]
-    name: Identifier
-    description: Optional[str]
-    created_at: datetime
+    id: NodeID[int]
+    db_record: strawberry.Private[Optional[models.Prompt]] = None
+
+    def __post_init__(self) -> None:
+        if self.db_record and self.id != self.db_record.id:
+            raise ValueError("Prompt ID mismatch")
+
+    @strawberry.field
+    async def source_prompt_id(
+        self,
+        info: Info[Context, None],
+    ) -> Optional[GlobalID]:
+        if self.db_record:
+            source_id = self.db_record.source_prompt_id
+        else:
+            source_id = await info.context.data_loaders.prompt_fields.load(
+                (self.id, models.Prompt.source_prompt_id),
+            )
+        if not source_id:
+            return None
+        return GlobalID(Prompt.__name__, str(source_id))
+
+    @strawberry.field
+    async def name(
+        self,
+        info: Info[Context, None],
+    ) -> Identifier:
+        if self.db_record:
+            val = self.db_record.name
+        else:
+            val = await info.context.data_loaders.prompt_fields.load(
+                (self.id, models.Prompt.name),
+            )
+        return Identifier(val.root)
+
+    @strawberry.field
+    async def description(
+        self,
+        info: Info[Context, None],
+    ) -> Optional[str]:
+        if self.db_record:
+            val = self.db_record.description
+        else:
+            val = await info.context.data_loaders.prompt_fields.load(
+                (self.id, models.Prompt.description),
+            )
+        return val
+
+    @strawberry.field
+    async def metadata(
+        self,
+        info: Info[Context, None],
+    ) -> JSON:
+        if self.db_record:
+            val = self.db_record.metadata_
+        else:
+            val = await info.context.data_loaders.prompt_fields.load(
+                (self.id, models.Prompt.metadata_),
+            )
+        return val
+
+    @strawberry.field
+    async def created_at(
+        self,
+        info: Info[Context, None],
+    ) -> datetime:
+        if self.db_record:
+            val = self.db_record.created_at
+        else:
+            val = await info.context.data_loaders.prompt_fields.load(
+                (self.id, models.Prompt.created_at),
+            )
+        return val
 
     @strawberry.field
     async def version(
-        self, info: Info[Context, None], version_id: Optional[GlobalID] = None
+        self,
+        info: Info[Context, None],
+        version_id: Optional[GlobalID] = None,
+        tag_name: Optional[Identifier] = None,
     ) -> PromptVersion:
         async with info.context.db() as session:
             if version_id:
@@ -44,15 +118,28 @@ class Prompt(Node):
                 version = await session.scalar(
                     select(models.PromptVersion).where(
                         models.PromptVersion.id == v_id,
-                        models.PromptVersion.prompt_id == self.id_attr,
+                        models.PromptVersion.prompt_id == self.id,
                     )
                 )
                 if not version:
                     raise NotFound(f"Prompt version not found: {version_id}")
+            elif tag_name:
+                try:
+                    name = IdentifierModel(tag_name)
+                except ValueError:
+                    raise NotFound(f"Prompt version tag not found: {tag_name}")
+                version = await session.scalar(
+                    select(models.PromptVersion)
+                    .where(models.PromptVersion.prompt_id == self.id)
+                    .join_from(models.PromptVersion, models.PromptVersionTag)
+                    .where(models.PromptVersionTag.name == name)
+                )
+                if not version:
+                    raise NotFound(f"This prompt has no associated versions by tag {tag_name}")
             else:
                 stmt = (
                     select(models.PromptVersion)
-                    .where(models.PromptVersion.prompt_id == self.id_attr)
+                    .where(models.PromptVersion.prompt_id == self.id)
                     .order_by(models.PromptVersion.id.desc())
                     .limit(1)
                 )
@@ -65,10 +152,11 @@ class Prompt(Node):
     async def version_tags(self, info: Info[Context, None]) -> list[PromptVersionTag]:
         async with info.context.db() as session:
             stmt = select(models.PromptVersionTag).where(
-                models.PromptVersionTag.prompt_id == self.id_attr
+                models.PromptVersionTag.prompt_id == self.id
             )
             return [
-                to_gql_prompt_version_tag(tag) async for tag in await session.stream_scalars(stmt)
+                PromptVersionTag(id=tag.id, db_record=tag)
+                async for tag in await session.stream_scalars(stmt)
             ]
 
     @strawberry.field
@@ -89,7 +177,7 @@ class Prompt(Node):
         row_number = func.row_number().over(order_by=models.PromptVersion.id).label("row_number")
         stmt = (
             select(models.PromptVersion, row_number)
-            .where(models.PromptVersion.prompt_id == self.id_attr)
+            .where(models.PromptVersion.prompt_id == self.id)
             .order_by(models.PromptVersion.id.desc())
         )
         async with info.context.db() as session:
@@ -101,34 +189,26 @@ class Prompt(Node):
 
     @strawberry.field
     async def source_prompt(self, info: Info[Context, None]) -> Optional["Prompt"]:
-        if not self.source_prompt_id:
+        if self.db_record:
+            id_ = self.db_record.source_prompt_id
+        else:
+            id_ = await info.context.data_loaders.prompt_fields.load(
+                (self.id, models.Prompt.source_prompt_id),
+            )
+        if not id_:
             return None
+        async with info.context.db() as session:
+            source_prompt = await session.get(models.Prompt, id_)
+        if not source_prompt:
+            raise NotFound(f"Source prompt not found: {id_}")
+        return Prompt(id=source_prompt.id, db_record=source_prompt)
 
-        source_prompt_id = from_global_id_with_expected_type(
-            global_id=self.source_prompt_id, expected_type_name=Prompt.__name__
-        )
-
+    @strawberry.field
+    async def labels(self, info: Info[Context, None]) -> list["PromptLabel"]:
         async with info.context.db() as session:
-            source_prompt = await session.scalar(
-                select(models.Prompt).where(models.Prompt.id == source_prompt_id)
+            labels = await session.scalars(
+                select(models.PromptLabel)
+                .join(models.PromptPromptLabel)
+                .where(models.PromptPromptLabel.prompt_id == self.id)
             )
-            if not source_prompt:
-                raise NotFound(f"Source prompt not found: {self.source_prompt_id}")
-            return to_gql_prompt_from_orm(source_prompt)
-
-
-def to_gql_prompt_from_orm(orm_model: "models.Prompt") -> Prompt:
-    if not orm_model.source_prompt_id:
-        source_prompt_gid = None
-    else:
-        source_prompt_gid = GlobalID(
-            Prompt.__name__,
-            str(orm_model.source_prompt_id),
-        )
-    return Prompt(
-        id_attr=orm_model.id,
-        source_prompt_id=source_prompt_gid,
-        name=Identifier(orm_model.name.root),
-        description=orm_model.description,
-        created_at=orm_model.created_at,
-    )
+            return [PromptLabel(id=label.id, db_record=label) for label in labels]