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

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

@@ -0,0 +1,308 @@
+from datetime import datetime
+from random import getrandbits
+from typing import Any, Optional
+
+from fastapi import APIRouter, HTTPException, Path
+from pydantic import Field
+from sqlalchemy import select
+from starlette.requests import Request
+from starlette.status import HTTP_404_NOT_FOUND
+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.server.dml_event import ExperimentInsertEvent
+
+from .pydantic_compat import V1RoutesBaseModel
+from .utils import ResponseBody, add_errors_to_responses
+
+router = APIRouter(tags=["experiments"], include_in_schema=True)
+
+
+def _short_uuid() -> str:
+    return str(getrandbits(32).to_bytes(4, "big").hex())
+
+
+def _generate_experiment_name(dataset_name: str) -> str:
+    """
+    Generate a semi-unique name for the experiment.
+    """
+    short_ds_name = dataset_name[:8].replace(" ", "-")
+    return f"{short_ds_name}-{_short_uuid()}"
+
+
+class Experiment(V1RoutesBaseModel):
+    id: str = Field(description="The ID of the experiment")
+    dataset_id: str = Field(description="The ID of the dataset associated with the experiment")
+    dataset_version_id: str = Field(
+        description="The ID of the dataset version associated with the experiment"
+    )
+    repetitions: int = Field(description="Number of times the experiment is repeated")
+    metadata: dict[str, Any] = Field(description="Metadata of the experiment")
+    project_name: Optional[str] = Field(
+        description="The name of the project associated with the experiment"
+    )
+    created_at: datetime = Field(description="The creation timestamp of the experiment")
+    updated_at: datetime = Field(description="The last update timestamp of the experiment")
+
+
+class CreateExperimentRequestBody(V1RoutesBaseModel):
+    """
+    Details of the experiment to be created
+    """
+
+    name: Optional[str] = Field(
+        default=None,
+        description=("Name of the experiment (if omitted, a random name will be generated)"),
+    )
+    description: Optional[str] = Field(
+        default=None, description="An optional description of the experiment"
+    )
+    metadata: Optional[dict[str, Any]] = Field(
+        default=None, description="Metadata for the experiment"
+    )
+    version_id: Optional[str] = Field(
+        default=None,
+        description=(
+            "ID of the dataset version over which the experiment will be run "
+            "(if omitted, the latest version will be used)"
+        ),
+    )
+    repetitions: int = Field(
+        default=1, description="Number of times the experiment should be repeated for each example"
+    )
+
+
+class CreateExperimentResponseBody(ResponseBody[Experiment]):
+    pass
+
+
+@router.post(
+    "/datasets/{dataset_id}/experiments",
+    operation_id="createExperiment",
+    summary="Create experiment on a dataset",
+    responses=add_errors_to_responses(
+        [{"status_code": HTTP_404_NOT_FOUND, "description": "Dataset or DatasetVersion not found"}]
+    ),
+    response_description="Experiment retrieved successfully",
+)
+async def create_experiment(
+    request: Request,
+    request_body: CreateExperimentRequestBody,
+    dataset_id: str = Path(..., title="Dataset ID"),
+) -> CreateExperimentResponseBody:
+    dataset_globalid = GlobalID.from_id(dataset_id)
+    try:
+        dataset_rowid = from_global_id_with_expected_type(dataset_globalid, "Dataset")
+    except ValueError:
+        raise HTTPException(
+            detail="Dataset with ID {dataset_globalid} does not exist",
+            status_code=HTTP_404_NOT_FOUND,
+        )
+
+    dataset_version_globalid_str = request_body.version_id
+    if dataset_version_globalid_str is not None:
+        try:
+            dataset_version_globalid = GlobalID.from_id(dataset_version_globalid_str)
+            dataset_version_id = from_global_id_with_expected_type(
+                dataset_version_globalid, "DatasetVersion"
+            )
+        except ValueError:
+            raise HTTPException(
+                detail=f"DatasetVersion with ID {dataset_version_globalid_str} does not exist",
+                status_code=HTTP_404_NOT_FOUND,
+            )
+
+    async with request.app.state.db() as session:
+        result = (
+            await session.execute(select(models.Dataset).where(models.Dataset.id == dataset_rowid))
+        ).scalar()
+        if result is None:
+            raise HTTPException(
+                detail=f"Dataset with ID {dataset_globalid} does not exist",
+                status_code=HTTP_404_NOT_FOUND,
+            )
+        dataset_name = result.name
+        if dataset_version_globalid_str is None:
+            dataset_version_result = await session.execute(
+                select(models.DatasetVersion)
+                .where(models.DatasetVersion.dataset_id == dataset_rowid)
+                .order_by(models.DatasetVersion.id.desc())
+            )
+            dataset_version = dataset_version_result.scalar()
+            if not dataset_version:
+                raise HTTPException(
+                    detail=f"Dataset {dataset_globalid} does not have any versions",
+                    status_code=HTTP_404_NOT_FOUND,
+                )
+            dataset_version_id = dataset_version.id
+            dataset_version_globalid = GlobalID("DatasetVersion", str(dataset_version_id))
+        else:
+            dataset_version = await session.execute(
+                select(models.DatasetVersion).where(models.DatasetVersion.id == dataset_version_id)
+            )
+            dataset_version = dataset_version.scalar()
+            if not dataset_version:
+                raise HTTPException(
+                    detail=f"DatasetVersion with ID {dataset_version_globalid} does not exist",
+                    status_code=HTTP_404_NOT_FOUND,
+                )
+
+        # generate a semi-unique name for the experiment
+        experiment_name = request_body.name or _generate_experiment_name(dataset_name)
+        project_name = f"Experiment-{getrandbits(96).to_bytes(12, 'big').hex()}"
+        project_description = (
+            f"dataset_id: {dataset_globalid}\ndataset_version_id: {dataset_version_globalid}"
+        )
+        experiment = models.Experiment(
+            dataset_id=int(dataset_rowid),
+            dataset_version_id=int(dataset_version_id),
+            name=experiment_name,
+            description=request_body.description,
+            repetitions=request_body.repetitions,
+            metadata_=request_body.metadata or {},
+            project_name=project_name,
+        )
+        session.add(experiment)
+        await session.flush()
+
+        dialect = SupportedSQLDialect(session.bind.dialect.name)
+        project_rowid = await session.scalar(
+            insert_on_conflict(
+                dict(
+                    name=project_name,
+                    description=project_description,
+                    created_at=experiment.created_at,
+                    updated_at=experiment.updated_at,
+                ),
+                dialect=dialect,
+                table=models.Project,
+                unique_by=("name",),
+            ).returning(models.Project.id)
+        )
+        assert project_rowid is not None
+
+        experiment_globalid = GlobalID("Experiment", str(experiment.id))
+        if dataset_version_globalid_str is None:
+            dataset_version_globalid = GlobalID(
+                "DatasetVersion", str(experiment.dataset_version_id)
+            )
+    request.state.event_queue.put(ExperimentInsertEvent((experiment.id,)))
+    return CreateExperimentResponseBody(
+        data=Experiment(
+            id=str(experiment_globalid),
+            dataset_id=str(dataset_globalid),
+            dataset_version_id=str(dataset_version_globalid),
+            repetitions=experiment.repetitions,
+            metadata=experiment.metadata_,
+            project_name=experiment.project_name,
+            created_at=experiment.created_at,
+            updated_at=experiment.updated_at,
+        )
+    )
+
+
+class GetExperimentResponseBody(ResponseBody[Experiment]):
+    pass
+
+
+@router.get(
+    "/experiments/{experiment_id}",
+    operation_id="getExperiment",
+    summary="Get experiment by ID",
+    responses=add_errors_to_responses(
+        [{"status_code": HTTP_404_NOT_FOUND, "description": "Experiment not found"}]
+    ),
+    response_description="Experiment retrieved successfully",
+)
+async def get_experiment(request: Request, experiment_id: str) -> GetExperimentResponseBody:
+    experiment_globalid = GlobalID.from_id(experiment_id)
+    try:
+        experiment_rowid = from_global_id_with_expected_type(experiment_globalid, "Experiment")
+    except ValueError:
+        raise HTTPException(
+            detail="Experiment with ID {experiment_globalid} does not exist",
+            status_code=HTTP_404_NOT_FOUND,
+        )
+
+    async with request.app.state.db() as session:
+        experiment = await session.execute(
+            select(models.Experiment).where(models.Experiment.id == experiment_rowid)
+        )
+        experiment = experiment.scalar()
+        if not experiment:
+            raise HTTPException(
+                detail=f"Experiment with ID {experiment_globalid} does not exist",
+                status_code=HTTP_404_NOT_FOUND,
+            )
+
+        dataset_globalid = GlobalID("Dataset", str(experiment.dataset_id))
+        dataset_version_globalid = GlobalID("DatasetVersion", str(experiment.dataset_version_id))
+    return GetExperimentResponseBody(
+        data=Experiment(
+            id=str(experiment_globalid),
+            dataset_id=str(dataset_globalid),
+            dataset_version_id=str(dataset_version_globalid),
+            repetitions=experiment.repetitions,
+            metadata=experiment.metadata_,
+            project_name=experiment.project_name,
+            created_at=experiment.created_at,
+            updated_at=experiment.updated_at,
+        )
+    )
+
+
+class ListExperimentsResponseBody(ResponseBody[list[Experiment]]):
+    pass
+
+
+@router.get(
+    "/datasets/{dataset_id}/experiments",
+    operation_id="listExperiments",
+    summary="List experiments by dataset",
+    response_description="Experiments retrieved successfully",
+)
+async def list_experiments(
+    request: Request,
+    dataset_id: str = Path(..., title="Dataset ID"),
+) -> ListExperimentsResponseBody:
+    dataset_gid = GlobalID.from_id(dataset_id)
+    try:
+        dataset_rowid = from_global_id_with_expected_type(dataset_gid, "Dataset")
+    except ValueError:
+        raise HTTPException(
+            detail=f"Dataset with ID {dataset_gid} does not exist",
+            status_code=HTTP_404_NOT_FOUND,
+        )
+    async with request.app.state.db() as session:
+        query = (
+            select(models.Experiment)
+            .where(models.Experiment.dataset_id == dataset_rowid)
+            .order_by(models.Experiment.id.desc())
+        )
+
+        result = await session.execute(query)
+        experiments = result.scalars().all()
+
+        if not experiments:
+            return ListExperimentsResponseBody(data=[])
+
+        data = [
+            Experiment(
+                id=str(GlobalID("Experiment", str(experiment.id))),
+                dataset_id=str(GlobalID("Dataset", str(experiment.dataset_id))),
+                dataset_version_id=str(
+                    GlobalID("DatasetVersion", str(experiment.dataset_version_id))
+                ),
+                repetitions=experiment.repetitions,
+                metadata=experiment.metadata_,
+                project_name=None,
+                created_at=experiment.created_at,
+                updated_at=experiment.updated_at,
+            )
+            for experiment in experiments
+        ]
+
+        return ListExperimentsResponseBody(data=data)

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

