arize-phoenix 8.32.1__py3-none-any.whl → 9.0.0__py3-none-any.whl

phoenix/experiments/functions.py

@@ -95,6 +95,7 @@ def run_experiment(
     dry_run: Union[bool, int] = False,
     print_summary: bool = True,
     concurrency: int = 3,
+    timeout: Optional[int] = None,
 ) -> RanExperiment:
     """
     Runs an experiment using a given set of dataset of examples.
@@ -148,6 +149,8 @@ def run_experiment(
         concurrency (int): Specifies the concurrency for task execution. In order to enable
             concurrent task execution, the task callable must be a coroutine function.
             Defaults to 3.
+        timeout (Optional[int]): The timeout for the task execution in seconds. Use this to run
+            longer tasks to avoid re-queuing the same task multiple times. Defaults to None.
 
     Returns:
         RanExperiment: The results of the experiment and evaluation. Additional evaluations can be
@@ -380,6 +383,7 @@ def run_experiment(
         fallback_return_value=None,
         tqdm_bar_format=get_tqdm_progress_bar_formatter("running tasks"),
         concurrency=concurrency,
+        timeout=timeout,
     )
 
     test_cases = [
@@ -752,7 +756,7 @@ def _print_experiment_error(
     Prints an experiment error.
     """
     display_error = RuntimeError(
-        f"{kind} failed for example id {repr(example_id)}, " f"repetition {repr(repetition_number)}"
+        f"{kind} failed for example id {repr(example_id)}, repetition {repr(repetition_number)}"
     )
     display_error.__cause__ = error
     formatted_exception = "".join(

phoenix/server/api/auth.py

@@ -42,3 +42,12 @@ class IsAdmin(Authorization):
         if not info.context.auth_enabled:
             return False
         return isinstance((user := info.context.user), PhoenixUser) and user.is_admin
+
+
+class IsAdminIfAuthEnabled(Authorization):
+    message = MSG_ADMIN_ONLY
+
+    def has_permission(self, source: Any, info: Info, **kwargs: Any) -> bool:
+        if not info.context.auth_enabled:
+            return True
+        return isinstance((user := info.context.user), PhoenixUser) and user.is_admin

phoenix/server/api/context.py

@@ -32,6 +32,7 @@ from phoenix.server.api.dataloaders import (
     NumChildSpansDataLoader,
     NumSpansPerTraceDataLoader,
     ProjectByNameDataLoader,
+    ProjectIdsByTraceRetentionPolicyIdDataLoader,
     PromptVersionSequenceNumberDataLoader,
     RecordCountDataLoader,
     SessionIODataLoader,
@@ -47,6 +48,7 @@ from phoenix.server.api.dataloaders import (
     TableFieldsDataLoader,
     TokenCountDataLoader,
     TraceByTraceIdsDataLoader,
+    TraceRetentionPolicyIdByProjectIdDataLoader,
     TraceRootSpansDataLoader,
     UserRolesDataLoader,
     UsersDataLoader,
@@ -82,6 +84,7 @@ class DataLoaders:
     num_child_spans: NumChildSpansDataLoader
     num_spans_per_trace: NumSpansPerTraceDataLoader
     project_fields: TableFieldsDataLoader
+    projects_by_trace_retention_policy_id: ProjectIdsByTraceRetentionPolicyIdDataLoader
     prompt_version_sequence_number: PromptVersionSequenceNumberDataLoader
     record_counts: RecordCountDataLoader
     session_first_inputs: SessionIODataLoader
@@ -99,6 +102,8 @@ class DataLoaders:
     token_counts: TokenCountDataLoader
     trace_by_trace_ids: TraceByTraceIdsDataLoader
     trace_fields: TableFieldsDataLoader
+    trace_retention_policy_id_by_project_id: TraceRetentionPolicyIdByProjectIdDataLoader
+    project_trace_retention_policy_fields: TableFieldsDataLoader
     trace_root_spans: TraceRootSpansDataLoader
     project_by_name: ProjectByNameDataLoader
     users: UsersDataLoader

phoenix/server/api/dataloaders/__init__.py

@@ -20,6 +20,7 @@ from .min_start_or_max_end_times import MinStartOrMaxEndTimeCache, MinStartOrMax
 from .num_child_spans import NumChildSpansDataLoader
 from .num_spans_per_trace import NumSpansPerTraceDataLoader
 from .project_by_name import ProjectByNameDataLoader
+from .project_ids_by_trace_retention_policy_id import ProjectIdsByTraceRetentionPolicyIdDataLoader
 from .prompt_version_sequence_number import PromptVersionSequenceNumberDataLoader
 from .record_counts import RecordCountCache, RecordCountDataLoader
 from .session_io import SessionIODataLoader
@@ -35,6 +36,7 @@ from .span_projects import SpanProjectsDataLoader
 from .table_fields import TableFieldsDataLoader
 from .token_counts import TokenCountCache, TokenCountDataLoader
 from .trace_by_trace_ids import TraceByTraceIdsDataLoader
+from .trace_retention_policy_id_by_project_id import TraceRetentionPolicyIdByProjectIdDataLoader
 from .trace_root_spans import TraceRootSpansDataLoader
 from .user_roles import UserRolesDataLoader
 from .users import UsersDataLoader
@@ -57,6 +59,7 @@ __all__ = [
     "MinStartOrMaxEndTimeDataLoader",
     "NumChildSpansDataLoader",
     "NumSpansPerTraceDataLoader",
+    "ProjectIdsByTraceRetentionPolicyIdDataLoader",
     "PromptVersionSequenceNumberDataLoader",
     "RecordCountDataLoader",
     "SessionIODataLoader",
@@ -71,6 +74,7 @@ __all__ = [
     "TableFieldsDataLoader",
     "TokenCountDataLoader",
     "TraceByTraceIdsDataLoader",
+    "TraceRetentionPolicyIdByProjectIdDataLoader",
     "TraceRootSpansDataLoader",
     "ProjectByNameDataLoader",
     "SpanAnnotationsDataLoader",

phoenix/server/api/dataloaders/annotation_summaries.py

@@ -1,11 +1,11 @@
 from collections import defaultdict
 from datetime import datetime
-from typing import Any, Literal, Optional
+from typing import Any, Literal, Optional, Type, Union, cast
 
 import pandas as pd
 from aioitertools.itertools import groupby
 from cachetools import LFUCache, TTLCache
-from sqlalchemy import Select, func, or_, select
+from sqlalchemy import Select, and_, case, distinct, func, or_, select
 from strawberry.dataloader import AbstractCache, DataLoader
 from typing_extensions import TypeAlias, assert_never
 
@@ -92,7 +92,7 @@ class AnnotationSummaryDataLoader(DataLoader[Key, Result]):
             async with self._db() as session:
                 data = await session.stream(stmt)
                 async for annotation_name, group in groupby(data, lambda row: row.name):
-                    summary = AnnotationSummary(pd.DataFrame(group))
+                    summary = AnnotationSummary(name=annotation_name, df=pd.DataFrame(group))
                     for position in params[annotation_name]:
                         results[position] = summary
         return results
@@ -103,23 +103,64 @@ def _get_stmt(
     *annotation_names: Param,
 ) -> Select[Any]:
     kind, project_rowid, (start_time, end_time), filter_condition = segment
-    stmt = select()
+
+    annotation_model: Union[Type[models.SpanAnnotation], Type[models.TraceAnnotation]]
+    entity_model: Union[Type[models.Span], Type[models.Trace]]
+    entity_join_model: Optional[Type[models.Base]]
+    entity_id_column: Any
+
     if kind == "span":
-        msa = models.SpanAnnotation
-        name_column, label_column, score_column = msa.name, msa.label, msa.score
-        time_column = models.Span.start_time
-        stmt = stmt.join(models.Span).join_from(models.Span, models.Trace)
-        if filter_condition:
-            sf = SpanFilter(filter_condition)
-            stmt = sf(stmt)
+        annotation_model = models.SpanAnnotation
+        entity_model = models.Span
+        entity_join_model = models.Trace
+        entity_id_column = models.Span.id.label("entity_id")
     elif kind == "trace":
-        mta = models.TraceAnnotation
-        name_column, label_column, score_column = mta.name, mta.label, mta.score
-        time_column = models.Trace.start_time
-        stmt = stmt.join(models.Trace)
+        annotation_model = models.TraceAnnotation
+        entity_model = models.Trace
+        entity_join_model = None
+        entity_id_column = models.Trace.id.label("entity_id")
     else:
         assert_never(kind)
-    stmt = stmt.add_columns(
+
+    name_column = annotation_model.name
+    label_column = annotation_model.label
+    score_column = annotation_model.score
+    time_column = entity_model.start_time
+
+    # First query: count distinct entities per annotation name
+    # This is used later to calculate accurate fractions that account for entities without labels
+    entity_count_query = select(
+        name_column, func.count(distinct(entity_id_column)).label("entity_count")
+    )
+
+    if kind == "span":
+        entity_count_query = entity_count_query.join(cast(Type[models.Span], entity_model))
+        entity_count_query = entity_count_query.join_from(
+            cast(Type[models.Span], entity_model), cast(Type[models.Trace], entity_join_model)
+        )
+        entity_count_query = entity_count_query.where(models.Trace.project_rowid == project_rowid)
+    elif kind == "trace":
+        entity_count_query = entity_count_query.join(cast(Type[models.Trace], entity_model))
+        entity_count_query = entity_count_query.where(
+            cast(Type[models.Trace], entity_model).project_rowid == project_rowid
+        )
+
+    entity_count_query = entity_count_query.where(
+        or_(score_column.is_not(None), label_column.is_not(None))
+    )
+    entity_count_query = entity_count_query.where(name_column.in_(annotation_names))
+
+    if start_time:
+        entity_count_query = entity_count_query.where(start_time <= time_column)
+    if end_time:
+        entity_count_query = entity_count_query.where(time_column < end_time)
+
+    entity_count_query = entity_count_query.group_by(name_column)
+    entity_count_subquery = entity_count_query.subquery()
+
+    # Main query: gets raw annotation data with counts per (span/trace)+name+label
+    base_stmt = select(
+        entity_id_column,
         name_column,
         label_column,
         func.count().label("record_count"),
@@ -127,13 +168,151 @@ def _get_stmt(
         func.count(score_column).label("score_count"),
         func.sum(score_column).label("score_sum"),
     )
-    stmt = stmt.group_by(name_column, label_column)
-    stmt = stmt.order_by(name_column, label_column)
-    stmt = stmt.where(models.Trace.project_rowid == project_rowid)
-    stmt = stmt.where(or_(score_column.is_not(None), label_column.is_not(None)))
-    stmt = stmt.where(name_column.in_(annotation_names))
+
+    if kind == "span":
+        base_stmt = base_stmt.join(cast(Type[models.Span], entity_model))
+        base_stmt = base_stmt.join_from(
+            cast(Type[models.Span], entity_model), cast(Type[models.Trace], entity_join_model)
+        )
+        base_stmt = base_stmt.where(models.Trace.project_rowid == project_rowid)
+        if filter_condition:
+            sf = SpanFilter(filter_condition)
+            base_stmt = sf(base_stmt)
+    elif kind == "trace":
+        base_stmt = base_stmt.join(cast(Type[models.Trace], entity_model))
+        base_stmt = base_stmt.where(
+            cast(Type[models.Trace], entity_model).project_rowid == project_rowid
+        )
+    else:
+        assert_never(kind)
+
+    base_stmt = base_stmt.where(or_(score_column.is_not(None), label_column.is_not(None)))
+    base_stmt = base_stmt.where(name_column.in_(annotation_names))
+
     if start_time:
-        stmt = stmt.where(start_time <= time_column)
+        base_stmt = base_stmt.where(start_time <= time_column)
     if end_time:
-        stmt = stmt.where(time_column < end_time)
-    return stmt
+        base_stmt = base_stmt.where(time_column < end_time)
+
+    # Group to get one row per (span/trace)+name+label combination
+    base_stmt = base_stmt.group_by(entity_id_column, name_column, label_column)
+
+    base_subquery = base_stmt.subquery()
+
+    # Calculate total counts per (span/trace)+name for computing fractions
+    entity_totals = (
+        select(
+            base_subquery.c.entity_id,
+            base_subquery.c.name,
+            func.sum(base_subquery.c.label_count).label("total_label_count"),
+            func.sum(base_subquery.c.score_count).label("total_score_count"),
+            func.sum(base_subquery.c.score_sum).label("entity_score_sum"),
+        )
+        .group_by(base_subquery.c.entity_id, base_subquery.c.name)
+        .subquery()
+    )
+
+    per_entity_fractions = (
+        select(
+            base_subquery.c.entity_id,
+            base_subquery.c.name,
+            base_subquery.c.label,
+            base_subquery.c.record_count,
+            base_subquery.c.label_count,
+            base_subquery.c.score_count,
+            base_subquery.c.score_sum,
+            # Calculate label fraction, avoiding division by zero when total_label_count is 0
+            case(
+                (
+                    entity_totals.c.total_label_count > 0,
+                    base_subquery.c.label_count * 1.0 / entity_totals.c.total_label_count,
+                ),
+                else_=None,
+            ).label("label_fraction"),
+            # Calculate average score for the entity (if there are any scores)
+            case(
+                (
+                    entity_totals.c.total_score_count > 0,
+                    entity_totals.c.entity_score_sum * 1.0 / entity_totals.c.total_score_count,
+                ),
+                else_=None,
+            ).label("entity_avg_score"),
+        )
+        .join(
+            entity_totals,
+            and_(
+                base_subquery.c.entity_id == entity_totals.c.entity_id,
+                base_subquery.c.name == entity_totals.c.name,
+            ),
+        )
+        .subquery()
+    )
+
+    # Aggregate metrics across (spans/traces) for each name+label combination.
+    label_entity_metrics = (
+        select(
+            per_entity_fractions.c.name,
+            per_entity_fractions.c.label,
+            func.count(distinct(per_entity_fractions.c.entity_id)).label("entities_with_label"),
+            func.sum(per_entity_fractions.c.label_count).label("total_label_count"),
+            func.sum(per_entity_fractions.c.score_count).label("total_score_count"),
+            func.sum(per_entity_fractions.c.score_sum).label("total_score_sum"),
+            # Average of label fractions for entities that have this label
+            func.avg(per_entity_fractions.c.label_fraction).label("avg_label_fraction_present"),
+            # Average of per-entity average scores (but we handle overall aggregation separately)
+        )
+        .group_by(per_entity_fractions.c.name, per_entity_fractions.c.label)
+        .subquery()
+    )
+
+    # Compute distinct per-entity average scores to ensure each entity counts only once.
+    distinct_entity_scores = (
+        select(
+            per_entity_fractions.c.entity_id,
+            per_entity_fractions.c.name,
+            per_entity_fractions.c.entity_avg_score,
+        )
+        .distinct()
+        .subquery()
+    )
+
+    overall_score_aggregates = (
+        select(
+            distinct_entity_scores.c.name,
+            func.avg(distinct_entity_scores.c.entity_avg_score).label("overall_avg_score"),
+        )
+        .group_by(distinct_entity_scores.c.name)
+        .subquery()
+    )
+
+    # Final result: adjust label fractions by the proportion of entities reporting this label
+    # and include the overall average score per annotation name.
+    final_stmt = (
+        select(
+            label_entity_metrics.c.name,
+            label_entity_metrics.c.label,
+            # Adjust label fraction, guarding against division by zero in entity_count
+            case(
+                (
+                    entity_count_subquery.c.entity_count > 0,
+                    label_entity_metrics.c.avg_label_fraction_present
+                    * label_entity_metrics.c.entities_with_label
+                    / entity_count_subquery.c.entity_count,
+                ),
+                else_=None,
+            ).label("avg_label_fraction"),
+            overall_score_aggregates.c.overall_avg_score.label("avg_score"),  # same for all labels
+            label_entity_metrics.c.total_label_count.label("label_count"),
+            label_entity_metrics.c.total_score_count.label("score_count"),
+            label_entity_metrics.c.total_score_sum.label("score_sum"),
+            label_entity_metrics.c.entities_with_label.label("record_count"),
+        )
+        .join(entity_count_subquery, label_entity_metrics.c.name == entity_count_subquery.c.name)
+        .join(
+            overall_score_aggregates,
+            label_entity_metrics.c.name == overall_score_aggregates.c.name,
+        )
+        .order_by(label_entity_metrics.c.name, label_entity_metrics.c.label)
+    )
+
+    return final_stmt

phoenix/server/api/dataloaders/project_ids_by_trace_retention_policy_id.py

@@ -0,0 +1,42 @@
+from collections import defaultdict
+
+from sqlalchemy import or_, select
+from strawberry.dataloader import DataLoader
+from typing_extensions import TypeAlias
+
+from phoenix.db.constants import DEFAULT_PROJECT_TRACE_RETENTION_POLICY_ID
+from phoenix.db.models import Project
+from phoenix.server.types import DbSessionFactory
+
+PolicyRowId: TypeAlias = int
+ProjectRowId: TypeAlias = int
+
+Key: TypeAlias = PolicyRowId
+Result: TypeAlias = list[ProjectRowId]
+
+
+class ProjectIdsByTraceRetentionPolicyIdDataLoader(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]:
+        ids = set(keys)
+        stmt = select(Project.trace_retention_policy_id, Project.id)
+        if DEFAULT_PROJECT_TRACE_RETENTION_POLICY_ID in ids:
+            stmt = stmt.where(
+                or_(
+                    Project.trace_retention_policy_id.in_(ids),
+                    Project.trace_retention_policy_id.is_(None),
+                )
+            )
+        else:
+            stmt = stmt.where(Project.trace_retention_policy_id.in_(ids))
+        projects: defaultdict[Key, Result] = defaultdict(list)
+        async with self._db() as session:
+            data = await session.stream(stmt)
+            async for policy_rowid, project_rowid in data:
+                projects[policy_rowid or DEFAULT_PROJECT_TRACE_RETENTION_POLICY_ID].append(
+                    project_rowid
+                )
+        return [projects.get(project_name, []).copy() for project_name in keys]

phoenix/server/api/dataloaders/trace_retention_policy_id_by_project_id.py

@@ -0,0 +1,34 @@
+from sqlalchemy import select
+from strawberry.dataloader import DataLoader
+from typing_extensions import TypeAlias
+
+from phoenix.db.constants import DEFAULT_PROJECT_TRACE_RETENTION_POLICY_ID
+from phoenix.db.models import Project
+from phoenix.server.types import DbSessionFactory
+
+PolicyRowId: TypeAlias = int
+ProjectRowId: TypeAlias = int
+
+Key: TypeAlias = ProjectRowId
+Result: TypeAlias = PolicyRowId
+
+
+class TraceRetentionPolicyIdByProjectIdDataLoader(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]:
+        ids = set(keys)
+        stmt = (
+            select(Project.id, Project.trace_retention_policy_id)
+            .where(Project.trace_retention_policy_id.isnot(None))
+            .where(Project.id.in_(ids))
+        )
+        async with self._db() as session:
+            data = await session.execute(stmt)
+        result = {project_rowid: policy_id for project_rowid, policy_id in data.all()}
+        return [
+            result.get(project_rowid, DEFAULT_PROJECT_TRACE_RETENTION_POLICY_ID)
+            for project_rowid in keys
+        ]

phoenix/server/api/helpers/annotations.py

@@ -0,0 +1,9 @@
+from strawberry.relay import GlobalID
+
+
+def get_user_identifier(user_id: int) -> str:
+    """
+    Generates an annotation identifier unique to the user.
+    """
+    user_gid = str(GlobalID(type_name="User", node_id=str(user_id)))
+    return f"px-app:{user_gid}"