arize-phoenix 3.16.0__py3-none-any.whl → 7.7.0__py3-none-any.whl

phoenix/server/api/types/TimeSeries.py

@@ -1,21 +1,25 @@
+from collections.abc import Iterable
 from dataclasses import replace
 from datetime import datetime, timedelta
 from functools import total_ordering
-from typing import Iterable, List, Optional, Tuple, Union, cast
+from typing import Optional, Union, cast
 
 import pandas as pd
 import strawberry
 from strawberry import UNSET
 
-from phoenix.core.model_schema import CONTINUOUS, PRIMARY, REFERENCE, Column, Dataset, Dimension
+from phoenix.core.model_schema import CONTINUOUS, PRIMARY, REFERENCE, Column, Dimension, Inferences
 from phoenix.metrics import Metric, binning
 from phoenix.metrics.mixins import UnaryOperator
 from phoenix.metrics.timeseries import timeseries
 from phoenix.server.api.input_types.Granularity import Granularity, to_timestamps
 from phoenix.server.api.input_types.TimeRange import TimeRange
 from phoenix.server.api.interceptor import GqlValueMediator
-from phoenix.server.api.types.DataQualityMetric import DataQualityMetric
-from phoenix.server.api.types.DatasetRole import DatasetRole
+from phoenix.server.api.types.DataQualityMetric import (
+    DATA_QUALITY_METRIC_FACTORIES,
+    DataQualityMetric,
+)
+from phoenix.server.api.types.InferencesRole import InferencesRole
 from phoenix.server.api.types.ScalarDriftMetricEnum import ScalarDriftMetric
 from phoenix.server.api.types.VectorDriftMetricEnum import VectorDriftMetric
 
@@ -39,7 +43,7 @@ def to_gql_datapoints(
     df: pd.DataFrame,
     metric: Metric,
     timestamps: Iterable[datetime],
-) -> List[TimeSeriesDataPoint]:
+) -> list[TimeSeriesDataPoint]:
     data = []
     for timestamp in timestamps:
         try:
