arize-phoenix 4.5.0__py3-none-any.whl → 4.6.2__py3-none-any.whl

phoenix/db/bulk_inserter.py

@@ -1,5 +1,6 @@
 import asyncio
 import logging
+from asyncio import Queue
 from dataclasses import dataclass, field
 from datetime import datetime, timezone
 from itertools import islice
@@ -14,6 +15,7 @@ from typing import (
     Optional,
     Set,
     Tuple,
+    cast,
 )
 
 from cachetools import LRUCache
@@ -22,10 +24,11 @@ from typing_extensions import TypeAlias
 
 import phoenix.trace.v1 as pb
 from phoenix.db.insertion.evaluation import (
-    EvaluationInsertionResult,
+    EvaluationInsertionEvent,
     InsertEvaluationError,
     insert_evaluation,
 )
+from phoenix.db.insertion.helpers import DataManipulation, DataManipulationEvent
 from phoenix.db.insertion.span import SpanInsertionEvent, insert_span
 from phoenix.server.api.dataloaders import CacheForDataLoaders
 from phoenix.trace.schemas import Span
@@ -46,23 +49,29 @@ class BulkInserter:
         db: Callable[[], AsyncContextManager[AsyncSession]],
         *,
         cache_for_dataloaders: Optional[CacheForDataLoaders] = None,
+        initial_batch_of_operations: Iterable[DataManipulation] = (),
         initial_batch_of_spans: Optional[Iterable[Tuple[Span, str]]] = None,
         initial_batch_of_evaluations: Optional[Iterable[pb.Evaluation]] = None,
         sleep: float = 0.1,
-        max_num_per_transaction: int = 1000,
+        max_ops_per_transaction: int = 1000,
+        max_queue_size: int = 1000,
         enable_prometheus: bool = False,
     ) -> None:
         """
         :param db: A function to initiate a new database session.
         :param initial_batch_of_spans: Initial batch of spans to insert.
         :param sleep: The time to sleep between bulk insertions
-        :param max_num_per_transaction: The maximum number of items to insert in a single
-        transaction. Multiple transactions will be used if there are more items in the batch.
+        :param max_ops_per_transaction: The maximum number of operations to dequeue from
+        the operations queue for each transaction.
+        :param max_queue_size: The maximum length of the operations queue.
+        :param enable_prometheus: Whether Prometheus is enabled.
         """
         self._db = db
         self._running = False
         self._sleep = sleep
-        self._max_num_per_transaction = max_num_per_transaction
+        self._max_ops_per_transaction = max_ops_per_transaction
+        self._operations: Optional[Queue[DataManipulation]] = None
+        self._max_queue_size = max_queue_size
         self._spans: List[Tuple[Span, str]] = (
             [] if initial_batch_of_spans is None else list(initial_batch_of_spans)
         )
@@ -81,27 +90,58 @@ class BulkInserter:
 
     async def __aenter__(
         self,
-    ) -> Tuple[Callable[[Span, str], Awaitable[None]], Callable[[pb.Evaluation], Awaitable[None]]]:
+    ) -> Tuple[
+        Callable[[Span, str], Awaitable[None]],
+        Callable[[pb.Evaluation], Awaitable[None]],
+        Callable[[DataManipulation], None],
+    ]:
         self._running = True
+        self._operations = Queue(maxsize=self._max_queue_size)
         self._task = asyncio.create_task(self._bulk_insert())
-        return self._queue_span, self._queue_evaluation
+        return (
+            self._queue_span,
+            self._queue_evaluation,
+            self._enqueue_operation,
+        )
 
     async def __aexit__(self, *args: Any) -> None:
+        self._operations = None
         self._running = False
 
+    def _enqueue_operation(self, operation: DataManipulation) -> None:
+        cast("Queue[DataManipulation]", self._operations).put_nowait(operation)
+
     async def _queue_span(self, span: Span, project_name: str) -> None:
         self._spans.append((span, project_name))
 
     async def _queue_evaluation(self, evaluation: pb.Evaluation) -> None:
         self._evaluations.append(evaluation)
 
+    async def _process_events(self, events: Iterable[Optional[DataManipulationEvent]]) -> None: ...
+
     async def _bulk_insert(self) -> None:
