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

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

@@ -1,10 +1,8 @@
 from datetime import datetime
-from typing import Any, List, Optional
 
-from fastapi import APIRouter, HTTPException
-from pydantic import BaseModel, 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
 
@@ -12,128 +10,188 @@ 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 .utils import ResponseBody, add_errors_to_responses
 
-router = APIRouter(tags=["experiments"], include_in_schema=False)
-
-
-class ExperimentRun(BaseModel):
-    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(BaseModel):
-    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)
+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"])
     try:
-        experiment_rowid = from_global_id_with_expected_type(experiment_gid, "Experiment")
+        experiment_id = from_global_id_with_expected_type(experiment_gid, "Experiment")
     except ValueError:
-        raise HTTPException(
-            detail=f"Experiment with ID {experiment_gid} does not exist",
+        return Response(
+            content=f"Experiment with ID {experiment_gid} does not exist",
             status_code=HTTP_404_NOT_FOUND,
         )
 
-    example_gid = GlobalID.from_id(request_body.dataset_example_id)
+    payload = await request.json()
+
+    example_gid = GlobalID.from_id(payload["dataset_example_id"])
     try:
         dataset_example_id = from_global_id_with_expected_type(example_gid, "DatasetExample")
     except ValueError:
-        raise HTTPException(
-            detail=f"DatasetExample with ID {example_gid} does not exist",
+        return Response(
+            content=f"DatasetExample with ID {example_gid} does not exist",
             status_code=HTTP_404_NOT_FOUND,
         )
 
-    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
+    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")
 
     async with request.app.state.db() as session:
         exp_run = models.ExperimentRun(
-            experiment_id=experiment_rowid,
+            experiment_id=experiment_id,
             dataset_example_id=dataset_example_id,
             trace_id=trace_id,
             output=ExperimentRunOutput(task_output=task_output),
             repetition_number=repetition_number,
-            start_time=start_time,
-            end_time=end_time,
+            start_time=datetime.fromisoformat(start_time),
+            end_time=datetime.fromisoformat(end_time),
             error=error,
         )
         session.add(exp_run)
         await session.flush()
     run_gid = GlobalID("ExperimentRun", str(exp_run.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)
+    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"])
     try:
-        experiment_rowid = from_global_id_with_expected_type(experiment_gid, "Experiment")
+        experiment_id = from_global_id_with_expected_type(experiment_gid, "Experiment")
     except ValueError:
-        raise HTTPException(
-            detail=f"Experiment with ID {experiment_gid} does not exist",
+        return Response(
+            content=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_rowid)
+            .where(models.ExperimentRun.experiment_id == experiment_id)
             # order by dataset_example_id to be consistent with `list_dataset_examples`
             .order_by(models.ExperimentRun.dataset_example_id.asc())
         )
@@ -144,9 +202,9 @@ async def list_experiment_runs(
             experiment_gid = GlobalID("Experiment", str(exp_run.experiment_id))
             example_gid = GlobalID("DatasetExample", str(exp_run.dataset_example_id))
             runs.append(
-                ExperimentRun(
-                    start_time=exp_run.start_time,
-                    end_time=exp_run.end_time,
+                dict(
+                    start_time=exp_run.start_time.isoformat(),
+                    end_time=exp_run.end_time.isoformat(),
                     experiment_id=str(experiment_gid),
                     dataset_example_id=str(example_gid),
                     repetition_number=exp_run.repetition_number,
@@ -156,4 +214,4 @@ async def list_experiment_runs(
                     trace_id=exp_run.trace_id,
                 )
             )
-    return ListExperimentRunsResponseBody(data=runs)
+    return JSONResponse(content={"data": runs}, status_code=200)