arize-phoenix 4.20.0__py3-none-any.whl → 4.20.2__py3-none-any.whl

phoenix/db/insertion/types.py

@@ -21,11 +21,12 @@ from typing import (
 )
 
 from sqlalchemy.ext.asyncio import AsyncSession
-from sqlalchemy.sql.dml import ReturningInsert
+from sqlalchemy.sql.dml import Insert
 
 from phoenix.db import models
 from phoenix.db.insertion.constants import DEFAULT_RETRY_ALLOWANCE, DEFAULT_RETRY_DELAY_SEC
-from phoenix.db.insertion.helpers import as_kv, insert_on_conflict
+from phoenix.db.insertion.helpers import insert_on_conflict
+from phoenix.server.dml_event import DmlEvent
 from phoenix.server.types import DbSessionFactory
 
 logger = logging.getLogger("__name__")
@@ -40,6 +41,7 @@ _AnyT = TypeVar("_AnyT")
 _PrecursorT = TypeVar("_PrecursorT")
 _InsertableT = TypeVar("_InsertableT", bound=Insertable)
 _RowT = TypeVar("_RowT", bound=models.Base)
+_DmlEventT = TypeVar("_DmlEventT", bound=DmlEvent)
 
 
 @dataclass(frozen=True)
@@ -56,7 +58,7 @@ class Postponed(Received[_AnyT]):
     retries_left: int = field(default=DEFAULT_RETRY_ALLOWANCE)
 
 
-class QueueInserter(ABC, Generic[_PrecursorT, _InsertableT, _RowT]):
+class QueueInserter(ABC, Generic[_PrecursorT, _InsertableT, _RowT, _DmlEventT]):
     table: Type[_RowT]
     unique_by: Sequence[str]
 
@@ -97,59 +99,63 @@ class QueueInserter(ABC, Generic[_PrecursorT, _InsertableT, _RowT]):
         List[Received[_PrecursorT]],
     ]: ...
 
-    async def insert(self) -> Tuple[Type[_RowT], List[int]]:
+    async def insert(self) -> Optional[List[_DmlEventT]]:
         if not self._queue:
-            return self.table, []
-        parcels = self._queue
-        self._queue = []
-        inserted_ids: List[int] = []
+            return None
+        self._queue, parcels = [], self._queue
+        events: List[_DmlEventT] = []
         async with self._db() as session:
             to_insert, to_postpone, _ = await self._partition(session, *parcels)
             if to_insert:
-                inserted_ids, to_retry, _ = await self._insert(session, *to_insert)
-                to_postpone.extend(to_retry)
+                events, to_retry, _ = await self._insert(session, *to_insert)
+                if to_retry:
+                    to_postpone.extend(to_retry)
         if to_postpone:
             loop = asyncio.get_running_loop()
             loop.call_later(self._retry_delay_sec, self._queue.extend, to_postpone)
-        return self.table, inserted_ids
+        return events
 
-    def _stmt(self, *records: Mapping[str, Any]) -> ReturningInsert[Tuple[int]]:
-        pk = next(c for c in self.table.__table__.c if c.primary_key)
+    def _insert_on_conflict(self, *records: Mapping[str, Any]) -> Insert:
         return insert_on_conflict(
             *records,
             table=self.table,
             unique_by=self.unique_by,
             dialect=self._db.dialect,
-        ).returning(pk)
+        )
+
+    @abstractmethod
+    async def _events(
+        self,
+        session: AsyncSession,
+        *insertions: _InsertableT,
+    ) -> List[_DmlEventT]: ...
 
     async def _insert(
         self,
         session: AsyncSession,
-        *insertions: Received[_InsertableT],
-    ) -> Tuple[List[int], List[Postponed[_PrecursorT]], List[Received[_InsertableT]]]:
-        records = [dict(as_kv(ins.item.row)) for ins in insertions]
-        inserted_ids: List[int] = []
+        *parcels: Received[_InsertableT],
+    ) -> Tuple[
+        List[_DmlEventT],
+        List[Postponed[_PrecursorT]],
+        List[Received[_InsertableT]],
+    ]:
         to_retry: List[Postponed[_PrecursorT]] = []
         failures: List[Received[_InsertableT]] = []
-        stmt = self._stmt(*records)
+        events: List[_DmlEventT] = []
         try:
             async with session.begin_nested():