+        assert isinstance(self._operations, Queue)
         spans_buffer, evaluations_buffer = None, None
         # start first insert immediately if the inserter has not run recently
-        while self._spans or self._evaluations or self._running:
-            if not (self._spans or self._evaluations):
+        while self._running or not self._operations.empty() or self._spans or self._evaluations:
+            if self._operations.empty() and not (self._spans or self._evaluations):
                 await asyncio.sleep(self._sleep)
                 continue
+            ops_remaining, events = self._max_ops_per_transaction, []
+            async with self._db() as session:
+                while ops_remaining and not self._operations.empty():
+                    ops_remaining -= 1
+                    op = await self._operations.get()
+                    try:
+                        async with session.begin_nested():
+                            events.append(await op(session))
+                    except Exception as e:
+                        if self._enable_prometheus:
+                            from phoenix.server.prometheus import BULK_LOADER_EXCEPTIONS
+
+                            BULK_LOADER_EXCEPTIONS.inc()
+                        logger.exception(str(e))
+            await self._process_events(events)
             # It's important to grab the buffers at the same time so there's
             # no race condition, since an eval insertion will fail if the span
             # it references doesn't exist. Grabbing the eval buffer later may
@@ -130,11 +170,11 @@ class BulkInserter:
 
     async def _insert_spans(self, spans: List[Tuple[Span, str]]) -> TransactionResult:
         transaction_result = TransactionResult()
-        for i in range(0, len(spans), self._max_num_per_transaction):
+        for i in range(0, len(spans), self._max_ops_per_transaction):
             try:
                 start = perf_counter()
                 async with self._db() as session:
-                    for span, project_name in islice(spans, i, i + self._max_num_per_transaction):
+                    for span, project_name in islice(spans, i, i + self._max_ops_per_transaction):
                         if self._enable_prometheus:
                             from phoenix.server.prometheus import BULK_LOADER_SPAN_INSERTIONS
 
@@ -169,16 +209,16 @@ class BulkInserter:
 
     async def _insert_evaluations(self, evaluations: List[pb.Evaluation]) -> TransactionResult:
         transaction_result = TransactionResult()
-        for i in range(0, len(evaluations), self._max_num_per_transaction):
+        for i in range(0, len(evaluations), self._max_ops_per_transaction):
             try:
                 start = perf_counter()
                 async with self._db() as session:
-                    for evaluation in islice(evaluations, i, i + self._max_num_per_transaction):
+                    for evaluation in islice(evaluations, i, i + self._max_ops_per_transaction):
                         if self._enable_prometheus:
                             from phoenix.server.prometheus import BULK_LOADER_EVALUATION_INSERTIONS
 
                             BULK_LOADER_EVALUATION_INSERTIONS.inc()
-                        result: Optional[EvaluationInsertionResult] = None
+                        result: Optional[EvaluationInsertionEvent] = None
                         try:
                             async with session.begin_nested():
                                 result = await insert_evaluation(session, evaluation)

phoenix/db/insertion/dataset.py

