arize-phoenix 4.4.2__py3-none-any.whl → 4.4.4rc0__py3-none-any.whl

phoenix/server/api/types/pagination.py

@@ -2,60 +2,18 @@ import base64
 from dataclasses import dataclass
 from datetime import datetime
 from enum import Enum, auto
-from typing import ClassVar, Generic, List, Optional, Tuple, TypeVar, Union
+from typing import Any, ClassVar, List, Optional, Tuple, Union
 
-import strawberry
 from strawberry import UNSET
+from strawberry.relay.types import Connection, Edge, NodeType, PageInfo
 from typing_extensions import TypeAlias, assert_never
 
 ID: TypeAlias = int
-GenericType = TypeVar("GenericType")
 CursorSortColumnValue: TypeAlias = Union[str, int, float, datetime]
 
-
-@strawberry.type
-class Connection(Generic[GenericType]):
-    """Represents a paginated relationship between two entities
-
-    This pattern is used when the relationship itself has attributes.
-    """
-
-    page_info: "PageInfo"
-    edges: List["Edge[GenericType]"]
-
-
-@strawberry.type
-class PageInfo:
-    """Pagination context to navigate objects with cursor-based pagination
-
-    Instead of classic offset pagination via `page` and `limit` parameters,
-    here we have a cursor of the last object and we fetch items starting from that one
-
-    Read more at:
-        - https://graphql.org/learn/pagination/#pagination-and-edges
-        - https://relay.dev/graphql/connections.htm
-    """
-
-    has_next_page: bool
-    has_previous_page: bool
-    start_cursor: Optional[str]
-    end_cursor: Optional[str]
-
-
 # A type alias for the connection cursor implementation
 CursorString = str
 
-
-@strawberry.type
-class Edge(Generic[GenericType]):
-    """
-    An edge may contain additional information of the relationship. This is the trivial case
-    """
-
-    node: GenericType
-    cursor: str
-
-
 # The hashing prefix for a connection cursor
 CURSOR_PREFIX = "connection:"
 
@@ -218,9 +176,9 @@ class ConnectionArgs:
 
 
 def connection_from_list(
-    data: List[GenericType],
+    data: List[NodeType],
     args: ConnectionArgs,
-) -> Connection[GenericType]:
+) -> Connection[NodeType]:
     """
     A simple function that accepts a list and connection arguments, and returns
     a connection object for use in GraphQL. It uses list offsets as pagination,
@@ -230,11 +188,11 @@ def connection_from_list(
 
 
 def connection_from_list_slice(
-    list_slice: List[GenericType],
+    list_slice: List[NodeType],
     args: ConnectionArgs,
     slice_start: int,
     list_length: int,
-) -> Connection[GenericType]:
+) -> Connection[NodeType]:
     """
     Given a slice (subset) of a list, returns a connection object for use in
     GraphQL.
@@ -295,12 +253,12 @@ def connection_from_list_slice(
     )
 
 
-def connections(
-    data: List[Tuple[Cursor, GenericType]],
+def connection_from_cursors_and_nodes(
+    cursors_and_nodes: List[Tuple[Any, NodeType]],
     has_previous_page: bool,
     has_next_page: bool,
-) -> Connection[GenericType]:
-    edges = [Edge(node=node, cursor=str(cursor)) for cursor, node in data]
+) -> Connection[NodeType]:
+    edges = [Edge(node=node, cursor=str(cursor)) for cursor, node in cursors_and_nodes]
     has_edges = len(edges) > 0
     first_edge = edges[0] if has_edges else None
     last_edge = edges[-1] if has_edges else None

phoenix/server/app.py

@@ -33,7 +33,6 @@ from starlette.middleware.base import BaseHTTPMiddleware, RequestResponseEndpoin
 from starlette.requests import Request
 from starlette.responses import FileResponse, PlainTextResponse, Response
 from starlette.routing import Mount, Route
-from starlette.schemas import SchemaGenerator
 from starlette.staticfiles import StaticFiles
 from starlette.templating import Jinja2Templates
 from starlette.types import Scope, StatefulLifespan
