arize-phoenix 4.4.3__py3-none-any.whl → 4.4.4rc1__py3-none-any.whl

phoenix/server/templates/index.html

@@ -31,6 +31,7 @@
             // injected into the client before React runs
             value: Object.freeze({
                 basename: "{{basename}}",
+                platformVersion: "{{platform_version}}",
                 hasInferences: Boolean("{{has_inferences}}" == "True"),
                 hasCorpus: Boolean("{{has_corpus}}" == "True"),
                 UMAP: {

phoenix/services.py

@@ -99,10 +99,10 @@ class AppService(Service):
 
     working_dir = SERVER_DIR
 
-    # Internal references to the name / directory of the dataset(s)
-    __primary_dataset_name: str
-    __reference_dataset_name: Optional[str]
-    __corpus_dataset_name: Optional[str]
+    # Internal references to the name / directory of the inferences(s)
+    __primary_inferences_name: str
+    __reference_inferences_name: Optional[str]
+    __corpus_inferences_name: Optional[str]
     __trace_dataset_name: Optional[str]
 
     def __init__(
@@ -112,10 +112,10 @@ class AppService(Service):
         host: str,
         port: int,
         root_path: str,
-        primary_dataset_name: str,
+        primary_inferences_name: str,
         umap_params: str,
-        reference_dataset_name: Optional[str],
-        corpus_dataset_name: Optional[str],
+        reference_inferences_name: Optional[str],
+        corpus_inferences_name: Optional[str],
         trace_dataset_name: Optional[str],
     ):
         self.database_url = database_url
@@ -123,10 +123,10 @@ class AppService(Service):
         self.host = host
         self.port = port
         self.root_path = root_path  # TODO(mikeldking): Add support for root_path
-        self.__primary_dataset_name = primary_dataset_name
+        self.__primary_inferences_name = primary_inferences_name
         self.__umap_params = umap_params
-        self.__reference_dataset_name = reference_dataset_name
-        self.__corpus_dataset_name = corpus_dataset_name
+        self.__reference_inferences_name = reference_inferences_name
+        self.__corpus_inferences_name = corpus_inferences_name
         self.__trace_dataset_name = trace_dataset_name
         super().__init__()
 
@@ -147,12 +147,12 @@ class AppService(Service):
             self.__umap_params,
             "datasets",
             "--primary",
-            str(self.__primary_dataset_name),
+            str(self.__primary_inferences_name),
         ]
-        if self.__reference_dataset_name is not None:
-            command.extend(["--reference", str(self.__reference_dataset_name)])
-        if self.__corpus_dataset_name is not None:
-            command.extend(["--corpus", str(self.__corpus_dataset_name)])
+        if self.__reference_inferences_name is not None:
+            command.extend(["--reference", str(self.__reference_inferences_name)])
+        if self.__corpus_inferences_name is not None:
+            command.extend(["--corpus", str(self.__corpus_inferences_name)])
         if self.__trace_dataset_name is not None:
             command.extend(["--trace", str(self.__trace_dataset_name)])
         logger.info(f"command: {' '.join(command)}")

phoenix/session/client.py

@@ -1,27 +1,46 @@
+import csv
 import gzip
 import logging
 import weakref
+from collections import Counter
 from datetime import datetime
-from io import BytesIO
-from typing import Any, List, Optional, Union, cast
-from urllib.parse import urljoin
+from io import BytesIO, StringIO
+from pathlib import Path
+from typing import (
+    Any,
+    BinaryIO,
+    Dict,
+    Iterable,
+    List,
+    Literal,
+    Mapping,
+    Optional,
+    Tuple,
+    Union,
+    cast,
+)
+from urllib.parse import quote, urljoin
 
+import httpx
 import pandas as pd
 import pyarrow as pa
 from opentelemetry.proto.collector.trace.v1.trace_service_pb2 import ExportTraceServiceRequest
 from opentelemetry.proto.common.v1.common_pb2 import AnyValue, KeyValue
 from opentelemetry.proto.resource.v1.resource_pb2 import Resource
 from opentelemetry.proto.trace.v1.trace_pb2 import ResourceSpans, ScopeSpans
-from pyarrow import ArrowInvalid
-from requests import Session
+from pyarrow import ArrowInvalid, Table
+from typing_extensions import TypeAlias, assert_never
 
 from phoenix.config import (
+    get_env_client_headers,
     get_env_collector_endpoint,
     get_env_host,
     get_env_port,
     get_env_project_name,
 )
+from phoenix.datasets.types import Dataset, Example
 from phoenix.datetime_utils import normalize_datetime
+from phoenix.db.insertion.dataset import DatasetKeys
 from phoenix.session.data_extractor import DEFAULT_SPAN_LIMIT, TraceDataExtractor
 from phoenix.trace import Evaluations, TraceDataset
 from phoenix.trace.dsl import SpanQuery
@@ -36,14 +55,20 @@ class Client(TraceDataExtractor):
         *,
         endpoint: Optional[str] = None,
         warn_if_server_not_running: bool = True,
+        headers: Optional[Mapping[str, str]] = None,
         **kwargs: Any,  # for backward-compatibility
     ):
         """
         Client for connecting to a Phoenix server.
 
         Args:
-            endpoint (str, optional): Phoenix server endpoint, e.g. http://localhost:6006. If not
-                provided, the endpoint will be inferred from the environment variables.
+            endpoint (str, optional): Phoenix server endpoint, e.g.
+            http://localhost:6006. If not provided, the endpoint will be
+            inferred from the environment variables.
+
+            headers (Mapping[str, str], optional): Headers to include in each
+            network request. If not provided, the headers will be inferred from
+            the environment variables (if present).
         """
         if kwargs.pop("use_active_session_if_available", None) is not None:
             print(
@@ -52,14 +77,14 @@ class Client(TraceDataExtractor):
             )
         if kwargs:
             raise TypeError(f"Unexpected keyword arguments: {', '.join(kwargs)}")
+        headers = headers or get_env_client_headers()
         host = get_env_host()
         if host == "0.0.0.0":
             host = "127.0.0.1"
         base_url = endpoint or get_env_collector_endpoint() or f"http://{host}:{get_env_port()}"
         self._base_url = base_url if base_url.endswith("/") else base_url + "/"
-
-        self._session = Session()
-        weakref.finalize(self, self._session.close)
+        self._client = httpx.Client(headers=headers)
+        weakref.finalize(self, self._client.close)
         if warn_if_server_not_running:
             self._warn_if_phoenix_is_not_running()
 
@@ -98,7 +123,7 @@ class Client(TraceDataExtractor):
                 "stop_time is deprecated. Use end_time instead.",
             )
             end_time = end_time or stop_time
