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

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

@@ -0,0 +1,142 @@
+from typing import Optional
+
+from fastapi import APIRouter, Depends, HTTPException, Query
+from pydantic import Field
+from sqlalchemy import select
+from starlette.requests import Request
+from strawberry.relay import GlobalID
+
+from phoenix.db import models
+from phoenix.db.helpers import SupportedSQLDialect
+from phoenix.db.insertion.helpers import as_kv, insert_on_conflict
+from phoenix.server.api.routers.v1.annotations import SpanDocumentAnnotationData
+from phoenix.server.api.types.DocumentAnnotation import DocumentAnnotation
+from phoenix.server.authorization import is_not_locked
+from phoenix.server.bearer_auth import PhoenixUser
+from phoenix.server.dml_event import DocumentAnnotationInsertEvent
+
+from .models import V1RoutesBaseModel
+from .utils import RequestBody, ResponseBody, add_errors_to_responses
+
+# Since the document annotations are spans related, we place it under spans
+router = APIRouter(tags=["spans"])
+
+
+class AnnotateSpanDocumentsRequestBody(RequestBody[list[SpanDocumentAnnotationData]]):
+    pass
+
+
+class InsertedSpanDocumentAnnotation(V1RoutesBaseModel):
+    id: str = Field(description="The ID of the inserted span document annotation")
+
+
+class AnnotateSpanDocumentsResponseBody(ResponseBody[list[InsertedSpanDocumentAnnotation]]):
+    pass
+
+
+@router.post(
+    "/document_annotations",
+    dependencies=[Depends(is_not_locked)],
+    operation_id="annotateSpanDocuments",
+    responses=add_errors_to_responses(
+        [
+            {
+                "status_code": 404,
+                "description": "Span not found",
+            },
+            {
+                "status_code": 422,
+                "description": "Invalid request - non-empty identifier not supported",
+            },
+        ]
+    ),
+    response_description="Span document annotation inserted successfully",
+    include_in_schema=True,
+)
+async def annotate_span_documents(
+    request: Request,
+    request_body: AnnotateSpanDocumentsRequestBody,
+    sync: bool = Query(
+        default=False, description="If set to true, the annotations are inserted synchronously."
+    ),
+) -> AnnotateSpanDocumentsResponseBody:
+    if not request_body.data:
+        return AnnotateSpanDocumentsResponseBody(data=[])
+
+    # Validate that identifiers are empty or only whitespace
+    for annotation in request_body.data:
+        if annotation.identifier.strip():
+            raise HTTPException(
+                detail=f"Non-empty identifier '{annotation.identifier}' is not supported",
+                status_code=422,  # Unprocessable Entity
+            )
+
+    user_id: Optional[int] = None
+    if request.app.state.authentication_enabled and isinstance(request.user, PhoenixUser):
+        user_id = int(request.user.identity)
+
+    span_document_annotations = request_body.data
+
+    precursors = [
+        annotation.as_precursor(user_id=user_id) for annotation in span_document_annotations
+    ]
+    if not sync:
+        await request.state.enqueue_annotations(*precursors)
+        return AnnotateSpanDocumentsResponseBody(data=[])
+
+    span_ids = {p.span_id for p in precursors}
+    # Account for the fact that the spans could arrive after the annotation
+    async with request.app.state.db() as session:
+        existing_spans = {
+            span_id: (id_, num_docs)
+            async for span_id, id_, num_docs in await session.stream(
+                select(models.Span.span_id, models.Span.id, models.Span.num_documents).filter(
+                    models.Span.span_id.in_(span_ids)
+                )
+            )
+        }
+
+        missing_span_ids = span_ids - set(existing_spans.keys())
+        # We prefer to fail the entire operation if there are missing spans in sync mode
+        if missing_span_ids:
+            raise HTTPException(
+                detail=f"Spans with IDs {', '.join(missing_span_ids)} do not exist.",
+                status_code=404,
+            )
+
+        # Validate that document positions are within bounds
+        for annotation in span_document_annotations:
+            _, num_docs = existing_spans[annotation.span_id]
+            if annotation.document_position not in range(num_docs):
+                raise HTTPException(
+                    detail=f"Document position {annotation.document_position} is out of bounds for "
+                    f"span {annotation.span_id} (max: {num_docs - 1})",
+                    status_code=422,  # Unprocessable Entity
+                )
+
+        inserted_document_annotation_ids = []
+        dialect = SupportedSQLDialect(session.bind.dialect.name)
+        for anno in precursors:
+            span_rowid, _ = existing_spans[anno.span_id]
+            values = dict(as_kv(anno.as_insertable(span_rowid).row))
+            span_document_annotation_id = await session.scalar(
+                insert_on_conflict(
+                    values,
+                    dialect=dialect,
+                    table=models.DocumentAnnotation,
+                    unique_by=("name", "span_rowid", "identifier", "document_position"),
+                    constraint_name="uq_document_annotations_name_span_rowid_document_pos_identifier",
+                ).returning(models.DocumentAnnotation.id)
+            )
+            inserted_document_annotation_ids.append(span_document_annotation_id)
+
+    # We queue an event to let the application know that annotations have changed
+    request.state.event_queue.put(
+        DocumentAnnotationInsertEvent(tuple(inserted_document_annotation_ids))
+    )
+    return AnnotateSpanDocumentsResponseBody(
+        data=[
+            InsertedSpanDocumentAnnotation(id=str(GlobalID(DocumentAnnotation.__name__, str(id_))))
+            for id_ in inserted_document_annotation_ids
+        ]
+    )

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