@@ -58,18 +57,27 @@ from phoenix.pointcloud.umap_parameters import UMAPParameters
 from phoenix.server.api.context import Context, DataLoaders
 from phoenix.server.api.dataloaders import (
     CacheForDataLoaders,
+    DatasetExampleRevisionsDataLoader,
+    DatasetExampleSpansDataLoader,
     DocumentEvaluationsDataLoader,
     DocumentEvaluationSummaryDataLoader,
     DocumentRetrievalMetricsDataLoader,
     EvaluationSummaryDataLoader,
+    ExperimentAnnotationSummaryDataLoader,
+    ExperimentErrorRatesDataLoader,
+    ExperimentSequenceNumberDataLoader,
     LatencyMsQuantileDataLoader,
     MinStartOrMaxEndTimeDataLoader,
+    ProjectByNameDataLoader,
     RecordCountDataLoader,
     SpanDescendantsDataLoader,
     SpanEvaluationsDataLoader,
+    SpanProjectsDataLoader,
     TokenCountDataLoader,
     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.schema import schema
 from phoenix.server.grpc_server import GrpcServer
@@ -84,10 +92,6 @@ logger = logging.getLogger(__name__)
 
 templates = Jinja2Templates(directory=SERVER_DIR / "templates")
 
-schemas = SchemaGenerator(
-    {"openapi": "3.0.0", "info": {"title": "ArizePhoenix API", "version": "1.0"}}
-)
-
 
 class AppConfig(NamedTuple):
     has_inferences: bool
@@ -126,6 +130,7 @@ class Static(StaticFiles):
                     "n_neighbors": self._app_config.n_neighbors,
                     "n_samples": self._app_config.n_samples,
                     "basename": request.scope.get("root_path", ""),
+                    "platform_version": phoenix.__version__,
                     "request": request,
                 },
             )
@@ -185,6 +190,8 @@ class GraphQLWithContext(GraphQL):  # type: ignore
             export_path=self.export_path,
             streaming_last_updated_at=self.streaming_last_updated_at,
             data_loaders=DataLoaders(
+                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
@@ -199,6 +206,9 @@ class GraphQLWithContext(GraphQL):  # type: ignore
                     if self.cache_for_dataloaders
                     else None,
                 ),
+                experiment_annotation_summaries=ExperimentAnnotationSummaryDataLoader(self.db),
+                experiment_error_rates=ExperimentErrorRatesDataLoader(self.db),
+                experiment_sequence_number=ExperimentSequenceNumberDataLoader(self.db),
                 latency_ms_quantile=LatencyMsQuantileDataLoader(
                     self.db,
                     cache_map=self.cache_for_dataloaders.latency_ms_quantile
@@ -219,6 +229,7 @@ class GraphQLWithContext(GraphQL):  # type: ignore
                 ),
                 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
@@ -226,6 +237,8 @@ class GraphQLWithContext(GraphQL):  # type: ignore
                     else None,
                 ),
                 trace_evaluations=TraceEvaluationsDataLoader(self.db),
+                trace_row_ids=TraceRowIdsDataLoader(self.db),
+                project_by_name=ProjectByNameDataLoader(self.db),
             ),
             cache_for_dataloaders=self.cache_for_dataloaders,
             read_only=self.read_only,
@@ -272,7 +285,11 @@ def _lifespan(
 ) -> StatefulLifespan[Starlette]:
     @contextlib.asynccontextmanager
     async def lifespan(_: Starlette) -> AsyncIterator[Dict[str, Any]]:
-        async with bulk_inserter as (queue_span, queue_evaluation), GrpcServer(
+        async with bulk_inserter as (
+            queue_span,
+            queue_evaluation,
+            enqueue_operation,
+        ), GrpcServer(
             queue_span,
             disabled=read_only,
             tracer_provider=tracer_provider,
@@ -281,6 +298,7 @@ def _lifespan(
             yield {
                 "queue_span_for_bulk_insert": queue_span,
                 "queue_evaluation_for_bulk_insert": queue_evaluation,
+                "enqueue_operation": enqueue_operation,
             }
         for clean_up in clean_ups:
             clean_up()
@@ -293,15 +311,63 @@ async def check_healthz(_: Request) -> PlainTextResponse:
 
 
 async def openapi_schema(request: Request) -> Response:
-    return schemas.OpenAPIResponse(request=request)
+    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")
 
 
-def create_app(
+class SessionFactory:
+    def __init__(
+        self,
+        session_factory: Callable[[], AsyncContextManager[AsyncSession]],
+        dialect: str,
+    ):
+        self.session_factory = session_factory
+        self.dialect = SupportedSQLDialect(dialect)
+
+    def __call__(self) -> AsyncContextManager[AsyncSession]:
+        return self.session_factory()
+
+
+def create_engine_and_run_migrations(
     database_url: str,
+) -> AsyncEngine:
+    try:
+        return create_engine(database_url)
+    except PhoenixMigrationError as e:
+        msg = (
+            "\n\n⚠️⚠️ Phoenix failed to migrate the database to the latest version. ⚠️⚠️\n\n"
+            "The database may be in a dirty state. To resolve this, the Alembic CLI can be used\n"
+            "from the `src/phoenix/db` directory inside the Phoenix project root. From here,\n"
+            "revert any partial migrations and run `alembic stamp` to reset the migration state,\n"
+            "then try starting Phoenix again.\n\n"
+            "If issues persist, please reach out for support in the Arize community Slack:\n"
+            "https://arize-ai.slack.com\n\n"
+            "You can also refer to the Alembic documentation for more information:\n"
+            "https://alembic.sqlalchemy.org/en/latest/tutorial.html\n\n"
+            ""
+        )
+        raise PhoenixMigrationError(msg) from e
+
+
+def instrument_engine_if_enabled(engine: AsyncEngine) -> List[Callable[[], None]]:
+    instrumentation_cleanups = []
+    if server_instrumentation_is_enabled():
+        from opentelemetry.instrumentation.sqlalchemy import SQLAlchemyInstrumentor
+
+        tracer_provider = initialize_opentelemetry_tracer_provider()
+        SQLAlchemyInstrumentor().instrument(
+            engine=engine.sync_engine,
+            tracer_provider=tracer_provider,
+        )
+        instrumentation_cleanups.append(SQLAlchemyInstrumentor().uninstrument)
+    return instrumentation_cleanups
+
+
+def create_app(
+    db: SessionFactory,
     export_path: Path,
     model: Model,
     umap_params: UMAPParameters,
@@ -311,8 +377,10 @@ def create_app(
     enable_prometheus: bool = False,
     initial_spans: Optional[Iterable[Union[Span, Tuple[Span, str]]]] = None,
     initial_evaluations: Optional[Iterable[pb.Evaluation]] = None,
+    serve_ui: bool = True,
+    clean_up_callbacks: List[Callable[[], None]] = [],
 ) -> Starlette:
-    clean_ups: List[Callable[[], None]] = []  # To be called at app shutdown.
+    clean_ups: List[Callable[[], None]] = clean_up_callbacks  # To be called at app shutdown.
     initial_batch_of_spans: Iterable[Tuple[Span, str]] = (
         ()
         if initial_spans is None
@@ -322,28 +390,10 @@ def create_app(
         )
     )
     initial_batch_of_evaluations = () if initial_evaluations is None else initial_evaluations
-    try:
-        engine = create_engine(database_url)
-    except PhoenixMigrationError as e:
-        msg = (
-            "\n\n⚠️⚠️ Phoenix failed to migrate the database to the latest version. ⚠️⚠️\n\n"
-            "The database may be in a dirty state. To resolve this, the Alembic CLI can be used\n"
-            "from the `src/phoenix/db` directory inside the Phoenix project root. From here,\n"
-            "revert any partial migrations and run `alembic stamp` to reset the migration state,\n"
-            "then try starting Phoenix again.\n\n"
-            "If issues persist, please reach out for support in the Arize community Slack:\n"
-            "https://arize-ai.slack.com\n\n"
-            "You can also refer to the Alembic documentation for more information:\n"
-            "https://alembic.sqlalchemy.org/en/latest/tutorial.html\n\n"
-            ""
-        )
-        raise PhoenixMigrationError(msg) from e
     cache_for_dataloaders = (
-        CacheForDataLoaders()
-        if SupportedSQLDialect(engine.dialect.name) is SupportedSQLDialect.SQLITE
-        else None
+        CacheForDataLoaders() if db.dialect is SupportedSQLDialect.SQLITE else None
     )
-    db = _db(engine)
+
     bulk_inserter = BulkInserter(
         db,
         enable_prometheus=enable_prometheus,
@@ -354,16 +404,9 @@ def create_app(
     tracer_provider = None
     strawberry_extensions = schema.get_extensions()
     if server_instrumentation_is_enabled():
-        from opentelemetry.instrumentation.sqlalchemy import SQLAlchemyInstrumentor
         from opentelemetry.trace import TracerProvider
         from strawberry.extensions.tracing import OpenTelemetryExtension
 
-        tracer_provider = initialize_opentelemetry_tracer_provider()
-        SQLAlchemyInstrumentor().instrument(
-            engine=engine.sync_engine,
-            tracer_provider=tracer_provider,
-        )
-        clean_ups.append(SQLAlchemyInstrumentor().uninstrument)
         if TYPE_CHECKING:
             # Type-check the class before monkey-patching its private attribute.
             assert OpenTelemetryExtension._tracer
@@ -377,6 +420,7 @@ def create_app(
                 self._tracer = cast(TracerProvider, tracer_provider).get_tracer("strawberry")
 
         strawberry_extensions.append(_OpenTelemetryExtension)
+
     graphql = GraphQLWithContext(
         db=db,
         schema=strawberry.Schema(
@@ -433,21 +477,27 @@ def create_app(
                 "/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,
+        ]
+        + (
+            [
+                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",
                 ),
-                name="static",
-            ),
-        ],
+            ]
+            if serve_ui
+            else []
+        ),
     )
     app.state.read_only = read_only
     app.state.db = db

phoenix/server/main.py

@@ -22,9 +22,9 @@ from phoenix.config import (
     get_pids_path,
     get_working_dir,
 )
-from phoenix.core.model_schema_adapter import create_model_from_datasets
+from phoenix.core.model_schema_adapter import create_model_from_inferences
 from phoenix.db import get_printable_db_url
-from phoenix.inferences.fixtures import FIXTURES, get_datasets
+from phoenix.inferences.fixtures import FIXTURES, get_inferences
 from phoenix.inferences.inferences import EMPTY_INFERENCES, Inferences
 from phoenix.pointcloud.umap_parameters import (
     DEFAULT_MIN_DIST,
@@ -32,14 +32,22 @@ from phoenix.pointcloud.umap_parameters import (
     DEFAULT_N_SAMPLES,
     UMAPParameters,
 )
-from phoenix.server.app import create_app
+from phoenix.server.app import (
+    SessionFactory,
+    _db,
+    create_app,
+    create_engine_and_run_migrations,
+    instrument_engine_if_enabled,
+)
 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,
     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
@@ -99,14 +107,14 @@ def _get_pid_file() -> Path:
 DEFAULT_UMAP_PARAMS_STR = f"{DEFAULT_MIN_DIST},{DEFAULT_N_NEIGHBORS},{DEFAULT_N_SAMPLES}"
 
 if __name__ == "__main__":
-    primary_dataset_name: str
-    reference_dataset_name: Optional[str]
+    primary_inferences_name: str
+    reference_inferences_name: Optional[str]
     trace_dataset_name: Optional[str] = None
     simulate_streaming: Optional[bool] = None
 
-    primary_dataset: Inferences = EMPTY_INFERENCES
-    reference_dataset: Optional[Inferences] = None
-    corpus_dataset: Optional[Inferences] = None
+    primary_inferences: Inferences = EMPTY_INFERENCES
+    reference_inferences: Optional[Inferences] = None
+    corpus_inferences: Optional[Inferences] = None
 
     # Initialize the settings for the Server
     Settings.log_migrations = True
@@ -150,34 +158,34 @@ if __name__ == "__main__":
     )
     export_path = Path(args.export_path) if args.export_path else EXPORT_DIR
     if args.command == "datasets":
-        primary_dataset_name = args.primary
-        reference_dataset_name = args.reference
-        corpus_dataset_name = args.corpus
-        primary_dataset = Inferences.from_name(primary_dataset_name)
-        reference_dataset = (
-            Inferences.from_name(reference_dataset_name)
-            if reference_dataset_name is not None
+        primary_inferences_name = args.primary
+        reference_inferences_name = args.reference
+        corpus_inferences_name = args.corpus
+        primary_inferences = Inferences.from_name(primary_inferences_name)
+        reference_inferences = (
+            Inferences.from_name(reference_inferences_name)
+            if reference_inferences_name is not None
             else None
         )
-        corpus_dataset = (
-            None if corpus_dataset_name is None else Inferences.from_name(corpus_dataset_name)
+        corpus_inferences = (
+            None if corpus_inferences_name is None else Inferences.from_name(corpus_inferences_name)
         )
     elif args.command == "fixture":
         fixture_name = args.fixture
         primary_only = args.primary_only
-        primary_dataset, reference_dataset, corpus_dataset = get_datasets(
+        primary_inferences, reference_inferences, corpus_inferences = get_inferences(
             fixture_name,
             args.no_internet,
         )
         if primary_only:
-            reference_dataset_name = None
-            reference_dataset = None
+            reference_inferences_name = None
+            reference_inferences = None
     elif args.command == "trace-fixture":
         trace_dataset_name = args.fixture
         simulate_streaming = args.simulate_streaming
     elif args.command == "demo":
         fixture_name = args.fixture
-        primary_dataset, reference_dataset, corpus_dataset = get_datasets(
+        primary_inferences, reference_inferences, corpus_inferences = get_inferences(
             fixture_name,
             args.no_internet,
         )
@@ -197,9 +205,11 @@ if __name__ == "__main__":
 
     port = args.port or get_env_port()
     host_root_path = get_env_host_root_path()
-    model = create_model_from_datasets(
-        primary_dataset,
-        reference_dataset,
+    read_only = args.read_only
+
+    model = create_model_from_inferences(
+        primary_inferences,
+        reference_inferences,
     )
 
     fixture_spans: List[Span] = []
@@ -216,13 +226,19 @@ if __name__ == "__main__":
             ),
             get_evals_from_fixture(trace_dataset_name),
         )
+        dataset_fixtures = list(get_dataset_fixtures(trace_dataset_name))
+        if not read_only:
+            Thread(
+                target=send_dataset_fixtures,
+                args=(f"http://{host}:{port}", dataset_fixtures),
+            ).start()
     umap_params_list = args.umap_params.split(",")
     umap_params = UMAPParameters(
         min_dist=float(umap_params_list[0]),
         n_neighbors=int(umap_params_list[1]),
         n_samples=int(umap_params_list[2]),
     )
-    read_only = args.read_only
+
     logger.info(f"Server umap params: {umap_params}")
     if enable_prometheus := get_env_enable_prometheus():
         from phoenix.server.prometheus import start_prometheus
@@ -230,17 +246,23 @@ if __name__ == "__main__":
         start_prometheus()
 
     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)
     app = create_app(
-        database_url=db_connection_str,
+        db=factory,
         export_path=export_path,
         model=model,
         umap_params=umap_params,
-        corpus=None if corpus_dataset is None else create_model_from_datasets(corpus_dataset),
+        corpus=None
+        if corpus_inferences is None
+        else create_model_from_inferences(corpus_inferences),
         debug=args.debug,
         read_only=read_only,
         enable_prometheus=enable_prometheus,
         initial_spans=fixture_spans,
         initial_evaluations=fixture_evals,
+        clean_up_callbacks=instrumentation_cleanups,
     )
     server = Server(config=Config(app, host=host, port=port, root_path=host_root_path))  # type: ignore
     Thread(target=_write_pid_file_when_ready, args=(server,), daemon=True).start()

phoenix/server/openapi/docs.py

@@ -43,6 +43,9 @@ def get_swagger_ui_html(
     <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({{