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

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

@@ -0,0 +1,317 @@
+"""datasets
+
+Revision ID: 10460e46d750
+Revises: cf03bd6bae1d
+Create Date: 2024-05-10 11:24:23.985834
+
+"""
+
+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(),  # type: ignore
+        "postgresql",
+    )
+    .with_variant(
+        JSONB(),
+        "sqlite",
+    )
+)
+
+# 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", ondelete="SET NULL"),
+            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=False),
+        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/3be8647b87d8_add_token_columns_to_spans_table.py

@@ -0,0 +1,126 @@
+"""add token columns to spans table
+
+Revision ID: 3be8647b87d8
+Revises: 10460e46d750
+Create Date: 2024-08-03 22:11:28.733133
+
+"""
+
+from typing import Any, Optional, Sequence, TypedDict, Union
+
+import sqlalchemy as sa
+from alembic import op
+from openinference.semconv.trace import SpanAttributes
+from sqlalchemy import (
+    JSON,
+    Dialect,
+    MetaData,
+    TypeDecorator,
+    update,
+)
+from sqlalchemy.dialects import postgresql
+from sqlalchemy.ext.compiler import compiles
+from sqlalchemy.orm import (
+    DeclarativeBase,
+    Mapped,
+    mapped_column,
+)
+
+
+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(),  # type: ignore
+        "postgresql",
+    )
+    .with_variant(
+        JSONB(),
+        "sqlite",
+    )
+)
+
+
+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 ExperimentRunOutput(TypedDict, total=False):
+    task_output: Any
+
+
+class Base(DeclarativeBase):
+    # Enforce best practices for naming constraints
+    # https://alembic.sqlalchemy.org/en/latest/naming.html#integration-of-naming-conventions-into-operations-autogenerate
+    metadata = MetaData(
+        naming_convention={
+            "ix": "ix_%(table_name)s_%(column_0_N_name)s",
+            "uq": "uq_%(table_name)s_%(column_0_N_name)s",
+            "ck": "ck_%(table_name)s_`%(constraint_name)s`",
+            "fk": "fk_%(table_name)s_%(column_0_name)s_%(referred_table_name)s",
+            "pk": "pk_%(table_name)s",
+        }
+    )
+    type_annotation_map = {
+        dict[str, Any]: JsonDict,
+        list[dict[str, Any]]: JsonList,
+        ExperimentRunOutput: JsonDict,
+    }
+
+
+class Span(Base):
+    __tablename__ = "spans"
+    id: Mapped[int] = mapped_column(primary_key=True)
+    attributes: Mapped[dict[str, Any]]
+    llm_token_count_prompt: Mapped[Optional[int]]
+    llm_token_count_completion: Mapped[Optional[int]]
+
+
+# revision identifiers, used by Alembic.
+revision: str = "3be8647b87d8"
+down_revision: Union[str, None] = "10460e46d750"
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+LLM_TOKEN_COUNT_PROMPT = SpanAttributes.LLM_TOKEN_COUNT_PROMPT.split(".")
+LLM_TOKEN_COUNT_COMPLETION = SpanAttributes.LLM_TOKEN_COUNT_COMPLETION.split(".")
+
+
+def upgrade() -> None:
+    op.add_column("spans", sa.Column("llm_token_count_prompt", sa.Integer, nullable=True))
+    op.add_column("spans", sa.Column("llm_token_count_completion", sa.Integer, nullable=True))
+    op.execute(
+        update(Span).values(
+            llm_token_count_prompt=Span.attributes[LLM_TOKEN_COUNT_PROMPT].as_float(),
+            llm_token_count_completion=Span.attributes[LLM_TOKEN_COUNT_COMPLETION].as_float(),
+        )
+    )
+
+
+def downgrade() -> None:
+    op.drop_column("spans", "llm_token_count_completion")
+    op.drop_column("spans", "llm_token_count_prompt")

phoenix/db/migrations/versions/4ded9e43755f_create_project_sessions_table.py

@@ -0,0 +1,66 @@
+"""create project_session table
+
+Revision ID: 4ded9e43755f
+Revises: cd164e83824f
+Create Date: 2024-10-08 22:53:24.539786
+
+"""
+
+from typing import Sequence, Union
+
+import sqlalchemy as sa
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision: str = "4ded9e43755f"
+down_revision: Union[str, None] = "cd164e83824f"
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+
+def upgrade() -> None:
+    op.create_table(
+        "project_sessions",
+        sa.Column("id", sa.Integer, primary_key=True),
+        sa.Column("session_id", sa.String, unique=True, nullable=False),
+        sa.Column(
+            "project_id",
+            sa.Integer,
+            sa.ForeignKey("projects.id", ondelete="CASCADE"),
+            nullable=False,
+            index=True,
+        ),
+        sa.Column(
+            "start_time",
+            sa.TIMESTAMP(timezone=True),
+            index=True,
+            nullable=False,
+        ),
+        sa.Column(
+            "end_time",
+            sa.TIMESTAMP(timezone=True),
+            index=True,
+            nullable=False,
+        ),
+    )
+    with op.batch_alter_table("traces") as batch_op:
+        batch_op.add_column(
+            sa.Column(
+                "project_session_rowid",
+                sa.Integer,
+                sa.ForeignKey("project_sessions.id", ondelete="CASCADE"),
+                nullable=True,
+            ),
+        )
+    op.create_index(
+        "ix_traces_project_session_rowid",
+        "traces",
+        ["project_session_rowid"],
+    )
+
+
+def downgrade() -> None:
+    op.drop_index("ix_traces_project_session_rowid")
+    with op.batch_alter_table("traces") as batch_op:
+        batch_op.drop_column("project_session_rowid")
+    op.drop_table("project_sessions")

