arize-phoenix 3.16.1__py3-none-any.whl → 7.7.1__py3-none-any.whl

phoenix/server/api/types/pagination.py

@@ -1,69 +1,149 @@
 import base64
 from dataclasses import dataclass
-from typing import Generic, List, Optional, TypeVar
+from datetime import datetime
+from enum import Enum, auto
+from typing import Any, ClassVar, Optional, Union
 
-import strawberry
 from strawberry import UNSET
+from strawberry.relay.types import Connection, Edge, NodeType, PageInfo
+from typing_extensions import TypeAlias, assert_never
 
-GenericType = TypeVar("GenericType")
-
-
-@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]
-    total_count: int
-
+ID: TypeAlias = int
+CursorSortColumnValue: TypeAlias = Union[str, int, float, datetime]
 
 # A type alias for the connection cursor implementation
-Cursor = 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
-
+CursorString = str
 
 # The hashing prefix for a connection cursor
 CURSOR_PREFIX = "connection:"
 
 
-def offset_to_cursor(offset: int) -> Cursor:
+class CursorSortColumnDataType(Enum):
+    STRING = auto()
+    INT = auto()
+    FLOAT = auto()
+    DATETIME = auto()
+
+
+@dataclass
+class CursorSortColumn:
+    type: CursorSortColumnDataType
+    value: CursorSortColumnValue
+
+    def __str__(self) -> str:
+        if isinstance(self.value, str):
+            return self.value
+        if isinstance(self.value, (int, float)):
+            return str(self.value)
+        if isinstance(self.value, datetime):
+            return self.value.isoformat()
+        assert_never(self.type)
+
+    @classmethod
+    def from_string(cls, type: CursorSortColumnDataType, cursor_string: str) -> "CursorSortColumn":
+        value: CursorSortColumnValue
+        if type is CursorSortColumnDataType.STRING:
+            value = cursor_string
+        elif type is CursorSortColumnDataType.INT:
+            value = int(cursor_string)
+        elif type is CursorSortColumnDataType.FLOAT:
+            value = float(cursor_string)
+        elif type is CursorSortColumnDataType.DATETIME:
+            value = datetime.fromisoformat(cursor_string)
+        else:
+            assert_never(type)
+        return cls(type=type, value=value)
+
+
+@dataclass
+class Cursor:
+    """
+    Serializes and deserializes cursor strings for ID-based pagination.
+
+    In the simplest case, a cursor encodes the rowid of a record. In the case
+    that a sort has been applied, the cursor additionally encodes the data type
+    and value of the column indexed for sorting so that the sort position can be
+    efficiently found. The encoding ensures that the cursor string is opaque to
+    the client and discourages the client from making use of the encoded
+    content.
+
+    Examples:
+        # encodes "10"
+        Cursor(rowid=10)
+
+        # encodes "11:STRING:abc"
+        Cursor(
+            rowid=11,
+            sort_column=CursorSortColumn(
+                type=CursorSortColumnDataType.STRING,
+                value="abc"
+            )
+        )
+
+        # encodes "10:INT:5"
+        Cursor(
+            rowid=10,
+            sort_column=CursorSortColumn(
+                type=CursorSortColumnDataType.INT,
+                value=5
+            )
+        )
+
+        # encodes "17:FLOAT:5.7"
+        Cursor(
+            rowid=17,
+            sort_column=CursorSortColumn(
+                type=CursorSortColumnDataType.FLOAT,
+                value=5.7
+            )
+        )
+
+        # encodes "20:DATETIME:2024-05-05T04:25:29.911245+00:00"
+        Cursor(
+            rowid=20,
+            sort_column=CursorSortColumn(
+                type=CursorSortColumnDataType.DATETIME,
+                value=datetime.fromisoformat("2024-05-05T04:25:29.911245+00:00")
+            )
+        )
+    """
+
+    rowid: int
+    sort_column: Optional[CursorSortColumn] = None
+
+    _DELIMITER: ClassVar[str] = ":"
+
+    def __str__(self) -> str:
+        cursor_parts = [str(self.rowid)]
+        if (sort_column := self.sort_column) is not None:
+            cursor_parts.extend([sort_column.type.name, str(sort_column)])
+        return base64.b64encode(self._DELIMITER.join(cursor_parts).encode()).decode()
+
+    @classmethod
+    def from_string(cls, cursor: str) -> "Cursor":
+        decoded = base64.b64decode(cursor).decode()
+        rowid_string = decoded
+        sort_column = None
+        if (first_delimiter_index := decoded.find(cls._DELIMITER)) > -1:
+            rowid_string = decoded[:first_delimiter_index]
+            second_delimiter_index = decoded.index(cls._DELIMITER, first_delimiter_index + 1)
+            sort_column = CursorSortColumn.from_string(
+                type=CursorSortColumnDataType[
+                    decoded[first_delimiter_index + 1 : second_delimiter_index]
+                ],
+                cursor_string=decoded[second_delimiter_index + 1 :],
+            )
+        return cls(rowid=int(rowid_string), sort_column=sort_column)
+
+
+def offset_to_cursor(offset: int) -> CursorString:
     """
     Creates the cursor string from an offset.
     """
     return base64.b64encode(f"{CURSOR_PREFIX}{offset}".encode("utf-8")).decode()
 
 
