feat: Add default UMAP parameters in launch_app()

src/phoenix/core/umap_parameters.py

@@ -0,0 +1,7 @@
+from dataclasses import dataclass
+from typing import Union, Dict
+
+
+@dataclass
+class UMAPParameters:
+    default_umap_params: Dict[str, Union[int, float]]

src/phoenix/server/api/context.py

@@ -8,6 +8,7 @@
 
 from phoenix.core.model_schema import Model
 from phoenix.core.traces import Traces
+from phoenix.core.umap_parameters import UMAPParameters
 
 
 @dataclass
@@ -18,3 +19,4 @@ class Context:
     export_path: Path
     corpus: Optional[Model] = None
     traces: Optional[Traces] = None
+    config: Optional[UMAPParameters] = None

src/phoenix/server/app.py

@@ -20,10 +20,12 @@
 from phoenix.config import SERVER_DIR
 from phoenix.core.model_schema import Model
 from phoenix.core.traces import Traces
+from phoenix.core.umap_parameters import UMAPParameters
 from phoenix.server.api.context import Context
 from phoenix.server.api.schema import schema
 from phoenix.server.span_handler import SpanHandler
 
+
 logger = logging.getLogger(__name__)
 
 
@@ -67,10 +69,12 @@ def __init__(
         graphiql: bool = False,
         corpus: Optional[Model] = None,
         traces: Optional[Traces] = None,
+        config: Optional[UMAPParameters] = None,
     ) -> None:
         self.model = model
         self.corpus = corpus
         self.traces = traces
+        self.config = config
         self.export_path = export_path
         super().__init__(schema, graphiql=graphiql)
 
@@ -85,6 +89,7 @@ async def get_context(
             model=self.model,
             corpus=self.corpus,
             traces=self.traces,
+            config=self.config,
             export_path=self.export_path,
         )
 
@@ -109,13 +114,15 @@ def create_app(
     model: Model,
     corpus: Optional[Model] = None,
     traces: Optional[Traces] = None,
+    config: Optional[UMAPParameters] = None,
     debug: bool = False,
 ) -> Starlette:
     graphql = GraphQLWithContext(
         schema=schema,
         model=model,
         corpus=corpus,
         traces=traces,
+        config=config,
         export_path=export_path,
         graphiql=True,
     )

src/phoenix/session/session.py

@@ -5,14 +5,15 @@
 from datetime import datetime
 from pathlib import Path
 from tempfile import TemporaryDirectory
-from typing import TYPE_CHECKING, Iterable, List, Optional, Set
+from typing import TYPE_CHECKING, Dict, Iterable, List, Optional, Set, Union
 
 import pandas as pd
 from portpicker import pick_unused_port
 
 from phoenix.config import HOST, PORT, get_exported_files
 from phoenix.core.model_schema_adapter import create_model_from_datasets
 from phoenix.core.traces import Traces
+from phoenix.core.umap_parameters import UMAPParameters
 from phoenix.datasets.dataset import EMPTY_DATASET, Dataset
 from phoenix.server.app import create_app
 from phoenix.server.thread_server import ThreadServer
@@ -70,13 +71,19 @@ def __init__(
         reference_dataset: Optional[Dataset] = None,
         corpus_dataset: Optional[Dataset] = None,
         trace_dataset: Optional[TraceDataset] = None,
+        default_umap_parameters: Optional[Dict[str, int]] = None,
         host: str = HOST,
         port: int = PORT,
     ):
         self.primary_dataset = primary_dataset
         self.reference_dataset = reference_dataset
         self.corpus_dataset = corpus_dataset
         self.trace_dataset = trace_dataset
+        self.config = (
+            UMAPParameters(default_umap_parameters)
+            if default_umap_parameters is not None
+            else None,
+        )
         self.model = create_model_from_datasets(
             primary_dataset,
             reference_dataset,
@@ -175,6 +182,7 @@ def __init__(
         reference_dataset: Optional[Dataset] = None,
         corpus_dataset: Optional[Dataset] = None,
         trace_dataset: Optional[TraceDataset] = None,
+        default_umap_parameters: Optional[Dict[str, int]] = None,
         host: str = HOST,
         port: int = PORT,
     ) -> None:
@@ -183,6 +191,7 @@ def __init__(
             reference_dataset=reference_dataset,
             corpus_dataset=corpus_dataset,
             trace_dataset=trace_dataset,
+            default_umap_parameters=default_umap_parameters,
             host=host,
             port=port or PORT,
         )
@@ -203,6 +212,7 @@ def __init__(
             corpus_dataset_name=(
                 self.corpus_dataset.name if self.corpus_dataset is not None else None
             ),
+            # TO-DO default_umap_params for ProcessSession
             trace_dataset_name=(
                 self.trace_dataset.name if self.trace_dataset is not None else None
             ),
@@ -224,6 +234,7 @@ def __init__(
         reference_dataset: Optional[Dataset] = None,
         corpus_dataset: Optional[Dataset] = None,
         trace_dataset: Optional[TraceDataset] = None,
+        default_umap_parameters: Optional[Dict[str, int]] = None,
         host: str = HOST,
         port: int = PORT,
     ):
@@ -232,6 +243,7 @@ def __init__(
             reference_dataset=reference_dataset,
             corpus_dataset=corpus_dataset,
             trace_dataset=trace_dataset,
+            default_umap_parameters=default_umap_parameters,
             host=host,
             port=port or pick_unused_port(),
         )
@@ -241,6 +253,7 @@ def __init__(
             model=self.model,
             corpus=self.corpus,
             traces=self.traces,
+            config=self.config,
         )
         self.server = ThreadServer(
             app=self.app,
@@ -264,6 +277,7 @@ def launch_app(
     reference: Optional[Dataset] = None,
     corpus: Optional[Dataset] = None,
     trace: Optional[TraceDataset] = None,
+    default_umap_parameters: Optional[Dict[str, Union[int, float]]] = None,
     host: str = HOST,
     port: int = PORT,
     run_in_thread: bool = True,
@@ -288,6 +302,8 @@ def launch_app(
         The port on which the server listens.
     run_in_thread: bool, optional, default=True
         Whether the server should run in a Thread or Process.
+    default_umap_parameters: Dict[str, Union[int, float]], optional, default=None
+        User specified default UMAP parameters eg: {nNeighbors: 10, nSamples: 5, minDist: 0.5}
 
     Returns
     -------
@@ -316,7 +332,9 @@ def launch_app(
                 "it down and starting a new instance..."
             )
             _session.end()
-        _session = ThreadSession(primary, reference, corpus, trace, host=host, port=port)
+        _session = ThreadSession(
+          primary, reference, corpus, trace, default_umap_parameters, host=host, port=port
+        )
         # TODO: catch exceptions from thread
         if not _session.active:
             logger.error(
@@ -325,7 +343,9 @@ def launch_app(
             )
             return None
     else:
-        _session = ProcessSession(primary, reference, corpus, trace, host=host, port=port)
+        _session = ProcessSession(
+          primary, reference, corpus, trace, default_umap_parameters, host=host, port=port
+        )
 
     print(f"🌍 To view the Phoenix app in your browser, visit {_session.url}")
     print("📺 To view the Phoenix app in a notebook, run `px.active_session().view()`")