arize-phoenix 8.32.1__py3-none-any.whl → 9.0.1__py3-none-any.whl

phoenix/db/insertion/helpers.py

@@ -36,6 +36,7 @@ def insert_on_conflict(
     unique_by: Sequence[str],
     on_conflict: OnConflict = OnConflict.DO_UPDATE,
     set_: Optional[Mapping[str, Any]] = None,
+    constraint_name: Optional[str] = None,
 ) -> Insert:
     """
     Dialect specific insertion statement using ON CONFLICT DO syntax.
@@ -50,7 +51,7 @@ def insert_on_conflict(
             unique_records.append(v)
             seen.add(k)
         records = tuple(reversed(unique_records))
-    constraint = "_".join(("uq", table.__tablename__, *unique_by))
+    constraint = constraint_name or "_".join(("uq", table.__tablename__, *unique_by))
     if dialect is SupportedSQLDialect.POSTGRESQL:
         stmt_postgresql = insert_postgresql(table).values(records)
         if on_conflict is OnConflict.DO_NOTHING:

phoenix/db/insertion/span_annotation.py

@@ -21,15 +21,23 @@ from phoenix.server.dml_event import SpanAnnotationDmlEvent
 _Name: TypeAlias = str
 _SpanId: TypeAlias = str
 _SpanRowId: TypeAlias = int
+_Identifier: TypeAlias = str
 _AnnoRowId: TypeAlias = int
 
-_Key: TypeAlias = tuple[_Name, _SpanId]
-_UniqueBy: TypeAlias = tuple[_Name, _SpanRowId]
+
+class _Key(NamedTuple):
+    annotation_name: _Name
+    annotation_identifier: _Identifier
+    span_id: _SpanId
+
+
+_UniqueBy: TypeAlias = tuple[_Name, _SpanRowId, _Identifier]
 _Existing: TypeAlias = tuple[
     _SpanRowId,
     _SpanId,
     Optional[_AnnoRowId],
     Optional[_Name],
+    Optional[_Identifier],
     Optional[datetime],
 ]
 
@@ -42,7 +50,7 @@ class SpanAnnotationQueueInserter(
         SpanAnnotationDmlEvent,
     ],
     table=models.SpanAnnotation,
-    unique_by=("name", "span_rowid"),
+    unique_by=("name", "span_rowid", "identifier"),
 ):
     async def _events(
         self,
@@ -73,7 +81,11 @@ class SpanAnnotationQueueInserter(
             e.span_id: _SpanAttr(e.span_rowid) for e in existing
         }
         existing_annos: Mapping[_Key, _AnnoAttr] = {
-            (e.name, e.span_id): _AnnoAttr(e.span_rowid, e.id, e.updated_at)
+            _Key(
+                annotation_name=e.name,
+                annotation_identifier=e.identifier,
+                span_id=e.span_id,
+            ): _AnnoAttr(e.span_rowid, e.id, e.updated_at)
             for e in existing
             if e.id is not None and e.name is not None and e.updated_at is not None
         }
@@ -119,19 +131,20 @@ class SpanAnnotationQueueInserter(
         anno = self.table
         span = (
             select(models.Span.id, models.Span.span_id)
-            .where(models.Span.span_id.in_({span_id for _, span_id in keys}))
+            .where(models.Span.span_id.in_({k.span_id for k in keys}))
             .cte()
         )
         onclause = and_(
             span.c.id == anno.span_rowid,
-            anno.name.in_({name for name, _ in keys}),
-            tuple_(anno.name, span.c.span_id).in_(keys),
+            anno.name.in_({k.annotation_name for k in keys}),
+            tuple_(anno.name, anno.identifier, span.c.span_id).in_(keys),
         )
         return select(
             span.c.id.label("span_rowid"),
             span.c.span_id,
             anno.id,
             anno.name,
+            anno.identifier,
             anno.updated_at,
         ).outerjoin_from(span, anno, onclause)
 
@@ -147,11 +160,15 @@ class _AnnoAttr(NamedTuple):
 
 
 def _key(p: Received[Precursors.SpanAnnotation]) -> _Key:
-    return p.item.obj.name, p.item.span_id
+    return _Key(
+        annotation_name=p.item.obj.name,
+        annotation_identifier=p.item.obj.identifier,
+        span_id=p.item.span_id,
+    )
 
 
 def _unique_by(p: Received[Insertables.SpanAnnotation]) -> _UniqueBy:
-    return p.item.obj.name, p.item.span_rowid
+    return p.item.obj.name, p.item.span_rowid, p.item.identifier
 
 
 def _time(p: Received[Any]) -> datetime:

phoenix/db/insertion/trace_annotation.py

@@ -22,9 +22,16 @@ _Name: TypeAlias = str
 _TraceId: TypeAlias = str
 _TraceRowId: TypeAlias = int
 _AnnoRowId: TypeAlias = int
+_Identifier: TypeAlias = str
 
-_Key: TypeAlias = tuple[_Name, _TraceId]
-_UniqueBy: TypeAlias = tuple[_Name, _TraceRowId]
+
+class _Key(NamedTuple):
+    annotation_name: _Name
+    annotation_identifier: _Identifier
+    trace_id: _TraceId
+
+
+_UniqueBy: TypeAlias = tuple[_Name, _TraceRowId, _Identifier]
 _Existing: TypeAlias = tuple[
     _TraceRowId,
     _TraceId,
@@ -42,7 +49,7 @@ class TraceAnnotationQueueInserter(
         TraceAnnotationDmlEvent,
     ],
     table=models.TraceAnnotation,
-    unique_by=("name", "trace_rowid"),
+    unique_by=("name", "trace_rowid", "identifier"),
 ):
     async def _events(
         self,
@@ -73,7 +80,11 @@ class TraceAnnotationQueueInserter(
             e.trace_id: _TraceAttr(e.trace_rowid) for e in existing
         }
         existing_annos: Mapping[_Key, _AnnoAttr] = {
-            (e.name, e.trace_id): _AnnoAttr(e.trace_rowid, e.id, e.updated_at)
+            _Key(
+                annotation_name=e.name,
+                annotation_identifier=e.identifier,
+                trace_id=e.trace_id,
+            ): _AnnoAttr(e.trace_rowid, e.id, e.updated_at)
             for e in existing
             if e.id is not None and e.name is not None and e.updated_at is not None
         }
@@ -119,19 +130,20 @@ class TraceAnnotationQueueInserter(
         anno = self.table
         trace = (
             select(models.Trace.id, models.Trace.trace_id)
-            .where(models.Trace.trace_id.in_({trace_id for _, trace_id in keys}))
+            .where(models.Trace.trace_id.in_({k.trace_id for k in keys}))
             .cte()
         )
         onclause = and_(
             trace.c.id == anno.trace_rowid,
-            anno.name.in_({name for name, _ in keys}),
-            tuple_(anno.name, trace.c.trace_id).in_(keys),
+            anno.name.in_({k.annotation_name for k in keys}),
+            tuple_(anno.name, anno.identifier, trace.c.trace_id).in_(keys),
         )
         return select(
             trace.c.id.label("trace_rowid"),
             trace.c.trace_id,
             anno.id,
             anno.name,
+            anno.identifier,
             anno.updated_at,
         ).outerjoin_from(trace, anno, onclause)
 
@@ -147,11 +159,15 @@ class _AnnoAttr(NamedTuple):
 
 
 def _key(p: Received[Precursors.TraceAnnotation]) -> _Key:
-    return p.item.obj.name, p.item.trace_id
+    return _Key(
+        annotation_name=p.item.obj.name,
+        annotation_identifier=p.item.obj.identifier,
+        trace_id=p.item.trace_id,
+    )
 
 
 def _unique_by(p: Received[Insertables.TraceAnnotation]) -> _UniqueBy:
-    return p.item.obj.name, p.item.trace_rowid
+    return p.item.obj.name, p.item.trace_rowid, p.item.identifier
 
 
 def _time(p: Received[Any]) -> datetime:

phoenix/db/insertion/types.py

@@ -50,14 +50,17 @@ class Postponed(Received[_AnyT]):
 class QueueInserter(ABC, Generic[_PrecursorT, _InsertableT, _RowT, _DmlEventT]):
     table: type[_RowT]
     unique_by: Sequence[str]
+    constraint_name: Optional[str] = None
 
     def __init_subclass__(
         cls,
         table: type[_RowT],
         unique_by: Sequence[str],
+        constraint_name: Optional[str] = None,
     ) -> None:
         cls.table = table
         cls.unique_by = unique_by
+        cls.constraint_name = constraint_name
 
     def __init__(
         self,
@@ -109,6 +112,7 @@ class QueueInserter(ABC, Generic[_PrecursorT, _InsertableT, _RowT, _DmlEventT]):
             *records,
             table=self.table,
             unique_by=self.unique_by,
+            constraint_name=self.constraint_name,
             dialect=self._db.dialect,
         )
 
@@ -219,6 +223,7 @@ class Insertables(ABC):
     @dataclass(frozen=True)
     class SpanAnnotation(Precursors.SpanAnnotation):
         span_rowid: int
+        identifier: str = ""
         id_: Optional[int] = None
 
         @property
@@ -232,6 +237,7 @@ class Insertables(ABC):
     @dataclass(frozen=True)
     class TraceAnnotation(Precursors.TraceAnnotation):
         trace_rowid: int
+        identifier: str = ""
         id_: Optional[int] = None
 
         @property
@@ -245,6 +251,7 @@ class Insertables(ABC):
     @dataclass(frozen=True)
     class DocumentAnnotation(Precursors.DocumentAnnotation):
         span_rowid: int
+        identifier: str = ""
         id_: Optional[int] = None
 
         @property

phoenix/db/migrations/versions/2f9d1a65945f_annotation_config_migration.py

@@ -0,0 +1,322 @@
+"""Annotation config migrations
+
+Revision ID: 2f9d1a65945f
+Revises: bc8fea3c2bc8
+Create Date: 2025-02-06 10:17:15.726197
+
+"""
+
+from typing import Any, Sequence, Union
+
+import sqlalchemy as sa
+from alembic import op
+from sqlalchemy import JSON, text
+from sqlalchemy.dialects import postgresql
+from sqlalchemy.ext.compiler import compiles
+
+# revision identifiers, used by Alembic.
+revision: str = "2f9d1a65945f"
+down_revision: Union[str, None] = "bc8fea3c2bc8"
+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",
+    )
+)
+
+
+def upgrade() -> None:
+    with op.batch_alter_table("span_annotations") as batch_op:
+        batch_op.drop_index("ix_span_annotations_score")
+        batch_op.drop_index("ix_span_annotations_label")
+        batch_op.add_column(
+            sa.Column(
+                "user_id",
+                sa.Integer,
+                sa.ForeignKey("users.id", ondelete="SET NULL"),
+                nullable=True,
+            ),
+        )
+        batch_op.add_column(
+            sa.Column(
+                "identifier",
+                sa.String,
+                server_default="",
+                nullable=False,
+            ),
+        )
+        batch_op.add_column(
+            sa.Column(
+                "source",
+                sa.String,
+                nullable=True,
+            ),
+        )
+        batch_op.drop_constraint(
+            constraint_name="valid_annotator_kind",
+            type_="check",
+        )
+        batch_op.create_check_constraint(
+            constraint_name="valid_annotator_kind",
+            condition="annotator_kind IN ('LLM', 'CODE', 'HUMAN')",
+        )
+        batch_op.drop_constraint("uq_span_annotations_name_span_rowid", type_="unique")
+        batch_op.create_unique_constraint(
+            "uq_span_annotations_name_span_rowid_identifier",
+            ["name", "span_rowid", "identifier"],
+        )
+    with op.batch_alter_table("span_annotations") as batch_op:
+        batch_op.execute(
+            text(
+                """
+                UPDATE span_annotations
+                SET source = CASE
+                    WHEN annotator_kind = 'HUMAN' THEN 'APP'
+                    ELSE 'API'
+                END
+                """
+            )
+        )
+        batch_op.alter_column(
+            "source",
+            nullable=False,
+            existing_nullable=True,
+        )
+        batch_op.create_check_constraint(
+            constraint_name="valid_source",
+            condition="source IN ('API', 'APP')",
+        )
+
+    with op.batch_alter_table("trace_annotations") as batch_op:
+        batch_op.drop_index("ix_trace_annotations_score")
+        batch_op.drop_index("ix_trace_annotations_label")
+        batch_op.add_column(
+            sa.Column(
+                "user_id",
+                sa.Integer,
+                sa.ForeignKey("users.id", ondelete="SET NULL"),
+                nullable=True,
+            ),
+        )
+        batch_op.add_column(
+            sa.Column(
+                "identifier",
+                sa.String,
+                server_default="",
+                nullable=False,
+            ),
+        )
+        batch_op.add_column(
+            sa.Column(
+                "source",
+                sa.String,
+                nullable=True,  # must initially be nullable before backfill
+            ),
+        )
+        batch_op.drop_constraint(
+            constraint_name="valid_annotator_kind",
+            type_="check",
+        )
+        batch_op.create_check_constraint(
+            constraint_name="valid_annotator_kind",
+            condition="annotator_kind IN ('LLM', 'CODE', 'HUMAN')",
+        )
+        batch_op.drop_constraint("uq_trace_annotations_name_trace_rowid", type_="unique")
+        batch_op.create_unique_constraint(
+            "uq_trace_annotations_name_trace_rowid_identifier",
+            ["name", "trace_rowid", "identifier"],
+        )
+    with op.batch_alter_table("trace_annotations") as batch_op:
+        batch_op.execute(
+            text(
+                """
+                UPDATE trace_annotations
+                SET source = CASE
+                    WHEN annotator_kind = 'HUMAN' THEN 'APP'
+                    ELSE 'API'
+                END
+                """
+            )
+        )
+        batch_op.alter_column(
+            "source",
+            nullable=False,
+            existing_nullable=True,
+        )
+        batch_op.create_check_constraint(
+            constraint_name="valid_source",
+            condition="source IN ('API', 'APP')",
+        )
+
+    with op.batch_alter_table("document_annotations") as batch_op:
+        batch_op.drop_index("ix_document_annotations_score")
+        batch_op.drop_index("ix_document_annotations_label")
+        batch_op.add_column(
+            sa.Column(
+                "user_id",
+                sa.Integer,
+                sa.ForeignKey("users.id", ondelete="SET NULL"),
+                nullable=True,
+            ),
+        )
+        batch_op.add_column(
+            sa.Column(
+                "identifier",
+                sa.String,
+                server_default="",
+                nullable=False,
+            ),
+        )
+        batch_op.add_column(
+            sa.Column(
+                "source",
+                sa.String,
+                nullable=True,
+            ),
+        )
+        batch_op.drop_constraint(
+            constraint_name="valid_annotator_kind",
+            type_="check",
+        )
+        batch_op.create_check_constraint(
+            constraint_name="valid_annotator_kind",
+            condition="annotator_kind IN ('LLM', 'CODE', 'HUMAN')",
+        )
+        batch_op.drop_constraint(
+            "uq_document_annotations_name_span_rowid_document_position",
+            type_="unique",
+        )
+        batch_op.create_unique_constraint(
+            "uq_document_annotations_name_span_rowid_document_pos_identifier",  # this name does not conform to the auto-generated pattern, which results in a name longer than the Postgres limit of 63 characters  # noqa: E501
+            ["name", "span_rowid", "document_position", "identifier"],
+        )
+    with op.batch_alter_table("document_annotations") as batch_op:
+        batch_op.execute(
+            text(
+                """
+                UPDATE document_annotations
+                SET source = CASE
+                    WHEN annotator_kind = 'HUMAN' THEN 'APP'
+                    ELSE 'API'
+                END
+                """
+            )
+        )
+        batch_op.alter_column(
+            "source",
+            nullable=False,
+            existing_nullable=True,
+        )
+        batch_op.create_check_constraint(
+            constraint_name="valid_source",
+            condition="source IN ('API', 'APP')",
+        )
+
+    op.create_table(
+        "annotation_configs",
+        sa.Column("id", sa.Integer, primary_key=True),
+        sa.Column("name", sa.String, nullable=False, unique=True),
+        sa.Column("config", JSON_, nullable=False),
+    )
+
+    op.create_table(
+        "project_annotation_configs",
+        sa.Column("id", sa.Integer, primary_key=True),
+        sa.Column(
+            "project_id",
+            sa.Integer,
+            sa.ForeignKey("projects.id", ondelete="CASCADE"),
+            nullable=False,
+            index=True,
+        ),
+        sa.Column(
+            "annotation_config_id",
+            sa.Integer,
+            sa.ForeignKey("annotation_configs.id", ondelete="CASCADE"),
+            nullable=False,
+            index=True,
+        ),
+        sa.UniqueConstraint(
+            "project_id",
+            "annotation_config_id",
+        ),
+    )
+
+
+def downgrade() -> None:
+    op.drop_table("project_annotation_configs")
+    op.drop_table("annotation_configs")
+
+    with op.batch_alter_table("document_annotations") as batch_op:
+        batch_op.create_index("ix_document_annotations_score", ["score"])
+        batch_op.create_index("ix_document_annotations_label", ["label"])
+        batch_op.drop_constraint(
+            "uq_document_annotations_name_span_rowid_document_pos_identifier", type_="unique"
+        )
+        batch_op.create_unique_constraint(
+            "uq_document_annotations_name_span_rowid_document_position",
+            ["name", "span_rowid", "document_position"],
+        )
+        batch_op.drop_constraint("valid_annotator_kind", type_="check")
+        batch_op.create_check_constraint(
+            "valid_annotator_kind",
+            condition="annotator_kind IN ('LLM', 'HUMAN')",
+        )
+        batch_op.drop_constraint("valid_source", type_="check")
+        batch_op.drop_column("source")
+        batch_op.drop_column("identifier")
+        batch_op.drop_column("user_id")
+
+    with op.batch_alter_table("trace_annotations") as batch_op:
+        batch_op.create_index("ix_trace_annotations_score", ["score"])
+        batch_op.create_index("ix_trace_annotations_label", ["label"])
+        batch_op.drop_constraint("uq_trace_annotations_name_trace_rowid_identifier", type_="unique")
+        batch_op.create_unique_constraint(
+            "uq_trace_annotations_name_trace_rowid", ["name", "trace_rowid"]
+        )
+        batch_op.drop_constraint("valid_annotator_kind", type_="check")
+        batch_op.create_check_constraint(
+            "valid_annotator_kind",
+            condition="annotator_kind IN ('LLM', 'HUMAN')",
+        )
+        batch_op.drop_constraint("valid_source", type_="check")
+        batch_op.drop_column("source")
+        batch_op.drop_column("identifier")
+        batch_op.drop_column("user_id")
+
+    with op.batch_alter_table("span_annotations") as batch_op:
+        batch_op.create_index("ix_span_annotations_score", ["score"])
+        batch_op.create_index("ix_span_annotations_label", ["label"])
+        batch_op.drop_constraint("uq_span_annotations_name_span_rowid_identifier", type_="unique")
+        batch_op.create_unique_constraint(
+            "uq_span_annotations_name_span_rowid", ["name", "span_rowid"]
+        )
+        batch_op.drop_constraint("valid_annotator_kind", type_="check")
+        batch_op.create_check_constraint(
+            "valid_annotator_kind",
+            condition="annotator_kind IN ('LLM', 'HUMAN')",
+        )
+        batch_op.drop_constraint("valid_source", type_="check")
+        batch_op.drop_column("source")
+        batch_op.drop_column("identifier")
+        batch_op.drop_column("user_id")

phoenix/db/migrations/versions/8a3764fe7f1a_change_jsonb_to_json_for_prompts.py

@@ -0,0 +1,76 @@
+"""change jsonb to json for prompts
+
+Revision ID: 8a3764fe7f1a
+Revises: bb8139330879
+Create Date: 2025-04-25 07:04:26.102957
+
+"""
+
+from typing import Any, Sequence, Union
+
+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 = "8a3764fe7f1a"
+down_revision: Union[str, None] = "bb8139330879"
+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",
+    )
+)
+
+
+def upgrade() -> None:
+    with op.batch_alter_table("prompt_versions") as batch_op:
+        batch_op.alter_column(
+            "tools",
+            type_=JSON,
+            existing_type=JSON_,
+            postgresql_using="tools::json",
+        )
+        batch_op.alter_column(
+            "response_format",
+            type_=JSON,
+            existing_type=JSON_,
+            postgresql_using="response_format::json",
+        )
+
+
+def downgrade() -> None:
+    with op.batch_alter_table("prompt_versions") as batch_op:
+        batch_op.alter_column(
+            "tools",
+            type_=JSON_,
+            existing_type=JSON,
+            postgresql_using="tools::jsonb",
+        )
+        batch_op.alter_column(
+            "response_format",
+            type_=JSON_,
+            existing_type=JSON,
+            postgresql_using="response_format::jsonb",
+        )