arize-phoenix 4.12.0rc1__py3-none-any.whl → 4.12.1rc1__py3-none-any.whl

phoenix/server/app.py

@@ -1,8 +1,8 @@
 import contextlib
-from functools import cached_property
 import json
 import logging
 from datetime import datetime
+from functools import cached_property
 from pathlib import Path
 from typing import (
     TYPE_CHECKING,
@@ -21,25 +21,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
 
@@ -82,11 +81,10 @@ 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.trace.schemas import Span
 
@@ -95,6 +93,8 @@ if TYPE_CHECKING:
 
 logger = logging.getLogger(__name__)
 
+router = APIRouter(include_in_schema=False)
+
 templates = Jinja2Templates(directory=SERVER_DIR / "templates")
 
 
@@ -125,12 +125,9 @@ class Static(StaticFiles):
                 return cast(Dict[str, Any], json.load(f))
         except FileNotFoundError as e:
             if self._app_config.is_development:
-                logger.warning(
-                    f"Web manifest not found at {self._app_config.web_manifest_path}"
-                )
                 return {}
             raise e
-    
+
     def _sanitize_basename(self, basename: str) -> str:
         return basename[:-1] if basename.endswith("/") else basename
 
@@ -156,7 +153,7 @@ class Static(StaticFiles):
                     "platform_version": phoenix.__version__,
                     "request": request,
                     "is_development": self._app_config.is_development,
-                    "manifest": self._web_manifest
+                    "manifest": self._web_manifest,
                 },
             )
         except Exception as e:
@@ -179,116 +176,20 @@ class HeadersMiddleware(BaseHTTPMiddleware):
 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,
-        )
-
-
-class Download(HTTPEndpoint):
-    path: Path
-
-    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",
-        )
+@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",
+    )
 
 
-async def version(_: Request) -> PlainTextResponse:
+@router.get("/arize_phoenix_version")
+async def version() -> PlainTextResponse:
     return PlainTextResponse(f"{phoenix.__version__}")
 
 
@@ -310,9 +211,9 @@ def _lifespan(
     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]]:
         async with bulk_inserter as (
             queue_span,
             queue_evaluation,
@@ -334,16 +235,90 @@ 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)
-
+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]
+    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_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 api_docs(request: Request) -> Response:
-    return get_swagger_ui_html(openapi_url="/schema", title="arize-phoenix API")
+    return GraphQLRouter(
+        schema,
+        graphiql=True,
+        context_getter=get_context,
+        include_in_schema=False,
+        prefix="/graphql",
+    )
 
 
 class SessionFactory:
@@ -394,6 +369,18 @@ 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,
@@ -408,7 +395,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]] = (
         ()
@@ -451,7 +438,7 @@ def create_app(
 
         strawberry_extensions.append(_OpenTelemetryExtension)
 
-    graphql = GraphQLWithContext(
+    graphql_router = create_graphql_router(
         db=db,
         schema=strawberry.Schema(
             query=schema.query,
@@ -462,7 +449,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,
@@ -473,7 +459,9 @@ 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(
             read_only=read_only,
             bulk_inserter=bulk_inserter,
@@ -485,58 +473,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/templates/index.html

@@ -1,43 +1,47 @@
-{% set rendered_files = [] %}
-{% set rendered_chunks = [] %}
-
-{% macro render_chunk(chunk_name, level=0) %}
-  {% if chunk_name not in rendered_chunks %}
-    {% set _ = rendered_chunks.append(chunk_name) %}
-    {% set chunk = manifest[chunk_name] %}
-    
-    {% if chunk.css %}
-      {% for css_file in chunk.css %}
-        {% if css_file not in rendered_files %}
-          <link rel="stylesheet" href="{{ basename }}/{{ css_file }}">
-          {% set _ = rendered_files.append(css_file) %}
+{%- set rendered_chunks = [] -%}
+{%- set css_links = [] -%}
+{%- set js_scripts = [] -%}
+{%- macro collect_assets(chunk_name, level=0) -%}
+    {% if chunk_name not in rendered_chunks %}
+        {% set _ = rendered_chunks.append(chunk_name) %}
+        {% set chunk = manifest[chunk_name] %}
+        {% if chunk.css %}
+            {% for css_file in chunk.css %}
+                {% if css_file not in css_links %}
+                    {% set _ = css_links.append((basename, css_file)) %}
+                {% endif %}
+            {% endfor %}
+        {% endif %}
+        {% if chunk.imports %}
+            {% for import in chunk.imports %}
+                {% set _ = collect_assets(import, level + 1) %}
+            {% endfor %}
+        {% endif %}
+        {% if chunk.file.endswith('.js') %}
+            {% if chunk.file not in js_scripts %}
+                {% set _ = js_scripts.append((basename, chunk.file, level == 0)) %}
+            {% endif %}
+        {% elif chunk.file.endswith('.css') %}
+            {% if chunk.file not in css_links %}
+                {% set _ = css_links.append((basename, chunk.file)) %}
+            {% endif %}
         {% endif %}
-      {% endfor %}
-    {% endif %}
-    
-    {% if chunk.imports %}
-      {% for import in chunk.imports %}
-        {{ render_chunk(import, level + 1) }}
-      {% endfor %}
-    {% endif %}
-    
-    {% if level == 0 %}
-      {% if chunk.file.endswith('.js') and chunk.file not in rendered_files %}
-        <script type="module" src="{{ basename }}/{{ chunk.file }}"></script>
-        {% set _ = rendered_files.append(chunk.file) %}
-      {% elif chunk.file.endswith('.css') and chunk.file not in rendered_files %}
-        <link rel="stylesheet" href="{{ basename }}/{{ chunk.file }}">
-        {% set _ = rendered_files.append(chunk.file) %}
-      {% endif %}
-    {% endif %}
-    
-    {% if chunk.file.endswith('.js') and chunk.file not in rendered_files %}
-      <link rel="modulepreload" href="{{ basename }}/{{ chunk.file }}">
-      {% set _ = rendered_files.append(chunk.file) %}
     {% endif %}
-  {% endif %}
-{% endmacro %}
-
+{%- endmacro -%}
+{%- macro render_css() -%}
+  {%- for basename, css_file in css_links -%}
+    <link rel="stylesheet" href="{{ basename }}/{{ css_file }}">
+  {% endfor -%}
+{%- endmacro -%}
+{%- macro render_js() -%}
+  {%- for basename, js_file, is_entry in js_scripts -%}
+    {%- if is_entry -%}
+      <script type="module" src="{{ basename }}/{{ js_file }}"></script>
+    {% else -%}
+      <link rel="modulepreload" href="{{ basename }}/{{ js_file }}">
+    {% endif -%}
+  {%- endfor -%}
+{%- endmacro -%}
 <!DOCTYPE html>
 <html>
   <head>
@@ -55,6 +59,10 @@
       rel="stylesheet"
       href="https://fonts.googleapis.com/css2?family=Roboto:wght@100;300;400;500;700;900&display=swap"
     />
+    {% if not is_development -%}
+      {% set _ = collect_assets('index.tsx') -%}
+      {{- render_css() -}}
+    {%- endif -%}
     <script src="{{basename}}/modernizr.js"></script>
   </head>
   <body>
@@ -82,7 +90,7 @@
           writable: false
       });
   })()</script>
-    {% if is_development %}
+    {% if is_development -%}
       <script type="module">
         import RefreshRuntime from 'http://localhost:5173/@react-refresh'
         RefreshRuntime.injectIntoGlobalHook(window)
@@ -92,8 +100,8 @@
       </script>
       <script type="module" src="http://localhost:5173/@vite/client"></script>
       <script type="module" src="http://localhost:5173/index.tsx"></script>
-    {% else %}
-      {{ render_chunk('index.tsx') }}
-    {% endif %}
-      </body>
+    {%- else -%}
+      {{- render_js() -}}
+    {%- endif -%}
+  </body>
 </html>

phoenix/server/thread_server.py

@@ -4,7 +4,7 @@ from threading import Thread
 from time import sleep, time
 from typing import Generator
 
-from starlette.applications import Starlette
+from fastapi import FastAPI
 from uvicorn import Config, Server
 from uvicorn.config import LoopSetupType
 
@@ -24,7 +24,7 @@ class ThreadServer(Server):
 
     def __init__(
         self,
-        app: Starlette,
+        app: FastAPI,
         host: str,
         port: int,
         root_path: str,

phoenix/session/client.py

@@ -300,14 +300,15 @@ class Client(TraceDataExtractor):
         for otlp_span in otlp_spans:
             serialized = otlp_span.SerializeToString()
             content = gzip.compress(serialized)
-            self._client.post(
+            response = self._client.post(
                 url=urljoin(self._base_url, "v1/traces"),
                 content=content,
                 headers={
                     "content-type": "application/x-protobuf",
                     "content-encoding": "gzip",
                 },
-            ).raise_for_status()
+            )
+            response.raise_for_status()
 
     def _get_dataset_id_by_name(self, name: str) -> str:
         """

phoenix/version.py

@@ -1 +1 @@
-__version__ = "4.12.0rc1"
+__version__ = "4.12.1rc1"