arize-phoenix 4.4.4rc4__py3-none-any.whl → 4.4.4rc5__py3-none-any.whl

phoenix/db/migrations/versions/10460e46d750_datasets.py

@@ -72,7 +72,7 @@ def upgrade() -> None:
         sa.Column(
             "span_rowid",
             sa.Integer,
-            sa.ForeignKey("spans.id"),
+            sa.ForeignKey("spans.id", ondelete="SET NULL"),
             nullable=True,
             index=True,
         ),
@@ -198,7 +198,7 @@ def upgrade() -> None:
             sa.String,
             nullable=True,
         ),
-        sa.Column("output", JSON_, nullable=True),
+        sa.Column("output", JSON_, nullable=False),
         sa.Column("start_time", sa.TIMESTAMP(timezone=True), nullable=False),
         sa.Column("end_time", sa.TIMESTAMP(timezone=True), nullable=False),
         sa.Column(

phoenix/db/models.py

@@ -1,5 +1,5 @@
 from datetime import datetime, timezone
-from typing import Any, Dict, List, Optional
+from typing import Any, Dict, List, Optional, TypedDict
 
 from sqlalchemy import (
     JSON,
@@ -91,6 +91,10 @@ class UtcTimeStamp(TypeDecorator[datetime]):
         return normalize_datetime(value, timezone.utc)
 
 
+class ExperimentResult(TypedDict, total=False):
+    result: Dict[str, Any]
+
+
 class Base(DeclarativeBase):
     # Enforce best practices for naming constraints
     # https://alembic.sqlalchemy.org/en/latest/naming.html#integration-of-naming-conventions-into-operations-autogenerate
@@ -106,6 +110,7 @@ class Base(DeclarativeBase):
     type_annotation_map = {
         Dict[str, Any]: JsonDict,
         List[Dict[str, Any]]: JsonList,
+        ExperimentResult: JsonDict,
     }
 
 
@@ -483,7 +488,7 @@ class DatasetExample(Base):
         index=True,
     )
     span_rowid: Mapped[Optional[int]] = mapped_column(
-        ForeignKey("spans.id"),
+        ForeignKey("spans.id", ondelete="SET NULL"),
         index=True,
         nullable=True,
     )
@@ -556,7 +561,7 @@ class ExperimentRun(Base):
     )
     repetition_number: Mapped[int]
     trace_id: Mapped[Optional[str]]
-    output: Mapped[Optional[Dict[str, Any]]]
+    output: Mapped[ExperimentResult]
     start_time: Mapped[datetime] = mapped_column(UtcTimeStamp)
     end_time: Mapped[datetime] = mapped_column(UtcTimeStamp)
     prompt_token_count: Mapped[Optional[int]]

phoenix/server/api/context.py