-def cursor_to_offset(cursor: Cursor) -> int:
+def cursor_to_offset(cursor: CursorString) -> int:
     """
     Extracts the offset from the cursor string.
     """
@@ -71,13 +151,13 @@ def cursor_to_offset(cursor: Cursor) -> int:
     return int(offset)
 
 
-def get_offset_with_default(cursor: Optional[Cursor], default_offset: int) -> int:
+def get_offset_with_default(cursor: Optional[CursorString], default_offset: int) -> int:
     """
     Given an optional cursor and a default offset, returns the offset
     to use; if the cursor contains a valid offset, that will be used,
     otherwise it will be the default.
     """
-    if not isinstance(cursor, Cursor):
+    if not isinstance(cursor, CursorString):
         return default_offset
     offset = cursor_to_offset(cursor)
     return offset if isinstance(offset, int) else default_offset
@@ -90,15 +170,15 @@ class ConnectionArgs:
     """
 
     first: Optional[int] = UNSET
-    after: Optional[Cursor] = UNSET
+    after: Optional[CursorString] = UNSET
     last: Optional[int] = UNSET
-    before: Optional[Cursor] = UNSET
+    before: Optional[CursorString] = UNSET
 
 
 def connection_from_list(
-    data: List[GenericType],
+    data: list[NodeType],
     args: ConnectionArgs,
-) -> Connection[GenericType]:
+) -> Connection[NodeType]:
     """
     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,
@@ -108,11 +188,11 @@ def connection_from_list(
 
 
 def connection_from_list_slice(
-    list_slice: List[GenericType],
+    list_slice: list[NodeType],
     args: ConnectionArgs,
     slice_start: int,
     list_length: int,
-) -> Connection[GenericType]:
+) -> Connection[NodeType]:
     """
     Given a slice (subset) of a list, returns a connection object for use in
     GraphQL.
@@ -169,6 +249,25 @@ def connection_from_list_slice(
             end_cursor=last_edge.cursor if last_edge else None,
             has_previous_page=start_offset > lower_bound if isinstance(args.last, int) else False,
             has_next_page=end_offset < upper_bound if isinstance(args.first, int) else False,
-            total_count=list_length,
+        ),
+    )
+
+
+def connection_from_cursors_and_nodes(
+    cursors_and_nodes: list[tuple[Any, NodeType]],
+    has_previous_page: bool,
+    has_next_page: bool,
+) -> Connection[NodeType]:
+    edges = [Edge(node=node, cursor=str(cursor)) for cursor, node in cursors_and_nodes]
+    has_edges = len(edges) > 0
+    first_edge = edges[0] if has_edges else None
+    last_edge = edges[-1] if has_edges else None
+    return Connection(
+        edges=edges,
+        page_info=PageInfo(
+            start_cursor=first_edge.cursor if first_edge else None,
+            end_cursor=last_edge.cursor if last_edge else None,
+            has_previous_page=has_previous_page,
+            has_next_page=has_next_page,
         ),
     )

phoenix/server/api/utils.py

@@ -0,0 +1,34 @@
+from sqlalchemy import delete
+
+from phoenix.db import models
+from phoenix.server.types import DbSessionFactory
+
+
+async def delete_projects(
+    db: DbSessionFactory,
+    *project_names: str,
+) -> list[int]:
+    if not project_names:
+        return []
+    stmt = (
+        delete(models.Project)
+        .where(models.Project.name.in_(set(project_names)))
+        .returning(models.Project.id)
+    )
+    async with db() as session:
+        return list(await session.scalars(stmt))
+
+
+async def delete_traces(
+    db: DbSessionFactory,
+    *trace_ids: str,
+) -> list[int]:
+    if not trace_ids:
+        return []
+    stmt = (
+        delete(models.Trace)
+        .where(models.Trace.trace_id.in_(set(trace_ids)))
+        .returning(models.Trace.id)
+    )
+    async with db() as session:
+        return list(await session.scalars(stmt))