@@ -1,5 +1,6 @@
 import gzip
 from collections.abc import Callable
+from datetime import datetime, timezone
 from itertools import chain
 from typing import Any, Iterator, Optional, Union, cast
 
@@ -14,12 +15,6 @@ from starlette.background import BackgroundTask
 from starlette.datastructures import State
 from starlette.requests import Request
 from starlette.responses import Response, StreamingResponse
-from starlette.status import (
-    HTTP_204_NO_CONTENT,
-    HTTP_404_NOT_FOUND,
-    HTTP_415_UNSUPPORTED_MEDIA_TYPE,
-    HTTP_422_UNPROCESSABLE_ENTITY,
-)
 from typing_extensions import TypeAlias
 
 import phoenix.trace.v1 as pb
@@ -49,16 +44,16 @@ router = APIRouter(tags=["traces"], include_in_schema=True)
     dependencies=[Depends(is_not_locked)],
     operation_id="addEvaluations",
     summary="Add span, trace, or document evaluations",
-    status_code=HTTP_204_NO_CONTENT,
+    status_code=204,
     responses=add_errors_to_responses(
         [
             {
-                "status_code": HTTP_415_UNSUPPORTED_MEDIA_TYPE,
+                "status_code": 415,
                 "description": (
                     "Unsupported content type, only gzipped protobuf and pandas-arrow are supported"
                 ),
             },
-            HTTP_422_UNPROCESSABLE_ENTITY,
+            422,
         ]
     ),
     openapi_extra={
@@ -79,29 +74,23 @@ async def post_evaluations(
     if content_type == "application/x-pandas-arrow":
         return await _process_pyarrow(request)
     if content_type != "application/x-protobuf":
-        raise HTTPException(
-            detail="Unsupported content type", status_code=HTTP_415_UNSUPPORTED_MEDIA_TYPE
-        )
+        raise HTTPException(detail="Unsupported content type", status_code=415)
     body = await request.body()
     if content_encoding == "gzip":
         body = gzip.decompress(body)
     elif content_encoding:
-        raise HTTPException(
-            detail="Unsupported content encoding", status_code=HTTP_415_UNSUPPORTED_MEDIA_TYPE
-        )
+        raise HTTPException(detail="Unsupported content encoding", status_code=415)
     evaluation = pb.Evaluation()
     try:
         evaluation.ParseFromString(body)
     except DecodeError:
-        raise HTTPException(
-            detail="Request body is invalid", status_code=HTTP_422_UNPROCESSABLE_ENTITY
-        )
+        raise HTTPException(detail="Request body is invalid", status_code=422)
     if not evaluation.name.strip():
         raise HTTPException(
             detail="Evaluation name must not be blank/empty",
-            status_code=HTTP_422_UNPROCESSABLE_ENTITY,
+            status_code=422,
         )
-    await request.state.queue_evaluation_for_bulk_insert(evaluation)
+    await request.state.enqueue_evaluation(evaluation)
     return Response()
 
 
@@ -109,7 +98,7 @@ async def post_evaluations(
     "/evaluations",
     operation_id="getEvaluations",
     summary="Get span, trace, or document evaluations from a project",
-    responses=add_errors_to_responses([HTTP_404_NOT_FOUND]),
+    responses=add_errors_to_responses([404]),
 )
 async def get_evaluations(
     request: Request,
@@ -148,7 +137,7 @@ async def get_evaluations(
         and span_evals_dataframe.empty
         and document_evals_dataframe.empty
     ):
-        return Response(status_code=HTTP_404_NOT_FOUND)
+        return Response(status_code=404)
 
     evals = chain(
         map(
@@ -178,7 +167,7 @@ async def _process_pyarrow(request: Request) -> Response:
     except pa.ArrowInvalid:
         raise HTTPException(
             detail="Request body is not valid pyarrow",
-            status_code=HTTP_422_UNPROCESSABLE_ENTITY,
+            status_code=422,
         )
     try:
         evaluations = Evaluations.from_pyarrow_reader(reader)
@@ -186,11 +175,11 @@ async def _process_pyarrow(request: Request) -> Response:
         if isinstance(e, PhoenixEvaluationNameIsMissing):
             raise HTTPException(
                 detail="Evaluation name must not be blank/empty",
-                status_code=HTTP_422_UNPROCESSABLE_ENTITY,
+                status_code=422,
             )
         raise HTTPException(
             detail="Invalid data in request body",
-            status_code=HTTP_422_UNPROCESSABLE_ENTITY,
+            status_code=422,
         )
     return Response(background=BackgroundTask(_add_evaluations, request.state, evaluations))
 
@@ -220,7 +209,7 @@ async def _add_evaluations(state: State, evaluations: Evaluations) -> None:
                 explanation=explanation,
                 metadata_={},
             )
-            await state.enqueue(document_annotation)
+            await state.enqueue_annotations(document_annotation)
     elif len(names) == 1 and names[0] in ("context.span_id", "span_id"):
         for index, row in dataframe.iterrows():
             score, label, explanation = _get_annotation_result(row)
@@ -234,7 +223,7 @@ async def _add_evaluations(state: State, evaluations: Evaluations) -> None:
                 explanation=explanation,
                 metadata_={},
             )
-            await state.enqueue(span_annotation)
+            await state.enqueue_annotations(span_annotation)
     elif len(names) == 1 and names[0] in ("context.trace_id", "trace_id"):
         for index, row in dataframe.iterrows():
             score, label, explanation = _get_annotation_result(row)
@@ -248,7 +237,7 @@ async def _add_evaluations(state: State, evaluations: Evaluations) -> None:
                 explanation=explanation,
                 metadata_={},
             )
-            await state.enqueue(trace_annotation)
+            await state.enqueue_annotations(trace_annotation)
 
 
 def _get_annotation_result(
@@ -269,6 +258,7 @@ def _document_annotation_factory(
     Callable[..., Precursors.DocumentAnnotation],
 ]:
     return lambda index: lambda **kwargs: Precursors.DocumentAnnotation(
+        datetime.now(timezone.utc),
         span_id=str(index[span_id_idx]),
         document_position=int(index[document_position_idx]),
         obj=models.DocumentAnnotation(
@@ -280,6 +270,7 @@ def _document_annotation_factory(
 
 def _span_annotation_factory(span_id: str) -> Callable[..., Precursors.SpanAnnotation]:
     return lambda **kwargs: Precursors.SpanAnnotation(
+        datetime.now(timezone.utc),
         span_id=str(span_id),
         obj=models.SpanAnnotation(**kwargs),
     )
@@ -287,6 +278,7 @@ def _span_annotation_factory(span_id: str) -> Callable[..., Precursors.SpanAnnot
 
 def _trace_annotation_factory(trace_id: str) -> Callable[..., Precursors.TraceAnnotation]:
     return lambda **kwargs: Precursors.TraceAnnotation(
+        datetime.now(timezone.utc),
         trace_id=str(trace_id),
         obj=models.TraceAnnotation(**kwargs),
     )

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

@@ -3,10 +3,10 @@ from typing import Any, Literal, Optional
 
 from dateutil.parser import isoparse
 from fastapi import APIRouter, HTTPException
-from pydantic import Field
+from pydantic import Field, model_validator
 from starlette.requests import Request
-from starlette.status import HTTP_404_NOT_FOUND
 from strawberry.relay import GlobalID
+from typing_extensions import Self
 
 from phoenix.db import models
 from phoenix.db.helpers import SupportedSQLDialect
@@ -36,15 +36,25 @@ class UpsertExperimentEvaluationRequestBody(V1RoutesBaseModel):
     )
     start_time: datetime = Field(description="The start time of the evaluation in ISO format")
     end_time: datetime = Field(description="The end time of the evaluation in ISO format")
-    result: ExperimentEvaluationResult = Field(description="The result of the evaluation")
+    result: Optional[ExperimentEvaluationResult] = Field(
+        None, description="The result of the evaluation. Either result or error must be provided."
+    )
     error: Optional[str] = Field(
-        None, description="Optional error message if the evaluation encountered an error"
+        None,
+        description="Error message if the evaluation encountered an error. "
+        "Either result or error must be provided.",
     )
     metadata: Optional[dict[str, Any]] = Field(
         default=None, description="Metadata for the evaluation"
     )
     trace_id: Optional[str] = Field(default=None, description="Optional trace ID for tracking")
 
+    @model_validator(mode="after")
+    def validate_result_or_error(self) -> Self:
+        if self.result is None and self.error is None:
+            raise ValueError("Either 'result' or 'error' must be provided")
+        return self
+
 
 class UpsertExperimentEvaluationResponseBodyData(V1RoutesBaseModel):
     id: str = Field(description="The ID of the upserted experiment evaluation")
@@ -61,7 +71,7 @@ class UpsertExperimentEvaluationResponseBody(
     operation_id="upsertExperimentEvaluation",
     summary="Create or update evaluation for an experiment run",
     responses=add_errors_to_responses(
-        [{"status_code": HTTP_404_NOT_FOUND, "description": "Experiment run not found"}]
+        [{"status_code": 404, "description": "Experiment run not found"}]
     ),
 )
 async def upsert_experiment_evaluation(
@@ -74,7 +84,7 @@ async def upsert_experiment_evaluation(
     except ValueError:
         raise HTTPException(
             detail=f"ExperimentRun with ID {experiment_run_gid} does not exist",
-            status_code=HTTP_404_NOT_FOUND,
+            status_code=404,
         )
     name = request_body.name
     annotator_kind = request_body.annotator_kind