arize-phoenix 3.16.1__py3-none-any.whl → 7.7.1__py3-none-any.whl

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

@@ -0,0 +1,208 @@
+import gzip
+import zlib
+from typing import Any, Literal, Optional
+
+from fastapi import APIRouter, BackgroundTasks, Header, HTTPException, Query
+from google.protobuf.json_format import MessageToJson
+from google.protobuf.message import DecodeError
+from opentelemetry.proto.collector.trace.v1.trace_service_pb2 import (
+    ExportTraceServiceRequest,
+    ExportTraceServiceResponse,
+)
+from pydantic import Field
+from sqlalchemy import select
+from starlette.concurrency import run_in_threadpool
+from starlette.datastructures import State
+from starlette.requests import Request
+from starlette.responses import JSONResponse
+from starlette.status import (
+    HTTP_404_NOT_FOUND,
+    HTTP_415_UNSUPPORTED_MEDIA_TYPE,
+    HTTP_422_UNPROCESSABLE_ENTITY,
+)
+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.db.insertion.types import Precursors
+from phoenix.server.dml_event import TraceAnnotationInsertEvent
+from phoenix.trace.otel import decode_otlp_span
+from phoenix.utilities.project import get_project_name
+
+from .pydantic_compat import V1RoutesBaseModel
+from .utils import RequestBody, ResponseBody, add_errors_to_responses
+
+router = APIRouter(tags=["traces"])
+
+
+@router.post(
+    "/traces",
+    operation_id="addTraces",
+    summary="Send traces",
+    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"}}
+            },
+        }
+    },
+    include_in_schema=False,
+)
+async def post_traces(
+    request: Request,
+    background_tasks: BackgroundTasks,
+    content_type: Optional[str] = Header(default=None),
+    content_encoding: Optional[str] = Header(default=None),
+) -> JSONResponse:
+    if content_type != "application/x-protobuf":
+        raise HTTPException(
+            detail=f"Unsupported content type: {content_type}",
+            status_code=HTTP_415_UNSUPPORTED_MEDIA_TYPE,
+        )
+    if content_encoding and content_encoding not in ("gzip", "deflate"):
+        raise HTTPException(
+            detail=f"Unsupported content encoding: {content_encoding}",
+            status_code=HTTP_415_UNSUPPORTED_MEDIA_TYPE,
+        )
+    body = await request.body()
+    if content_encoding == "gzip":
+        body = await run_in_threadpool(gzip.decompress, body)
+    elif content_encoding == "deflate":
+        body = await run_in_threadpool(zlib.decompress, body)
+    req = ExportTraceServiceRequest()
+    try:
+        await run_in_threadpool(req.ParseFromString, body)
+    except DecodeError:
+        raise HTTPException(
+            detail="Request body is invalid ExportTraceServiceRequest",
+            status_code=HTTP_422_UNPROCESSABLE_ENTITY,
+        )
+    background_tasks.add_task(_add_spans, req, request.state)
+    return JSONResponse(MessageToJson(ExportTraceServiceResponse()))
+
+
+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"}]
+    ),
+    include_in_schema=False,
+)
+async def annotate_traces(
+    request: Request,
+    request_body: AnnotateTracesRequestBody,
+    sync: bool = Query(default=True, description="If true, fulfill request synchronously."),
+) -> AnnotateTracesResponseBody:
+    if not request_body.data:
+        return AnnotateTracesResponseBody(data=[])
+    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:
+        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 = trace_ids - set(existing_traces.keys())
+        if missing_trace_ids:
+            raise HTTPException(
+                detail=f"Traces with IDs {', '.join(missing_trace_ids)} do not exist.",
+                status_code=HTTP_404_NOT_FOUND,
+            )
+        inserted_ids = []
+        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,
+                    dialect=dialect,
+                    table=models.TraceAnnotation,
+                    unique_by=("name", "trace_rowid"),
+                ).returning(models.TraceAnnotation.id)
+            )
+            inserted_ids.append(trace_annotation_id)
+    request.state.event_queue.put(TraceAnnotationInsertEvent(tuple(inserted_ids)))
+    return AnnotateTracesResponseBody(
+        data=[
+            InsertedTraceAnnotation(id=str(GlobalID("TraceAnnotation", str(id_))))
+            for id_ in inserted_ids
+        ]
+    )
+
+
+async def _add_spans(req: ExportTraceServiceRequest, state: State) -> None:
+    for resource_spans in req.resource_spans:
+        project_name = get_project_name(resource_spans.resource.attributes)
+        for scope_span in resource_spans.scope_spans:
+            for otlp_span in scope_span.spans:
+                span = await run_in_threadpool(decode_otlp_span, otlp_span)
+                await state.queue_span_for_bulk_insert(span, project_name)

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

