arize-phoenix 4.30.2__py3-none-any.whl → 4.31.0__py3-none-any.whl

phoenix/trace/fixtures.py

@@ -8,7 +8,18 @@ from io import StringIO
 from random import getrandbits
 from tempfile import NamedTemporaryFile
 from time import sleep, time
-from typing import Dict, Iterable, Iterator, List, NamedTuple, Optional, Sequence, Tuple, cast
+from typing import (
+    DefaultDict,
+    Dict,
+    Iterable,
+    Iterator,
+    List,
+    NamedTuple,
+    Optional,
+    Sequence,
+    Tuple,
+    cast,
+)
 from urllib.parse import urljoin
 
 import httpx
@@ -20,7 +31,11 @@ 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 download_json_traces_fixture, is_jsonl_file, json_lines_to_df
+from phoenix.trace.utils import (
+    download_json_traces_fixture,
+    json_lines_to_df,
+    parse_file_extension,
+)
 
 logger = logging.getLogger(__name__)
 logger.setLevel(logging.INFO)
@@ -58,7 +73,23 @@ class DatasetFixture:
 
     def load(self) -> "DatasetFixture":
         if self._df is None:
-            df = pd.read_csv(_url(self.file_name))
+            url = _url(self.file_name)
+
+            if parse_file_extension(self.file_name) == ".jsonl":
+                df = json_lines_to_df(download_json_traces_fixture(url))
+            elif parse_file_extension(self.file_name) == ".csv":
+                df = pd.read_csv(_url(self.file_name))
+            else:
+                try:
+                    df = pd.read_parquet(url)
+                except Exception:
+                    logger.warning(
+                        f"Failed to download example traces from {url=} "
+                        "due to exception {e=}. "
+                        "Returning empty dataframe for DatasetFixture"
+                    )
+                    df = pd.DataFrame()
+
             object.__setattr__(self, "_df", df)
         return self
 
@@ -105,6 +136,15 @@ demo_llama_index_rag_fixture = TracesFixture(
             file_name="demo_llama_index_rag_doc_relevance_eval.parquet",
         ),
     ),
+    dataset_fixtures=(
+        DatasetFixture(
+            file_name="demo_llama_index_finetune_dataset.jsonl",
+            input_keys=("messages",),
+            output_keys=("messages",),
+            name="Demo LlamaIndex: RAG Q&A",
+            description="OpenAI GPT-3.5 LLM dataset for LlamaIndex demo",
+        ),
+    ),
 )
 
 demo_llama_index_rag_llm_fixture = TracesFixture(
@@ -225,7 +265,7 @@ TRACES_FIXTURES: List[TracesFixture] = [
 NAME_TO_TRACES_FIXTURE: Dict[str, TracesFixture] = {
     fixture.name: fixture for fixture in TRACES_FIXTURES
 }
-PROJ_NAME_TO_TRACES_FIXTURE = defaultdict(list)
+PROJ_NAME_TO_TRACES_FIXTURE: DefaultDict[str, List[TracesFixture]] = defaultdict(list)
 for fixture in TRACES_FIXTURES:
     if fixture.project_name:
         PROJ_NAME_TO_TRACES_FIXTURE[fixture.project_name].append(fixture)
@@ -268,14 +308,10 @@ def load_example_traces(fixture_name: str) -> TraceDataset:
     """
     Loads a trace dataframe by name.
     """
-    host = "https://storage.googleapis.com/"
-    bucket = "arize-phoenix-assets"
-    prefix = "traces/"
     fixture = get_trace_fixture_by_name(fixture_name)
+    url = _url(fixture.file_name)
 
-    url = f"{host}{bucket}/{prefix}{fixture.file_name}"
-
-    if is_jsonl_file(fixture.file_name):
+    if parse_file_extension(fixture.file_name) == ".jsonl":
         return TraceDataset(json_lines_to_df(download_json_traces_fixture(url)))
 
     try:

phoenix/trace/utils.py

@@ -8,14 +8,8 @@ from urllib import request
 import pandas as pd
 
 
-def is_jsonl_file(file_path: str) -> bool:
-    """
-    Check if the given file is a Parquet file.
-    """
-    file_extension = os.path.splitext(file_path)[-1]
-    if file_extension == ".jsonl":
-        return True
-    return False
+def parse_file_extension(file_path: str) -> str:
+    return os.path.splitext(file_path)[-1]
 
 
 def download_json_traces_fixture(

phoenix/version.py

@@ -1 +1 @@
-__version__ = "4.30.2"
+__version__ = "4.31.0"