@@ -20,6 +20,7 @@ from phoenix.server.api.dataloaders import (
     EvaluationSummaryDataLoader,
     ExperimentAnnotationSummaryDataLoader,
     ExperimentErrorRatesDataLoader,
+    ExperimentRunCountsDataLoader,
     ExperimentSequenceNumberDataLoader,
     LatencyMsQuantileDataLoader,
     MinStartOrMaxEndTimeDataLoader,
@@ -44,6 +45,7 @@ class DataLoaders:
     evaluation_summaries: EvaluationSummaryDataLoader
     experiment_annotation_summaries: ExperimentAnnotationSummaryDataLoader
     experiment_error_rates: ExperimentErrorRatesDataLoader
+    experiment_run_counts: ExperimentRunCountsDataLoader
     experiment_sequence_number: ExperimentSequenceNumberDataLoader
     latency_ms_quantile: LatencyMsQuantileDataLoader
     min_start_or_max_end_times: MinStartOrMaxEndTimeDataLoader

phoenix/server/api/dataloaders/__init__.py

@@ -19,6 +19,7 @@ from .document_retrieval_metrics import DocumentRetrievalMetricsDataLoader
 from .evaluation_summaries import EvaluationSummaryCache, EvaluationSummaryDataLoader
 from .experiment_annotation_summaries import ExperimentAnnotationSummaryDataLoader
 from .experiment_error_rates import ExperimentErrorRatesDataLoader
+from .experiment_run_counts import ExperimentRunCountsDataLoader
 from .experiment_sequence_number import ExperimentSequenceNumberDataLoader
 from .latency_ms_quantile import LatencyMsQuantileCache, LatencyMsQuantileDataLoader
 from .min_start_or_max_end_times import MinStartOrMaxEndTimeCache, MinStartOrMaxEndTimeDataLoader
@@ -41,6 +42,7 @@ __all__ = [
     "EvaluationSummaryDataLoader",
     "ExperimentAnnotationSummaryDataLoader",
     "ExperimentErrorRatesDataLoader",
+    "ExperimentRunCountsDataLoader",
     "ExperimentSequenceNumberDataLoader",
     "LatencyMsQuantileDataLoader",
     "MinStartOrMaxEndTimeDataLoader",

phoenix/server/api/dataloaders/experiment_run_counts.py

@@ -0,0 +1,42 @@
+from typing import (
+    AsyncContextManager,
+    Callable,
+    List,
+)
+
+from sqlalchemy import func, select
+from sqlalchemy.ext.asyncio import AsyncSession
+from strawberry.dataloader import DataLoader
+from typing_extensions import TypeAlias
+
+from phoenix.db import models
+
+ExperimentID: TypeAlias = int
+RunCount: TypeAlias = int
+Key: TypeAlias = ExperimentID
+Result: TypeAlias = RunCount
+
+
+class ExperimentRunCountsDataLoader(DataLoader[Key, Result]):
+    def __init__(
+        self,
+        db: Callable[[], AsyncContextManager[AsyncSession]],
+    ) -> None:
+        super().__init__(load_fn=self._load_fn)
+        self._db = db
+
+    async def _load_fn(self, keys: List[Key]) -> List[Result]:
+        experiment_ids = keys
+        async with self._db() as session:
+            run_counts = {
+                experiment_id: run_count
+                async for experiment_id, run_count in await session.stream(
+                    select(models.ExperimentRun.experiment_id, func.count())
+                    .where(models.ExperimentRun.experiment_id.in_(set(experiment_ids)))
+                    .group_by(models.ExperimentRun.experiment_id)
+                )
+            }
+        return [
+            run_counts.get(experiment_id, ValueError(f"Unknown experiment: {experiment_id}"))
+            for experiment_id in experiment_ids
+        ]

phoenix/server/api/helpers/dataset_helpers.py

@@ -1,3 +1,4 @@
+import json
 from typing import Any, Dict, Literal, Mapping, Optional, Protocol
 
 from openinference.semconv.trace import (
@@ -128,14 +129,14 @@ def _get_generic_io_value(
     Makes a best-effort attempt to extract the input or output value from a span
     and returns it as a dictionary.
     """
-    if isinstance(io_value, str) and (
-        mime_type == OpenInferenceMimeTypeValues.TEXT.value or mime_type is None
-    ):
+    if mime_type == OpenInferenceMimeTypeValues.JSON.value:
+        parsed_value = json.loads(io_value)
+        if isinstance(parsed_value, dict):
+            return parsed_value
+        else:
+            return {kind: parsed_value}
+    if isinstance(io_value, str):
         return {kind: io_value}
-    if isinstance(io_value, dict) and (
-        mime_type == OpenInferenceMimeTypeValues.JSON.value or mime_type is None
-    ):
-        return io_value
     return {}
 
 

phoenix/server/api/input_types/ClearProjectInput.py

@@ -0,0 +1,15 @@
+from datetime import datetime
+from typing import Optional
+
+import strawberry
+from strawberry import UNSET
+from strawberry.relay import GlobalID
+
+
+@strawberry.input
+class ClearProjectInput:
+    id: GlobalID
+    end_time: Optional[datetime] = strawberry.field(
+        default=UNSET,
+        description="The time up to which to purge data. Time is right-open /non-inclusive.",
+    )

phoenix/server/api/mutations/project_mutations.py

@@ -8,6 +8,7 @@ from phoenix.config import DEFAULT_PROJECT_NAME
 from phoenix.db import models
 from phoenix.db.insertion.span import ClearProjectSpansEvent
 from phoenix.server.api.context import Context
+from phoenix.server.api.input_types.ClearProjectInput import ClearProjectInput
 from phoenix.server.api.mutations.auth import IsAuthenticated
 from phoenix.server.api.queries import Query
 from phoenix.server.api.types.node import from_global_id_with_expected_type
@@ -32,11 +33,15 @@ class ProjectMutationMixin:
         return Query()
 
     @strawberry.mutation(permission_classes=[IsAuthenticated])  # type: ignore
-    async def clear_project(self, info: Info[Context, None], id: GlobalID) -> Query:
-        project_id = from_global_id_with_expected_type(global_id=id, expected_type_name="Project")
+    async def clear_project(self, info: Info[Context, None], input: ClearProjectInput) -> Query:
+        project_id = from_global_id_with_expected_type(
+            global_id=input.id, expected_type_name="Project"
+        )
         delete_statement = delete(models.Trace).where(models.Trace.project_rowid == project_id)
+        if input.end_time is not None:
+            delete_statement = delete_statement.where(models.Trace.start_time < input.end_time)
         async with info.context.db() as session:
             await session.execute(delete_statement)
-            if cache := info.context.cache_for_dataloaders:
-                cache.invalidate(ClearProjectSpansEvent(project_rowid=project_id))
+        if cache := info.context.cache_for_dataloaders:
+            cache.invalidate(ClearProjectSpansEvent(project_rowid=project_id))
         return Query()

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

@@ -13,11 +13,12 @@ from typing import (
     Awaitable,
     Callable,
     Coroutine,
-    Dict,
     FrozenSet,
     Iterator,
     List,
+    Mapping,
     Optional,
+    Sequence,
     Tuple,
     Union,
     cast,
@@ -32,8 +33,8 @@ from starlette.datastructures import FormData, UploadFile
 from starlette.requests import Request
 from starlette.responses import JSONResponse, Response
 from starlette.status import (
-    HTTP_403_FORBIDDEN,
     HTTP_404_NOT_FOUND,
+    HTTP_409_CONFLICT,
     HTTP_422_UNPROCESSABLE_ENTITY,
     HTTP_429_TOO_MANY_REQUESTS,
 )
@@ -44,6 +45,7 @@ from phoenix.db import models
 from phoenix.db.insertion.dataset import (
     DatasetAction,
     DatasetExampleAdditionEvent,
+    ExampleContent,
     add_dataset_examples,
 )
 from phoenix.server.api.types.Dataset import Dataset
@@ -350,7 +352,7 @@ async def get_dataset_versions(request: Request) -> Response:
 
 async def post_datasets_upload(request: Request) -> Response:
     """
-    summary: Upload CSV or PyArrow file as dataset
+    summary: Upload dataset as either JSON or file (CSV or PyArrow)
     operationId: uploadDataset
     tags:
       - datasets
@@ -362,6 +364,32 @@ async def post_datasets_upload(request: Request) -> Response:
           type: boolean
     requestBody:
       content:
+        application/json:
+          schema:
+            type: object
+            required:
+              - name
+              - inputs
+            properties:
+              action:
+                type: string
+                enum: [create, append]
+              name:
+                type: string
+              description:
+                type: string
+              inputs:
+                type: array
+                items:
+                  type: object
+              outputs:
+                type: array
+                items:
+                  type: object
+              metadata:
+                type: array
+                items:
+                  type: object
         multipart/form-data:
           schema:
             type: object
@@ -401,22 +429,18 @@ async def post_datasets_upload(request: Request) -> Response:
         description: Success
       403:
         description: Forbidden
+      409:
+        description: Dataset of the same name already exists
       422:
         description: Request body is invalid
     """
-    if request.app.state.read_only:
-        return Response(status_code=HTTP_403_FORBIDDEN)
-    async with request.form() as form:
+    request_content_type = request.headers["content-type"]
+    examples: Union[Examples, Awaitable[Examples]]
+    if request_content_type.startswith("application/json"):
         try:
-            (
-                action,
-                name,
-                description,
-                input_keys,
-                output_keys,
-                metadata_keys,
-                file,
-            ) = await _parse_form_data(form)
+            examples, action, name, description = await run_in_threadpool(
+                _process_json, await request.json()
+            )
         except ValueError as e:
             return Response(
                 content=str(e),
@@ -426,24 +450,53 @@ async def post_datasets_upload(request: Request) -> Response:
             async with request.app.state.db() as session:
                 if await _check_table_exists(session, name):
                     return Response(
-                        content=f"Dataset already exists: {name=}",
-                        status_code=HTTP_422_UNPROCESSABLE_ENTITY,
+                        content=f"Dataset with the same name already exists: {name=}",
+                        status_code=HTTP_409_CONFLICT,
                     )
-        content = await file.read()
-    try:
-        examples: Union[Examples, Awaitable[Examples]]
-        content_type = FileContentType(file.content_type)
-        if content_type is FileContentType.CSV:
-            encoding = FileContentEncoding(file.headers.get("content-encoding"))
-            examples, column_headers = await _process_csv(content, encoding)
-        elif content_type is FileContentType.PYARROW:
-            examples, column_headers = await _process_pyarrow(content)
-        else:
-            assert_never(content_type)
-        _check_keys_exist(column_headers, input_keys, output_keys, metadata_keys)
-    except ValueError as e:
+    elif request_content_type.startswith("multipart/form-data"):
+        async with request.form() as form:
+            try:
+                (
+                    action,
+                    name,
+                    description,
+                    input_keys,
+                    output_keys,
+                    metadata_keys,
+                    file,
+                ) = await _parse_form_data(form)
+            except ValueError as e:
+                return Response(
+                    content=str(e),
+                    status_code=HTTP_422_UNPROCESSABLE_ENTITY,
+                )
+            if action is DatasetAction.CREATE:
+                async with request.app.state.db() as session:
+                    if await _check_table_exists(session, name):
+                        return Response(
+                            content=f"Dataset with the same name already exists: {name=}",
+                            status_code=HTTP_409_CONFLICT,
+                        )
+            content = await file.read()
+        try:
+            file_content_type = FileContentType(file.content_type)
+            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
+                )
+            elif file_content_type is FileContentType.PYARROW:
+                examples = await _process_pyarrow(content, input_keys, output_keys, metadata_keys)
+            else:
+                assert_never(file_content_type)
+        except ValueError as e:
+            return Response(
+                content=str(e),
+                status_code=HTTP_422_UNPROCESSABLE_ENTITY,
+            )
+    else:
         return Response(
-            content=str(e),
+            content=str("Invalid request Content-Type"),
             status_code=HTTP_422_UNPROCESSABLE_ENTITY,
         )
     operation = cast(
@@ -454,9 +507,6 @@ async def post_datasets_upload(request: Request) -> Response:
             action=action,
             name=name,
             description=description,
-            input_keys=input_keys,
-            output_keys=output_keys,
-            metadata_keys=metadata_keys,
         ),
     )
     if request.query_params.get("sync") == "true":
@@ -505,13 +555,46 @@ InputKeys: TypeAlias = FrozenSet[str]
 OutputKeys: TypeAlias = FrozenSet[str]
 MetadataKeys: TypeAlias = FrozenSet[str]
 DatasetId: TypeAlias = int
-Examples: TypeAlias = Iterator[Dict[str, Any]]
+Examples: TypeAlias = Iterator[ExampleContent]
+
+
+def _process_json(
+    data: Mapping[str, Any],
+) -> Tuple[Examples, DatasetAction, Name, Description]:
+    name = data.get("name")
+    if not name:
+        raise ValueError("Dataset name is required")
+    description = data.get("description") or ""
+    inputs = data.get("inputs")
+    if not inputs:
+        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")
+    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"
+            )
+    examples: List[ExampleContent] = []
+    for i, obj in enumerate(inputs):
+        example = ExampleContent(
+            input=obj,
+            output=outputs[i] if outputs else {},
+            metadata=metadata[i] if metadata else {},
+        )
+        examples.append(example)
+    action = DatasetAction(cast(Optional[str], data.get("action")) or "create")
+    return iter(examples), action, name, description
 
 
 async def _process_csv(
     content: bytes,
     content_encoding: FileContentEncoding,
-) -> Tuple[Examples, FrozenSet[str]]:
+    input_keys: InputKeys,
+    output_keys: OutputKeys,
+    metadata_keys: MetadataKeys,
+) -> Examples:
     if content_encoding is FileContentEncoding.GZIP:
         content = await run_in_threadpool(gzip.decompress, content)
     elif content_encoding is FileContentEncoding.DEFLATE:
@@ -525,22 +608,39 @@ async def _process_csv(
     if freq > 1:
         raise ValueError(f"Duplicated column header in CSV file: {header}")
     column_headers = frozenset(reader.fieldnames)
-    return reader, column_headers
+    _check_keys_exist(column_headers, input_keys, output_keys, metadata_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},
+        )
+        for row in iter(reader)
+    )
 
 
 async def _process_pyarrow(
     content: bytes,
-) -> Tuple[Awaitable[Examples], FrozenSet[str]]:
+    input_keys: InputKeys,
+    output_keys: OutputKeys,
+    metadata_keys: MetadataKeys,
+) -> 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)
+
+    def get_examples() -> Iterator[ExampleContent]:
+        for row in reader.read_pandas().to_dict(orient="records"):
+            yield 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},
+            )
 
-    def get_examples() -> Iterator[Dict[str, Any]]:
-        yield from reader.read_pandas().to_dict(orient="records")
-
-    return run_in_threadpool(get_examples), column_headers
+    return run_in_threadpool(get_examples)
 
 
 async def _check_table_exists(session: AsyncSession, name: str) -> bool:
@@ -859,3 +959,7 @@ async def _get_db_examples(request: Request) -> Tuple[str, List[models.DatasetEx
             raise ValueError("Dataset does not exist.")
         examples = [r async for r in await session.stream_scalars(stmt)]
     return dataset_name, examples
+
+
+def _is_all_dict(seq: Sequence[Any]) -> bool:
+    return all(map(lambda obj: isinstance(obj, dict), seq))

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

@@ -43,6 +43,7 @@ async def create_experiment_evaluation(request: Request) -> Response:
             metadata_=metadata,
             start_time=datetime.fromisoformat(start_time),
             end_time=datetime.fromisoformat(end_time),
+            trace_id=payload.get("trace_id"),
         )
         session.add(exp_eval_run)
         await session.flush()

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

@@ -61,7 +61,7 @@ async def create_experiment_run(request: Request) -> Response:
             experiment_id=str(experiment_gid),
             dataset_example_id=str(example_gid),
             repetition_number=exp_run.repetition_number,
-            output=ExperimentResult(result=exp_run.output),
+            output=ExperimentResult.from_dict(exp_run.output) if exp_run.output else None,
             error=exp_run.error,
             id=str(run_gid),
             trace_id=exp_run.trace_id,
@@ -99,7 +99,7 @@ async def list_experiment_runs(request: Request) -> Response:
                     experiment_id=str(experiment_gid),
                     dataset_example_id=str(example_gid),
                     repetition_number=exp_run.repetition_number,
-                    output=ExperimentResult(result=exp_run.output),
+                    output=ExperimentResult.from_dict(exp_run.output) if exp_run.output else None,
                     error=exp_run.error,
                     id=str(run_gid),
                     trace_id=exp_run.trace_id,

phoenix/server/api/types/Experiment.py

@@ -75,6 +75,11 @@ class Experiment(Node):
             ).all()
         return connection_from_list([to_gql_experiment_run(run) for run in runs], args)
 
+    @strawberry.field
+    async def run_count(self, info: Info[Context, None]) -> int:
+        experiment_id = self.id_attr
+        return await info.context.data_loaders.experiment_run_counts.load(experiment_id)
+
     @strawberry.field
     async def annotation_summaries(
         self, info: Info[Context, None]

phoenix/server/api/types/ExperimentRun.py

@@ -84,7 +84,7 @@ def to_gql_experiment_run(run: models.ExperimentRun) -> ExperimentRun:
         trace_id=trace_id
         if (trace := run.trace) and (trace_id := trace.trace_id) is not None
         else None,
-        output=run.output,
+        output=run.output.get("result"),
         start_time=run.start_time,
         end_time=run.end_time,
         error=run.error,

phoenix/server/api/types/ExperimentRunAnnotation.py

@@ -33,7 +33,7 @@ class ExperimentRunAnnotation(Node):
         if (trace := await dataloader.load(self.trace_id)) is None:
             return None
         trace_row_id, project_row_id = trace
-        return Trace(id_attr=trace_row_id, trace_id=trace.trace_id, project_rowid=project_row_id)
+        return Trace(id_attr=trace_row_id, trace_id=self.trace_id, project_rowid=project_row_id)
 
 
 def to_gql_experiment_run_annotation(

phoenix/server/app.py

@@ -65,6 +65,7 @@ from phoenix.server.api.dataloaders import (
     EvaluationSummaryDataLoader,
     ExperimentAnnotationSummaryDataLoader,
     ExperimentErrorRatesDataLoader,
+    ExperimentRunCountsDataLoader,
     ExperimentSequenceNumberDataLoader,
     LatencyMsQuantileDataLoader,
     MinStartOrMaxEndTimeDataLoader,
@@ -208,6 +209,7 @@ class GraphQLWithContext(GraphQL):  # type: ignore
                 ),
                 experiment_annotation_summaries=ExperimentAnnotationSummaryDataLoader(self.db),
                 experiment_error_rates=ExperimentErrorRatesDataLoader(self.db),
+                experiment_run_counts=ExperimentRunCountsDataLoader(self.db),
                 experiment_sequence_number=ExperimentSequenceNumberDataLoader(self.db),
                 latency_ms_quantile=LatencyMsQuantileDataLoader(
                     self.db,