@@ -0,0 +1,95 @@
+from typing import Any, Generic, Optional, TypedDict, TypeVar, Union
+
+from typing_extensions import TypeAlias, 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/schema.py

@@ -1,246 +1,49 @@
-from collections import defaultdict
-from typing import Dict, List, Optional, Set, Union
+from itertools import chain
+from typing import Any, Iterable, Iterator, Optional, Union
 
-import numpy as np
-import numpy.typing as npt
 import strawberry
-from strawberry import ID, UNSET
-from strawberry.types import Info
-from typing_extensions import Annotated
-
-from phoenix.pointcloud.clustering import Hdbscan
-from phoenix.server.api.helpers import ensure_list
-from phoenix.server.api.input_types.ClusterInput import ClusterInput
-from phoenix.server.api.input_types.Coordinates import (
-    InputCoordinate2D,
-    InputCoordinate3D,
-)
-from phoenix.server.api.types.Cluster import Cluster, to_gql_clusters
-from phoenix.server.api.types.Project import Project
-
-from .context import Context
-from .types.DatasetRole import AncillaryDatasetRole, DatasetRole
-from .types.Dimension import to_gql_dimension
-from .types.EmbeddingDimension import (
-    DEFAULT_CLUSTER_SELECTION_EPSILON,
-    DEFAULT_MIN_CLUSTER_SIZE,
-    DEFAULT_MIN_SAMPLES,
-    to_gql_embedding_dimension,
+from strawberry.extensions import SchemaExtension
+from strawberry.types.base import StrawberryObjectDefinition, StrawberryType
+
+from phoenix.server.api.exceptions import get_mask_errors_extension
+from phoenix.server.api.mutations import Mutation
+from phoenix.server.api.queries import Query
+from phoenix.server.api.subscriptions import Subscription
+from phoenix.server.api.types.ChatCompletionSubscriptionPayload import (
+    ChatCompletionSubscriptionPayload,
 )
-from .types.Event import create_event_id, unpack_event_id
-from .types.ExportEventsMutation import ExportEventsMutation
-from .types.Functionality import Functionality
-from .types.Model import Model
-from .types.node import GlobalID, Node, from_global_id, from_global_id_with_expected_type
-from .types.pagination import Connection, ConnectionArgs, Cursor, connection_from_list
-
-
-@strawberry.type
-class Query:
-    @strawberry.field
-    def projects(
-        self,
-        info: Info[Context, None],
-        first: Optional[int] = 50,
-        last: Optional[int] = UNSET,
-        after: Optional[Cursor] = UNSET,
-        before: Optional[Cursor] = UNSET,
-    ) -> Connection[Project]:
-        args = ConnectionArgs(
-            first=first,
-            after=after if isinstance(after, Cursor) else None,
-            last=last,
-            before=before if isinstance(before, Cursor) else None,
-        )
-        data = (
-            []
-            if (traces := info.context.traces) is None
-            else [
-                Project(id_attr=project_id, name=project_name, project=project)
-                for project_id, project_name, project in traces.get_projects()
-            ]
-        )
-        return connection_from_list(data=data, args=args)
-
-    @strawberry.field
-    def functionality(self, info: Info[Context, None]) -> "Functionality":
-        has_model_inferences = not info.context.model.is_empty
-        has_traces = info.context.traces is not None
-        return Functionality(
-            model_inferences=has_model_inferences,
-            tracing=has_traces,
-        )
-
-    @strawberry.field
-    def model(self) -> Model:
-        return Model()
-
-    @strawberry.field
-    def node(self, id: GlobalID, info: Info[Context, None]) -> Node:
-        type_name, node_id = from_global_id(str(id))
-        if type_name == "Dimension":
-            dimension = info.context.model.scalar_dimensions[node_id]
-            return to_gql_dimension(node_id, dimension)
-        elif type_name == "EmbeddingDimension":
-            embedding_dimension = info.context.model.embedding_dimensions[node_id]
-            return to_gql_embedding_dimension(node_id, embedding_dimension)
-        elif type_name == "Project":
-            if (traces := info.context.traces) is not None:
-                projects = {
-                    project_id: (project_name, project)
-                    for project_id, project_name, project in traces.get_projects()
-                }
-                if node_id in projects:
-                    name, project = projects[node_id]
-                    return Project(id_attr=node_id, name=name, project=project)
-            raise Exception(f"Unknown project: {id}")
-
-        raise Exception(f"Unknown node type: {type}")
-
-    @strawberry.field
-    def clusters(
-        self,
-        clusters: List[ClusterInput],
-    ) -> List[Cluster]:
-        clustered_events: Dict[str, Set[ID]] = defaultdict(set)
-        for i, cluster in enumerate(clusters):
-            clustered_events[cluster.id or str(i)].update(cluster.event_ids)
-        return to_gql_clusters(
-            clustered_events=clustered_events,
-        )
-
-    @strawberry.field
-    def hdbscan_clustering(
-        self,
-        info: Info[Context, None],
-        event_ids: Annotated[
-            List[ID],
-            strawberry.argument(
-                description="Event ID of the coordinates",
-            ),
-        ],
-        coordinates_2d: Annotated[
-            Optional[List[InputCoordinate2D]],
-            strawberry.argument(
-                description="Point coordinates. Must be either 2D or 3D.",
-            ),
-        ] = UNSET,
-        coordinates_3d: Annotated[
-            Optional[List[InputCoordinate3D]],
-            strawberry.argument(
-                description="Point coordinates. Must be either 2D or 3D.",
-            ),
-        ] = UNSET,
-        min_cluster_size: Annotated[
-            int,
-            strawberry.argument(
-                description="HDBSCAN minimum cluster size",
-            ),
-        ] = DEFAULT_MIN_CLUSTER_SIZE,
-        cluster_min_samples: Annotated[
-            int,
-            strawberry.argument(
-                description="HDBSCAN minimum samples",
-            ),
-        ] = DEFAULT_MIN_SAMPLES,
-        cluster_selection_epsilon: Annotated[
-            float,
-            strawberry.argument(
-                description="HDBSCAN cluster selection epsilon",
-            ),
-        ] = DEFAULT_CLUSTER_SELECTION_EPSILON,
-    ) -> List[Cluster]:
-        coordinates_3d = ensure_list(coordinates_3d)
-        coordinates_2d = ensure_list(coordinates_2d)
-
-        if len(coordinates_3d) > 0 and len(coordinates_2d) > 0:
-            raise ValueError("must specify only one of 2D or 3D coordinates")
-
-        if len(coordinates_3d) > 0:
-            coordinates = list(
-                map(
-                    lambda coord: np.array(
-                        [coord.x, coord.y, coord.z],
-                    ),
-                    coordinates_3d,
-                )
-            )
-        else:
-            coordinates = list(
-                map(
-                    lambda coord: np.array(
-                        [coord.x, coord.y],
-                    ),
-                    coordinates_2d,
-                )
-            )
-
-        if len(event_ids) != len(coordinates):
-            raise ValueError(
-                f"length mismatch between "
-                f"event_ids ({len(event_ids)}) "
-                f"and coordinates ({len(coordinates)})"
-            )
-
-        if len(event_ids) == 0:
-            return []
-
-        grouped_event_ids: Dict[
-            Union[DatasetRole, AncillaryDatasetRole],
-            List[ID],
-        ] = defaultdict(list)
-        grouped_coordinates: Dict[
-            Union[DatasetRole, AncillaryDatasetRole],
-            List[npt.NDArray[np.float64]],
-        ] = defaultdict(list)
-
-        for event_id, coordinate in zip(event_ids, coordinates):
-            row_id, dataset_role = unpack_event_id(event_id)
-            grouped_coordinates[dataset_role].append(coordinate)
-            grouped_event_ids[dataset_role].append(create_event_id(row_id, dataset_role))
-
-        stacked_event_ids = (
-            grouped_event_ids[DatasetRole.primary]
-            + grouped_event_ids[DatasetRole.reference]
-            + grouped_event_ids[AncillaryDatasetRole.corpus]
-        )
-        stacked_coordinates = np.stack(
-            grouped_coordinates[DatasetRole.primary]
-            + grouped_coordinates[DatasetRole.reference]
-            + grouped_coordinates[AncillaryDatasetRole.corpus]
-        )
-
-        clusters = Hdbscan(
-            min_cluster_size=min_cluster_size,
-            min_samples=cluster_min_samples,
-            cluster_selection_epsilon=cluster_selection_epsilon,
-        ).find_clusters(stacked_coordinates)
-
-        clustered_events = {
-            str(i): {stacked_event_ids[row_idx] for row_idx in cluster}
-            for i, cluster in enumerate(clusters)
-        }
-
-        return to_gql_clusters(
-            clustered_events=clustered_events,
-        )
-
-
-@strawberry.type
-class Mutation(ExportEventsMutation):
-    @strawberry.mutation
-    def delete_project(self, info: Info[Context, None], id: GlobalID) -> Query:
-        if (traces := info.context.traces) is not None:
-            node_id = from_global_id_with_expected_type(str(id), "Project")
-            traces.archive_project(node_id)
-        return Query()
-
-    @strawberry.mutation
-    def archive_project(self, info: Info[Context, None], id: GlobalID) -> Query:
-        if (traces := info.context.traces) is not None:
-            node_id = from_global_id_with_expected_type(str(id), "Project")
-            traces.archive_project(node_id)
-        return Query()
 
 
-schema = strawberry.Schema(query=Query, mutation=Mutation)
+def build_graphql_schema(
+    extensions: Optional[Iterable[Union[type[SchemaExtension], SchemaExtension]]] = None,
+) -> strawberry.Schema:
+    """
+    Builds a strawberry schema.
+    """
+    return strawberry.Schema(
+        query=Query,
+        mutation=Mutation,
+        extensions=list(chain(extensions or [], [get_mask_errors_extension()])),
+        subscription=Subscription,
+        types=_implementing_types(ChatCompletionSubscriptionPayload),
+    )
+
+
+def _implementing_types(interface: Any) -> Iterator[StrawberryType]:
+    """
+    Iterates over strawberry types implementing the given strawberry interface.
+    """
+    assert isinstance(
+        strawberry_definition := getattr(interface, "__strawberry_definition__", None),
+        StrawberryObjectDefinition,
+    )
+    assert strawberry_definition.is_interface
+    for subcls in interface.__subclasses__():
+        if isinstance(
+            getattr(subcls, "__strawberry_definition__", None),
+            StrawberryObjectDefinition,
+        ):
+            yield subcls
+
+
+_EXPORTED_GRAPHQL_SCHEMA = build_graphql_schema()  # used to export the GraphQL schema to file