arize-phoenix 4.19.0__py3-none-any.whl → 4.20.1__py3-none-any.whl

phoenix/server/api/routers/v1/spans.py

@@ -16,6 +16,7 @@ from phoenix.db.helpers import SupportedSQLDialect
 from phoenix.db.insertion.helpers import as_kv, insert_on_conflict
 from phoenix.db.insertion.types import Precursors
 from phoenix.server.api.routers.utils import df_to_bytes
+from phoenix.server.dml_event import SpanAnnotationInsertEvent
 from phoenix.trace.dsl import SpanQuery as SpanQuery_
 
 from .pydantic_compat import V1RoutesBaseModel
@@ -23,7 +24,7 @@ from .utils import RequestBody, ResponseBody, add_errors_to_responses
 
 DEFAULT_SPAN_LIMIT = 1000
 
-router = APIRouter(tags=["traces"], include_in_schema=False)
+router = APIRouter(tags=["spans"])
 
 
 class SpanQuery(V1RoutesBaseModel):
@@ -65,6 +66,7 @@ class QuerySpansRequestBody(V1RoutesBaseModel):
     operation_id="querySpans",
     summary="Query spans with query DSL",
     responses=add_errors_to_responses([HTTP_404_NOT_FOUND, HTTP_422_UNPROCESSABLE_ENTITY]),
+    include_in_schema=False,
 )
 async def query_spans_handler(
     request: Request,
@@ -189,12 +191,15 @@ class AnnotateSpansResponseBody(ResponseBody[List[InsertedSpanAnnotation]]):
         [{"status_code": HTTP_404_NOT_FOUND, "description": "Span not found"}]
     ),
     response_description="Span annotations inserted successfully",
+    include_in_schema=True,
 )
 async def annotate_spans(
     request: Request,
     request_body: AnnotateSpansRequestBody,
     sync: bool = Query(default=True, description="If true, fulfill request synchronously."),
 ) -> AnnotateSpansResponseBody:
+    if not request_body.data:
+        return AnnotateSpansResponseBody(data=[])
     precursors = [d.as_precursor() for d in request_body.data]
     if not sync:
         await request.state.enqueue(*precursors)
@@ -215,9 +220,7 @@ async def annotate_spans(
                 detail=f"Spans with IDs {', '.join(missing_span_ids)} do not exist.",
                 status_code=HTTP_404_NOT_FOUND,
             )
-
-        inserted_annotations = []
-
+        inserted_ids = []
         dialect = SupportedSQLDialect(session.bind.dialect.name)
         for p in precursors:
             values = dict(as_kv(p.as_insertable(existing_spans[p.span_id]).row))
@@ -229,8 +232,11 @@ async def annotate_spans(
                     unique_by=("name", "span_rowid"),
                 ).returning(models.SpanAnnotation.id)
             )
-            inserted_annotations.append(
-                InsertedSpanAnnotation(id=str(GlobalID("SpanAnnotation", str(span_annotation_id))))
-            )
-
-    return AnnotateSpansResponseBody(data=inserted_annotations)
+            inserted_ids.append(span_annotation_id)
+    request.state.event_queue.put(SpanAnnotationInsertEvent(tuple(inserted_ids)))
+    return AnnotateSpansResponseBody(
+        data=[
+            InsertedSpanAnnotation(id=str(GlobalID("SpanAnnotation", str(id_))))
+            for id_ in inserted_ids
+        ]
+    )

phoenix/server/api/routers/v1/traces.py

@@ -24,13 +24,14 @@ from phoenix.db import models
 from phoenix.db.helpers import SupportedSQLDialect
 from phoenix.db.insertion.helpers import as_kv, insert_on_conflict
 from phoenix.db.insertion.types import Precursors
+from phoenix.server.dml_event import TraceAnnotationInsertEvent
 from phoenix.trace.otel import decode_otlp_span
 from phoenix.utilities.project import get_project_name
 
 from .pydantic_compat import V1RoutesBaseModel
 from .utils import RequestBody, ResponseBody, add_errors_to_responses
 
