arize-phoenix 4.10.1__py3-none-any.whl → 4.10.2rc1__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
 from google.protobuf.message import DecodeError
 from opentelemetry.proto.collector.trace.v1.trace_service_pb2 import (
     ExportTraceServiceRequest,
 )
+from pydantic import BaseModel, 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,
@@ -26,40 +27,49 @@ from phoenix.server.api.types.node import from_global_id_with_expected_type
 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 .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,
+    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,96 +81,69 @@ 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]
+    BackgroundTasks().add_task(_add_spans, req, request.state)
+    return None
+
+
+class AnnotationResult(BaseModel):
+    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(BaseModel):
+    trace_id: str = Field(description="The ID of the trace being annotated")
+    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[AnnotationResult] = Field(
+        default=None, description="The result of the annotation"
+    )
+    metadata: Optional[Dict[str, Any]] = Field(
+        default=None, description="Metadata for the annotation"
+    )
+
+
+class AnnotateTracesRequestBody(RequestBody[List[TraceAnnotation]]):
+    data: List[TraceAnnotation] = Field(description="The trace annotations to be upserted")
+
+
+class InsertedTraceAnnotation(BaseModel):
+    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
+) -> AnnotateTracesResponseBody:
+    trace_annotations = request_body.data
+    trace_gids = [GlobalID.from_id(annotation.trace_id) for annotation in trace_annotations]
 
     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",
+            raise HTTPException(
+                detail="Trace with ID {trace_gid} does not exist",
                 status_code=HTTP_404_NOT_FOUND,
             )
 
@@ -175,24 +158,24 @@ async def annotate_traces(request: Request) -> Response:
             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_gids)} do not exist.",
                 status_code=HTTP_404_NOT_FOUND,
             )
 
         inserted_annotations = []
 
-        for annotation in payload:
-            trace_gid = GlobalID.from_id(annotation["trace_id"])
+        for annotation in trace_annotations:
+            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 {}
+            name = annotation.name
+            annotator_kind = annotation.annotator_kind
+            result = annotation.result
+            label = result.label if result else None
+            score = result.score if result else None
+            explanation = result.explanation if result else None
+            metadata = annotation.metadata or {}
 
             values = dict(
                 trace_rowid=trace_id,
@@ -213,10 +196,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,94 @@
+from typing import Any, Dict, Generic, List, Optional, TypedDict, Union
+
+from pydantic import BaseModel
+from typing_extensions import TypeAlias, TypeVar, assert_never
+
+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(BaseModel, 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(BaseModel, 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(BaseModel, 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/Annotation.py

@@ -0,0 +1,21 @@
+from typing import Optional
+
+import strawberry
+
+
+@strawberry.interface
+class Annotation:
+    name: str = strawberry.field(
+        description="Name of the annotation, e.g. 'helpfulness' or 'relevance'."
+    )
+    score: Optional[float] = strawberry.field(
+        description="Value of the annotation in the form of a numeric score."
+    )
+    label: Optional[str] = strawberry.field(
+        description="Value of the annotation in the form of a string, e.g. "
+        "'helpful' or 'not helpful'. Note that the label is not necessarily binary."
+    )
+    explanation: Optional[str] = strawberry.field(
+        description="The annotator's explanation for the annotation result (i.e. "
+        "score or label, or both) given to the subject."
+    )

phoenix/server/api/types/Evaluation.py

@@ -1,31 +1,13 @@
-from typing import Optional
-
 import strawberry
 
 import phoenix.trace.v1 as pb
 from phoenix.db.models import DocumentAnnotation, SpanAnnotation, TraceAnnotation
 
-
-@strawberry.interface
-class Evaluation:
-    name: str = strawberry.field(
-        description="Name of the evaluation, e.g. 'helpfulness' or 'relevance'."
-    )
-    score: Optional[float] = strawberry.field(
-        description="Result of the evaluation in the form of a numeric score."
-    )
-    label: Optional[str] = strawberry.field(
-        description="Result of the evaluation in the form of a string, e.g. "
-        "'helpful' or 'not helpful'. Note that the label is not necessarily binary."
-    )
-    explanation: Optional[str] = strawberry.field(
-        description="The evaluator's explanation for the evaluation result (i.e. "
-        "score or label, or both) given to the subject."
-    )
+from .Annotation import Annotation
 
 
 @strawberry.type
-class TraceEvaluation(Evaluation):
+class TraceEvaluation(Annotation):
     @staticmethod
     def from_pb_evaluation(evaluation: pb.Evaluation) -> "TraceEvaluation":
         result = evaluation.result
@@ -50,7 +32,7 @@ class TraceEvaluation(Evaluation):
 
 
 @strawberry.type
-class SpanEvaluation(Evaluation):
+class SpanEvaluation(Annotation):
     @staticmethod
     def from_pb_evaluation(evaluation: pb.Evaluation) -> "SpanEvaluation":
         result = evaluation.result
@@ -75,7 +57,7 @@ class SpanEvaluation(Evaluation):
 
 
 @strawberry.type
-class DocumentEvaluation(Evaluation):
+class DocumentEvaluation(Annotation):
     document_position: int = strawberry.field(
         description="The zero-based index among retrieved documents, which "
         "is collected as a list (even when ordering is not inherently meaningful)."

phoenix/server/api/types/Span.py

@@ -19,12 +19,14 @@ from phoenix.server.api.helpers.dataset_helpers import (
     get_dataset_example_input,
     get_dataset_example_output,
 )
-from phoenix.server.api.types.DocumentRetrievalMetrics import DocumentRetrievalMetrics
-from phoenix.server.api.types.Evaluation import DocumentEvaluation, SpanEvaluation
-from phoenix.server.api.types.ExampleRevisionInterface import ExampleRevision
-from phoenix.server.api.types.MimeType import MimeType
 from phoenix.trace.attributes import get_attribute_value
 
+from .DocumentRetrievalMetrics import DocumentRetrievalMetrics
+from .Evaluation import DocumentEvaluation, SpanEvaluation
+from .ExampleRevisionInterface import ExampleRevision
+from .MimeType import MimeType
+from .SpanAnnotation import SpanAnnotation
+
 if TYPE_CHECKING:
     from phoenix.server.api.types.Project import Project
 
@@ -172,6 +174,15 @@ class Span(Node):
     async def span_evaluations(self, info: Info[Context, None]) -> List[SpanEvaluation]:
         return await info.context.data_loaders.span_evaluations.load(self.id_attr)
 
+    @strawberry.field(
+        description=(
+            "Annotations of the span's parent 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)
+
     @strawberry.field(
         description="Evaluations of the documents associated with the span, e.g. "
         "if the span is a RETRIEVER with a list of documents in its RETRIEVAL_DOCUMENTS "

phoenix/server/api/types/SpanAnnotation.py

@@ -6,17 +6,15 @@ from strawberry.relay import GlobalID, Node, NodeID
 from strawberry.scalars import JSON
 
 from phoenix.db import models
-from phoenix.server.api.types.AnnotatorKind import AnnotatorKind
+
+from .Annotation import Annotation
+from .AnnotatorKind import AnnotatorKind
 
 
 @strawberry.type
-class SpanAnnotation(Node):
+class SpanAnnotation(Node, Annotation):
     id_attr: NodeID[int]
-    name: str
     annotator_kind: AnnotatorKind
-    label: Optional[str]
-    score: Optional[float]
-    explanation: Optional[str]
     metadata: JSON
     span_rowid: Private[Optional[int]]