arize-phoenix 11.8.0__py3-none-any.whl → 11.9.0__py3-none-any.whl

phoenix/server/api/types/Project.py

@@ -11,6 +11,7 @@ from sqlalchemy import and_, desc, distinct, exists, func, or_, select
 from sqlalchemy.dialects import postgresql, sqlite
 from sqlalchemy.sql.elements import ColumnElement
 from sqlalchemy.sql.expression import tuple_
+from sqlalchemy.sql.functions import percentile_cont
 from strawberry import ID, UNSET, Private, lazy
 from strawberry.relay import Connection, Edge, Node, NodeID, PageInfo
 from strawberry.types import Info
@@ -841,6 +842,475 @@ class Project(Node):
                 data[timestamp] = TimeSeriesDataPoint(timestamp=timestamp)
         return TraceCountTimeSeries(data=sorted(data.values(), key=lambda x: x.timestamp))
 
+    @strawberry.field
+    async def trace_count_by_status_time_series(
+        self,
+        info: Info[Context, None],
+        time_range: TimeRange,
+        time_bin_config: Optional[TimeBinConfig] = UNSET,
+    ) -> TraceCountByStatusTimeSeries:
+        if time_range.start is None:
+            raise BadRequest("Start time is required")
+
+        dialect = info.context.db.dialect
+        utc_offset_minutes = 0
+        field: Literal["minute", "hour", "day", "week", "month", "year"] = "hour"
+        if time_bin_config:
+            utc_offset_minutes = time_bin_config.utc_offset_minutes
+            if time_bin_config.scale is TimeBinScale.MINUTE:
+                field = "minute"
+            elif time_bin_config.scale is TimeBinScale.HOUR:
+                field = "hour"
+            elif time_bin_config.scale is TimeBinScale.DAY:
+                field = "day"
+            elif time_bin_config.scale is TimeBinScale.WEEK:
+                field = "week"
+            elif time_bin_config.scale is TimeBinScale.MONTH:
+                field = "month"
+            elif time_bin_config.scale is TimeBinScale.YEAR:
+                field = "year"
+        bucket = date_trunc(dialect, field, models.Trace.start_time, utc_offset_minutes)
+        trace_error_status_counts = (
+            select(
+                models.Span.trace_rowid,
+            )
+            .where(models.Span.parent_id.is_(None))
+            .group_by(models.Span.trace_rowid)
+            .having(func.max(models.Span.cumulative_error_count) > 0)
+        ).subquery()
+        stmt = (
+            select(
+                bucket,
+                func.count(models.Trace.id).label("total_count"),
+                func.coalesce(func.count(trace_error_status_counts.c.trace_rowid), 0).label(
+                    "error_count"
+                ),
+            )
+            .join_from(
+                models.Trace,
+                trace_error_status_counts,
+                onclause=trace_error_status_counts.c.trace_rowid == models.Trace.id,
+                isouter=True,
+            )
+            .where(models.Trace.project_rowid == self.project_rowid)
+            .group_by(bucket)
+            .order_by(bucket)
+        )
+        if time_range:
+            if time_range.start:
+                stmt = stmt.where(time_range.start <= models.Trace.start_time)
+            if time_range.end:
+                stmt = stmt.where(models.Trace.start_time < time_range.end)
+        data: dict[datetime, TraceCountByStatusTimeSeriesDataPoint] = {}
+        async with info.context.db() as session:
+            async for t, total_count, error_count in await session.stream(stmt):
+                timestamp = _as_datetime(t)
+                data[timestamp] = TraceCountByStatusTimeSeriesDataPoint(
+                    timestamp=timestamp,
+                    ok_count=total_count - error_count,
+                    error_count=error_count,
+                    total_count=total_count,
+                )
+
+        data_timestamps: list[datetime] = [data_point.timestamp for data_point in data.values()]
+        min_time = min([*data_timestamps, time_range.start])
+        max_time = max(
+            [
+                *data_timestamps,
+                *([time_range.end] if time_range.end else []),
+            ],
+            default=datetime.now(timezone.utc),
+        )
+        for timestamp in get_timestamp_range(
+            start_time=min_time,
+            end_time=max_time,
+            stride=field,
+            utc_offset_minutes=utc_offset_minutes,
+        ):
+            if timestamp not in data:
+                data[timestamp] = TraceCountByStatusTimeSeriesDataPoint(
+                    timestamp=timestamp,
+                    ok_count=0,
+                    error_count=0,
+                    total_count=0,
+                )
+        return TraceCountByStatusTimeSeries(data=sorted(data.values(), key=lambda x: x.timestamp))
+
+    @strawberry.field
+    async def trace_latency_ms_percentile_time_series(
+        self,
+        info: Info[Context, None],
+        time_range: TimeRange,
+        time_bin_config: Optional[TimeBinConfig] = UNSET,
+    ) -> TraceLatencyPercentileTimeSeries:
+        if time_range.start is None:
+            raise BadRequest("Start time is required")
+
+        dialect = info.context.db.dialect
+        utc_offset_minutes = 0
+        field: Literal["minute", "hour", "day", "week", "month", "year"] = "hour"
+        if time_bin_config:
+            utc_offset_minutes = time_bin_config.utc_offset_minutes
+            if time_bin_config.scale is TimeBinScale.MINUTE:
+                field = "minute"
+            elif time_bin_config.scale is TimeBinScale.HOUR:
+                field = "hour"
+            elif time_bin_config.scale is TimeBinScale.DAY:
+                field = "day"
+            elif time_bin_config.scale is TimeBinScale.WEEK:
+                field = "week"
+            elif time_bin_config.scale is TimeBinScale.MONTH:
+                field = "month"
+            elif time_bin_config.scale is TimeBinScale.YEAR:
+                field = "year"
+        bucket = date_trunc(dialect, field, models.Trace.start_time, utc_offset_minutes)
+
+        stmt = select(bucket).where(models.Trace.project_rowid == self.project_rowid)
+        if time_range.start:
+            stmt = stmt.where(time_range.start <= models.Trace.start_time)
+        if time_range.end:
+            stmt = stmt.where(models.Trace.start_time < time_range.end)
+
+        if dialect is SupportedSQLDialect.POSTGRESQL:
+            stmt = stmt.add_columns(
+                percentile_cont(0.50).within_group(models.Trace.latency_ms.asc()).label("p50"),
+                percentile_cont(0.75).within_group(models.Trace.latency_ms.asc()).label("p75"),
+                percentile_cont(0.90).within_group(models.Trace.latency_ms.asc()).label("p90"),
+                percentile_cont(0.95).within_group(models.Trace.latency_ms.asc()).label("p95"),
+                percentile_cont(0.99).within_group(models.Trace.latency_ms.asc()).label("p99"),
+                percentile_cont(0.999).within_group(models.Trace.latency_ms.asc()).label("p999"),
+                func.max(models.Trace.latency_ms).label("max"),
+            )
+        elif dialect is SupportedSQLDialect.SQLITE:
+            stmt = stmt.add_columns(
+                func.percentile(models.Trace.latency_ms, 50).label("p50"),
+                func.percentile(models.Trace.latency_ms, 75).label("p75"),
+                func.percentile(models.Trace.latency_ms, 90).label("p90"),
+                func.percentile(models.Trace.latency_ms, 95).label("p95"),
+                func.percentile(models.Trace.latency_ms, 99).label("p99"),
+                func.percentile(models.Trace.latency_ms, 99.9).label("p999"),
+                func.max(models.Trace.latency_ms).label("max"),
+            )
+        else:
+            assert_never(dialect)
+
+        stmt = stmt.group_by(bucket).order_by(bucket)
+
+        data: dict[datetime, TraceLatencyMsPercentileTimeSeriesDataPoint] = {}
+        async with info.context.db() as session:
+            async for (
+                bucket_time,
+                p50,
+                p75,
+                p90,
+                p95,
+                p99,
+                p999,
+                max_latency,
+            ) in await session.stream(stmt):
+                timestamp = _as_datetime(bucket_time)
+                data[timestamp] = TraceLatencyMsPercentileTimeSeriesDataPoint(
+                    timestamp=timestamp,
+                    p50=p50,
+                    p75=p75,
+                    p90=p90,
+                    p95=p95,
+                    p99=p99,
+                    p999=p999,
+                    max=max_latency,
+                )
+
+        data_timestamps: list[datetime] = [data_point.timestamp for data_point in data.values()]
+        min_time = min([*data_timestamps, time_range.start])
+        max_time = max(
+            [
+                *data_timestamps,
+                *([time_range.end] if time_range.end else []),
+            ],
+            default=datetime.now(timezone.utc),
+        )
+        for timestamp in get_timestamp_range(
+            start_time=min_time,
+            end_time=max_time,
+            stride=field,
+            utc_offset_minutes=utc_offset_minutes,
+        ):
+            if timestamp not in data:
+                data[timestamp] = TraceLatencyMsPercentileTimeSeriesDataPoint(timestamp=timestamp)
+        return TraceLatencyPercentileTimeSeries(
+            data=sorted(data.values(), key=lambda x: x.timestamp)
+        )
+
+    @strawberry.field
+    async def trace_token_count_time_series(
+        self,
+        info: Info[Context, None],
+        time_range: TimeRange,
+        time_bin_config: Optional[TimeBinConfig] = UNSET,
+    ) -> TraceTokenCountTimeSeries:
+        if time_range.start is None:
+            raise BadRequest("Start time is required")
+
+        dialect = info.context.db.dialect
+        utc_offset_minutes = 0
+        field: Literal["minute", "hour", "day", "week", "month", "year"] = "hour"
+        if time_bin_config:
+            utc_offset_minutes = time_bin_config.utc_offset_minutes
+            if time_bin_config.scale is TimeBinScale.MINUTE:
+                field = "minute"
+            elif time_bin_config.scale is TimeBinScale.HOUR:
+                field = "hour"
+            elif time_bin_config.scale is TimeBinScale.DAY:
+                field = "day"
+            elif time_bin_config.scale is TimeBinScale.WEEK:
+                field = "week"
+            elif time_bin_config.scale is TimeBinScale.MONTH:
+                field = "month"
+            elif time_bin_config.scale is TimeBinScale.YEAR:
+                field = "year"
+        bucket = date_trunc(dialect, field, models.Trace.start_time, utc_offset_minutes)
+        stmt = (
+            select(
+                bucket,
+                func.sum(models.SpanCost.total_tokens),
+                func.sum(models.SpanCost.prompt_tokens),
+                func.sum(models.SpanCost.completion_tokens),
+            )
+            .join_from(
+                models.Trace,
+                models.SpanCost,
+                onclause=models.SpanCost.trace_rowid == models.Trace.id,
+            )
+            .where(models.Trace.project_rowid == self.project_rowid)
+            .group_by(bucket)
+            .order_by(bucket)
+        )
+        if time_range:
+            if time_range.start:
+                stmt = stmt.where(time_range.start <= models.Trace.start_time)
+            if time_range.end:
+                stmt = stmt.where(models.Trace.start_time < time_range.end)
+        data: dict[datetime, TraceTokenCountTimeSeriesDataPoint] = {}
+        async with info.context.db() as session:
+            async for (
+                t,
+                total_tokens,
+                prompt_tokens,
+                completion_tokens,
+            ) in await session.stream(stmt):
+                timestamp = _as_datetime(t)
+                data[timestamp] = TraceTokenCountTimeSeriesDataPoint(
+                    timestamp=timestamp,
+                    prompt_token_count=prompt_tokens,
+                    completion_token_count=completion_tokens,
+                    total_token_count=total_tokens,
+                )
+
+        data_timestamps: list[datetime] = [data_point.timestamp for data_point in data.values()]
+        min_time = min([*data_timestamps, time_range.start])
+        max_time = max(
+            [
+                *data_timestamps,
+                *([time_range.end] if time_range.end else []),
+            ],
+            default=datetime.now(timezone.utc),
+        )
+        for timestamp in get_timestamp_range(
+            start_time=min_time,
+            end_time=max_time,
+            stride=field,
+            utc_offset_minutes=utc_offset_minutes,
+        ):
+            if timestamp not in data:
+                data[timestamp] = TraceTokenCountTimeSeriesDataPoint(timestamp=timestamp)
+        return TraceTokenCountTimeSeries(data=sorted(data.values(), key=lambda x: x.timestamp))
+
+    @strawberry.field
+    async def trace_token_cost_time_series(
+        self,
+        info: Info[Context, None],
+        time_range: TimeRange,
+        time_bin_config: Optional[TimeBinConfig] = UNSET,
+    ) -> TraceTokenCostTimeSeries:
+        if time_range.start is None:
+            raise BadRequest("Start time is required")
+
+        dialect = info.context.db.dialect
+        utc_offset_minutes = 0
+        field: Literal["minute", "hour", "day", "week", "month", "year"] = "hour"
+        if time_bin_config:
+            utc_offset_minutes = time_bin_config.utc_offset_minutes
+            if time_bin_config.scale is TimeBinScale.MINUTE:
+                field = "minute"
+            elif time_bin_config.scale is TimeBinScale.HOUR:
+                field = "hour"
+            elif time_bin_config.scale is TimeBinScale.DAY:
+                field = "day"
+            elif time_bin_config.scale is TimeBinScale.WEEK:
+                field = "week"
+            elif time_bin_config.scale is TimeBinScale.MONTH:
+                field = "month"
+            elif time_bin_config.scale is TimeBinScale.YEAR:
+                field = "year"
+        bucket = date_trunc(dialect, field, models.Trace.start_time, utc_offset_minutes)
+        stmt = (
+            select(
+                bucket,
+                func.sum(models.SpanCost.total_cost),
+                func.sum(models.SpanCost.prompt_cost),
+                func.sum(models.SpanCost.completion_cost),
+            )
+            .join_from(
+                models.Trace,
+                models.SpanCost,
+                onclause=models.SpanCost.trace_rowid == models.Trace.id,
+            )
+            .where(models.Trace.project_rowid == self.project_rowid)
+            .group_by(bucket)
+            .order_by(bucket)
+        )
+        if time_range:
+            if time_range.start:
+                stmt = stmt.where(time_range.start <= models.Trace.start_time)
+            if time_range.end:
+                stmt = stmt.where(models.Trace.start_time < time_range.end)
+        data: dict[datetime, TraceTokenCostTimeSeriesDataPoint] = {}
+        async with info.context.db() as session:
+            async for (
+                t,
+                total_cost,
+                prompt_cost,
+                completion_cost,
+            ) in await session.stream(stmt):
+                timestamp = _as_datetime(t)
+                data[timestamp] = TraceTokenCostTimeSeriesDataPoint(
+                    timestamp=timestamp,
+                    prompt_cost=prompt_cost,
+                    completion_cost=completion_cost,
+                    total_cost=total_cost,
+                )
+
+        data_timestamps: list[datetime] = [data_point.timestamp for data_point in data.values()]
+        min_time = min([*data_timestamps, time_range.start])
+        max_time = max(
+            [
+                *data_timestamps,
+                *([time_range.end] if time_range.end else []),
+            ],
+            default=datetime.now(timezone.utc),
+        )
+        for timestamp in get_timestamp_range(
+            start_time=min_time,
+            end_time=max_time,
+            stride=field,
+            utc_offset_minutes=utc_offset_minutes,
+        ):
+            if timestamp not in data:
+                data[timestamp] = TraceTokenCostTimeSeriesDataPoint(timestamp=timestamp)
+        return TraceTokenCostTimeSeries(data=sorted(data.values(), key=lambda x: x.timestamp))
+
+    @strawberry.field
+    async def span_annotation_score_time_series(
+        self,
+        info: Info[Context, None],
+        time_range: TimeRange,
+        time_bin_config: Optional[TimeBinConfig] = UNSET,
+    ) -> SpanAnnotationScoreTimeSeries:
+        if time_range.start is None:
+            raise BadRequest("Start time is required")
+
+        dialect = info.context.db.dialect
+        utc_offset_minutes = 0
+        field: Literal["minute", "hour", "day", "week", "month", "year"] = "hour"
+        if time_bin_config:
+            utc_offset_minutes = time_bin_config.utc_offset_minutes
+            if time_bin_config.scale is TimeBinScale.MINUTE:
+                field = "minute"
+            elif time_bin_config.scale is TimeBinScale.HOUR:
+                field = "hour"
+            elif time_bin_config.scale is TimeBinScale.DAY:
+                field = "day"
+            elif time_bin_config.scale is TimeBinScale.WEEK:
+                field = "week"
+            elif time_bin_config.scale is TimeBinScale.MONTH:
+                field = "month"
+            elif time_bin_config.scale is TimeBinScale.YEAR:
+                field = "year"
+        bucket = date_trunc(dialect, field, models.Trace.start_time, utc_offset_minutes)
+        stmt = (
+            select(
+                bucket,
+                models.SpanAnnotation.name,
+                func.avg(models.SpanAnnotation.score).label("average_score"),
+            )
+            .join_from(
+                models.SpanAnnotation,
+                models.Span,
+                onclause=models.SpanAnnotation.span_rowid == models.Span.id,
+            )
+            .join_from(
+                models.Span,
+                models.Trace,
+                onclause=models.Span.trace_rowid == models.Trace.id,
+            )
+            .where(models.Trace.project_rowid == self.project_rowid)
+            .group_by(bucket, models.SpanAnnotation.name)
+            .order_by(bucket)
+        )
+        if time_range:
+            if time_range.start:
+                stmt = stmt.where(time_range.start <= models.Trace.start_time)
+            if time_range.end:
+                stmt = stmt.where(models.Trace.start_time < time_range.end)
+        scores: dict[datetime, dict[str, float]] = {}
+        unique_names: set[str] = set()
+        async with info.context.db() as session:
+            async for (
+                t,
+                name,
+                average_score,
+            ) in await session.stream(stmt):
+                timestamp = _as_datetime(t)
+                if timestamp not in scores:
+                    scores[timestamp] = {}
+                scores[timestamp][name] = average_score
+                unique_names.add(name)
+
+        score_timestamps: list[datetime] = [timestamp for timestamp in scores]
+        min_time = min([*score_timestamps, time_range.start])
+        max_time = max(
+            [
+                *score_timestamps,
+                *([time_range.end] if time_range.end else []),
+            ],
+            default=datetime.now(timezone.utc),
+        )
+        data: dict[datetime, SpanAnnotationScoreTimeSeriesDataPoint] = {
+            timestamp: SpanAnnotationScoreTimeSeriesDataPoint(
+                timestamp=timestamp,
+                scores_with_labels=[
+                    SpanAnnotationScoreWithLabel(label=label, score=scores[timestamp][label])
+                    for label in scores[timestamp]
+                ],
+            )
+            for timestamp in score_timestamps
+        }
+        for timestamp in get_timestamp_range(
+            start_time=min_time,
+            end_time=max_time,
+            stride=field,
+            utc_offset_minutes=utc_offset_minutes,
+        ):
+            if timestamp not in data:
+                data[timestamp] = SpanAnnotationScoreTimeSeriesDataPoint(
+                    timestamp=timestamp,
+                    scores_with_labels=[],
+                )
+        return SpanAnnotationScoreTimeSeries(
+            data=sorted(data.values(), key=lambda x: x.timestamp),
+            names=sorted(list(unique_names)),
+        )
+
 
 @strawberry.type
 class SpanCountTimeSeries(TimeSeries):