-router = APIRouter(tags=["traces"], include_in_schema=False)
+router = APIRouter(tags=["traces"])
 
 
 @router.post(
@@ -57,6 +58,7 @@ router = APIRouter(tags=["traces"], include_in_schema=False)
             },
         }
     },
+    include_in_schema=False,
 )
 async def post_traces(
     request: Request,
@@ -145,12 +147,15 @@ class AnnotateTracesResponseBody(ResponseBody[List[InsertedTraceAnnotation]]):
     responses=add_errors_to_responses(
         [{"status_code": HTTP_404_NOT_FOUND, "description": "Trace not found"}]
     ),
+    include_in_schema=False,
 )
 async def annotate_traces(
     request: Request,
     request_body: AnnotateTracesRequestBody,
     sync: bool = Query(default=True, description="If true, fulfill request synchronously."),
 ) -> AnnotateTracesResponseBody:
+    if not request_body.data:
+        return AnnotateTracesResponseBody(data=[])
     precursors = [d.as_precursor() for d in request_body.data]
     if not sync:
         await request.state.enqueue(*precursors)
@@ -171,9 +176,7 @@ async def annotate_traces(
                 detail=f"Traces with IDs {', '.join(missing_trace_ids)} do not exist.",
                 status_code=HTTP_404_NOT_FOUND,
             )
-
-        inserted_annotations = []
-
+        inserted_ids = []
         dialect = SupportedSQLDialect(session.bind.dialect.name)
         for p in precursors:
             values = dict(as_kv(p.as_insertable(existing_traces[p.trace_id]).row))
@@ -185,13 +188,14 @@ async def annotate_traces(
                     unique_by=("name", "trace_rowid"),
                 ).returning(models.TraceAnnotation.id)
             )
-            inserted_annotations.append(
-                InsertedTraceAnnotation(
-                    id=str(GlobalID("TraceAnnotation", str(trace_annotation_id)))
-                )
-            )
-
-    return AnnotateTracesResponseBody(data=inserted_annotations)
+            inserted_ids.append(trace_annotation_id)
+    request.state.event_queue.put(TraceAnnotationInsertEvent(tuple(inserted_ids)))
+    return AnnotateTracesResponseBody(
+        data=[
+            InsertedTraceAnnotation(id=str(GlobalID("TraceAnnotation", str(id_))))
+            for id_ in inserted_ids
+        ]
+    )
 
 
 async def _add_spans(req: ExportTraceServiceRequest, state: State) -> None:

phoenix/server/api/types/Dataset.py

@@ -1,5 +1,5 @@
 from datetime import datetime
-from typing import AsyncIterable, List, Optional, Tuple, cast
+from typing import AsyncIterable, ClassVar, List, Optional, Tuple, Type, cast
 
 import strawberry
 from sqlalchemy import and_, func, select
@@ -27,6 +27,7 @@ from phoenix.server.api.types.SortDir import SortDir
 
 @strawberry.type
 class Dataset(Node):
+    _table: ClassVar[Type[models.Base]] = models.Experiment
     id_attr: NodeID[int]
     name: str
     description: Optional[str]
@@ -284,6 +285,10 @@ class Dataset(Node):
                 ) in await session.stream(query)
             ]
 
