arize-phoenix 4.14.1__py3-none-any.whl → 4.16.0__py3-none-any.whl

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

@@ -1,8 +1,10 @@
 from datetime import datetime
+from typing import Any, List, Optional
 
+from fastapi import APIRouter, HTTPException
+from pydantic import Field
 from sqlalchemy import select
 from starlette.requests import Request
-from starlette.responses import JSONResponse, Response
 from starlette.status import HTTP_404_NOT_FOUND
 from strawberry.relay import GlobalID
 
@@ -10,188 +12,129 @@ from phoenix.db import models
 from phoenix.db.models import ExperimentRunOutput
 from phoenix.server.api.types.node import from_global_id_with_expected_type
 
+from .pydantic_compat import V1RoutesBaseModel
+from .utils import ResponseBody, add_errors_to_responses
 
-async def create_experiment_run(request: Request) -> Response:
-    """
-    summary: Create a new experiment run for a specific experiment
-    operationId: createExperimentRun
-    tags:
-      - private
-    parameters:
-      - in: path
-        name: experiment_id
-        required: true
-        description: The ID of the experiment for which the run is being created
-        schema:
-          type: string
-    requestBody:
-      description: Details of the experiment run to be created
-      required: true
-      content:
-        application/json:
-          schema:
-            type: object
-            properties:
-              dataset_example_id:
-                type: string
-                description: The ID of the dataset example used in the experiment run
-              trace_id:
-                type: string
-                description: Optional trace ID for tracking
-              output:
-                description: The output of the experiment task
-              repetition_number:
-                type: integer
-                description: The repetition number of the experiment run
-              start_time:
-                type: string
-                format: date-time
-                description: The start time of the experiment run in ISO format
-              end_time:
-                type: string
-                format: date-time
-                description: The end time of the experiment run in ISO format
-              error:
-                type: string
-                description: Optional error message if the experiment run encountered an error
-                nullable: true
-            required:
-              - dataset_example_id
-              - output
-              - repetition_number
-              - start_time
-              - end_time
-    responses:
-      200:
-        description: Experiment run created successfully
-        content:
-          application/json:
-            schema:
-              type: object
-              properties:
-                data:
-                  type: object
-                  properties:
-                    id:
-                      type: string
-                      description: The ID of the created experiment run
-      404:
-        description: Experiment or DatasetExample not found
-    """
-    experiment_gid = GlobalID.from_id(request.path_params["experiment_id"])
+router = APIRouter(tags=["experiments"], include_in_schema=False)
+
+
+class ExperimentRun(V1RoutesBaseModel):
+    dataset_example_id: str = Field(
+        description="The ID of the dataset example used in the experiment run"
+    )
+    output: Any = Field(description="The output of the experiment task")
+    repetition_number: int = Field(description="The repetition number of the experiment run")
+    start_time: datetime = Field(description="The start time of the experiment run")
+    end_time: datetime = Field(description="The end time of the experiment run")
+    trace_id: Optional[str] = Field(
+        default=None, description="The ID of the corresponding trace (if one exists)"
+    )
+    error: Optional[str] = Field(
+        default=None,
+        description="Optional error message if the experiment run encountered an error",
+    )
+
+
+class CreateExperimentRunRequestBody(ExperimentRun):
+    pass
+
+
+class CreateExperimentRunResponseBodyData(V1RoutesBaseModel):
+    id: str = Field(description="The ID of the newly created experiment run")
+
+
+class CreateExperimentResponseBody(ResponseBody[CreateExperimentRunResponseBodyData]):
+    pass
+
+
+@router.post(
+    "/experiments/{experiment_id}/runs",
+    operation_id="createExperimentRun",
+    summary="Create run for an experiment",
+    response_description="Experiment run created successfully",
+    responses=add_errors_to_responses(
+        [
+            {
+                "status_code": HTTP_404_NOT_FOUND,
+                "description": "Experiment or dataset example not found",
+            }
+        ]
+    ),
+)
+async def create_experiment_run(
+    request: Request, experiment_id: str, request_body: CreateExperimentRunRequestBody
+) -> CreateExperimentResponseBody:
+    experiment_gid = GlobalID.from_id(experiment_id)
     try:
-        experiment_id = from_global_id_with_expected_type(experiment_gid, "Experiment")
+        experiment_rowid = from_global_id_with_expected_type(experiment_gid, "Experiment")
     except ValueError:
-        return Response(
-            content=f"Experiment with ID {experiment_gid} does not exist",
+        raise HTTPException(
+            detail=f"Experiment with ID {experiment_gid} does not exist",
             status_code=HTTP_404_NOT_FOUND,
         )
 
-    payload = await request.json()
-
-    example_gid = GlobalID.from_id(payload["dataset_example_id"])
+    example_gid = GlobalID.from_id(request_body.dataset_example_id)
     try:
         dataset_example_id = from_global_id_with_expected_type(example_gid, "DatasetExample")
     except ValueError:
-        return Response(
-            content=f"DatasetExample with ID {example_gid} does not exist",
+        raise HTTPException(
+            detail=f"DatasetExample with ID {example_gid} does not exist",
             status_code=HTTP_404_NOT_FOUND,
         )
 
-    trace_id = payload.get("trace_id", None)
-    task_output = payload["output"]
-    repetition_number = payload["repetition_number"]
-    start_time = payload["start_time"]
-    end_time = payload["end_time"]
-    error = payload.get("error")
+    trace_id = request_body.trace_id
+    task_output = request_body.output
+    repetition_number = request_body.repetition_number
+    start_time = request_body.start_time
+    end_time = request_body.end_time
+    error = request_body.error
 
     async with request.app.state.db() as session:
         exp_run = models.ExperimentRun(
-            experiment_id=experiment_id,
+            experiment_id=experiment_rowid,
             dataset_example_id=dataset_example_id,
             trace_id=trace_id,
             output=ExperimentRunOutput(task_output=task_output),
             repetition_number=repetition_number,
-            start_time=datetime.fromisoformat(start_time),
-            end_time=datetime.fromisoformat(end_time),
+            start_time=start_time,
+            end_time=end_time,
             error=error,
         )
         session.add(exp_run)
         await session.flush()
     run_gid = GlobalID("ExperimentRun", str(exp_run.id))
-    return JSONResponse(content={"data": {"id": str(run_gid)}})
-
-
-async def list_experiment_runs(request: Request) -> Response:
-    """
-    summary: List all runs for a specific experiment
-    operationId: listExperimentRuns
-    tags:
-      - private
-    parameters:
-      - in: path
-        name: experiment_id
-        required: true
-        description: The ID of the experiment to list runs for
-        schema:
-          type: string
-    responses:
-      200:
-        description: Experiment runs retrieved successfully
-        content:
-          application/json:
-            schema:
-              type: object
-              properties:
-                data:
-                  type: array
-                  items:
-                    type: object
-                    properties:
-                      id:
-                        type: string
-                        description: The ID of the experiment run
-                      experiment_id:
-                        type: string
-                        description: The ID of the experiment
-                      dataset_example_id:
-                        type: string
-                        description: The ID of the dataset example
-                      repetition_number:
-                        type: integer
-                        description: The repetition number of the experiment run
-                      start_time:
-                        type: string
-                        format: date-time
-                        description: The start time of the experiment run in ISO format
-                      end_time:
-                        type: string
-                        format: date-time
-                        description: The end time of the experiment run in ISO format
-                      output:
-                        description: The output of the experiment task
-                      error:
-                        type: string
-                        description: Error message if the experiment run encountered an error
-                      trace_id:
-                        type: string
-                        description: Optional trace ID for tracking
-      404:
-        description: Experiment not found
-    """
-    experiment_gid = GlobalID.from_id(request.path_params["experiment_id"])
+    return CreateExperimentResponseBody(data=CreateExperimentRunResponseBodyData(id=str(run_gid)))
+
+
+class ListExperimentRunsResponseBody(ResponseBody[List[ExperimentRun]]):
+    pass
+
+
+@router.get(
+    "/experiments/{experiment_id}/runs",
+    operation_id="listExperimentRuns",
+    summary="List runs for an experiment",
+    response_description="Experiment runs retrieved successfully",
+    responses=add_errors_to_responses(
+        [{"status_code": HTTP_404_NOT_FOUND, "description": "Experiment not found"}]
+    ),
+)
+async def list_experiment_runs(
+    request: Request, experiment_id: str
+) -> ListExperimentRunsResponseBody:
+    experiment_gid = GlobalID.from_id(experiment_id)
     try:
-        experiment_id = from_global_id_with_expected_type(experiment_gid, "Experiment")
+        experiment_rowid = from_global_id_with_expected_type(experiment_gid, "Experiment")
     except ValueError:
-        return Response(
-            content=f"Experiment with ID {experiment_gid} does not exist",
+        raise HTTPException(
+            detail=f"Experiment with ID {experiment_gid} does not exist",
             status_code=HTTP_404_NOT_FOUND,
         )
 
     async with request.app.state.db() as session:
         experiment_runs = await session.execute(
             select(models.ExperimentRun)
-            .where(models.ExperimentRun.experiment_id == experiment_id)
+            .where(models.ExperimentRun.experiment_id == experiment_rowid)
             # order by dataset_example_id to be consistent with `list_dataset_examples`
             .order_by(models.ExperimentRun.dataset_example_id.asc())
         )
@@ -202,9 +145,9 @@ async def list_experiment_runs(request: Request) -> Response:
             experiment_gid = GlobalID("Experiment", str(exp_run.experiment_id))
             example_gid = GlobalID("DatasetExample", str(exp_run.dataset_example_id))
             runs.append(
-                dict(
-                    start_time=exp_run.start_time.isoformat(),
-                    end_time=exp_run.end_time.isoformat(),
+                ExperimentRun(
+                    start_time=exp_run.start_time,
+                    end_time=exp_run.end_time,
                     experiment_id=str(experiment_gid),
                     dataset_example_id=str(example_gid),
                     repetition_number=exp_run.repetition_number,
@@ -214,4 +157,4 @@ async def list_experiment_runs(request: Request) -> Response:
                     trace_id=exp_run.trace_id,
                 )
             )
-    return JSONResponse(content={"data": runs}, status_code=200)
+    return ListExperimentRunsResponseBody(data=runs)