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

phoenix/db/iam_auth.py

@@ -0,0 +1,64 @@
+from __future__ import annotations
+
+import logging
+
+logger = logging.getLogger(__name__)
+
+
+def generate_aws_rds_token(
+    host: str,
+    port: int,
+    user: str,
+) -> str:
+    """Generate an AWS RDS IAM authentication token.
+
+    This function creates a short-lived (15 minutes) authentication token for connecting
+    to AWS RDS/Aurora PostgreSQL instances using IAM database authentication.
+
+    The AWS region is automatically resolved using boto3.
+
+    Args:
+        host: The database hostname (e.g., 'mydb.abc123.us-west-2.rds.amazonaws.com')
+        port: The database port (typically 5432 for PostgreSQL)
+        user: The database username (must match an IAM-enabled database user)
+
+    Returns:
+        A temporary authentication token string to use as the database password
+
+    Raises:
+        ImportError: If boto3 is not installed
+        Exception: If AWS credentials/region are not configured or token generation fails
+
+    Example:
+        >>> token = generate_aws_rds_token(
+        ...     host='mydb.us-west-2.rds.amazonaws.com',
+        ...     port=5432,
+        ...     user='myuser'
+        ... )
+    """
+    try:
+        import boto3  # type: ignore
+    except ImportError as e:
+        raise ImportError(
+            "boto3 is required for AWS RDS IAM authentication. "
+            "Install it with: pip install 'arize-phoenix[aws]'"
+        ) from e
+
+    try:
+        client = boto3.client("rds")
+
+        logger.debug(f"Generating AWS RDS IAM auth token for user '{user}' at {host}:{port}")
+        token = client.generate_db_auth_token(  # pyright: ignore
+            DBHostname=host,
+            Port=port,
+            DBUsername=user,
+        )
+
+        return str(token)  # pyright: ignore
+
+    except Exception as e:
+        logger.error(
+            f"Failed to generate AWS RDS IAM authentication token: {e}. "
+            "Ensure AWS credentials are configured and have 'rds-db:connect' permission."
+        )
+        raise

phoenix/db/insertion/dataset.py

@@ -11,7 +11,8 @@ from sqlalchemy.ext.asyncio import AsyncSession
 from typing_extensions import TypeAlias
 
 from phoenix.db import models
-from phoenix.db.insertion.helpers import DataManipulationEvent
+from phoenix.db.helpers import SupportedSQLDialect
+from phoenix.db.insertion.helpers import DataManipulationEvent, OnConflict, insert_on_conflict
 
 logger = logging.getLogger(__name__)
 
@@ -27,6 +28,7 @@ class ExampleContent:
     input: dict[str, Any] = field(default_factory=dict)
     output: dict[str, Any] = field(default_factory=dict)
     metadata: dict[str, Any] = field(default_factory=dict)
+    splits: frozenset[str] = field(default_factory=frozenset)  # Set of split names
 
 
 Examples: TypeAlias = Iterable[ExampleContent]
@@ -44,6 +46,7 @@ async def insert_dataset(
     description: Optional[str] = None,
     metadata: Optional[Mapping[str, Any]] = None,
     created_at: Optional[datetime] = None,
+    user_id: Optional[int] = None,
 ) -> DatasetId:
     id_ = await session.scalar(
         insert(models.Dataset)
@@ -52,6 +55,7 @@ async def insert_dataset(
             description=description,
             metadata_=metadata,
             created_at=created_at,
+            user_id=user_id,
         )
         .returning(models.Dataset.id)
     )
@@ -64,6 +68,7 @@ async def insert_dataset_version(
     description: Optional[str] = None,
     metadata: Optional[Mapping[str, Any]] = None,
     created_at: Optional[datetime] = None,
+    user_id: Optional[int] = None,
 ) -> DatasetVersionId:
     id_ = await session.scalar(
         insert(models.DatasetVersion)
@@ -72,6 +77,7 @@ async def insert_dataset_version(
             description=description,
             metadata_=metadata,
             created_at=created_at,
+            user_id=user_id,
         )
         .returning(models.DatasetVersion.id)
     )
