arize-phoenix 4.4.3__py3-none-any.whl → 4.4.4rc1__py3-none-any.whl

phoenix/db/migrations/versions/10460e46d750_datasets.py

@@ -0,0 +1,291 @@
+"""datasets
+
+Revision ID: 10460e46d750
+Revises: cf03bd6bae1d
+Create Date: 2024-05-10 11:24:23.985834
+
+"""
+
+from typing import Sequence, Union
+
+import sqlalchemy as sa
+from alembic import op
+from phoenix.db.migrations.types import JSON_
+
+# revision identifiers, used by Alembic.
+revision: str = "10460e46d750"
+down_revision: Union[str, None] = "cf03bd6bae1d"
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+
+def upgrade() -> None:
+    op.create_table(
+        "datasets",
+        sa.Column("id", sa.Integer, primary_key=True),
+        sa.Column("name", sa.String, nullable=False, unique=True),
+        sa.Column("description", sa.String, nullable=True),
+        sa.Column("metadata", JSON_, nullable=False),
+        sa.Column(
+            "created_at",
+            sa.TIMESTAMP(timezone=True),
+            nullable=False,
+            server_default=sa.func.now(),
+        ),
+        sa.Column(
+            "updated_at",
+            sa.TIMESTAMP(timezone=True),
+            nullable=False,
+            server_default=sa.func.now(),
+            onupdate=sa.func.now(),
+        ),
+    )
+    op.create_table(
+        "dataset_versions",
+        sa.Column("id", sa.Integer, primary_key=True),
+        sa.Column(
+            "dataset_id",
+            sa.Integer,
+            sa.ForeignKey("datasets.id", ondelete="CASCADE"),
+            nullable=False,
+            index=True,
+        ),
+        sa.Column("description", sa.String, nullable=True),
+        sa.Column("metadata", JSON_, nullable=False),
+        sa.Column(
+            "created_at",
+            sa.TIMESTAMP(timezone=True),
+            nullable=False,
+            server_default=sa.func.now(),
+        ),
+    )
+    op.create_table(
+        "dataset_examples",
+        sa.Column("id", sa.Integer, primary_key=True),
+        sa.Column(
+            "dataset_id",
+            sa.Integer,
+            sa.ForeignKey("datasets.id", ondelete="CASCADE"),
+            nullable=False,
+            index=True,
+        ),
+        sa.Column(
+            "span_rowid",
+            sa.Integer,
+            sa.ForeignKey("spans.id"),
+            nullable=True,
+            index=True,
+        ),
+        sa.Column(
+            "created_at",
+            sa.TIMESTAMP(timezone=True),
+            nullable=False,
+            server_default=sa.func.now(),
+        ),
+    )
+    op.create_table(
+        "dataset_example_revisions",
+        sa.Column("id", sa.Integer, primary_key=True),
+        sa.Column(
+            "dataset_example_id",
+            sa.Integer,
+            sa.ForeignKey("dataset_examples.id", ondelete="CASCADE"),
+            nullable=False,
+            index=True,
+        ),
+        sa.Column(
+            "dataset_version_id",
+            sa.Integer,
+            sa.ForeignKey("dataset_versions.id", ondelete="CASCADE"),
+            nullable=False,
+            index=True,
+        ),
+        sa.Column("input", JSON_, nullable=False),
+        sa.Column("output", JSON_, nullable=False),
+        sa.Column("metadata", JSON_, nullable=False),
+        sa.Column(
+            "revision_kind",
+            sa.String,
+            sa.CheckConstraint(
+                "revision_kind IN ('CREATE', 'PATCH', 'DELETE')",
+                name="valid_revision_kind",
+            ),
+            nullable=False,
+        ),
+        sa.Column(
+            "created_at",
+            sa.TIMESTAMP(timezone=True),
+            nullable=False,
+            server_default=sa.func.now(),
+        ),
+        sa.UniqueConstraint(
+            "dataset_example_id",
+            "dataset_version_id",
+        ),
+    )
+    op.create_table(
+        "experiments",
+        sa.Column("id", sa.Integer, primary_key=True),
+        sa.Column(
+            "dataset_id",
+            sa.Integer,
+            sa.ForeignKey("datasets.id", ondelete="CASCADE"),
+            nullable=False,
+            index=True,
+        ),
+        sa.Column(
+            "dataset_version_id",
+            sa.Integer,
+            sa.ForeignKey("dataset_versions.id", ondelete="CASCADE"),
+            nullable=False,
+            index=True,
+        ),
+        sa.Column(
+            "name",
+            sa.String,
+            nullable=False,
+        ),
+        sa.Column(
+            "description",
+            sa.String,
+            nullable=True,
+        ),
+        sa.Column(
+            "repetitions",
+            sa.Integer,
+            nullable=False,
+        ),
+        sa.Column("metadata", JSON_, nullable=False),
+        sa.Column("project_name", sa.String, index=True),
+        sa.Column(
+            "created_at",
+            sa.TIMESTAMP(timezone=True),
+            nullable=False,
+            server_default=sa.func.now(),
+        ),
+        sa.Column(
+            "updated_at",
+            sa.TIMESTAMP(timezone=True),
+            nullable=False,
+            server_default=sa.func.now(),
+            onupdate=sa.func.now(),
+        ),
+    )
+    op.create_table(
+        "experiment_runs",
+        sa.Column("id", sa.Integer, primary_key=True),
+        sa.Column(
+            "experiment_id",
+            sa.Integer,
+            sa.ForeignKey("experiments.id", ondelete="CASCADE"),
+            nullable=False,
+            index=True,
+        ),
+        sa.Column(
+            "dataset_example_id",
+            sa.Integer,
+            sa.ForeignKey("dataset_examples.id", ondelete="CASCADE"),
+            nullable=False,
+            index=True,
+        ),
+        sa.Column(
+            "repetition_number",
+            sa.Integer,
+            nullable=False,
+        ),
+        sa.Column(
+            "trace_id",
+            sa.String,
+            nullable=True,
+        ),
+        sa.Column("output", JSON_, nullable=True),
+        sa.Column("start_time", sa.TIMESTAMP(timezone=True), nullable=False),
+        sa.Column("end_time", sa.TIMESTAMP(timezone=True), nullable=False),
+        sa.Column(
+            "prompt_token_count",
+            sa.Integer,
+            nullable=True,
+        ),
+        sa.Column(
+            "completion_token_count",
+            sa.Integer,
+            nullable=True,
+        ),
+        sa.Column(
+            "error",
+            sa.String,
+            nullable=True,
+        ),
+        sa.UniqueConstraint(
+            "experiment_id",
+            "dataset_example_id",
+            "repetition_number",
+        ),
+    )
+    op.create_table(
+        "experiment_run_annotations",
+        sa.Column("id", sa.Integer, primary_key=True),
+        sa.Column(
+            "experiment_run_id",
+            sa.Integer,
+            sa.ForeignKey("experiment_runs.id", ondelete="CASCADE"),
+            nullable=False,
+            index=True,
+        ),
+        sa.Column(
+            "name",
+            sa.String,
+            nullable=False,
+        ),
+        sa.Column(
+            "annotator_kind",
+            sa.String,
+            sa.CheckConstraint(
+                "annotator_kind IN ('LLM', 'CODE', 'HUMAN')",
+                name="valid_annotator_kind",
+            ),
+            nullable=False,
+        ),
+        sa.Column(
+            "label",
+            sa.String,
+            nullable=True,
+        ),
+        sa.Column(
+            "score",
+            sa.Float,
+            nullable=True,
+        ),
+        sa.Column(
+            "explanation",
+            sa.String,
+            nullable=True,
+        ),
+        sa.Column(
+            "trace_id",
+            sa.String,
+            nullable=True,
+        ),
+        sa.Column(
+            "error",
+            sa.String,
+            nullable=True,
+        ),
+        sa.Column("metadata", JSON_, nullable=False),
+        sa.Column("start_time", sa.TIMESTAMP(timezone=True), nullable=False),
+        sa.Column("end_time", sa.TIMESTAMP(timezone=True), nullable=False),
+        sa.UniqueConstraint(
+            "experiment_run_id",
+            "name",
+        ),
+    )
+
+
+def downgrade() -> None:
+    op.drop_table("experiment_run_annotations")
+    op.drop_table("experiment_runs")
+    op.drop_table("experiments")
+    op.drop_table("dataset_example_revisions")
+    op.drop_table("dataset_examples")
+    op.drop_table("dataset_versions")
+    op.drop_table("datasets")

