arize-phoenix 12.8.0__py3-none-any.whl → 12.9.0__py3-none-any.whl

phoenix/server/api/types/ProjectSession.py

@@ -1,13 +1,13 @@
 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 import UNSET, Info, lazy
 from strawberry.relay import Connection, Node, NodeID
 
 from phoenix.db import models
@@ -30,12 +30,51 @@ if TYPE_CHECKING:
 
 @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 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(
@@ -44,28 +83,34 @@ class ProjectSession(Node):
     ) -> Annotated["Project", lazy(".Project")]:
         from phoenix.server.api.types.Project import Project
 
-        return Project(project_rowid=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(
@@ -78,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(
@@ -91,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,
@@ -116,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
@@ -131,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
@@ -140,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,
@@ -162,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,
@@ -181,15 +226,14 @@ class ProjectSession(Node):
         info: Info[Context, None],
     ) -> list[Annotated["ProjectSessionAnnotation", lazy(".ProjectSessionAnnotation")]]:
         """Get all annotations for this session."""
-        from phoenix.server.api.types.ProjectSessionAnnotation import (
-            to_gql_project_session_annotation,
-        )
+        from .ProjectSessionAnnotation import ProjectSessionAnnotation
 
-        stmt = select(models.ProjectSessionAnnotation).filter_by(project_session_id=self.id_attr)
+        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 [
-                to_gql_project_session_annotation(annotation) async for annotation in annotations
+                ProjectSessionAnnotation(id=annotation.id, db_record=annotation)
+                async for annotation in annotations
             ]
 
     @strawberry.field(
@@ -217,9 +261,7 @@ class ProjectSession(Node):
             - 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_attr
-        )
+        annotations = await info.context.data_loaders.session_annotations_by_session.load(self.id)
 
         # Apply filter if provided to narrow down the annotations
         if filter:
@@ -251,16 +293,6 @@ class ProjectSession(Node):
         return result
 
 
-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,
-    )
-
-
 INPUT_VALUE = SpanAttributes.INPUT_VALUE.split(".")
 INPUT_MIME_TYPE = SpanAttributes.INPUT_MIME_TYPE.split(".")
 OUTPUT_VALUE = SpanAttributes.OUTPUT_VALUE.split(".")

phoenix/server/api/types/ProjectSessionAnnotation.py

@@ -1,7 +1,7 @@
-from typing import Optional
+from math import isfinite
+from typing import TYPE_CHECKING, Annotated, Optional
 
 import strawberry
-from strawberry import Private
 from strawberry.relay import GlobalID, Node, NodeID
 from strawberry.scalars import JSON
 from strawberry.types import Info
@@ -10,59 +10,178 @@ 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
-from .User import User, to_gql_user
+
+if TYPE_CHECKING:
+    from .ProjectSession import ProjectSession
+    from .User import User
 
 
 @strawberry.type
-class ProjectSessionAnnotation(Node):
-    id_attr: NodeID[int]
-    user_id: Private[Optional[int]]
-    name: str
-    annotator_kind: AnnotatorKind
-    label: Optional[str]
-    score: Optional[float]
-    explanation: Optional[str]
-    metadata: JSON
-    _project_session_id: Private[Optional[int]]
-    identifier: str
-    source: AnnotationSource
-
-    @strawberry.field
-    async def project_session_id(self) -> GlobalID:
+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
 
-        return GlobalID(type_name=ProjectSession.__name__, node_id=str(self._project_session_id))
+        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
+    @strawberry.field(description="The user that produced the annotation.")  # type: ignore
     async def user(
         self,
         info: Info[Context, None],
-    ) -> Optional[User]:
-        if self.user_id is 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
-        user = await info.context.data_loaders.users.load(self.user_id)
-        if user is None:
-            return None
-        return to_gql_user(user)
-
-
-def to_gql_project_session_annotation(
-    annotation: models.ProjectSessionAnnotation,
-) -> ProjectSessionAnnotation:
-    """
-    Converts an ORM projectSession annotation to a GraphQL ProjectSessionAnnotation.
-    """
-    return ProjectSessionAnnotation(
-        id_attr=annotation.id,
-        user_id=annotation.user_id,
-        _project_session_id=annotation.project_session_id,
-        name=annotation.name,
-        annotator_kind=AnnotatorKind(annotation.annotator_kind),
-        label=annotation.label,
-        score=annotation.score,
-        explanation=annotation.explanation,
-        metadata=JSON(annotation.metadata_),
-        identifier=annotation.identifier,
-        source=AnnotationSource(annotation.source),
-    )
+        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

@@ -20,21 +20,76 @@ from phoenix.server.api.types.pagination import (
     connection_from_list,
 )
 
-from .PromptLabel import PromptLabel, to_gql_prompt_label
+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 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(
@@ -49,7 +104,7 @@ 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:
@@ -61,7 +116,7 @@ class Prompt(Node):
                     raise NotFound(f"Prompt version tag not found: {tag_name}")
                 version = await session.scalar(
                     select(models.PromptVersion)
-                    .where(models.PromptVersion.prompt_id == self.id_attr)
+                    .where(models.PromptVersion.prompt_id == self.id)
                     .join_from(models.PromptVersion, models.PromptVersionTag)
                     .where(models.PromptVersionTag.name == name)
                 )
@@ -70,7 +125,7 @@ class Prompt(Node):
             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)
                 )
@@ -83,10 +138,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
@@ -107,7 +163,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:
@@ -119,20 +175,19 @@ 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
-
-        source_prompt_id = from_global_id_with_expected_type(
-            global_id=self.source_prompt_id, expected_type_name=Prompt.__name__
-        )
-
         async with info.context.db() as session:
-            source_prompt = await session.scalar(
-                select(models.Prompt).where(models.Prompt.id == source_prompt_id)
-            )
-            if not source_prompt:
-                raise NotFound(f"Source prompt not found: {self.source_prompt_id}")
-            return to_gql_prompt_from_orm(source_prompt)
+            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)
 
     @strawberry.field
     async def labels(self, info: Info[Context, None]) -> list["PromptLabel"]:
@@ -140,23 +195,6 @@ class Prompt(Node):
             labels = await session.scalars(
                 select(models.PromptLabel)
                 .join(models.PromptPromptLabel)
-                .where(models.PromptPromptLabel.prompt_id == self.id_attr)
+                .where(models.PromptPromptLabel.prompt_id == self.id)
             )
-            return [to_gql_prompt_label(label) for label in labels]
-
-
-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]