-                ids = [id_ async for id_ in await session.stream_scalars(stmt)]
-                inserted_ids.extend(ids)
+                events.extend(await self._events(session, *(p.item for p in parcels)))
         except BaseException:
             logger.exception(
                 f"Failed to bulk insert for {self.table.__name__}. "
-                f"Will try to insert ({len(records)} records) individually instead."
+                f"Will try to insert ({len(parcels)} records) individually instead."
             )
-            for i, record in enumerate(records):
-                stmt = self._stmt(record)
+            for p in parcels:
                 try:
                     async with session.begin_nested():
-                        ids = [id_ async for id_ in await session.stream_scalars(stmt)]
-                        inserted_ids.extend(ids)
+                        events.extend(await self._events(session, p.item))
                 except BaseException:
                     logger.exception(f"Failed to insert for {self.table.__name__}.")
-                    p = insertions[i]
                     if isinstance(p, Postponed) and p.retries_left == 1:
                         failures.append(p)
                     else:
@@ -162,7 +168,7 @@ class QueueInserter(ABC, Generic[_PrecursorT, _InsertableT, _RowT]):
                                 else self._retry_allowance,
                             )
                         )
-        return inserted_ids, to_retry, failures
+        return events, to_retry, failures
 
 
 class Precursors(ABC):

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

@@ -6,11 +6,37 @@ Create Date: 2024-05-10 11:24:23.985834
 
 """
 
-from typing import Sequence, Union
+from typing import Any, Sequence, Union
 
 import sqlalchemy as sa
 from alembic import op
-from phoenix.db.migrations.types import JSON_
+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",
+    )
+)
 
 # revision identifiers, used by Alembic.
 revision: str = "10460e46d750"

phoenix/db/migrations/versions/3be8647b87d8_add_token_columns_to_spans_table.py

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

phoenix/db/migrations/versions/cf03bd6bae1d_init.py

@@ -6,11 +6,37 @@ Create Date: 2024-04-03 19:41:48.871555
 
 """
 
-from typing import Sequence, Union
+from typing import Any, Sequence, Union
 
 import sqlalchemy as sa
 from alembic import op
-from phoenix.db.migrations.types import JSON_
+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",
+    )
+)
 
 # revision identifiers, used by Alembic.
 revision: str = "cf03bd6bae1d"

phoenix/db/models.py

@@ -197,7 +197,7 @@ class Span(Base):
         ForeignKey("traces.id", ondelete="CASCADE"),
         index=True,
     )
-    span_id: Mapped[str]
+    span_id: Mapped[str] = mapped_column(index=True)
     parent_id: Mapped[Optional[str]] = mapped_column(index=True)
     name: Mapped[str]
     span_kind: Mapped[str]
@@ -214,6 +214,8 @@ class Span(Base):
     cumulative_error_count: Mapped[int]
     cumulative_llm_token_count_prompt: Mapped[int]
     cumulative_llm_token_count_completion: Mapped[int]
+    llm_token_count_prompt: Mapped[Optional[int]]
+    llm_token_count_completion: Mapped[Optional[int]]
 
     @hybrid_property
     def latency_ms(self) -> float:
@@ -230,6 +232,12 @@ class Span(Base):
     def cumulative_llm_token_count_total(self) -> int:
         return self.cumulative_llm_token_count_prompt + self.cumulative_llm_token_count_completion
 
+    @hybrid_property
+    def llm_token_count_total(self) -> Optional[int]:
+        if self.llm_token_count_prompt is None and self.llm_token_count_completion is None:
+            return None
+        return (self.llm_token_count_prompt or 0) + (self.llm_token_count_completion or 0)
+
     trace: Mapped["Trace"] = relationship("Trace", back_populates="spans")
     document_annotations: Mapped[List["DocumentAnnotation"]] = relationship(back_populates="span")
     dataset_examples: Mapped[List["DatasetExample"]] = relationship(back_populates="span")

phoenix/server/api/context.py

@@ -1,10 +1,8 @@
 from dataclasses import dataclass
-from datetime import datetime
 from pathlib import Path
-from typing import Callable, Optional
+from typing import Any, Optional
 
 from strawberry.fastapi import BaseContext
-from typing_extensions import TypeAlias
 
 from phoenix.core.model_schema import Model
 from phoenix.server.api.dataloaders import (
@@ -34,7 +32,8 @@ from phoenix.server.api.dataloaders import (
     TraceEvaluationsDataLoader,
     TraceRowIdsDataLoader,
 )
-from phoenix.server.types import DbSessionFactory
+from phoenix.server.dml_event import DmlEvent
+from phoenix.server.types import CanGetLastUpdatedAt, CanPutItem, DbSessionFactory
 
 
 @dataclass
@@ -65,7 +64,9 @@ class DataLoaders:
     project_by_name: ProjectByNameDataLoader
 
 
-ProjectRowId: TypeAlias = int
+class _NoOp:
+    def get(self, *args: Any, **kwargs: Any) -> Any: ...
+    def put(self, *args: Any, **kwargs: Any) -> Any: ...
 
 
 @dataclass
@@ -75,6 +76,7 @@ class Context(BaseContext):
     cache_for_dataloaders: Optional[CacheForDataLoaders]
     model: Model
     export_path: Path
+    last_updated_at: CanGetLastUpdatedAt = _NoOp()
+    event_queue: CanPutItem[DmlEvent] = _NoOp()
     corpus: Optional[Model] = None
-    streaming_last_updated_at: Callable[[ProjectRowId], Optional[datetime]] = lambda _: None
     read_only: bool = False

phoenix/server/api/dataloaders/__init__.py

@@ -1,12 +1,4 @@
 from dataclasses import dataclass, field
-from functools import singledispatchmethod
-
-from phoenix.db.insertion.evaluation import (
-    DocumentEvaluationInsertionEvent,
-    SpanEvaluationInsertionEvent,
-    TraceEvaluationInsertionEvent,
-)
-from phoenix.db.insertion.span import ClearProjectSpansEvent, SpanInsertionEvent
 
 from .annotation_summaries import AnnotationSummaryCache, AnnotationSummaryDataLoader
 from .average_experiment_run_latency import AverageExperimentRunLatencyDataLoader
@@ -88,42 +80,3 @@ class CacheForDataLoaders:
     token_count: TokenCountCache = field(
         default_factory=TokenCountCache,
     )
-
-    def _update_spans(self, project_rowid: int) -> None:
-        self.latency_ms_quantile.invalidate(project_rowid)
-        self.token_count.invalidate(project_rowid)
-        self.record_count.invalidate(project_rowid)
-        self.min_start_or_max_end_time.invalidate(project_rowid)
-
-    def _clear_spans(self, project_rowid: int) -> None:
-        self._update_spans(project_rowid)
-        self.annotation_summary.invalidate_project(project_rowid)
-        self.evaluation_summary.invalidate_project(project_rowid)
-        self.document_evaluation_summary.invalidate_project(project_rowid)
-
-    @singledispatchmethod
-    def invalidate(self, event: SpanInsertionEvent) -> None:
-        project_rowid, *_ = event
-        self._update_spans(project_rowid)
-
-    @invalidate.register
-    def _(self, event: ClearProjectSpansEvent) -> None:
-        project_rowid, *_ = event
-        self._clear_spans(project_rowid)
-
-    @invalidate.register
-    def _(self, event: DocumentEvaluationInsertionEvent) -> None:
-        project_rowid, evaluation_name = event
-        self.document_evaluation_summary.invalidate((project_rowid, evaluation_name))
-
-    @invalidate.register
-    def _(self, event: SpanEvaluationInsertionEvent) -> None:
-        project_rowid, evaluation_name = event
-        self.annotation_summary.invalidate((project_rowid, evaluation_name, "span"))
-        self.evaluation_summary.invalidate((project_rowid, evaluation_name, "span"))
-
-    @invalidate.register
-    def _(self, event: TraceEvaluationInsertionEvent) -> None:
-        project_rowid, evaluation_name = event
-        self.annotation_summary.invalidate((project_rowid, evaluation_name, "trace"))
-        self.evaluation_summary.invalidate((project_rowid, evaluation_name, "trace"))

phoenix/server/api/dataloaders/token_counts.py

@@ -10,7 +10,6 @@ from typing import (
 )
 
 from cachetools import LFUCache, TTLCache
-from openinference.semconv.trace import SpanAttributes
 from sqlalchemy import Select, func, select
 from sqlalchemy.sql.functions import coalesce
 from strawberry.dataloader import AbstractCache, DataLoader
@@ -107,8 +106,8 @@ def _get_stmt(
     *params: Param,
 ) -> Select[Any]:
     (start_time, end_time), filter_condition = segment
-    prompt = func.sum(models.Span.attributes[_LLM_TOKEN_COUNT_PROMPT].as_float())
-    completion = func.sum(models.Span.attributes[_LLM_TOKEN_COUNT_COMPLETION].as_float())
+    prompt = func.sum(models.Span.llm_token_count_prompt)
+    completion = func.sum(models.Span.llm_token_count_completion)
     total = coalesce(prompt, 0) + coalesce(completion, 0)
     pid = models.Trace.project_rowid
     stmt: Select[Any] = (
@@ -130,7 +129,3 @@ def _get_stmt(
         stmt = sf(stmt)
     stmt = stmt.where(pid.in_([rowid for rowid, _ in params]))
     return stmt
-
-
-_LLM_TOKEN_COUNT_PROMPT = SpanAttributes.LLM_TOKEN_COUNT_PROMPT.split(".")
-_LLM_TOKEN_COUNT_COMPLETION = SpanAttributes.LLM_TOKEN_COUNT_COMPLETION.split(".")

phoenix/server/api/input_types/SpanSort.py

@@ -3,7 +3,6 @@ from enum import Enum, auto
 from typing import Any, Optional, Protocol
 
 import strawberry
-from openinference.semconv.trace import SpanAttributes
 from sqlalchemy import and_, desc, nulls_last
 from sqlalchemy.orm import InstrumentedAttribute
 from sqlalchemy.sql.expression import Select
@@ -16,10 +15,6 @@ from phoenix.server.api.types.pagination import CursorSortColumnDataType
 from phoenix.server.api.types.SortDir import SortDir
 from phoenix.trace.schemas import SpanID
 
-LLM_TOKEN_COUNT_PROMPT = SpanAttributes.LLM_TOKEN_COUNT_PROMPT.split(".")
-LLM_TOKEN_COUNT_COMPLETION = SpanAttributes.LLM_TOKEN_COUNT_COMPLETION.split(".")
-LLM_TOKEN_COUNT_TOTAL = SpanAttributes.LLM_TOKEN_COUNT_TOTAL.split(".")
-
 
 @strawberry.enum
 class SpanColumn(Enum):
@@ -47,11 +42,11 @@ class SpanColumn(Enum):
         elif self is SpanColumn.latencyMs:
             expr = models.Span.latency_ms
         elif self is SpanColumn.tokenCountTotal:
-            expr = models.Span.attributes[LLM_TOKEN_COUNT_TOTAL].as_float()
+            expr = models.Span.llm_token_count_total
         elif self is SpanColumn.tokenCountPrompt:
-            expr = models.Span.attributes[LLM_TOKEN_COUNT_PROMPT].as_float()
+            expr = models.Span.llm_token_count_prompt
         elif self is SpanColumn.tokenCountCompletion:
-            expr = models.Span.attributes[LLM_TOKEN_COUNT_COMPLETION].as_float()
+            expr = models.Span.llm_token_count_completion
         elif self is SpanColumn.cumulativeTokenCountTotal:
             expr = (
                 models.Span.cumulative_llm_token_count_prompt

phoenix/server/api/mutations/dataset_mutations.py

@@ -33,6 +33,7 @@ from phoenix.server.api.types.DatasetExample import DatasetExample
 from phoenix.server.api.types.node import from_global_id_with_expected_type
 from phoenix.server.api.types.Span import Span
 from phoenix.server.api.utils import delete_projects, delete_traces
+from phoenix.server.dml_event import DatasetDeleteEvent, DatasetInsertEvent
 
 
 @strawberry.type
@@ -62,6 +63,7 @@ class DatasetMutationMixin:
                 .returning(models.Dataset)
             )
             assert dataset is not None
+        info.context.event_queue.put(DatasetInsertEvent((dataset.id,)))
         return DatasetMutationPayload(dataset=to_gql_dataset(dataset))
 
     @strawberry.mutation(permission_classes=[IsAuthenticated])  # type: ignore
@@ -90,6 +92,7 @@ class DatasetMutationMixin:
                 .values(**patch)
             )
             assert dataset is not None
+        info.context.event_queue.put(DatasetInsertEvent((dataset.id,)))
         return DatasetMutationPayload(dataset=to_gql_dataset(dataset))
 
     @strawberry.mutation(permission_classes=[IsAuthenticated])  # type: ignore
@@ -218,6 +221,7 @@ class DatasetMutationMixin:
                     for dataset_example_rowid, span in zip(dataset_example_rowids, spans)
                 ],
             )
