arize-phoenix 4.12.0rc1__py3-none-any.whl → 4.12.1rc1__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 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,51 @@ 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 .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,96 +83,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]
+    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="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[TraceAnnotationResult] = 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(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
+) -> 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 +160,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 +198,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