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

phoenix/server/app.py

@@ -19,24 +19,25 @@ 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 PlainTextResponse, Response
+from starlette.responses import FileResponse, PlainTextResponse, Response
+from starlette.routing import Mount, Route
 from starlette.staticfiles import StaticFiles
 from starlette.templating import Jinja2Templates
 from starlette.types import Scope, StatefulLifespan
-from strawberry.fastapi import GraphQLRouter
+from starlette.websockets import WebSocket
+from strawberry.asgi import GraphQL
 from strawberry.schema import BaseSchema
 from typing_extensions import TypeAlias
 
@@ -79,10 +80,11 @@ from phoenix.server.api.dataloaders import (
     TraceEvaluationsDataLoader,
     TraceRowIdsDataLoader,
 )
-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.openapi.schema import OPENAPI_SCHEMA_GENERATOR
+from phoenix.server.api.routers.v1 import V1_ROUTES
 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.trace.schemas import Span
 
@@ -91,8 +93,6 @@ if TYPE_CHECKING:
 
 logger = logging.getLogger(__name__)
 
-router = APIRouter(include_in_schema=False)
-
 templates = Jinja2Templates(directory=SERVER_DIR / "templates")
 
 
@@ -157,20 +157,116 @@ class HeadersMiddleware(BaseHTTPMiddleware):
 ProjectRowId: TypeAlias = int
 
 
-@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 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,
+        )
+
 
+class Download(HTTPEndpoint):
+    path: Path
 
-@router.get("/arize_phoenix_version")
-async def version() -> PlainTextResponse:
+    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__}")
 
 
@@ -192,9 +288,9 @@ def _lifespan(
     enable_prometheus: bool = False,
     clean_ups: Iterable[Callable[[], None]] = (),
     read_only: bool = False,
-) -> StatefulLifespan[FastAPI]:
+) -> StatefulLifespan[Starlette]:
     @contextlib.asynccontextmanager
-    async def lifespan(_: FastAPI) -> AsyncIterator[Dict[str, Any]]:
+    async def lifespan(_: Starlette) -> AsyncIterator[Dict[str, Any]]:
         async with bulk_inserter as (
             queue_span,
             queue_evaluation,
@@ -216,89 +312,16 @@ def _lifespan(
     return lifespan
 
 
-@router.get("/healthz")
 async def check_healthz(_: Request) -> PlainTextResponse:
     return PlainTextResponse("OK")
 
 
-def create_graphql_router(
-    *,
-    schema: BaseSchema,
-    db: Callable[[], AsyncContextManager[AsyncSession]],
-    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]
-    context = 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_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,
-    )
+async def openapi_schema(request: Request) -> Response:
+    return OPENAPI_SCHEMA_GENERATOR.OpenAPIResponse(request=request)
 
-    return GraphQLRouter(
-        schema,
-        graphiql=True,
-        context_getter=lambda: context,
-        include_in_schema=False,
-        prefix="/graphql",
-    )
+
+async def api_docs(request: Request) -> Response:
+    return get_swagger_ui_html(openapi_url="/schema", title="arize-phoenix API")
 
 
 class SessionFactory:
@@ -349,18 +372,6 @@ 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,
     export_path: Path,
@@ -374,7 +385,7 @@ def create_app(
     initial_evaluations: Optional[Iterable[pb.Evaluation]] = None,
     serve_ui: bool = True,
     clean_up_callbacks: List[Callable[[], None]] = [],
-) -> FastAPI:
+) -> Starlette:
     clean_ups: List[Callable[[], None]] = clean_up_callbacks  # To be called at app shutdown.
     initial_batch_of_spans: Iterable[Tuple[Span, str]] = (
         ()
@@ -399,6 +410,7 @@ def create_app(
     tracer_provider = None
     strawberry_extensions = schema.get_extensions()
     if server_instrumentation_is_enabled():
+        tracer_provider = initialize_opentelemetry_tracer_provider()
         from opentelemetry.trace import TracerProvider
         from strawberry.extensions.tracing import OpenTelemetryExtension
 
@@ -416,7 +428,7 @@ def create_app(
 
         strawberry_extensions.append(_OpenTelemetryExtension)
 
-    graphql_router = create_graphql_router(
+    graphql = GraphQLWithContext(
         db=db,
         schema=strawberry.Schema(
             query=schema.query,
@@ -427,6 +439,7 @@ 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,
@@ -437,9 +450,7 @@ def create_app(
         prometheus_middlewares = [Middleware(PrometheusMiddleware)]
     else:
         prometheus_middlewares = []
-    app = FastAPI(
-        title="Arize-Phoenix REST API",
-        version=REST_API_VERSION,
+    app = Starlette(
         lifespan=_lifespan(
             read_only=read_only,
             bulk_inserter=bulk_inserter,
@@ -451,39 +462,56 @@ def create_app(
             Middleware(HeadersMiddleware),
             *prometheus_middlewares,
         ],
-        exception_handlers={HTTPException: plain_text_http_exception_handler},
         debug=debug,
-        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,
+        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},
                 ),
             ),
-            name="static",
-        )
-
+            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,
+                        ),
+                    ),
+                    name="static",
+                ),
+            ]
+            if serve_ui
+            else []
+        ),
+    )
+    app.state.read_only = read_only
     app.state.db = db
     if tracer_provider:
-        from opentelemetry.instrumentation.fastapi import FastAPIInstrumentor
+        from opentelemetry.instrumentation.starlette import StarletteInstrumentor
 
-        FastAPIInstrumentor().instrument(tracer_provider=tracer_provider)
-        FastAPIInstrumentor.instrument_app(app, tracer_provider=tracer_provider)
-        clean_ups.append(FastAPIInstrumentor().uninstrument)
+        StarletteInstrumentor().instrument(tracer_provider=tracer_provider)
+        StarletteInstrumentor.instrument_app(app, tracer_provider=tracer_provider)
+        clean_ups.append(StarletteInstrumentor().uninstrument)
     return app

phoenix/server/openapi/docs.py

@@ -0,0 +1,221 @@
+import json
+from typing import Any, Dict, Optional
+
+from starlette.responses import HTMLResponse
+
+swagger_ui_default_parameters: Dict[str, Any] = {
+    "dom_id": "#swagger-ui",
+    "layout": "BaseLayout",
+    "deepLinking": True,
+    "showExtensions": True,
+    "showCommonExtensions": True,
+}
+
+
+def get_swagger_ui_html(
+    *,
+    openapi_url: str = "/schema",
+    title: str,
+    swagger_js_url: str = "https://cdn.jsdelivr.net/npm/swagger-ui-dist@5.9.0/swagger-ui-bundle.js",
+    swagger_css_url: str = "https://cdn.jsdelivr.net/npm/swagger-ui-dist@5.9.0/swagger-ui.css",
+    swagger_favicon_url: str = "/favicon.ico",
+    oauth2_redirect_url: Optional[str] = None,
+    init_oauth: Optional[str] = None,
+    swagger_ui_parameters: Optional[Dict[str, Any]] = None,
+) -> HTMLResponse:
+    """
+    Generate and return the HTML  that loads Swagger UI for the interactive API
+    docs (normally served at `/docs`).
+    """
+    current_swagger_ui_parameters = swagger_ui_default_parameters.copy()
+    if swagger_ui_parameters:
+        current_swagger_ui_parameters.update(swagger_ui_parameters)
+
+    html = f"""
+    <!DOCTYPE html>
+    <html>
+    <head>
+    <link type="text/css" rel="stylesheet" href="{swagger_css_url}">
+    <link rel="shortcut icon" href="{swagger_favicon_url}">
+    <title>{title}</title>
+    </head>
+    <body>
+    <div id="swagger-ui">
+    </div>
+    <script src="{swagger_js_url}"></script>
+    <style type="text/css">
+    div[id^="operations-private"]{{display:none}} #operations-tag-private{{display:none}}
+    </style>
+    <!-- `SwaggerUIBundle` is now available on the page -->
+    <script>
+    const ui = SwaggerUIBundle({{
+        url: '{openapi_url}',
+    """
+
+    for key, value in current_swagger_ui_parameters.items():
+        html += f"{json.dumps(key)}: {json.dumps(value)},\n"
+
+    if oauth2_redirect_url:
+        html += f"oauth2RedirectUrl: window.location.origin + '{oauth2_redirect_url}',"
+
+    html += """
+    presets: [
+        SwaggerUIBundle.presets.apis,
+        SwaggerUIBundle.SwaggerUIStandalonePreset
+        ],
+    })"""
+
+    if init_oauth:
+        html += f"""
+        ui.initOAuth({json.dumps(init_oauth)})
+        """
+
+    html += """
+    </script>
+    </body>
+    </html>
+    """
+    return HTMLResponse(html)
+
+
+def get_redoc_html(
+    *,
+    openapi_url: str,
+    title: str,
+    redoc_js_url: str = "https://cdn.jsdelivr.net/npm/redoc@next/bundles/redoc.standalone.js",
+    redoc_favicon_url: str = "/favicon.ico",
+    with_google_fonts: bool = True,
+) -> HTMLResponse:
+    """
+    Generate and return the HTML response that loads ReDoc for the alternative
+    API docs (normally served at `/redoc`).
+
+
+    """
+    html = f"""
+    <!DOCTYPE html>
+    <html>
+    <head>
+    <title>{title}</title>
+    <!-- needed for adaptive design -->
+    <meta charset="utf-8"/>
+    <meta name="viewport" content="width=device-width, initial-scale=1">
+    """
+    if with_google_fonts:
+        html += """
+    <link href="https://fonts.googleapis.com/css?family=Montserrat:300,400,700|Roboto:300,400,700" rel="stylesheet">
+    """  # noqa: E501
+    html += f"""
+    <link rel="shortcut icon" href="{redoc_favicon_url}">
+    <!--
+    ReDoc doesn't change outer page styles
+    -->
+    <style>
+      body {{
+        margin: 0;
+        padding: 0;
+      }}
+    </style>
+    </head>
+    <body>
+    <noscript>
+        ReDoc requires Javascript to function. Please enable it to browse the documentation.
+    </noscript>
+    <redoc spec-url="{openapi_url}"></redoc>
+    <script src="{redoc_js_url}"> </script>
+    </body>
+    </html>
+    """
+    return HTMLResponse(html)
+
+
+# Not needed now but copy-pasting for future reference
+def get_swagger_ui_oauth2_redirect_html() -> HTMLResponse:
+    """
+    Generate the HTML response with the OAuth2 redirection for Swagger UI.
+
+    You normally don't need to use or change this.
+    """
+    # copied from https://github.com/swagger-api/swagger-ui/blob/v4.14.0/dist/oauth2-redirect.html
+    html = """
+    <!doctype html>
+    <html lang="en-US">
+    <head>
+        <title>Swagger UI: OAuth2 Redirect</title>
+    </head>
+    <body>
+    <script>
+        'use strict';
+        function run () {
+            var oauth2 = window.opener.swaggerUIRedirectOauth2;
+            var sentState = oauth2.state;
+            var redirectUrl = oauth2.redirectUrl;
+            var isValid, qp, arr;
+
+            if (/code|token|error/.test(window.location.hash)) {
+                qp = window.location.hash.substring(1).replace('?', '&');
+            } else {
+                qp = location.search.substring(1);
+            }
+
+            arr = qp.split("&");
+            arr.forEach(function (v,i,_arr) { _arr[i] = '"' + v.replace('=', '":"') + '"';});
+            qp = qp ? JSON.parse('{' + arr.join() + '}',
+                    function (key, value) {
+                        return key === "" ? value : decodeURIComponent(value);
+                    }
+            ) : {};
+
+            isValid = qp.state === sentState;
+
+            if ((
+              oauth2.auth.schema.get("flow") === "accessCode" ||
+              oauth2.auth.schema.get("flow") === "authorizationCode" ||
+              oauth2.auth.schema.get("flow") === "authorization_code"
+            ) && !oauth2.auth.code) {
+                if (!isValid) {
+                    oauth2.errCb({
+                        authId: oauth2.auth.name,
+                        source: "auth",
+                        level: "warning",
+                        message: "Authorization may be unsafe, passed state was changed in server. The passed state wasn't returned from auth server."
+                    });
+                }
+
+                if (qp.code) {
+                    delete oauth2.state;
+                    oauth2.auth.code = qp.code;
+                    oauth2.callback({auth: oauth2.auth, redirectUrl: redirectUrl});
+                } else {
+                    let oauthErrorMsg;
+                    if (qp.error) {
+                        oauthErrorMsg = "["+qp.error+"]: " +
+                            (qp.error_description ? qp.error_description+ ". " : "no accessCode received from the server. ") +
+                            (qp.error_uri ? "More info: "+qp.error_uri : "");
+                    }
+
+                    oauth2.errCb({
+                        authId: oauth2.auth.name,
+                        source: "auth",
+                        level: "error",
+                        message: oauthErrorMsg || "[Authorization failed]: no accessCode received from the server."
+                    });
+                }
+            } else {
+                oauth2.callback({auth: oauth2.auth, token: qp, isValid: isValid, redirectUrl: redirectUrl});
+            }
+            window.close();
+        }
+
+        if (document.readyState !== 'loading') {
+            run();
+        } else {
+            document.addEventListener('DOMContentLoaded', function () {
+                run();
+            });
+        }
+    </script>
+    </body>
+    </html>
+        """  # noqa: E501
+    return HTMLResponse(content=html)