+        info.context.event_queue.put(DatasetInsertEvent((dataset.id,)))
         return DatasetMutationPayload(dataset=to_gql_dataset(dataset))
 
     @strawberry.mutation(permission_classes=[IsAuthenticated])  # type: ignore
@@ -303,6 +307,7 @@ class DatasetMutationMixin:
                     )
                 ],
             )
+        info.context.event_queue.put(DatasetInsertEvent((dataset.id,)))
         return DatasetMutationPayload(dataset=to_gql_dataset(dataset))
 
     @strawberry.mutation(permission_classes=[IsAuthenticated])  # type: ignore
@@ -333,6 +338,7 @@ class DatasetMutationMixin:
             delete_traces(info.context.db, *eval_trace_ids),
             return_exceptions=True,
         )
+        info.context.event_queue.put(DatasetDeleteEvent((dataset.id,)))
         return DatasetMutationPayload(dataset=to_gql_dataset(dataset))
 
     @strawberry.mutation(permission_classes=[IsAuthenticated])  # type: ignore
@@ -424,7 +430,7 @@ class DatasetMutationMixin:
                     for revision, patch, example_id in zip(revisions, patches, example_ids)
                 ],
             )
-
+        info.context.event_queue.put(DatasetInsertEvent((dataset.id,)))
         return DatasetMutationPayload(dataset=to_gql_dataset(dataset))
 
     @strawberry.mutation(permission_classes=[IsAuthenticated])  # type: ignore
