arize-phoenix 4.14.1__py3-none-any.whl → 4.16.0__py3-none-any.whl

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

@@ -1,18 +1,19 @@
 import gzip
 import zlib
-from typing import Any, Dict, List
+from typing import Any, Dict, List, Literal, Optional
 
+from fastapi import APIRouter, BackgroundTasks, Header, HTTPException, Query
 from google.protobuf.message import DecodeError
 from opentelemetry.proto.collector.trace.v1.trace_service_pb2 import (
     ExportTraceServiceRequest,
 )
+from pydantic import Field
 from sqlalchemy import select
-from starlette.background import BackgroundTask
 from starlette.concurrency import run_in_threadpool
 from starlette.datastructures import State
 from starlette.requests import Request
-from starlette.responses import JSONResponse, Response
 from starlette.status import (
+    HTTP_204_NO_CONTENT,
     HTTP_404_NOT_FOUND,
     HTTP_415_UNSUPPORTED_MEDIA_TYPE,
     HTTP_422_UNPROCESSABLE_ENTITY,
@@ -21,45 +22,56 @@ from strawberry.relay import GlobalID
 
 from phoenix.db import models
 from phoenix.db.helpers import SupportedSQLDialect
-from phoenix.db.insertion.helpers import insert_on_conflict
-from phoenix.server.api.types.node import from_global_id_with_expected_type
+from phoenix.db.insertion.helpers import as_kv, insert_on_conflict
+from phoenix.db.insertion.types import Precursors
 from phoenix.trace.otel import decode_otlp_span
 from phoenix.utilities.project import get_project_name
 
-
-async def post_traces(request: Request) -> Response:
-    """
-    summary: Send traces to Phoenix
-    operationId: addTraces
-    tags:
-      - private
-    requestBody:
-      required: true
-      content:
-        application/x-protobuf:
-          schema:
-            type: string
-            format: binary
-    responses:
-      200:
-        description: Success
-      403:
-        description: Forbidden
-      415:
-        description: Unsupported content type, only gzipped protobuf
-      422:
-        description: Request body is invalid
-    """
-    content_type = request.headers.get("content-type")
+from .pydantic_compat import V1RoutesBaseModel
+from .utils import RequestBody, ResponseBody, add_errors_to_responses
+
+router = APIRouter(tags=["traces"], include_in_schema=False)
+
+
+@router.post(
+    "/traces",
+    operation_id="addTraces",
+    summary="Send traces",
+    status_code=HTTP_204_NO_CONTENT,
+    responses=add_errors_to_responses(
+        [
+            {
+                "status_code": HTTP_415_UNSUPPORTED_MEDIA_TYPE,
+                "description": (
+                    "Unsupported content type (only `application/x-protobuf` is supported)"
+                ),
+            },
+            {"status_code": HTTP_422_UNPROCESSABLE_ENTITY, "description": "Invalid request body"},
+        ]
+    ),
+    openapi_extra={
+        "requestBody": {
+            "required": True,
+            "content": {
+                "application/x-protobuf": {"schema": {"type": "string", "format": "binary"}}
+            },
+        }
+    },
+)
+async def post_traces(
+    request: Request,
+    background_tasks: BackgroundTasks,
+    content_type: Optional[str] = Header(default=None),
+    content_encoding: Optional[str] = Header(default=None),
+) -> None:
     if content_type != "application/x-protobuf":
-        return Response(
-            content=f"Unsupported content type: {content_type}",
+        raise HTTPException(
+            detail=f"Unsupported content type: {content_type}",
             status_code=HTTP_415_UNSUPPORTED_MEDIA_TYPE,
         )
-    content_encoding = request.headers.get("content-encoding")
     if content_encoding and content_encoding not in ("gzip", "deflate"):
-        return Response(
-            content=f"Unsupported content encoding: {content_encoding}",
+        raise HTTPException(
+            detail=f"Unsupported content encoding: {content_encoding}",
             status_code=HTTP_415_UNSUPPORTED_MEDIA_TYPE,
         )
     body = await request.body()
@@ -71,139 +83,100 @@ async def post_traces(request: Request) -> Response:
     try:
         await run_in_threadpool(req.ParseFromString, body)
     except DecodeError:
-        return Response(
-            content="Request body is invalid ExportTraceServiceRequest",
+        raise HTTPException(
+            detail="Request body is invalid ExportTraceServiceRequest",
             status_code=HTTP_422_UNPROCESSABLE_ENTITY,
         )
-    return Response(background=BackgroundTask(_add_spans, req, request.state))
-
-
-async def annotate_traces(request: Request) -> Response:
-    """
-    summary: Upsert annotations for traces
-    operationId: annotateTraces
-    tags:
-      - private
-    requestBody:
-      description: List of trace annotations to be inserted
-      required: true
-      content:
-        application/json:
-          schema:
-            type: object
-            properties:
-              data:
-                type: array
-                items:
-                  type: object
-                  properties:
-                    trace_id:
-                      type: string
-                      description: The ID of the trace being annotated
-                    name:
-                      type: string
-                      description: The name of the annotation
-                    annotator_kind:
-                      type: string
-                      description: The kind of annotator used for the annotation ("LLM" or "HUMAN")
-                    result:
-                      type: object
-                      description: The result of the annotation
-                      properties:
-                        label:
-                          type: string
-                          description: The label assigned by the annotation
-                        score:
-                          type: number
-                          format: float
-                          description: The score assigned by the annotation
-                        explanation:
-                          type: string
-                          description: Explanation of the annotation result
-                    error:
-                      type: string
-                      description: Optional error message if the annotation encountered an error
-                    metadata:
-                      type: object
-                      description: Metadata for the annotation
-                      additionalProperties:
-                        type: string
-                  required:
-                    - trace_id
-                    - name
-                    - annotator_kind
-    responses:
-      200:
-        description: Trace annotations inserted successfully
-        content:
-          application/json:
-            schema:
-              type: object
-              properties:
-                data:
-                  type: array
-                  items:
-                    type: object
-                    properties:
-                      id:
-                        type: string
-                        description: The ID of the inserted trace annotation
-      404:
-        description: Trace not found
-    """
-    payload: List[Dict[str, Any]] = (await request.json()).get("data", [])
-    trace_gids = [GlobalID.from_id(annotation["trace_id"]) for annotation in payload]
-
-    resolved_trace_ids = []
-    for trace_gid in trace_gids:
-        try:
-            resolved_trace_ids.append(from_global_id_with_expected_type(trace_gid, "Trace"))
-        except ValueError:
-            return Response(
-                content="Trace with ID {trace_gid} does not exist",
-                status_code=HTTP_404_NOT_FOUND,
-            )
+    background_tasks.add_task(_add_spans, req, request.state)
+    return None
+
+
+class TraceAnnotationResult(V1RoutesBaseModel):
+    label: Optional[str] = Field(default=None, description="The label assigned by the annotation")
+    score: Optional[float] = Field(default=None, description="The score assigned by the annotation")
+    explanation: Optional[str] = Field(
+        default=None, description="Explanation of the annotation result"
+    )
+
+
+class TraceAnnotation(V1RoutesBaseModel):
+    trace_id: str = Field(description="OpenTelemetry Trace ID (hex format w/o 0x prefix)")
+    name: str = Field(description="The name of the annotation")
+    annotator_kind: Literal["LLM", "HUMAN"] = Field(
+        description="The kind of annotator used for the annotation"
+    )
+    result: Optional[TraceAnnotationResult] = Field(
+        default=None, description="The result of the annotation"
+    )
+    metadata: Optional[Dict[str, Any]] = Field(
+        default=None, description="Metadata for the annotation"
+    )
+
+    def as_precursor(self) -> Precursors.TraceAnnotation:
+        return Precursors.TraceAnnotation(
+            self.trace_id,
+            models.TraceAnnotation(
+                name=self.name,
+                annotator_kind=self.annotator_kind,
+                score=self.result.score if self.result else None,
+                label=self.result.label if self.result else None,
+                explanation=self.result.explanation if self.result else None,
+                metadata_=self.metadata or {},
+            ),
+        )
+
+
+class AnnotateTracesRequestBody(RequestBody[List[TraceAnnotation]]):
+    data: List[TraceAnnotation] = Field(description="The trace annotations to be upserted")
+
 
+class InsertedTraceAnnotation(V1RoutesBaseModel):
+    id: str = Field(description="The ID of the inserted trace annotation")
+
+
+class AnnotateTracesResponseBody(ResponseBody[List[InsertedTraceAnnotation]]):
+    pass
+
+
+@router.post(
+    "/trace_annotations",
+    operation_id="annotateTraces",
+    summary="Create or update trace annotations",
+    responses=add_errors_to_responses(
+        [{"status_code": HTTP_404_NOT_FOUND, "description": "Trace not found"}]
+    ),
+)
+async def annotate_traces(
+    request: Request,
+    request_body: AnnotateTracesRequestBody,
+    sync: bool = Query(default=True, description="If true, fulfill request synchronously."),
+) -> AnnotateTracesResponseBody:
+    precursors = [d.as_precursor() for d in request_body.data]
+    if not sync:
+        await request.state.enqueue(*precursors)
+        return AnnotateTracesResponseBody(data=[])
+
+    trace_ids = {p.trace_id for p in precursors}
     async with request.app.state.db() as session:
-        traces = await session.execute(
-            select(models.Trace).filter(models.Trace.id.in_(resolved_trace_ids))
-        )
-        existing_trace_ids = {trace.id for trace in traces.scalars()}
+        existing_traces = {
+            trace.trace_id: trace.id
+            async for trace in await session.stream_scalars(
+                select(models.Trace).filter(models.Trace.trace_id.in_(trace_ids))
+            )
+        }
 
-        missing_trace_ids = set(resolved_trace_ids) - existing_trace_ids
+        missing_trace_ids = trace_ids - set(existing_traces.keys())
         if missing_trace_ids:
-            missing_trace_gids = [
-                str(GlobalID("Trace", str(trace_gid))) for trace_gid in missing_trace_ids
-            ]
-            return Response(
-                content=f"Traces with IDs {', '.join(missing_trace_gids)} do not exist.",
+            raise HTTPException(
+                detail=f"Traces with IDs {', '.join(missing_trace_ids)} do not exist.",
                 status_code=HTTP_404_NOT_FOUND,
             )
 
         inserted_annotations = []
 
-        for annotation in payload:
-            trace_gid = GlobalID.from_id(annotation["trace_id"])
-            trace_id = from_global_id_with_expected_type(trace_gid, "Trace")
-
-            name = annotation["name"]
-            annotator_kind = annotation["annotator_kind"]
-            result = annotation.get("result")
-            label = result.get("label") if result else None
-            score = result.get("score") if result else None
-            explanation = result.get("explanation") if result else None
-            metadata = annotation.get("metadata") or {}
-
-            values = dict(
-                trace_rowid=trace_id,
-                name=name,
-                label=label,
-                score=score,
-                explanation=explanation,
-                annotator_kind=annotator_kind,
-                metadata_=metadata,
-            )
-            dialect = SupportedSQLDialect(session.bind.dialect.name)
+        dialect = SupportedSQLDialect(session.bind.dialect.name)
+        for p in precursors:
+            values = dict(as_kv(p.as_insertable(existing_traces[p.trace_id]).row))
             trace_annotation_id = await session.scalar(
                 insert_on_conflict(
                     values,
@@ -213,10 +186,12 @@ async def annotate_traces(request: Request) -> Response:
                 ).returning(models.TraceAnnotation.id)
             )
             inserted_annotations.append(
-                {"id": str(GlobalID("TraceAnnotation", str(trace_annotation_id)))}
+                InsertedTraceAnnotation(
+                    id=str(GlobalID("TraceAnnotation", str(trace_annotation_id)))
+                )
             )
 
-    return JSONResponse(content={"data": inserted_annotations})
+    return AnnotateTracesResponseBody(data=inserted_annotations)
 
 
 async def _add_spans(req: ExportTraceServiceRequest, state: State) -> None:

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

@@ -0,0 +1,95 @@
+from typing import Any, Dict, Generic, List, Optional, TypedDict, Union
+
+from typing_extensions import TypeAlias, TypeVar, assert_never
+
+from .pydantic_compat import V1RoutesBaseModel
+
+StatusCode: TypeAlias = int
+DataType = TypeVar("DataType")
+Responses: TypeAlias = Dict[
+    Union[int, str], Dict[str, Any]
+]  # input type for the `responses` parameter of a fastapi route
+
+
+class StatusCodeWithDescription(TypedDict):
+    """
+    A duck type for a status code with a description detailing under what
+    conditions the status code is raised.
+    """
+
+    status_code: StatusCode
+    description: str
+
+
+class RequestBody(V1RoutesBaseModel, Generic[DataType]):
+    # A generic request type accepted by V1 routes.
+    #
+    # Don't use """ for this docstring or it will be included as a description
+    # in the generated OpenAPI schema.
+    data: DataType
+
+
+class ResponseBody(V1RoutesBaseModel, Generic[DataType]):
+    # A generic response type returned by V1 routes.
+    #
+    # Don't use """ for this docstring or it will be included as a description
+    # in the generated OpenAPI schema.
+
+    data: DataType
+
+
+class PaginatedResponseBody(V1RoutesBaseModel, Generic[DataType]):
+    # A generic paginated response type returned by V1 routes.
+    #
+    # Don't use """ for this docstring or it will be included as a description
+    # in the generated OpenAPI schema.
+
+    data: List[DataType]
+    next_cursor: Optional[str]
+
+
+def add_errors_to_responses(
+    errors: List[Union[StatusCode, StatusCodeWithDescription]],
+    /,
+    *,
+    responses: Optional[Responses] = None,
+) -> Responses:
+    """
+    Creates or updates a patch for an OpenAPI schema's `responses` section to
+    include status codes in the generated OpenAPI schema.
+    """
+    output_responses: Responses = responses or {}
+    for error in errors:
+        status_code: int
+        description: Optional[str] = None
+        if isinstance(error, StatusCode):
+            status_code = error
+        elif isinstance(error, dict):
+            status_code = error["status_code"]
+            description = error["description"]
+        else:
+            assert_never(error)
+        if status_code not in output_responses:
+            output_responses[status_code] = {
+                "content": {"text/plain": {"schema": {"type": "string"}}}
+            }
+        if description:
+            output_responses[status_code]["description"] = description
+    return output_responses
+
+
+def add_text_csv_content_to_responses(
+    status_code: StatusCode, /, *, responses: Optional[Responses] = None
+) -> Responses:
+    """
+    Creates or updates a patch for an OpenAPI schema's `responses` section to
+    ensure that the response for the given status code is marked as text/csv in
+    the generated OpenAPI schema.
+    """
+    output_responses: Responses = responses or {}
+    if status_code not in output_responses:
+        output_responses[status_code] = {}
+    output_responses[status_code]["content"] = {
+        "text/csv": {"schema": {"type": "string", "contentMediaType": "text/csv"}}
+    }
+    return output_responses

phoenix/server/api/types/Span.py

@@ -7,6 +7,7 @@ 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
@@ -19,6 +20,9 @@ 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.types.SortDir import SortDir
+from phoenix.server.api.types.SpanAnnotation import to_gql_span_annotation
 from phoenix.trace.attributes import get_attribute_value
 
 from .DocumentRetrievalMetrics import DocumentRetrievalMetrics
@@ -177,12 +181,27 @@ class Span(Node):
 
     @strawberry.field(
         description=(
-            "Annotations of the span's parent span. This encompasses both "
+            "Annotations associated with the span. This encompasses both "
             "LLM and human annotations."
         )
     )  # type: ignore
-    async def span_annotations(self, info: Info[Context, None]) -> List[SpanAnnotation]:
-        return await info.context.data_loaders.span_annotations.load(self.id_attr)
+    async def span_annotations(
+        self,
+        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)
+        return [to_gql_span_annotation(annotation) for annotation in annotations]
 
     @strawberry.field(
         description="Evaluations of the documents associated with the span, e.g. "
@@ -258,6 +277,11 @@ class Span(Node):
         project = await info.context.data_loaders.span_projects.load(span_id)
         return to_gql_project(project)
 
+    @strawberry.field(description="Indicates if the span is contained in any dataset")  # type: ignore
+    async def contained_in_dataset(self, info: Info[Context, None]) -> bool:
+        examples = await info.context.data_loaders.span_dataset_examples.load(self.id_attr)
+        return bool(examples)
+
 
 def to_gql_span(span: models.Span) -> Span:
     events: List[SpanEvent] = list(map(SpanEvent.from_dict, span.events))

phoenix/server/api/types/Trace.py

@@ -11,13 +11,15 @@ from strawberry.types import Info
 
 from phoenix.db import models
 from phoenix.server.api.context import Context
-from phoenix.server.api.types.Evaluation import TraceEvaluation
+from phoenix.server.api.input_types.TraceAnnotationSort import TraceAnnotationSort
 from phoenix.server.api.types.pagination import (
     ConnectionArgs,
     CursorString,
     connection_from_list,
 )
+from phoenix.server.api.types.SortDir import SortDir
 from phoenix.server.api.types.Span import Span, to_gql_span
+from phoenix.server.api.types.TraceAnnotation import TraceAnnotation, to_gql_trace_annotation
 
 
 @strawberry.type
@@ -62,6 +64,21 @@ class Trace(Node):
             data = [to_gql_span(span) async for span in spans]
         return connection_from_list(data=data, args=args)
 
-    @strawberry.field(description="Evaluations associated with the trace")  # type: ignore
-    async def trace_evaluations(self, info: Info[Context, None]) -> List[TraceEvaluation]:
-        return await info.context.data_loaders.trace_evaluations.load(self.id_attr)
+    @strawberry.field(description="Annotations associated with the trace.")  # type: ignore
+    async def span_annotations(
+        self,
+        info: Info[Context, None],
+        sort: Optional[TraceAnnotationSort] = None,
+    ) -> List[TraceAnnotation]:
+        async with info.context.db() as session:
+            stmt = select(models.TraceAnnotation).filter_by(span_rowid=self.id_attr)
+            if sort:
+                sort_col = getattr(models.TraceAnnotation, sort.col.value)
+                if sort.dir is SortDir.desc:
+                    stmt = stmt.order_by(sort_col.desc(), models.TraceAnnotation.id.desc())
+                else:
+                    stmt = stmt.order_by(sort_col.asc(), models.TraceAnnotation.id.asc())
+            else:
+                stmt = stmt.order_by(models.TraceAnnotation.created_at.desc())
+            annotations = await session.scalars(stmt)
+        return [to_gql_trace_annotation(annotation) for annotation in annotations]

phoenix/server/api/utils.py

@@ -1,13 +1,13 @@
-from typing import AsyncContextManager, Callable, List
+from typing import List
 
 from sqlalchemy import delete
-from sqlalchemy.ext.asyncio import AsyncSession
 
 from phoenix.db import models
+from phoenix.server.types import DbSessionFactory
 
 
 async def delete_projects(
-    db: Callable[[], AsyncContextManager[AsyncSession]],
+    db: DbSessionFactory,
     *project_names: str,
 ) -> List[int]:
     if not project_names:
@@ -22,7 +22,7 @@ async def delete_projects(
 
 
 async def delete_traces(
-    db: Callable[[], AsyncContextManager[AsyncSession]],
+    db: DbSessionFactory,
     *trace_ids: str,
 ) -> List[int]:
     if not trace_ids: