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

phoenix/db/migrations/versions/a20694b15f82_cost.py

@@ -0,0 +1,196 @@
+"""Cost-related tables
+
+Revision ID: a20694b15f82
+Revises: migrations/versions/6a88424799fe_update_users_with_auth_method.py
+Create Date: 2025-05-30 17:15:12.663565
+
+"""
+
+from typing import Sequence, Union
+
+import sqlalchemy as sa
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision: str = "a20694b15f82"
+down_revision: Union[str, None] = "6a88424799fe"
+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(
+        "generative_models",
+        sa.Column(
+            "id",
+            _Integer,
+            primary_key=True,
+        ),
+        sa.Column(
+            "name",
+            sa.String,
+            nullable=False,
+        ),
+        sa.Column(
+            "provider",
+            sa.String,
+            nullable=False,
+        ),
+        sa.Column(
+            "name_pattern",
+            sa.String,
+            nullable=False,
+        ),
+        sa.Column(
+            "is_built_in",
+            sa.Boolean,
+            nullable=False,
+        ),
+        sa.Column(
+            "start_time",
+            sa.TIMESTAMP(timezone=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(),
+        ),
+        sa.Column(
+            "deleted_at",
+            sa.TIMESTAMP(timezone=True),
+            nullable=True,
+        ),
+    )
+    op.create_index(
+        "ix_generative_models_match_criteria",
+        "generative_models",
+        ["name_pattern", "provider", "is_built_in"],
+        unique=True,
+        postgresql_where=sa.text("deleted_at IS NULL"),
+        sqlite_where=sa.text("deleted_at IS NULL"),
+    )
+    op.create_index(
+        "ix_generative_models_name_is_built_in",
+        "generative_models",
+        ["name", "is_built_in"],
+        unique=True,
+        postgresql_where=sa.text("deleted_at IS NULL"),
+        sqlite_where=sa.text("deleted_at IS NULL"),
+    )
+    op.create_table(
+        "token_prices",
+        sa.Column(
+            "id",
+            _Integer,
+            primary_key=True,
+        ),
+        sa.Column(
+            "model_id",
+            _Integer,
+            sa.ForeignKey("generative_models.id", ondelete="CASCADE"),
+            nullable=False,
+            index=True,
+        ),
+        sa.Column("token_type", sa.String, nullable=False),
+        sa.Column("is_prompt", sa.Boolean, nullable=False),
+        sa.Column("base_rate", sa.Float, nullable=False),
+        sa.Column("customization", sa.JSON),
+        sa.UniqueConstraint(
+            "model_id",
+            "token_type",
+            "is_prompt",
+        ),
+    )
+    op.create_table(
+        "span_costs",
+        sa.Column(
+            "id",
+            _Integer,
+            primary_key=True,
+        ),
+        sa.Column(
+            "span_rowid",
+            _Integer,
+            sa.ForeignKey("spans.id", ondelete="CASCADE"),
+            nullable=False,
+            index=True,
+        ),
+        sa.Column(
+            "trace_rowid",
+            _Integer,
+            sa.ForeignKey("traces.id", ondelete="CASCADE"),
+            nullable=False,
+            index=True,
+        ),
+        sa.Column(
+            "model_id",
+            _Integer,
+            sa.ForeignKey(
+                "generative_models.id",
+                ondelete="RESTRICT",
+            ),
+            nullable=True,
+        ),
+        sa.Column(
+            "span_start_time",
+            sa.TIMESTAMP(timezone=True),
+            nullable=False,
+            index=True,
+        ),
+        sa.Column("total_cost", sa.Float),
+        sa.Column("total_tokens", sa.Float),
+        sa.Column("prompt_cost", sa.Float),
+        sa.Column("prompt_tokens", sa.Float),
+        sa.Column("completion_cost", sa.Float),
+        sa.Column("completion_tokens", sa.Float),
+    )
+    op.create_index(
+        "ix_span_costs_model_id_span_start_time",
+        "span_costs",
+        ["model_id", "span_start_time"],
+    )
+    op.create_table(
+        "span_cost_details",
+        sa.Column(
+            "id",
+            _Integer,
+            primary_key=True,
+        ),
+        sa.Column(
+            "span_cost_id",
+            _Integer,
+            sa.ForeignKey("span_costs.id", ondelete="CASCADE"),
+            nullable=False,
+            index=True,
+        ),
+        sa.Column("token_type", sa.String, nullable=False, index=True),
+        sa.Column("is_prompt", sa.Boolean, nullable=False),
+        sa.Column("cost", sa.Float),
+        sa.Column("tokens", sa.Float),
+        sa.Column("cost_per_token", sa.Float),
+        sa.UniqueConstraint(
+            "span_cost_id",
+            "token_type",
+            "is_prompt",
+        ),
+    )
+
+
+def downgrade() -> None:
+    op.drop_table("span_cost_details")
+    op.drop_table("span_costs")
+    op.drop_table("token_prices")
+    op.drop_table("generative_models")

phoenix/db/migrations/versions/ab513d89518b_add_user_id_on_dataset_versions.py

@@ -0,0 +1,40 @@
+"""add user_id on dataset_versions
+
+Revision ID: ab513d89518b
+Revises: 01a8342c9cdf
+Create Date: 2025-09-26 11:00:06.961920
+
+"""
+
+from typing import Sequence, Union
+
+import sqlalchemy as sa
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision: str = "ab513d89518b"
+down_revision: Union[str, None] = "01a8342c9cdf"
+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("dataset_versions") 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("dataset_versions") as batch_op:
+        batch_op.drop_column("user_id")

phoenix/db/migrations/versions/d0690a79ea51_users_on_experiments.py

@@ -0,0 +1,40 @@
+"""add user_id on experiments
+
+Revision ID: d0690a79ea51
+Revises: ab513d89518b
+Create Date: 2025-08-26 19:12:47.849806
+
+"""
+
+from typing import Sequence, Union
+
+import sqlalchemy as sa
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision: str = "d0690a79ea51"
+down_revision: Union[str, None] = "ab513d89518b"
+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("experiments") 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("experiments") as batch_op:
+        batch_op.drop_column("user_id")

phoenix/db/migrations/versions/deb2c81c0bb2_dataset_splits.py

@@ -0,0 +1,139 @@
+"""dataset_splits
+
+Revision ID: deb2c81c0bb2
+Revises: 58228d933c91
+Create Date: 2025-09-08 15:50:12.066217
+
+"""
+
+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
+
+
+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",
+)
+
+
+# revision identifiers, used by Alembic.
+revision: str = "deb2c81c0bb2"
+down_revision: Union[str, None] = "e76cbd66ffc3"
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+
+def upgrade() -> None:
+    # Create dataset_splits table
+    op.create_table(
+        "dataset_splits",
+        sa.Column("id", _Integer, primary_key=True),
+        sa.Column(
+            "user_id",
+            _Integer,
+            sa.ForeignKey("users.id", ondelete="SET NULL"),
+            nullable=True,
+            index=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("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(),
+        ),
+    )
+
+    # Create crosswalk table: dataset_splits_dataset_examples
+    op.create_table(
+        "dataset_splits_dataset_examples",
+        sa.Column(
+            "dataset_split_id",
+            _Integer,
+            sa.ForeignKey("dataset_splits.id", ondelete="CASCADE"),
+            nullable=False,
+        ),
+        sa.Column(
+            "dataset_example_id",
+            _Integer,
+            sa.ForeignKey("dataset_examples.id", ondelete="CASCADE"),
+            nullable=False,
+            # index on the second element of the composite primary key
+            index=True,
+        ),
+        sa.PrimaryKeyConstraint(
+            "dataset_split_id",
+            "dataset_example_id",
+        ),
+    )
+
+    # Create experiments_dataset_splits table The rational of this table is to
+    # gather examples for a specific dataset split for a specific experiment.
+    # Select all dataset examples where examples belong to a dataset split and
+    # examples belong to a experiment.
+
+    op.create_table(
+        "experiments_dataset_splits",
+        sa.Column(
+            "experiment_id",
+            _Integer,
+            sa.ForeignKey("experiments.id", ondelete="CASCADE"),
+            nullable=False,
+        ),
+        sa.Column(
+            "dataset_split_id",
+            _Integer,
+            sa.ForeignKey("dataset_splits.id", ondelete="CASCADE"),
+            nullable=False,
+            # index on the second element of the composite primary key
+            index=True,
+        ),
+        sa.PrimaryKeyConstraint(
+            "experiment_id",
+            "dataset_split_id",
+        ),
+    )
+
+
+def downgrade() -> None:
+    op.drop_table("experiments_dataset_splits")
+    op.drop_table("dataset_splits_dataset_examples")
+    op.drop_table("dataset_splits")

phoenix/db/migrations/versions/e76cbd66ffc3_add_experiments_dataset_examples.py

@@ -0,0 +1,87 @@
+"""add experiments_dataset_examples junction table
+
+Revision ID: e76cbd66ffc3
+Revises: deb2c81c0bb2
+Create Date: 2025-09-23 12:33:13.554164
+
+"""
+
+from typing import Sequence, Union
+
+import sqlalchemy as sa
+from alembic import op
+
+_Integer = sa.Integer().with_variant(
+    sa.BigInteger(),
+    "postgresql",
+)
+
+# revision identifiers, used by Alembic.
+revision: str = "e76cbd66ffc3"
+down_revision: Union[str, None] = "58228d933c91"
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+BACKFILL = """\
+INSERT INTO experiments_dataset_examples (
+    experiment_id,
+    dataset_example_id,
+    dataset_example_revision_id
+)
+SELECT
+    ranked.experiment_id,
+    ranked.dataset_example_id,
+    ranked.dataset_example_revision_id
+FROM (
+    SELECT
+        e.id as experiment_id,
+        der.dataset_example_id,
+        der.id as dataset_example_revision_id,
+        der.revision_kind,
+        ROW_NUMBER() OVER (
+            PARTITION BY e.id, der.dataset_example_id
+            ORDER BY der.dataset_version_id DESC
+        ) as rn
+    FROM experiments e
+        JOIN dataset_examples de ON de.dataset_id = e.dataset_id
+        JOIN dataset_example_revisions der ON der.dataset_example_id = de.id
+    WHERE der.dataset_version_id <= e.dataset_version_id
+) ranked
+WHERE ranked.rn = 1
+    AND ranked.revision_kind != 'DELETE'
+"""
+
+
+def upgrade() -> None:
+    op.create_table(
+        "experiments_dataset_examples",
+        sa.Column(
+            "experiment_id",
+            _Integer,
+            sa.ForeignKey("experiments.id", ondelete="CASCADE"),
+            nullable=False,
+        ),
+        sa.Column(
+            "dataset_example_id",
+            _Integer,
+            sa.ForeignKey("dataset_examples.id", ondelete="CASCADE"),
+            nullable=False,
+            index=True,
+        ),
+        sa.Column(
+            "dataset_example_revision_id",
+            _Integer,
+            sa.ForeignKey("dataset_example_revisions.id", ondelete="CASCADE"),
+            nullable=False,
+            index=True,
+        ),
+        sa.PrimaryKeyConstraint(
+            "experiment_id",
+            "dataset_example_id",
+        ),
+    )
+    op.execute(BACKFILL)
+
+
+def downgrade() -> None:
+    op.drop_table("experiments_dataset_examples")