@@ -134,6 +140,92 @@ async def insert_dataset_example_revision(
     return cast(DatasetExampleRevisionId, id_)
 
 
+async def bulk_create_dataset_splits(
+    session: AsyncSession,
+    split_names: set[str],
+    user_id: Optional[int] = None,
+) -> dict[str, int]:
+    """
+    Bulk create dataset splits using upsert pattern.
+    Returns a mapping of split name to split ID.
+    """
+    if not split_names:
+        return {}
+
+    dialect = SupportedSQLDialect(session.bind.dialect.name)
+    records = [
+        {
+            "name": name,
+            "color": "#808080",  # Default gray color
+            "metadata_": {},
+            "user_id": user_id,
+        }
+        for name in split_names
+    ]
+
+    # Bulk upsert all splits - uses ON CONFLICT DO NOTHING to handle race conditions
+    stmt = insert_on_conflict(
+        *records,
+        table=models.DatasetSplit,
+        dialect=dialect,
+        unique_by=["name"],
+        on_conflict=OnConflict.DO_NOTHING,
+    )
+    await session.execute(stmt)
+
+    # Fetch all split IDs by name
+    result = await session.execute(
+        select(models.DatasetSplit.name, models.DatasetSplit.id).where(
+            models.DatasetSplit.name.in_(split_names)
+        )
+    )
+    return {name: split_id for name, split_id in result.all()}
+
+
+async def bulk_assign_examples_to_splits(
+    session: AsyncSession,
+    assignments: list[tuple[DatasetExampleId, int]],
+) -> None:
+    """
+    Bulk assign examples to splits.
+    assignments is a list of (dataset_example_id, dataset_split_id) tuples.
+    """
+    if not assignments:
+        return
+
+    from sqlalchemy.dialects.postgresql import insert as pg_insert
+    from sqlalchemy.dialects.sqlite import insert as sqlite_insert
+    from typing_extensions import assert_never
+
+    dialect = SupportedSQLDialect(session.bind.dialect.name)
+    records = [
+        {
+            "dataset_example_id": example_id,
+            "dataset_split_id": split_id,
+        }
+        for example_id, split_id in assignments
+    ]
+
+    # Use index_elements instead of constraint name because the table uses
+    # a PrimaryKeyConstraint, not a unique constraint
+    if dialect is SupportedSQLDialect.POSTGRESQL:
+        pg_stmt = pg_insert(models.DatasetSplitDatasetExample).values(records)
+        await session.execute(
+            pg_stmt.on_conflict_do_nothing(
+                index_elements=["dataset_split_id", "dataset_example_id"]
+            )
+        )
+    elif dialect is SupportedSQLDialect.SQLITE:
+        sqlite_stmt = sqlite_insert(models.DatasetSplitDatasetExample).values(records)
+        await session.execute(
+            sqlite_stmt.on_conflict_do_nothing(
+                index_elements=["dataset_split_id", "dataset_example_id"]
+            )
+        )
+    else:
+        assert_never(dialect)
+
+
 class DatasetAction(Enum):
     CREATE = "create"
     APPEND = "append"
@@ -152,6 +244,7 @@ async def add_dataset_examples(
     description: Optional[str] = None,
     metadata: Optional[Mapping[str, Any]] = None,
     action: DatasetAction = DatasetAction.CREATE,
+    user_id: Optional[int] = None,
 ) -> Optional[DatasetExampleAdditionEvent]:
     created_at = datetime.now(timezone.utc)
     dataset_id: Optional[DatasetId] = None
@@ -167,6 +260,7 @@ async def add_dataset_examples(
                 description=description,
                 metadata=metadata,
                 created_at=created_at,
+                user_id=user_id,
             )
         except Exception:
             logger.exception(f"Failed to insert dataset: {name=}")
@@ -176,10 +270,14 @@ async def add_dataset_examples(
             session=session,
             dataset_id=dataset_id,
             created_at=created_at,
+            user_id=user_id,
         )
     except Exception:
         logger.exception(f"Failed to insert dataset version for {dataset_id=}")
         raise
+
+    # Process examples and collect split assignments (by name, resolved to IDs after iteration)
+    split_assignments: list[tuple[DatasetExampleId, str]] = []
     for example in (await examples) if isinstance(examples, Awaitable) else examples:
         try:
             dataset_example_id = await insert_dataset_example(
@@ -206,6 +304,40 @@ async def add_dataset_examples(
                 f"{dataset_example_id=}"
             )
             raise
+
+        # Collect split assignments by name for bulk insert later
+        for split_name in example.splits:
+            split_assignments.append((dataset_example_id, split_name))
+
+    # Bulk create splits and assign examples after iteration
+    if split_assignments:
+        # Collect all unique split names
+        all_split_names = {name for _, name in split_assignments}
+        try:
+            split_name_to_id = await bulk_create_dataset_splits(
+                session=session,
+                split_names=all_split_names,
+                user_id=user_id,
+            )
+        except Exception:
+            logger.exception(f"Failed to bulk create dataset splits: {all_split_names}")
+            raise
+
+        # Convert name-based assignments to ID-based assignments
+        id_assignments = [
+            (example_id, split_name_to_id[split_name])
+            for example_id, split_name in split_assignments
+        ]
+
+        try:
+            await bulk_assign_examples_to_splits(
+                session=session,
+                assignments=id_assignments,
+            )
+        except Exception:
+            logger.exception("Failed to bulk assign examples to splits")
+            raise
+
     return DatasetExampleAdditionEvent(dataset_id=dataset_id, dataset_version_id=dataset_version_id)
 
 

phoenix/db/insertion/document_annotation.py

@@ -7,7 +7,7 @@ from sqlalchemy.ext.asyncio import AsyncSession
 from typing_extensions import TypeAlias
 
 from phoenix.db import models
-from phoenix.db.helpers import dedup, num_docs_col
+from phoenix.db.helpers import dedup
 from phoenix.db.insertion.helpers import as_kv
 from phoenix.db.insertion.types import (
     Insertables,
@@ -63,7 +63,7 @@ class DocumentAnnotationQueueInserter(
         session: AsyncSession,
         *insertions: Insertables.DocumentAnnotation,
     ) -> list[DocumentAnnotationDmlEvent]:
-        records = [dict(as_kv(ins.row)) for ins in insertions]
+        records = [{**dict(as_kv(ins.row)), "updated_at": ins.row.updated_at} for ins in insertions]
         stmt = self._insert_on_conflict(*records).returning(self.table.id)
         ids = tuple([_ async for _ in await session.stream_scalars(stmt)])
         return [DocumentAnnotationDmlEvent(ids)]
@@ -99,7 +99,7 @@ class DocumentAnnotationQueueInserter(
 
         for p in parcels:
             if (anno := existing_annos.get(_key(p))) is not None:
-                if p.received_at <= anno.updated_at:
+                if p.item.updated_at <= anno.updated_at:
                     to_discard.append(p)
                 else:
                     to_insert.append(
@@ -107,7 +107,6 @@ class DocumentAnnotationQueueInserter(
                             received_at=p.received_at,
                             item=p.item.as_insertable(
                                 span_rowid=anno.span_rowid,
-                                id_=anno.id_,
                             ),
                         )
                     )
@@ -140,7 +139,11 @@ class DocumentAnnotationQueueInserter(
     def _select_existing(self, *keys: _Key) -> Select[_Existing]:
         anno = self.table
         span = (
-            select(models.Span.id, models.Span.span_id, num_docs_col(self._db.dialect))
+            select(
+                models.Span.id,
+                models.Span.span_id,
+                models.Span.num_documents.label("num_docs"),
+            )
             .where(models.Span.span_id.in_({k.span_id for k in keys}))
             .cte()
         )
@@ -182,7 +185,7 @@ def _key(p: Received[Precursors.DocumentAnnotation]) -> _Key:
 
 
 def _unique_by(p: Received[Insertables.DocumentAnnotation]) -> _UniqueBy:
-    return p.item.obj.name, p.item.span_rowid, p.item.document_position, p.item.identifier
+    return p.item.obj.name, p.item.span_rowid, p.item.document_position, p.item.obj.identifier
 
 
 def _time(p: Received[Any]) -> datetime:

phoenix/db/insertion/evaluation.py

@@ -5,7 +5,7 @@ from sqlalchemy.ext.asyncio import AsyncSession
 from typing_extensions import assert_never
 
 from phoenix.db import models
-from phoenix.db.helpers import SupportedSQLDialect, num_docs_col
+from phoenix.db.helpers import SupportedSQLDialect
 from phoenix.db.insertion.helpers import insert_on_conflict
 from phoenix.exceptions import PhoenixException
 from phoenix.trace import v1 as pb
@@ -153,12 +153,11 @@ async def _insert_document_evaluation(
     score: Optional[float],
     explanation: Optional[str],
 ) -> EvaluationInsertionEvent:
-    dialect = SupportedSQLDialect(session.bind.dialect.name)
     stmt = (
         select(
             models.Trace.project_rowid,
             models.Span.id,
-            num_docs_col(dialect),
+            models.Span.num_documents,
         )
         .join_from(models.Span, models.Trace)
         .where(models.Span.span_id == span_id)

phoenix/db/insertion/helpers.py

@@ -12,7 +12,7 @@ from sqlalchemy.sql.elements import KeyedColumnElement
 from typing_extensions import TypeAlias, assert_never
 
 from phoenix.db import models
-from phoenix.db.helpers import SupportedSQLDialect
+from phoenix.db.helpers import SupportedSQLDialect, truncate_name
 from phoenix.db.models import Base
 from phoenix.trace.attributes import get_attribute_value
 
@@ -53,7 +53,7 @@ def insert_on_conflict(
             unique_records.append(v)
             seen.add(k)
         records = tuple(reversed(unique_records))
-    constraint = constraint_name or "_".join(("uq", table.__tablename__, *unique_by))
+    constraint = constraint_name or truncate_name("_".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/session_annotation.py

@@ -0,0 +1,176 @@
+from collections.abc import Mapping
+from datetime import datetime
+from typing import Any, NamedTuple, Optional
+
+from sqlalchemy import Row, Select, and_, select, tuple_
+from sqlalchemy.ext.asyncio import AsyncSession
+from typing_extensions import TypeAlias
+
+from phoenix.db import models
+from phoenix.db.helpers import dedup
+from phoenix.db.insertion.helpers import as_kv
+from phoenix.db.insertion.types import (
+    Insertables,
+    Postponed,
+    Precursors,
+    QueueInserter,
+    Received,
+)
+from phoenix.server.dml_event import ProjectSessionAnnotationDmlEvent
+
+# Type alias for consistency with other annotation patterns
+SessionAnnotationDmlEvent = ProjectSessionAnnotationDmlEvent
+
+_Name: TypeAlias = str
+_SessionId: TypeAlias = str
+_SessionRowId: TypeAlias = int
+_AnnoRowId: TypeAlias = int
+_Identifier: TypeAlias = str
+
+
+class _Key(NamedTuple):
+    annotation_name: _Name
+    annotation_identifier: _Identifier
+    session_id: _SessionId
+
+
+_UniqueBy: TypeAlias = tuple[_Name, _SessionRowId, _Identifier]
+_Existing: TypeAlias = tuple[
+    _SessionRowId,
+    _SessionId,
+    Optional[_AnnoRowId],
+    Optional[_Name],
+    Optional[datetime],
+]
+
+
+class SessionAnnotationQueueInserter(
+    QueueInserter[
+        Precursors.SessionAnnotation,
+        Insertables.SessionAnnotation,
+        models.ProjectSessionAnnotation,
+        SessionAnnotationDmlEvent,
+    ],
+    table=models.ProjectSessionAnnotation,
+    unique_by=("name", "project_session_id", "identifier"),
+):
+    async def _events(
+        self,
+        session: AsyncSession,
+        *insertions: Insertables.SessionAnnotation,
+    ) -> list[SessionAnnotationDmlEvent]:
+        records = [{**dict(as_kv(ins.row)), "updated_at": ins.row.updated_at} for ins in insertions]
+        stmt = self._insert_on_conflict(*records).returning(self.table.id)
+        ids = tuple([_ async for _ in await session.stream_scalars(stmt)])
+        return [SessionAnnotationDmlEvent(ids)]
+
+    async def _partition(
+        self,
+        session: AsyncSession,
+        *parcels: Received[Precursors.SessionAnnotation],
+    ) -> tuple[
+        list[Received[Insertables.SessionAnnotation]],
+        list[Postponed[Precursors.SessionAnnotation]],
+        list[Received[Precursors.SessionAnnotation]],
+    ]:
+        to_insert: list[Received[Insertables.SessionAnnotation]] = []
+        to_postpone: list[Postponed[Precursors.SessionAnnotation]] = []
+        to_discard: list[Received[Precursors.SessionAnnotation]] = []
+
+        stmt = self._select_existing(*map(_key, parcels))
+        existing: list[Row[_Existing]] = [_ async for _ in await session.stream(stmt)]
+        existing_sessions: Mapping[str, _SessionAttr] = {
+            e.session_id: _SessionAttr(e.session_rowid) for e in existing
+        }
+        existing_annos: Mapping[_Key, _AnnoAttr] = {
+            _Key(
+                annotation_name=e.name,
+                annotation_identifier=e.identifier,
+                session_id=e.session_id,
+            ): _AnnoAttr(e.session_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
+        }
+
+        for p in parcels:
+            if (anno := existing_annos.get(_key(p))) is not None:
+                if p.item.updated_at <= anno.updated_at:
+                    to_discard.append(p)
+                else:
+                    to_insert.append(
+                        Received(
+                            received_at=p.received_at,
+                            item=p.item.as_insertable(
+                                project_session_rowid=anno.session_rowid,
+                            ),
+                        )
+                    )
+            elif (existing_session := existing_sessions.get(p.item.session_id)) is not None:
+                to_insert.append(
+                    Received(
+                        received_at=p.received_at,
+                        item=p.item.as_insertable(
+                            project_session_rowid=existing_session.session_rowid,
+                        ),
+                    )
+                )
+            elif isinstance(p, Postponed):
+                if p.retries_left > 1:
+                    to_postpone.append(p.postpone(p.retries_left - 1))
+                else:
+                    to_discard.append(p)
+            elif isinstance(p, Received):
+                to_postpone.append(p.postpone(self._retry_allowance))
+            else:
+                to_discard.append(p)
+
+        assert len(to_insert) + len(to_postpone) + len(to_discard) == len(parcels)
+        to_insert = dedup(sorted(to_insert, key=_time, reverse=True), _unique_by)[::-1]
+        return to_insert, to_postpone, to_discard
+
+    def _select_existing(self, *keys: _Key) -> Select[_Existing]:
+        anno = self.table
+        session = (
+            select(models.ProjectSession.id, models.ProjectSession.session_id)
+            .where(models.ProjectSession.session_id.in_({k.session_id for k in keys}))
+            .cte()
+        )
+        onclause = and_(
+            session.c.id == anno.project_session_id,
+            anno.name.in_({k.annotation_name for k in keys}),
+            tuple_(anno.name, anno.identifier, session.c.session_id).in_(keys),
+        )
+        return select(
+            session.c.id.label("session_rowid"),
+            session.c.session_id,
+            anno.id,
+            anno.name,
+            anno.identifier,
+            anno.updated_at,
+        ).outerjoin_from(session, anno, onclause)
+
+
+class _SessionAttr(NamedTuple):
+    session_rowid: _SessionRowId
+
+
+class _AnnoAttr(NamedTuple):
+    session_rowid: _SessionRowId
+    id_: _AnnoRowId
+    updated_at: datetime
+
+
+def _key(p: Received[Precursors.SessionAnnotation]) -> _Key:
+    return _Key(
+        annotation_name=p.item.obj.name,
+        annotation_identifier=p.item.obj.identifier,
+        session_id=p.item.session_id,
+    )
+
+
+def _unique_by(p: Received[Insertables.SessionAnnotation]) -> _UniqueBy:
+    return p.item.obj.name, p.item.project_session_rowid, p.item.obj.identifier
+
+
+def _time(p: Received[Any]) -> datetime:
+    return p.received_at

phoenix/db/insertion/span_annotation.py

@@ -57,7 +57,7 @@ class SpanAnnotationQueueInserter(
         session: AsyncSession,
         *insertions: Insertables.SpanAnnotation,
     ) -> list[SpanAnnotationDmlEvent]:
-        records = [dict(as_kv(ins.row)) for ins in insertions]
+        records = [{**dict(as_kv(ins.row)), "updated_at": ins.row.updated_at} for ins in insertions]
         stmt = self._insert_on_conflict(*records).returning(self.table.id)
         ids = tuple([_ async for _ in await session.stream_scalars(stmt)])
         return [SpanAnnotationDmlEvent(ids)]
@@ -92,7 +92,7 @@ class SpanAnnotationQueueInserter(
 
         for p in parcels:
             if (anno := existing_annos.get(_key(p))) is not None:
-                if p.received_at <= anno.updated_at:
+                if p.item.updated_at <= anno.updated_at:
                     to_discard.append(p)
                 else:
                     to_insert.append(
@@ -100,7 +100,6 @@ class SpanAnnotationQueueInserter(
                             received_at=p.received_at,
                             item=p.item.as_insertable(
                                 span_rowid=anno.span_rowid,
-                                id_=anno.id_,
                             ),
                         )
                     )
@@ -168,7 +167,7 @@ def _key(p: Received[Precursors.SpanAnnotation]) -> _Key:
 
 
 def _unique_by(p: Received[Insertables.SpanAnnotation]) -> _UniqueBy:
-    return p.item.obj.name, p.item.span_rowid, p.item.identifier
+    return p.item.obj.name, p.item.span_rowid, p.item.obj.identifier
 
 
 def _time(p: Received[Any]) -> datetime:

phoenix/db/insertion/trace_annotation.py

@@ -56,7 +56,7 @@ class TraceAnnotationQueueInserter(
         session: AsyncSession,
         *insertions: Insertables.TraceAnnotation,
     ) -> list[TraceAnnotationDmlEvent]:
-        records = [dict(as_kv(ins.row)) for ins in insertions]
+        records = [{**dict(as_kv(ins.row)), "updated_at": ins.row.updated_at} for ins in insertions]
         stmt = self._insert_on_conflict(*records).returning(self.table.id)
         ids = tuple([_ async for _ in await session.stream_scalars(stmt)])
         return [TraceAnnotationDmlEvent(ids)]
@@ -91,7 +91,7 @@ class TraceAnnotationQueueInserter(
 
         for p in parcels:
             if (anno := existing_annos.get(_key(p))) is not None:
-                if p.received_at <= anno.updated_at:
+                if p.item.updated_at <= anno.updated_at:
                     to_discard.append(p)
                 else:
                     to_insert.append(
@@ -99,7 +99,6 @@ class TraceAnnotationQueueInserter(
                             received_at=p.received_at,
                             item=p.item.as_insertable(
                                 trace_rowid=anno.trace_rowid,
-                                id_=anno.id_,
                             ),
                         )
                     )
@@ -167,7 +166,7 @@ def _key(p: Received[Precursors.TraceAnnotation]) -> _Key:
 
 
 def _unique_by(p: Received[Insertables.TraceAnnotation]) -> _UniqueBy:
-    return p.item.obj.name, p.item.trace_rowid, p.item.identifier
+    return p.item.obj.name, p.item.trace_rowid, p.item.obj.identifier
 
 
 def _time(p: Received[Any]) -> datetime: