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

phoenix/server/api/types/TraceAnnotation.py

@@ -1,7 +1,7 @@
+from math import isfinite
 from typing import TYPE_CHECKING, Annotated, Optional
 
 import strawberry
-from strawberry import Private
 from strawberry.relay import Node, NodeID
 from strawberry.scalars import JSON
 from strawberry.types import Info
@@ -11,61 +11,157 @@ from phoenix.server.api.context import Context
 from phoenix.server.api.types.AnnotatorKind import AnnotatorKind
 
 from .AnnotationSource import AnnotationSource
-from .User import User, to_gql_user
 
 if TYPE_CHECKING:
     from .Trace import Trace
+    from .User import User
 
 
 @strawberry.type
 class TraceAnnotation(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
-    trace_rowid: Private[int]
-    identifier: str
-    source: AnnotationSource
-
-    @strawberry.field
-    async def trace(self) -> Annotated["Trace", strawberry.lazy(".Trace")]:
-        from phoenix.server.api.types.Trace import Trace
-
-        return Trace(trace_rowid=self.trace_rowid)
-
-    @strawberry.field
+    id: NodeID[int]
+    db_record: strawberry.Private[Optional[models.TraceAnnotation]] = None
+
+    def __post_init__(self) -> None:
+        if self.db_record and self.id != self.db_record.id:
+            raise ValueError("TraceAnnotation 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.trace_annotation_fields.load(
+                (self.id, models.TraceAnnotation.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.trace_annotation_fields.load(
+                (self.id, models.TraceAnnotation.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.trace_annotation_fields.load(
+                (self.id, models.TraceAnnotation.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.trace_annotation_fields.load(
+                (self.id, models.TraceAnnotation.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.trace_annotation_fields.load(
+                (self.id, models.TraceAnnotation.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.trace_annotation_fields.load(
+                (self.id, models.TraceAnnotation.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.trace_annotation_fields.load(
+                (self.id, models.TraceAnnotation.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.trace_annotation_fields.load(
+                (self.id, models.TraceAnnotation.source),
+            )
+        return AnnotationSource(val)
+
+    @strawberry.field(description="The trace associated with the annotation.")  # type: ignore
+    async def trace(
+        self,
+        info: Info[Context, None],
+    ) -> Annotated["Trace", strawberry.lazy(".Trace")]:
+        if self.db_record:
+            trace_rowid = self.db_record.trace_rowid
+        else:
+            trace_rowid = await info.context.data_loaders.trace_annotation_fields.load(
+                (self.id, models.TraceAnnotation.trace_rowid),
+            )
+        from .Trace import Trace
+
+        return Trace(id=trace_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.trace_annotation_fields.load(
+                (self.id, models.TraceAnnotation.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_trace_annotation(
-    annotation: models.TraceAnnotation,
-) -> TraceAnnotation:
-    """
-    Converts an ORM trace annotation to a GraphQL TraceAnnotation.
-    """
-    return TraceAnnotation(
-        id_attr=annotation.id,
-        user_id=annotation.user_id,
-        trace_rowid=annotation.trace_rowid,
-        name=annotation.name,
-        annotator_kind=AnnotatorKind(annotation.annotator_kind),
-        label=annotation.label,
-        score=annotation.score,
-        explanation=annotation.explanation,
-        metadata=annotation.metadata_,
-        identifier=annotation.identifier,
-        source=AnnotationSource(annotation.source),
-    )
+        from .User import User
+
+        return User(id=user_id)

phoenix/server/api/types/User.py

@@ -3,7 +3,6 @@ from typing import Optional
 
 import strawberry
 from sqlalchemy import select
-from strawberry import Private
 from strawberry.relay import Node, NodeID
 from strawberry.types import Info
 
@@ -12,59 +11,130 @@ from phoenix.db import models
 from phoenix.server.api.context import Context
 from phoenix.server.api.exceptions import NotFound
 from phoenix.server.api.types.AuthMethod import AuthMethod
-from phoenix.server.api.types.UserApiKey import UserApiKey, to_gql_api_key
+from phoenix.server.api.types.UserApiKey import UserApiKey
 
 from .UserRole import UserRole, to_gql_user_role
 
 
 @strawberry.type
 class User(Node):
-    id_attr: NodeID[int]
-    password_needs_reset: bool
-    email: str
-    username: str
-    profile_picture_url: Optional[str]
-    created_at: datetime
-    user_role_id: Private[int]
-    auth_method: AuthMethod
+    id: NodeID[int]
+    db_record: strawberry.Private[Optional[models.User]] = None
+
+    def __post_init__(self) -> None:
+        if self.db_record and self.id != self.db_record.id:
+            raise ValueError("User ID mismatch")
+
+    @strawberry.field
+    async def password_needs_reset(
+        self,
+        info: Info[Context, None],
+    ) -> bool:
+        if self.db_record:
+            val = self.db_record.reset_password
+        else:
+            val = await info.context.data_loaders.user_fields.load(
+                (self.id, models.User.reset_password),
+            )
+        return val
+
+    @strawberry.field
+    async def email(
+        self,
+        info: Info[Context, None],
+    ) -> str:
+        if self.db_record:
+            val = self.db_record.email
+        else:
+            val = await info.context.data_loaders.user_fields.load(
+                (self.id, models.User.email),
+            )
+        return val
+
+    @strawberry.field
+    async def username(
+        self,
+        info: Info[Context, None],
+    ) -> str:
+        if self.db_record:
+            val = self.db_record.username
+        else:
+            val = await info.context.data_loaders.user_fields.load(
+                (self.id, models.User.username),
+            )
+        return val
+
+    @strawberry.field
+    async def profile_picture_url(
+        self,
+        info: Info[Context, None],
+    ) -> Optional[str]:
+        if self.db_record:
+            val = self.db_record.profile_picture_url
+        else:
+            val = await info.context.data_loaders.user_fields.load(
+                (self.id, models.User.profile_picture_url),
+            )
+        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.user_fields.load(
+                (self.id, models.User.created_at),
+            )
+        return val
+
+    @strawberry.field
+    async def auth_method(
+        self,
+        info: Info[Context, None],
+    ) -> AuthMethod:
+        if self.db_record:
+            val = self.db_record.auth_method
+        else:
+            val = await info.context.data_loaders.user_fields.load(
+                (self.id, models.User.auth_method),
+            )
+        return AuthMethod(val)
 
     @strawberry.field
     async def role(self, info: Info[Context, None]) -> UserRole:
-        role = await info.context.data_loaders.user_roles.load(self.user_role_id)
+        if self.db_record:
+            user_role_id = self.db_record.user_role_id
+        else:
+            user_role_id = await info.context.data_loaders.user_fields.load(
+                (self.id, models.User.user_role_id),
+            )
+        role = await info.context.data_loaders.user_roles.load(user_role_id)
         if role is None:
-            raise NotFound(f"User role with id {self.user_role_id} not found")
+            raise NotFound(f"User role with id {user_role_id} not found")
         return to_gql_user_role(role)
 
     @strawberry.field
     async def api_keys(self, info: Info[Context, None]) -> list[UserApiKey]:
         async with info.context.db() as session:
             api_keys = await session.scalars(
-                select(models.ApiKey).where(models.ApiKey.user_id == self.id_attr)
+                select(models.ApiKey).where(models.ApiKey.user_id == self.id)
             )
-        return [to_gql_api_key(api_key) for api_key in api_keys]
+        return [UserApiKey(id=api_key.id, db_record=api_key) for api_key in api_keys]
 
     @strawberry.field
-    async def is_management_user(self) -> bool:
+    async def is_management_user(self, info: Info[Context, None]) -> bool:
         initial_admins = get_env_admins()
         # this field is only visible to initial admins as they are the ones likely to have access to
         # a management interface / the phoenix environment.
-        if self.email in initial_admins or self.email == "admin@localhost":
+        if self.db_record:
+            email = self.db_record.email
+        else:
+            email = await info.context.data_loaders.user_fields.load(
+                (self.id, models.User.email),
+            )
+        if email in initial_admins or email == "admin@localhost":
             return True
         return False
-
-
-def to_gql_user(user: models.User, api_keys: Optional[list[models.ApiKey]] = None) -> User:
-    """
-    Converts an ORM user to a GraphQL user.
-    """
-    assert user.auth_method is not None
-    return User(
-        id_attr=user.id,
-        password_needs_reset=user.reset_password,
-        username=user.username,
-        email=user.email,
-        profile_picture_url=user.profile_picture_url,
-        created_at=user.created_at,
-        user_role_id=user.user_role_id,
-        auth_method=AuthMethod(user.auth_method),
-    )

phoenix/server/api/types/UserApiKey.py

@@ -1,14 +1,13 @@
-from typing import TYPE_CHECKING
+from datetime import datetime
+from typing import TYPE_CHECKING, Optional
 
 import strawberry
-from strawberry import Private
 from strawberry.relay import Node, NodeID
 from strawberry.types import Info
 from typing_extensions import Annotated
 
 from phoenix.db.models import ApiKey as OrmApiKey
 from phoenix.server.api.context import Context
-from phoenix.server.api.exceptions import NotFound
 
 from .ApiKey import ApiKey
 
@@ -18,28 +17,76 @@ if TYPE_CHECKING:
 
 @strawberry.type
 class UserApiKey(ApiKey, Node):
-    id_attr: NodeID[int]
-    user_id: Private[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("UserApiKey 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
 
     @strawberry.field
-    async def user(self, info: Info[Context, None]) -> Annotated["User", strawberry.lazy(".User")]:
-        user = await info.context.data_loaders.users.load(self.user_id)
-        if user is None:
-            raise NotFound(f"User with id {self.user_id} not found")
-        from .User import to_gql_user
-
-        return to_gql_user(user)
-
-
-def to_gql_api_key(api_key: OrmApiKey) -> UserApiKey:
-    """
-    Converts an ORM API key to a GraphQL UserApiKey type.
-    """
-    return UserApiKey(
-        id_attr=api_key.id,
-        user_id=api_key.user_id,
-        name=api_key.name,
-        description=api_key.description,
-        created_at=api_key.created_at,
-        expires_at=api_key.expires_at,
-    )
+    async def user(
+        self,
+        info: Info[Context, None],
+    ) -> Annotated["User", strawberry.lazy(".User")]:
+        if self.db_record:
+            user_id = self.db_record.user_id
+        else:
+            user_id = await info.context.data_loaders.user_api_key_fields.load(
+                (self.id, OrmApiKey.user_id),
+            )
+        from .User import User
+
+        return User(id=user_id)

phoenix/server/app.py

@@ -103,6 +103,7 @@ from phoenix.server.api.dataloaders import (
     ExperimentRepeatedRunGroupsDataLoader,
     ExperimentRunAnnotations,
     ExperimentRunCountsDataLoader,
+    ExperimentRunsByExperimentAndExampleDataLoader,
     ExperimentSequenceNumberDataLoader,
     LastUsedTimesByGenerativeModelIdDataLoader,
     LatencyMsQuantileDataLoader,
@@ -139,6 +140,7 @@ from phoenix.server.api.dataloaders import (
     SpanProjectsDataLoader,
     TableFieldsDataLoader,
     TokenCountDataLoader,
+    TokenPricesByModelDataLoader,
     TraceAnnotationsByTraceDataLoader,
     TraceByTraceIdsDataLoader,
     TraceRetentionPolicyIdByProjectIdDataLoader,
@@ -712,13 +714,18 @@ def create_graphql_router(
                 ),
                 average_experiment_run_latency=AverageExperimentRunLatencyDataLoader(db),
                 dataset_dataset_splits=DatasetDatasetSplitsDataLoader(db),
+                dataset_example_fields=TableFieldsDataLoader(db, models.DatasetExample),
                 dataset_example_revisions=DatasetExampleRevisionsDataLoader(db),
                 dataset_example_spans=DatasetExampleSpansDataLoader(db),
                 dataset_examples_and_versions_by_experiment_run=DatasetExamplesAndVersionsByExperimentRunDataLoader(
                     db
                 ),
                 dataset_example_splits=DatasetExampleSplitsDataLoader(db),
+                dataset_fields=TableFieldsDataLoader(db, models.Dataset),
+                dataset_split_fields=TableFieldsDataLoader(db, models.DatasetSplit),
+                dataset_version_fields=TableFieldsDataLoader(db, models.DatasetVersion),
                 dataset_labels=DatasetLabelsDataLoader(db),
+                dataset_label_fields=TableFieldsDataLoader(db, models.DatasetLabel),
                 document_evaluation_summaries=DocumentEvaluationSummaryDataLoader(
                     db,
                     cache_map=(
@@ -727,6 +734,7 @@ def create_graphql_router(
                         else None
                     ),
                 ),
+                document_annotation_fields=TableFieldsDataLoader(db, models.DocumentAnnotation),
                 document_evaluations=DocumentEvaluationsDataLoader(db),
                 document_retrieval_metrics=DocumentRetrievalMetricsDataLoader(db),
                 annotation_summaries=AnnotationSummaryDataLoader(
@@ -738,13 +746,22 @@ def create_graphql_router(
                 experiment_annotation_summaries=ExperimentAnnotationSummaryDataLoader(db),
                 experiment_dataset_splits=ExperimentDatasetSplitsDataLoader(db),
                 experiment_error_rates=ExperimentErrorRatesDataLoader(db),
+                experiment_fields=TableFieldsDataLoader(db, models.Experiment),
                 experiment_repeated_run_group_annotation_summaries=ExperimentRepeatedRunGroupAnnotationSummariesDataLoader(
                     db
                 ),
                 experiment_repeated_run_groups=ExperimentRepeatedRunGroupsDataLoader(db),
+                experiment_run_annotation_fields=TableFieldsDataLoader(
+                    db, models.ExperimentRunAnnotation
+                ),
                 experiment_run_annotations=ExperimentRunAnnotations(db),
                 experiment_run_counts=ExperimentRunCountsDataLoader(db),
+                experiment_run_fields=TableFieldsDataLoader(db, models.ExperimentRun),
+                experiment_runs_by_experiment_and_example=ExperimentRunsByExperimentAndExampleDataLoader(
+                    db
+                ),
                 experiment_sequence_number=ExperimentSequenceNumberDataLoader(db),
+                generative_model_fields=TableFieldsDataLoader(db, models.GenerativeModel),
                 last_used_times_by_generative_model_id=LastUsedTimesByGenerativeModelIdDataLoader(
                     db
                 ),
@@ -768,7 +785,14 @@ def create_graphql_router(
                 projects_by_trace_retention_policy_id=ProjectIdsByTraceRetentionPolicyIdDataLoader(
                     db
                 ),
+                prompt_fields=TableFieldsDataLoader(db, models.Prompt),
+                prompt_label_fields=TableFieldsDataLoader(db, models.PromptLabel),
                 prompt_version_sequence_number=PromptVersionSequenceNumberDataLoader(db),
+                prompt_version_tag_fields=TableFieldsDataLoader(db, models.PromptVersionTag),
+                project_session_annotation_fields=TableFieldsDataLoader(
+                    db, models.ProjectSessionAnnotation
+                ),
+                project_session_fields=TableFieldsDataLoader(db, models.ProjectSession),
                 record_counts=RecordCountDataLoader(
                     db,
                     cache_map=cache_for_dataloaders.record_count if cache_for_dataloaders else None,
@@ -780,6 +804,7 @@ def create_graphql_router(
                 session_num_traces_with_error=SessionNumTracesWithErrorDataLoader(db),
                 session_token_usages=SessionTokenUsagesDataLoader(db),
                 session_trace_latency_ms_quantile=SessionTraceLatencyMsQuantileDataLoader(db),
+                span_annotation_fields=TableFieldsDataLoader(db, models.SpanAnnotation),
                 span_annotations=SpanAnnotationsDataLoader(db),
                 span_fields=TableFieldsDataLoader(db, models.Span),
                 span_by_id=SpanByIdDataLoader(db),
@@ -820,6 +845,8 @@ def create_graphql_router(
                     db,
                     cache_map=cache_for_dataloaders.token_count if cache_for_dataloaders else None,
                 ),
+                token_prices_by_model=TokenPricesByModelDataLoader(db),
+                trace_annotation_fields=TableFieldsDataLoader(db, models.TraceAnnotation),
                 trace_annotations_by_trace=TraceAnnotationsByTraceDataLoader(db),
                 trace_by_trace_ids=TraceByTraceIdsDataLoader(db),
                 trace_fields=TableFieldsDataLoader(db, models.Trace),
@@ -832,6 +859,8 @@ def create_graphql_router(
                 trace_root_spans=TraceRootSpansDataLoader(db),
                 project_by_name=ProjectByNameDataLoader(db),
                 users=UsersDataLoader(db),
+                user_api_key_fields=TableFieldsDataLoader(db, models.ApiKey),
+                user_fields=TableFieldsDataLoader(db, models.User),
                 user_roles=UserRolesDataLoader(db),
             ),
             cache_for_dataloaders=cache_for_dataloaders,

phoenix/server/cost_tracking/model_cost_manifest.json

@@ -784,7 +784,7 @@
           "token_type": "output"
         },
         {
-          "base_rate": 7.5e-8,
+          "base_rate": 3e-8,
           "is_prompt": true,
           "token_type": "cache_read"
         },
@@ -1027,7 +1027,7 @@
           "token_type": "output"
         },
         {
-          "base_rate": 3.125e-7,
+          "base_rate": 1.25e-7,
           "is_prompt": true,
           "token_type": "cache_read"
         }