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

phoenix/server/app.py

@@ -1,3 +1,4 @@
+import asyncio
 import contextlib
 import json
 import logging
@@ -21,25 +22,24 @@ from typing import (
 )
 
 import strawberry
+from fastapi import APIRouter, FastAPI
+from fastapi.middleware.gzip import GZipMiddleware
+from fastapi.responses import FileResponse
+from fastapi.utils import is_body_allowed_for_status_code
 from sqlalchemy.ext.asyncio import (
     AsyncEngine,
     AsyncSession,
     async_sessionmaker,
 )
-from starlette.applications import Starlette
-from starlette.datastructures import QueryParams
-from starlette.endpoints import HTTPEndpoint
 from starlette.exceptions import HTTPException
 from starlette.middleware import Middleware
 from starlette.middleware.base import BaseHTTPMiddleware, RequestResponseEndpoint
 from starlette.requests import Request
-from starlette.responses import FileResponse, PlainTextResponse, Response
-from starlette.routing import Mount, Route
+from starlette.responses import PlainTextResponse, Response
 from starlette.staticfiles import StaticFiles
 from starlette.templating import Jinja2Templates
 from starlette.types import Scope, StatefulLifespan
-from starlette.websockets import WebSocket
-from strawberry.asgi import GraphQL
+from strawberry.fastapi import GraphQLRouter
 from strawberry.schema import BaseSchema
 from typing_extensions import TypeAlias
 
@@ -75,6 +75,7 @@ from phoenix.server.api.dataloaders import (
     ProjectByNameDataLoader,
     RecordCountDataLoader,
     SpanAnnotationsDataLoader,
+    SpanDatasetExamplesDataLoader,
     SpanDescendantsDataLoader,
     SpanEvaluationsDataLoader,
     SpanProjectsDataLoader,
@@ -82,19 +83,22 @@ from phoenix.server.api.dataloaders import (
     TraceEvaluationsDataLoader,
     TraceRowIdsDataLoader,
 )
-from phoenix.server.api.openapi.schema import OPENAPI_SCHEMA_GENERATOR
-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
 from phoenix.server.api.schema import schema
 from phoenix.server.grpc_server import GrpcServer
-from phoenix.server.openapi.docs import get_swagger_ui_html
 from phoenix.server.telemetry import initialize_opentelemetry_tracer_provider
+from phoenix.server.types import DbSessionFactory
 from phoenix.trace.schemas import Span
+from phoenix.utilities.client import PHOENIX_SERVER_VERSION_HEADER
 
 if TYPE_CHECKING:
     from opentelemetry.trace import TracerProvider
 
 logger = logging.getLogger(__name__)
 
+router = APIRouter(include_in_schema=False)
+
 templates = Jinja2Templates(directory=SERVER_DIR / "templates")
 
 
@@ -167,125 +171,35 @@ class HeadersMiddleware(BaseHTTPMiddleware):
         request: Request,
         call_next: RequestResponseEndpoint,
     ) -> Response:
+        from phoenix import __version__ as phoenix_version
+
         response = await call_next(request)
         response.headers["x-colab-notebook-cache-control"] = "no-cache"
+        response.headers[PHOENIX_SERVER_VERSION_HEADER] = phoenix_version
         return response
 
 
 ProjectRowId: TypeAlias = int
 
 