phoenix/db/migrations/versions/cf03bd6bae1d_init.py

@@ -6,13 +6,11 @@ Create Date: 2024-04-03 19:41:48.871555
 
 """
 
-from typing import Any, Sequence, Union
+from typing import 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
+from phoenix.db.migrations.types import JSON_
 
 # revision identifiers, used by Alembic.
 revision: str = "cf03bd6bae1d"
@@ -21,30 +19,6 @@ 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")  # type: ignore
-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(),  # type: ignore
-        "postgresql",
-    )
-    .with_variant(
-        JSONB(),
-        "sqlite",
-    )
-)
-
-
 def upgrade() -> None:
     projects_table = op.create_table(
         "projects",

phoenix/db/models.py

@@ -15,12 +15,14 @@ from sqlalchemy import (
     String,
     TypeDecorator,
     UniqueConstraint,
+    case,
     func,
     insert,
+    select,
     text,
 )
 from sqlalchemy.dialects import postgresql
-from sqlalchemy.ext.asyncio import AsyncEngine
+from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession
 from sqlalchemy.ext.compiler import compiles
 from sqlalchemy.ext.hybrid import hybrid_property
 from sqlalchemy.orm import (
@@ -59,6 +61,24 @@ JSON_ = (
 )
 
 
+class JsonDict(TypeDecorator[Dict[str, Any]]):
+    # See # See https://docs.sqlalchemy.org/en/20/core/custom_types.html
+    cache_ok = True
+    impl = JSON_
+
+    def process_bind_param(self, value: Optional[Dict[str, Any]], _: Dialect) -> Dict[str, Any]:
+        return value if isinstance(value, dict) else {}
+
+
+class JsonList(TypeDecorator[List[Any]]):
+    # See # See https://docs.sqlalchemy.org/en/20/core/custom_types.html
+    cache_ok = True
+    impl = JSON_
+
+    def process_bind_param(self, value: Optional[List[Any]], _: Dialect) -> List[Any]:
+        return value if isinstance(value, list) else []
+
+
 class UtcTimeStamp(TypeDecorator[datetime]):
     # See # See https://docs.sqlalchemy.org/en/20/core/custom_types.html
     cache_ok = True
@@ -84,8 +104,8 @@ class Base(DeclarativeBase):
         }
     )
     type_annotation_map = {
-        Dict[str, Any]: JSON_,
-        List[Dict[str, Any]]: JSON_,
+        Dict[str, Any]: JsonDict,
+        List[Dict[str, Any]]: JsonList,
     }
 
 
@@ -154,6 +174,10 @@ class Trace(Base):
         cascade="all, delete-orphan",
         uselist=True,
     )
+    experiment_runs: Mapped[List["ExperimentRun"]] = relationship(
+        primaryjoin="foreign(ExperimentRun.trace_id) == Trace.trace_id",
+        back_populates="trace",
+    )
     __table_args__ = (
         UniqueConstraint(
             "trace_id",
@@ -203,6 +227,7 @@ class Span(Base):
 
     trace: Mapped["Trace"] = relationship("Trace", back_populates="spans")
     document_annotations: Mapped[List["DocumentAnnotation"]] = relationship(back_populates="span")
+    dataset_examples: Mapped[List["DatasetExample"]] = relationship(back_populates="span")
 
     __table_args__ = (
         UniqueConstraint(
@@ -376,3 +401,205 @@ class DocumentAnnotation(Base):
             "document_position",
         ),
     )
+
+
+class Dataset(Base):
+    __tablename__ = "datasets"
+    id: Mapped[int] = mapped_column(primary_key=True)
+    name: Mapped[str] = mapped_column(unique=True)
+    description: Mapped[Optional[str]]
+    metadata_: Mapped[Dict[str, Any]] = mapped_column("metadata")
+    created_at: Mapped[datetime] = mapped_column(UtcTimeStamp, server_default=func.now())
+    updated_at: Mapped[datetime] = mapped_column(
+        UtcTimeStamp, server_default=func.now(), onupdate=func.now()
+    )
+
+    @hybrid_property
+    def example_count(self) -> Optional[int]:
+        if hasattr(self, "_example_count_value"):
+            assert isinstance(self._example_count_value, int)
+            return self._example_count_value
+        return None
+
+    @example_count.inplace.expression
+    def _example_count(cls) -> ColumnElement[int]:
+        return (
+            select(
+                func.sum(
+                    case(
+                        (DatasetExampleRevision.revision_kind == "CREATE", 1),
+                        (DatasetExampleRevision.revision_kind == "DELETE", -1),
+                        else_=0,
+                    )
+                )
+            )
+            .select_from(DatasetExampleRevision)
+            .join(
+                DatasetExample,
+                onclause=DatasetExample.id == DatasetExampleRevision.dataset_example_id,
+            )
+            .filter(DatasetExample.dataset_id == cls.id)
+            .label("example_count")
+        )
+
+    async def load_example_count(self, session: AsyncSession) -> None:
+        if not hasattr(self, "_example_count_value"):
+            self._example_count_value = await session.scalar(
+                select(
+                    func.sum(
+                        case(
+                            (DatasetExampleRevision.revision_kind == "CREATE", 1),
+                            (DatasetExampleRevision.revision_kind == "DELETE", -1),
+                            else_=0,
+                        )
+                    )
+                )
+                .select_from(DatasetExampleRevision)
+                .join(
+                    DatasetExample,
+                    onclause=DatasetExample.id == DatasetExampleRevision.dataset_example_id,
+                )
+                .filter(DatasetExample.dataset_id == self.id)
+            )
+
+
+class DatasetVersion(Base):
+    __tablename__ = "dataset_versions"
+    id: Mapped[int] = mapped_column(primary_key=True)
+    dataset_id: Mapped[int] = mapped_column(
+        ForeignKey("datasets.id", ondelete="CASCADE"),
+        index=True,
+    )
+    description: Mapped[Optional[str]]
+    metadata_: Mapped[Dict[str, Any]] = mapped_column("metadata")
+    created_at: Mapped[datetime] = mapped_column(UtcTimeStamp, server_default=func.now())
+
+
+class DatasetExample(Base):
+    __tablename__ = "dataset_examples"
+    id: Mapped[int] = mapped_column(primary_key=True)
+    dataset_id: Mapped[int] = mapped_column(
+        ForeignKey("datasets.id", ondelete="CASCADE"),
+        index=True,
+    )
+    span_rowid: Mapped[Optional[int]] = mapped_column(
+        ForeignKey("spans.id"),
+        index=True,
+        nullable=True,
+    )
+    created_at: Mapped[datetime] = mapped_column(UtcTimeStamp, server_default=func.now())
+
+    span: Mapped[Optional[Span]] = relationship(back_populates="dataset_examples")
+
+
+class DatasetExampleRevision(Base):
+    __tablename__ = "dataset_example_revisions"
+    id: Mapped[int] = mapped_column(primary_key=True)
+    dataset_example_id: Mapped[int] = mapped_column(
+        ForeignKey("dataset_examples.id", ondelete="CASCADE"),
+        index=True,
+    )
+    dataset_version_id: Mapped[int] = mapped_column(
+        ForeignKey("dataset_versions.id", ondelete="CASCADE"),
+        index=True,
+    )
+    input: Mapped[Dict[str, Any]]
+    output: Mapped[Dict[str, Any]]
+    metadata_: Mapped[Dict[str, Any]] = mapped_column("metadata")
+    revision_kind: Mapped[str] = mapped_column(
+        CheckConstraint(
+            "revision_kind IN ('CREATE', 'PATCH', 'DELETE')", name="valid_revision_kind"
+        ),
+    )
+    created_at: Mapped[datetime] = mapped_column(UtcTimeStamp, server_default=func.now())
+
+    __table_args__ = (
+        UniqueConstraint(
+            "dataset_example_id",
+            "dataset_version_id",
+        ),
+    )
+
+
+class Experiment(Base):
+    __tablename__ = "experiments"
+    id: Mapped[int] = mapped_column(primary_key=True)
+    dataset_id: Mapped[int] = mapped_column(
+        ForeignKey("datasets.id", ondelete="CASCADE"),
+        index=True,
+    )
+    dataset_version_id: Mapped[int] = mapped_column(
+        ForeignKey("dataset_versions.id", ondelete="CASCADE"),
+        index=True,
+    )
+    name: Mapped[str]
+    description: Mapped[Optional[str]]
+    repetitions: Mapped[int]
+    metadata_: Mapped[Dict[str, Any]] = mapped_column("metadata")
+    project_name: Mapped[Optional[str]] = mapped_column(index=True)
+    created_at: Mapped[datetime] = mapped_column(UtcTimeStamp, server_default=func.now())
+    updated_at: Mapped[datetime] = mapped_column(
+        UtcTimeStamp, server_default=func.now(), onupdate=func.now()
+    )
+
+
+class ExperimentRun(Base):
+    __tablename__ = "experiment_runs"
+    id: Mapped[int] = mapped_column(primary_key=True)
+    experiment_id: Mapped[int] = mapped_column(
+        ForeignKey("experiments.id", ondelete="CASCADE"),
+        index=True,
+    )
+    dataset_example_id: Mapped[int] = mapped_column(
+        ForeignKey("dataset_examples.id", ondelete="CASCADE"),
+        index=True,
+    )
+    repetition_number: Mapped[int]
+    trace_id: Mapped[Optional[str]]
+    output: Mapped[Optional[Dict[str, Any]]]
+    start_time: Mapped[datetime] = mapped_column(UtcTimeStamp)
+    end_time: Mapped[datetime] = mapped_column(UtcTimeStamp)
+    prompt_token_count: Mapped[Optional[int]]
+    completion_token_count: Mapped[Optional[int]]
+    error: Mapped[Optional[str]]
+
+    trace: Mapped["Trace"] = relationship(
+        primaryjoin="foreign(ExperimentRun.trace_id) == Trace.trace_id",
+        back_populates="experiment_runs",
+    )
+
+    __table_args__ = (
+        UniqueConstraint(
+            "experiment_id",
+            "dataset_example_id",
+            "repetition_number",
+        ),
+    )
+
+
+class ExperimentRunAnnotation(Base):
+    __tablename__ = "experiment_run_annotations"
+    id: Mapped[int] = mapped_column(primary_key=True)
+    experiment_run_id: Mapped[int] = mapped_column(
+        ForeignKey("experiment_runs.id", ondelete="CASCADE"),
+        index=True,
+    )
+    name: Mapped[str]
+    annotator_kind: Mapped[str] = mapped_column(
+        CheckConstraint("annotator_kind IN ('LLM', 'CODE', 'HUMAN')", name="valid_annotator_kind"),
+    )
+    label: Mapped[Optional[str]]
+    score: Mapped[Optional[float]]
+    explanation: Mapped[Optional[str]]
+    trace_id: Mapped[Optional[str]]
+    error: Mapped[Optional[str]]
+    metadata_: Mapped[Dict[str, Any]] = mapped_column("metadata")
+    start_time: Mapped[datetime] = mapped_column(UtcTimeStamp)
+    end_time: Mapped[datetime] = mapped_column(UtcTimeStamp)
+
+    __table_args__ = (
+        UniqueConstraint(
+            "experiment_run_id",
+            "name",
+        ),
+    )