-        response = self._session.post(
+        response = self._client.post(
             url=urljoin(self._base_url, "v1/spans"),
             params={"project-name": project_name},
             json={
@@ -145,8 +170,8 @@ class Client(TraceDataExtractor):
                 empty list if no evaluations are found.
         """
         project_name = project_name or get_env_project_name()
-        response = self._session.get(
-            urljoin(self._base_url, "v1/evaluations"),
+        response = self._client.get(
+            url=urljoin(self._base_url, "v1/evaluations"),
             params={"project-name": project_name},
         )
         if response.status_code == 404:
@@ -167,7 +192,7 @@ class Client(TraceDataExtractor):
 
     def _warn_if_phoenix_is_not_running(self) -> None:
         try:
-            self._session.get(urljoin(self._base_url, "arize_phoenix_version")).raise_for_status()
+            self._client.get(urljoin(self._base_url, "arize_phoenix_version")).raise_for_status()
         except Exception:
             logger.warning(
                 f"Arize Phoenix is not running on {self._base_url}. Launch Phoenix "
@@ -197,9 +222,9 @@ class Client(TraceDataExtractor):
             headers = {"content-type": "application/x-pandas-arrow"}
             with pa.ipc.new_stream(sink, table.schema) as writer:
                 writer.write_table(table)
-            self._session.post(
-                urljoin(self._base_url, "v1/evaluations"),
-                data=cast(bytes, sink.getvalue().to_pybytes()),
+            self._client.post(
+                url=urljoin(self._base_url, "v1/evaluations"),
+                content=cast(bytes, sink.getvalue().to_pybytes()),
                 headers=headers,
             ).raise_for_status()
 
@@ -239,16 +264,286 @@ class Client(TraceDataExtractor):
         ]
         for otlp_span in otlp_spans:
             serialized = otlp_span.SerializeToString()
-            data = gzip.compress(serialized)
-            self._session.post(
-                urljoin(self._base_url, "v1/traces"),
-                data=data,
+            content = gzip.compress(serialized)
+            self._client.post(
+                url=urljoin(self._base_url, "v1/traces"),
+                content=content,
                 headers={
                     "content-type": "application/x-protobuf",
                     "content-encoding": "gzip",
                 },
             ).raise_for_status()
 
+    def _get_dataset_id_by_name(self, name: str) -> str:
+        """
+         Gets a dataset by name.
+
+         Args:
+             name (str): The name of the dataset.
+             version_id (Optional[str]): The version ID of the dataset. Default None.
+
+        Returns:
+             Dataset: The dataset object.
+        """
+        response = self._client.get(
+            urljoin(self._base_url, "/v1/datasets"),
+            params={"name": name},
+        )
+        response.raise_for_status()
+        if not (records := response.json()["data"]):
+            raise ValueError(f"Failed to query dataset by name: {name}")
+        if len(records) > 1 or not records[0]:
+            raise ValueError(f"Failed to find a single dataset with the given name: {name}")
+        dataset = records[0]
+        return str(dataset["id"])
+
+    def get_dataset(
+        self,
+        *,
+        id: Optional[str] = None,
+        name: Optional[str] = None,
+        version_id: Optional[str] = None,
+    ) -> Dataset:
+        """
+        Gets the dataset for a specific version, or gets the latest version of
+        the dataset if no version is specified.
+
+        Args:
+
+            id (Optional[str]): An ID for the dataset.
+
+            name (Optional[str]): the name for the dataset. If provided, the ID
+            is ignored and the dataset is retrieved by name.
+
+            version_id (Optional[str]): An ID for the version of the dataset, or
+            None.
+
+        Returns:
+            A dataset object.
+        """
+        if name:
+            id = self._get_dataset_id_by_name(name)
+
+        if not id:
+            raise ValueError("Dataset id or name must be provided.")
+
+        response = self._client.get(
+            urljoin(self._base_url, f"/v1/datasets/{quote(id)}/examples"),
+            params={"version-id": version_id} if version_id else None,
+        )
+        response.raise_for_status()
+        data = response.json()["data"]
+        examples = [
+            Example(
+                id=example["id"],
+                input=example["input"],
+                output=example["output"],
+                metadata=example["metadata"],
+                updated_at=datetime.fromisoformat(example["updated_at"]),
+            )
+            for example in data["examples"]
+        ]
+        resolved_dataset_id = data["dataset_id"]
+        resolved_version_id = data["version_id"]
+        return Dataset(
+            id=resolved_dataset_id,
+            version_id=resolved_version_id,
+            examples=examples,
+        )
+
+    def get_dataset_versions(
+        self,
+        dataset_id: str,
+        /,
+        *,
+        limit: Optional[int] = 100,
+    ) -> pd.DataFrame:
+        """
+        Get dataset versions as pandas DataFrame.
+
+        Args:
+            dataset_id (str): dataset ID
+            limit (Optional[int]): maximum number of versions to return,
+                starting from the most recent version
+
+        Returns:
+            pandas DataFrame
+        """
+        url = urljoin(self._base_url, f"v1/datasets/{dataset_id}/versions")
+        response = httpx.get(url=url, params={"limit": limit})
+        response.raise_for_status()
+        if not (records := response.json()["data"]):
+            return pd.DataFrame()
+        df = pd.DataFrame.from_records(records, index="version_id")
+        df["created_at"] = pd.to_datetime(df.created_at)
+        return df
+
+    def download_dataset_examples(
+        self,
+        dataset_id: str,
+        /,
+        *,
+        dataset_version_id: Optional[str] = None,
+    ) -> pd.DataFrame:
+        """
+        Download dataset examples as pandas DataFrame.
+
+        Args:
+            dataset_id (str): dataset ID
+            dataset_version_id (Optional[str]): dataset version ID, if omitted,
+               the latest version is returned.
+
+        Returns:
+            pandas DataFrame
+        """
+        url = f"v1/datasets/{dataset_id}/csv"
+        response = httpx.get(
+            url=urljoin(self._base_url, url),
+            params={"version": dataset_version_id} if dataset_version_id else {},
+        )
+        response.raise_for_status()
+        return pd.read_csv(
+            StringIO(response.content.decode()),
+            index_col="example_id",
+        )
+
+    def upload_dataset(
+        self,
+        table: Union[str, Path, pd.DataFrame],
+        /,
+        *,
+        name: str,
+        input_keys: Iterable[str],
+        output_keys: Iterable[str],
+        metadata_keys: Iterable[str] = (),
+        description: Optional[str] = None,
+        action: Literal["create", "append"] = "create",
+    ) -> Dataset:
+        """
+        Upload examples as dataset to the Phoenix server.
+
+        Args:
+            table (str | Path | pd.DataFrame): Location of a CSV text file, or
+                pandas DataFrame.
+            name: (str): Name of the dataset. Required if action=append.
+            input_keys (Iterable[str]): List of column names used as input keys.
+                input_keys, output_keys, metadata_keys must be disjoint, and must
+                exist in CSV column headers.
+            output_keys (Iterable[str]): List of column names used as output keys.
+                input_keys, output_keys, metadata_keys must be disjoint, and must
+                exist in CSV column headers.
+            metadata_keys (Iterable[str]): List of column names used as metadata keys.
+                input_keys, output_keys, metadata_keys must be disjoint, and must
+                exist in CSV column headers.
+            description: (Optional[str]): Description of the dataset.
+            action: (Literal["create", "append"): Create new dataset or append to an
+                existing dataset. If action=append, dataset name is required.
+
+        Returns:
+            A Dataset object with the uploaded examples.
+        """
+        if action not in ("create", "append"):
+            raise ValueError(f"Invalid action: {action}")
+        if not name:
+            raise ValueError("Dataset name must not be blank")
+        keys = DatasetKeys(
+            frozenset(input_keys),
+            frozenset(output_keys),
+            frozenset(metadata_keys),
+        )
+        if isinstance(table, pd.DataFrame):
+            file = _prepare_pyarrow(table, keys)
+        elif isinstance(table, (str, Path)):
+            file = _prepare_csv(Path(table), keys)
+        else:
+            assert_never(table)
+        response = self._client.post(
+            url=urljoin(self._base_url, "v1/datasets/upload"),
+            files={"file": file},
+            data={
+                "action": action,
+                "name": name,
+                "description": description,
+                "input_keys[]": sorted(keys.input),
+                "output_keys[]": sorted(keys.output),
+                "metadata_keys[]": sorted(keys.metadata),
+            },
+            params={"sync": True},
+        )
+        response.raise_for_status()
+        data = response.json()["data"]
+        dataset_id = data["dataset_id"]
+        response = self._client.get(
+            url=urljoin(self._base_url, f"v1/datasets/{dataset_id}/examples")
+        )
+        response.raise_for_status()
+        data = response.json()["data"]
+        version_id = data["version_id"]
+        examples = data["examples"]
+        return Dataset(
+            id=dataset_id,
+            version_id=version_id,
+            examples=[
+                Example(
+                    id=example["id"],
+                    input=example["input"],
+                    output=example["output"],
+                    metadata=example["metadata"],
+                    updated_at=datetime.fromisoformat(example["updated_at"]),
+                )
+                for example in examples
+            ],
+        )
+
+
+FileName: TypeAlias = str
+FilePointer: TypeAlias = BinaryIO
+FileType: TypeAlias = str
+FileHeaders: TypeAlias = Dict[str, str]
+
+
+def _prepare_csv(
+    path: Path,
+    keys: DatasetKeys,
+) -> Tuple[FileName, FilePointer, FileType, FileHeaders]:
+    path = path.resolve()
+    if not path.is_file():
+        raise FileNotFoundError(f"File does not exist: {path}")
+    with open(path, "r") as f:
+        rows = csv.reader(f)
+        try:
+            column_headers = next(rows)
+            _ = next(rows)
+        except StopIteration:
+            raise ValueError("csv file has no data")
+    (header, freq), *_ = Counter(column_headers).most_common(1)
+    if freq > 1:
+        raise ValueError(f"Duplicated column header in CSV file: {header}")
+    keys.check_differences(frozenset(column_headers))
+    file = BytesIO()
+    with open(path, "rb") as f:
+        file.write(gzip.compress(f.read()))
+    return path.name, file, "text/csv", {"Content-Encoding": "gzip"}
+
+
+def _prepare_pyarrow(
+    df: pd.DataFrame,
+    keys: DatasetKeys,
+) -> Tuple[FileName, FilePointer, FileType, FileHeaders]:
+    if df.empty:
+        raise ValueError("dataframe has no data")
+    (header, freq), *_ = Counter(df.columns).most_common(1)
+    if freq > 1:
+        raise ValueError(f"Duplicated column header in file: {header}")
+    keys.check_differences(frozenset(df.columns))
+    table = Table.from_pandas(df.loc[:, list(keys)])
+    sink = pa.BufferOutputStream()
+    options = pa.ipc.IpcWriteOptions(compression="lz4")
+    with pa.ipc.new_stream(sink, table.schema, options=options) as writer:
+        writer.write_table(table)
+    file = BytesIO(sink.getvalue().to_pybytes())
+    return "pandas", file, "application/x-pandas-pyarrow", {}
+
 
 def _to_iso_format(value: Optional[datetime]) -> Optional[str]:
     return value.isoformat() if value else None

phoenix/session/session.py

@@ -37,10 +37,16 @@ from phoenix.config import (
     get_exported_files,
     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.inferences.inferences import EMPTY_INFERENCES, Inferences
 from phoenix.pointcloud.umap_parameters import get_umap_parameters
-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.server.thread_server import ThreadServer
 from phoenix.services import AppService
 from phoenix.session.client import Client
@@ -108,9 +114,9 @@ class Session(TraceDataExtractor, ABC):
     def __init__(
         self,
         database_url: str,
-        primary_dataset: Inferences,
-        reference_dataset: Optional[Inferences] = None,
-        corpus_dataset: Optional[Inferences] = None,
+        primary_inferences: Inferences,
+        reference_inferences: Optional[Inferences] = None,
+        corpus_inferences: Optional[Inferences] = None,
         trace_dataset: Optional[TraceDataset] = None,
         default_umap_parameters: Optional[Mapping[str, Any]] = None,
         host: Optional[str] = None,
@@ -118,9 +124,9 @@ class Session(TraceDataExtractor, ABC):
         notebook_env: Optional[NotebookEnvironment] = None,
     ):
         self._database_url = database_url
-        self.primary_dataset = primary_dataset
-        self.reference_dataset = reference_dataset
-        self.corpus_dataset = corpus_dataset
+        self.primary_inferences = primary_inferences
+        self.reference_inferences = reference_inferences
+        self.corpus_inferences = corpus_inferences
         self.trace_dataset = trace_dataset
         self.umap_parameters = get_umap_parameters(default_umap_parameters)
         self.host = host or get_env_host()
@@ -264,9 +270,9 @@ class ProcessSession(Session):
     def __init__(
         self,
         database_url: str,
-        primary_dataset: Inferences,
-        reference_dataset: Optional[Inferences] = None,
-        corpus_dataset: Optional[Inferences] = None,
+        primary_inferences: Inferences,
+        reference_inferences: Optional[Inferences] = None,
+        corpus_inferences: Optional[Inferences] = None,
         trace_dataset: Optional[TraceDataset] = None,
         default_umap_parameters: Optional[Mapping[str, Any]] = None,
         host: Optional[str] = None,
@@ -276,20 +282,20 @@ class ProcessSession(Session):
     ) -> None:
         super().__init__(
             database_url=database_url,
-            primary_dataset=primary_dataset,
-            reference_dataset=reference_dataset,
-            corpus_dataset=corpus_dataset,
+            primary_inferences=primary_inferences,
+            reference_inferences=reference_inferences,
+            corpus_inferences=corpus_inferences,
             trace_dataset=trace_dataset,
             default_umap_parameters=default_umap_parameters,
             host=host,
             port=port,
             notebook_env=notebook_env,
         )
-        primary_dataset.to_disc()
-        if isinstance(reference_dataset, Inferences):
-            reference_dataset.to_disc()
-        if isinstance(corpus_dataset, Inferences):
-            corpus_dataset.to_disc()
+        primary_inferences.to_disc()
+        if isinstance(reference_inferences, Inferences):
+            reference_inferences.to_disc()
+        if isinstance(corpus_inferences, Inferences):
+            corpus_inferences.to_disc()
         if isinstance(trace_dataset, TraceDataset):
             trace_dataset.to_disc()
         umap_params_str = (
@@ -304,13 +310,13 @@ class ProcessSession(Session):
             host=self.host,
             port=self.port,
             root_path=self.root_path,
-            primary_dataset_name=self.primary_dataset.name,
+            primary_inferences_name=self.primary_inferences.name,
             umap_params=umap_params_str,
-            reference_dataset_name=(
-                self.reference_dataset.name if self.reference_dataset is not None else None
+            reference_inferences_name=(
+                self.reference_inferences.name if self.reference_inferences is not None else None
             ),
-            corpus_dataset_name=(
-                self.corpus_dataset.name if self.corpus_dataset is not None else None
+            corpus_inferences_name=(
+                self.corpus_inferences.name if self.corpus_inferences is not None else None
             ),
             trace_dataset_name=(
                 self.trace_dataset.name if self.trace_dataset is not None else None
@@ -330,9 +336,9 @@ class ThreadSession(Session):
     def __init__(
         self,
         database_url: str,
-        primary_dataset: Inferences,
-        reference_dataset: Optional[Inferences] = None,
-        corpus_dataset: Optional[Inferences] = None,
+        primary_inferences: Inferences,
+        reference_inferences: Optional[Inferences] = None,
+        corpus_inferences: Optional[Inferences] = None,
         trace_dataset: Optional[TraceDataset] = None,
         default_umap_parameters: Optional[Mapping[str, Any]] = None,
         host: Optional[str] = None,
@@ -342,29 +348,32 @@ class ThreadSession(Session):
     ):
         super().__init__(
             database_url=database_url,
-            primary_dataset=primary_dataset,
-            reference_dataset=reference_dataset,
-            corpus_dataset=corpus_dataset,
+            primary_inferences=primary_inferences,
+            reference_inferences=reference_inferences,
+            corpus_inferences=corpus_inferences,
             trace_dataset=trace_dataset,
             default_umap_parameters=default_umap_parameters,
             host=host,
             port=port,
             notebook_env=notebook_env,
         )
-        self.model = create_model_from_datasets(
-            primary_dataset,
-            reference_dataset,
+        self.model = create_model_from_inferences(
+            primary_inferences,
+            reference_inferences,
         )
         self.corpus = (
-            create_model_from_datasets(
-                corpus_dataset,
+            create_model_from_inferences(
+                corpus_inferences,
             )
-            if corpus_dataset is not None
+            if corpus_inferences is not None
             else None
         )
         # Initialize an app service that keeps the server running
+        engine = create_engine_and_run_migrations(database_url)
+        instrumentation_cleanups = instrument_engine_if_enabled(engine)
+        factory = SessionFactory(session_factory=_db(engine), dialect=engine.dialect.name)
         self.app = create_app(
-            database_url=database_url,
+            db=factory,
             export_path=self.export_path,
             model=self.model,
             corpus=self.corpus,
@@ -375,6 +384,7 @@ class ThreadSession(Session):
                 if (trace_dataset and (initial_evaluations := trace_dataset.evaluations))
                 else None
             ),
+            clean_up_callbacks=instrumentation_cleanups,
         )
         self.server = ThreadServer(
             app=self.app,