+    @strawberry.field
+    def last_updated_at(self, info: Info[Context, None]) -> Optional[datetime]:
+        return info.context.last_updated_at.get(self._table, self.id_attr)
+
 
 def to_gql_dataset(dataset: models.Dataset) -> Dataset:
     """

phoenix/server/api/types/Experiment.py

@@ -1,5 +1,5 @@
 from datetime import datetime
-from typing import List, Optional
+from typing import ClassVar, List, Optional, Type
 
 import strawberry
 from sqlalchemy import select
@@ -23,6 +23,7 @@ from phoenix.server.api.types.Project import Project
 
 @strawberry.type
 class Experiment(Node):
+    _table: ClassVar[Type[models.Base]] = models.Experiment
     cached_sequence_number: Private[Optional[int]] = None
     id_attr: NodeID[int]
     name: str
@@ -127,6 +128,10 @@ class Experiment(Node):
             gradient_end_color=db_project.gradient_end_color,
         )
 
+    @strawberry.field
+    def last_updated_at(self, info: Info[Context, None]) -> Optional[datetime]:
+        return info.context.last_updated_at.get(self._table, self.id_attr)
+
 
 def to_gql_experiment(
     experiment: models.Experiment,

phoenix/server/api/types/Project.py

@@ -2,8 +2,10 @@ import operator
 from datetime import datetime
 from typing import (
     Any,
+    ClassVar,
     List,
     Optional,
+    Type,
 )
 
 import strawberry
@@ -38,6 +40,7 @@ from phoenix.trace.dsl import SpanFilter
 
 @strawberry.type
 class Project(Node):
+    _table: ClassVar[Type[models.Base]] = models.Project
     id_attr: NodeID[int]
     name: str
     gradient_start_color: str
@@ -397,7 +400,7 @@ class Project(Node):
         self,
         info: Info[Context, None],
     ) -> Optional[datetime]:
-        return info.context.streaming_last_updated_at(self.id_attr)
+        return info.context.last_updated_at.get(self._table, self.id_attr)
 
     @strawberry.field
     async def validate_span_filter_condition(self, condition: str) -> ValidationResult:

phoenix/server/api/types/Span.py

@@ -7,7 +7,6 @@ from typing import TYPE_CHECKING, Any, List, Mapping, Optional, Sized, cast
 import numpy as np
 import strawberry
 from openinference.semconv.trace import EmbeddingAttributes, SpanAttributes
-from sqlalchemy import select
 from strawberry import ID, UNSET
 from strawberry.relay import Node, NodeID
 from strawberry.types import Info
@@ -20,7 +19,10 @@ from phoenix.server.api.helpers.dataset_helpers import (
     get_dataset_example_input,
     get_dataset_example_output,
 )
-from phoenix.server.api.input_types.SpanAnnotationSort import SpanAnnotationSort
+from phoenix.server.api.input_types.SpanAnnotationSort import (
+    SpanAnnotationColumn,
+    SpanAnnotationSort,
+)
 from phoenix.server.api.types.SortDir import SortDir
 from phoenix.server.api.types.SpanAnnotation import to_gql_span_annotation
 from phoenix.trace.attributes import get_attribute_value
@@ -190,17 +192,16 @@ class Span(Node):
         info: Info[Context, None],
         sort: Optional[SpanAnnotationSort] = UNSET,
     ) -> List[SpanAnnotation]:
-        async with info.context.db() as session:
-            stmt = select(models.SpanAnnotation).filter_by(span_rowid=self.id_attr)
-            if sort:
-                sort_col = getattr(models.SpanAnnotation, sort.col.value)
-                if sort.dir is SortDir.desc:
-                    stmt = stmt.order_by(sort_col.desc(), models.SpanAnnotation.id.desc())
-                else:
-                    stmt = stmt.order_by(sort_col.asc(), models.SpanAnnotation.id.asc())
-            else:
-                stmt = stmt.order_by(models.SpanAnnotation.created_at.desc())
-            annotations = await session.scalars(stmt)
+        span_id = self.id_attr
+        annotations = await info.context.data_loaders.span_annotations.load(span_id)
+        sort_key = SpanAnnotationColumn.name.value
+        sort_descending = False
+        if sort:
+            sort_key = sort.col.value
+            sort_descending = sort.dir is SortDir.desc
+        annotations.sort(
+            key=lambda annotation: getattr(annotation, sort_key), reverse=sort_descending
+        )
         return [to_gql_span_annotation(annotation) for annotation in annotations]
 
     @strawberry.field(

phoenix/server/app.py

@@ -2,7 +2,6 @@ import asyncio
 import contextlib
 import json
 import logging
-from datetime import datetime
 from functools import cached_property
 from pathlib import Path
 from typing import (
@@ -87,9 +86,16 @@ from phoenix.server.api.dataloaders import (
 from phoenix.server.api.routers.v1 import REST_API_VERSION
 from phoenix.server.api.routers.v1 import router as v1_router
 from phoenix.server.api.schema import schema
+from phoenix.server.dml_event import DmlEvent
+from phoenix.server.dml_event_handler import DmlEventHandler
 from phoenix.server.grpc_server import GrpcServer
 from phoenix.server.telemetry import initialize_opentelemetry_tracer_provider
-from phoenix.server.types import DbSessionFactory
+from phoenix.server.types import (
+    CanGetLastUpdatedAt,
+    CanPutItem,
+    DbSessionFactory,
+    LastUpdatedAt,
+)
 from phoenix.trace.schemas import Span
 from phoenix.utilities.client import PHOENIX_SERVER_VERSION_HEADER
 
@@ -220,6 +226,7 @@ def _lifespan(
     *,
     dialect: SupportedSQLDialect,
     bulk_inserter: BulkInserter,
+    dml_event_handler: DmlEventHandler,
     tracer_provider: Optional["TracerProvider"] = None,
     enable_prometheus: bool = False,
     clean_ups: Iterable[Callable[[], None]] = (),
@@ -239,8 +246,9 @@ def _lifespan(
             disabled=read_only,
             tracer_provider=tracer_provider,
             enable_prometheus=enable_prometheus,
-        ):
+        ), dml_event_handler:
             yield {
+                "event_queue": dml_event_handler,
                 "enqueue": enqueue,
                 "queue_span_for_bulk_insert": queue_span,
                 "queue_evaluation_for_bulk_insert": queue_evaluation,
@@ -263,9 +271,10 @@ def create_graphql_router(
     db: DbSessionFactory,
     model: Model,
     export_path: Path,
+    last_updated_at: CanGetLastUpdatedAt,
     corpus: Optional[Model] = None,
-    streaming_last_updated_at: Callable[[ProjectRowId], Optional[datetime]] = lambda _: None,
     cache_for_dataloaders: Optional[CacheForDataLoaders] = None,
+    event_queue: CanPutItem[DmlEvent],
     read_only: bool = False,
 ) -> GraphQLRouter:  # type: ignore[type-arg]
     def get_context() -> Context:
@@ -274,7 +283,8 @@ def create_graphql_router(
             model=model,
             corpus=corpus,
             export_path=export_path,
-            streaming_last_updated_at=streaming_last_updated_at,
+            last_updated_at=last_updated_at,
+            event_queue=event_queue,
             data_loaders=DataLoaders(
                 average_experiment_run_latency=AverageExperimentRunLatencyDataLoader(db),
                 dataset_example_revisions=DatasetExampleRevisionsDataLoader(db),
@@ -420,11 +430,16 @@ def create_app(
     cache_for_dataloaders = (
         CacheForDataLoaders() if db.dialect is SupportedSQLDialect.SQLITE else None
     )
-
+    last_updated_at = LastUpdatedAt()
+    dml_event_handler = DmlEventHandler(
+        db=db,
+        cache_for_dataloaders=cache_for_dataloaders,
+        last_updated_at=last_updated_at,
+    )
     bulk_inserter = BulkInserter(
         db,
         enable_prometheus=enable_prometheus,
-        cache_for_dataloaders=cache_for_dataloaders,
+        event_queue=dml_event_handler,
         initial_batch_of_spans=initial_batch_of_spans,
         initial_batch_of_evaluations=initial_batch_of_evaluations,
     )
@@ -460,7 +475,8 @@ def create_app(
         model=model,
         corpus=corpus,
         export_path=export_path,
-        streaming_last_updated_at=bulk_inserter.last_updated_at,
+        last_updated_at=last_updated_at,
+        event_queue=dml_event_handler,
         cache_for_dataloaders=cache_for_dataloaders,
         read_only=read_only,
     )
@@ -477,6 +493,7 @@ def create_app(
             dialect=db.dialect,
             read_only=read_only,
             bulk_inserter=bulk_inserter,
+            dml_event_handler=dml_event_handler,
             tracer_provider=tracer_provider,
             enable_prometheus=enable_prometheus,
             clean_ups=clean_ups,

phoenix/server/dml_event.py

@@ -0,0 +1,136 @@
+from __future__ import annotations
+
+from abc import ABC
+from dataclasses import dataclass, field
+from typing import ClassVar, Tuple, Type
+
+from phoenix.db import models
+
+
+@dataclass(frozen=True)
+class DmlEvent(ABC):
+    """
+    Event corresponding to a Data Manipulation Language (DML)
+    operation, e.g. insertion, update, or deletion.
+    """
+
+    table: ClassVar[Type[models.Base]]
+    ids: Tuple[int, ...] = field(default_factory=tuple)
+
+    def __bool__(self) -> bool:
+        return bool(self.ids)
+
+    def __hash__(self) -> int:
+        return id(self)
+
+
+@dataclass(frozen=True)
+class ProjectDmlEvent(DmlEvent):
+    table = models.Project
+
+
+@dataclass(frozen=True)
+class ProjectDeleteEvent(ProjectDmlEvent): ...
+
+
+@dataclass(frozen=True)
+class SpanDmlEvent(ProjectDmlEvent): ...
+
+
+@dataclass(frozen=True)
+class SpanInsertEvent(SpanDmlEvent): ...
+
+
+@dataclass(frozen=True)
+class SpanDeleteEvent(SpanDmlEvent): ...
+
+
+@dataclass(frozen=True)
+class DatasetDmlEvent(DmlEvent):
+    table = models.Dataset
+
+
+@dataclass(frozen=True)
+class DatasetInsertEvent(DatasetDmlEvent): ...
+
+
+@dataclass(frozen=True)
+class DatasetDeleteEvent(DatasetDmlEvent): ...
+
+
+@dataclass(frozen=True)
+class ExperimentDmlEvent(DmlEvent):
+    table = models.Experiment
+
+
+@dataclass(frozen=True)
+class ExperimentInsertEvent(ExperimentDmlEvent): ...
+
+
+@dataclass(frozen=True)
+class ExperimentDeleteEvent(ExperimentDmlEvent): ...
+
+
+@dataclass(frozen=True)
+class ExperimentRunDmlEvent(DmlEvent):
+    table = models.ExperimentRun
+
+
+@dataclass(frozen=True)
+class ExperimentRunInsertEvent(ExperimentRunDmlEvent): ...
+
+
+@dataclass(frozen=True)
+class ExperimentRunDeleteEvent(ExperimentRunDmlEvent): ...
+
+
+@dataclass(frozen=True)
+class ExperimentRunAnnotationDmlEvent(DmlEvent):
+    table = models.ExperimentRunAnnotation
+
+
+@dataclass(frozen=True)
+class ExperimentRunAnnotationInsertEvent(ExperimentRunAnnotationDmlEvent): ...
+
+
+@dataclass(frozen=True)
+class ExperimentRunAnnotationDeleteEvent(ExperimentRunAnnotationDmlEvent): ...
+
+
+@dataclass(frozen=True)
+class SpanAnnotationDmlEvent(DmlEvent):
+    table = models.SpanAnnotation
+
+
+@dataclass(frozen=True)
+class SpanAnnotationInsertEvent(SpanAnnotationDmlEvent): ...
+
+
+@dataclass(frozen=True)
+class SpanAnnotationDeleteEvent(SpanAnnotationDmlEvent): ...
+
+
+@dataclass(frozen=True)
+class TraceAnnotationDmlEvent(DmlEvent):
+    table = models.TraceAnnotation
+
+
+@dataclass(frozen=True)
+class TraceAnnotationInsertEvent(TraceAnnotationDmlEvent): ...
+
+
+@dataclass(frozen=True)
+class TraceAnnotationDeleteEvent(TraceAnnotationDmlEvent): ...
+
+
+@dataclass(frozen=True)
+class DocumentAnnotationDmlEvent(DmlEvent):
+    table = models.DocumentAnnotation
+
+
+@dataclass(frozen=True)
+class DocumentAnnotationInsertEvent(DocumentAnnotationDmlEvent): ...
+
+
+@dataclass(frozen=True)
+class DocumentAnnotationDeleteEvent(DocumentAnnotationDmlEvent): ...