@@ -0,0 +1,237 @@
+import logging
+from dataclasses import dataclass, field
+from datetime import datetime, timezone
+from enum import Enum
+from itertools import chain
+from typing import (
+    Any,
+    Awaitable,
+    Dict,
+    FrozenSet,
+    Iterable,
+    Iterator,
+    Mapping,
+    Optional,
+    Union,
+    cast,
+)
+
+from sqlalchemy import insert, select
+from sqlalchemy.ext.asyncio import AsyncSession
+from typing_extensions import TypeAlias
+
+from phoenix.db import models
+from phoenix.db.insertion.helpers import DataManipulationEvent
+
+logger = logging.getLogger(__name__)
+
+DatasetId: TypeAlias = int
+DatasetVersionId: TypeAlias = int
+DatasetExampleId: TypeAlias = int
+DatasetExampleRevisionId: TypeAlias = int
+SpanRowId: TypeAlias = int
+
+
+@dataclass(frozen=True)
+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)
+
+
+Examples: TypeAlias = Iterable[ExampleContent]
+
+
+@dataclass(frozen=True)
+class DatasetExampleAdditionEvent(DataManipulationEvent):
+    dataset_id: DatasetId
+
+
+async def insert_dataset(
+    session: AsyncSession,
+    name: str,
+    description: Optional[str] = None,
+    metadata: Optional[Mapping[str, Any]] = None,
+    created_at: Optional[datetime] = None,
+) -> DatasetId:
+    id_ = await session.scalar(
+        insert(models.Dataset)
+        .values(
+            name=name,
+            description=description,
+            metadata_=metadata,
+            created_at=created_at,
+        )
+        .returning(models.Dataset.id)
+    )
+    return cast(DatasetId, id_)
+
+
+async def insert_dataset_version(
+    session: AsyncSession,
+    dataset_id: DatasetId,
+    description: Optional[str] = None,
+    metadata: Optional[Mapping[str, Any]] = None,
+    created_at: Optional[datetime] = None,
+) -> DatasetVersionId:
+    id_ = await session.scalar(
+        insert(models.DatasetVersion)
+        .values(
+            dataset_id=dataset_id,
+            description=description,
+            metadata_=metadata,
+            created_at=created_at,
+        )
+        .returning(models.DatasetVersion.id)
+    )
+    return cast(DatasetVersionId, id_)
+
+
+async def insert_dataset_example(
+    session: AsyncSession,
+    dataset_id: DatasetId,
+    span_rowid: Optional[SpanRowId] = None,
+    created_at: Optional[datetime] = None,
+) -> DatasetExampleId:
+    id_ = await session.scalar(
+        insert(models.DatasetExample)
+        .values(
+            dataset_id=dataset_id,
+            span_rowid=span_rowid,
+            created_at=created_at,
+        )
+        .returning(models.DatasetExample.id)
+    )
+    return cast(DatasetExampleId, id_)
+
+
+class RevisionKind(Enum):
+    CREATE = "CREATE"
+    PATCH = "PATCH"
+    DELETE = "DELETE"
+
+    @classmethod
+    def _missing_(cls, v: Any) -> "RevisionKind":
+        if isinstance(v, str) and v and v.isascii() and not v.isupper():
+            return cls(v.upper())
+        raise ValueError(f"Invalid revision kind: {v}")
+
+
+async def insert_dataset_example_revision(
+    session: AsyncSession,
+    dataset_version_id: DatasetVersionId,
+    dataset_example_id: DatasetExampleId,
+    input: Mapping[str, Any],
+    output: Mapping[str, Any],
+    metadata: Optional[Mapping[str, Any]] = None,
+    revision_kind: RevisionKind = RevisionKind.CREATE,
+    created_at: Optional[datetime] = None,
+) -> DatasetExampleRevisionId:
+    id_ = await session.scalar(
+        insert(models.DatasetExampleRevision)
+        .values(
+            dataset_version_id=dataset_version_id,
+            dataset_example_id=dataset_example_id,
+            input=input,
+            output=output,
+            metadata_=metadata,
+            revision_kind=revision_kind.value,
+            created_at=created_at,
+        )
+        .returning(models.DatasetExampleRevision.id)
+    )
+    return cast(DatasetExampleRevisionId, id_)
+
+
+class DatasetAction(Enum):
+    CREATE = "create"
+    APPEND = "append"
+
+    @classmethod
+    def _missing_(cls, v: Any) -> "DatasetAction":
+        if isinstance(v, str) and v and v.isascii() and not v.islower():
+            return cls(v.lower())
+        raise ValueError(f"Invalid dateset action: {v}")
+
+
+async def add_dataset_examples(
+    session: AsyncSession,
+    name: str,
+    examples: Union[Examples, Awaitable[Examples]],
+    description: Optional[str] = None,
+    metadata: Optional[Mapping[str, Any]] = None,
+    action: DatasetAction = DatasetAction.CREATE,
+) -> Optional[DatasetExampleAdditionEvent]:
+    created_at = datetime.now(timezone.utc)
+    dataset_id: Optional[DatasetId] = None
+    if action is DatasetAction.APPEND and name:
+        dataset_id = await session.scalar(
+            select(models.Dataset.id).where(models.Dataset.name == name)
+        )
+    if action is DatasetAction.CREATE or dataset_id is None:
+        try:
+            dataset_id = await insert_dataset(
+                session=session,
+                name=name,
+                description=description,
+                metadata=metadata,
+                created_at=created_at,
+            )
+        except Exception:
+            logger.exception(f"Failed to insert dataset: {name=}")
+            raise
+    try:
+        dataset_version_id = await insert_dataset_version(
+            session=session,
+            dataset_id=dataset_id,
+            created_at=created_at,
+        )
+    except Exception:
+        logger.exception(f"Failed to insert dataset version for {dataset_id=}")
+        raise
+    for example in (await examples) if isinstance(examples, Awaitable) else examples:
+        try:
+            dataset_example_id = await insert_dataset_example(
+                session=session,
+                dataset_id=dataset_id,
+                created_at=created_at,
+            )
+        except Exception:
+            logger.exception(f"Failed to insert dataset example for {dataset_id=}")
+            raise
+        try:
+            await insert_dataset_example_revision(
+                session=session,
+                dataset_version_id=dataset_version_id,
+                dataset_example_id=dataset_example_id,
+                input=example.input,
+                output=example.output,
+                metadata=example.metadata,
+                created_at=created_at,
+            )
+        except Exception:
+            logger.exception(
+                f"Failed to insert dataset example revision for {dataset_version_id=}, "
+                f"{dataset_example_id=}"
+            )
+            raise
+    return DatasetExampleAdditionEvent(dataset_id=dataset_id)
+
+
+@dataclass(frozen=True)
+class DatasetKeys:
+    input: FrozenSet[str]
+    output: FrozenSet[str]
+    metadata: FrozenSet[str]
+
+    def __iter__(self) -> Iterator[str]:
+        yield from sorted(set(chain(self.input, self.output, self.metadata)))
+
+    def check_differences(self, column_headers_set: FrozenSet[str]) -> None:
+        for category, keys in (
+            ("input", self.input),
+            ("output", self.output),
+            ("metadata", self.metadata),
+        ):
+            if diff := keys.difference(column_headers_set):
+                raise ValueError(f"{category} keys not found in table column headers: {diff}")

