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

phoenix/db/insertion/trace_annotation.py

@@ -56,7 +56,7 @@ class TraceAnnotationQueueInserter(
         session: AsyncSession,
         *insertions: Insertables.TraceAnnotation,
     ) -> list[TraceAnnotationDmlEvent]:
-        records = [dict(as_kv(ins.row)) for ins in insertions]
+        records = [{**dict(as_kv(ins.row)), "updated_at": ins.row.updated_at} for ins in insertions]
         stmt = self._insert_on_conflict(*records).returning(self.table.id)
         ids = tuple([_ async for _ in await session.stream_scalars(stmt)])
         return [TraceAnnotationDmlEvent(ids)]
@@ -91,7 +91,7 @@ class TraceAnnotationQueueInserter(
 
         for p in parcels:
             if (anno := existing_annos.get(_key(p))) is not None:
-                if p.received_at <= anno.updated_at:
+                if p.item.updated_at <= anno.updated_at:
                     to_discard.append(p)
                 else:
                     to_insert.append(
@@ -99,7 +99,6 @@ class TraceAnnotationQueueInserter(
                             received_at=p.received_at,
                             item=p.item.as_insertable(
                                 trace_rowid=anno.trace_rowid,
-                                id_=anno.id_,
                             ),
                         )
                     )
@@ -167,7 +166,7 @@ def _key(p: Received[Precursors.TraceAnnotation]) -> _Key:
 
 
 def _unique_by(p: Received[Insertables.TraceAnnotation]) -> _UniqueBy:
-    return p.item.obj.name, p.item.trace_rowid, p.item.identifier
+    return p.item.obj.name, p.item.trace_rowid, p.item.obj.identifier
 
 
 def _time(p: Received[Any]) -> datetime:

phoenix/db/insertion/types.py

@@ -94,7 +94,10 @@ class QueueInserter(ABC, Generic[_PrecursorT, _InsertableT, _RowT, _DmlEventT]):
     async def insert(self) -> Optional[list[_DmlEventT]]:
         if not self._queue:
             return None
-        self._queue, parcels = [], self._queue
+        parcels = self._queue.copy()
+        # IMPORTANT: Use .clear() instead of reassignment, i.e. self._queue = [], to
+        # avoid potential race conditions when appending postponed items to the queue.
+        self._queue.clear()
         events: list[_DmlEventT] = []
         async with self._db() as session:
             to_insert, to_postpone, _ = await self._partition(session, *parcels)
@@ -104,9 +107,13 @@ class QueueInserter(ABC, Generic[_PrecursorT, _InsertableT, _RowT, _DmlEventT]):
                     to_postpone.extend(to_retry)
         if to_postpone:
             loop = asyncio.get_running_loop()
-            loop.call_later(self._retry_delay_sec, self._queue.extend, to_postpone)
+            loop.call_later(self._retry_delay_sec, self._add_postponed_to_queue, to_postpone)
         return events
 
+    def _add_postponed_to_queue(self, items: list[Postponed[_PrecursorT]]) -> None:
+        """Add postponed items back to the queue for retry."""
+        self._queue.extend(items)
+
     def _insert_on_conflict(self, *records: Mapping[str, Any]) -> Insert:
         return insert_on_conflict(
             *records,
@@ -167,40 +174,41 @@ class QueueInserter(ABC, Generic[_PrecursorT, _InsertableT, _RowT, _DmlEventT]):
 class Precursors(ABC):
     @dataclass(frozen=True)
     class SpanAnnotation:
+        updated_at: datetime
         span_id: str
         obj: models.SpanAnnotation
 
         def as_insertable(
             self,
             span_rowid: int,
-            id_: Optional[int] = None,
         ) -> Insertables.SpanAnnotation:
             return Insertables.SpanAnnotation(
+                updated_at=self.updated_at,
                 span_id=self.span_id,
                 obj=self.obj,
                 span_rowid=span_rowid,
-                id_=id_,
             )
 
     @dataclass(frozen=True)
     class TraceAnnotation:
+        updated_at: datetime
         trace_id: str
         obj: models.TraceAnnotation
 
         def as_insertable(
             self,
             trace_rowid: int,
-            id_: Optional[int] = None,
         ) -> Insertables.TraceAnnotation:
             return Insertables.TraceAnnotation(
+                updated_at=self.updated_at,
                 trace_id=self.trace_id,
                 obj=self.obj,
                 trace_rowid=trace_rowid,
-                id_=id_,
             )
 
     @dataclass(frozen=True)
     class DocumentAnnotation:
+        updated_at: datetime
         span_id: str
         document_position: int
         obj: models.DocumentAnnotation
@@ -208,56 +216,78 @@ class Precursors(ABC):
         def as_insertable(
             self,
             span_rowid: int,
-            id_: Optional[int] = None,
         ) -> Insertables.DocumentAnnotation:
             return Insertables.DocumentAnnotation(
+                updated_at=self.updated_at,
                 span_id=self.span_id,
                 document_position=self.document_position,
                 obj=self.obj,
                 span_rowid=span_rowid,
-                id_=id_,
+            )
+
+    @dataclass(frozen=True)
+    class SessionAnnotation:
+        updated_at: datetime
+        session_id: str
+        obj: models.ProjectSessionAnnotation
+
+        def as_insertable(
+            self,
+            project_session_rowid: int,
+        ) -> Insertables.SessionAnnotation:
+            return Insertables.SessionAnnotation(
+                updated_at=self.updated_at,
+                session_id=self.session_id,
+                obj=self.obj,
+                project_session_rowid=project_session_rowid,
             )
 
 
 class Insertables(ABC):
     @dataclass(frozen=True)
     class SpanAnnotation(Precursors.SpanAnnotation):
+        updated_at: datetime
         span_rowid: int
-        identifier: str = ""
-        id_: Optional[int] = None
 
         @property
         def row(self) -> models.SpanAnnotation:
             obj = copy(self.obj)
             obj.span_rowid = self.span_rowid
-            if self.id_ is not None:
-                obj.id = self.id_
+            obj.updated_at = self.updated_at
             return obj
 
     @dataclass(frozen=True)
     class TraceAnnotation(Precursors.TraceAnnotation):
+        updated_at: datetime
         trace_rowid: int
-        identifier: str = ""
-        id_: Optional[int] = None
 
         @property
         def row(self) -> models.TraceAnnotation:
             obj = copy(self.obj)
             obj.trace_rowid = self.trace_rowid
-            if self.id_ is not None:
-                obj.id = self.id_
+            obj.updated_at = self.updated_at
             return obj
 
     @dataclass(frozen=True)
     class DocumentAnnotation(Precursors.DocumentAnnotation):
+        updated_at: datetime
         span_rowid: int
-        identifier: str = ""
-        id_: Optional[int] = None
 
         @property
         def row(self) -> models.DocumentAnnotation:
             obj = copy(self.obj)
             obj.span_rowid = self.span_rowid
-            if self.id_ is not None:
-                obj.id = self.id_
+            obj.updated_at = self.updated_at
+            return obj
+
+    @dataclass(frozen=True)
+    class SessionAnnotation(Precursors.SessionAnnotation):
+        updated_at: datetime
+        project_session_rowid: int
+
+        @property
+        def row(self) -> models.ProjectSessionAnnotation:
+            obj = copy(self.obj)
+            obj.project_session_id = self.project_session_rowid
+            obj.updated_at = self.updated_at
             return obj

phoenix/db/migrations/versions/01a8342c9cdf_add_user_id_on_datasets.py

@@ -0,0 +1,40 @@
+"""add user_id on datasets
+
+Revision ID: 01a8342c9cdf
+Revises: 0df286449799
+Create Date: 2025-09-25 16:08:51.254947
+
+"""
+
+from typing import Sequence, Union
+
+import sqlalchemy as sa
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision: str = "01a8342c9cdf"
+down_revision: Union[str, None] = "0df286449799"
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+_Integer = sa.Integer().with_variant(
+    sa.BigInteger(),
+    "postgresql",
+)
+
+
+def upgrade() -> None:
+    with op.batch_alter_table("datasets") as batch_op:
+        batch_op.add_column(
+            sa.Column(
+                "user_id",
+                _Integer,
+                sa.ForeignKey("users.id", ondelete="SET NULL"),
+                nullable=True,
+            ),
+        )
+
+
+def downgrade() -> None:
+    with op.batch_alter_table("datasets") as batch_op:
+        batch_op.drop_column("user_id")

phoenix/db/migrations/versions/0df286449799_add_session_annotations_table.py

@@ -0,0 +1,105 @@
+"""add session annotations table
+
+Revision ID: 0df286449799
+Revises: 735d3d93c33e
+Create Date: 2025-08-06 11:27:01.479664
+"""
+
+from typing import Any, Sequence, Union
+
+import sqlalchemy as sa
+from alembic import op
+from sqlalchemy import JSON
+from sqlalchemy.dialects import postgresql
+from sqlalchemy.ext.compiler import compiles
+
+# revision identifiers, used by Alembic.
+revision: str = "0df286449799"
+down_revision: Union[str, None] = "735d3d93c33e"
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+
+class JSONB(JSON):
+    # See https://docs.sqlalchemy.org/en/20/core/custom_types.html
+    __visit_name__ = "JSONB"
+
+
+@compiles(JSONB, "sqlite")
+def _(*args: Any, **kwargs: Any) -> str:
+    # See https://docs.sqlalchemy.org/en/20/core/custom_types.html
+    return "JSONB"
+
+
+JSON_ = (
+    JSON()
+    .with_variant(
+        postgresql.JSONB(),
+        "postgresql",
+    )
+    .with_variant(
+        JSONB(),
+        "sqlite",
+    )
+)
+
+_Integer = sa.Integer().with_variant(
+    sa.BigInteger(),
+    "postgresql",
+)
+
+
+def upgrade() -> None:
+    op.create_table(
+        "project_session_annotations",
+        sa.Column("id", _Integer, primary_key=True),
+        sa.Column(
+            "project_session_id",
+            _Integer,
+            sa.ForeignKey("project_sessions.id", ondelete="CASCADE"),
+            nullable=False,
+            index=True,
+        ),
+        sa.Column("name", sa.String, nullable=False),
+        sa.Column("label", sa.String),
+        sa.Column("score", sa.Float),
+        sa.Column("explanation", sa.String),
+        sa.Column("metadata", JSON_, nullable=False),
+        sa.Column(
+            "annotator_kind",
+            sa.String,
+            sa.CheckConstraint(
+                "annotator_kind IN ('LLM', 'CODE', 'HUMAN')",
+                name="valid_annotator_kind",
+            ),
+            nullable=False,
+        ),
+        sa.Column(
+            "user_id",
+            _Integer,
+            sa.ForeignKey("users.id", ondelete="SET NULL"),
+            nullable=True,
+        ),
+        sa.Column("identifier", sa.String, server_default="", nullable=False),
+        sa.Column(
+            "source",
+            sa.String,
+            sa.CheckConstraint("source IN ('API', 'APP')", name="valid_source"),
+            nullable=False,
+        ),
+        sa.Column(
+            "created_at", sa.TIMESTAMP(timezone=True), server_default=sa.func.now(), nullable=False
+        ),
+        sa.Column(
+            "updated_at",
+            sa.TIMESTAMP(timezone=True),
+            server_default=sa.func.now(),
+            onupdate=sa.func.now(),
+            nullable=False,
+        ),
+        sa.UniqueConstraint("name", "project_session_id", "identifier"),
+    )
+
+
+def downgrade() -> None:
+    op.drop_table("project_session_annotations")

phoenix/db/migrations/versions/272b66ff50f8_drop_single_indices.py

@@ -0,0 +1,119 @@
+"""drop single indices from traces, project_sessions, and experiment_runs
+
+Revision ID: 272b66ff50f8
+Revises: a20694b15f82
+Create Date: 2025-08-11 20:37:46.941940
+
+"""
+
+from typing import Sequence, Union
+
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision: str = "272b66ff50f8"
+down_revision: Union[str, None] = "a20694b15f82"
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+
+def upgrade() -> None:
+    op.drop_index(
+        "ix_traces_project_rowid",
+        table_name="traces",
+        if_exists=True,
+    )
+    op.drop_index(
+        "ix_traces_start_time",
+        table_name="traces",
+        if_exists=True,
+    )
+
+    op.drop_index(
+        "ix_project_sessions_project_id",
+        table_name="project_sessions",
+        if_exists=True,
+    )
+    op.drop_index(
+        "ix_project_sessions_start_time",
+        table_name="project_sessions",
+        if_exists=True,
+    )
+
+    op.drop_index(
+        "ix_experiment_runs_experiment_id",
+        table_name="experiment_runs",
+        if_exists=True,
+    )
+    op.drop_index(
+        "ix_experiment_run_annotations_experiment_run_id",
+        table_name="experiment_run_annotations",
+        if_exists=True,
+    )
+
+    op.drop_index(
+        "ix_dataset_example_revisions_dataset_example_id",
+        table_name="dataset_example_revisions",
+        if_exists=True,
+    )
+
+    op.drop_index(
+        "ix_span_cost_details_span_cost_id",
+        table_name="span_cost_details",
+        if_exists=True,
+    )
+
+
+def downgrade() -> None:
+    op.create_index(
+        "ix_traces_project_rowid",
+        "traces",
+        ["project_rowid"],
+        if_not_exists=True,
+    )
+    op.create_index(
+        "ix_traces_start_time",
+        "traces",
+        ["start_time"],
+        if_not_exists=True,
+    )
+
+    op.create_index(
+        "ix_project_sessions_project_id",
+        "project_sessions",
+        ["project_id"],
+        if_not_exists=True,
+    )
+    op.create_index(
+        "ix_project_sessions_start_time",
+        "project_sessions",
+        ["start_time"],
+        if_not_exists=True,
+    )
+
+    op.create_index(
+        "ix_experiment_runs_experiment_id",
+        "experiment_runs",
+        ["experiment_id"],
+        if_not_exists=True,
+    )
+    op.create_index(
+        "ix_experiment_run_annotations_experiment_run_id",
+        "experiment_run_annotations",
+        ["experiment_run_id"],
+        if_not_exists=True,
+    )
+
+    op.create_index(
+        "ix_dataset_example_revisions_dataset_example_id",
+        "dataset_example_revisions",
+        ["dataset_example_id"],
+        if_not_exists=True,
+    )
+
+    op.create_index(
+        "ix_span_cost_details_span_cost_id",
+        "span_cost_details",
+        ["span_cost_id"],
+        if_not_exists=True,
+    )

phoenix/db/migrations/versions/58228d933c91_dataset_labels.py

@@ -0,0 +1,67 @@
+"""dataset_labels
+
+Revision ID: 58228d933c91
+Revises: 699f655af132
+Create Date: 2025-09-05 17:47:34.637329
+
+"""
+
+from typing import Sequence, Union
+
+import sqlalchemy as sa
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision: str = "58228d933c91"
+down_revision: Union[str, None] = "699f655af132"
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+_Integer = sa.Integer().with_variant(
+    sa.BigInteger(),
+    "postgresql",
+)
+
+
+def upgrade() -> None:
+    op.create_table(
+        "dataset_labels",
+        sa.Column("id", _Integer, primary_key=True),
+        sa.Column("name", sa.String, nullable=False, unique=True),
+        sa.Column("description", sa.String, nullable=True),
+        sa.Column("color", sa.String, nullable=False),
+        sa.Column(
+            "user_id",
+            _Integer,
+            sa.ForeignKey("users.id", ondelete="SET NULL"),
+            nullable=True,
+            index=True,
+        ),
+    )
+
+    op.create_table(
+        "datasets_dataset_labels",
+        sa.Column(
+            "dataset_id",
+            _Integer,
+            sa.ForeignKey("datasets.id", ondelete="CASCADE"),
+            nullable=False,
+        ),
+        sa.Column(
+            "dataset_label_id",
+            _Integer,
+            sa.ForeignKey("dataset_labels.id", ondelete="CASCADE"),
+            nullable=False,
+            # index on the second element of the composite primary key
+            index=True,
+        ),
+        sa.PrimaryKeyConstraint(
+            "dataset_id",
+            "dataset_label_id",
+        ),
+    )
+
+
+def downgrade() -> None:
+    op.drop_table("datasets_dataset_labels")
+    op.drop_table("dataset_labels")

phoenix/db/migrations/versions/699f655af132_experiment_tags.py

@@ -0,0 +1,57 @@
+"""experiment_tags
+
+Revision ID: 699f655af132
+Revises: d0690a79ea51
+Create Date: 2025-09-05 13:14:22.676233
+
+"""
+
+from typing import Sequence, Union
+
+import sqlalchemy as sa
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision: str = "699f655af132"
+down_revision: Union[str, None] = "d0690a79ea51"
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+_Integer = sa.Integer().with_variant(
+    sa.BigInteger(),
+    "postgresql",
+)
+
+
+def upgrade() -> None:
+    op.create_table(
+        "experiment_tags",
+        sa.Column("id", _Integer, primary_key=True),
+        sa.Column(
+            "experiment_id",
+            _Integer,
+            sa.ForeignKey("experiments.id", ondelete="CASCADE"),
+            nullable=False,
+            index=True,
+        ),
+        sa.Column(
+            "dataset_id",
+            _Integer,
+            sa.ForeignKey("datasets.id", ondelete="CASCADE"),
+            nullable=False,
+        ),
+        sa.Column(
+            "user_id",
+            _Integer,
+            sa.ForeignKey("users.id", ondelete="SET NULL"),
+            nullable=True,
+            index=True,
+        ),
+        sa.Column("name", sa.String, nullable=False),
+        sa.Column("description", sa.String, nullable=True),
+        sa.UniqueConstraint("dataset_id", "name"),
+    )
+
+
+def downgrade() -> None:
+    op.drop_table("experiment_tags")

phoenix/db/migrations/versions/735d3d93c33e_add_composite_indices.py

@@ -0,0 +1,41 @@
+"""add composite indices to traces and project_sessions
+
+Revision ID: 735d3d93c33e
+Revises: 272b66ff50f8
+Create Date: 2025-08-11 20:52:47.477712
+
+"""
+
+from typing import Sequence, Union
+
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision: str = "735d3d93c33e"
+down_revision: Union[str, None] = "272b66ff50f8"
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+
+def upgrade() -> None:
+    op.execute(
+        "CREATE INDEX IF NOT EXISTS ix_traces_project_rowid_start_time "
+        "ON traces (project_rowid, start_time DESC)"
+    )
+    op.execute(
+        "CREATE INDEX IF NOT EXISTS ix_project_sessions_project_id_start_time "
+        "ON project_sessions (project_id, start_time DESC)"
+    )
+
+
+def downgrade() -> None:
+    op.drop_index(
+        "ix_traces_project_rowid_start_time",
+        table_name="traces",
+        if_exists=True,
+    )
+    op.drop_index(
+        "ix_project_sessions_project_id_start_time",
+        table_name="project_sessions",
+        if_exists=True,
+    )