arize-phoenix 4.4.3__py3-none-any.whl → 4.4.4rc0__py3-none-any.whl

phoenix/datasets/types.py

@@ -0,0 +1,212 @@
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from datetime import datetime
+from types import MappingProxyType
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Dict,
+    List,
+    Mapping,
+    Optional,
+    Protocol,
+    Sequence,
+    Union,
+    runtime_checkable,
+)
+
+from typing_extensions import TypeAlias
+
+JSONSerializable: TypeAlias = Optional[Union[Dict[str, Any], List[Any], str, int, float, bool]]
+
+ExperimentId: TypeAlias = str
+DatasetId: TypeAlias = str
+DatasetVersionId: TypeAlias = str
+ExampleId: TypeAlias = str
+RepetitionNumber: TypeAlias = int
+ExperimentRunId: TypeAlias = str
+TraceId: TypeAlias = str
+
+
+@dataclass(frozen=True)
+class Example:
+    id: ExampleId
+    updated_at: datetime
+    input: Mapping[str, JSONSerializable]
+    output: Mapping[str, JSONSerializable]
+    metadata: Mapping[str, JSONSerializable] = field(default_factory=lambda: MappingProxyType({}))
+
+    @classmethod
+    def from_dict(cls, obj: Mapping[str, Any]) -> Example:
+        return cls(
+            input=obj["input"],
+            output=obj["output"],
+            metadata=obj.get("metadata") or {},
+            id=obj["id"],
+            updated_at=obj["updated_at"],
+        )
+
+
+@dataclass(frozen=True)
+class Dataset:
+    id: DatasetId
+    version_id: DatasetVersionId
+    examples: Sequence[Example]
+
+
+@dataclass(frozen=True)
+class TestCase:
+    example: Example
+    repetition_number: RepetitionNumber
+
+
+@dataclass(frozen=True)
+class Experiment:
+    id: ExperimentId
+    dataset_id: DatasetId
+    dataset_version_id: DatasetVersionId
+    project_name: Optional[str] = None
+
+
+@dataclass(frozen=True)
+class ExperimentResult:
+    result: JSONSerializable
+
+    @classmethod
+    def from_dict(cls, obj: Optional[Mapping[str, Any]]) -> Optional[ExperimentResult]:
+        if not obj:
+            return None
+        return cls(result=obj["result"])
+
+
+@dataclass(frozen=True)
+class ExperimentRun:
+    start_time: datetime
+    end_time: datetime
+    experiment_id: ExperimentId
+    dataset_example_id: ExampleId
+    repetition_number: RepetitionNumber
+    output: Optional[ExperimentResult] = None
+    error: Optional[str] = None
+    id: Optional[ExperimentRunId] = None
+    trace_id: Optional[TraceId] = None
+
+    @classmethod
+    def from_dict(cls, obj: Mapping[str, Any]) -> ExperimentRun:
+        return cls(
+            start_time=obj["start_time"],
+            end_time=obj["end_time"],
+            experiment_id=obj["experiment_id"],
+            dataset_example_id=obj["dataset_example_id"],
+            repetition_number=obj.get("repetition_number") or 1,
+            output=ExperimentResult.from_dict(obj["output"]),
+            error=obj.get("error"),
+            id=obj.get("id"),
+            trace_id=obj.get("trace_id"),
+        )
+
+    def __post_init__(self) -> None:
+        if bool(self.output) == bool(self.error):
+            ValueError("Must specify either result or error")
+
+
+@dataclass(frozen=True)
+class EvaluationResult:
+    score: Optional[float] = None
+    label: Optional[str] = None
+    explanation: Optional[str] = None
+    metadata: Mapping[str, JSONSerializable] = field(default_factory=lambda: MappingProxyType({}))
+
+    @classmethod
+    def from_dict(cls, obj: Optional[Mapping[str, Any]]) -> Optional[EvaluationResult]:
+        if not obj:
+            return None
+        return cls(
+            score=obj.get("score"),
+            label=obj.get("label"),
+            explanation=obj.get("explanation"),
+            metadata=obj.get("metadata") or {},
+        )
+
+    def __post_init__(self) -> None:
+        if self.score is None and not self.label and not self.explanation:
+            ValueError("Must specify one of score, label, or explanation")
+
+
+@dataclass(frozen=True)
+class ExperimentEvaluationRun:
+    experiment_run_id: ExperimentRunId
+    start_time: datetime
+    end_time: datetime
+    name: str
+    annotator_kind: str
+    error: Optional[str] = None
+    result: Optional[EvaluationResult] = None
+    id: Optional[str] = None
+    trace_id: Optional[TraceId] = None
+
+    @classmethod
+    def from_dict(cls, obj: Mapping[str, Any]) -> ExperimentEvaluationRun:
+        return cls(
+            experiment_run_id=obj["experiment_run_id"],
+            start_time=obj["start_time"],
+            end_time=obj["end_time"],
+            name=obj["name"],
+            annotator_kind=obj["annotator_kind"],
+            error=obj.get("error"),
+            result=EvaluationResult.from_dict(obj.get("result")),
+            id=obj.get("id"),
+            trace_id=obj.get("trace_id"),
+        )
+
+    def __post_init__(self) -> None:
+        if bool(self.result) == bool(self.error):
+            ValueError("Must specify either result or error")
+
+
+class _HasName(Protocol):
+    name: str
+
+
+class _HasKind(Protocol):
+    @property
+    def annotator_kind(self) -> str: ...
+
+
+@runtime_checkable
+class CanEvaluate(_HasName, _HasKind, Protocol):
+    def evaluate(
+        self,
+        example: Example,
+        experiment_run: ExperimentRun,
+    ) -> EvaluationResult: ...
+
+
+@runtime_checkable
+class CanAsyncEvaluate(_HasName, _HasKind, Protocol):
+    async def async_evaluate(
+        self,
+        example: Example,
+        experiment_run: ExperimentRun,
+    ) -> EvaluationResult: ...
+
+
+ExperimentEvaluator: TypeAlias = Union[CanEvaluate, CanAsyncEvaluate]
+
+
+# Someday we'll do type checking in unit tests.
+if TYPE_CHECKING:
+
+    class _EvaluatorDummy:
+        annotator_kind: str
+        name: str
+
+        def evaluate(self, _: Example, __: ExperimentRun) -> EvaluationResult:
+            raise NotImplementedError
+
+        async def async_evaluate(self, _: Example, __: ExperimentRun) -> EvaluationResult:
+            raise NotImplementedError
+
+    _: ExperimentEvaluator
+    _ = _EvaluatorDummy()

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,234 @@
+import logging
+from dataclasses import dataclass
+from datetime import datetime, timezone
+from enum import Enum
+from itertools import chain
+from typing import (
+    Any,
+    Awaitable,
+    FrozenSet,
+    Iterable,
+    Iterator,
+    Mapping,
+    Optional,
+    Sequence,
+    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
+Examples: TypeAlias = Iterable[Mapping[str, Any]]
+
+
+@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]],
+    input_keys: Sequence[str],
+    output_keys: Sequence[str],
+    metadata_keys: Sequence[str] = (),
+    description: Optional[str] = None,
+    metadata: Optional[Mapping[str, Any]] = None,
+    action: DatasetAction = DatasetAction.CREATE,
+) -> Optional[DatasetExampleAdditionEvent]:
+    keys = DatasetKeys(frozenset(input_keys), frozenset(output_keys), frozenset(metadata_keys))
+    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"Fail to insert dataset: {input_keys=}, {output_keys=}, {metadata_keys=}"
+            )
+            raise
+    try:
+        dataset_version_id = await insert_dataset_version(
+            session=session,
+            dataset_id=dataset_id,
+            created_at=created_at,
+        )
+    except Exception:
+        logger.exception(f"Fail 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"Fail 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={key: example.get(key) for key in keys.input},
+                output={key: example.get(key) for key in keys.output},
+                metadata={key: example.get(key) for key in keys.metadata},
+                created_at=created_at,
+            )
+        except Exception:
+            logger.exception(
+                f"Fail 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

@@ -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
@@ -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(

phoenix/db/insertion/helpers.py

@@ -1,14 +1,25 @@
+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.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()

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",
+    )
+)