arize-phoenix 4.10.2rc2__py3-none-any.whl → 4.12.0__py3-none-any.whl

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

@@ -1,19 +1,18 @@
 import gzip
 import zlib
-from typing import Any, Dict, List, Literal, Optional
+from typing import Any, Dict, List
 
-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,
@@ -27,50 +26,40 @@ 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
 
-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:
+
+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")
     if content_type != "application/x-protobuf":
-        raise HTTPException(
-            detail=f"Unsupported content type: {content_type}",
+        return Response(
+            content=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"):
-        raise HTTPException(
-            detail=f"Unsupported content encoding: {content_encoding}",
+        return Response(
+            content=f"Unsupported content encoding: {content_encoding}",
             status_code=HTTP_415_UNSUPPORTED_MEDIA_TYPE,
         )
     body = await request.body()
@@ -82,69 +71,96 @@ async def post_traces(
     try:
         await run_in_threadpool(req.ParseFromString, body)
     except DecodeError:
-        raise HTTPException(
-            detail="Request body is invalid ExportTraceServiceRequest",
+        return Response(
+            content="Request body is invalid ExportTraceServiceRequest",
             status_code=HTTP_422_UNPROCESSABLE_ENTITY,
         )
-    background_tasks.add_task(_add_spans, req, request.state)
-    return None
-
-
-class TraceAnnotationResult(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[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(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]
+    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:
-            raise HTTPException(
-                detail="Trace with ID {trace_gid} does not exist",
+            return Response(
+                content="Trace with ID {trace_gid} does not exist",
                 status_code=HTTP_404_NOT_FOUND,
             )
 
@@ -159,24 +175,24 @@ async def annotate_traces(
             missing_trace_gids = [
                 str(GlobalID("Trace", str(trace_gid))) for trace_gid in missing_trace_ids
             ]
-            raise HTTPException(
-                detail=f"Traces with IDs {', '.join(missing_trace_gids)} do not exist.",
+            return Response(
+                content=f"Traces with IDs {', '.join(missing_trace_gids)} do not exist.",
                 status_code=HTTP_404_NOT_FOUND,
             )
 
         inserted_annotations = []
 
-        for annotation in trace_annotations:
-            trace_gid = GlobalID.from_id(annotation.trace_id)
+        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.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 {}
+            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,
@@ -197,12 +213,10 @@ async def annotate_traces(
                 ).returning(models.TraceAnnotation.id)
             )
             inserted_annotations.append(
-                InsertedTraceAnnotation(
-                    id=str(GlobalID("TraceAnnotation", str(trace_annotation_id)))
-                )
+                {"id": str(GlobalID("TraceAnnotation", str(trace_annotation_id)))}
             )
 
-    return AnnotateTracesResponseBody(data=inserted_annotations)
+    return JSONResponse(content={"data": inserted_annotations})
 
 
 async def _add_spans(req: ExportTraceServiceRequest, state: State) -> None:

phoenix/server/api/types/Span.py

@@ -62,6 +62,7 @@ class SpanKind(Enum):
     agent = "AGENT"
     reranker = "RERANKER"
     evaluator = "EVALUATOR"
+    guardrail = "GUARDRAIL"
     unknown = "UNKNOWN"
 
     @classmethod