@@ -0,0 +1,78 @@
+from datetime import datetime
+from enum import Enum
+from importlib.metadata import PackageNotFoundError, version
+
+from pydantic import BaseModel
+from typing_extensions import assert_never
+
+
+def datetime_encoder(dt: datetime) -> str:
+    """
+    Encodes a `datetime` object to an ISO-formatted timestamp string.
+
+    By default, Pydantic v2 serializes `datetime` objects in a format that
+    cannot be parsed by `datetime.fromisoformat`. Adding this encoder to the
+    `json_encoders` config for a Pydantic model ensures that the serialized
+    `datetime` objects are parseable.
+    """
+    return dt.isoformat()
+
+
+class PydanticMajorVersion(Enum):
+    """
+    The major version of `pydantic`.
+    """
+
+    V1 = "v1"
+    V2 = "v2"
+
+
+def get_pydantic_major_version() -> PydanticMajorVersion:
+    """
+    Returns the major version of `pydantic` or raises an error if `pydantic` is
+    not installed.
+    """
+    try:
+        pydantic_version = version("pydantic")
+    except PackageNotFoundError:
+        raise RuntimeError("Please install pydantic with `pip install pydantic`.")
+    if pydantic_version.startswith("1"):
+        return PydanticMajorVersion.V1
+    elif pydantic_version.startswith("2"):
+        return PydanticMajorVersion.V2
+    raise ValueError(f"Unsupported Pydantic version: {pydantic_version}")
+
+
+if (pydantic_major_version := get_pydantic_major_version()) is PydanticMajorVersion.V1:
+
+    class V1RoutesBaseModel(BaseModel):
+        class Config:
+            json_encoders = {datetime: datetime_encoder}
+
+elif pydantic_major_version is PydanticMajorVersion.V2:
+    from pydantic import ConfigDict
+
+    # `json_encoders` is a configuration setting from Pydantic v1 that was
+    # removed in Pydantic v2.0.* but restored in Pydantic v2.1.0 with a
+    # deprecation warning. At this time, it remains the simplest way to
+    # configure custom JSON serialization for specific data types in a manner
+    # that is consistent between Pydantic v1 and v2.
+    #
+    # For details, see:
+    # - https://github.com/pydantic/pydantic/pull/6811
+    # - https://github.com/pydantic/pydantic/releases/tag/v2.1.0
+    #
+    # The assertion below is added in case a future release of Pydantic v2 fully
+    # removes the `json_encoders` parameter.
+    assert "json_encoders" in ConfigDict.__annotations__, (
+        "If you encounter this error with `pydantic==2.0.*`, "
+        "please upgrade `pydantic` with `pip install -U pydantic>=2.1.0`. "
+        "If you encounter this error with `pydantic>=2.1.0`, "
+        "please upgrade `arize-phoenix` with `pip install -U arize-phoenix`, "
+        "or downgrade `pydantic` to a version that supports the `json_encoders` config setting."
+    )
+
+    class V1RoutesBaseModel(BaseModel):  # type: ignore[no-redef]
+        model_config = ConfigDict({"json_encoders": {datetime: datetime_encoder}})
+else:
+    assert_never(pydantic_major_version)

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

