arize-phoenix 4.4.4rc5__py3-none-any.whl → 4.5.0__py3-none-any.whl

phoenix/server/api/types/Span.py

@@ -1,33 +1,23 @@
 import json
-from dataclasses import dataclass
 from datetime import datetime
 from enum import Enum
-from typing import TYPE_CHECKING, Any, List, Mapping, Optional, Sized, cast
+from typing import Any, List, Mapping, Optional, Sized, cast
 
 import numpy as np
 import strawberry
 from openinference.semconv.trace import EmbeddingAttributes, SpanAttributes
 from strawberry import ID, UNSET
-from strawberry.relay import Node, NodeID
 from strawberry.types import Info
-from typing_extensions import Annotated
 
 import phoenix.trace.schemas as trace_schema
 from phoenix.db import models
 from phoenix.server.api.context import Context
-from phoenix.server.api.helpers.dataset_helpers import (
-    get_dataset_example_input,
-    get_dataset_example_output,
-)
 from phoenix.server.api.types.DocumentRetrievalMetrics import DocumentRetrievalMetrics
 from phoenix.server.api.types.Evaluation import DocumentEvaluation, SpanEvaluation
-from phoenix.server.api.types.ExampleRevisionInterface import ExampleRevision
 from phoenix.server.api.types.MimeType import MimeType
+from phoenix.server.api.types.node import Node
 from phoenix.trace.attributes import get_attribute_value
 
-if TYPE_CHECKING:
-    from phoenix.server.api.types.Project import Project
-
 EMBEDDING_EMBEDDINGS = SpanAttributes.EMBEDDING_EMBEDDINGS
 EMBEDDING_VECTOR = EmbeddingAttributes.EMBEDDING_VECTOR
 INPUT_MIME_TYPE = SpanAttributes.INPUT_MIME_TYPE
@@ -35,9 +25,6 @@ INPUT_VALUE = SpanAttributes.INPUT_VALUE
 LLM_TOKEN_COUNT_COMPLETION = SpanAttributes.LLM_TOKEN_COUNT_COMPLETION
 LLM_TOKEN_COUNT_PROMPT = SpanAttributes.LLM_TOKEN_COUNT_PROMPT
 LLM_TOKEN_COUNT_TOTAL = SpanAttributes.LLM_TOKEN_COUNT_TOTAL
-LLM_PROMPT_TEMPLATE_VARIABLES = SpanAttributes.LLM_PROMPT_TEMPLATE_VARIABLES
-LLM_INPUT_MESSAGES = SpanAttributes.LLM_INPUT_MESSAGES
-LLM_OUTPUT_MESSAGES = SpanAttributes.LLM_OUTPUT_MESSAGES
 METADATA = SpanAttributes.METADATA
 OUTPUT_MIME_TYPE = SpanAttributes.OUTPUT_MIME_TYPE
 OUTPUT_VALUE = SpanAttributes.OUTPUT_VALUE
@@ -59,7 +46,6 @@ class SpanKind(Enum):
     embedding = "EMBEDDING"
     agent = "AGENT"
     reranker = "RERANKER"
-    evaluator = "EVALUATOR"
     unknown = "UNKNOWN"
 
     @classmethod
@@ -115,14 +101,8 @@ class SpanEvent:
         )
 
 
-@strawberry.type
-class SpanAsExampleRevision(ExampleRevision): ...
-
-
 @strawberry.type
 class Span(Node):
-    id_attr: NodeID[int]
-    db_span: strawberry.Private[models.Span]
     name: str
     status_code: SpanStatusCode
     status_message: str
@@ -208,44 +188,6 @@ class Span(Node):
         spans = await info.context.data_loaders.span_descendants.load(span_id)
         return [to_gql_span(span) for span in spans]
 
-    @strawberry.field(
-        description="The span's attributes translated into an example revision for a dataset",
-    )  # type: ignore
-    def as_example_revision(self) -> SpanAsExampleRevision:
-        db_span = self.db_span
-        attributes = db_span.attributes
-        span_io = _SpanIO(
-            span_kind=db_span.span_kind,
-            input_value=get_attribute_value(attributes, INPUT_VALUE),
-            input_mime_type=get_attribute_value(attributes, INPUT_MIME_TYPE),
-            output_value=get_attribute_value(attributes, OUTPUT_VALUE),
-            output_mime_type=get_attribute_value(attributes, OUTPUT_MIME_TYPE),
-            llm_prompt_template_variables=get_attribute_value(
-                attributes, LLM_PROMPT_TEMPLATE_VARIABLES
-            ),
-            llm_input_messages=get_attribute_value(attributes, LLM_INPUT_MESSAGES),
-            llm_output_messages=get_attribute_value(attributes, LLM_OUTPUT_MESSAGES),
-            retrieval_documents=get_attribute_value(attributes, RETRIEVAL_DOCUMENTS),
-        )
-        return SpanAsExampleRevision(
-            input=get_dataset_example_input(span_io),
-            output=get_dataset_example_output(span_io),
-            metadata=attributes,
-        )
-
-    @strawberry.field(description="The project that this span belongs to.")  # type: ignore
-    async def project(
-        self,
-        info: Info[Context, None],
-    ) -> Annotated[
-        "Project", strawberry.lazy("phoenix.server.api.types.Project")
-    ]:  # use lazy types to avoid circular import: https://strawberry.rocks/docs/types/lazy
-        from phoenix.server.api.types.Project import to_gql_project
-
-        span_id = self.id_attr
-        project = await info.context.data_loaders.span_projects.load(span_id)
-        return to_gql_project(project)
-
 
 def to_gql_span(span: models.Span) -> Span:
     events: List[SpanEvent] = list(map(SpanEvent.from_dict, span.events))
