arize-phoenix 10.0.4__py3-none-any.whl → 12.28.1__py3-none-any.whl

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

@@ -15,21 +15,16 @@ from typing import Any, Optional, Union, cast
 
 import pandas as pd
 import pyarrow as pa
-from fastapi import APIRouter, BackgroundTasks, HTTPException, Path, Query
+from fastapi import APIRouter, BackgroundTasks, Depends, HTTPException, Path, Query
 from fastapi.responses import PlainTextResponse, StreamingResponse
-from sqlalchemy import and_, delete, func, select
+from sqlalchemy import and_, case, delete, func, select
 from sqlalchemy.ext.asyncio import AsyncSession
 from starlette.concurrency import run_in_threadpool
 from starlette.datastructures import FormData, UploadFile
 from starlette.requests import Request
 from starlette.responses import Response
 from starlette.status import (
-    HTTP_200_OK,
-    HTTP_204_NO_CONTENT,
     HTTP_404_NOT_FOUND,
-    HTTP_409_CONFLICT,
-    HTTP_422_UNPROCESSABLE_ENTITY,
-    HTTP_429_TOO_MANY_REQUESTS,
 )
 from strawberry.relay import GlobalID
 from typing_extensions import TypeAlias, assert_never
@@ -42,11 +37,15 @@ from phoenix.db.insertion.dataset import (
     ExampleContent,
     add_dataset_examples,
 )
+from phoenix.db.types.db_models import UNDEFINED
 from phoenix.server.api.types.Dataset import Dataset as DatasetNodeType
 from phoenix.server.api.types.DatasetExample import DatasetExample as DatasetExampleNodeType
+from phoenix.server.api.types.DatasetSplit import DatasetSplit as DatasetSplitNodeType
 from phoenix.server.api.types.DatasetVersion import DatasetVersion as DatasetVersionNodeType
 from phoenix.server.api.types.node import from_global_id_with_expected_type
 from phoenix.server.api.utils import delete_projects, delete_traces
+from phoenix.server.authorization import is_not_locked
+from phoenix.server.bearer_auth import PhoenixUser
 from phoenix.server.dml_event import DatasetInsertEvent
 
 from .models import V1RoutesBaseModel
@@ -57,6 +56,11 @@ from .utils import (
     add_text_csv_content_to_responses,
 )
 
+csv.field_size_limit(
+    1_000_000_000  # allows large field sizes for CSV upload (1GB)
+)
+
+
 logger = logging.getLogger(__name__)
 
 DATASET_NODE_NAME = DatasetNodeType.__name__
@@ -73,6 +77,7 @@ class Dataset(V1RoutesBaseModel):
     metadata: dict[str, Any]
     created_at: datetime
     updated_at: datetime
+    example_count: int
 
 
 class ListDatasetsResponseBody(PaginatedResponseBody[Dataset]):
@@ -83,7 +88,7 @@ class ListDatasetsResponseBody(PaginatedResponseBody[Dataset]):
     "/datasets",
     operation_id="listDatasets",
     summary="List datasets",
-    responses=add_errors_to_responses([HTTP_422_UNPROCESSABLE_ENTITY]),
+    responses=add_errors_to_responses([422]),
 )
 async def list_datasets(
     request: Request,
@@ -97,7 +102,18 @@ async def list_datasets(
     ),
 ) -> ListDatasetsResponseBody:
     async with request.app.state.db() as session:
-        query = select(models.Dataset).order_by(models.Dataset.id.desc())
+        value = case(
+            (models.DatasetExampleRevision.revision_kind == "CREATE", 1),
+            (models.DatasetExampleRevision.revision_kind == "DELETE", -1),
+        )
+        query = (
+            select(models.Dataset)
+            .add_columns(func.coalesce(func.sum(value), 0).label("example_count"))
+            .outerjoin_from(models.Dataset, models.DatasetExample)
+            .outerjoin_from(models.DatasetExample, models.DatasetExampleRevision)
+            .group_by(models.Dataset.id)
+            .order_by(models.Dataset.id.desc())
+        )
 
         if cursor:
             try:
@@ -106,25 +122,26 @@ async def list_datasets(
             except ValueError:
                 raise HTTPException(
                     detail=f"Invalid cursor format: {cursor}",
-                    status_code=HTTP_422_UNPROCESSABLE_ENTITY,
+                    status_code=422,
                 )
         if name:
             query = query.filter(models.Dataset.name == name)
 
         query = query.limit(limit + 1)
         result = await session.execute(query)
-        datasets = result.scalars().all()
-
+        datasets = result.all()
         if not datasets:
             return ListDatasetsResponseBody(next_cursor=None, data=[])
 
         next_cursor = None
         if len(datasets) == limit + 1:
-            next_cursor = str(GlobalID(DATASET_NODE_NAME, str(datasets[-1].id)))
+            dataset = datasets[-1][0]
+            next_cursor = str(GlobalID(DATASET_NODE_NAME, str(dataset.id)))
             datasets = datasets[:-1]
 
         data = []
-        for dataset in datasets:
+        for row in datasets:
+            dataset = row[0]
             data.append(
                 Dataset(
                     id=str(GlobalID(DATASET_NODE_NAME, str(dataset.id))),
@@ -133,6 +150,7 @@ async def list_datasets(
                     metadata=dataset.metadata_,
                     created_at=dataset.created_at,
                     updated_at=dataset.updated_at,
+                    example_count=row[1],
                 )
             )
 
@@ -143,11 +161,11 @@ async def list_datasets(
     "/datasets/{id}",
     operation_id="deleteDatasetById",
     summary="Delete dataset by ID",
-    status_code=HTTP_204_NO_CONTENT,
+    status_code=204,
     responses=add_errors_to_responses(
         [
-            {"status_code": HTTP_404_NOT_FOUND, "description": "Dataset not found"},
-            {"status_code": HTTP_422_UNPROCESSABLE_ENTITY, "description": "Invalid dataset ID"},
+            {"status_code": 404, "description": "Dataset not found"},
+            {"status_code": 422, "description": "Invalid dataset ID"},
         ]
     ),
 )
@@ -161,11 +179,9 @@ async def delete_dataset(
                 DATASET_NODE_NAME,
             )
         except ValueError:
-            raise HTTPException(
-                detail=f"Invalid Dataset ID: {id}", status_code=HTTP_422_UNPROCESSABLE_ENTITY
-            )
+            raise HTTPException(detail=f"Invalid Dataset ID: {id}", status_code=422)
     else:
-        raise HTTPException(detail="Missing Dataset ID", status_code=HTTP_422_UNPROCESSABLE_ENTITY)
+        raise HTTPException(detail="Missing Dataset ID", status_code=422)
     project_names_stmt = get_project_names_for_datasets(dataset_id)
     eval_trace_ids_stmt = get_eval_trace_ids_for_datasets(dataset_id)
     stmt = (
@@ -175,7 +191,7 @@ async def delete_dataset(
         project_names = await session.scalars(project_names_stmt)
         eval_trace_ids = await session.scalars(eval_trace_ids_stmt)
         if (await session.scalar(stmt)) is None:
-            raise HTTPException(detail="Dataset does not exist", status_code=HTTP_404_NOT_FOUND)
+            raise HTTPException(detail="Dataset does not exist", status_code=404)
     tasks = BackgroundTasks()
     tasks.add_task(delete_projects, request.app.state.db, *project_names)
     tasks.add_task(delete_traces, request.app.state.db, *eval_trace_ids)
@@ -193,17 +209,21 @@ class GetDatasetResponseBody(ResponseBody[DatasetWithExampleCount]):
     "/datasets/{id}",
     operation_id="getDataset",
     summary="Get dataset by ID",
-    responses=add_errors_to_responses([HTTP_404_NOT_FOUND]),
+    responses=add_errors_to_responses([404]),
 )
 async def get_dataset(
     request: Request, id: str = Path(description="The ID of the dataset")
 ) -> GetDatasetResponseBody:
-    dataset_id = GlobalID.from_id(id)
+    try:
+        dataset_id = GlobalID.from_id(id)
+    except Exception as e:
+        raise HTTPException(
+            detail=f"Invalid dataset ID format: {id}",
+            status_code=422,
+        ) from e
 
     if (type_name := dataset_id.type_name) != DATASET_NODE_NAME:
-        raise HTTPException(
-            detail=f"ID {dataset_id} refers to a f{type_name}", status_code=HTTP_404_NOT_FOUND
-        )
+        raise HTTPException(detail=f"ID {dataset_id} refers to a f{type_name}", status_code=404)
     async with request.app.state.db() as session:
         result = await session.execute(
             select(models.Dataset, models.Dataset.example_count).filter(
@@ -214,9 +234,7 @@ async def get_dataset(
         dataset = dataset_query[0] if dataset_query else None
         example_count = dataset_query[1] if dataset_query else 0
         if dataset is None:
-            raise HTTPException(
-                detail=f"Dataset with ID {dataset_id} not found", status_code=HTTP_404_NOT_FOUND
-            )
+            raise HTTPException(detail=f"Dataset with ID {dataset_id} not found", status_code=404)
 
         dataset = DatasetWithExampleCount(
             id=str(dataset_id),
@@ -245,7 +263,7 @@ class ListDatasetVersionsResponseBody(PaginatedResponseBody[DatasetVersion]):
     "/datasets/{id}/versions",
     operation_id="listDatasetVersionsByDatasetId",
     summary="List dataset versions",
-    responses=add_errors_to_responses([HTTP_422_UNPROCESSABLE_ENTITY]),
+    responses=add_errors_to_responses([422]),
 )
 async def list_dataset_versions(
     request: Request,
@@ -267,12 +285,12 @@ async def list_dataset_versions(
         except ValueError:
             raise HTTPException(
                 detail=f"Invalid Dataset ID: {id}",
-                status_code=HTTP_422_UNPROCESSABLE_ENTITY,
+                status_code=422,
             )
     else:
         raise HTTPException(
             detail="Missing Dataset ID",
-            status_code=HTTP_422_UNPROCESSABLE_ENTITY,
+            status_code=422,
         )
     stmt = (
         select(models.DatasetVersion)
@@ -288,7 +306,7 @@ async def list_dataset_versions(
         except ValueError:
             raise HTTPException(
                 detail=f"Invalid cursor: {cursor}",
-                status_code=HTTP_422_UNPROCESSABLE_ENTITY,
+                status_code=422,
             )
         max_dataset_version_id = (
             select(models.DatasetVersion.id)
@@ -312,6 +330,7 @@ async def list_dataset_versions(
 
 class UploadDatasetData(V1RoutesBaseModel):
     dataset_id: str
+    version_id: str
 
 
 class UploadDatasetResponseBody(ResponseBody[UploadDatasetData]):
@@ -320,15 +339,16 @@ class UploadDatasetResponseBody(ResponseBody[UploadDatasetData]):
 
 @router.post(
     "/datasets/upload",
+    dependencies=[Depends(is_not_locked)],
     operation_id="uploadDataset",
-    summary="Upload dataset from JSON, CSV, or PyArrow",
+    summary="Upload dataset from JSON, JSONL, CSV, or PyArrow",
     responses=add_errors_to_responses(
         [
             {
-                "status_code": HTTP_409_CONFLICT,
+                "status_code": 409,
                 "description": "Dataset of the same name already exists",
             },
-            {"status_code": HTTP_422_UNPROCESSABLE_ENTITY, "description": "Invalid request body"},
+            {"status_code": 422, "description": "Invalid request body"},
         ]
     ),
     # FastAPI cannot generate the request body portion of the OpenAPI schema for
@@ -350,6 +370,17 @@ class UploadDatasetResponseBody(ResponseBody[UploadDatasetData]):
                             "inputs": {"type": "array", "items": {"type": "object"}},
                             "outputs": {"type": "array", "items": {"type": "object"}},
                             "metadata": {"type": "array", "items": {"type": "object"}},
+                            "splits": {
+                                "type": "array",
+                                "items": {
+                                    "oneOf": [
+                                        {"type": "string"},
+                                        {"type": "array", "items": {"type": "string"}},
+                                        {"type": "null"},
+                                    ]
+                                },
+                                "description": "Split per example: string, string array, or null",
+                            },
                         },
                     }
                 },
@@ -376,6 +407,12 @@ class UploadDatasetResponseBody(ResponseBody[UploadDatasetData]):
                                 "items": {"type": "string"},
                                 "uniqueItems": True,
                             },
+                            "split_keys[]": {
+                                "type": "array",
+                                "items": {"type": "string"},
+                                "uniqueItems": True,
+                                "description": "Column names for auto-assigning examples to splits",
+                            },
                             "file": {"type": "string", "format": "binary"},
                         },
                     }
@@ -391,7 +428,12 @@ async def upload_dataset(
         description="If true, fulfill request synchronously and return JSON containing dataset_id.",
     ),
 ) -> Optional[UploadDatasetResponseBody]:
-    request_content_type = request.headers["content-type"]
+    request_content_type = request.headers.get("content-type")
+    if not request_content_type:
+        raise HTTPException(
+            detail="Missing content-type header",
+            status_code=400,
+        )
     examples: Union[Examples, Awaitable[Examples]]
     if request_content_type.startswith("application/json"):
         try:
@@ -401,14 +443,14 @@ async def upload_dataset(
         except ValueError as e:
             raise HTTPException(
                 detail=str(e),
-                status_code=HTTP_422_UNPROCESSABLE_ENTITY,
+                status_code=422,
             )
         if action is DatasetAction.CREATE:
             async with request.app.state.db() as session:
                 if await _check_table_exists(session, name):
                     raise HTTPException(
                         detail=f"Dataset with the same name already exists: {name=}",
-                        status_code=HTTP_409_CONFLICT,
+                        status_code=409,
                     )
     elif request_content_type.startswith("multipart/form-data"):
         async with request.form() as form:
@@ -420,19 +462,20 @@ async def upload_dataset(
                     input_keys,
                     output_keys,
                     metadata_keys,
+                    split_keys,
                     file,
                 ) = await _parse_form_data(form)
             except ValueError as e:
                 raise HTTPException(
                     detail=str(e),
-                    status_code=HTTP_422_UNPROCESSABLE_ENTITY,
+                    status_code=422,
                 )
             if action is DatasetAction.CREATE:
                 async with request.app.state.db() as session:
                     if await _check_table_exists(session, name):
                         raise HTTPException(
                             detail=f"Dataset with the same name already exists: {name=}",
-                            status_code=HTTP_409_CONFLICT,
+                            status_code=409,
                         )
             content = await file.read()
         try:
@@ -440,22 +483,32 @@ async def upload_dataset(
             if file_content_type is FileContentType.CSV:
                 encoding = FileContentEncoding(file.headers.get("content-encoding"))
                 examples = await _process_csv(
-                    content, encoding, input_keys, output_keys, metadata_keys
+                    content, encoding, input_keys, output_keys, metadata_keys, split_keys
                 )
             elif file_content_type is FileContentType.PYARROW:
-                examples = await _process_pyarrow(content, input_keys, output_keys, metadata_keys)
+                examples = await _process_pyarrow(
+                    content, input_keys, output_keys, metadata_keys, split_keys
+                )
+            elif file_content_type is FileContentType.JSONL:
+                encoding = FileContentEncoding(file.headers.get("content-encoding"))
+                examples = await _process_jsonl(
+                    content, encoding, input_keys, output_keys, metadata_keys, split_keys
+                )
             else:
                 assert_never(file_content_type)
         except ValueError as e:
             raise HTTPException(
                 detail=str(e),
-                status_code=HTTP_422_UNPROCESSABLE_ENTITY,
+                status_code=422,
             )
     else:
         raise HTTPException(
             detail="Invalid request Content-Type",
-            status_code=HTTP_422_UNPROCESSABLE_ENTITY,
+            status_code=422,
         )
+    user_id: Optional[int] = None
+    if request.app.state.authentication_enabled and isinstance(request.user, PhoenixUser):
+        user_id = int(request.user.identity)
     operation = cast(
         Callable[[AsyncSession], Awaitable[DatasetExampleAdditionEvent]],
         partial(
@@ -464,27 +517,34 @@ async def upload_dataset(
             action=action,
             name=name,
             description=description,
+            user_id=user_id,
         ),
     )
     if sync:
         async with request.app.state.db() as session:
-            dataset_id = (await operation(session)).dataset_id
+            event = await operation(session)
+            dataset_id = event.dataset_id
+            version_id = event.dataset_version_id
         request.state.event_queue.put(DatasetInsertEvent((dataset_id,)))
         return UploadDatasetResponseBody(
-            data=UploadDatasetData(dataset_id=str(GlobalID(Dataset.__name__, str(dataset_id))))
+            data=UploadDatasetData(
+                dataset_id=str(GlobalID(Dataset.__name__, str(dataset_id))),
+                version_id=str(GlobalID(DatasetVersion.__name__, str(version_id))),
+            )
         )
     try:
         request.state.enqueue_operation(operation)
     except QueueFull:
         if isinstance(examples, Coroutine):
             examples.close()
-        raise HTTPException(detail="Too many requests.", status_code=HTTP_429_TOO_MANY_REQUESTS)
+        raise HTTPException(detail="Too many requests.", status_code=429)
     return None
 
 
 class FileContentType(Enum):
     CSV = "text/csv"
     PYARROW = "application/x-pandas-pyarrow"
+    JSONL = "application/jsonl"
 
     @classmethod
     def _missing_(cls, v: Any) -> "FileContentType":
@@ -512,6 +572,7 @@ Description: TypeAlias = Optional[str]
 InputKeys: TypeAlias = frozenset[str]
 OutputKeys: TypeAlias = frozenset[str]
 MetadataKeys: TypeAlias = frozenset[str]
+SplitKeys: TypeAlias = frozenset[str]
 DatasetId: TypeAlias = int
 Examples: TypeAlias = Iterator[ExampleContent]
 
@@ -528,18 +589,55 @@ def _process_json(
         raise ValueError("input is required")
     if not isinstance(inputs, list) or not _is_all_dict(inputs):
         raise ValueError("Input should be a list containing only dictionary objects")
-    outputs, metadata = data.get("outputs"), data.get("metadata")
+    outputs, metadata, splits = data.get("outputs"), data.get("metadata"), data.get("splits")
     for k, v in {"outputs": outputs, "metadata": metadata}.items():
         if v and not (isinstance(v, list) and len(v) == len(inputs) and _is_all_dict(v)):
             raise ValueError(
                 f"{k} should be a list of same length as input containing only dictionary objects"
             )
+
+    # Validate splits format if provided
+    if splits is not None:
+        if not isinstance(splits, list):
+            raise ValueError("splits must be a list")
+        if len(splits) != len(inputs):
+            raise ValueError(
+                f"splits must have same length as inputs ({len(splits)} != {len(inputs)})"
+            )
     examples: list[ExampleContent] = []
     for i, obj in enumerate(inputs):
+        # Extract split values, validating they're non-empty strings
+        split_set: set[str] = set()
+        if splits:
+            split_value = splits[i]
+            if split_value is None:
+                # Sparse assignment: None means no splits for this example
+                pass
+            elif isinstance(split_value, str):
+                # Format 1: Single string value
+                if split_value.strip():
+                    split_set.add(split_value.strip())
+            elif isinstance(split_value, list):
+                # Format 2: List of strings (multiple splits)
+                for v in split_value:
+                    if v is None:
+                        continue  # Skip None values in the list
+                    if not isinstance(v, str):
+                        raise ValueError(
+                            f"Split value must be a string or None, got {type(v).__name__}"
+                        )
+                    if v.strip():
+                        split_set.add(v.strip())
+            else:
+                raise ValueError(
+                    f"Split value must be a string, list of strings, or None, "
+                    f"got {type(split_value).__name__}"
+                )
         example = ExampleContent(
             input=obj,
             output=outputs[i] if outputs else {},
             metadata=metadata[i] if metadata else {},
+            splits=frozenset(split_set),
         )
         examples.append(example)
     action = DatasetAction(cast(Optional[str], data.get("action")) or "create")
@@ -552,6 +650,7 @@ async def _process_csv(
     input_keys: InputKeys,
     output_keys: OutputKeys,
     metadata_keys: MetadataKeys,
+    split_keys: SplitKeys,
 ) -> Examples:
     if content_encoding is FileContentEncoding.GZIP:
         content = await run_in_threadpool(gzip.decompress, content)
@@ -566,12 +665,15 @@ async def _process_csv(
     if freq > 1:
         raise ValueError(f"Duplicated column header in CSV file: {header}")
     column_headers = frozenset(reader.fieldnames)
-    _check_keys_exist(column_headers, input_keys, output_keys, metadata_keys)
+    _check_keys_exist(column_headers, input_keys, output_keys, metadata_keys, split_keys)
     return (
         ExampleContent(
             input={k: row.get(k) for k in input_keys},
             output={k: row.get(k) for k in output_keys},
             metadata={k: row.get(k) for k in metadata_keys},
+            splits=frozenset(
+                str(v).strip() for k in split_keys if (v := row.get(k)) and str(v).strip()
+            ),  # Only include non-empty, non-whitespace split values
         )
         for row in iter(reader)
     )
@@ -582,13 +684,14 @@ async def _process_pyarrow(
     input_keys: InputKeys,
     output_keys: OutputKeys,
     metadata_keys: MetadataKeys,
+    split_keys: SplitKeys,
 ) -> Awaitable[Examples]:
     try:
         reader = pa.ipc.open_stream(content)
     except pa.ArrowInvalid as e:
         raise ValueError("File is not valid pyarrow") from e
     column_headers = frozenset(reader.schema.names)
-    _check_keys_exist(column_headers, input_keys, output_keys, metadata_keys)
+    _check_keys_exist(column_headers, input_keys, output_keys, metadata_keys, split_keys)
 
     def get_examples() -> Iterator[ExampleContent]:
         for row in reader.read_pandas().to_dict(orient="records"):
@@ -596,11 +699,48 @@ async def _process_pyarrow(
                 input={k: row.get(k) for k in input_keys},
                 output={k: row.get(k) for k in output_keys},
                 metadata={k: row.get(k) for k in metadata_keys},
+                splits=frozenset(
+                    str(v).strip() for k in split_keys if (v := row.get(k)) and str(v).strip()
+                ),  # Only include non-empty, non-whitespace split values
             )
 
     return run_in_threadpool(get_examples)
 
 
+async def _process_jsonl(
+    content: bytes,
+    encoding: FileContentEncoding,
+    input_keys: InputKeys,
+    output_keys: OutputKeys,
+    metadata_keys: MetadataKeys,
+    split_keys: SplitKeys,
+) -> Examples:
+    if encoding is FileContentEncoding.GZIP:
+        content = await run_in_threadpool(gzip.decompress, content)
+    elif encoding is FileContentEncoding.DEFLATE:
+        content = await run_in_threadpool(zlib.decompress, content)
+    elif encoding is not FileContentEncoding.NONE:
+        assert_never(encoding)
+    # content is a newline delimited list of JSON objects
+    # parse within a threadpool
+    reader = await run_in_threadpool(
+        lambda c: [json.loads(line) for line in c.decode().splitlines()], content
+    )
+
+    examples: list[ExampleContent] = []
+    for obj in reader:
+        example = ExampleContent(
+            input={k: obj.get(k) for k in input_keys},
+            output={k: obj.get(k) for k in output_keys},
+            metadata={k: obj.get(k) for k in metadata_keys},
+            splits=frozenset(
+                str(v).strip() for k in split_keys if (v := obj.get(k)) and str(v).strip()
+            ),  # Only include non-empty, non-whitespace split values
+        )
+        examples.append(example)
+    return iter(examples)
+
+
 async def _check_table_exists(session: AsyncSession, name: str) -> bool:
     return bool(
         await session.scalar(
@@ -614,11 +754,13 @@ def _check_keys_exist(
     input_keys: InputKeys,
     output_keys: OutputKeys,
     metadata_keys: MetadataKeys,
+    split_keys: SplitKeys,
 ) -> None:
     for desc, keys in (
         ("input", input_keys),
         ("output", output_keys),
         ("metadata", metadata_keys),
+        ("split", split_keys),
     ):
         if keys and (diff := keys.difference(column_headers)):
             raise ValueError(f"{desc} keys not found in column headers: {diff}")
@@ -633,6 +775,7 @@ async def _parse_form_data(
     InputKeys,
     OutputKeys,
     MetadataKeys,
+    SplitKeys,
     UploadFile,
 ]:
     name = cast(Optional[str], form.get("name"))
@@ -646,6 +789,7 @@ async def _parse_form_data(
     input_keys = frozenset(filter(bool, cast(list[str], form.getlist("input_keys[]"))))
     output_keys = frozenset(filter(bool, cast(list[str], form.getlist("output_keys[]"))))
     metadata_keys = frozenset(filter(bool, cast(list[str], form.getlist("metadata_keys[]"))))
+    split_keys = frozenset(filter(bool, cast(list[str], form.getlist("split_keys[]"))))
     return (
         action,
         name,
@@ -653,6 +797,7 @@ async def _parse_form_data(
         input_keys,
         output_keys,
         metadata_keys,
+        split_keys,
         file,
     )
 
@@ -668,6 +813,7 @@ class DatasetExample(V1RoutesBaseModel):
 class ListDatasetExamplesData(V1RoutesBaseModel):
     dataset_id: str
     version_id: str
+    filtered_splits: list[str] = UNDEFINED
     examples: list[DatasetExample]
 
 
@@ -679,7 +825,7 @@ class ListDatasetExamplesResponseBody(ResponseBody[ListDatasetExamplesData]):
     "/datasets/{id}/examples",
     operation_id="getDatasetExamples",
     summary="Get examples from a dataset",
-    responses=add_errors_to_responses([HTTP_404_NOT_FOUND]),
+    responses=add_errors_to_responses([404]),
 )
 async def get_dataset_examples(
     request: Request,
@@ -687,22 +833,38 @@ async def get_dataset_examples(
     version_id: Optional[str] = Query(
         default=None,
         description=(
-            "The ID of the dataset version " "(if omitted, returns data from the latest version)"
+            "The ID of the dataset version (if omitted, returns data from the latest version)"
         ),
     ),
+    split: Optional[list[str]] = Query(
+        default=None,
+        description="List of dataset split identifiers (GlobalIDs or names) to filter by",
+    ),
 ) -> ListDatasetExamplesResponseBody:
-    dataset_gid = GlobalID.from_id(id)
-    version_gid = GlobalID.from_id(version_id) if version_id else None
+    try:
+        dataset_gid = GlobalID.from_id(id)
+    except Exception as e:
+        raise HTTPException(
+            detail=f"Invalid dataset ID format: {id}",
+            status_code=422,
+        ) from e
+
+    if version_id:
+        try:
+            version_gid = GlobalID.from_id(version_id)
+        except Exception as e:
+            raise HTTPException(
+                detail=f"Invalid dataset version ID format: {version_id}",
+                status_code=422,
+            ) from e
+    else:
+        version_gid = None
 
     if (dataset_type := dataset_gid.type_name) != "Dataset":
-        raise HTTPException(
-            detail=f"ID {dataset_gid} refers to a {dataset_type}", status_code=HTTP_404_NOT_FOUND
-        )
+        raise HTTPException(detail=f"ID {dataset_gid} refers to a {dataset_type}", status_code=404)
 
     if version_gid and (version_type := version_gid.type_name) != "DatasetVersion":
-        raise HTTPException(
-            detail=f"ID {version_gid} refers to a {version_type}", status_code=HTTP_404_NOT_FOUND
-        )
+        raise HTTPException(detail=f"ID {version_gid} refers to a {version_type}", status_code=404)
 
     async with request.app.state.db() as session:
         if (
@@ -712,7 +874,7 @@ async def get_dataset_examples(
         ) is None:
             raise HTTPException(
                 detail=f"No dataset with id {dataset_gid} can be found.",
-                status_code=HTTP_404_NOT_FOUND,
+                status_code=404,
             )
 
         # Subquery to find the maximum created_at for each dataset_example_id
@@ -734,7 +896,7 @@ async def get_dataset_examples(
             ) is None:
                 raise HTTPException(
                     detail=f"No dataset version with id {version_id} can be found.",
-                    status_code=HTTP_404_NOT_FOUND,
+                    status_code=404,
                 )
             # if a version_id is provided, filter the subquery to only include revisions from that
             partial_subquery = partial_subquery.filter(
@@ -750,13 +912,17 @@ async def get_dataset_examples(
             ) is None:
                 raise HTTPException(
                     detail="Dataset has no versions.",
-                    status_code=HTTP_404_NOT_FOUND,
+                    status_code=404,
                 )
 
         subquery = partial_subquery.subquery()
+
         # Query for the most recent example revisions that are not deleted
         query = (
-            select(models.DatasetExample, models.DatasetExampleRevision)
+            select(
+                models.DatasetExample,
+                models.DatasetExampleRevision,
+            )
             .join(
                 models.DatasetExampleRevision,
                 models.DatasetExample.id == models.DatasetExampleRevision.dataset_example_id,
@@ -769,6 +935,28 @@ async def get_dataset_examples(
             .filter(models.DatasetExampleRevision.revision_kind != "DELETE")
             .order_by(models.DatasetExample.id.asc())
         )
+
+        # If splits are provided, filter by dataset splits
+        resolved_split_names: list[str] = []
+        if split:
+            # Resolve split identifiers (IDs or names) to IDs and names
+            resolved_split_ids, resolved_split_names = await _resolve_split_identifiers(
+                session, split
+            )
+
+            # Add filter for splits (join with the association table)
+            # Use distinct() to prevent duplicates when an example belongs to
+            # multiple splits
+            query = (
+                query.join(
+                    models.DatasetSplitDatasetExample,
+                    models.DatasetExample.id
+                    == models.DatasetSplitDatasetExample.dataset_example_id,
+                )
+                .filter(models.DatasetSplitDatasetExample.dataset_split_id.in_(resolved_split_ids))
+                .distinct()
+            )
+
         examples = [
             DatasetExample(
                 id=str(GlobalID("DatasetExample", str(example.id))),
@@ -783,6 +971,7 @@ async def get_dataset_examples(
         data=ListDatasetExamplesData(
             dataset_id=str(GlobalID("Dataset", str(resolved_dataset_id))),
             version_id=str(GlobalID("DatasetVersion", str(resolved_version_id))),
+            filtered_splits=resolved_split_names,
             examples=examples,
         )
     )
@@ -793,10 +982,10 @@ async def get_dataset_examples(
     operation_id="getDatasetCsv",
     summary="Download dataset examples as CSV file",
     response_class=StreamingResponse,
-    status_code=HTTP_200_OK,
+    status_code=200,
     responses={
-        **add_errors_to_responses([HTTP_422_UNPROCESSABLE_ENTITY]),
-        **add_text_csv_content_to_responses(HTTP_200_OK),
+        **add_errors_to_responses([422]),
+        **add_text_csv_content_to_responses(200),
     },
 )
 async def get_dataset_csv(
@@ -806,7 +995,7 @@ async def get_dataset_csv(
     version_id: Optional[str] = Query(
         default=None,
         description=(
-            "The ID of the dataset version " "(if omitted, returns data from the latest version)"
+            "The ID of the dataset version (if omitted, returns data from the latest version)"
         ),
     ),
 ) -> Response:
@@ -816,7 +1005,7 @@ async def get_dataset_csv(
                 session=session, id=id, version_id=version_id
             )
     except ValueError as e:
-        raise HTTPException(detail=str(e), status_code=HTTP_422_UNPROCESSABLE_ENTITY)
+        raise HTTPException(detail=str(e), status_code=422)
     content = await run_in_threadpool(_get_content_csv, examples)
     encoded_dataset_name = urllib.parse.quote(dataset_name)
     return Response(
@@ -836,7 +1025,7 @@ async def get_dataset_csv(
     responses=add_errors_to_responses(
         [
             {
-                "status_code": HTTP_422_UNPROCESSABLE_ENTITY,
+                "status_code": 422,
                 "description": "Invalid dataset or version ID",
             }
         ]
@@ -849,7 +1038,7 @@ async def get_dataset_jsonl_openai_ft(
     version_id: Optional[str] = Query(
         default=None,
         description=(
-            "The ID of the dataset version " "(if omitted, returns data from the latest version)"
+            "The ID of the dataset version (if omitted, returns data from the latest version)"
         ),
     ),
 ) -> bytes:
@@ -859,7 +1048,7 @@ async def get_dataset_jsonl_openai_ft(
                 session=session, id=id, version_id=version_id
             )
     except ValueError as e:
-        raise HTTPException(detail=str(e), status_code=HTTP_422_UNPROCESSABLE_ENTITY)
+        raise HTTPException(detail=str(e), status_code=422)
     content = await run_in_threadpool(_get_content_jsonl_openai_ft, examples)
     encoded_dataset_name = urllib.parse.quote(dataset_name)
     response.headers["content-disposition"] = (
@@ -876,7 +1065,7 @@ async def get_dataset_jsonl_openai_ft(
     responses=add_errors_to_responses(
         [
             {
-                "status_code": HTTP_422_UNPROCESSABLE_ENTITY,
+                "status_code": 422,
                 "description": "Invalid dataset or version ID",
             }
         ]
@@ -889,7 +1078,7 @@ async def get_dataset_jsonl_openai_evals(
     version_id: Optional[str] = Query(
         default=None,
         description=(
-            "The ID of the dataset version " "(if omitted, returns data from the latest version)"
+            "The ID of the dataset version (if omitted, returns data from the latest version)"
         ),
     ),
 ) -> bytes:
@@ -899,7 +1088,7 @@ async def get_dataset_jsonl_openai_evals(
                 session=session, id=id, version_id=version_id
             )
     except ValueError as e:
-        raise HTTPException(detail=str(e), status_code=HTTP_422_UNPROCESSABLE_ENTITY)
+        raise HTTPException(detail=str(e), status_code=422)
     content = await run_in_threadpool(_get_content_jsonl_openai_evals, examples)
     encoded_dataset_name = urllib.parse.quote(dataset_name)
     response.headers["content-disposition"] = (
@@ -978,12 +1167,25 @@ def _get_content_jsonl_openai_evals(examples: list[models.DatasetExampleRevision
 async def _get_db_examples(
     *, session: Any, id: str, version_id: Optional[str]
 ) -> tuple[str, list[models.DatasetExampleRevision]]:
-    dataset_id = from_global_id_with_expected_type(GlobalID.from_id(id), DATASET_NODE_NAME)
+    try:
+        dataset_id = from_global_id_with_expected_type(GlobalID.from_id(id), DATASET_NODE_NAME)
+    except Exception as e:
+        raise HTTPException(
+            detail=f"Invalid dataset ID format: {id}",
+            status_code=422,
+        ) from e
+
     dataset_version_id: Optional[int] = None
     if version_id:
-        dataset_version_id = from_global_id_with_expected_type(
-            GlobalID.from_id(version_id), DATASET_VERSION_NODE_NAME
-        )
+        try:
+            dataset_version_id = from_global_id_with_expected_type(
+                GlobalID.from_id(version_id), DATASET_VERSION_NODE_NAME
+            )
+        except Exception as e:
+            raise HTTPException(
+                detail=f"Invalid dataset version ID format: {version_id}",
+                status_code=422,
+            ) from e
     latest_version = (
         select(
             models.DatasetExampleRevision.dataset_example_id,
@@ -1026,3 +1228,115 @@ async def _get_db_examples(
 
 def _is_all_dict(seq: Sequence[Any]) -> bool:
     return all(map(lambda obj: isinstance(obj, dict), seq))
+
+
+# Split identifier helper types and functions
+class _SplitId(int): ...
+
+
+_SplitIdentifier: TypeAlias = Union[_SplitId, str]
+
+
+def _parse_split_identifier(split_identifier: str) -> _SplitIdentifier:
+    """
+    Parse a split identifier as either a GlobalID or a name.
+
+    Args:
+        split_identifier: The identifier string (GlobalID or name)
+
+    Returns:
+        Either a _SplitId or an Identifier
+
+    Raises:
+        HTTPException: If the identifier format is invalid
+    """
+    if not split_identifier:
+        raise HTTPException(422, "Invalid split identifier")
+    try:
+        split_id = from_global_id_with_expected_type(
+            GlobalID.from_id(split_identifier),
+            DatasetSplitNodeType.__name__,
+        )
+    except ValueError:
+        return split_identifier
+    return _SplitId(split_id)
+
+
+async def _resolve_split_identifiers(
+    session: AsyncSession,
+    split_identifiers: list[str],
+) -> tuple[list[int], list[str]]:
+    """
+    Resolve a list of split identifiers (IDs or names) to split IDs and names.
+
+    Args:
+        session: The database session
+        split_identifiers: List of split identifiers (GlobalIDs or names)
+
+    Returns:
+        Tuple of (list of split IDs, list of split names)
+
+    Raises:
+        HTTPException: If any split identifier is invalid or not found
+    """
+    split_ids: list[int] = []
+    split_names: list[str] = []
+
+    # Parse all identifiers first
+    parsed_identifiers: list[_SplitIdentifier] = []
+    for identifier_str in split_identifiers:
+        parsed_identifiers.append(_parse_split_identifier(identifier_str.strip()))
+
+    # Separate IDs and names
+    requested_ids: list[int] = []
+    requested_names: list[str] = []
+    for identifier in parsed_identifiers:
+        if isinstance(identifier, _SplitId):
+            requested_ids.append(int(identifier))
+        elif isinstance(identifier, str):
+            requested_names.append(identifier)
+        else:
+            assert_never(identifier)
+
+    # Query for splits by ID
+    if requested_ids:
+        id_results = await session.stream(
+            select(models.DatasetSplit.id, models.DatasetSplit.name).where(
+                models.DatasetSplit.id.in_(requested_ids)
+            )
+        )
+        async for split_id, split_name in id_results:
+            split_ids.append(split_id)
+            split_names.append(split_name)
+
+        # Check if all requested IDs were found
+        found_ids = set(split_ids[-len(requested_ids) :] if requested_ids else [])
+        missing_ids = [sid for sid in requested_ids if sid not in found_ids]
+        if missing_ids:
+            raise HTTPException(
+                status_code=HTTP_404_NOT_FOUND,
+                detail=f"Dataset splits not found for IDs: {', '.join(map(str, missing_ids))}",
+            )
+
+    # Query for splits by name
+    if requested_names:
+        name_results = await session.stream(
+            select(models.DatasetSplit.id, models.DatasetSplit.name).where(
+                models.DatasetSplit.name.in_(requested_names)
+            )
+        )
+        name_to_id: dict[str, int] = {}
+        async for split_id, split_name in name_results:
+            split_ids.append(split_id)
+            split_names.append(split_name)
+            name_to_id[split_name] = split_id
+
+        # Check if all requested names were found
+        missing_names = [name for name in requested_names if name not in name_to_id]
+        if missing_names:
+            raise HTTPException(
+                status_code=HTTP_404_NOT_FOUND,
+                detail=f"Dataset splits not found: {', '.join(missing_names)}",
+            )
+
+    return split_ids, split_names