-class GraphQLWithContext(GraphQL):  # type: ignore
-    def __init__(
-        self,
-        schema: BaseSchema,
-        db: Callable[[], AsyncContextManager[AsyncSession]],
-        model: Model,
-        export_path: Path,
-        graphiql: bool = False,
-        corpus: Optional[Model] = None,
-        streaming_last_updated_at: Callable[[ProjectRowId], Optional[datetime]] = lambda _: None,
-        cache_for_dataloaders: Optional[CacheForDataLoaders] = None,
-        read_only: bool = False,
-    ) -> None:
-        self.db = db
-        self.model = model
-        self.corpus = corpus
-        self.export_path = export_path
-        self.streaming_last_updated_at = streaming_last_updated_at
-        self.cache_for_dataloaders = cache_for_dataloaders
-        self.read_only = read_only
-        super().__init__(schema, graphiql=graphiql)
-
-    async def get_context(
-        self,
-        request: Union[Request, WebSocket],
-        response: Optional[Response] = None,
-    ) -> Context:
-        return Context(
-            request=request,
-            response=response,
-            db=self.db,
-            model=self.model,
-            corpus=self.corpus,
-            export_path=self.export_path,
-            streaming_last_updated_at=self.streaming_last_updated_at,
-            data_loaders=DataLoaders(
-                average_experiment_run_latency=AverageExperimentRunLatencyDataLoader(self.db),
-                dataset_example_revisions=DatasetExampleRevisionsDataLoader(self.db),
-                dataset_example_spans=DatasetExampleSpansDataLoader(self.db),
-                document_evaluation_summaries=DocumentEvaluationSummaryDataLoader(
-                    self.db,
-                    cache_map=self.cache_for_dataloaders.document_evaluation_summary
-                    if self.cache_for_dataloaders
-                    else None,
-                ),
-                document_evaluations=DocumentEvaluationsDataLoader(self.db),
-                document_retrieval_metrics=DocumentRetrievalMetricsDataLoader(self.db),
-                evaluation_summaries=EvaluationSummaryDataLoader(
-                    self.db,
-                    cache_map=self.cache_for_dataloaders.evaluation_summary
-                    if self.cache_for_dataloaders
-                    else None,
-                ),
-                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,
-                    cache_map=self.cache_for_dataloaders.latency_ms_quantile
-                    if self.cache_for_dataloaders
-                    else None,
-                ),
-                min_start_or_max_end_times=MinStartOrMaxEndTimeDataLoader(
-                    self.db,
-                    cache_map=self.cache_for_dataloaders.min_start_or_max_end_time
-                    if self.cache_for_dataloaders
-                    else None,
-                ),
-                record_counts=RecordCountDataLoader(
-                    self.db,
-                    cache_map=self.cache_for_dataloaders.record_count
-                    if self.cache_for_dataloaders
-                    else None,
-                ),
-                span_descendants=SpanDescendantsDataLoader(self.db),
-                span_evaluations=SpanEvaluationsDataLoader(self.db),
-                span_projects=SpanProjectsDataLoader(self.db),
-                token_counts=TokenCountDataLoader(
-                    self.db,
-                    cache_map=self.cache_for_dataloaders.token_count
-                    if self.cache_for_dataloaders
-                    else None,
-                ),
-                trace_evaluations=TraceEvaluationsDataLoader(self.db),
-                trace_row_ids=TraceRowIdsDataLoader(self.db),
-                project_by_name=ProjectByNameDataLoader(self.db),
-                span_annotations=SpanAnnotationsDataLoader(self.db),
-            ),
-            cache_for_dataloaders=self.cache_for_dataloaders,
-            read_only=self.read_only,
-        )
+@router.get("/exports")
+async def download_exported_file(request: Request, filename: str) -> FileResponse:
+    file = request.app.state.export_path / (filename + ".parquet")
+    if not file.is_file():
+        raise HTTPException(status_code=404)
+    return FileResponse(
+        path=file,
+        filename=file.name,
+        media_type="application/x-octet-stream",
+    )
 
 
-class Download(HTTPEndpoint):
-    path: Path
+@router.get("/arize_phoenix_version")
+async def version() -> PlainTextResponse:
+    return PlainTextResponse(f"{phoenix.__version__}")
 
-    async def get(self, request: Request) -> FileResponse:
-        params = QueryParams(request.query_params)
-        file = self.path / (params.get("filename", "") + ".parquet")
-        if not file.is_file():
-            raise HTTPException(status_code=404)
-        return FileResponse(
-            path=file,
-            filename=file.name,
-            media_type="application/x-octet-stream",
-        )
 
-
-async def version(_: Request) -> PlainTextResponse:
-    return PlainTextResponse(f"{phoenix.__version__}")
+DB_MUTEX: Optional[asyncio.Lock] = None
 
 
 def _db(engine: AsyncEngine) -> Callable[[], AsyncContextManager[AsyncSession]]:
@@ -293,23 +207,29 @@ def _db(engine: AsyncEngine) -> Callable[[], AsyncContextManager[AsyncSession]]:
 
     @contextlib.asynccontextmanager
     async def factory() -> AsyncIterator[AsyncSession]:
-        async with Session.begin() as session:
-            yield session
+        async with contextlib.AsyncExitStack() as stack:
+            if DB_MUTEX:
+                await stack.enter_async_context(DB_MUTEX)
+            yield await stack.enter_async_context(Session.begin())
 
     return factory
 
 
 def _lifespan(
     *,
+    dialect: SupportedSQLDialect,
     bulk_inserter: BulkInserter,
     tracer_provider: Optional["TracerProvider"] = None,
     enable_prometheus: bool = False,
     clean_ups: Iterable[Callable[[], None]] = (),
     read_only: bool = False,
-) -> StatefulLifespan[Starlette]:
+) -> StatefulLifespan[FastAPI]:
     @contextlib.asynccontextmanager