@@ -852,6 +1322,80 @@ class TraceCountTimeSeries(TimeSeries):
     """A time series of trace count"""
 
 
+@strawberry.type
+class TraceCountByStatusTimeSeriesDataPoint:
+    timestamp: datetime
+    ok_count: int
+    error_count: int
+    total_count: int
+
+
+@strawberry.type
+class TraceCountByStatusTimeSeries:
+    data: list[TraceCountByStatusTimeSeriesDataPoint]
+
+
+@strawberry.type
+class TraceLatencyMsPercentileTimeSeriesDataPoint:
+    timestamp: datetime
+    p50: Optional[float] = None
+    p75: Optional[float] = None
+    p90: Optional[float] = None
+    p95: Optional[float] = None
+    p99: Optional[float] = None
+    p999: Optional[float] = None
+    max: Optional[float] = None
+
+
+@strawberry.type
+class TraceLatencyPercentileTimeSeries:
+    data: list[TraceLatencyMsPercentileTimeSeriesDataPoint]
+
+
+@strawberry.type
+class TraceTokenCountTimeSeriesDataPoint:
+    timestamp: datetime
+    prompt_token_count: Optional[float] = None
+    completion_token_count: Optional[float] = None
+    total_token_count: Optional[float] = None
+
+
+@strawberry.type
+class TraceTokenCountTimeSeries:
+    data: list[TraceTokenCountTimeSeriesDataPoint]
+
+
+@strawberry.type
+class TraceTokenCostTimeSeriesDataPoint:
+    timestamp: datetime
+    prompt_cost: Optional[float] = None
+    completion_cost: Optional[float] = None
+    total_cost: Optional[float] = None
+
+
+@strawberry.type
+class TraceTokenCostTimeSeries:
+    data: list[TraceTokenCostTimeSeriesDataPoint]
+
+
+@strawberry.type
+class SpanAnnotationScoreWithLabel:
+    label: str
+    score: float
+
+
+@strawberry.type
+class SpanAnnotationScoreTimeSeriesDataPoint:
+    timestamp: datetime
+    scores_with_labels: list[SpanAnnotationScoreWithLabel]
+
+
+@strawberry.type
+class SpanAnnotationScoreTimeSeries:
+    data: list[SpanAnnotationScoreTimeSeriesDataPoint]
+    names: list[str]
+
+
 INPUT_VALUE = SpanAttributes.INPUT_VALUE.split(".")
 OUTPUT_VALUE = SpanAttributes.OUTPUT_VALUE.split(".")
 
@@ -1060,3 +1604,13 @@ async def _paginate_span_by_trace_start_time(
             has_next_page=has_next_page,
         ),
     )
+
+
+def to_gql_project(project: models.Project) -> Project:
+    """
+    Converts an ORM project to a GraphQL project.
+    """
+    return Project(
+        project_rowid=project.id,
+        db_project=project,
+    )