@@ -59,7 +63,7 @@ def to_gql_datapoints(
 class TimeSeries:
     """A collection of data points over time"""
 
-    data: List[TimeSeriesDataPoint]
+    data: list[TimeSeriesDataPoint]
 
 
 def get_timeseries_data(
@@ -67,7 +71,7 @@ def get_timeseries_data(
     metric: Metric,
     time_range: TimeRange,
     granularity: Granularity,
-) -> List[TimeSeriesDataPoint]:
+) -> list[TimeSeriesDataPoint]:
     return df.pipe(
         timeseries(
             start_time=time_range.start,
@@ -97,16 +101,16 @@ def get_data_quality_timeseries_data(
     metric: DataQualityMetric,
     time_range: TimeRange,
     granularity: Granularity,
-    dataset_role: DatasetRole,
-) -> List[TimeSeriesDataPoint]:
-    metric_instance = metric.value()
+    inferences_role: InferencesRole,
+) -> list[TimeSeriesDataPoint]:
+    metric_instance = DATA_QUALITY_METRIC_FACTORIES[metric]()
     if isinstance(metric_instance, UnaryOperator):
         metric_instance = replace(
             metric_instance,
             operand=Column(dimension.name),
         )
     df = pd.DataFrame(
-        {dimension.name: dimension[dataset_role.value]},
+        {dimension.name: dimension[inferences_role.value]},
         copy=False,
     )
     return get_timeseries_data(
@@ -128,7 +132,7 @@ def get_drift_timeseries_data(
     time_range: TimeRange,
     granularity: Granularity,
     reference_data: pd.DataFrame,
-) -> List[TimeSeriesDataPoint]:
+) -> list[TimeSeriesDataPoint]:
     metric_instance = metric.value()
     metric_instance = replace(
         metric_instance,
@@ -160,12 +164,12 @@ class PerformanceTimeSeries(TimeSeries):
 
 
 def ensure_timeseries_parameters(
-    dataset: Dataset,
+    inferences: Inferences,
     time_range: Optional[TimeRange] = UNSET,
     granularity: Optional[Granularity] = UNSET,
-) -> Tuple[TimeRange, Granularity]:
+) -> tuple[TimeRange, Granularity]:
     if not isinstance(time_range, TimeRange):
-        start, stop = dataset.time_range
+        start, stop = inferences.time_range
         time_range = TimeRange(start=start, end=stop)
     if not isinstance(granularity, Granularity):
         total_minutes = int((time_range.end - time_range.start).total_seconds()) // 60

phoenix/server/api/types/TokenUsage.py

@@ -0,0 +1,11 @@
+import strawberry
+
+
+@strawberry.type
+class TokenUsage:
+    prompt: int = 0
+    completion: int = 0
+
+    @strawberry.field
+    async def total(self) -> int:
+        return self.prompt + self.completion

phoenix/server/api/types/Trace.py

@@ -0,0 +1,154 @@
+from __future__ import annotations
+
+from datetime import datetime
+from typing import TYPE_CHECKING, Annotated, Optional, Union
+
+import strawberry
+from openinference.semconv.trace import SpanAttributes
+from sqlalchemy import desc, select
+from sqlalchemy.orm import contains_eager
+from strawberry import UNSET, Private, lazy
+from strawberry.relay import Connection, GlobalID, Node, NodeID
+from strawberry.types import Info
+
+from phoenix.db import models
+from phoenix.server.api.context import Context
+from phoenix.server.api.input_types.TraceAnnotationSort import TraceAnnotationSort
+from phoenix.server.api.types.pagination import (
+    ConnectionArgs,
+    CursorString,
+    connection_from_list,
+)
+from phoenix.server.api.types.SortDir import SortDir
+from phoenix.server.api.types.Span import Span, to_gql_span
+from phoenix.server.api.types.TraceAnnotation import TraceAnnotation, to_gql_trace_annotation
+
+if TYPE_CHECKING:
+    from phoenix.server.api.types.ProjectSession import ProjectSession
+
+
+@strawberry.type
+class Trace(Node):
+    id_attr: NodeID[int]
+    project_rowid: Private[int]
+    project_session_rowid: Private[Optional[int]]
+    trace_id: str
+    start_time: datetime
+    end_time: datetime
+
+    @strawberry.field
+    async def latency_ms(
+        self,
+        info: Info[Context, None],
+    ) -> Optional[float]:
+        async with info.context.db() as session:
+            latency = await session.scalar(
+                select(
+                    models.Trace.latency_ms,
+                ).where(models.Trace.id == self.id_attr)
+            )
+        return latency
+
+    @strawberry.field
+    async def project_id(self) -> GlobalID:
+        from phoenix.server.api.types.Project import Project
+
+        return GlobalID(type_name=Project.__name__, node_id=str(self.project_rowid))
+
+    @strawberry.field
+    async def project_session_id(self) -> Optional[GlobalID]:
+        if self.project_session_rowid is None:
+            return None
+        from phoenix.server.api.types.ProjectSession import ProjectSession
+
+        return GlobalID(type_name=ProjectSession.__name__, node_id=str(self.project_session_rowid))
+
+    @strawberry.field
+    async def session(
+        self,
+        info: Info[Context, None],
+    ) -> Union[Annotated["ProjectSession", lazy(".ProjectSession")], None]:
+        if self.project_session_rowid is None:
+            return None
+        from phoenix.server.api.types.ProjectSession import to_gql_project_session
+
+        stmt = select(models.ProjectSession).filter_by(id=self.project_session_rowid)
+        async with info.context.db() as session:
+            project_session = await session.scalar(stmt)
+        if project_session is None:
+            return None
+        return to_gql_project_session(project_session)
+
+    @strawberry.field
+    async def root_span(
+        self,
+        info: Info[Context, None],
+    ) -> Optional[Span]:
+        span = await info.context.data_loaders.trace_root_spans.load(self.id_attr)
+        if span is None:
+            return None
+        return to_gql_span(span)
+
+    @strawberry.field
+    async def spans(
+        self,
+        info: Info[Context, None],
+        first: Optional[int] = 50,
+        last: Optional[int] = UNSET,
+        after: Optional[CursorString] = UNSET,
+        before: Optional[CursorString] = UNSET,
+    ) -> Connection[Span]:
+        args = ConnectionArgs(
+            first=first,
+            after=after if isinstance(after, CursorString) else None,
+            last=last,
+            before=before if isinstance(before, CursorString) else None,
+        )
+        stmt = (
+            select(models.Span)
+            .join(models.Trace)
+            .where(models.Trace.id == self.id_attr)
+            .options(contains_eager(models.Span.trace).load_only(models.Trace.trace_id))
+            # Sort descending because the root span tends to show up later
+            # in the ingestion process.
+            .order_by(desc(models.Span.id))
+            .limit(first)
+        )
+        async with info.context.db() as session:
+            spans = await session.stream_scalars(stmt)
+            data = [to_gql_span(span) async for span in spans]
+        return connection_from_list(data=data, args=args)
+
+    @strawberry.field(description="Annotations associated with the trace.")  # type: ignore
+    async def span_annotations(
+        self,
+        info: Info[Context, None],
+        sort: Optional[TraceAnnotationSort] = None,
+    ) -> list[TraceAnnotation]:
+        async with info.context.db() as session:
+            stmt = select(models.TraceAnnotation).filter_by(span_rowid=self.id_attr)
+            if sort:
+                sort_col = getattr(models.TraceAnnotation, sort.col.value)
+                if sort.dir is SortDir.desc:
+                    stmt = stmt.order_by(sort_col.desc(), models.TraceAnnotation.id.desc())
+                else:
+                    stmt = stmt.order_by(sort_col.asc(), models.TraceAnnotation.id.asc())
+            else:
+                stmt = stmt.order_by(models.TraceAnnotation.created_at.desc())
+            annotations = await session.scalars(stmt)
+        return [to_gql_trace_annotation(annotation) for annotation in annotations]
+
+
+def to_gql_trace(trace: models.Trace) -> Trace:
+    return Trace(
+        id_attr=trace.id,
+        project_rowid=trace.project_rowid,
+        project_session_rowid=trace.project_session_rowid,
+        trace_id=trace.trace_id,
+        start_time=trace.start_time,
+        end_time=trace.end_time,
+    )
+
+
+INPUT_VALUE = SpanAttributes.INPUT_VALUE.split(".")
+OUTPUT_VALUE = SpanAttributes.OUTPUT_VALUE.split(".")

phoenix/server/api/types/TraceAnnotation.py

@@ -0,0 +1,45 @@
+from typing import Optional
+
+import strawberry
+from strawberry import Private
+from strawberry.relay import GlobalID, Node, NodeID
+from strawberry.scalars import JSON
+
+from phoenix.db import models
+from phoenix.server.api.types.AnnotatorKind import AnnotatorKind
+
+
+@strawberry.type
+class TraceAnnotation(Node):
+    id_attr: NodeID[int]
+    name: str
+    annotator_kind: AnnotatorKind
+    label: Optional[str]
+    score: Optional[float]
+    explanation: Optional[str]
+    metadata: JSON
+    trace_rowid: Private[Optional[int]]
+
+    @strawberry.field
+    async def trace_id(self) -> GlobalID:
+        from phoenix.server.api.types.Trace import Trace
+
+        return GlobalID(type_name=Trace.__name__, node_id=str(self.trace_rowid))
+
+
+def to_gql_trace_annotation(
+    annotation: models.TraceAnnotation,
+) -> TraceAnnotation:
+    """
+    Converts an ORM trace annotation to a GraphQL TraceAnnotation.
+    """
+    return TraceAnnotation(
+        id_attr=annotation.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_,
+    )

phoenix/server/api/types/UMAPPoints.py

@@ -1,15 +1,15 @@
-from typing import List, Union
+from typing import Union
 
 import numpy as np
 import numpy.typing as npt
 import strawberry
+from strawberry.relay.types import GlobalID
 from strawberry.scalars import ID
 
 from phoenix.server.api.types.Cluster import Cluster
 
 from .EmbeddingMetadata import EmbeddingMetadata
 from .EventMetadata import EventMetadata
-from .node import GlobalID
 from .Retrieval import Retrieval
 
 
@@ -57,8 +57,8 @@ class UMAPPoint:
 
 @strawberry.type
 class UMAPPoints:
-    data: List[UMAPPoint]
-    reference_data: List[UMAPPoint]
-    clusters: List[Cluster]
-    corpus_data: List[UMAPPoint] = strawberry.field(default_factory=list)
-    context_retrievals: List[Retrieval] = strawberry.field(default_factory=list)
+    data: list[UMAPPoint]
+    reference_data: list[UMAPPoint]
+    clusters: list[Cluster]
+    corpus_data: list[UMAPPoint] = strawberry.field(default_factory=list)
+    context_retrievals: list[Retrieval] = strawberry.field(default_factory=list)

phoenix/server/api/types/User.py

@@ -0,0 +1,60 @@
+from datetime import datetime
+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
+
+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 .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
+
+    @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 role is None:
+            raise NotFound(f"User role with id {self.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)
+            )
+        return [to_gql_api_key(api_key) for api_key in api_keys]
+
+
+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

@@ -0,0 +1,45 @@
+from typing import TYPE_CHECKING
+
+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
+
+if TYPE_CHECKING:
+    from .User import User
+
+
+@strawberry.type
+class UserApiKey(ApiKey, Node):
+    id_attr: NodeID[int]
+    user_id: Private[int]
+
+    @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,
+    )

phoenix/server/api/types/UserRole.py

@@ -0,0 +1,15 @@
+import strawberry
+from strawberry.relay import Node, NodeID
+
+from phoenix.db import models
+
+
+@strawberry.type
+class UserRole(Node):
+    id_attr: NodeID[int]
+    name: str
+
+
+def to_gql_user_role(role: models.UserRole) -> UserRole:
+    """Convert an ORM user role to a GraphQL user role."""
+    return UserRole(id_attr=role.id, name=role.name)

phoenix/server/api/types/node.py

@@ -1,119 +1,25 @@
-import base64
-import dataclasses
-from typing import Tuple, Union
+from strawberry.relay import GlobalID
 
-import strawberry
-from graphql import GraphQLID
-from strawberry.custom_scalar import ScalarDefinition
-from strawberry.schema.types.scalar import DEFAULT_SCALAR_REGISTRY
 
-
-def to_global_id(type_name: str, node_id: int) -> str:
-    """
-    Encode the given id into a global id.
-
-    :param type_name: The type of the node.
-    :param node_id: The id of the node.
-    :return: A global id.
-    """
-    return base64.b64encode(f"{type_name}:{node_id}".encode("utf-8")).decode()
-
-
-def from_global_id(global_id: str) -> Tuple[str, int]:
+def from_global_id(global_id: GlobalID) -> tuple[str, int]:
     """
     Decode the given global id into a type and id.
 
     :param global_id: The global id to decode.
     :return: A tuple of type and id.
     """
-    type_name, node_id = base64.b64decode(global_id).decode().split(":")
-    return type_name, int(node_id)
+    return global_id.type_name, int(global_id.node_id)
 
 
-class GlobalIDValueError(ValueError):
-    """GlobalID value error, usually related to parsing or serialization."""
-
-
-@dataclasses.dataclass(frozen=True)
-class GlobalID:
-    """Global ID for relay types.
-    Different from `strawberry.ID`, this ID wraps the original object ID in a string
-    that contains both its GraphQL type name and the ID itself, and encodes it
-    to a base64_ string.
-    This object contains helpers to work with that, including method to retrieve
-    the python object type or even the encoded node itself.
-    Attributes:
-        type_name:
-            The type name part of the id
-        node_id:
-            The node id part of the id
-    .. _base64:
-        https://en.wikipedia.org/wiki/Base64
+def from_global_id_with_expected_type(global_id: GlobalID, expected_type_name: str) -> int:
     """
-
-    type_name: str
-    node_id: int
-
-    def __post_init__(self) -> None:
-        if not isinstance(self.type_name, str):
-            raise GlobalIDValueError(
-                f"type_name is expected to be a string, found {self.type_name}"
-            )
-        try:
-            # node_id could be numpy.int64, hence the need for coercion
-            object.__setattr__(self, "node_id", int(self.node_id))
-        except ValueError:
-            raise GlobalIDValueError(f"node_id is expected to be an int, found {self.node_id}")
-
-    def __str__(self) -> str:
-        return to_global_id(self.type_name, self.node_id)
-
-    @classmethod
-    def from_id(cls, value: Union[str, strawberry.ID]) -> "GlobalID":
-        """Create a new GlobalID from parsing the given value.
-        Args:
-            value:
-                The value to be parsed, as a base64 string in the "TypeName:NodeID" format
-        Returns:
-            An instance of GLobalID
-        Raises:
-            GlobalIDValueError:
-                If the value is not in a GLobalID format
-        """
-        try:
-            type_name, node_id = from_global_id(value)
-        except ValueError as e:
-            raise GlobalIDValueError(str(e)) from e
-
-        return cls(type_name=type_name, node_id=node_id)
-
-
-@strawberry.interface(description="A node in the graph with a globally unique ID")
-class Node:
-    """
-    All types that are relay ready should inherit from this interface and
-    implement the following methods.
-
-    Attributes:
-        id_attr:
-            The raw id field of node. Typically a database id or index
+    Decodes the given global id and return the id, checking that the type
+    matches the expected type.
     """
-
-    id_attr: strawberry.Private[int]
-
-    @strawberry.field
-    def id(self) -> GlobalID:
-        return GlobalID(type(self).__name__, self.id_attr)
-
-
-# Register our GlobalID scalar
-DEFAULT_SCALAR_REGISTRY[GlobalID] = ScalarDefinition(
-    # Use the same name/description/parse_literal from GraphQLID
-    # specs expect this type to be "ID".
-    name="GlobalID",
-    description=GraphQLID.description,
-    parse_literal=lambda v, vars=None: GlobalID.from_id(GraphQLID.parse_literal(v, vars)),
-    parse_value=GlobalID.from_id,
-    serialize=str,
-    specified_by_url="https://relay.dev/graphql/objectidentification.htm",
-)
+    type_name, node_id = from_global_id(global_id)
+    if type_name != expected_type_name:
+        raise ValueError(
+            f"The node id must correspond to a node of type {expected_type_name}, "
+            f"but instead corresponds to a node of type: {type_name}"
+        )
+    return node_id