-    async def lifespan(_: Starlette) -> AsyncIterator[Dict[str, Any]]:
+    async def lifespan(_: FastAPI) -> AsyncIterator[Dict[str, Any]]:
+        global DB_MUTEX
+        DB_MUTEX = asyncio.Lock() if dialect is SupportedSQLDialect.SQLITE else None
         async with bulk_inserter as (
+            enqueue,
             queue_span,
             queue_evaluation,
             enqueue_operation,
@@ -320,6 +240,7 @@ def _lifespan(
             enable_prometheus=enable_prometheus,
         ):
             yield {
+                "enqueue": enqueue,
                 "queue_span_for_bulk_insert": queue_span,
                 "queue_evaluation_for_bulk_insert": queue_evaluation,
                 "enqueue_operation": enqueue_operation,
@@ -330,29 +251,91 @@ def _lifespan(
     return lifespan
 
 
+@router.get("/healthz")
 async def check_healthz(_: Request) -> PlainTextResponse:
     return PlainTextResponse("OK")
 
 
-async def openapi_schema(request: Request) -> Response:
-    return OPENAPI_SCHEMA_GENERATOR.OpenAPIResponse(request=request)
-
-
-async def api_docs(request: Request) -> Response:
-    return get_swagger_ui_html(openapi_url="/schema", title="arize-phoenix API")
-
-
-class SessionFactory:
-    def __init__(
-        self,
-        session_factory: Callable[[], AsyncContextManager[AsyncSession]],
-        dialect: str,
-    ):
-        self.session_factory = session_factory
-        self.dialect = SupportedSQLDialect(dialect)
+def create_graphql_router(
+    *,
+    schema: BaseSchema,
+    db: DbSessionFactory,
+    model: Model,
+    export_path: Path,
+    corpus: Optional[Model] = None,
+    streaming_last_updated_at: Callable[[ProjectRowId], Optional[datetime]] = lambda _: None,
+    cache_for_dataloaders: Optional[CacheForDataLoaders] = None,
+    read_only: bool = False,
+) -> GraphQLRouter:  # type: ignore[type-arg]
+    def get_context() -> Context:
+        return Context(
+            db=db,
+            model=model,
+            corpus=corpus,
+            export_path=export_path,
+            streaming_last_updated_at=streaming_last_updated_at,
+            data_loaders=DataLoaders(
+                average_experiment_run_latency=AverageExperimentRunLatencyDataLoader(db),
+                dataset_example_revisions=DatasetExampleRevisionsDataLoader(db),
+                dataset_example_spans=DatasetExampleSpansDataLoader(db),
+                document_evaluation_summaries=DocumentEvaluationSummaryDataLoader(
+                    db,
+                    cache_map=cache_for_dataloaders.document_evaluation_summary
+                    if cache_for_dataloaders
+                    else None,
+                ),
+                document_evaluations=DocumentEvaluationsDataLoader(db),
+                document_retrieval_metrics=DocumentRetrievalMetricsDataLoader(db),
+                evaluation_summaries=EvaluationSummaryDataLoader(
+                    db,
+                    cache_map=cache_for_dataloaders.evaluation_summary
+                    if cache_for_dataloaders
+                    else None,
+                ),
+                experiment_annotation_summaries=ExperimentAnnotationSummaryDataLoader(db),
+                experiment_error_rates=ExperimentErrorRatesDataLoader(db),
+                experiment_run_counts=ExperimentRunCountsDataLoader(db),
+                experiment_sequence_number=ExperimentSequenceNumberDataLoader(db),
+                latency_ms_quantile=LatencyMsQuantileDataLoader(
+                    db,
+                    cache_map=cache_for_dataloaders.latency_ms_quantile
+                    if cache_for_dataloaders
+                    else None,
+                ),
+                min_start_or_max_end_times=MinStartOrMaxEndTimeDataLoader(
+                    db,
+                    cache_map=cache_for_dataloaders.min_start_or_max_end_time
+                    if cache_for_dataloaders
+                    else None,
+                ),
+                record_counts=RecordCountDataLoader(
+                    db,
+                    cache_map=cache_for_dataloaders.record_count if cache_for_dataloaders else None,
+                ),
+                span_annotations=SpanAnnotationsDataLoader(db),
+                span_dataset_examples=SpanDatasetExamplesDataLoader(db),
+                span_descendants=SpanDescendantsDataLoader(db),
+                span_evaluations=SpanEvaluationsDataLoader(db),
+                span_projects=SpanProjectsDataLoader(db),
+                token_counts=TokenCountDataLoader(
+                    db,
+                    cache_map=cache_for_dataloaders.token_count if cache_for_dataloaders else None,
+                ),
+                trace_evaluations=TraceEvaluationsDataLoader(db),
+                trace_row_ids=TraceRowIdsDataLoader(db),
+                project_by_name=ProjectByNameDataLoader(db),
+            ),
+            cache_for_dataloaders=cache_for_dataloaders,
+            read_only=read_only,
+        )
 
-    def __call__(self) -> AsyncContextManager[AsyncSession]:
-        return self.session_factory()
+    return GraphQLRouter(
+        schema,
+        graphiql=True,
+        context_getter=get_context,
+        include_in_schema=False,
+        prefix="/graphql",
+    )
 
 
 def create_engine_and_run_migrations(
@@ -390,8 +373,20 @@ def instrument_engine_if_enabled(engine: AsyncEngine) -> List[Callable[[], None]
     return instrumentation_cleanups
 
 
+async def plain_text_http_exception_handler(request: Request, exc: HTTPException) -> Response:
+    """
+    Overrides the default handler for HTTPExceptions to return a plain text
+    response instead of a JSON response. For the original source code, see
+    https://github.com/tiangolo/fastapi/blob/d3cdd3bbd14109f3b268df7ca496e24bb64593aa/fastapi/exception_handlers.py#L11
+    """
+    headers = getattr(exc, "headers", None)
+    if not is_body_allowed_for_status_code(exc.status_code):
+        return Response(status_code=exc.status_code, headers=headers)
+    return PlainTextResponse(str(exc.detail), status_code=exc.status_code, headers=headers)
+
+
 def create_app(
-    db: SessionFactory,
+    db: DbSessionFactory,
     export_path: Path,
     model: Model,
     umap_params: UMAPParameters,
@@ -404,7 +399,7 @@ def create_app(
     initial_evaluations: Optional[Iterable[pb.Evaluation]] = None,
     serve_ui: bool = True,
     clean_up_callbacks: List[Callable[[], None]] = [],
-) -> Starlette:
+) -> FastAPI:
     clean_ups: List[Callable[[], None]] = clean_up_callbacks  # To be called at app shutdown.
     initial_batch_of_spans: Iterable[Tuple[Span, str]] = (
         ()
@@ -447,7 +442,7 @@ def create_app(
 
         strawberry_extensions.append(_OpenTelemetryExtension)
 
-    graphql = GraphQLWithContext(
+    graphql_router = create_graphql_router(
         db=db,
         schema=strawberry.Schema(
             query=schema.query,
@@ -458,7 +453,6 @@ def create_app(
         model=model,
         corpus=corpus,
         export_path=export_path,
-        graphiql=True,
         streaming_last_updated_at=bulk_inserter.last_updated_at,
         cache_for_dataloaders=cache_for_dataloaders,
         read_only=read_only,
@@ -469,8 +463,11 @@ def create_app(
         prometheus_middlewares = [Middleware(PrometheusMiddleware)]
     else:
         prometheus_middlewares = []
-    app = Starlette(
+    app = FastAPI(
+        title="Arize-Phoenix REST API",
+        version=REST_API_VERSION,
         lifespan=_lifespan(
+            dialect=db.dialect,
             read_only=read_only,
             bulk_inserter=bulk_inserter,
             tracer_provider=tracer_provider,
@@ -481,58 +478,41 @@ def create_app(
             Middleware(HeadersMiddleware),
             *prometheus_middlewares,
         ],
+        exception_handlers={HTTPException: plain_text_http_exception_handler},
         debug=debug,
-        routes=V1_ROUTES
-        + [
-            Route("/schema", endpoint=openapi_schema, include_in_schema=False),
-            Route("/arize_phoenix_version", version),
-            Route("/healthz", check_healthz),
-            Route(
-                "/exports",
-                type(
-                    "DownloadExports",
-                    (Download,),
-                    {"path": export_path},
-                ),
-            ),
-            Route(
-                "/docs",
-                api_docs,
-            ),
-            Route(
-                "/graphql",
-                graphql,
-            ),
-        ]
-        + (
-            [
-                Mount(
-                    "/",
-                    app=Static(
-                        directory=SERVER_DIR / "static",
-                        app_config=AppConfig(
-                            has_inferences=model.is_empty is not True,
-                            has_corpus=corpus is not None,
-                            min_dist=umap_params.min_dist,
-                            n_neighbors=umap_params.n_neighbors,
-                            n_samples=umap_params.n_samples,
-                            is_development=dev,
-                            web_manifest_path=SERVER_DIR / "static" / ".vite" / "manifest.json",
-                        ),
-                    ),
-                    name="static",
-                ),
-            ]
-            if serve_ui
-            else []
-        ),
+        swagger_ui_parameters={
+            "defaultModelsExpandDepth": -1,  # hides the schema section in the Swagger UI
+        },
     )
     app.state.read_only = read_only
+    app.state.export_path = export_path
+    app.include_router(v1_router)
+    app.include_router(router)
+    app.include_router(graphql_router)
+    app.add_middleware(GZipMiddleware)
+    if serve_ui:
+        app.mount(
+            "/",
+            app=Static(
+                directory=SERVER_DIR / "static",
+                app_config=AppConfig(
+                    has_inferences=model.is_empty is not True,
+                    has_corpus=corpus is not None,
+                    min_dist=umap_params.min_dist,
+                    n_neighbors=umap_params.n_neighbors,
+                    n_samples=umap_params.n_samples,
+                    is_development=dev,
+                    web_manifest_path=SERVER_DIR / "static" / ".vite" / "manifest.json",
+                ),
+            ),
+            name="static",
+        )
+
     app.state.db = db
     if tracer_provider:
-        from opentelemetry.instrumentation.starlette import StarletteInstrumentor
+        from opentelemetry.instrumentation.fastapi import FastAPIInstrumentor
 
-        StarletteInstrumentor().instrument(tracer_provider=tracer_provider)
-        StarletteInstrumentor.instrument_app(app, tracer_provider=tracer_provider)
-        clean_ups.append(StarletteInstrumentor().uninstrument)
+        FastAPIInstrumentor().instrument(tracer_provider=tracer_provider)
+        FastAPIInstrumentor.instrument_app(app, tracer_provider=tracer_provider)
+        clean_ups.append(FastAPIInstrumentor().uninstrument)
     return app

phoenix/server/grpc_server.py

@@ -23,7 +23,7 @@ if TYPE_CHECKING:
 ProjectName: TypeAlias = str
 
 
-class Servicer(TraceServiceServicer):
+class Servicer(TraceServiceServicer):  # type:ignore
     def __init__(
         self,
         callback: Callable[[Span, ProjectName], Awaitable[None]],
@@ -78,7 +78,7 @@ class GrpcServer:
             interceptors=interceptors,
         )
         server.add_insecure_port(f"[::]:{get_env_grpc_port()}")
-        add_TraceServiceServicer_to_server(Servicer(self._callback), server)  # type: ignore
+        add_TraceServiceServicer_to_server(Servicer(self._callback), server)
         await server.start()
         self._server = server
 

phoenix/server/main.py

@@ -33,25 +33,23 @@ from phoenix.pointcloud.umap_parameters import (
     UMAPParameters,
 )
 from phoenix.server.app import (
-    SessionFactory,
     _db,
     create_app,
     create_engine_and_run_migrations,
     instrument_engine_if_enabled,
 )
+from phoenix.server.types import DbSessionFactory
 from phoenix.settings import Settings
 from phoenix.trace.fixtures import (
     TRACES_FIXTURES,
-    download_traces_fixture,
     get_dataset_fixtures,
     get_evals_from_fixture,
-    get_trace_fixture_by_name,
+    load_example_traces,
     reset_fixture_span_ids_and_timestamps,
     send_dataset_fixtures,
 )
 from phoenix.trace.otel import decode_otlp_span, encode_span_to_otlp
 from phoenix.trace.schemas import Span
-from phoenix.trace.span_json_decoder import json_string_to_span
 
 logger = logging.getLogger(__name__)
 
@@ -221,10 +219,8 @@ if __name__ == "__main__":
             (
                 # Apply `encode` here because legacy jsonl files contains UUIDs as strings.
                 # `encode` removes the hyphens in the UUIDs.
-                decode_otlp_span(encode_span_to_otlp(json_string_to_span(json_span)))
-                for json_span in download_traces_fixture(
-                    get_trace_fixture_by_name(trace_dataset_name)
-                )
+                decode_otlp_span(encode_span_to_otlp(span))
+                for span in load_example_traces(trace_dataset_name).to_spans()
             ),
             get_evals_from_fixture(trace_dataset_name),
         )
@@ -250,7 +246,7 @@ if __name__ == "__main__":
     working_dir = get_working_dir().resolve()
     engine = create_engine_and_run_migrations(db_connection_str)
     instrumentation_cleanups = instrument_engine_if_enabled(engine)
-    factory = SessionFactory(session_factory=_db(engine), dialect=engine.dialect.name)
+    factory = DbSessionFactory(db=_db(engine), dialect=engine.dialect.name)
     app = create_app(
         db=factory,
         export_path=export_path,