@@ -0,0 +1,267 @@
+from asyncio import get_running_loop
+from collections.abc import AsyncIterator
+from datetime import datetime, timezone
+from secrets import token_urlsafe
+from typing import Any, Literal, Optional
+
+import pandas as pd
+from fastapi import APIRouter, Header, HTTPException, Query
+from pydantic import Field
+from sqlalchemy import select
+from starlette.requests import Request
+from starlette.responses import Response, StreamingResponse
+from starlette.status import HTTP_404_NOT_FOUND, HTTP_422_UNPROCESSABLE_ENTITY
+from strawberry.relay import GlobalID
+
+from phoenix.config import DEFAULT_PROJECT_NAME
+from phoenix.datetime_utils import normalize_datetime
+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.api.routers.utils import df_to_bytes
+from phoenix.server.dml_event import SpanAnnotationInsertEvent
+from phoenix.trace.dsl import SpanQuery as SpanQuery_
+from phoenix.utilities.json import encode_df_as_json_string
+
+from .pydantic_compat import V1RoutesBaseModel
+from .utils import RequestBody, ResponseBody, add_errors_to_responses
+
+DEFAULT_SPAN_LIMIT = 1000
+
+router = APIRouter(tags=["spans"])
+
+
+class SpanQuery(V1RoutesBaseModel):
+    select: Optional[dict[str, Any]] = None
+    filter: Optional[dict[str, Any]] = None
+    explode: Optional[dict[str, Any]] = None
+    concat: Optional[dict[str, Any]] = None
+    rename: Optional[dict[str, Any]] = None
+    index: Optional[dict[str, Any]] = None
+
+
+class QuerySpansRequestBody(V1RoutesBaseModel):
+    queries: list[SpanQuery]
+    start_time: Optional[datetime] = None
+    end_time: Optional[datetime] = None
+    limit: int = DEFAULT_SPAN_LIMIT
+    root_spans_only: Optional[bool] = None
+    project_name: Optional[str] = Field(
+        default=None,
+        description=(
+            "The name of the project to query. "
+            "This parameter has been deprecated, use the project_name query parameter instead."
+        ),
+        deprecated=True,
+    )
+    stop_time: Optional[datetime] = Field(
+        default=None,
+        description=(
+            "An upper bound on the time to query for. "
+            "This parameter has been deprecated, use the end_time parameter instead."
+        ),
+        deprecated=True,
+    )
+
+
+# TODO: Add property details to SpanQuery schema
+@router.post(
+    "/spans",
+    operation_id="querySpans",
+    summary="Query spans with query DSL",
+    responses=add_errors_to_responses([HTTP_404_NOT_FOUND, HTTP_422_UNPROCESSABLE_ENTITY]),
+    include_in_schema=False,
+)
+async def query_spans_handler(
+    request: Request,
+    request_body: QuerySpansRequestBody,
+    accept: Optional[str] = Header(None),
+    project_name: Optional[str] = Query(
+        default=None, description="The project name to get evaluations from"
+    ),
+) -> Response:
+    queries = request_body.queries
+    project_name = (
+        project_name
+        or request.query_params.get("project-name")  # for backward compatibility
+        or request.headers.get(
+            "project-name"
+        )  # read from headers/payload for backward-compatibility
+        or request_body.project_name
+        or DEFAULT_PROJECT_NAME
+    )
+    end_time = request_body.end_time or request_body.stop_time
+    try:
+        span_queries = [SpanQuery_.from_dict(query.dict()) for query in queries]
+    except Exception as e:
+        raise HTTPException(
+            detail=f"Invalid query: {e}",
+            status_code=HTTP_422_UNPROCESSABLE_ENTITY,
+        )
+    async with request.app.state.db() as session:
+        results = []
+        for query in span_queries:
+            results.append(
+                await session.run_sync(
+                    query,
+                    project_name=project_name,
+                    start_time=normalize_datetime(
+                        request_body.start_time,
+                        timezone.utc,
+                    ),
+                    end_time=normalize_datetime(
+                        end_time,
+                        timezone.utc,
+                    ),
+                    limit=request_body.limit,
+                    root_spans_only=request_body.root_spans_only,
+                )
+            )
+    if not results:
+        raise HTTPException(status_code=HTTP_404_NOT_FOUND)
+
+    if accept == "application/json":
+        boundary_token = token_urlsafe(64)
+        return StreamingResponse(
+            content=_json_multipart(results, boundary_token),
+            media_type=f"multipart/mixed; boundary={boundary_token}",
+        )
+
+    async def content() -> AsyncIterator[bytes]:
+        for result in results:
+            yield df_to_bytes(result)
+
+    return StreamingResponse(
+        content=content(),
+        media_type="application/x-pandas-arrow",
+    )
+
+
+async def _json_multipart(
+    results: list[pd.DataFrame],
+    boundary_token: str,
+) -> AsyncIterator[str]:
+    for df in results:
+        yield f"--{boundary_token}\r\n"
+        yield "Content-Type: application/json\r\n\r\n"
+        yield await get_running_loop().run_in_executor(None, encode_df_as_json_string, df)
+        yield "\r\n"
+    yield f"--{boundary_token}--\r\n"
+
+
+@router.get("/spans", include_in_schema=False, deprecated=True)
+async def get_spans_handler(
+    request: Request,
+    request_body: QuerySpansRequestBody,
+    project_name: Optional[str] = Query(
+        default=None, description="The project name to get evaluations from"
+    ),
+) -> Response:
+    return await query_spans_handler(request, request_body, project_name)
+
+
+class SpanAnnotationResult(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 SpanAnnotation(V1RoutesBaseModel):
+    span_id: str = Field(description="OpenTelemetry Span 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[SpanAnnotationResult] = 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.SpanAnnotation:
+        return Precursors.SpanAnnotation(
+            self.span_id,
+            models.SpanAnnotation(
+                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 AnnotateSpansRequestBody(RequestBody[list[SpanAnnotation]]):
+    data: list[SpanAnnotation]
+
+
+class InsertedSpanAnnotation(V1RoutesBaseModel):
+    id: str = Field(description="The ID of the inserted span annotation")
+
+
+class AnnotateSpansResponseBody(ResponseBody[list[InsertedSpanAnnotation]]):
+    pass
+
+
+@router.post(
+    "/span_annotations",
+    operation_id="annotateSpans",
+    summary="Create or update span annotations",
+    responses=add_errors_to_responses(
+        [{"status_code": HTTP_404_NOT_FOUND, "description": "Span not found"}]
+    ),
+    response_description="Span annotations inserted successfully",
+    include_in_schema=True,
+)
+async def annotate_spans(
+    request: Request,
+    request_body: AnnotateSpansRequestBody,
+    sync: bool = Query(default=False, description="If true, fulfill request synchronously."),
+) -> AnnotateSpansResponseBody:
+    if not request_body.data:
+        return AnnotateSpansResponseBody(data=[])
+    precursors = [d.as_precursor() for d in request_body.data]
+    if not sync:
+        await request.state.enqueue(*precursors)
+        return AnnotateSpansResponseBody(data=[])
+
+    span_ids = {p.span_id for p in precursors}
+    async with request.app.state.db() as session:
+        existing_spans = {
+            span.span_id: span.id
+            async for span in await session.stream_scalars(
+                select(models.Span).filter(models.Span.span_id.in_(span_ids))
+            )
+        }
+
+        missing_span_ids = span_ids - set(existing_spans.keys())
+        if missing_span_ids:
+            raise HTTPException(
+                detail=f"Spans with IDs {', '.join(missing_span_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_spans[p.span_id]).row))
+            span_annotation_id = await session.scalar(
+                insert_on_conflict(
+                    values,
+                    dialect=dialect,
+                    table=models.SpanAnnotation,
+                    unique_by=("name", "span_rowid"),
+                ).returning(models.SpanAnnotation.id)
+            )
+            inserted_ids.append(span_annotation_id)
+    request.state.event_queue.put(SpanAnnotationInsertEvent(tuple(inserted_ids)))
+    return AnnotateSpansResponseBody(
+        data=[
+            InsertedSpanAnnotation(id=str(GlobalID("SpanAnnotation", str(id_))))
+            for id_ in inserted_ids
+        ]
+    )