phoenix/db/insertion/evaluation.py

@@ -6,7 +6,7 @@ from typing_extensions import assert_never
 
 from phoenix.db import models
 from phoenix.db.helpers import SupportedSQLDialect, num_docs_col
-from phoenix.db.insertion.helpers import OnConflict, insert_stmt
+from phoenix.db.insertion.helpers import OnConflict, insert_on_conflict
 from phoenix.exceptions import PhoenixException
 from phoenix.trace import v1 as pb
 
@@ -15,24 +15,24 @@ class InsertEvaluationError(PhoenixException):
     pass
 
 
-class EvaluationInsertionResult(NamedTuple):
+class EvaluationInsertionEvent(NamedTuple):
     project_rowid: int
     evaluation_name: str
 
 
-class SpanEvaluationInsertionEvent(EvaluationInsertionResult): ...
+class SpanEvaluationInsertionEvent(EvaluationInsertionEvent): ...
 
 
-class TraceEvaluationInsertionEvent(EvaluationInsertionResult): ...
+class TraceEvaluationInsertionEvent(EvaluationInsertionEvent): ...
 
 
-class DocumentEvaluationInsertionEvent(EvaluationInsertionResult): ...
+class DocumentEvaluationInsertionEvent(EvaluationInsertionEvent): ...
 
 
 async def insert_evaluation(
     session: AsyncSession,
     evaluation: pb.Evaluation,
-) -> Optional[EvaluationInsertionResult]:
+) -> Optional[EvaluationInsertionEvent]:
     evaluation_name = evaluation.name
     result = evaluation.result
     label = result.label.value if result.HasField("label") else None
