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

phoenix/server/api/types/SpanAnnotation.py

@@ -1,7 +1,8 @@
-from typing import Optional
+from datetime import datetime
+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
@@ -12,56 +13,202 @@ from phoenix.server.api.context import Context
 from .Annotation import Annotation
 from .AnnotationSource import AnnotationSource
 from .AnnotatorKind import AnnotatorKind
-from .User import User, to_gql_user
+
+if TYPE_CHECKING:
+    from .Span import Span
+    from .User import User
 
 
 @strawberry.type
 class SpanAnnotation(Node, Annotation):
-    id_attr: NodeID[int]
-    user_id: Private[Optional[int]]
-    annotator_kind: AnnotatorKind
-    metadata: JSON
-    span_rowid: Private[Optional[int]]
-    source: AnnotationSource
-    identifier: str
+    id: NodeID[int]
+    db_record: strawberry.Private[Optional[models.SpanAnnotation]] = None
+
+    def __post_init__(self) -> None:
+        if self.db_record and self.id != self.db_record.id:
+            raise ValueError("SpanAnnotation 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.span_annotation_fields.load(
+                (self.id, models.SpanAnnotation.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.span_annotation_fields.load(
+                (self.id, models.SpanAnnotation.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."
+    )  # 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.span_annotation_fields.load(
+                (self.id, models.SpanAnnotation.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.span_annotation_fields.load(
+                (self.id, models.SpanAnnotation.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."
+    )  # 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.span_annotation_fields.load(
+                (self.id, models.SpanAnnotation.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.span_annotation_fields.load(
+                (self.id, models.SpanAnnotation.metadata_),
+            )
+        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.span_annotation_fields.load(
+                (self.id, models.SpanAnnotation.source),
+            )
+        return AnnotationSource(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.span_annotation_fields.load(
+                (self.id, models.SpanAnnotation.identifier),
+            )
+        return val
+
+    @strawberry.field(description="The date and time when the annotation was created.")  # type: ignore
+    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.span_annotation_fields.load(
+                (self.id, models.SpanAnnotation.created_at),
+            )
+        return val
+
+    @strawberry.field(description="The date and time when the annotation was last updated.")  # type: ignore
+    async def updated_at(
+        self,
+        info: Info[Context, None],
+    ) -> datetime:
+        if self.db_record:
+            val = self.db_record.updated_at
+        else:
+            val = await info.context.data_loaders.span_annotation_fields.load(
+                (self.id, models.SpanAnnotation.updated_at),
+            )
+        return val
 
     @strawberry.field
-    async def span_id(self) -> GlobalID:
+    async def span_id(
+        self,
+        info: Info[Context, None],
+    ) -> GlobalID:
         from phoenix.server.api.types.Span import Span
 
-        return GlobalID(type_name=Span.__name__, node_id=str(self.span_rowid))
+        if self.db_record:
+            span_rowid = self.db_record.span_rowid
+        else:
+            span_rowid = await info.context.data_loaders.span_annotation_fields.load(
+                (self.id, models.SpanAnnotation.span_rowid),
+            )
+        return GlobalID(type_name=Span.__name__, node_id=str(span_rowid))
 
-    @strawberry.field
+    @strawberry.field(description="The span associated with the annotation.")  # type: ignore
+    async def span(
+        self,
+        info: Info[Context, None],
+    ) -> Annotated["Span", strawberry.lazy(".Span")]:
+        if self.db_record:
+            span_rowid = self.db_record.span_rowid
+        else:
+            span_rowid = await info.context.data_loaders.span_annotation_fields.load(
+                (self.id, models.SpanAnnotation.span_rowid),
+            )
+        from .Span import Span
+
+        return Span(id=span_rowid)
+
+    @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.span_annotation_fields.load(
+                (self.id, models.SpanAnnotation.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_span_annotation(
-    annotation: models.SpanAnnotation,
-) -> SpanAnnotation:
-    """
-    Converts an ORM span annotation to a GraphQL SpanAnnotation.
-    """
-    return SpanAnnotation(
-        id_attr=annotation.id,
-        user_id=annotation.user_id,
-        span_rowid=annotation.span_rowid,
-        name=annotation.name,
-        annotator_kind=AnnotatorKind(annotation.annotator_kind),
-        label=annotation.label,
-        score=annotation.score,
-        explanation=annotation.explanation,
-        metadata=annotation.metadata_,
-        source=AnnotationSource(annotation.source),
-        identifier=annotation.identifier,
-        created_at=annotation.created_at,
-        updated_at=annotation.updated_at,
-    )
+        from .User import User
+
+        return User(id=user_id)

phoenix/server/api/types/SystemApiKey.py

@@ -1,9 +1,73 @@
+from datetime import datetime
+from typing import Optional
+
 import strawberry
 from strawberry.relay import Node, NodeID
+from strawberry.types import Info
+
+from phoenix.db.models import ApiKey as OrmApiKey
+from phoenix.server.api.context import Context
 
 from .ApiKey import ApiKey
 
 
 @strawberry.type
 class SystemApiKey(ApiKey, Node):
-    id_attr: NodeID[int]
+    id: NodeID[int]
+    db_record: strawberry.Private[Optional[OrmApiKey]] = None
+
+    def __post_init__(self) -> None:
+        if self.db_record and self.id != self.db_record.id:
+            raise ValueError("SystemApiKey ID mismatch")
+
+    @strawberry.field(description="Name of the API key.")  # 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.user_api_key_fields.load(
+                (self.id, OrmApiKey.name),
+            )
+        return val
+
+    @strawberry.field(description="Description of the API key.")  # type: ignore
+    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.user_api_key_fields.load(
+                (self.id, OrmApiKey.description),
+            )
+        return val
+
+    @strawberry.field(description="The date and time the API key was created.")  # type: ignore
+    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.user_api_key_fields.load(
+                (self.id, OrmApiKey.created_at),
+            )
+        return val
+
+    @strawberry.field(description="The date and time the API key will expire.")  # type: ignore
+    async def expires_at(
+        self,
+        info: Info[Context, None],
+    ) -> Optional[datetime]:
+        if self.db_record:
+            val = self.db_record.expires_at
+        else:
+            val = await info.context.data_loaders.user_api_key_fields.load(
+                (self.id, OrmApiKey.expires_at),
+            )
+        return val