phoenix/db/migrations/versions/cd164e83824f_users_and_tokens.py

@@ -0,0 +1,157 @@
+"""users and tokens
+
+Revision ID: cd164e83824f
+Revises: 10460e46d750
+Create Date: 2024-08-01 18:36:52.157604
+
+"""
+
+from typing import Sequence, Union
+
+import sqlalchemy as sa
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision: str = "cd164e83824f"
+down_revision: Union[str, None] = "3be8647b87d8"
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+
+def upgrade() -> None:
+    op.create_table(
+        "user_roles",
+        sa.Column("id", sa.Integer, primary_key=True),
+        sa.Column(
+            "name",
+            sa.String,
+            nullable=False,
+            unique=True,
+            index=True,
+        ),
+    )
+    op.create_table(
+        "users",
+        sa.Column("id", sa.Integer, primary_key=True),
+        sa.Column(
+            "user_role_id",
+            sa.Integer,
+            sa.ForeignKey("user_roles.id", ondelete="CASCADE"),
+            nullable=False,
+            index=True,
+        ),
+        sa.Column("username", sa.String, nullable=False, unique=True, index=True),
+        sa.Column("email", sa.String, nullable=False, unique=True, index=True),
+        sa.Column("profile_picture_url", sa.String, nullable=True),
+        sa.Column("password_hash", sa.LargeBinary, nullable=True),
+        sa.Column("password_salt", sa.LargeBinary, nullable=True),
+        sa.Column("reset_password", sa.Boolean, nullable=False),
+        sa.Column("oauth2_client_id", sa.String, nullable=True, index=True),
+        sa.Column("oauth2_user_id", sa.String, nullable=True, 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(),
+        ),
+        sa.CheckConstraint(
+            "(password_hash IS NULL) = (password_salt IS NULL)",
+            name="password_hash_and_salt",
+        ),
+        sa.CheckConstraint(
+            "(oauth2_client_id IS NULL) = (oauth2_user_id IS NULL)",
+            name="oauth2_client_id_and_user_id",
+        ),
+        sa.CheckConstraint(
+            "(password_hash IS NULL) != (oauth2_client_id IS NULL)",
+            name="exactly_one_auth_method",
+        ),
+        sa.UniqueConstraint(
+            "oauth2_client_id",
+            "oauth2_user_id",
+        ),
+        sqlite_autoincrement=True,
+    )
+    op.create_table(
+        "password_reset_tokens",
+        sa.Column("id", sa.Integer, primary_key=True),
+        sa.Column(
+            "user_id",
+            sa.Integer,
+            sa.ForeignKey("users.id", ondelete="CASCADE"),
+            unique=True,
+            index=True,
+        ),
+        sa.Column(
+            "created_at",
+            sa.TIMESTAMP(timezone=True),
+            nullable=False,
+            server_default=sa.func.now(),
+        ),
+        sa.Column("expires_at", sa.TIMESTAMP(timezone=True), nullable=False, index=True),
+        sqlite_autoincrement=True,
+    )
+    op.create_table(
+        "refresh_tokens",
+        sa.Column("id", sa.Integer, primary_key=True),
+        sa.Column("user_id", sa.Integer, sa.ForeignKey("users.id", ondelete="CASCADE"), index=True),
+        sa.Column(
+            "created_at",
+            sa.TIMESTAMP(timezone=True),
+            nullable=False,
+            server_default=sa.func.now(),
+        ),
+        sa.Column("expires_at", sa.TIMESTAMP(timezone=True), nullable=False, index=True),
+        sqlite_autoincrement=True,
+    )
+    op.create_table(
+        "access_tokens",
+        sa.Column("id", sa.Integer, primary_key=True),
+        sa.Column("user_id", sa.Integer, sa.ForeignKey("users.id", ondelete="CASCADE"), index=True),
+        sa.Column(
+            "created_at",
+            sa.TIMESTAMP(timezone=True),
+            nullable=False,
+            server_default=sa.func.now(),
+        ),
+        sa.Column("expires_at", sa.TIMESTAMP(timezone=True), nullable=False, index=True),
+        sa.Column(
+            "refresh_token_id",
+            sa.Integer,
+            sa.ForeignKey("refresh_tokens.id", ondelete="CASCADE"),
+            index=True,
+            unique=True,
+        ),
+        sqlite_autoincrement=True,
+    )
+    op.create_table(
+        "api_keys",
+        sa.Column("id", sa.Integer, primary_key=True),
+        sa.Column("user_id", sa.Integer, sa.ForeignKey("users.id", ondelete="CASCADE"), index=True),
+        sa.Column("name", sa.String, nullable=False),
+        sa.Column("description", sa.String, nullable=True),
+        sa.Column(
+            "created_at",
+            sa.TIMESTAMP(timezone=True),
+            nullable=False,
+            server_default=sa.func.now(),
+        ),
+        sa.Column("expires_at", sa.TIMESTAMP(timezone=True), nullable=True, index=True),
+        sqlite_autoincrement=True,
+    )
+
+
+def downgrade() -> None:
+    op.drop_table("api_keys")
+    op.drop_table("access_tokens")
+    op.drop_table("refresh_tokens")
+    op.drop_table("password_reset_tokens")
+    op.drop_table("users")
+    op.drop_table("user_roles")