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

phoenix/trace/exporter.py

@@ -7,13 +7,17 @@ from types import MethodType
 from typing import Any, Optional
 from urllib.parse import urljoin
 
-import requests
+import httpx
 from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter
-from requests import Session
 from typing_extensions import TypeAlias, assert_never
 
 import phoenix.trace.v1 as pb
-from phoenix.config import get_env_collector_endpoint, get_env_host, get_env_port
+from phoenix.config import (
+    get_env_client_headers,
+    get_env_collector_endpoint,
+    get_env_host,
+    get_env_port,
+)
 
 logger = logging.getLogger(__name__)
 logger.addHandler(logging.NullHandler())
@@ -75,9 +79,10 @@ class HttpExporter:
         )
         self._base_url = base_url if base_url.endswith("/") else base_url + "/"
         _warn_if_phoenix_is_not_running(self._base_url)
-        self._session = Session()
-        weakref.finalize(self, self._session.close)
-        self._session.headers.update(
+        headers = get_env_client_headers()
+        self._client = httpx.Client(headers=headers)
+        weakref.finalize(self, self._client.close)
+        self._client.headers.update(
             {
                 "content-type": "application/x-protobuf",
                 "content-encoding": "gzip",
@@ -110,9 +115,9 @@ class HttpExporter:
 
     def _send(self, message: Message) -> None:
         serialized = message.SerializeToString()
-        data = gzip.compress(serialized)
+        content = gzip.compress(serialized)
         try:
-            self._session.post(self._url(message), data=data).raise_for_status()
+            self._client.post(self._url(message), content=content).raise_for_status()
         except Exception as e:
             logger.exception(e)
 
@@ -125,7 +130,7 @@ class HttpExporter:
 
 def _warn_if_phoenix_is_not_running(base_url: str) -> None:
     try:
-        requests.get(urljoin(base_url, "arize_phoenix_version")).raise_for_status()
+        httpx.get(urljoin(base_url, "arize_phoenix_version")).raise_for_status()
     except Exception:
         logger.warning(
             f"Arize Phoenix is not running on {base_url}. Launch Phoenix "

phoenix/trace/fixtures.py

@@ -1,18 +1,29 @@
+import logging
+import shutil
 from binascii import hexlify
 from dataclasses import dataclass, field, replace
 from datetime import datetime, timezone
+from io import StringIO
 from random import getrandbits
-from typing import Dict, Iterable, Iterator, List, NamedTuple, Optional, Tuple, cast
+from tempfile import NamedTemporaryFile
+from time import sleep, time
+from typing import Dict, Iterable, Iterator, List, NamedTuple, Optional, Sequence, Tuple, cast
 from urllib import request
+from urllib.parse import urljoin
 
+import httpx
 import pandas as pd
 from google.protobuf.wrappers_pb2 import DoubleValue, StringValue
+from httpx import ConnectError, HTTPStatusError
 
 import phoenix.trace.v1 as pb
+from phoenix import Client
 from phoenix.trace.schemas import Span
 from phoenix.trace.trace_dataset import TraceDataset
 from phoenix.trace.utils import json_lines_to_df
 
+logger = logging.getLogger(__name__)
+
 
 class EvaluationResultSchema(NamedTuple):
     label: Optional[str] = "label"
@@ -32,12 +43,44 @@ class DocumentEvaluationFixture(EvaluationFixture):
     document_position: str = "document_position"
 
 
+@dataclass(frozen=True)
+class DatasetFixture:
+    file_name: str
+    name: str
+    input_keys: Sequence[str]
+    output_keys: Sequence[str]
+    metadata_keys: Sequence[str] = ()
+    description: Optional[str] = field(default=None)
+    _df: Optional[pd.DataFrame] = field(default=None, init=False, repr=False)
+    _csv: Optional[str] = field(default=None, init=False, repr=False)
+
+    def load(self) -> "DatasetFixture":
+        if self._df is None:
+            df = pd.read_csv(_url(self.file_name))
+            object.__setattr__(self, "_df", df)
+        return self
+
+    @property
+    def dataframe(self) -> pd.DataFrame:
+        self.load()
+        return cast(pd.DataFrame, self._df).copy(deep=False)
+
+    @property
+    def csv(self) -> StringIO:
+        if self._csv is None:
+            with StringIO() as buffer:
+                self.dataframe.to_csv(buffer, index=False)
+                object.__setattr__(self, "_csv", buffer.getvalue())
+        return StringIO(self._csv)
+
+
 @dataclass(frozen=True)
 class TracesFixture:
     name: str
     description: str
     file_name: str
     evaluation_fixtures: Iterable[EvaluationFixture] = ()
+    dataset_fixtures: Iterable[DatasetFixture] = ()
 
 
 llama_index_rag_fixture = TracesFixture(
@@ -58,6 +101,36 @@ llama_index_rag_fixture = TracesFixture(
             file_name="llama_index_rag_v8.retrieved_documents_eval.parquet",
         ),
     ),
+    dataset_fixtures=(
+        DatasetFixture(
+            file_name="hybridial_samples.csv.gz",
+            input_keys=("messages", "ctxs"),
+            output_keys=("answers",),
+            name="ChatRAG-Bench: Hybrid Dialogue (samples)",
+            description="https://huggingface.co/datasets/nvidia/ChatRAG-Bench/viewer/hybridial",
+        ),
+        DatasetFixture(
+            file_name="sqa_samples.csv.gz",
+            input_keys=("messages", "ctxs"),
+            output_keys=("answers",),
+            name="ChatRAG-Bench: SQA (samples)",
+            description="https://huggingface.co/datasets/nvidia/ChatRAG-Bench/viewer/sqa",
+        ),
+        DatasetFixture(
+            file_name="doqa_cooking_samples.csv.gz",
+            input_keys=("messages", "ctxs"),
+            output_keys=("answers",),
+            name="ChatRAG-Bench: DoQA Cooking (samples)",
+            description="https://huggingface.co/datasets/nvidia/ChatRAG-Bench/viewer/doqa_cooking",
+        ),
+        DatasetFixture(
+            file_name="synthetic_convqa_samples.csv.gz",
+            input_keys=("messages", "document"),
+            output_keys=("answers",),
+            name="ChatQA-Train: Synthetic ConvQA (samples)",
+            description="https://huggingface.co/datasets/nvidia/ChatQA-Training-Data/viewer/synthetic_convqa",
+        ),
+    ),
 )
 
 llama_index_calculator_agent_fixture = TracesFixture(
@@ -138,16 +211,69 @@ def download_traces_fixture(
         return cast(List[str], f.readlines())
 
 
-def load_example_traces(use_case: str) -> TraceDataset:
+def load_example_traces(fixture_name: str) -> TraceDataset:
     """
     Loads a trace dataframe by name.
     """
-    fixture = get_trace_fixture_by_name(use_case)
+    fixture = get_trace_fixture_by_name(fixture_name)
     return TraceDataset(json_lines_to_df(download_traces_fixture(fixture)))
 
 
-def get_evals_from_fixture(use_case: str) -> Iterator[pb.Evaluation]:
-    fixture = get_trace_fixture_by_name(use_case)
+def get_dataset_fixtures(fixture_name: str) -> Iterable[DatasetFixture]:
+    return (fixture.load() for fixture in get_trace_fixture_by_name(fixture_name).dataset_fixtures)
+
+
+def send_dataset_fixtures(
+    endpoint: str,
+    fixtures: Iterable[DatasetFixture],
+) -> None:
+    expiration = time() + 5
+    while time() < expiration:
+        try:
+            url = urljoin(endpoint, "/healthz")
+            httpx.get(url=url).raise_for_status()
+        except ConnectError:
+            sleep(0.1)
+            continue
+        except Exception as e:
+            print(str(e))
+            raise
+        break
+    client = Client(endpoint=endpoint)
+    for i, fixture in enumerate(fixtures):
+        try:
+            if i % 2:
+                client.upload_dataset(
+                    fixture.dataframe,
+                    name=fixture.name,
+                    input_keys=fixture.input_keys,
+                    output_keys=fixture.output_keys,
+                    metadata_keys=fixture.metadata_keys,
+                    description=fixture.description,
+                )
+            else:
+                with NamedTemporaryFile() as tf:
+                    with open(tf.name, "w") as f:
+                        shutil.copyfileobj(fixture.csv, f)
+                        f.flush()
+                    client.upload_dataset(
+                        tf.name,
+                        name=fixture.name,
+                        input_keys=fixture.input_keys,
+                        output_keys=fixture.output_keys,
+                        metadata_keys=fixture.metadata_keys,
+                        description=fixture.description,
+                    )
+        except HTTPStatusError as e:
+            print(e.response.content.decode())
+            pass
+        else:
+            name, df = fixture.name, fixture.dataframe
+            print(f"Dataset sent: {name=}, {len(df)=}")
+
+
+def get_evals_from_fixture(fixture_name: str) -> Iterator[pb.Evaluation]:
+    fixture = get_trace_fixture_by_name(fixture_name)
     for eval_fixture in fixture.evaluation_fixtures:
         yield from _read_eval_fixture(eval_fixture)
 
@@ -195,8 +321,8 @@ def _read_eval_fixture(eval_fixture: EvaluationFixture) -> Iterator[pb.Evaluatio
 def _url(
     file_name: str,
     host: Optional[str] = "https://storage.googleapis.com/",
-    bucket: Optional[str] = "arize-assets",
-    prefix: Optional[str] = "phoenix/traces/",
+    bucket: Optional[str] = "arize-phoenix-assets",
+    prefix: Optional[str] = "traces/",
 ) -> str:
     return f"{host}{bucket}/{prefix}{file_name}"
 

phoenix/trace/span_evaluations.py

@@ -11,7 +11,7 @@ import pandas as pd
 from pandas.api.types import is_integer_dtype, is_numeric_dtype, is_string_dtype
 from pyarrow import RecordBatchStreamReader, Schema, Table, parquet
 
-from phoenix.config import TRACE_DATASET_DIR
+from phoenix.config import TRACE_DATASETS_DIR
 from phoenix.exceptions import PhoenixEvaluationNameIsMissing
 from phoenix.trace.errors import InvalidParquetMetadataError
 
@@ -201,7 +201,7 @@ class Evaluations(NeedsNamedIndex, NeedsResultColumns, ABC):
             UUID: The ID of the evaluations, which can be used as a key to load
             the evaluations from disk using `load`.
         """
-        directory = Path(directory) if directory else TRACE_DATASET_DIR
+        directory = Path(directory) if directory else TRACE_DATASETS_DIR
         path = directory / EVAL_PARQUET_FILE_NAME.format(id=self.id)
         table = self.to_pyarrow_table()
         parquet.write_table(table, path)
@@ -229,7 +229,7 @@ class Evaluations(NeedsNamedIndex, NeedsResultColumns, ABC):
         """
         if not isinstance(id, UUID):
             id = UUID(id)
-        path = Path(directory or TRACE_DATASET_DIR) / EVAL_PARQUET_FILE_NAME.format(id=id)
+        path = Path(directory or TRACE_DATASETS_DIR) / EVAL_PARQUET_FILE_NAME.format(id=id)
         schema = parquet.read_schema(path)
         eval_id, eval_name, evaluations_cls = _parse_schema_metadata(schema)
         if id != eval_id:

phoenix/trace/trace_dataset.py

@@ -14,7 +14,7 @@ from openinference.semconv.trace import (
 from pandas import DataFrame, read_parquet
 from pyarrow import Schema, Table, parquet
 
-from phoenix.config import DATASET_DIR, GENERATED_DATASET_NAME_PREFIX, TRACE_DATASET_DIR
+from phoenix.config import GENERATED_INFERENCES_NAME_PREFIX, INFERENCES_DIR, TRACE_DATASETS_DIR
 from phoenix.datetime_utils import normalize_timestamps
 from phoenix.trace.attributes import unflatten
 from phoenix.trace.errors import InvalidParquetMetadataError
@@ -138,7 +138,7 @@ class TraceDataset:
         self._id = uuid4()
         self.dataframe = normalize_dataframe(dataframe)
         # TODO: This is not used in any meaningful way. Should remove
-        self.name = name or f"{GENERATED_DATASET_NAME_PREFIX}{str(self._id)}"
+        self.name = name or f"{GENERATED_INFERENCES_NAME_PREFIX}{str(self._id)}"
         self.evaluations = list(evaluations)
 
     @classmethod
@@ -201,13 +201,13 @@ class TraceDataset:
     @classmethod
     def from_name(cls, name: str) -> "TraceDataset":
         """Retrieves a dataset by name from the file system"""
-        directory = DATASET_DIR / name
+        directory = INFERENCES_DIR / name
         df = read_parquet(directory / cls._data_file_name)
         return cls(df, name)
 
     def to_disc(self) -> None:
         """writes the data to disc"""
-        directory = DATASET_DIR / self.name
+        directory = INFERENCES_DIR / self.name
         directory.mkdir(parents=True, exist_ok=True)
         get_serializable_spans_dataframe(self.dataframe).to_parquet(
             directory / self._data_file_name,
@@ -230,7 +230,7 @@ class TraceDataset:
             UUID: The id of the trace dataset, which can be used as key to load
             the dataset from disk using `load`.
         """
-        directory = Path(directory or TRACE_DATASET_DIR)
+        directory = Path(directory or TRACE_DATASETS_DIR)
         for evals in self.evaluations:
             evals.save(directory)
         path = directory / TRACE_DATASET_PARQUET_FILE_NAME.format(id=self._id)
@@ -280,7 +280,7 @@ class TraceDataset:
         """
         if not isinstance(id, UUID):
             id = UUID(id)
-        path = Path(directory or TRACE_DATASET_DIR) / TRACE_DATASET_PARQUET_FILE_NAME.format(id=id)
+        path = Path(directory or TRACE_DATASETS_DIR) / TRACE_DATASET_PARQUET_FILE_NAME.format(id=id)
         schema = parquet.read_schema(path)
         dataset_id, dataset_name, eval_ids = _parse_schema_metadata(schema)
         if id != dataset_id:

phoenix/utilities/json.py

@@ -0,0 +1,61 @@
+import dataclasses
+import datetime
+from enum import Enum
+from pathlib import Path
+from typing import Any, Mapping, Sequence, SupportsFloat, Union, get_args, get_origin
+
+import numpy as np
+
+
+def jsonify(obj: Any) -> Any:
+    """
+    Coerce object to be json serializable.
+    """
+    if isinstance(obj, Enum):
+        return jsonify(obj.value)
+    if isinstance(obj, (str, int, float, bool)) or obj is None:
+        return obj
+    if isinstance(obj, np.ndarray):
+        return [jsonify(v) for v in obj]
+    if isinstance(obj, SupportsFloat):
+        return float(obj)
+    if dataclasses.is_dataclass(obj):
+        return {
+            k: jsonify(v)
+            for field in dataclasses.fields(obj)
+            if not (
+                (v := getattr(obj, (k := field.name))) is None
+                and get_origin(field) is Union
+                and type(None) in get_args(field)
+            )
+        }
+    if isinstance(obj, (Sequence, set, frozenset)):
+        return [jsonify(v) for v in obj]
+    if isinstance(obj, Mapping):
+        return {jsonify(k): jsonify(v) for k, v in obj.items()}
+    if isinstance(obj, (datetime.date, datetime.datetime, datetime.time)):
+        return obj.isoformat()
+    if isinstance(obj, datetime.timedelta):
+        return obj.total_seconds()
+    if isinstance(obj, Path):
+        return str(obj)
+    if isinstance(obj, BaseException):
+        return str(obj)
+    if hasattr(obj, "model_dump") and callable(obj.model_dump):
+        # pydantic v2
+        try:
+            assert isinstance(d := obj.model_dump(), dict)
+        except BaseException:
+            pass
+        else:
+            return jsonify(d)
+    if hasattr(obj, "dict") and callable(obj.dict):
+        # pydantic v1
+        try:
+            assert isinstance(d := obj.dict(), dict)
+        except BaseException:
+            pass
+        else:
+            return jsonify(d)
+    cls = obj.__class__
+    return f"<{cls.__module__}.{cls.__name__} object>"

phoenix/utilities/re.py

@@ -0,0 +1,50 @@
+from logging import getLogger
+from re import compile, split
+from typing import Dict, List
+from urllib.parse import unquote
+
+_logger = getLogger(__name__)
+
+# Optional whitespace
+_OWS = r"[ \t]*"
+# A key contains printable US-ASCII characters except: SP and "(),/:;<=>?@[\]{}
+_KEY_FORMAT = r"[\x21\x23-\x27\x2a\x2b\x2d\x2e\x30-\x39\x41-\x5a\x5e-\x7a\x7c\x7e]+"
+# A value contains a URL-encoded UTF-8 string. The encoded form can contain any
+# printable US-ASCII characters (0x20-0x7f) other than SP, DEL, and ",;/
+_VALUE_FORMAT = r"[\x21\x23-\x2b\x2d-\x3a\x3c-\x5b\x5d-\x7e]*"
+# A key-value is key=value, with optional whitespace surrounding key and value
+_KEY_VALUE_FORMAT = rf"{_OWS}{_KEY_FORMAT}{_OWS}={_OWS}{_VALUE_FORMAT}{_OWS}"
+
+_HEADER_PATTERN = compile(_KEY_VALUE_FORMAT)
+_DELIMITER_PATTERN = compile(r"[ \t]*,[ \t]*")
+
+
+def parse_env_headers(s: str) -> Dict[str, str]:
+    """
+    Parse ``s``, which is a ``str`` instance containing HTTP headers encoded
+    for use in ENV variables per the W3C Baggage HTTP header format at
+    https://www.w3.org/TR/baggage/#baggage-http-header-format, except that
+    additional semi-colon delimited metadata is not supported.
+
+    src: https://github.com/open-telemetry/opentelemetry-python/blob/2d5cd58f33bd8a16f45f30be620a96699bc14297/opentelemetry-api/src/opentelemetry/util/re.py#L52
+    """
+    headers: Dict[str, str] = {}
+    headers_list: List[str] = split(_DELIMITER_PATTERN, s)
+    for header in headers_list:
+        if not header:  # empty string
+            continue
+        match = _HEADER_PATTERN.fullmatch(header.strip())
+        if not match:
+            _logger.warning(
+                "Header format invalid! Header values in environment variables must be "
+                "URL encoded: %s",
+                header,
+            )
+            continue
+        # value may contain any number of `=`
+        name, value = match.string.split("=", 1)
+        name = unquote(name).strip().lower()
+        value = unquote(value).strip()
+        headers[name] = value
+
+    return headers

phoenix/version.py

@@ -1 +1 @@
-__version__ = "4.4.2"
+__version__ = "4.4.4rc0"

phoenix/server/api/types/DatasetRole.py

@@ -1,23 +0,0 @@
-from enum import Enum
-from typing import Dict, Union
-
-import strawberry
-
-from phoenix.core.model_schema import PRIMARY, REFERENCE
-
-
-@strawberry.enum
-class DatasetRole(Enum):
-    primary = PRIMARY
-    reference = REFERENCE
-
-
-class AncillaryDatasetRole(Enum):
-    corpus = "DatasetRole.CORPUS"
-
-
-STR_TO_DATASET_ROLE: Dict[str, Union[DatasetRole, AncillaryDatasetRole]] = {
-    str(DatasetRole.primary.value): DatasetRole.primary,
-    str(DatasetRole.reference.value): DatasetRole.reference,
-    str(AncillaryDatasetRole.corpus.value): AncillaryDatasetRole.corpus,
-}