arize-phoenix 12.8.0__tar.gz → 12.9.0__tar.gz

{arize_phoenix-12.8.0 → arize_phoenix-12.9.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: arize-phoenix
-Version: 12.8.0
+Version: 12.9.0
 Summary: AI Observability and Evaluation
 Project-URL: Documentation, https://arize.com/docs/phoenix/
 Project-URL: Issues, https://github.com/Arize-ai/phoenix/issues
@@ -59,6 +59,8 @@ Requires-Dist: tqdm
 Requires-Dist: typing-extensions>=4.6
 Requires-Dist: uvicorn
 Requires-Dist: wrapt<2,>=1.17.2
+Provides-Extra: aws
+Requires-Dist: boto3; extra == 'aws'
 Provides-Extra: container
 Requires-Dist: aiohttp; extra == 'container'
 Requires-Dist: anthropic>=0.49.0; extra == 'container'

{arize_phoenix-12.8.0 → arize_phoenix-12.9.0}/pyproject.toml

@@ -120,6 +120,9 @@ pg = [
   "asyncpg",
   "psycopg[binary,pool]",
 ]
+aws = [
+  "boto3",
+]
 container = [
   "anthropic>=0.49.0",
   "google-genai",

{arize_phoenix-12.8.0 → arize_phoenix-12.9.0}/src/phoenix/config.py

@@ -91,16 +91,44 @@ Used with PHOENIX_POSTGRES_HOST to specify the port to use for the PostgreSQL da
 ENV_PHOENIX_POSTGRES_USER = "PHOENIX_POSTGRES_USER"
 """
 Used with PHOENIX_POSTGRES_HOST to specify the user to use for the PostgreSQL database (required).
+
+When using AWS RDS IAM authentication (PHOENIX_POSTGRES_USE_AWS_IAM_AUTH=true), this should be
+set to the IAM-enabled database username configured in your RDS/Aurora instance.
 """
 ENV_PHOENIX_POSTGRES_PASSWORD = "PHOENIX_POSTGRES_PASSWORD"
 """
 Used with PHOENIX_POSTGRES_HOST to specify the password to use for the PostgreSQL database
-(required).
+(required, unless PHOENIX_POSTGRES_USE_AWS_IAM_AUTH is enabled).
+
+When using AWS RDS IAM authentication (PHOENIX_POSTGRES_USE_AWS_IAM_AUTH=true), this password
+is NOT used. Instead, authentication tokens are generated dynamically using AWS IAM credentials.
 """
 ENV_PHOENIX_POSTGRES_DB = "PHOENIX_POSTGRES_DB"
 """
 Used with PHOENIX_POSTGRES_HOST to specify the database to use for the PostgreSQL database.
 """
+ENV_PHOENIX_POSTGRES_USE_AWS_IAM_AUTH = "PHOENIX_POSTGRES_USE_AWS_IAM_AUTH"
+"""
+Enable AWS RDS IAM database authentication. When enabled, Phoenix will use AWS IAM credentials
+to generate short-lived authentication tokens instead of using a static password.
+
+This requires:
+- boto3 to be installed: pip install 'arize-phoenix[aws]'
+- AWS credentials configured (via environment, ~/.aws/credentials, or IAM role)
+- AWS region configured via standard AWS methods
+- The database user to be configured for IAM authentication in RDS/Aurora
+- SSL to be enabled (required by AWS RDS IAM auth)
+
+When enabled, PHOENIX_POSTGRES_PASSWORD should NOT be set.
+"""
+ENV_PHOENIX_POSTGRES_AWS_IAM_TOKEN_LIFETIME_SECONDS = (
+    "PHOENIX_POSTGRES_AWS_IAM_TOKEN_LIFETIME_SECONDS"
+)
+"""
+Token lifetime in seconds for connection pool recycling when using AWS RDS IAM authentication.
+AWS RDS auth tokens are valid for 15 minutes. This should be set slightly lower to ensure
+tokens are refreshed before expiration. Defaults to 840 seconds (14 minutes).
+"""
 ENV_PHOENIX_SQL_DATABASE_SCHEMA = "PHOENIX_SQL_DATABASE_SCHEMA"
 """
 The schema to use for the PostgresSQL database. (This is ignored for SQLite.)
@@ -1400,18 +1428,36 @@ def get_env_postgres_connection_str() -> Optional[str]:
     """
     Build PostgreSQL connection string from environment variables.
     """
-    if not (
-        (pg_host := getenv(ENV_PHOENIX_POSTGRES_HOST, "").rstrip("/"))
-        and (pg_user := getenv(ENV_PHOENIX_POSTGRES_USER))
-        and (pg_password := getenv(ENV_PHOENIX_POSTGRES_PASSWORD))
-    ):
+    pg_host = getenv(ENV_PHOENIX_POSTGRES_HOST, "").rstrip("/")
+    pg_user = getenv(ENV_PHOENIX_POSTGRES_USER)
+    pg_password = getenv(ENV_PHOENIX_POSTGRES_PASSWORD)
+    use_iam_auth = _bool_val(ENV_PHOENIX_POSTGRES_USE_AWS_IAM_AUTH, False)
+
+    if not (pg_host and pg_user):
         return None
+
+    if use_iam_auth:
+        if pg_password:
+            raise ValueError(
+                f"The environment variable {ENV_PHOENIX_POSTGRES_PASSWORD} is set but will be "
+                "ignored when using AWS RDS IAM authentication "
+                f"({ENV_PHOENIX_POSTGRES_USE_AWS_IAM_AUTH}=true). Authentication tokens will be "
+                "generated using AWS credentials."
+            )
+        connection_str = f"postgresql://{quote(pg_user)}@{pg_host}"
+    else:
+        if not pg_password:
+            raise ValueError(
+                f"The environment variable {ENV_PHOENIX_POSTGRES_PASSWORD} is not set. "
+                "Please set it to the password for the PostgreSQL database."
+            )
+        encoded_user = quote(pg_user)
+        encoded_password = quote(pg_password)
+        connection_str = f"postgresql://{encoded_user}:{encoded_password}@{pg_host}"
+
     pg_port = getenv(ENV_PHOENIX_POSTGRES_PORT)
     pg_db = getenv(ENV_PHOENIX_POSTGRES_DB)
 
-    encoded_user = quote(pg_user)
-    encoded_password = quote(pg_password)
-    connection_str = f"postgresql://{encoded_user}:{encoded_password}@{pg_host}"
     if pg_port:
         connection_str = f"{connection_str}:{pg_port}"
     if pg_db:
@@ -1968,6 +2014,7 @@ def verify_server_environment_variables() -> None:
     get_env_database_usage_insertion_blocking_threshold_percentage()
     get_env_max_spans_queue_size()
     validate_env_support_email()
+    _validate_iam_auth_config()
 
     # Notify users about deprecated environment variables if they are being used.
     if os.getenv("PHOENIX_ENABLE_WEBSOCKETS") is not None:
@@ -2026,3 +2073,78 @@ def get_env_allow_external_resources() -> bool:
     Defaults to True if not set.
     """
     return _bool_val(ENV_PHOENIX_ALLOW_EXTERNAL_RESOURCES, True)
+
+
+def get_env_postgres_use_iam_auth() -> bool:
+    """
+    Gets whether AWS RDS IAM authentication is enabled for PostgreSQL connections.
+
+    Returns:
+        bool: True if IAM authentication should be used, False otherwise (default)
+    """
+    return _bool_val(ENV_PHOENIX_POSTGRES_USE_AWS_IAM_AUTH, False)
+
+
+def get_env_postgres_iam_token_lifetime() -> int:
+    """
+    Gets the token lifetime in seconds for AWS RDS IAM authentication pool recycling.
+
+    AWS RDS IAM tokens are valid for 15 minutes (900 seconds). This value should be
+    set slightly lower to ensure connections are recycled before token expiration.
+
+    Returns:
+        int: Token lifetime in seconds (default: 840 = 14 minutes)
+    """
+    lifetime = _int_val(ENV_PHOENIX_POSTGRES_AWS_IAM_TOKEN_LIFETIME_SECONDS, 840)
+    if lifetime <= 0:
+        raise ValueError(
+            f"{ENV_PHOENIX_POSTGRES_AWS_IAM_TOKEN_LIFETIME_SECONDS} must be a positive integer. "
+            f"Got: {lifetime}"
+        )
+    if lifetime > 900:
+        logger.warning(
+            f"{ENV_PHOENIX_POSTGRES_AWS_IAM_TOKEN_LIFETIME_SECONDS} is set to {lifetime} seconds, "
+            f"which exceeds AWS RDS IAM token validity (900 seconds / 15 minutes). "
+            f"Consider setting it to 840 seconds (14 minutes) or less."
+        )
+    return lifetime
+
+
+def _validate_iam_auth_config() -> None:
+    """
+    Validate AWS RDS IAM authentication configuration if enabled.
+
+    Raises:
+        ImportError: If boto3 is not installed when IAM auth is enabled
+        ValueError: If configuration is invalid
+    """
+    if not get_env_postgres_use_iam_auth():
+        return
+
+    pg_host = getenv(ENV_PHOENIX_POSTGRES_HOST)
+    if not pg_host:
+        return
+
+    try:
+        import boto3  # type: ignore  # noqa: F401
+    except ImportError:
+        raise ImportError(
+            f"boto3 is required when {ENV_PHOENIX_POSTGRES_USE_AWS_IAM_AUTH} is enabled. "
+            "Install it with: pip install 'arize-phoenix[aws]'"
+        )
+
+    if not getenv(ENV_PHOENIX_POSTGRES_USER):
+        raise ValueError(
+            f"{ENV_PHOENIX_POSTGRES_USER} must be set when using AWS RDS IAM authentication"
+        )
+
+    try:
+        client = boto3.client("sts")  # pyright: ignore
+        client.get_caller_identity()  # pyright: ignore
+        logger.info("✓ AWS credentials validated for RDS IAM authentication")
+    except Exception as e:
+        raise ValueError(
+            f"Failed to validate AWS credentials for RDS IAM authentication: {e}. "
+            "Ensure AWS credentials are configured via environment variables, "
+            "~/.aws/credentials, or IAM role."
+        )

{arize_phoenix-12.8.0 → arize_phoenix-12.9.0}/src/phoenix/db/engines.py

@@ -5,7 +5,7 @@ import logging
 from collections.abc import Callable
 from enum import Enum
 from sqlite3 import Connection
-from typing import Any
+from typing import Any, Optional
 
 import aiosqlite
 import numpy as np
@@ -168,29 +168,142 @@ def aio_postgresql_engine(
     log_to_stdout: bool = False,
     log_migrations_to_stdout: bool = True,
 ) -> AsyncEngine:
-    asyncpg_url, asyncpg_args = get_pg_config(url, "asyncpg")
-    engine = create_async_engine(
-        url=asyncpg_url,
-        connect_args=asyncpg_args,
-        echo=log_to_stdout,
-        json_serializer=_dumps,
+    from phoenix.config import (
+        get_env_postgres_iam_token_lifetime,
+        get_env_postgres_use_iam_auth,
     )
+
+    use_iam_auth = get_env_postgres_use_iam_auth()
+
+    asyncpg_url, asyncpg_args = get_pg_config(url, "asyncpg", enforce_ssl=use_iam_auth)
+
+    iam_config: Optional[dict[str, Any]] = None
+    token_lifetime: int = 0
+    if use_iam_auth:
+        iam_config = _extract_iam_config_from_url(url)
+        token_lifetime = get_env_postgres_iam_token_lifetime()
+
+        async def iam_async_creator() -> Any:
+            import asyncpg  # type: ignore
+
+            from phoenix.db.iam_auth import generate_aws_rds_token
+
+            assert iam_config is not None
+            token = generate_aws_rds_token(
+                host=iam_config["host"],
+                port=iam_config["port"],
+                user=iam_config["user"],
+            )
+
+            conn_kwargs = {
+                "host": iam_config["host"],
+                "port": iam_config["port"],
+                "user": iam_config["user"],
+                "password": token,
+                "database": iam_config["database"],
+            }
+
+            if asyncpg_args:
+                conn_kwargs.update(asyncpg_args)
+
+            return await asyncpg.connect(**conn_kwargs)
+
+        engine = create_async_engine(
+            url=asyncpg_url,
+            async_creator=iam_async_creator,
+            echo=log_to_stdout,
+            json_serializer=_dumps,
+            pool_recycle=token_lifetime,
+        )
+    else:
+        engine = create_async_engine(
+            url=asyncpg_url,
+            connect_args=asyncpg_args,
+            echo=log_to_stdout,
+            json_serializer=_dumps,
+        )
+
     if not migrate:
         return engine
 
-    psycopg_url, psycopg_args = get_pg_config(url, "psycopg")
-    sync_engine = sqlalchemy.create_engine(
-        url=psycopg_url,
-        connect_args=psycopg_args,
-        echo=log_migrations_to_stdout,
-        json_serializer=_dumps,
-    )
+    psycopg_url, psycopg_args = get_pg_config(url, "psycopg", enforce_ssl=use_iam_auth)
+
+    if use_iam_auth:
+        assert iam_config is not None
+
+        def iam_sync_creator() -> Any:
+            import psycopg
+
+            from phoenix.db.iam_auth import generate_aws_rds_token
+
+            token = generate_aws_rds_token(
+                host=iam_config["host"],
+                port=iam_config["port"],
+                user=iam_config["user"],
+            )
+
+            conn_kwargs = {
+                "host": iam_config["host"],
+                "port": iam_config["port"],
+                "user": iam_config["user"],
+                "password": token,
+                "dbname": iam_config["database"],
+            }
+
+            if psycopg_args:
+                conn_kwargs.update(psycopg_args)
+
+            return psycopg.connect(**conn_kwargs)
+
+        sync_engine = sqlalchemy.create_engine(
+            url=psycopg_url,
+            creator=iam_sync_creator,
+            echo=log_migrations_to_stdout,
+            json_serializer=_dumps,
+            pool_recycle=token_lifetime,
+        )
+    else:
+        sync_engine = sqlalchemy.create_engine(
+            url=psycopg_url,
+            connect_args=psycopg_args,
+            echo=log_migrations_to_stdout,
+            json_serializer=_dumps,
+        )
+
     if schema := get_env_database_schema():
         event.listen(sync_engine, "connect", set_postgresql_search_path(schema))
     migrate_in_thread(sync_engine)
     return engine
 
 
+def _extract_iam_config_from_url(url: URL) -> dict[str, Any]:
+    """Extract connection parameters needed for IAM authentication from a SQLAlchemy URL.
+
+    Args:
+        url: SQLAlchemy database URL
+
+    Returns:
+        Dictionary with host, port, user, and database
+    """
+    host = url.host
+    if not host:
+        raise ValueError("Database host is required for IAM authentication")
+
+    port = url.port or 5432
+    user = url.username
+    if not user:
+        raise ValueError("Database user is required for IAM authentication")
+
+    database = url.database or "postgres"
+
+    return {
+        "host": host,
+        "port": port,
+        "user": user,
+        "database": database,
+    }
+
+
 def _dumps(obj: Any) -> str:
     return orjson.dumps(obj, default=_default).decode()
 

arize_phoenix-12.9.0/src/phoenix/db/iam_auth.py

@@ -0,0 +1,64 @@
+from __future__ import annotations
+
+import logging
+
+logger = logging.getLogger(__name__)
+
+
+def generate_aws_rds_token(
+    host: str,
+    port: int,
+    user: str,
+) -> str:
+    """Generate an AWS RDS IAM authentication token.
+
+    This function creates a short-lived (15 minutes) authentication token for connecting
+    to AWS RDS/Aurora PostgreSQL instances using IAM database authentication.
+
+    The AWS region is automatically resolved using boto3.
+
+    Args:
+        host: The database hostname (e.g., 'mydb.abc123.us-west-2.rds.amazonaws.com')
+        port: The database port (typically 5432 for PostgreSQL)
+        user: The database username (must match an IAM-enabled database user)
+
+    Returns:
+        A temporary authentication token string to use as the database password
+
+    Raises:
+        ImportError: If boto3 is not installed
+        Exception: If AWS credentials/region are not configured or token generation fails
+
+    Example:
+        >>> token = generate_aws_rds_token(
+        ...     host='mydb.us-west-2.rds.amazonaws.com',
+        ...     port=5432,
+        ...     user='myuser'
+        ... )
+    """
+    try:
+        import boto3  # type: ignore
+    except ImportError as e:
+        raise ImportError(
+            "boto3 is required for AWS RDS IAM authentication. "
+            "Install it with: pip install 'arize-phoenix[aws]'"
+        ) from e
+
+    try:
+        client = boto3.client("rds")
+
+        logger.debug(f"Generating AWS RDS IAM auth token for user '{user}' at {host}:{port}")
+        token = client.generate_db_auth_token(  # pyright: ignore
+            DBHostname=host,
+            Port=port,
+            DBUsername=user,
+        )
+
+        return str(token)  # pyright: ignore
+
+    except Exception as e:
+        logger.error(
+            f"Failed to generate AWS RDS IAM authentication token: {e}. "
+            "Ensure AWS credentials are configured and have 'rds-db:connect' permission."
+        )
+        raise

{arize_phoenix-12.8.0 → arize_phoenix-12.9.0}/src/phoenix/db/pg_config.py

@@ -10,12 +10,14 @@ from typing_extensions import assert_never
 def get_pg_config(
     url: URL,
     driver: Literal["psycopg", "asyncpg"],
+    enforce_ssl: bool = False,
 ) -> tuple[URL, dict[str, Any]]:
     """Convert SQLAlchemy URL to driver-specific configuration.
 
     Args:
         url: SQLAlchemy URL
         driver: "psycopg" or "asyncpg"
+        enforce_ssl: If True, ensure SSL is enabled (required for AWS RDS IAM auth)
 
     Returns:
         Tuple of (base_url, connect_args):
@@ -26,6 +28,14 @@ def get_pg_config(
     query = url.query
     ssl_args = _get_ssl_args(query)
 
+    if enforce_ssl and not ssl_args:
+        ssl_args = {"sslmode": "require"}
+    elif enforce_ssl and ssl_args.get("sslmode") == "disable":
+        raise ValueError(
+            "SSL cannot be disabled when using AWS RDS IAM authentication. "
+            "Remove 'sslmode=disable' from the connection string."
+        )
+
     # Create base URL without SSL parameters
     base_url = url.set(
         drivername=f"postgresql+{driver}",

{arize_phoenix-12.8.0 → arize_phoenix-12.9.0}/src/phoenix/server/api/context.py

@@ -35,6 +35,7 @@ from phoenix.server.api.dataloaders import (
     ExperimentRepeatedRunGroupsDataLoader,
     ExperimentRunAnnotations,
     ExperimentRunCountsDataLoader,
+    ExperimentRunsByExperimentAndExampleDataLoader,
     ExperimentSequenceNumberDataLoader,
     LastUsedTimesByGenerativeModelIdDataLoader,
     LatencyMsQuantileDataLoader,
@@ -71,6 +72,7 @@ from phoenix.server.api.dataloaders import (
     SpanProjectsDataLoader,
     TableFieldsDataLoader,
     TokenCountDataLoader,
+    TokenPricesByModelDataLoader,
     TraceAnnotationsByTraceDataLoader,
     TraceByTraceIdsDataLoader,
     TraceRetentionPolicyIdByProjectIdDataLoader,
@@ -100,27 +102,38 @@ class DataLoaders:
         AverageExperimentRepeatedRunGroupLatencyDataLoader
     )
     average_experiment_run_latency: AverageExperimentRunLatencyDataLoader
+    dataset_example_fields: TableFieldsDataLoader
     dataset_example_revisions: DatasetExampleRevisionsDataLoader
     dataset_example_spans: DatasetExampleSpansDataLoader
     dataset_labels: DatasetLabelsDataLoader
+    dataset_label_fields: TableFieldsDataLoader
     dataset_dataset_splits: DatasetDatasetSplitsDataLoader
     dataset_examples_and_versions_by_experiment_run: (
         DatasetExamplesAndVersionsByExperimentRunDataLoader
     )
     dataset_example_splits: DatasetExampleSplitsDataLoader
+    dataset_fields: TableFieldsDataLoader
+    dataset_split_fields: TableFieldsDataLoader
+    dataset_version_fields: TableFieldsDataLoader
+    document_annotation_fields: TableFieldsDataLoader
     document_evaluation_summaries: DocumentEvaluationSummaryDataLoader
     document_evaluations: DocumentEvaluationsDataLoader
     document_retrieval_metrics: DocumentRetrievalMetricsDataLoader
     experiment_annotation_summaries: ExperimentAnnotationSummaryDataLoader
     experiment_dataset_splits: ExperimentDatasetSplitsDataLoader
     experiment_error_rates: ExperimentErrorRatesDataLoader
+    experiment_fields: TableFieldsDataLoader
     experiment_repeated_run_group_annotation_summaries: (
         ExperimentRepeatedRunGroupAnnotationSummariesDataLoader
     )
     experiment_repeated_run_groups: ExperimentRepeatedRunGroupsDataLoader
+    experiment_run_annotation_fields: TableFieldsDataLoader
     experiment_run_annotations: ExperimentRunAnnotations
     experiment_run_counts: ExperimentRunCountsDataLoader
+    experiment_run_fields: TableFieldsDataLoader
+    experiment_runs_by_experiment_and_example: ExperimentRunsByExperimentAndExampleDataLoader
     experiment_sequence_number: ExperimentSequenceNumberDataLoader
+    generative_model_fields: TableFieldsDataLoader
     last_used_times_by_generative_model_id: LastUsedTimesByGenerativeModelIdDataLoader
     latency_ms_quantile: LatencyMsQuantileDataLoader
     min_start_or_max_end_times: MinStartOrMaxEndTimeDataLoader
@@ -130,7 +143,12 @@ class DataLoaders:
     project_fields: TableFieldsDataLoader
     project_trace_retention_policy_fields: TableFieldsDataLoader
     projects_by_trace_retention_policy_id: ProjectIdsByTraceRetentionPolicyIdDataLoader
+    prompt_fields: TableFieldsDataLoader
+    prompt_label_fields: TableFieldsDataLoader
     prompt_version_sequence_number: PromptVersionSequenceNumberDataLoader
+    prompt_version_tag_fields: TableFieldsDataLoader
+    project_session_annotation_fields: TableFieldsDataLoader
+    project_session_fields: TableFieldsDataLoader
     record_counts: RecordCountDataLoader
     session_annotations_by_session: SessionAnnotationsBySessionDataLoader
     session_first_inputs: SessionIODataLoader
@@ -139,6 +157,7 @@ class DataLoaders:
     session_num_traces_with_error: SessionNumTracesWithErrorDataLoader
     session_token_usages: SessionTokenUsagesDataLoader
     session_trace_latency_ms_quantile: SessionTraceLatencyMsQuantileDataLoader
+    span_annotation_fields: TableFieldsDataLoader
     span_annotations: SpanAnnotationsDataLoader
     span_by_id: SpanByIdDataLoader
     span_cost_by_span: SpanCostBySpanDataLoader
@@ -167,12 +186,16 @@ class DataLoaders:
     span_fields: TableFieldsDataLoader
     span_projects: SpanProjectsDataLoader
     token_counts: TokenCountDataLoader
+    token_prices_by_model: TokenPricesByModelDataLoader
+    trace_annotation_fields: TableFieldsDataLoader
     trace_annotations_by_trace: TraceAnnotationsByTraceDataLoader
     trace_by_trace_ids: TraceByTraceIdsDataLoader
     trace_fields: TableFieldsDataLoader
     trace_retention_policy_id_by_project_id: TraceRetentionPolicyIdByProjectIdDataLoader
     trace_root_spans: TraceRootSpansDataLoader
     user_roles: UserRolesDataLoader
+    user_api_key_fields: TableFieldsDataLoader
+    user_fields: TableFieldsDataLoader
     users: UsersDataLoader
 
 

{arize_phoenix-12.8.0 → arize_phoenix-12.9.0}/src/phoenix/server/api/dataloaders/__init__.py

@@ -33,6 +33,9 @@ from .experiment_repeated_run_group_annotation_summaries import (
 from .experiment_repeated_run_groups import ExperimentRepeatedRunGroupsDataLoader
 from .experiment_run_annotations import ExperimentRunAnnotations
 from .experiment_run_counts import ExperimentRunCountsDataLoader
+from .experiment_runs_by_experiment_and_example import (
+    ExperimentRunsByExperimentAndExampleDataLoader,
+)
 from .experiment_sequence_number import ExperimentSequenceNumberDataLoader
 from .last_used_times_by_generative_model_id import LastUsedTimesByGenerativeModelIdDataLoader
 from .latency_ms_quantile import LatencyMsQuantileCache, LatencyMsQuantileDataLoader
@@ -73,6 +76,7 @@ from .span_descendants import SpanDescendantsDataLoader
 from .span_projects import SpanProjectsDataLoader
 from .table_fields import TableFieldsDataLoader
 from .token_counts import TokenCountCache, TokenCountDataLoader
+from .token_prices_by_model import TokenPricesByModelDataLoader
 from .trace_annotations_by_trace import TraceAnnotationsByTraceDataLoader
 from .trace_by_trace_ids import TraceByTraceIdsDataLoader
 from .trace_retention_policy_id_by_project_id import TraceRetentionPolicyIdByProjectIdDataLoader
@@ -102,6 +106,7 @@ __all__ = [
     "ExperimentRepeatedRunGroupAnnotationSummariesDataLoader",
     "ExperimentRunAnnotations",
     "ExperimentRunCountsDataLoader",
+    "ExperimentRunsByExperimentAndExampleDataLoader",
     "ExperimentSequenceNumberDataLoader",
     "LastUsedTimesByGenerativeModelIdDataLoader",
     "LatencyMsQuantileDataLoader",
@@ -139,6 +144,7 @@ __all__ = [
     "SpanProjectsDataLoader",
     "TableFieldsDataLoader",
     "TokenCountDataLoader",
+    "TokenPricesByModelDataLoader",
     "TraceAnnotationsByTraceDataLoader",
     "TraceByTraceIdsDataLoader",
     "TraceRetentionPolicyIdByProjectIdDataLoader",

{arize_phoenix-12.8.0 → arize_phoenix-12.9.0}/src/phoenix/server/api/dataloaders/experiment_repeated_run_groups.py

@@ -1,7 +1,6 @@
 from dataclasses import dataclass
 
 from sqlalchemy import select, tuple_
-from sqlalchemy.orm import joinedload
 from strawberry.dataloader import DataLoader
 from typing_extensions import TypeAlias
 
@@ -38,7 +37,6 @@ class ExperimentRepeatedRunGroupsDataLoader(DataLoader[Key, Result]):
                 ).in_(set(keys))
             )
             .order_by(models.ExperimentRun.repetition_number)
-            .options(joinedload(models.ExperimentRun.trace).load_only(models.Trace.trace_id))
         )
 
         async with self._db() as session:

arize_phoenix-12.9.0/src/phoenix/server/api/dataloaders/experiment_runs_by_experiment_and_example.py

@@ -0,0 +1,44 @@
+from collections import defaultdict
+
+from sqlalchemy import select, tuple_
+from strawberry.dataloader import DataLoader
+from typing_extensions import TypeAlias
+
+from phoenix.db import models
+from phoenix.server.types import DbSessionFactory
+
+ExperimentId: TypeAlias = int
+DatasetExampleId: TypeAlias = int
+Key: TypeAlias = tuple[ExperimentId, DatasetExampleId]
+Result: TypeAlias = list[models.ExperimentRun]
+
+
+class ExperimentRunsByExperimentAndExampleDataLoader(DataLoader[Key, Result]):
+    def __init__(self, db: DbSessionFactory) -> None:
+        super().__init__(load_fn=self._load_fn)
+        self._db = db
+
+    async def _load_fn(self, keys: list[Key]) -> list[Result]:
+        runs_by_key: defaultdict[Key, Result] = defaultdict(list)
+
+        async with self._db() as session:
+            stmt = (
+                select(models.ExperimentRun)
+                .where(
+                    tuple_(
+                        models.ExperimentRun.experiment_id,
+                        models.ExperimentRun.dataset_example_id,
+                    ).in_(keys)
+                )
+                .order_by(
+                    models.ExperimentRun.experiment_id,
+                    models.ExperimentRun.dataset_example_id,
+                    models.ExperimentRun.repetition_number,
+                )
+            )
+            result = await session.stream_scalars(stmt)
+            async for run in result:
+                key = (run.experiment_id, run.dataset_example_id)
+                runs_by_key[key].append(run)
+
+        return [runs_by_key[key] for key in keys]

{arize_phoenix-12.8.0 → arize_phoenix-12.9.0}/src/phoenix/server/api/dataloaders/span_costs.py

@@ -1,7 +1,6 @@
 from typing import Optional
 
 from sqlalchemy import select
-from sqlalchemy.orm import joinedload, load_only
 from strawberry.dataloader import DataLoader
 from typing_extensions import TypeAlias
 
@@ -22,14 +21,9 @@ class SpanCostsDataLoader(DataLoader[Key, Result]):
         span_ids = list(set(keys))
         async with self._db() as session:
             costs = {
-                span.id: span.span_cost
-                async for span in await session.stream_scalars(
-                    select(models.Span)
-                    .where(models.Span.id.in_(span_ids))
-                    .options(
-                        load_only(models.Span.id),
-                        joinedload(models.Span.span_cost),
-                    )
+                span_cost.span_rowid: span_cost
+                async for span_cost in await session.stream_scalars(
+                    select(models.SpanCost).where(models.SpanCost.span_rowid.in_(span_ids))
                 )
             }
         return [costs.get(span_id) for span_id in keys]