@@ -507,8 +513,8 @@ class DatasetMutationMixin:
                     for dataset_example_rowid in example_db_ids
                 ],
             )
-
-            return DatasetMutationPayload(dataset=to_gql_dataset(dataset))
+        info.context.event_queue.put(DatasetInsertEvent((dataset.id,)))
+        return DatasetMutationPayload(dataset=to_gql_dataset(dataset))
 
 
 def _span_attribute(semconv: str) -> Any:

phoenix/server/api/mutations/experiment_mutations.py

@@ -14,6 +14,7 @@ from phoenix.server.api.mutations.auth import IsAuthenticated
 from phoenix.server.api.types.Experiment import Experiment, to_gql_experiment
 from phoenix.server.api.types.node import from_global_id_with_expected_type
 from phoenix.server.api.utils import delete_projects, delete_traces
+from phoenix.server.dml_event import ExperimentDeleteEvent
 
 
 @strawberry.type
@@ -66,6 +67,7 @@ class ExperimentMutationMixin:
             delete_traces(info.context.db, *eval_trace_ids),
             return_exceptions=True,
         )
+        info.context.event_queue.put(ExperimentDeleteEvent(tuple(experiments.keys())))
         return ExperimentMutationPayload(
             experiments=[
                 to_gql_experiment(experiments[experiment_id]) for experiment_id in experiment_ids

phoenix/server/api/mutations/project_mutations.py

@@ -6,23 +6,23 @@ from strawberry.types import Info
 
 from phoenix.config import DEFAULT_PROJECT_NAME
 from phoenix.db import models
-from phoenix.db.insertion.span import ClearProjectSpansEvent
 from phoenix.server.api.context import Context
 from phoenix.server.api.input_types.ClearProjectInput import ClearProjectInput
 from phoenix.server.api.mutations.auth import IsAuthenticated
 from phoenix.server.api.queries import Query
 from phoenix.server.api.types.node import from_global_id_with_expected_type
+from phoenix.server.dml_event import ProjectDeleteEvent, SpanDeleteEvent
 
 
 @strawberry.type
 class ProjectMutationMixin:
     @strawberry.mutation(permission_classes=[IsAuthenticated])  # type: ignore
     async def delete_project(self, info: Info[Context, None], id: GlobalID) -> Query:
-        node_id = from_global_id_with_expected_type(global_id=id, expected_type_name="Project")
+        project_id = from_global_id_with_expected_type(global_id=id, expected_type_name="Project")
         async with info.context.db() as session:
             project = await session.scalar(
                 select(models.Project)
-                .where(models.Project.id == node_id)
+                .where(models.Project.id == project_id)
                 .options(load_only(models.Project.name))
             )
             if project is None:
@@ -30,6 +30,7 @@ class ProjectMutationMixin:
             if project.name == DEFAULT_PROJECT_NAME:
                 raise ValueError(f"Cannot delete the {DEFAULT_PROJECT_NAME} project")
             await session.delete(project)
+        info.context.event_queue.put(ProjectDeleteEvent((project_id,)))
         return Query()
 
     @strawberry.mutation(permission_classes=[IsAuthenticated])  # type: ignore
@@ -42,6 +43,5 @@ class ProjectMutationMixin:
             delete_statement = delete_statement.where(models.Trace.start_time < input.end_time)
         async with info.context.db() as session:
             await session.execute(delete_statement)
-        if cache := info.context.cache_for_dataloaders:
-            cache.invalidate(ClearProjectSpansEvent(project_rowid=project_id))
+        info.context.event_queue.put(SpanDeleteEvent((project_id,)))
         return Query()