arize-phoenix 4.10.1__py3-none-any.whl → 4.10.2rc1__py3-none-any.whl

{arize_phoenix-4.10.1.dist-info → arize_phoenix-4.10.2rc1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: arize-phoenix
-Version: 4.10.1
+Version: 4.10.2rc1
 Summary: AI Observability and Evaluation
 Project-URL: Documentation, https://docs.arize.com/phoenix/
 Project-URL: Issues, https://github.com/Arize-ai/phoenix/issues
@@ -22,6 +22,7 @@ Requires-Dist: aiosqlite
 Requires-Dist: alembic<2,>=1.3.0
 Requires-Dist: arize-phoenix-evals>=0.13.1
 Requires-Dist: cachetools
+Requires-Dist: fastapi
 Requires-Dist: grpcio
 Requires-Dist: hdbscan>=0.8.33
 Requires-Dist: httpx
@@ -40,8 +41,8 @@ Requires-Dist: pandas>=1.0
 Requires-Dist: protobuf<6.0,>=3.20
 Requires-Dist: psutil
 Requires-Dist: pyarrow
+Requires-Dist: pydantic<3,>=1.0
 Requires-Dist: python-multipart
-Requires-Dist: pyyaml
 Requires-Dist: scikit-learn
 Requires-Dist: scipy
 Requires-Dist: sqlalchemy[asyncio]<3,>=2.0.4
@@ -56,9 +57,9 @@ Requires-Dist: uvicorn
 Requires-Dist: wrapt
 Provides-Extra: container
 Requires-Dist: opentelemetry-exporter-otlp; extra == 'container'
+Requires-Dist: opentelemetry-instrumentation-fastapi; extra == 'container'
 Requires-Dist: opentelemetry-instrumentation-grpc; extra == 'container'
 Requires-Dist: opentelemetry-instrumentation-sqlalchemy; extra == 'container'
-Requires-Dist: opentelemetry-instrumentation-starlette; extra == 'container'
 Requires-Dist: opentelemetry-proto>=1.12.0; extra == 'container'
 Requires-Dist: opentelemetry-sdk; extra == 'container'
 Requires-Dist: opentelemetry-semantic-conventions; extra == 'container'

{arize_phoenix-4.10.1.dist-info → arize_phoenix-4.10.2rc1.dist-info}/RECORD

@@ -5,7 +5,7 @@ phoenix/exceptions.py,sha256=n2L2KKuecrdflB9MsCdAYCiSEvGJptIsfRkXMoJle7A,169
 phoenix/py.typed,sha256=AbpHGcgLb-kRsJGnwFEktk7uzpZOCcBY74-YBdrKVGs,1
 phoenix/services.py,sha256=aTxhcOA1pZHB6U-B3TEcp6fqDF5oT0xCUvEUNMZVTUQ,5175
 phoenix/settings.py,sha256=cO-qgis_S27nHirTobYI9hHPfZH18R--WMmxNdsVUwc,273
-phoenix/version.py,sha256=ogZ1wJ10Itlh9CdbwovCzGjxyYPJlQvG3lxFz6oBJ-c,23
+phoenix/version.py,sha256=Zr3eSPGuwbAWbpKhNj_q-3yeFI5ZZ6HVj7RbDS_7Y14,26
 phoenix/core/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 phoenix/core/embedding_dimension.py,sha256=zKGbcvwOXgLf-yrJBpQyKtd-LEOPRKHnUToyAU8Owis,87
 phoenix/core/model.py,sha256=km_a--PBHOuA337ClRw9xqhOHhrUT6Rl9pz_zV0JYkQ,4843
@@ -60,19 +60,19 @@ phoenix/pointcloud/pointcloud.py,sha256=4zAIkKs2xOUbchpj4XDAV-iPMXrfAJ15TG6rlIYG
 phoenix/pointcloud/projectors.py,sha256=zO_RrtDYSv2rqVOfIP2_9Cv11Dc8EmcZR94xhFcBYPU,1057
 phoenix/pointcloud/umap_parameters.py,sha256=3UQSjrysVOvq2V4KNpTMqNqNiK0BsTZnPBHWZ4fyJtQ,1708
 phoenix/server/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-phoenix/server/app.py,sha256=Ld9NvW7sBT4aJn8CXAB_PIw4c5wlkxCcjT_hzjtn7dM,18478
+phoenix/server/app.py,sha256=2-zGgV1zyeWtnkLLmgqgIU0o7QQ7eHqGTb2bv1ydLXY,17555
 phoenix/server/grpc_server.py,sha256=faktLxEtWGlCB1bPR4QwwTsRoQloahKMx0hAWqRGI5s,3379
 phoenix/server/main.py,sha256=GpjmstVb4o-xjg0u7t4z1WXRmdxWMOw_XhQ7tPc7Fd8,11113
 phoenix/server/prometheus.py,sha256=j9DHB2fERuq_ZKmwVaqR-9wx5WcPPuU1Cm5Bhg5241Y,2996
 phoenix/server/telemetry.py,sha256=T_2OKrxNViAeaANlNspEekg_Y5uZIFWvKAnpz8Aoqvk,2762
-phoenix/server/thread_server.py,sha256=dP6cm6Cf08jNhDA1TRlVZpziu1YgtPDmaeIJMm725eI,2154
+phoenix/server/thread_server.py,sha256=RwXQGP_QhGD7le6WB7xEygEEuwBl5Ck_Zo8xGIYGi9M,2135
 phoenix/server/api/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-phoenix/server/api/context.py,sha256=4jcy203Gtx38399FP21iU3HmFsq-50EKFJlX4IW2Los,2878
+phoenix/server/api/context.py,sha256=G9bmB3h-a4Top6wGAu91ymS-YE6D1DYQyWtPIlZnbgc,2811
 phoenix/server/api/interceptor.py,sha256=ykDnoC_apUd-llVli3m1CW18kNSIgjz2qZ6m5JmPDu8,1294
 phoenix/server/api/queries.py,sha256=eq2xHaQF-x4k6AGSY6b6mU2pie9bj-AJML6P2Mr0_DM,19886
 phoenix/server/api/schema.py,sha256=BcxdqO5CSGqpKd-AAJHMjFlzaK9oJA8GJuxmMfcdjn4,434
 phoenix/server/api/utils.py,sha256=Y1lGu8J8r8BSBX9OzffgewI8QMziovbG-ePDvZrrwGI,949
-phoenix/server/api/dataloaders/__init__.py,sha256=qehXL37vGdw7v5PFs3kbZVIuhuzrVNVeZACDQjYpwyo,4847
+phoenix/server/api/dataloaders/__init__.py,sha256=F6-8dwb-aQ_T8LeRYg3LPR7T94__8Xe4ysM0VGQeQUQ,4936
 phoenix/server/api/dataloaders/average_experiment_run_latency.py,sha256=RiO0AKC6Y5byafsV0zTJEIOt8Nudjte73f1T78cBe1k,1817
 phoenix/server/api/dataloaders/dataset_example_revisions.py,sha256=Vpr5IEKSR4QnAVxE5NM7u92fPNgeHQV2ieYc6JakCj0,3788
 phoenix/server/api/dataloaders/dataset_example_spans.py,sha256=_jLlo0KdUS65d4PNTtE9aXVyG_NZWgA7VcpNC9udQ8U,1484
@@ -88,6 +88,7 @@ phoenix/server/api/dataloaders/latency_ms_quantile.py,sha256=pEc7QjB2iiNOQm_Fmo9
 phoenix/server/api/dataloaders/min_start_or_max_end_times.py,sha256=IoFX5PtSpvQdMk_7-oB8TpIse3Q4PMxep4qKggkHpzo,2902
 phoenix/server/api/dataloaders/project_by_name.py,sha256=W4q-ddeVnja2DNwgg5l8mV2GNADNJf1CTXNcJaBWdfE,1165
 phoenix/server/api/dataloaders/record_counts.py,sha256=mp3KlhwFw-Iy7o6bFxtJKC6B5kGinPIh5PxxNkrxf8o,4283
+phoenix/server/api/dataloaders/span_annotations.py,sha256=pJQHF92b-CgBJFfZOf0siLsa-jHzr9ByfLiNGpb5oic,1262
 phoenix/server/api/dataloaders/span_descendants.py,sha256=djqXXwupWaXkFdscoy-iILYijuxlqr8hJcv6JawsV6s,2162
 phoenix/server/api/dataloaders/span_evaluations.py,sha256=quXGyj_OYvHrcWfzXlYIv7C1OCZiL1K7cWU-YEAlsNA,1316
 phoenix/server/api/dataloaders/span_projects.py,sha256=KV--NTLpdXyFUGYN69KHeVVHDBJR2oupo1cgLvqK6Ko,1304
@@ -133,19 +134,21 @@ phoenix/server/api/mutations/project_mutations.py,sha256=d_xtYkYfZ5flpVgEkGknKB8
 phoenix/server/api/mutations/span_annotations_mutations.py,sha256=Pfaq4y-FGskPit4z_9GvsyWeBwK1g3CDi2UhGKxyjFE,4973
 phoenix/server/api/mutations/trace_annotations_mutations.py,sha256=4xm-zg8PEjwjfVRsWzBD_iiOS94JI6UYtK3fV2dtA4M,5013
 phoenix/server/api/openapi/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-phoenix/server/api/openapi/main.py,sha256=WY0pj3B7siQyyYqKyhqnzWC7P8MtEtiukOBUjGwLXfw,153
-phoenix/server/api/openapi/schema.py,sha256=uuSYe1Ecu72aXRgTNjyMu-9ZPE13DAHJPKtedS-MsSs,451
+phoenix/server/api/openapi/main.py,sha256=KNutA_7AvV_WlGX8cOkvvDujcJKQ7AD1HT6rTpCpR8A,616
+phoenix/server/api/openapi/schema.py,sha256=oVZoflWMfzOrLKMIrjr3iLnJ13rmN-t_DOe9g6KoN5s,471
 phoenix/server/api/routers/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 phoenix/server/api/routers/utils.py,sha256=M41BoH-fl37izhRuN2aX7lWm7jOC20A_3uClv9TVUUY,583
-phoenix/server/api/routers/v1/__init__.py,sha256=D1EFRWG4PcsTubeF3A3ENlNatCRq26AA52FxW11BGjM,3048
-phoenix/server/api/routers/v1/dataset_examples.py,sha256=XfqOvDKF1oxb0pkeYfBycwwGt3LnSyyGdMLKC5VKoGQ,6690
-phoenix/server/api/routers/v1/datasets.py,sha256=r0WcNxF8SKVa3-4rrTIg4Andwr4NmRmW1ybpKuxR9qw,33639
-phoenix/server/api/routers/v1/evaluations.py,sha256=8g6P_e2BweV3RDU0esFmpkb0L5fCwonQPXiJ0y6HLwg,9126
-phoenix/server/api/routers/v1/experiment_evaluations.py,sha256=TE1GMSOLN_96uAaJpnRpIH2u9x6_ebtkECgZRHvqt-w,5098
-phoenix/server/api/routers/v1/experiment_runs.py,sha256=jy4SynmzdtQMoUzlowmG6wsVU14SsLAzfcW4JOhXjeQ,8154
-phoenix/server/api/routers/v1/experiments.py,sha256=uVdmhyJgYI-UqOiRSJ-8OcVpL8a6Z02B5H2Rt_7yboY,11829
-phoenix/server/api/routers/v1/spans.py,sha256=tryWFoJVFRLALzt6dfPmBBhKMS0s3hhlYdTathxVEU4,9638
-phoenix/server/api/routers/v1/traces.py,sha256=PBIrpdJHVJ9gyiukCy1Ck1w0xts0VEHtRKaF7Noa248,8434
+phoenix/server/api/routers/v1/__init__.py,sha256=nb49zcOdAi3DSGuC9gUubN9Yri-o7-WFdlGak4jGuFw,1462
+phoenix/server/api/routers/v1/dataset_examples.py,sha256=HSPc4PTnGRqyHSXP7qAZtZ39OG3SVk7btU0HmJyrL14,5595
+phoenix/server/api/routers/v1/datasets.py,sha256=B2k8IiIkmXEPTdBruGlOUZQmgjj2_G2Bet0_h-X-yZY,32189
+phoenix/server/api/routers/v1/evaluations.py,sha256=M1PKCNL_5xB9EXY4EOnrOu9uDoeJan6V1tRkSbRJEgA,9673
+phoenix/server/api/routers/v1/experiment_evaluations.py,sha256=Wk6YnjvA3zSh8N0lRxS8V5uf17tLsNzB12LrNDry1I0,4622
+phoenix/server/api/routers/v1/experiment_runs.py,sha256=CJtTWSJni0xucDUQ-ouifcAfQFTD2xf-y7B_Xy7S6dA,6007
+phoenix/server/api/routers/v1/experiments.py,sha256=oEPCoH_OKnIn04iBYYQcyx9N1LXnZGF0BEHwdO7XbOw,9760
+phoenix/server/api/routers/v1/spans.py,sha256=ZGZ67Q5pmKTr2cPqc2Lf5fONmbp-Nr_A-pxLswATk_M,8724
+phoenix/server/api/routers/v1/traces.py,sha256=Wm6bVm-MpubOxZwLwS9Xsn93RvDlSwlNYSprXnoIVQo,7921
+phoenix/server/api/routers/v1/utils.py,sha256=R2WxEpTyi4Z94K8tyt-a4dMYGh4jP0a_lbJ33FG_dKc,3045
+phoenix/server/api/types/Annotation.py,sha256=7Ym7iuVcbwHlw2yIRylz4nATAF_Cm-Z17qcjiooj1cc,751
 phoenix/server/api/types/AnnotatorKind.py,sha256=rPgGdbN1Gvc109sGQ_ZH-gfJbp93V9wlarzTEJNtUwI,236
 phoenix/server/api/types/Cluster.py,sha256=ac4YfT1OH3xLVmex7EUmB6b9IpULnhLTt554LR0jglE,5689
 phoenix/server/api/types/CreateDatasetPayload.py,sha256=R-6zCmuD0f76RU9Giu78xwTHlASQs6Aq8yzvX1Kxc3g,140
@@ -164,7 +167,7 @@ phoenix/server/api/types/DocumentEvaluationSummary.py,sha256=EpgKn0uiH2Vbpyltc1M
 phoenix/server/api/types/DocumentRetrievalMetrics.py,sha256=amkpC3H5IU5-9GvO0telpbq00m6lIcv_2v446OpwFwc,1822
 phoenix/server/api/types/EmbeddingDimension.py,sha256=2OTVzAcfY-2xnJLl7122EAaqJyp0sob1v9-ryaNWknw,18961
 phoenix/server/api/types/EmbeddingMetadata.py,sha256=fJvNNYCbkf3SJalArLy9rcBq9Uj1SNac60zjqe1PFnM,461
-phoenix/server/api/types/Evaluation.py,sha256=pSoaJH8-WzPvmZKn5fButzkXSN3Rvuq9GJ7g9H1aaaU,3997
+phoenix/server/api/types/Evaluation.py,sha256=oW9UfUCs6Td60kb5s6_sJNvUkxk5TQXjG0pmQ0-gT4M,3298
 phoenix/server/api/types/EvaluationSummary.py,sha256=EFucuzAhcxR9sdEn6WNAtmAGJk-QR_3j3y7Mn2A9_bU,1528
 phoenix/server/api/types/Event.py,sha256=XdYgaIxcVIW-YFViCkxj5l9OaVNepyIrCtm5Iqg2le8,3989
 phoenix/server/api/types/EventMetadata.py,sha256=-J0tYF9eZTHwCjwxQHY7Gckr2_MNW5OoWT1mydweZNM,635
@@ -188,8 +191,8 @@ phoenix/server/api/types/Retrieval.py,sha256=OhMK2ncjoyp5h1yjKhjlKpoTbQrMHuxmgSF
 phoenix/server/api/types/ScalarDriftMetricEnum.py,sha256=IUAcRPpgL41WdoIgK6cNk2Te38SspXGyEs-S1fY23_A,232
 phoenix/server/api/types/Segments.py,sha256=m2yoegrxA1Tn7ZAy1rMjjD1isc752MaAXMoffkBlvrM,2921
 phoenix/server/api/types/SortDir.py,sha256=OUpXhlCzCxPoXSDkJJygEs9Rw9pMymfaZUG5zPTrw4Y,152
-phoenix/server/api/types/Span.py,sha256=W4Rsg85bgqbDhgYwpjgOTrIQKbkwpFQPpL6nqMyzhCs,13865
-phoenix/server/api/types/SpanAnnotation.py,sha256=n9pehpY-NqJD9KvK03UJYlHLKdies7mxrqdr9i1c5oE,1241
+phoenix/server/api/types/Span.py,sha256=VjFC5VXe3Incf9AhEvctYhRpUixhF82oN2iepxjU50Q,14180
+phoenix/server/api/types/SpanAnnotation.py,sha256=6b5G-b_OoRvDL2ayWk7MkbqarLK-F-pQMx21CpUuNGY,1168
 phoenix/server/api/types/TimeSeries.py,sha256=wjzuxHFqCey0O7Ys25qiXyuqXK8an-osyNWUE8A_8G4,5227
 phoenix/server/api/types/Trace.py,sha256=ep-mPexub1ijxAnBvc2KrGsNVXO2SfDR1WxqER2wcD8,2376
 phoenix/server/api/types/TraceAnnotation.py,sha256=OW6A2zr1gomOuG0XQe55dk15XXX2DSM0DzatRbHWH5A,1256
@@ -200,7 +203,6 @@ phoenix/server/api/types/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZ
 phoenix/server/api/types/node.py,sha256=V0Xh9U4cGkz3iMg-vzEXtcs6cumU29JFPiU-JuGzjWI,848
 phoenix/server/api/types/pagination.py,sha256=PcaJ0s4exsTKgCZC4aFm1cgZNrGpHSdo6PbkWzPcweg,9077
 phoenix/server/openapi/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-phoenix/server/openapi/docs.py,sha256=fTb9q2oOSKC5bLVQy2Dsg3Bs0mGkCOKX1ypX7731sE0,7044
 phoenix/server/static/apple-touch-icon-114x114.png,sha256=xtFVXAYQnJkpUApg2D1hltSTuyO4Is4sD4A0ZkikiVU,9486
 phoenix/server/static/apple-touch-icon-120x120.png,sha256=iqZVAk634BbjJMozA8aHYyw15JjhIlIrG41FA2DFFaE,9957
 phoenix/server/static/apple-touch-icon-144x144.png,sha256=VgARtkHKoU8zikb3_G83h_cb02kpPcoJqO78yRh1AfU,10047
@@ -256,8 +258,8 @@ phoenix/utilities/logging.py,sha256=lDXd6EGaamBNcQxL4vP1au9-i_SXe0OraUDiJOcszSw,
 phoenix/utilities/project.py,sha256=qWsvKnG1oKhOFUowXf9qiOL2ia7jaFe_ijFFHEt8GJo,431
 phoenix/utilities/re.py,sha256=PDve_OLjRTM8yQQJHC8-n3HdIONi7aNils3ZKRZ5uBM,2045
 phoenix/utilities/span_store.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-arize_phoenix-4.10.1.dist-info/METADATA,sha256=0m2W26GT35O9_Mm0_WUpcvE7Vr9oOq-p2DtR2v6dmkI,11452
-arize_phoenix-4.10.1.dist-info/WHEEL,sha256=1yFddiXMmvYK7QYTqtRNtX66WJ0Mz8PYEiEUoOUUxRY,87
-arize_phoenix-4.10.1.dist-info/licenses/IP_NOTICE,sha256=JBqyyCYYxGDfzQ0TtsQgjts41IJoa-hiwDrBjCb9gHM,469
-arize_phoenix-4.10.1.dist-info/licenses/LICENSE,sha256=HFkW9REuMOkvKRACuwLPT0hRydHb3zNg-fdFt94td18,3794
-arize_phoenix-4.10.1.dist-info/RECORD,,
+arize_phoenix-4.10.2rc1.dist-info/METADATA,sha256=kvMhdHvg67eriHuvnz3zfwCN-P9ISpjC_3LrYshITeY,11486
+arize_phoenix-4.10.2rc1.dist-info/WHEEL,sha256=1yFddiXMmvYK7QYTqtRNtX66WJ0Mz8PYEiEUoOUUxRY,87
+arize_phoenix-4.10.2rc1.dist-info/licenses/IP_NOTICE,sha256=JBqyyCYYxGDfzQ0TtsQgjts41IJoa-hiwDrBjCb9gHM,469
+arize_phoenix-4.10.2rc1.dist-info/licenses/LICENSE,sha256=HFkW9REuMOkvKRACuwLPT0hRydHb3zNg-fdFt94td18,3794
+arize_phoenix-4.10.2rc1.dist-info/RECORD,,

phoenix/server/api/context.py

@@ -1,12 +1,10 @@
 from dataclasses import dataclass
 from datetime import datetime
 from pathlib import Path
-from typing import AsyncContextManager, Callable, Optional, Union
+from typing import AsyncContextManager, Callable, Optional
 
 from sqlalchemy.ext.asyncio import AsyncSession
-from starlette.requests import Request
-from starlette.responses import Response
-from starlette.websockets import WebSocket
+from strawberry.fastapi import BaseContext
 from typing_extensions import TypeAlias
 
 from phoenix.core.model_schema import Model
@@ -27,6 +25,7 @@ from phoenix.server.api.dataloaders import (
     MinStartOrMaxEndTimeDataLoader,
     ProjectByNameDataLoader,
     RecordCountDataLoader,
+    SpanAnnotationsDataLoader,
     SpanDescendantsDataLoader,
     SpanEvaluationsDataLoader,
     SpanProjectsDataLoader,
@@ -59,15 +58,14 @@ class DataLoaders:
     trace_evaluations: TraceEvaluationsDataLoader
     trace_row_ids: TraceRowIdsDataLoader
     project_by_name: ProjectByNameDataLoader
+    span_annotations: SpanAnnotationsDataLoader
 
 
 ProjectRowId: TypeAlias = int
 
 
 @dataclass
-class Context:
-    request: Union[Request, WebSocket]
-    response: Optional[Response]
+class Context(BaseContext):
     db: Callable[[], AsyncContextManager[AsyncSession]]
     data_loaders: DataLoaders
     cache_for_dataloaders: Optional[CacheForDataLoaders]

phoenix/server/api/dataloaders/__init__.py

@@ -26,6 +26,7 @@ from .latency_ms_quantile import LatencyMsQuantileCache, LatencyMsQuantileDataLo
 from .min_start_or_max_end_times import MinStartOrMaxEndTimeCache, MinStartOrMaxEndTimeDataLoader
 from .project_by_name import ProjectByNameDataLoader
 from .record_counts import RecordCountCache, RecordCountDataLoader
+from .span_annotations import SpanAnnotationsDataLoader
 from .span_descendants import SpanDescendantsDataLoader
 from .span_evaluations import SpanEvaluationsDataLoader
 from .span_projects import SpanProjectsDataLoader
@@ -56,6 +57,7 @@ __all__ = [
     "TraceEvaluationsDataLoader",
     "TraceRowIdsDataLoader",
     "ProjectByNameDataLoader",
+    "SpanAnnotationsDataLoader",
 ]
 
 

phoenix/server/api/dataloaders/span_annotations.py

@@ -0,0 +1,35 @@
+from collections import defaultdict
+from typing import (
+    AsyncContextManager,
+    Callable,
+    DefaultDict,
+    List,
+)
+
+from sqlalchemy import select
+from sqlalchemy.ext.asyncio import AsyncSession
+from strawberry.dataloader import DataLoader
+from typing_extensions import TypeAlias
+
+from phoenix.db import models
+from phoenix.server.api.types.SpanAnnotation import SpanAnnotation, to_gql_span_annotation
+
+Key: TypeAlias = int
+Result: TypeAlias = List[SpanAnnotation]
+
+
+class SpanAnnotationsDataLoader(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]:
+        span_annotations_by_id: DefaultDict[Key, Result] = defaultdict(list)
+        msa = models.SpanAnnotation
+        async with self._db() as session:
+            data = await session.stream_scalars(select(msa).where(msa.span_rowid.in_(keys)))
+            async for span_annotation in data:
+                span_annotations_by_id[span_annotation.span_rowid].append(
+                    to_gql_span_annotation(span_annotation)
+                )
+        return [span_annotations_by_id[key] for key in keys]

phoenix/server/api/openapi/main.py

@@ -1,6 +1,22 @@
+import json
+from argparse import ArgumentParser
+from typing import Optional, Tuple
+
 from .schema import get_openapi_schema
 
 if __name__ == "__main__":
-    import yaml  # type: ignore
+    parser = ArgumentParser()
+    parser.add_argument(
+        "--compress",
+        action="store_true",
+        help="Whether to output a compressed version of the OpenAPI schema",
+    )
+    args = parser.parse_args()
 
-    print(yaml.dump(get_openapi_schema(), indent=2))
+    indent: Optional[int] = None
+    separator: Optional[Tuple[str, str]] = None
+    if args.compress:
+        separator = (",", ":")
+    else:
+        indent = 2
+    print(json.dumps(get_openapi_schema(), indent=indent, separators=separator))

phoenix/server/api/openapi/schema.py

@@ -1,16 +1,16 @@
-from typing import Any
+from typing import Any, Dict
 
-from starlette.schemas import SchemaGenerator
+from fastapi.openapi.utils import get_openapi
 
-from phoenix.server.api.routers.v1 import V1_ROUTES
+from phoenix.server.api.routers.v1 import REST_API_VERSION
+from phoenix.server.api.routers.v1 import router as v1_router
 
-OPENAPI_SCHEMA_GENERATOR = SchemaGenerator(
-    {"openapi": "3.0.0", "info": {"title": "Arize-Phoenix API", "version": "1.0"}}
-)
 
-
-def get_openapi_schema() -> Any:
-    """
-    Exports an OpenAPI schema for the Phoenix REST API as a JSON object.
-    """
-    return OPENAPI_SCHEMA_GENERATOR.get_schema(V1_ROUTES)  # type: ignore
+def get_openapi_schema() -> Dict[str, Any]:
+    return get_openapi(
+        title="Arize-Phoenix REST API",
+        version=REST_API_VERSION,
+        openapi_version="3.1.0",
+        description="Schema for Arize-Phoenix REST API",
+        routes=v1_router.routes,
+    )

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

@@ -1,89 +1,42 @@
-from typing import Any, Awaitable, Callable, Mapping, Tuple
-
-import wrapt
-from starlette import routing
-from starlette.requests import Request
-from starlette.responses import Response
+from fastapi import APIRouter, Depends, HTTPException, Request
 from starlette.status import HTTP_403_FORBIDDEN
 
-from . import (
-    datasets,
-    evaluations,
-    experiment_evaluations,
-    experiment_runs,
-    experiments,
-    spans,
-    traces,
-)
-from .dataset_examples import list_dataset_examples
-
+from .datasets import router as datasets_router
+from .evaluations import router as evaluations_router
+from .experiment_evaluations import router as experiment_evaluations_router
+from .experiment_runs import router as experiment_runs_router
+from .experiments import router as experiments_router
+from .spans import router as spans_router
+from .traces import router as traces_router
+from .utils import add_errors_to_responses
 
-@wrapt.decorator  # type: ignore
-async def forbid_if_readonly(
-    wrapped: Callable[[Request], Awaitable[Response]],
-    _: Any,
-    args: Tuple[Request],
-    kwargs: Mapping[str, Any],
-) -> Response:
-    request, *_ = args
-    if request.app.state.read_only:
-        return Response(status_code=HTTP_403_FORBIDDEN)
-    return await wrapped(*args, **kwargs)
+REST_API_VERSION = "1.0"
 
 
-class Route(routing.Route):
-    def __init__(self, path: str, endpoint: Callable[..., Any], **kwargs: Any) -> None:
-        super().__init__(path, forbid_if_readonly(endpoint), **kwargs)
-
-
-V1_ROUTES = [
-    Route("/v1/evaluations", evaluations.post_evaluations, methods=["POST"]),
-    Route("/v1/evaluations", evaluations.get_evaluations, methods=["GET"]),
-    Route("/v1/traces", traces.post_traces, methods=["POST"]),
-    Route("/v1/trace_annotations", traces.annotate_traces, methods=["POST"]),
-    Route("/v1/spans", spans.query_spans_handler, methods=["POST"]),
-    Route("/v1/spans", spans.get_spans_handler, methods=["GET"]),
-    Route("/v1/span_annotations", spans.annotate_spans, methods=["POST"]),
-    Route("/v1/datasets/upload", datasets.post_datasets_upload, methods=["POST"]),
-    Route("/v1/datasets", datasets.list_datasets, methods=["GET"]),
-    Route("/v1/datasets/{id:str}", datasets.delete_dataset_by_id, methods=["DELETE"]),
-    Route("/v1/datasets/{id:str}", datasets.get_dataset_by_id, methods=["GET"]),
-    Route("/v1/datasets/{id:str}/csv", datasets.get_dataset_csv, methods=["GET"]),
-    Route(
-        "/v1/datasets/{id:str}/jsonl/openai_ft",
-        datasets.get_dataset_jsonl_openai_ft,
-        methods=["GET"],
-    ),
-    Route(
-        "/v1/datasets/{id:str}/jsonl/openai_evals",
-        datasets.get_dataset_jsonl_openai_evals,
-        methods=["GET"],
-    ),
-    Route("/v1/datasets/{id:str}/examples", list_dataset_examples, methods=["GET"]),
-    Route("/v1/datasets/{id:str}/versions", datasets.get_dataset_versions, methods=["GET"]),
-    Route(
-        "/v1/datasets/{dataset_id:str}/experiments",
-        experiments.create_experiment,
-        methods=["POST"],
-    ),
-    Route(
-        "/v1/experiments/{experiment_id:str}",
-        experiments.read_experiment,
-        methods=["GET"],
-    ),
-    Route(
-        "/v1/experiments/{experiment_id:str}/runs",
-        experiment_runs.create_experiment_run,
-        methods=["POST"],
-    ),
-    Route(
-        "/v1/experiments/{experiment_id:str}/runs",
-        experiment_runs.list_experiment_runs,
-        methods=["GET"],
-    ),
-    Route(
-        "/v1/experiment_evaluations",
-        experiment_evaluations.upsert_experiment_evaluation,
-        methods=["POST"],
+async def prevent_access_in_read_only_mode(request: Request) -> None:
+    """
+    Prevents access to the REST API in read-only mode.
+    """
+    if request.app.state.read_only:
+        raise HTTPException(
+            detail="The Phoenix REST API is disabled in read-only mode.",
+            status_code=HTTP_403_FORBIDDEN,
+        )
+
+
+router = APIRouter(
+    prefix="/v1",
+    dependencies=[Depends(prevent_access_in_read_only_mode)],
+    responses=add_errors_to_responses(
+        [
+            HTTP_403_FORBIDDEN  # adds a 403 response to each route in the generated OpenAPI schema
+        ]
     ),
-]
+)
+router.include_router(datasets_router)
+router.include_router(experiments_router)
+router.include_router(experiment_runs_router)
+router.include_router(experiment_evaluations_router)
+router.include_router(traces_router)
+router.include_router(spans_router)
+router.include_router(evaluations_router)