@@ -91,7 +91,7 @@ async def _insert_trace_evaluation(
     set_.pop("metadata_")
     set_["metadata"] = values["metadata_"]  # `metadata` must match database
     await session.execute(
-        insert_stmt(
+        insert_on_conflict(
             dialect=dialect,
             table=models.TraceAnnotation,
             values=values,
@@ -139,7 +139,7 @@ async def _insert_span_evaluation(
     set_.pop("metadata_")
     set_["metadata"] = values["metadata_"]  # `metadata` must match database
     await session.execute(
-        insert_stmt(
+        insert_on_conflict(
             dialect=dialect,
             table=models.SpanAnnotation,
             values=values,
@@ -160,7 +160,7 @@ async def _insert_document_evaluation(
     label: Optional[str],
     score: Optional[float],
     explanation: Optional[str],
-) -> EvaluationInsertionResult:
+) -> EvaluationInsertionEvent:
     dialect = SupportedSQLDialect(session.bind.dialect.name)
     stmt = (
         select(
@@ -196,7 +196,7 @@ async def _insert_document_evaluation(
     set_.pop("metadata_")
     set_["metadata"] = values["metadata_"]  # `metadata` must match database
     await session.execute(
-        insert_stmt(
+        insert_on_conflict(
             dialect=dialect,
             table=models.DocumentAnnotation,
             values=values,

phoenix/db/insertion/helpers.py

@@ -1,39 +1,42 @@
+from abc import ABC
 from enum import Enum, auto
-from typing import Any, Mapping, Optional, Sequence
+from typing import Any, Awaitable, Callable, Mapping, Optional, Sequence
 
-from sqlalchemy import Insert, insert
+from sqlalchemy import Insert
 from sqlalchemy.dialects.postgresql import insert as insert_postgresql
 from sqlalchemy.dialects.sqlite import insert as insert_sqlite
-from typing_extensions import assert_never
+from sqlalchemy.ext.asyncio import AsyncSession
+from typing_extensions import TypeAlias, assert_never
 
 from phoenix.db.helpers import SupportedSQLDialect
 
 
+class DataManipulationEvent(ABC):
+    """
+    Execution of DML (Data Manipulation Language) statements.
+    """
+
+
+DataManipulation: TypeAlias = Callable[[AsyncSession], Awaitable[Optional[DataManipulationEvent]]]
+
+
 class OnConflict(Enum):
     DO_NOTHING = auto()
     DO_UPDATE = auto()
 
 
-def insert_stmt(
+def insert_on_conflict(
     dialect: SupportedSQLDialect,
     table: Any,
     values: Mapping[str, Any],
-    constraint: Optional[str] = None,
-    column_names: Sequence[str] = (),
+    constraint: str,
+    column_names: Sequence[str],
     on_conflict: OnConflict = OnConflict.DO_NOTHING,
     set_: Optional[Mapping[str, Any]] = None,
 ) -> Insert:
     """
     Dialect specific insertion statement using ON CONFLICT DO syntax.
     """
-    if bool(constraint) != bool(column_names):
-        raise ValueError(
-            "Both `constraint` and `column_names` must be provided or omitted at the same time."
-        )
-    if (dialect is SupportedSQLDialect.POSTGRESQL and constraint is None) or (
-        dialect is SupportedSQLDialect.SQLITE and not column_names
-    ):
-        return insert(table).values(values)
     if dialect is SupportedSQLDialect.POSTGRESQL:
         stmt_postgresql = insert_postgresql(table).values(values)
         if on_conflict is OnConflict.DO_NOTHING or not set_:

phoenix/db/insertion/span.py

@@ -7,7 +7,7 @@ from sqlalchemy.ext.asyncio import AsyncSession
 
 from phoenix.db import models
 from phoenix.db.helpers import SupportedSQLDialect
-from phoenix.db.insertion.helpers import OnConflict, insert_stmt
+from phoenix.db.insertion.helpers import OnConflict, insert_on_conflict
 from phoenix.trace.attributes import get_attribute_value
 from phoenix.trace.schemas import Span, SpanStatusCode
 
@@ -27,7 +27,7 @@ async def insert_span(
 ) -> Optional[SpanInsertionEvent]:
     dialect = SupportedSQLDialect(session.bind.dialect.name)
     project_rowid = await session.scalar(
-        insert_stmt(
+        insert_on_conflict(
             dialect=dialect,
             table=models.Project,
             constraint="uq_projects_name",
@@ -87,7 +87,7 @@ async def insert_span(
         cumulative_llm_token_count_prompt += cast(int, accumulation[1] or 0)
         cumulative_llm_token_count_completion += cast(int, accumulation[2] or 0)
     span_rowid = await session.scalar(
-        insert_stmt(
+        insert_on_conflict(
             dialect=dialect,
             table=models.Span,
             constraint="uq_spans_span_id",

phoenix/db/migrations/types.py

@@ -0,0 +1,29 @@
+from typing import Any
+
+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")  # 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",
+    )
+)