@@ -255,7 +197,6 @@ def to_gql_span(span: models.Span) -> Span:
     num_documents = len(retrieval_documents) if isinstance(retrieval_documents, Sized) else None
     return Span(
         id_attr=span.id,
-        db_span=span,
         name=span.name,
         status_code=SpanStatusCode(span.status_code),
         status_message=span.status_message,
@@ -361,21 +302,3 @@ def _convert_metadata_to_string(metadata: Any) -> Optional[str]:
         return json.dumps(metadata)
     except Exception:
         return str(metadata)
-
-
-@dataclass
-class _SpanIO:
-    """
-    An class that contains the information needed to extract dataset example
-    input and output values from a span.
-    """
-
-    span_kind: Optional[str]
-    input_value: Any
-    input_mime_type: Optional[str]
-    output_value: Any
-    output_mime_type: Optional[str]
-    llm_prompt_template_variables: Any
-    llm_input_messages: Any
-    llm_output_messages: Any
-    retrieval_documents: Any

phoenix/server/api/types/TimeSeries.py

@@ -7,7 +7,7 @@ import pandas as pd
 import strawberry
 from strawberry import UNSET
 
-from phoenix.core.model_schema import CONTINUOUS, PRIMARY, REFERENCE, Column, Dimension, Inferences
+from phoenix.core.model_schema import CONTINUOUS, PRIMARY, REFERENCE, Column, Dataset, Dimension
 from phoenix.metrics import Metric, binning
 from phoenix.metrics.mixins import UnaryOperator
 from phoenix.metrics.timeseries import timeseries
@@ -15,7 +15,7 @@ from phoenix.server.api.input_types.Granularity import Granularity, to_timestamp
 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.InferencesRole import InferencesRole
+from phoenix.server.api.types.DatasetRole import DatasetRole
 from phoenix.server.api.types.ScalarDriftMetricEnum import ScalarDriftMetric
 from phoenix.server.api.types.VectorDriftMetricEnum import VectorDriftMetric
 
@@ -97,7 +97,7 @@ def get_data_quality_timeseries_data(
     metric: DataQualityMetric,
     time_range: TimeRange,
     granularity: Granularity,
-    inferences_role: InferencesRole,
+    dataset_role: DatasetRole,
 ) -> List[TimeSeriesDataPoint]:
     metric_instance = metric.value()
     if isinstance(metric_instance, UnaryOperator):
@@ -106,7 +106,7 @@ def get_data_quality_timeseries_data(
             operand=Column(dimension.name),
         )
     df = pd.DataFrame(
-        {dimension.name: dimension[inferences_role.value]},
+        {dimension.name: dimension[dataset_role.value]},
         copy=False,
     )
     return get_timeseries_data(
@@ -160,12 +160,12 @@ class PerformanceTimeSeries(TimeSeries):
 
 
 def ensure_timeseries_parameters(
-    inferences: Inferences,
+    dataset: Dataset,
     time_range: Optional[TimeRange] = UNSET,
     granularity: Optional[Granularity] = UNSET,
 ) -> Tuple[TimeRange, Granularity]:
     if not isinstance(time_range, TimeRange):
-        start, stop = inferences.time_range
+        start, stop = dataset.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/Trace.py

@@ -1,18 +1,17 @@
-from __future__ import annotations
-
 from typing import List, Optional
 
 import strawberry
 from sqlalchemy import desc, select
 from sqlalchemy.orm import contains_eager
-from strawberry import UNSET, Private
-from strawberry.relay import Connection, GlobalID, Node, NodeID
+from strawberry import UNSET
 from strawberry.types import Info
 
 from phoenix.db import models
 from phoenix.server.api.context import Context
 from phoenix.server.api.types.Evaluation import TraceEvaluation
+from phoenix.server.api.types.node import Node
 from phoenix.server.api.types.pagination import (
+    Connection,
     ConnectionArgs,
     CursorString,
     connection_from_list,
@@ -22,16 +21,6 @@ from phoenix.server.api.types.Span import Span, to_gql_span
 
 @strawberry.type
 class Trace(Node):
-    id_attr: NodeID[int]
-    project_rowid: Private[int]
-    trace_id: str
-
-    @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 spans(
         self,
@@ -51,7 +40,7 @@ class Trace(Node):
             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))
+            .options(contains_eager(models.Span.trace))
             # Sort descending because the root span tends to show up later
             # in the ingestion process.
             .order_by(desc(models.Span.id))

phoenix/server/api/types/UMAPPoints.py

@@ -3,13 +3,13 @@ from typing import List, 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
 
 

phoenix/server/api/types/node.py

@@ -1,19 +1,36 @@
-from typing import Tuple
+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 from_global_id(global_id: GlobalID) -> Tuple[str, int]:
+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]:
     """
     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.
     """
-    return global_id.type_name, int(global_id.node_id)
+    type_name, node_id = base64.b64decode(global_id).decode().split(":")
+    return type_name, int(node_id)
 
 
-def from_global_id_with_expected_type(global_id: GlobalID, expected_type_name: str) -> int:
+def from_global_id_with_expected_type(global_id: str, expected_type_name: str) -> int:
     """
     Decodes the given global id and return the id, checking that the type
     matches the expected type.
@@ -25,3 +42,92 @@ def from_global_id_with_expected_type(global_id: GlobalID, expected_type_name: s
             f"but instead corresponds to a node of type: {type_name}"
         )
     return 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
+    """
+
+    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
+    """
+
+    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",
+)

phoenix/server/api/types/pagination.py

@@ -2,18 +2,60 @@ import base64
 from dataclasses import dataclass
 from datetime import datetime
 from enum import Enum, auto
-from typing import Any, ClassVar, List, Optional, Tuple, Union
+from typing import ClassVar, Generic, List, Optional, Tuple, TypeVar, Union
 
+import strawberry
 from strawberry import UNSET
-from strawberry.relay.types import Connection, Edge, NodeType, PageInfo
 from typing_extensions import TypeAlias, assert_never
 
 ID: TypeAlias = int
+GenericType = TypeVar("GenericType")
 CursorSortColumnValue: TypeAlias = Union[str, int, float, datetime]
 
+
+@strawberry.type
+class Connection(Generic[GenericType]):
+    """Represents a paginated relationship between two entities
+
+    This pattern is used when the relationship itself has attributes.
+    """
+
+    page_info: "PageInfo"
+    edges: List["Edge[GenericType]"]
+
+
+@strawberry.type
+class PageInfo:
+    """Pagination context to navigate objects with cursor-based pagination
+
+    Instead of classic offset pagination via `page` and `limit` parameters,
+    here we have a cursor of the last object and we fetch items starting from that one
+
+    Read more at:
+        - https://graphql.org/learn/pagination/#pagination-and-edges
+        - https://relay.dev/graphql/connections.htm
+    """
+
+    has_next_page: bool
+    has_previous_page: bool
+    start_cursor: Optional[str]
+    end_cursor: Optional[str]
+
+
 # A type alias for the connection cursor implementation
 CursorString = str
 
+
+@strawberry.type
+class Edge(Generic[GenericType]):
+    """
+    An edge may contain additional information of the relationship. This is the trivial case
+    """
+
+    node: GenericType
+    cursor: str
+
+
 # The hashing prefix for a connection cursor
 CURSOR_PREFIX = "connection:"
 
@@ -176,9 +218,9 @@ class ConnectionArgs:
 
 
 def connection_from_list(
-    data: List[NodeType],
+    data: List[GenericType],
     args: ConnectionArgs,
-) -> Connection[NodeType]:
+) -> Connection[GenericType]:
     """
     A simple function that accepts a list and connection arguments, and returns
     a connection object for use in GraphQL. It uses list offsets as pagination,
@@ -188,11 +230,11 @@ def connection_from_list(
 
 
 def connection_from_list_slice(
-    list_slice: List[NodeType],
+    list_slice: List[GenericType],
     args: ConnectionArgs,
     slice_start: int,
     list_length: int,
-) -> Connection[NodeType]:
+) -> Connection[GenericType]:
     """
     Given a slice (subset) of a list, returns a connection object for use in
     GraphQL.
@@ -253,12 +295,12 @@ def connection_from_list_slice(
     )
 
 
-def connection_from_cursors_and_nodes(
-    cursors_and_nodes: List[Tuple[Any, NodeType]],
+def connections(
+    data: List[Tuple[Cursor, GenericType]],
     has_previous_page: bool,
     has_next_page: bool,
-) -> Connection[NodeType]:
-    edges = [Edge(node=node, cursor=str(cursor)) for cursor, node in cursors_and_nodes]
+) -> Connection[GenericType]:
+    edges = [Edge(node=node, cursor=str(cursor)) for cursor, node in data]
     has_edges = len(edges) > 0
     first_edge = edges[0] if has_edges else None
     last_edge = edges[-1] if has_edges else None