arize-phoenix 4.4.4rc5__py3-none-any.whl → 4.5.0__py3-none-any.whl

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

@@ -1,291 +0,0 @@
-"""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", 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/server/api/dataloaders/dataset_example_revisions.py

@@ -1,100 +0,0 @@
-from typing import (
-    AsyncContextManager,
-    Callable,
-    List,
-    Optional,
-    Tuple,
-    Union,
-)
-
-from sqlalchemy import Integer, case, func, literal, or_, select, union
-from sqlalchemy.ext.asyncio import AsyncSession
-from strawberry.dataloader import DataLoader
-from typing_extensions import TypeAlias
-
-from phoenix.db import models
-from phoenix.server.api.types.DatasetExampleRevision import DatasetExampleRevision
-
-ExampleID: TypeAlias = int
-VersionID: TypeAlias = Optional[int]
-Key: TypeAlias = Tuple[ExampleID, Optional[VersionID]]
-Result: TypeAlias = DatasetExampleRevision
-
-
-class DatasetExampleRevisionsDataLoader(DataLoader[Key, Result]):
-    def __init__(self, db: Callable[[], AsyncContextManager[AsyncSession]]) -> None:
-        super().__init__(load_fn=self._load_fn)
-        self._db = db
-
-    async def _load_fn(self, keys: List[Key]) -> List[Union[Result, ValueError]]:
-        # sqlalchemy has limited SQLite support for VALUES, so use UNION ALL instead.
-        # For details, see https://github.com/sqlalchemy/sqlalchemy/issues/7228
-        keys_subquery = union(
-            *(
-                select(
-                    literal(example_id, Integer).label("example_id"),
-                    literal(version_id, Integer).label("version_id"),
-                )
-                for example_id, version_id in keys
-            )
-        ).subquery()
-        revision_ids = (
-            select(
-                keys_subquery.c.example_id,
-                keys_subquery.c.version_id,
-                func.max(models.DatasetExampleRevision.id).label("revision_id"),
-            )
-            .select_from(keys_subquery)
-            .join(
-                models.DatasetExampleRevision,
-                onclause=keys_subquery.c.example_id
-                == models.DatasetExampleRevision.dataset_example_id,
-            )
-            .where(
-                or_(
-                    keys_subquery.c.version_id.is_(None),
-                    models.DatasetExampleRevision.dataset_version_id <= keys_subquery.c.version_id,
-                )
-            )
-            .group_by(keys_subquery.c.example_id, keys_subquery.c.version_id)
-        ).subquery()
-        query = (
-            select(
-                revision_ids.c.example_id,
-                revision_ids.c.version_id,
-                case(
-                    (
-                        or_(
-                            revision_ids.c.version_id.is_(None),
-                            models.DatasetVersion.id.is_not(None),
-                        ),
-                        True,
-                    ),
-                    else_=False,
-                ).label("is_valid_version"),  # check that non-null versions exist
-                models.DatasetExampleRevision,
-            )
-            .select_from(revision_ids)
-            .join(
-                models.DatasetExampleRevision,
-                onclause=revision_ids.c.revision_id == models.DatasetExampleRevision.id,
-            )
-            .join(
-                models.DatasetVersion,
-                onclause=revision_ids.c.version_id == models.DatasetVersion.id,
-                isouter=True,  # keep rows where the version id is null
-            )
-            .where(models.DatasetExampleRevision.revision_kind != "DELETE")
-        )
-        async with self._db() as session:
-            results = {
-                (example_id, version_id): DatasetExampleRevision.from_orm_revision(revision)
-                async for (
-                    example_id,
-                    version_id,
-                    is_valid_version,
-                    revision,
-                ) in await session.stream(query)
-                if is_valid_version
-            }
-        return [results.get(key, ValueError("Could not find revision.")) for key in keys]

phoenix/server/api/dataloaders/dataset_example_spans.py

@@ -1,43 +0,0 @@
-from typing import (
-    AsyncContextManager,
-    Callable,
-    List,
-    Optional,
-)
-
-from sqlalchemy import select
-from sqlalchemy.ext.asyncio import AsyncSession
-from sqlalchemy.orm import joinedload
-from strawberry.dataloader import DataLoader
-from typing_extensions import TypeAlias
-
-from phoenix.db import models
-
-ExampleID: TypeAlias = int
-Key: TypeAlias = ExampleID
-Result: TypeAlias = Optional[models.Span]
-
-
-class DatasetExampleSpansDataLoader(DataLoader[Key, Result]):
-    def __init__(self, db: Callable[[], AsyncContextManager[AsyncSession]]) -> None:
-        super().__init__(load_fn=self._load_fn)
-        self._db = db
-
-    async def _load_fn(self, keys: List[Key]) -> List[Result]:
-        example_ids = keys
-        async with self._db() as session:
-            spans = {
-                example_id: span
-                async for example_id, span in await session.stream(
-                    select(models.DatasetExample.id, models.Span)
-                    .select_from(models.DatasetExample)
-                    .join(models.Span, models.DatasetExample.span_rowid == models.Span.id)
-                    .where(models.DatasetExample.id.in_(example_ids))
-                    .options(
-                        joinedload(models.Span.trace, innerjoin=True).load_only(
-                            models.Trace.trace_id
-                        )
-                    )
-                )
-            }
-        return [spans.get(example_id) for example_id in example_ids]

phoenix/server/api/dataloaders/experiment_annotation_summaries.py

@@ -1,85 +0,0 @@
-from collections import defaultdict
-from dataclasses import dataclass
-from typing import (
-    AsyncContextManager,
-    Callable,
-    DefaultDict,
-    List,
-    Optional,
-)
-
-from sqlalchemy import func, select
-from sqlalchemy.ext.asyncio import AsyncSession
-from strawberry.dataloader import AbstractCache, DataLoader
-from typing_extensions import TypeAlias
-
-from phoenix.db import models
-
-
-@dataclass
-class ExperimentAnnotationSummary:
-    annotation_name: str
-    min_score: float
-    max_score: float
-    mean_score: float
-    count: int
-    error_count: int
-
-
-ExperimentID: TypeAlias = int
-Key: TypeAlias = ExperimentID
-Result: TypeAlias = List[ExperimentAnnotationSummary]
-
-
-class ExperimentAnnotationSummaryDataLoader(DataLoader[Key, Result]):
-    def __init__(
-        self,
-        db: Callable[[], AsyncContextManager[AsyncSession]],
-        cache_map: Optional[AbstractCache[Key, Result]] = None,
-    ) -> None:
-        super().__init__(load_fn=self._load_fn)
-        self._db = db
-
-    async def _load_fn(self, keys: List[Key]) -> List[Result]:
-        experiment_ids = keys
-        summaries: DefaultDict[ExperimentID, Result] = defaultdict(list)
-        async with self._db() as session:
-            async for (
-                experiment_id,
-                annotation_name,
-                min_score,
-                max_score,
-                mean_score,
-                count,
-                error_count,
-            ) in await session.stream(
-                select(
-                    models.ExperimentRun.experiment_id,
-                    models.ExperimentRunAnnotation.name,
-                    func.min(models.ExperimentRunAnnotation.score),
-                    func.max(models.ExperimentRunAnnotation.score),
-                    func.avg(models.ExperimentRunAnnotation.score),
-                    func.count(),
-                    func.count(models.ExperimentRunAnnotation.error),
-                )
-                .join(
-                    models.ExperimentRun,
-                    models.ExperimentRunAnnotation.experiment_run_id == models.ExperimentRun.id,
-                )
-                .where(models.ExperimentRun.experiment_id.in_(experiment_ids))
-                .group_by(models.ExperimentRun.experiment_id, models.ExperimentRunAnnotation.name)
-            ):
-                summaries[experiment_id].append(
-                    ExperimentAnnotationSummary(
-                        annotation_name=annotation_name,
-                        min_score=min_score,
-                        max_score=max_score,
-                        mean_score=mean_score,
-                        count=count,
-                        error_count=error_count,
-                    )
-                )
-        return [
-            sorted(summaries[experiment_id], key=lambda summary: summary.annotation_name)
-            for experiment_id in experiment_ids
-        ]

phoenix/server/api/dataloaders/experiment_error_rates.py

@@ -1,43 +0,0 @@
-from typing import (
-    AsyncContextManager,
-    Callable,
-    List,
-    Optional,
-)
-
-from sqlalchemy import func, select
-from sqlalchemy.ext.asyncio import AsyncSession
-from strawberry.dataloader import DataLoader
-from typing_extensions import TypeAlias
-
-from phoenix.db import models
-
-ExperimentID: TypeAlias = int
-ErrorRate: TypeAlias = float
-Key: TypeAlias = ExperimentID
-Result: TypeAlias = Optional[ErrorRate]
-
-
-class ExperimentErrorRatesDataLoader(DataLoader[Key, Result]):
-    def __init__(
-        self,
-        db: Callable[[], AsyncContextManager[AsyncSession]],
-    ) -> None:
-        super().__init__(load_fn=self._load_fn)
-        self._db = db
-
-    async def _load_fn(self, keys: List[Key]) -> List[Result]:
-        experiment_ids = keys
-        async with self._db() as session:
-            error_rates = {
-                experiment_id: error_rate
-                async for experiment_id, error_rate in await session.stream(
-                    select(
-                        models.ExperimentRun.experiment_id,
-                        func.count(models.ExperimentRun.error) / func.count(),
-                    )
-                    .group_by(models.ExperimentRun.experiment_id)
-                    .where(models.ExperimentRun.experiment_id.in_(experiment_ids))
-                )
-            }
-        return [error_rates.get(experiment_id) for experiment_id in experiment_ids]

phoenix/server/api/dataloaders/experiment_run_counts.py

@@ -1,42 +0,0 @@
-from typing import (
-    AsyncContextManager,
-    Callable,
-    List,
-)
-
-from sqlalchemy import func, select
-from sqlalchemy.ext.asyncio import AsyncSession
-from strawberry.dataloader import DataLoader
-from typing_extensions import TypeAlias
-
-from phoenix.db import models
-
-ExperimentID: TypeAlias = int
-RunCount: TypeAlias = int
-Key: TypeAlias = ExperimentID
-Result: TypeAlias = RunCount
-
-
-class ExperimentRunCountsDataLoader(DataLoader[Key, Result]):
-    def __init__(
-        self,
-        db: Callable[[], AsyncContextManager[AsyncSession]],
-    ) -> None:
-        super().__init__(load_fn=self._load_fn)
-        self._db = db
-
-    async def _load_fn(self, keys: List[Key]) -> List[Result]:
-        experiment_ids = keys
-        async with self._db() as session:
-            run_counts = {
-                experiment_id: run_count
-                async for experiment_id, run_count in await session.stream(
-                    select(models.ExperimentRun.experiment_id, func.count())
-                    .where(models.ExperimentRun.experiment_id.in_(set(experiment_ids)))
-                    .group_by(models.ExperimentRun.experiment_id)
-                )
-            }
-        return [
-            run_counts.get(experiment_id, ValueError(f"Unknown experiment: {experiment_id}"))
-            for experiment_id in experiment_ids
-        ]

phoenix/server/api/dataloaders/experiment_sequence_number.py

@@ -1,49 +0,0 @@
-from typing import (
-    AsyncContextManager,
-    Callable,
-    List,
-    Optional,
-)
-
-from sqlalchemy import distinct, func, select
-from sqlalchemy.ext.asyncio import AsyncSession
-from strawberry.dataloader import DataLoader
-from typing_extensions import TypeAlias
-
-from phoenix.db import models
-
-ExperimentId: TypeAlias = int
-Key: TypeAlias = ExperimentId
-Result: TypeAlias = Optional[int]
-
-
-class ExperimentSequenceNumberDataLoader(DataLoader[Key, Result]):
-    def __init__(self, db: Callable[[], AsyncContextManager[AsyncSession]]) -> None:
-        super().__init__(load_fn=self._load_fn)
-        self._db = db
-
-    async def _load_fn(self, keys: List[Key]) -> List[Result]:
-        experiment_ids = keys
-        dataset_ids = (
-            select(distinct(models.Experiment.dataset_id))
-            .where(models.Experiment.id.in_(experiment_ids))
-            .scalar_subquery()
-        )
-        row_number = (
-            func.row_number().over(
-                partition_by=models.Experiment.dataset_id,
-                order_by=models.Experiment.id,
-            )
-        ).label("row_number")
-        subq = (
-            select(models.Experiment.id, row_number)
-            .where(models.Experiment.dataset_id.in_(dataset_ids))
-            .subquery()
-        )
-        stmt = select(subq).where(subq.c.id.in_(experiment_ids))
-        async with self._db() as session:
-            result = {
-                experiment_id: sequence_number
-                async for experiment_id, sequence_number in await session.stream(stmt)
-            }
-        return [result.get(experiment_id) for experiment_id in experiment_ids]

phoenix/server/api/dataloaders/project_by_name.py

@@ -1,31 +0,0 @@
-from collections import defaultdict
-from typing import AsyncContextManager, Callable, DefaultDict, List, Optional
-
-from sqlalchemy import select
-from sqlalchemy.ext.asyncio import AsyncSession
-from strawberry.dataloader import DataLoader
-from typing_extensions import TypeAlias
-
-from phoenix.db import models
-
-ProjectName: TypeAlias = str
-Key: TypeAlias = ProjectName
-Result: TypeAlias = Optional[models.Project]
-
-
-class ProjectByNameDataLoader(DataLoader[Key, Result]):
-    def __init__(self, db: Callable[[], AsyncContextManager[AsyncSession]]) -> None:
-        super().__init__(load_fn=self._load_fn)
-        self._db = db
-
-    async def _load_fn(self, keys: List[Key]) -> List[Result]:
-        project_names = list(set(keys))
-        projects_by_name: DefaultDict[Key, Result] = defaultdict(None)
-        async with self._db() as session:
-            data = await session.stream_scalars(
-                select(models.Project).where(models.Project.name.in_(project_names))
-            )
-            async for project in data:
-                projects_by_name[project.name] = project
-
-        return [projects_by_name[project_name] for project_name in project_names]