import base64
import json
import re
from dataclasses import dataclass, field
from datetime import datetime, timedelta, timezone
from random import random
from secrets import token_hex
from typing import Any, Literal, Optional

import httpx
import pandas as pd
import pytest
from faker import Faker
from sqlalchemy import insert
from sqlalchemy.ext.asyncio import AsyncSession
from strawberry.relay import GlobalID
from typing_extensions import assert_never

from phoenix.config import DEFAULT_PROJECT_NAME
from phoenix.db import models
from phoenix.server.api.input_types.TimeBinConfig import TimeBinConfig, TimeBinScale
from phoenix.server.api.input_types.TimeRange import TimeRange
from phoenix.server.api.types.pagination import Cursor, CursorSortColumn, CursorSortColumnDataType
from phoenix.server.api.types.Project import Project
from phoenix.server.types import DbSessionFactory
from tests.unit.graphql import AsyncGraphQLClient

from ...._helpers import _add_project, _add_project_session, _add_span, _add_trace, _gid, _node

fake = Faker()

PROJECT_ID = str(GlobalID(type_name="Project", node_id="1"))


async def _add_generative_model(
    session: AsyncSession,
    name: Optional[str] = None,
    provider: str = "openai",
    is_built_in: bool = False,
) -> models.GenerativeModel:
    """Helper function to create a GenerativeModel for testing."""
    name = name or f"test-model-{token_hex(4)}"
    model = models.GenerativeModel(
        name=name,
        provider=provider,
        name_pattern=re.compile(re.escape(name)),
        is_built_in=is_built_in,
    )
    session.add(model)
    await session.flush()
    return model


async def _add_span_cost(
    session: AsyncSession,
    span: models.Span,
    trace: models.Trace,
    model: models.GenerativeModel,
    total_cost: Optional[float] = None,
    total_tokens: Optional[float] = None,
    prompt_cost: Optional[float] = None,
    prompt_tokens: Optional[float] = None,
    completion_cost: Optional[float] = None,
    completion_tokens: Optional[float] = None,
    span_start_time: Optional[datetime] = None,
) -> models.SpanCost:
    """Helper function to create a SpanCost for testing."""
    span_cost = models.SpanCost(
        span_rowid=span.id,
        trace_rowid=trace.id,
        span_start_time=span_start_time or span.start_time,
        model_id=model.id,
        total_cost=total_cost,
        total_tokens=total_tokens,
        prompt_cost=prompt_cost,
        prompt_tokens=prompt_tokens,
        completion_cost=completion_cost,
        completion_tokens=completion_tokens,
    )
    session.add(span_cost)
    await session.flush()
    return span_cost


async def _add_span_cost_detail(
    session: AsyncSession,
    span_cost: models.SpanCost,
    token_type: str,
    is_prompt: bool,
    tokens: Optional[float],
    cost: Optional[float] = None,
) -> models.SpanCostDetail:
    span_cost_detail = models.SpanCostDetail(
        span_cost_id=span_cost.id,
        token_type=token_type,
        is_prompt=is_prompt,
        tokens=tokens,
        cost=cost,
        cost_per_token=None,
    )
    session.add(span_cost_detail)
    await session.flush()
    return span_cost_detail


@dataclass
class _CostTestData:
    project: models.Project
    generative_models: dict[str, models.GenerativeModel]
    spans: list[models.Span]
    traces: list[models.Trace]
    span_costs: list[models.SpanCost]
    base_time: datetime


@pytest.fixture
async def _cost_data(db: DbSessionFactory) -> _CostTestData:
    """Create comprehensive test data for cost and token testing.

    Creates:
    - 3 different models (gpt-4, gpt-3.5-turbo, claude)
    - Multiple spans with different costs and token counts
    - Data spread across different time periods for time range testing
    """
    async with db() as session:
        # Create project
        project = await _add_project(session, name="cost-test-project")

        # Create generative models with distinct characteristics
        # Model 1: High cost per token, moderate volume (gpt-4)
        gpt4 = await _add_generative_model(session, name="gpt-4", provider="openai")

        # Model 2: Lower cost per token, high volume (gpt-3.5-turbo)
        gpt35 = await _add_generative_model(session, name="gpt-3.5-turbo", provider="openai")

        # Model 3: Moderate cost and volume (claude)
        claude = await _add_generative_model(session, name="claude-3-sonnet", provider="anthropic")

        base_time = datetime.fromisoformat("2024-01-01T00:00:00+00:00")
        spans = []
        traces = []
        span_costs = []

        # Time period 1: 2024-01-01 (Hour 0-1)
        # Create spans for gpt-4 - High cost, lower token count
        for i in range(3):
            trace = await _add_trace(
                session,
                project,
                start_time=base_time + timedelta(minutes=i * 10),
                end_time=base_time + timedelta(minutes=i * 10 + 5),
            )
            traces.append(trace)

            span = await _add_span(
                session,
                trace=trace,
                start_time=trace.start_time,
                end_time=trace.end_time,
            )
            spans.append(span)

            # GPT-4: High cost per token, moderate token count
            cost = await _add_span_cost(
                session,
                span=span,
                trace=trace,
                model=gpt4,
                total_cost=1.50 + (i * 0.30),  # $1.50, $1.80, $2.10
                total_tokens=1000 + (i * 200),  # 1000, 1200, 1400 tokens
                prompt_cost=(1.50 + (i * 0.30)) * 0.75,  # 75% prompt cost
                prompt_tokens=(1000 + (i * 200)) * 0.8,  # 80% prompt tokens
                completion_cost=(1.50 + (i * 0.30)) * 0.25,  # 25% completion cost
                completion_tokens=(1000 + (i * 200)) * 0.2,  # 20% completion tokens
                span_start_time=span.start_time,
            )
            span_costs.append(cost)

        # Time period 2: 2024-01-01 (Hour 2-3)
        # Create spans for gpt-3.5-turbo - Lower cost, higher token count
        for i in range(4):  # More spans = higher total
            trace = await _add_trace(
                session,
                project,
                start_time=base_time + timedelta(hours=2, minutes=i * 10),
                end_time=base_time + timedelta(hours=2, minutes=i * 10 + 5),
            )
            traces.append(trace)

            span = await _add_span(
                session,
                trace=trace,
                start_time=trace.start_time,
                end_time=trace.end_time,
            )
            spans.append(span)

            # GPT-3.5: Lower cost per token, higher token count
            cost = await _add_span_cost(
                session,
                span=span,
                trace=trace,
                model=gpt35,
                total_cost=0.40 + (i * 0.20),  # $0.40, $0.60, $0.80, $1.00
                total_tokens=2000 + (i * 500),  # 2000, 2500, 3000, 3500 tokens
                prompt_cost=(0.40 + (i * 0.20)) * 0.70,  # 70% prompt cost
                prompt_tokens=(2000 + (i * 500)) * 0.85,  # 85% prompt tokens
                completion_cost=(0.40 + (i * 0.20)) * 0.30,  # 30% completion cost
                completion_tokens=(2000 + (i * 500)) * 0.15,  # 15% completion tokens
                span_start_time=span.start_time,
            )
            span_costs.append(cost)

        # Time period 3: 2024-01-01 (Hour 4-5)
        # Create spans for claude - Balanced cost and tokens
        for i in range(2):
            trace = await _add_trace(
                session,
                project,
                start_time=base_time + timedelta(hours=4, minutes=i * 15),
                end_time=base_time + timedelta(hours=4, minutes=i * 15 + 8),
            )
            traces.append(trace)

            span = await _add_span(
                session,
                trace=trace,
                start_time=trace.start_time,
                end_time=trace.end_time,
            )
            spans.append(span)

            # Claude: Balanced cost and token count
            cost = await _add_span_cost(
                session,
                span=span,
                trace=trace,
                model=claude,
                total_cost=1.00 + (i * 0.50),  # $1.00, $1.50
                total_tokens=1500 + (i * 300),  # 1500, 1800 tokens
                prompt_cost=(1.00 + (i * 0.50)) * 0.80,  # 80% prompt cost
                prompt_tokens=(1500 + (i * 300)) * 0.75,  # 75% prompt tokens
                completion_cost=(1.00 + (i * 0.50)) * 0.20,  # 20% completion cost
                completion_tokens=(1500 + (i * 300)) * 0.25,  # 25% completion tokens
                span_start_time=span.start_time,
            )
            span_costs.append(cost)

        await session.commit()

        return _CostTestData(
            project=project,
            generative_models={
                "gpt4": gpt4,
                "gpt35": gpt35,
                "claude": claude,
            },
            spans=spans,
            traces=traces,
            span_costs=span_costs,
            base_time=base_time,
        )


class TestTopModels:
    """Comprehensive test suite for top_models_by_token_count and top_models_by_cost GraphQL fields."""

    async def test_top_models_comprehensive(
        self,
        _cost_data: _CostTestData,
        gql_client: AsyncGraphQLClient,
    ) -> None:
        """Comprehensive test for both top_models_by_token_count and top_models_by_cost fields."""
        project = _cost_data.project
        base_time = _cost_data.base_time
        project_gid = str(GlobalID(type_name="Project", node_id=str(project.id)))

        # Full time range for comprehensive testing
        full_time_range = {
            "start": base_time.isoformat(),
            "end": (base_time + timedelta(days=1)).isoformat(),
        }

        # --- TEST 1: Basic ordering for both token count and cost ---

        token_query = """
            query ($projectId: ID!, $timeRange: TimeRange!) {
                node(id: $projectId) {
                    ... on Project {
                        topModelsByTokenCount(timeRange: $timeRange) {
                            name
                            costSummary(projectId: $projectId, timeRange: $timeRange) {
                                total { tokens cost }
                                prompt { tokens cost }
                                completion { tokens cost }
                            }
                        }
                    }
                }
            }
        """

        cost_query = """
            query ($projectId: ID!, $timeRange: TimeRange!) {
                node(id: $projectId) {
                    ... on Project {
                        topModelsByCost(timeRange: $timeRange) {
                            name
                            costSummary(projectId: $projectId, timeRange: $timeRange) {
                                total { tokens cost }
                                prompt { tokens cost }
                                completion { tokens cost }
                            }
                        }
                    }
                }
            }
        """

        # Test token-based ordering
        token_response = await gql_client.execute(
            query=token_query,
            variables={"projectId": project_gid, "timeRange": full_time_range},
        )
        assert not token_response.errors
        assert (token_data := token_response.data) is not None

        token_models = token_data["node"]["topModelsByTokenCount"]
        assert len(token_models) == 3

        # Expected token totals: gpt-3.5-turbo (11K), gpt-4 (3.6K), claude (3.3K)
        assert token_models[0]["name"] == "gpt-3.5-turbo"
        assert token_models[1]["name"] == "gpt-4"
        assert token_models[2]["name"] == "claude-3-sonnet"

        token_counts = [m["costSummary"]["total"]["tokens"] for m in token_models]
        assert token_counts == sorted(token_counts, reverse=True)
        assert token_counts[0] == 11000  # gpt-3.5-turbo
        assert token_counts[1] == 3600  # gpt-4
        assert token_counts[2] == 3300  # claude

        # Test cost-based ordering
        cost_response = await gql_client.execute(
            query=cost_query,
            variables={"projectId": project_gid, "timeRange": full_time_range},
        )
        assert not cost_response.errors
        assert (cost_data := cost_response.data) is not None

        cost_models = cost_data["node"]["topModelsByCost"]
        assert len(cost_models) == 3

        # Expected cost totals: gpt-4 ($5.40), gpt-3.5-turbo ($2.80), claude ($2.50)
        assert cost_models[0]["name"] == "gpt-4"
        assert cost_models[1]["name"] == "gpt-3.5-turbo"
        assert cost_models[2]["name"] == "claude-3-sonnet"

        costs = [m["costSummary"]["total"]["cost"] for m in cost_models]
        assert costs == sorted(costs, reverse=True)
        assert abs(costs[0] - 5.40) < 0.01  # gpt-4
        assert abs(costs[1] - 2.80) < 0.01  # gpt-3.5-turbo
        assert abs(costs[2] - 2.50) < 0.01  # claude

        # Verify ordering is different between token and cost
        token_order = [m["name"] for m in token_models]
        cost_order = [m["name"] for m in cost_models]
        assert token_order != cost_order

        # --- TEST 2: Cost summary calculations accuracy ---

        # Find gpt-4 model and verify its calculations
        gpt4_model = next(m for m in cost_models if m["name"] == "gpt-4")
        cost_summary = gpt4_model["costSummary"]

        # Verify totals match sum of prompt + completion
        total_cost = cost_summary["total"]["cost"]
        prompt_cost = cost_summary["prompt"]["cost"]
        completion_cost = cost_summary["completion"]["cost"]
        assert abs(total_cost - (prompt_cost + completion_cost)) < 0.01

        total_tokens = cost_summary["total"]["tokens"]
        prompt_tokens = cost_summary["prompt"]["tokens"]
        completion_tokens = cost_summary["completion"]["tokens"]
        assert abs(total_tokens - (prompt_tokens + completion_tokens)) < 1

        # --- TEST 3: Time range filtering ---

        # Test filtering to only hour 2-3 (gpt-3.5-turbo data only)
        filtered_time_range = {
            "start": (base_time + timedelta(hours=2)).isoformat(),
            "end": (base_time + timedelta(hours=4)).isoformat(),
        }

        token_filtered_response = await gql_client.execute(
            query=token_query,
            variables={"projectId": project_gid, "timeRange": filtered_time_range},
        )
        assert not token_filtered_response.errors
        assert (token_filtered_data := token_filtered_response.data) is not None

        filtered_models = token_filtered_data["node"]["topModelsByTokenCount"]
        assert len(filtered_models) == 1
        assert filtered_models[0]["name"] == "gpt-3.5-turbo"
        assert filtered_models[0]["costSummary"]["total"]["tokens"] == 11000

        # Test filtering to only hour 0-1 (gpt-4 data only)
        gpt4_time_range = {
            "start": base_time.isoformat(),
            "end": (base_time + timedelta(hours=2)).isoformat(),
        }

        cost_filtered_response = await gql_client.execute(
            query=cost_query,
            variables={"projectId": project_gid, "timeRange": gpt4_time_range},
        )
        assert not cost_filtered_response.errors
        assert (cost_filtered_data := cost_filtered_response.data) is not None

        gpt4_only_models = cost_filtered_data["node"]["topModelsByCost"]
        assert len(gpt4_only_models) == 1
        assert gpt4_only_models[0]["name"] == "gpt-4"
        assert abs(gpt4_only_models[0]["costSummary"]["total"]["cost"] - 5.40) < 0.01

        # --- TEST 4: Partial time ranges ---

        # Query that includes gpt-4 (hours 0-1) and gpt-3.5 (hours 2-3) but excludes claude (hours 4-5)
        partial_time_range = {
            "start": base_time.isoformat(),
            "end": (base_time + timedelta(hours=4)).isoformat(),
        }

        partial_response = await gql_client.execute(
            query=cost_query,
            variables={"projectId": project_gid, "timeRange": partial_time_range},
        )
        assert not partial_response.errors
        assert (partial_data := partial_response.data) is not None

        partial_models = partial_data["node"]["topModelsByCost"]
        assert len(partial_models) == 2  # Should include gpt-4 and gpt-3.5, exclude claude

        model_names = [m["name"] for m in partial_models]
        assert "gpt-4" in model_names
        assert "gpt-3.5-turbo" in model_names
        assert "claude-3-sonnet" not in model_names

        # gpt-4 should still rank higher by cost than gpt-3.5-turbo
        assert partial_models[0]["name"] == "gpt-4"
        assert partial_models[1]["name"] == "gpt-3.5-turbo"

        # --- TEST 5: No data scenarios ---

        empty_time_range = {
            "start": "2023-01-01T00:00:00+00:00",
            "end": "2024-01-01T00:00:00+00:00",
        }

        # Test token count with no data
        empty_token_response = await gql_client.execute(
            query=token_query,
            variables={"projectId": project_gid, "timeRange": empty_time_range},
        )
        assert not empty_token_response.errors
        assert (empty_token_data := empty_token_response.data) is not None
        assert len(empty_token_data["node"]["topModelsByTokenCount"]) == 0

        # Test cost with no data
        empty_cost_response = await gql_client.execute(
            query=cost_query,
            variables={"projectId": project_gid, "timeRange": empty_time_range},
        )
        assert not empty_cost_response.errors
        assert (empty_cost_data := empty_cost_response.data) is not None
        assert len(empty_cost_data["node"]["topModelsByCost"]) == 0

    @pytest.mark.parametrize(
        "bad_project_id",
        [
            pytest.param(
                str(GlobalID(type_name="Span", node_id="1")),
                id="wrong_type_numeric_payload",
            ),
            pytest.param(
                str(GlobalID(type_name="SandboxProvider", node_id="not-a-number")),
                id="wrong_type_non_numeric_payload",
            ),
        ],
    )
    async def test_token_details_reject_a_project_id_of_another_type(
        self,
        bad_project_id: str,
        _cost_data: _CostTestData,
        gql_client: AsyncGraphQLClient,
    ) -> None:
        """A non-Project ID is a bad request, whether or not its payload parses as an int."""
        project_gid = str(GlobalID(type_name="Project", node_id=str(_cost_data.project.id)))
        query = """
            query ($projectId: ID!, $scopeProjectId: ID!, $timeRange: TimeRange!) {
                node(id: $projectId) {
                    ... on Project {
                        topModelsByCost(timeRange: $timeRange) {
                            costDetailSummaryEntries(
                                projectId: $scopeProjectId
                                timeRange: $timeRange
                            ) {
                                tokenType
                            }
                        }
                    }
                }
            }
        """
        response = await gql_client.execute(
            query=query,
            variables={
                "projectId": project_gid,
                "scopeProjectId": bad_project_id,
                "timeRange": {
                    "start": _cost_data.base_time.isoformat(),
                    "end": (_cost_data.base_time + timedelta(minutes=15)).isoformat(),
                },
            },
        )
        assert response.errors
        assert any("Invalid Project ID" in error.message for error in response.errors)

    async def test_token_details_are_aggregated_by_model_project_and_time_range(
        self,
        _cost_data: _CostTestData,
        db: DbSessionFactory,
        gql_client: AsyncGraphQLClient,
    ) -> None:
        project = _cost_data.project
        base_time = _cost_data.base_time
        gpt4 = _cost_data.generative_models["gpt4"]
        gpt4_span_costs = _cost_data.span_costs[:3]

        async with db() as session:
            for span_cost, input_tokens, cache_read_tokens, output_tokens, cost in (
                (gpt4_span_costs[0], 700, 100, 200, 1.5),
                (gpt4_span_costs[1], 800, 160, 240, 1.8),
                (gpt4_span_costs[2], 900, 220, 280, 2.1),
            ):
                await _add_span_cost_detail(
                    session,
                    span_cost,
                    token_type="input",
                    is_prompt=True,
                    tokens=input_tokens,
                    cost=cost * 0.55,
                )
                await _add_span_cost_detail(
                    session,
                    span_cost,
                    token_type="cache_read",
                    is_prompt=True,
                    tokens=cache_read_tokens,
                    cost=cost * 0.20,
                )
                await _add_span_cost_detail(
                    session,
                    span_cost,
                    token_type="output",
                    is_prompt=False,
                    tokens=output_tokens,
                    cost=cost * 0.25,
                )

            other_project = await _add_project(session, name="other-cost-detail-project")
            other_trace = await _add_trace(
                session,
                other_project,
                start_time=base_time,
                end_time=base_time + timedelta(minutes=1),
            )
            other_span = await _add_span(
                session,
                trace=other_trace,
                start_time=other_trace.start_time,
                end_time=other_trace.end_time,
            )
            other_span_cost = await _add_span_cost(
                session,
                span=other_span,
                trace=other_trace,
                model=gpt4,
                total_cost=100,
                total_tokens=100_000,
                prompt_cost=100,
                prompt_tokens=100_000,
                completion_cost=0,
                completion_tokens=0,
            )
            await _add_span_cost_detail(
                session,
                other_span_cost,
                token_type="cache_read",
                is_prompt=True,
                tokens=100_000,
                cost=100,
            )
            await session.commit()

        query = """
            query ($projectId: ID!, $timeRange: TimeRange!) {
                node(id: $projectId) {
                    ... on Project {
                        topModelsByTokenCount(timeRange: $timeRange) {
                            name
                            costDetailSummaryEntries(
                                projectId: $projectId
                                timeRange: $timeRange
                            ) {
                                tokenType
                                isPrompt
                                value { tokens cost }
                            }
                        }
                        topModelsByCost(timeRange: $timeRange) {
                            name
                            costDetailSummaryEntries(
                                projectId: $projectId
                                timeRange: $timeRange
                            ) {
                                tokenType
                                isPrompt
                                value { tokens cost }
                            }
                        }
                    }
                }
            }
        """
        project_id = str(GlobalID(type_name="Project", node_id=str(project.id)))
        response = await gql_client.execute(
            query=query,
            variables={
                "projectId": project_id,
                "timeRange": {
                    "start": base_time.isoformat(),
                    "end": (base_time + timedelta(minutes=15)).isoformat(),
                },
            },
        )
        assert response.data is not None
        assert not response.errors

        for field_name in ("topModelsByTokenCount", "topModelsByCost"):
            model = next(
                model for model in response.data["node"][field_name] if model["name"] == "gpt-4"
            )
            details = {
                (entry["isPrompt"], entry["tokenType"]): entry["value"]
                for entry in model["costDetailSummaryEntries"]
            }
            assert details == {
                (True, "cache_read"): {"tokens": 260, "cost": pytest.approx(0.66)},
                (True, "input"): {"tokens": 1500, "cost": pytest.approx(1.815)},
                (False, "output"): {"tokens": 440, "cost": pytest.approx(0.825)},
            }


class TestTraceTokenCountTimeSeries:
    _QUERY = """
        query ($projectId: ID!, $timeRange: TimeRange!, $timeBinConfig: TimeBinConfig!) {
            node(id: $projectId) {
                ... on Project {
                    traceTokenCountTimeSeries(
                        timeRange: $timeRange
                        timeBinConfig: $timeBinConfig
                    ) {
                        data {
                            timestamp
                            promptTokenCount
                            completionTokenCount
                            totalTokenCount
                            promptTokenCountDetails {
                                tokenType
                                tokenCount
                            }
                            completionTokenCountDetails {
                                tokenType
                                tokenCount
                            }
                        }
                    }
                }
            }
        }
    """

    async def test_token_detail_breakdowns_match_prompt_and_completion_totals(
        self,
        db: DbSessionFactory,
        gql_client: AsyncGraphQLClient,
    ) -> None:
        base_time = datetime.fromisoformat("2024-01-01T00:00:00+00:00")
        async with db() as session:
            project = await _add_project(session, name="token-detail-time-series-project")
            model = await _add_generative_model(session, name="gpt-4o-mini", provider="openai")

            trace_with_details = await _add_trace(
                session,
                project,
                start_time=base_time + timedelta(minutes=10),
                end_time=base_time + timedelta(minutes=12),
            )
            span_with_details = await _add_span(
                session,
                trace=trace_with_details,
                start_time=trace_with_details.start_time,
                end_time=trace_with_details.end_time,
            )
            span_cost_with_details = await _add_span_cost(
                session,
                span=span_with_details,
                trace=trace_with_details,
                model=model,
                total_tokens=170,
                prompt_tokens=120,
                completion_tokens=50,
                span_start_time=span_with_details.start_time,
            )
            await _add_span_cost_detail(
                session,
                span_cost=span_cost_with_details,
                token_type="input",
                is_prompt=True,
                tokens=50,
            )
            await _add_span_cost_detail(
                session,
                span_cost=span_cost_with_details,
                token_type="cache_read",
                is_prompt=True,
                tokens=40,
            )
            await _add_span_cost_detail(
                session,
                span_cost=span_cost_with_details,
                token_type="cache_write",
                is_prompt=True,
                tokens=30,
            )
            await _add_span_cost_detail(
                session,
                span_cost=span_cost_with_details,
                token_type="output",
                is_prompt=False,
                tokens=45,
            )
            await _add_span_cost_detail(
                session,
                span_cost=span_cost_with_details,
                token_type="reasoning",
                is_prompt=False,
                tokens=5,
            )

            trace_without_details = await _add_trace(
                session,
                project,
                start_time=base_time + timedelta(minutes=35),
                end_time=base_time + timedelta(minutes=37),
            )
            span_without_details = await _add_span(
                session,
                trace=trace_without_details,
                start_time=trace_without_details.start_time,
                end_time=trace_without_details.end_time,
            )
            await _add_span_cost(
                session,
                span=span_without_details,
                trace=trace_without_details,
                model=model,
                total_tokens=80,
                prompt_tokens=60,
                completion_tokens=20,
                span_start_time=span_without_details.start_time,
            )
            await session.commit()

        response = await gql_client.execute(
            query=self._QUERY,
            variables={
                "projectId": str(GlobalID(type_name="Project", node_id=str(project.id))),
                "timeRange": {
                    "start": base_time.isoformat(),
                    "end": (base_time + timedelta(hours=2)).isoformat(),
                },
                "timeBinConfig": {"scale": "HOUR", "utcOffsetMinutes": 0},
            },
        )
        assert not response.errors
        assert response.data is not None

        data = response.data["node"]["traceTokenCountTimeSeries"]["data"]
        bucket = next(data_point for data_point in data if data_point["totalTokenCount"] == 250)
        assert bucket["promptTokenCount"] == 180
        assert bucket["completionTokenCount"] == 70

        prompt_details = {
            detail["tokenType"]: detail["tokenCount"]
            for detail in bucket["promptTokenCountDetails"]
        }
        completion_details = {
            detail["tokenType"]: detail["tokenCount"]
            for detail in bucket["completionTokenCountDetails"]
        }
        assert prompt_details == {
            "input": 110,
            "cache_read": 40,
            "cache_write": 30,
        }
        assert completion_details == {
            "output": 65,
            "reasoning": 5,
        }
        assert sum(prompt_details.values()) == bucket["promptTokenCount"]
        assert sum(completion_details.values()) == bucket["completionTokenCount"]
        assert (
            bucket["promptTokenCount"] + bucket["completionTokenCount"] == bucket["totalTokenCount"]
        )


@pytest.mark.parametrize(
    "variables, start_cursor, end_cursor, has_next_page",
    [
        pytest.param(
            {
                "projectId": PROJECT_ID,
                "first": 2,
            },
            Cursor(rowid=1),
            Cursor(rowid=2),
            True,
            id="basic-query",
        ),
        pytest.param(
            {
                "projectId": PROJECT_ID,
                "after": str(Cursor(rowid=13)),
                "first": 2,
            },
            Cursor(rowid=14),
            Cursor(rowid=15),
            False,
            id="page-ends-exactly-on-last-record",
        ),
        pytest.param(
            {
                "projectId": PROJECT_ID,
                "after": str(Cursor(14)),
                "first": 2,
            },
            Cursor(rowid=15),
            Cursor(rowid=15),
            False,
            id="page-ends-before-it-reaches-limit",
        ),
        pytest.param(
            {
                "projectId": PROJECT_ID,
                "first": 2,
                "filterCondition": "span_kind == 'LLM'",
            },
            Cursor(rowid=5),
            Cursor(rowid=10),
            True,
            id="filter-condition",
        ),
        pytest.param(
            {
                "projectId": PROJECT_ID,
                "first": 2,
                "after": str(Cursor(5)),  # skip the first span satisfying the filter condition
                "filterCondition": "span_kind == 'LLM'",
            },
            Cursor(rowid=10),
            Cursor(rowid=15),
            False,
            id="filter-condition-with-cursor",
        ),
        pytest.param(
            {
                "projectId": PROJECT_ID,
                "sort": {"col": "startTime", "dir": "desc"},
                "first": 2,
            },
            Cursor(
                rowid=15,
                sort_column=CursorSortColumn(
                    type=CursorSortColumnDataType.DATETIME,
                    value=datetime.fromisoformat("2023-12-11T17:43:26.706204+00:00"),
                ),
            ),
            Cursor(
                rowid=14,
                sort_column=CursorSortColumn(
                    type=CursorSortColumnDataType.DATETIME,
                    value=datetime.fromisoformat("2023-12-11T17:43:26.704532+00:00"),
                ),
            ),
            True,
            id="sort-by-descending-start-time",
        ),
        pytest.param(
            {
                "projectId": PROJECT_ID,
                "sort": {"col": "startTime", "dir": "asc"},
                "first": 2,
            },
            Cursor(
                rowid=1,
                sort_column=CursorSortColumn(
                    type=CursorSortColumnDataType.DATETIME,
                    value=datetime.fromisoformat("2023-12-11T17:43:23.306838+00:00"),
                ),
            ),
            Cursor(
                rowid=2,
                sort_column=CursorSortColumn(
                    type=CursorSortColumnDataType.DATETIME,
                    value=datetime.fromisoformat("2023-12-11T17:43:23.306945+00:00"),
                ),
            ),
            True,
            id="sort-by-ascending-start-time",
        ),
        pytest.param(
            {
                "projectId": PROJECT_ID,
                "sort": {"col": "cumulativeTokenCountTotal", "dir": "desc"},
                "first": 2,
            },
            Cursor(
                rowid=15,
                sort_column=CursorSortColumn(
                    type=CursorSortColumnDataType.INT,
                    value=382,
                ),
            ),
            Cursor(
                rowid=14,
                sort_column=CursorSortColumn(
                    type=CursorSortColumnDataType.INT,
                    value=382,
                ),
            ),
            True,
            id="sort-by-descending-cumulative-prompt-token-count-total",
        ),
        pytest.param(
            {
                "projectId": PROJECT_ID,
                "sort": {"col": "cumulativeTokenCountTotal", "dir": "asc"},
                "first": 2,
            },
            Cursor(
                rowid=2,
                sort_column=CursorSortColumn(
                    type=CursorSortColumnDataType.INT,
                    value=0,
                ),
            ),
            Cursor(
                rowid=3,
                sort_column=CursorSortColumn(
                    type=CursorSortColumnDataType.INT,
                    value=0,
                ),
            ),
            True,
            id="sort-by-ascending-cumulative-prompt-token-count-total",
        ),
        pytest.param(
            {
                "projectId": PROJECT_ID,
                "sort": {"col": "startTime", "dir": "desc"},
                "first": 2,
                "after": str(
                    Cursor(
                        3,
                        sort_column=CursorSortColumn.from_string(
                            type=CursorSortColumnDataType.DATETIME,
                            cursor_string="2023-12-11T17:43:23.307166+00:00",
                        ),
                    )
                ),
            },
            Cursor(
                rowid=2,
                sort_column=CursorSortColumn(
                    type=CursorSortColumnDataType.DATETIME,
                    value=datetime.fromisoformat("2023-12-11T17:43:23.306945+00:00"),
                ),
            ),
            Cursor(
                rowid=1,
                sort_column=CursorSortColumn(
                    type=CursorSortColumnDataType.DATETIME,
                    value=datetime.fromisoformat("2023-12-11T17:43:23.306838+00:00"),
                ),
            ),
            False,
            id="sort-by-descending-start-time-with-cursor",
        ),
        pytest.param(
            {
                "projectId": PROJECT_ID,
                "sort": {"col": "startTime", "dir": "asc"},
                "first": 2,
                "after": str(
                    Cursor(
                        3,
                        sort_column=CursorSortColumn.from_string(
                            type=CursorSortColumnDataType.DATETIME,
                            cursor_string="2023-12-11T17:43:23.307166+00:00",
                        ),
                    )
                ),
            },
            Cursor(
                rowid=4,
                sort_column=CursorSortColumn(
                    type=CursorSortColumnDataType.DATETIME,
                    value=datetime.fromisoformat("2023-12-11T17:43:23.710148+00:00"),
                ),
            ),
            Cursor(
                rowid=5,
                sort_column=CursorSortColumn(
                    type=CursorSortColumnDataType.DATETIME,
                    value=datetime.fromisoformat("2023-12-11T17:43:23.712144+00:00"),
                ),
            ),
            True,
            id="sort-by-ascending-start-time-with-cursor",
        ),
        pytest.param(
            {
                "projectId": PROJECT_ID,
                "sort": {"col": "cumulativeTokenCountTotal", "dir": "desc"},
                "first": 2,
                "after": str(
                    Cursor(
                        rowid=4,  # row 4 is in between rows 1 and 5, which also have 296 cumulative prompt tokens
                        sort_column=CursorSortColumn(
                            type=CursorSortColumnDataType.FLOAT, value=296
                        ),
                    )
                ),
            },
            Cursor(
                rowid=1,
                sort_column=CursorSortColumn(
                    type=CursorSortColumnDataType.INT,
                    value=296,
                ),
            ),
            Cursor(
                rowid=13,
                sort_column=CursorSortColumn(
                    type=CursorSortColumnDataType.INT,
                    value=0,
                ),
            ),
            True,
            id="sort-by-descending-cumulative-prompt-token-count-total-with-cursor",
        ),
        pytest.param(
            {
                "projectId": PROJECT_ID,
                "sort": {"col": "cumulativeTokenCountTotal", "dir": "asc"},
                "first": 2,
                "after": str(
                    Cursor(
                        rowid=4,  # row 4 is in between rows 1 and 5, which also have 296 cumulative prompt tokens
                        sort_column=CursorSortColumn(
                            type=CursorSortColumnDataType.FLOAT, value=296
                        ),
                    )
                ),
            },
            Cursor(
                rowid=5,
                sort_column=CursorSortColumn(
                    type=CursorSortColumnDataType.INT,
                    value=296,
                ),
            ),
            Cursor(
                rowid=6,
                sort_column=CursorSortColumn(
                    type=CursorSortColumnDataType.INT,
                    value=336,
                ),
            ),
            True,
            id="sort-by-ascending-cumulative-prompt-token-count-total-with-cursor",
        ),
        pytest.param(
            {
                "projectId": PROJECT_ID,
                "sort": {
                    "evalResultKey": {"name": "Hallucination", "attr": "label"},
                    "dir": "desc",
                },
                "first": 2,
            },
            Cursor(
                rowid=11,
                sort_column=CursorSortColumn(
                    type=CursorSortColumnDataType.STRING,
                    value="hallucinated",
                ),
            ),
            Cursor(
                rowid=1,
                sort_column=CursorSortColumn(
                    type=CursorSortColumnDataType.STRING,
                    value="hallucinated",
                ),
            ),
            True,
            id="sort-by-descending-hallucination-eval-label",
        ),
        pytest.param(
            {
                "projectId": PROJECT_ID,
                "sort": {
                    "evalResultKey": {"name": "Hallucination", "attr": "label"},
                    "dir": "asc",
                },
                "first": 2,
            },
            Cursor(
                rowid=6,
                sort_column=CursorSortColumn(
                    type=CursorSortColumnDataType.STRING,
                    value="factual",
                ),
            ),
            Cursor(
                rowid=1,
                sort_column=CursorSortColumn(
                    type=CursorSortColumnDataType.STRING,
                    value="hallucinated",
                ),
            ),
            True,
            id="sort-by-ascending-hallucination-eval-label",
        ),
        pytest.param(
            {
                "projectId": PROJECT_ID,
                "sort": {
                    "evalResultKey": {"name": "Hallucination", "attr": "label"},
                    "dir": "desc",
                },
                "first": 2,
                "after": str(
                    Cursor(
                        rowid=11,
                        sort_column=CursorSortColumn(
                            type=CursorSortColumnDataType.STRING, value="hallucinated"
                        ),
                    )
                ),
            },
            Cursor(
                rowid=1,
                sort_column=CursorSortColumn(
                    type=CursorSortColumnDataType.STRING,
                    value="hallucinated",
                ),
            ),
            Cursor(
                rowid=6,
                sort_column=CursorSortColumn(
                    type=CursorSortColumnDataType.STRING,
                    value="factual",
                ),
            ),
            False,
            id="sort-by-descending-hallucination-eval-label-with-cursor",
        ),
        pytest.param(
            {
                "projectId": PROJECT_ID,
                "sort": {
                    "evalResultKey": {"name": "Hallucination", "attr": "label"},
                    "dir": "asc",
                },
                "first": 2,
                "after": str(
                    Cursor(
                        rowid=6,
                        sort_column=CursorSortColumn(
                            type=CursorSortColumnDataType.STRING, value="factual"
                        ),
                    )
                ),
            },
            Cursor(
                rowid=1,
                sort_column=CursorSortColumn(
                    type=CursorSortColumnDataType.STRING,
                    value="hallucinated",
                ),
            ),
            Cursor(
                rowid=11,
                sort_column=CursorSortColumn(
                    type=CursorSortColumnDataType.STRING,
                    value="hallucinated",
                ),
            ),
            False,
            id="sort-by-ascending-hallucination-eval-label-with-cursor",
        ),
        pytest.param(
            {
                "projectId": PROJECT_ID,
                "sort": {
                    "evalResultKey": {"name": "Hallucination", "attr": "score"},
                    "dir": "desc",
                },
                "first": 2,
                "after": str(
                    Cursor(
                        rowid=11,
                        sort_column=CursorSortColumn(type=CursorSortColumnDataType.FLOAT, value=0),
                    )
                ),
            },
            Cursor(
                rowid=1,
                sort_column=CursorSortColumn(
                    type=CursorSortColumnDataType.FLOAT,
                    value=0.0,
                ),
            ),
            Cursor(
                rowid=1,
                sort_column=CursorSortColumn(
                    type=CursorSortColumnDataType.FLOAT,
                    value=0.0,
                ),
            ),
            False,
            id="sort-by-descending-hallucination-eval-score-with-cursor",
        ),
        pytest.param(
            {
                "projectId": PROJECT_ID,
                "sort": {
                    "evalResultKey": {"name": "Hallucination", "attr": "score"},
                    "dir": "asc",
                },
                "first": 5,
                "after": str(
                    Cursor(
                        rowid=1,
                        sort_column=CursorSortColumn(type=CursorSortColumnDataType.FLOAT, value=0),
                    )
                ),
            },
            Cursor(
                rowid=11,
                sort_column=CursorSortColumn(
                    type=CursorSortColumnDataType.FLOAT,
                    value=0.0,
                ),
            ),
            Cursor(
                rowid=6,
                sort_column=CursorSortColumn(
                    type=CursorSortColumnDataType.FLOAT,
                    value=1.0,
                ),
            ),
            False,
            id="sort-by-ascending-hallucination-eval-score-with-cursor",
        ),
    ],
)
async def test_project_spans(
    variables: dict[str, Any],
    start_cursor: Cursor,
    end_cursor: Cursor,
    has_next_page: bool,
    gql_client: AsyncGraphQLClient,
    llama_index_rag_spans: Any,
) -> None:
    query = """
      query ($projectId: ID!, $after: String = null, $before: String = null, $filterCondition: String = null, $first: Int!, $last: Int = null, $sort: SpanSort = null) {
        node(id: $projectId) {
          ... on Project {
            spans(
              after: $after
              before: $before
              filterCondition: $filterCondition
              first: $first
              last: $last
              rootSpansOnly: false
              sort: $sort
            ) {
              edges {
                cursor
              }
              pageInfo {
                hasNextPage
                startCursor
                endCursor
              }
            }
          }
        }
      }
    """
    response = await gql_client.execute(query=query, variables=variables)
    assert not response.errors
    assert (data := response.data) is not None
    spans = data["node"]["spans"]
    page_info = spans["pageInfo"]
    assert Cursor.from_string(page_info["startCursor"]) == start_cursor
    assert Cursor.from_string(page_info["endCursor"]) == end_cursor
    assert page_info["hasNextPage"] == has_next_page
    edges = spans["edges"]
    assert Cursor.from_string(edges[0]["cursor"]) == start_cursor
    assert Cursor.from_string(edges[-1]["cursor"]) == end_cursor


async def test_project_spans_require_first(
    monkeypatch: pytest.MonkeyPatch,
    gql_client: AsyncGraphQLClient,
    llama_index_rag_spans: Any,
) -> None:
    monkeypatch.setenv("PHOENIX_MASK_INTERNAL_SERVER_ERRORS", "false")
    query = """
      query ($projectId: ID!) {
        node(id: $projectId) {
          ... on Project {
            spans {
              edges {
                node {
                  id
                }
              }
            }
          }
        }
      }
    """
    response = await gql_client.execute(query=query, variables={"projectId": PROJECT_ID})
    assert response.errors
    assert len(response.errors) == 1
    assert (
        response.errors[0].message
        == "Field 'spans' argument 'first' of type 'Int!' is required, but it was not provided."
    )


async def test_project_spans_limit_first_page_size(
    monkeypatch: pytest.MonkeyPatch,
    gql_client: AsyncGraphQLClient,
    llama_index_rag_spans: Any,
) -> None:
    monkeypatch.setenv("PHOENIX_MASK_INTERNAL_SERVER_ERRORS", "false")
    query = """
      query ($projectId: ID!, $first: Int!) {
        node(id: $projectId) {
          ... on Project {
            spans(first: $first) {
              edges {
                node {
                  id
                }
              }
            }
          }
        }
      }
    """
    response = await gql_client.execute(
        query=query,
        variables={"projectId": PROJECT_ID, "first": 1001},
    )
    assert response.errors
    assert len(response.errors) == 1
    assert response.errors[0].message == "`first` must be less than or equal to 1000"


@pytest.fixture
async def llama_index_rag_spans(db: DbSessionFactory) -> None:
    # Inserts the first three traces from the llama-index-rag trace fixture
    # (minus embeddings) along with associated span evaluations.
    async with db() as session:
        project_row_id = await session.scalar(
            insert(models.Project).values(name=DEFAULT_PROJECT_NAME).returning(models.Project.id)
        )
        trace_rowids = (
            await session.scalars(
                insert(models.Trace).returning(models.Trace.id),
                [
                    {
                        "trace_id": "0f5bb2e69a0640de87b9d424622b9f13",
                        "project_rowid": project_row_id,
                        "start_time": datetime.fromisoformat("2023-12-11T17:43:23.306838+00:00"),
                        "end_time": datetime.fromisoformat("2023-12-11T17:43:25.534589+00:00"),
                    },
                    {
                        "trace_id": "a4083327f7d0400a9e99906242e71aa4",
                        "project_rowid": project_row_id,
                        "start_time": datetime.fromisoformat("2023-12-11T17:43:25.540371+00:00"),
                        "end_time": datetime.fromisoformat("2023-12-11T17:43:26.492242+00:00"),
                    },
                    {
                        "trace_id": "17f383d1c85648899368bde24b566411",
                        "project_rowid": project_row_id,
                        "start_time": datetime.fromisoformat("2023-12-11T17:43:26.495969+00:00"),
                        "end_time": datetime.fromisoformat("2023-12-11T17:43:27.336284+00:00"),
                    },
                ],
            )
        ).all()
        span_rowids = (
            await session.scalars(
                insert(models.Span).returning(models.Span.id),
                [
                    {
                        "trace_rowid": trace_rowids[0],
                        "span_id": "c0055a08295841ab946f2a16e5089fad",
                        "parent_id": None,
                        "name": "query",
                        "span_kind": "CHAIN",
                        "start_time": datetime.fromisoformat("2023-12-11T17:43:23.306838+00:00"),
                        "end_time": datetime.fromisoformat("2023-12-11T17:43:25.534589+00:00"),
                        "attributes": {
                            "openinference": {"span": {"kind": "CHAIN"}},
                            "output": {
                                "value": "To use the SDK to upload a ranking model, you can follow the documentation provided by the SDK. The documentation will guide you through the necessary steps to upload the model and integrate it into your system. Make sure to carefully follow the instructions to ensure a successful upload and integration process."
                            },
                            "input": {"value": "How do I use the SDK to upload a ranking model?"},
                        },
                        "events": [],
                        "status_code": "OK",
                        "status_message": "",
                        "cumulative_error_count": 0,
                        "cumulative_llm_token_count_prompt": 240,
                        "cumulative_llm_token_count_completion": 56,
                    },
                    {
                        "trace_rowid": trace_rowids[0],
                        "span_id": "edcd8a83c7b34fd2b83e946f58e9a9c0",
                        "parent_id": "c0055a08295841ab946f2a16e5089fad",
                        "name": "retrieve",
                        "span_kind": "RETRIEVER",
                        "start_time": datetime.fromisoformat("2023-12-11T17:43:23.306945+00:00"),
                        "end_time": datetime.fromisoformat("2023-12-11T17:43:23.710062+00:00"),
                        "attributes": {
                            "openinference": {"span": {"kind": "RETRIEVER"}},
                            "retrieval": {
                                "documents": [
                                    {
                                        "document": {
                                            "content": "\nRanking models are used by search engines to display query results ranked in the order of the highest relevance. These predictions seek to maximize user actions that are then used to evaluate model performance.&#x20;\n\nThe complexity within a ranking model makes failures challenging to pinpoint as a model\u2019s dimensions expand per recommendation. Notable challenges within ranking models include upstream data quality issues, poor-performing segments, the cold start problem, and more. &#x20;\n\n\n\n",
                                            "id": "ad17eeea-e339-4195-991b-8eef54b1db65",
                                            "score": 0.8022561073303223,
                                        }
                                    },
                                    {
                                        "document": {
                                            "content": "\n**Use the 'arize-demo-hotel-ranking' model, available in all free accounts, to follow along.**&#x20;\n\n",
                                            "id": "0ce66871-4a50-4d2f-94d2-1531924bf48a",
                                            "score": 0.7964192032814026,
                                        }
                                    },
                                ]
                            },
                            "input": {"value": "How do I use the SDK to upload a ranking model?"},
                        },
                        "events": [],
                        "status_code": "OK",
                        "status_message": "",
                        "cumulative_error_count": 0,
                        "cumulative_llm_token_count_prompt": 0,
                        "cumulative_llm_token_count_completion": 0,
                    },
                    {
                        "trace_rowid": trace_rowids[0],
                        "span_id": "a91ad81bb187489093afeb8f3f5816b5",
                        "parent_id": "edcd8a83c7b34fd2b83e946f58e9a9c0",
                        "name": "embedding",
                        "span_kind": "EMBEDDING",
                        "start_time": datetime.fromisoformat("2023-12-11T17:43:23.307166+00:00"),
                        "end_time": datetime.fromisoformat("2023-12-11T17:43:23.638792+00:00"),
                        "attributes": {
                            "openinference": {"span": {"kind": "EMBEDDING"}},
                            "embedding": {
                                "model_name": "text-embedding-ada-002",
                                "embeddings": [
                                    {
                                        "embedding": {
                                            "vector": [1.0],
                                            "text": "How do I use the SDK to upload a ranking model?",
                                        }
                                    }
                                ],
                            },
                        },
                        "events": [],
                        "status_code": "OK",
                        "status_message": "",
                        "cumulative_error_count": 0,
                        "cumulative_llm_token_count_prompt": 0,
                        "cumulative_llm_token_count_completion": 0,
                    },
                    {
                        "trace_rowid": trace_rowids[0],
                        "span_id": "78742859b73e427f90b43ec6cc8c42ba",
                        "parent_id": "c0055a08295841ab946f2a16e5089fad",
                        "name": "synthesize",
                        "span_kind": "CHAIN",
                        "start_time": datetime.fromisoformat("2023-12-11T17:43:23.710148+00:00"),
                        "end_time": datetime.fromisoformat("2023-12-11T17:43:25.534461+00:00"),
                        "attributes": {
                            "openinference": {"span": {"kind": "CHAIN"}},
                            "output": {
                                "value": "To use the SDK to upload a ranking model, you can follow the documentation provided by the SDK. The documentation will guide you through the necessary steps to upload the model and integrate it into your system. Make sure to carefully follow the instructions to ensure a successful upload and integration process."
                            },
                            "input": {"value": "How do I use the SDK to upload a ranking model?"},
                        },
                        "events": [],
                        "status_code": "OK",
                        "status_message": "",
                        "cumulative_error_count": 0,
                        "cumulative_llm_token_count_prompt": 240,
                        "cumulative_llm_token_count_completion": 56,
                    },
                    {
                        "trace_rowid": trace_rowids[0],
                        "span_id": "258bef0a3e384bcaaa5a388065af0d8f",
                        "parent_id": "78742859b73e427f90b43ec6cc8c42ba",
                        "name": "llm",
                        "span_kind": "LLM",
                        "start_time": datetime.fromisoformat("2023-12-11T17:43:23.712144+00:00"),
                        "end_time": datetime.fromisoformat("2023-12-11T17:43:25.532714+00:00"),
                        "attributes": {
                            "llm": {
                                "invocation_parameters": '{"model": "gpt-3.5-turbo", "temperature": 0.0, "max_tokens": None}',
                                "input_messages": [
                                    {
                                        "message": {
                                            "role": "system",
                                        }
                                    },
                                    {
                                        "message": {
                                            "content": "Context information is below.\n---------------------\nRanking models are used by search engines to display query results ranked in the order of the highest relevance. These predictions seek to maximize user actions that are then used to evaluate model performance.&#x20;\n\nThe complexity within a ranking model makes failures challenging to pinpoint as a model\u2019s dimensions expand per recommendation. Notable challenges within ranking models include upstream data quality issues, poor-performing segments, the cold start problem, and more. &#x20;\n\n**Use the 'arize-demo-hotel-ranking' model, available in all free accounts, to follow along.**&#x20;\n---------------------\nGiven the context information and not prior knowledge, answer the query.\nQuery: How do I use the SDK to upload a ranking model?\nAnswer: ",
                                            "role": "user",
                                        }
                                    },
                                ],
                                "model_name": "gpt-3.5-turbo",
                                "output_messages": [
                                    {
                                        "message": {
                                            "content": "To use the SDK to upload a ranking model, you can follow the documentation provided by the SDK. The documentation will guide you through the necessary steps to upload the model and integrate it into your system. Make sure to carefully follow the instructions to ensure a successful upload and integration process.",
                                            "role": "assistant",
                                        }
                                    }
                                ],
                                "prompt_template": {
                                    "template": "system: You are an expert Q&A system that is trusted around the world.\nAlways answer the query using the provided context information, and not prior knowledge.\nSome rules to follow:\n1. Never directly reference the given context in your answer.\n2. Avoid statements like 'Based on the context, ...' or 'The context information ...' or anything along those lines.\nuser: Context information is below.\n---------------------\n{context_str}\n---------------------\nGiven the context information and not prior knowledge, answer the query.\nQuery: {query_str}\nAnswer: \nassistant: ",
                                    "variables": {
                                        "context_str": "Ranking models are used by search engines to display query results ranked in the order of the highest relevance. These predictions seek to maximize user actions that are then used to evaluate model performance.&#x20;\n\nThe complexity within a ranking model makes failures challenging to pinpoint as a model\u2019s dimensions expand per recommendation. Notable challenges within ranking models include upstream data quality issues, poor-performing segments, the cold start problem, and more. &#x20;\n\n**Use the 'arize-demo-hotel-ranking' model, available in all free accounts, to follow along.**&#x20;",
                                        "query_str": "How do I use the SDK to upload a ranking model?",
                                    },
                                },
                                "token_count": {
                                    "prompt": 240.0,
                                    "total": 296.0,
                                    "completion": 56.0,
                                },
                            },
                            "output": {
                                "value": "To use the SDK to upload a ranking model, you can follow the documentation provided by the SDK. The documentation will guide you through the necessary steps to upload the model and integrate it into your system. Make sure to carefully follow the instructions to ensure a successful upload and integration process."
                            },
                            "openinference": {"span": {"kind": "LLM"}},
                        },
                        "events": [],
                        "status_code": "OK",
                        "status_message": "",
                        "cumulative_error_count": 0,
                        "cumulative_llm_token_count_prompt": 240,
                        "cumulative_llm_token_count_completion": 56,
                    },
                    {
                        "trace_rowid": trace_rowids[1],
                        "span_id": "094ae70b0e9c4dec83601b0f0b89e551",
                        "parent_id": None,
                        "name": "query",
                        "span_kind": "CHAIN",
                        "start_time": datetime.fromisoformat("2023-12-11T17:43:25.540371+00:00"),
                        "end_time": datetime.fromisoformat("2023-12-11T17:43:26.492242+00:00"),
                        "attributes": {
                            "openinference": {"span": {"kind": "CHAIN"}},
                            "output": {
                                "value": "Arize supports drift metrics such as Population Stability Index, KL Divergence, and Wasserstein Distance."
                            },
                            "input": {"value": "What drift metrics are supported in Arize?"},
                        },
                        "events": [],
                        "status_code": "OK",
                        "status_message": "",
                        "cumulative_error_count": 0,
                        "cumulative_llm_token_count_prompt": 315,
                        "cumulative_llm_token_count_completion": 21,
                    },
                    {
                        "trace_rowid": trace_rowids[1],
                        "span_id": "fc7f4cb067124f0abed01e5749a6aead",
                        "parent_id": "094ae70b0e9c4dec83601b0f0b89e551",
                        "name": "retrieve",
                        "span_kind": "RETRIEVER",
                        "start_time": datetime.fromisoformat("2023-12-11T17:43:25.540449+00:00"),
                        "end_time": datetime.fromisoformat("2023-12-11T17:43:25.842912+00:00"),
                        "attributes": {
                            "openinference": {"span": {"kind": "RETRIEVER"}},
                            "retrieval": {
                                "documents": [
                                    {
                                        "document": {
                                            "content": "\nDrift monitors measure distribution drift, which is the difference between two statistical distributions.&#x20;\n\nArize offers various distributional drift metrics to choose from when setting up a monitor. Each metric is tailored to a specific use case; refer to this guide to help choose the appropriate metric for various ML use cases.\n\n",
                                            "id": "60f3c900-dcee-43ef-816e-ae8f5289a544",
                                            "score": 0.8768844604492188,
                                        }
                                    },
                                    {
                                        "document": {
                                            "content": "\nArize calculates drift metrics such as Population Stability Index, KL Divergence, and Wasserstein Distance. Arize computes drift by measuring distribution changes between the model\u2019s production values and a baseline (reference dataset). Users can configure a baseline to be any time window of a:\n\n1. Pre-production dataset (training, test, validation) or\n2. Fixed or moving time period from production (e.g. last 30 days, last 60 days).&#x20;\n\nBaselines are saved in Arize so that users can compare several versions and/or environments against each other across moving or fixed time windows. For more details on baselines, visit here.\n\n",
                                            "id": "2e468875-ee22-4b5d-a7f4-57074eb5adfa",
                                            "score": 0.873500406742096,
                                        }
                                    },
                                ]
                            },
                            "input": {"value": "What drift metrics are supported in Arize?"},
                        },
                        "events": [],
                        "status_code": "OK",
                        "status_message": "",
                        "cumulative_error_count": 0,
                        "cumulative_llm_token_count_prompt": 0,
                        "cumulative_llm_token_count_completion": 0,
                    },
                    {
                        "trace_rowid": trace_rowids[1],
                        "span_id": "70252f342dcc496dac93404dcfbaa211",
                        "parent_id": "fc7f4cb067124f0abed01e5749a6aead",
                        "name": "embedding",
                        "span_kind": "EMBEDDING",
                        "start_time": datetime.fromisoformat("2023-12-11T17:43:25.540677+00:00"),
                        "end_time": datetime.fromisoformat("2023-12-11T17:43:25.768612+00:00"),
                        "attributes": {
                            "openinference": {"span": {"kind": "EMBEDDING"}},
                            "embedding": {
                                "model_name": "text-embedding-ada-002",
                                "embeddings": [
                                    {
                                        "embedding": {
                                            "vector": [1.0],
                                            "text": "What drift metrics are supported in Arize?",
                                        }
                                    }
                                ],
                            },
                        },
                        "events": [],
                        "status_code": "OK",
                        "status_message": "",
                        "cumulative_error_count": 0,
                        "cumulative_llm_token_count_prompt": 0,
                        "cumulative_llm_token_count_completion": 0,
                    },
                    {
                        "trace_rowid": trace_rowids[1],
                        "span_id": "c5ff03a4cf534b07a5ad6a00836acb1e",
                        "parent_id": "094ae70b0e9c4dec83601b0f0b89e551",
                        "name": "synthesize",
                        "span_kind": "CHAIN",
                        "start_time": datetime.fromisoformat("2023-12-11T17:43:25.842986+00:00"),
                        "end_time": datetime.fromisoformat("2023-12-11T17:43:26.492192+00:00"),
                        "attributes": {
                            "openinference": {"span": {"kind": "CHAIN"}},
                            "output": {
                                "value": "Arize supports drift metrics such as Population Stability Index, KL Divergence, and Wasserstein Distance."
                            },
                            "input": {"value": "What drift metrics are supported in Arize?"},
                        },
                        "events": [],
                        "status_code": "OK",
                        "status_message": "",
                        "cumulative_error_count": 0,
                        "cumulative_llm_token_count_prompt": 315,
                        "cumulative_llm_token_count_completion": 21,
                    },
                    {
                        "trace_rowid": trace_rowids[1],
                        "span_id": "0890b8716c4943c18b3ad45c6d9aaf5d",
                        "parent_id": "c5ff03a4cf534b07a5ad6a00836acb1e",
                        "name": "llm",
                        "span_kind": "LLM",
                        "start_time": datetime.fromisoformat("2023-12-11T17:43:25.844758+00:00"),
                        "end_time": datetime.fromisoformat("2023-12-11T17:43:26.491989+00:00"),
                        "attributes": {
                            "llm": {
                                "invocation_parameters": '{"model": "gpt-3.5-turbo", "temperature": 0.0, "max_tokens": None}',
                                "input_messages": [
                                    {
                                        "message": {
                                            "content": "You are an expert Q&A system that is trusted around the world.\nAlways answer the query using the provided context information, and not prior knowledge.\nSome rules to follow:\n1. Never directly reference the given context in your answer.\n2. Avoid statements like 'Based on the context, ...' or 'The context information ...' or anything along those lines.",
                                            "role": "system",
                                        }
                                    },
                                    {
                                        "message": {
                                            "content": "Context information is below.\n---------------------\nDrift monitors measure distribution drift, which is the difference between two statistical distributions.&#x20;\n\nArize offers various distributional drift metrics to choose from when setting up a monitor. Each metric is tailored to a specific use case; refer to this guide to help choose the appropriate metric for various ML use cases.\n\nArize calculates drift metrics such as Population Stability Index, KL Divergence, and Wasserstein Distance. Arize computes drift by measuring distribution changes between the model\u2019s production values and a baseline (reference dataset). Users can configure a baseline to be any time window of a:\n\n1. Pre-production dataset (training, test, validation) or\n2. Fixed or moving time period from production (e.g. last 30 days, last 60 days).&#x20;\n\nBaselines are saved in Arize so that users can compare several versions and/or environments against each other across moving or fixed time windows. For more details on baselines, visit here.\n---------------------\nGiven the context information and not prior knowledge, answer the query.\nQuery: What drift metrics are supported in Arize?\nAnswer: ",
                                            "role": "user",
                                        }
                                    },
                                ],
                                "model_name": "gpt-3.5-turbo",
                                "output_messages": [
                                    {
                                        "message": {
                                            "content": "Arize supports drift metrics such as Population Stability Index, KL Divergence, and Wasserstein Distance.",
                                            "role": "assistant",
                                        }
                                    }
                                ],
                                "prompt_template": {
                                    "template": "system: You are an expert Q&A system that is trusted around the world.\nAlways answer the query using the provided context information, and not prior knowledge.\nSome rules to follow:\n1. Never directly reference the given context in your answer.\n2. Avoid statements like 'Based on the context, ...' or 'The context information ...' or anything along those lines.\nuser: Context information is below.\n---------------------\n{context_str}\n---------------------\nGiven the context information and not prior knowledge, answer the query.\nQuery: {query_str}\nAnswer: \nassistant: ",
                                    "variables": {
                                        "context_str": "Drift monitors measure distribution drift, which is the difference between two statistical distributions.&#x20;\n\nArize offers various distributional drift metrics to choose from when setting up a monitor. Each metric is tailored to a specific use case; refer to this guide to help choose the appropriate metric for various ML use cases.\n\nArize calculates drift metrics such as Population Stability Index, KL Divergence, and Wasserstein Distance. Arize computes drift by measuring distribution changes between the model\u2019s production values and a baseline (reference dataset). Users can configure a baseline to be any time window of a:\n\n1. Pre-production dataset (training, test, validation) or\n2. Fixed or moving time period from production (e.g. last 30 days, last 60 days).&#x20;\n\nBaselines are saved in Arize so that users can compare several versions and/or environments against each other across moving or fixed time windows. For more details on baselines, visit here.",
                                        "query_str": "What drift metrics are supported in Arize?",
                                    },
                                },
                                "token_count": {
                                    "prompt": 315.0,
                                    "total": 336.0,
                                    "completion": 21.0,
                                },
                            },
                            "output": {
                                "value": "Arize supports drift metrics such as Population Stability Index, KL Divergence, and Wasserstein Distance."
                            },
                            "openinference": {"span": {"kind": "LLM"}},
                        },
                        "events": [],
                        "status_code": "OK",
                        "status_message": "",
                        "cumulative_error_count": 0,
                        "cumulative_llm_token_count_prompt": 315,
                        "cumulative_llm_token_count_completion": 21,
                    },
                    {
                        "trace_rowid": trace_rowids[2],
                        "span_id": "63b60ed12a61418ab9bd3757bd7eb09f",
                        "parent_id": None,
                        "name": "query",
                        "span_kind": "CHAIN",
                        "start_time": datetime.fromisoformat("2023-12-11T17:43:26.495969+00:00"),
                        "end_time": datetime.fromisoformat("2023-12-11T17:43:27.336284+00:00"),
                        "attributes": {
                            "openinference": {"span": {"kind": "CHAIN"}},
                            "output": {"value": "Yes, Arize supports batch models."},
                            "input": {"value": "Does Arize support batch models?"},
                        },
                        "events": [],
                        "status_code": "OK",
                        "status_message": "",
                        "cumulative_error_count": 0,
                        "cumulative_llm_token_count_prompt": 374,
                        "cumulative_llm_token_count_completion": 8,
                    },
                    {
                        "trace_rowid": trace_rowids[2],
                        "span_id": "2a3744dfdb954d6ea6a1dc0acb1e81d3",
                        "parent_id": "63b60ed12a61418ab9bd3757bd7eb09f",
                        "name": "retrieve",
                        "span_kind": "RETRIEVER",
                        "start_time": datetime.fromisoformat("2023-12-11T17:43:26.496043+00:00"),
                        "end_time": datetime.fromisoformat("2023-12-11T17:43:26.704469+00:00"),
                        "attributes": {
                            "openinference": {"span": {"kind": "RETRIEVER"}},
                            "retrieval": {
                                "documents": [
                                    {
                                        "document": {
                                            "content": "\nArize supports many model types - check out our various Model Types to learn more.&#x20;\n\n",
                                            "id": "b18c5cbd-ab0b-43d2-b0d3-55e755ee3561",
                                            "score": 0.8502153754234314,
                                        }
                                    },
                                    {
                                        "document": {
                                            "content": 'developers to create, train, and deploy machine-learning models in the cloud. Monitor and observe models deployed on SageMaker with Arize for data quality issues, performance checks, and drift.&#x20;\n\n{% content-ref url="spell.md" %}\nspell.md\n{% endcontent-ref %}\n\n> Spell is an end-to-end ML platform that provides infrastructure for company to deploy and train models. Visualize your model\'s performance, understand drift & data quality issues, and share insights learned from your models deployed on Spell.\n\n{% content-ref url="ubiops.md" %}\nubiops.md\n{% endcontent-ref %}\n\n> UbiOps is an MLOps platform with APIs to deploy and serve models. The Arize platform can easily integrate with UbiOps to enable model observability, explainability, and monitoring.\n\n{% content-ref url="weights-and-biases.md" %}\nweights-and-biases.md\n{% endcontent-ref %}\n\n> Weights and Biases helps you build better model by logging metrics and visualize your experiments before production. Arize helps you visualize your model performance, understand drift & data quality issues, and share insights learned from your models.\n\n\n\n',
                                            "id": "a975f095-94b3-4164-9483-dcf94864ee40",
                                            "score": 0.8405197262763977,
                                        }
                                    },
                                ]
                            },
                            "input": {"value": "Does Arize support batch models?"},
                        },
                        "events": [],
                        "status_code": "OK",
                        "status_message": "",
                        "cumulative_error_count": 0,
                        "cumulative_llm_token_count_prompt": 0,
                        "cumulative_llm_token_count_completion": 0,
                    },
                    {
                        "trace_rowid": trace_rowids[2],
                        "span_id": "11d1530f518c4f8cb8154d27a90c7023",
                        "parent_id": "2a3744dfdb954d6ea6a1dc0acb1e81d3",
                        "name": "embedding",
                        "span_kind": "EMBEDDING",
                        "start_time": datetime.fromisoformat("2023-12-11T17:43:26.496177+00:00"),
                        "end_time": datetime.fromisoformat("2023-12-11T17:43:26.644424+00:00"),
                        "attributes": {
                            "openinference": {"span": {"kind": "EMBEDDING"}},
                            "embedding": {
                                "model_name": "text-embedding-ada-002",
                                "embeddings": [
                                    {
                                        "embedding": {
                                            "vector": [1.0],
                                            "text": "Does Arize support batch models?",
                                        }
                                    }
                                ],
                            },
                        },
                        "events": [],
                        "status_code": "OK",
                        "status_message": "",
                        "cumulative_error_count": 0,
                        "cumulative_llm_token_count_prompt": 0,
                        "cumulative_llm_token_count_completion": 0,
                    },
                    {
                        "trace_rowid": trace_rowids[2],
                        "span_id": "5f763ad643f2458181062f5a815004e6",
                        "parent_id": "63b60ed12a61418ab9bd3757bd7eb09f",
                        "name": "synthesize",
                        "span_kind": "CHAIN",
                        "start_time": datetime.fromisoformat("2023-12-11T17:43:26.704532+00:00"),
                        "end_time": datetime.fromisoformat("2023-12-11T17:43:27.336235+00:00"),
                        "attributes": {
                            "openinference": {"span": {"kind": "CHAIN"}},
                            "output": {"value": "Yes, Arize supports batch models."},
                            "input": {"value": "Does Arize support batch models?"},
                        },
                        "events": [],
                        "status_code": "OK",
                        "status_message": "",
                        "cumulative_error_count": 0,
                        "cumulative_llm_token_count_prompt": 374,
                        "cumulative_llm_token_count_completion": 8,
                    },
                    {
                        "trace_rowid": trace_rowids[2],
                        "span_id": "d0c8e22b54f1499db8d2b006d4425508",
                        "parent_id": "5f763ad643f2458181062f5a815004e6",
                        "name": "llm",
                        "span_kind": "LLM",
                        "start_time": datetime.fromisoformat("2023-12-11T17:43:26.706204+00:00"),
                        "end_time": datetime.fromisoformat("2023-12-11T17:43:27.336029+00:00"),
                        "attributes": {
                            "llm": {
                                "invocation_parameters": '{"model": "gpt-3.5-turbo", "temperature": 0.0, "max_tokens": None}',
                                "input_messages": [
                                    {
                                        "message": {
                                            "content": "You are an expert Q&A system that is trusted around the world.\nAlways answer the query using the provided context information, and not prior knowledge.\nSome rules to follow:\n1. Never directly reference the given context in your answer.\n2. Avoid statements like 'Based on the context, ...' or 'The context information ...' or anything along those lines.",
                                            "role": "system",
                                        }
                                    },
                                    {
                                        "message": {
                                            "content": 'Context information is below.\n---------------------\nArize supports many model types - check out our various Model Types to learn more.&#x20;\n\ndevelopers to create, train, and deploy machine-learning models in the cloud. Monitor and observe models deployed on SageMaker with Arize for data quality issues, performance checks, and drift.&#x20;\n\n{% content-ref url="spell.md" %}\nspell.md\n{% endcontent-ref %}\n\n> Spell is an end-to-end ML platform that provides infrastructure for company to deploy and train models. Visualize your model\'s performance, understand drift & data quality issues, and share insights learned from your models deployed on Spell.\n\n{% content-ref url="ubiops.md" %}\nubiops.md\n{% endcontent-ref %}\n\n> UbiOps is an MLOps platform with APIs to deploy and serve models. The Arize platform can easily integrate with UbiOps to enable model observability, explainability, and monitoring.\n\n{% content-ref url="weights-and-biases.md" %}\nweights-and-biases.md\n{% endcontent-ref %}\n\n> Weights and Biases helps you build better model by logging metrics and visualize your experiments before production. Arize helps you visualize your model performance, understand drift & data quality issues, and share insights learned from your models.\n---------------------\nGiven the context information and not prior knowledge, answer the query.\nQuery: Does Arize support batch models?\nAnswer: ',
                                            "role": "user",
                                        }
                                    },
                                ],
                                "model_name": "gpt-3.5-turbo",
                                "output_messages": [
                                    {
                                        "message": {
                                            "content": "Yes, Arize supports batch models.",
                                            "role": "assistant",
                                        }
                                    }
                                ],
                                "prompt_template": {
                                    "template": "system: You are an expert Q&A system that is trusted around the world.\nAlways answer the query using the provided context information, and not prior knowledge.\nSome rules to follow:\n1. Never directly reference the given context in your answer.\n2. Avoid statements like 'Based on the context, ...' or 'The context information ...' or anything along those lines.\nuser: Context information is below.\n---------------------\n{context_str}\n---------------------\nGiven the context information and not prior knowledge, answer the query.\nQuery: {query_str}\nAnswer: \nassistant: ",
                                    "variables": {
                                        "context_str": 'Arize supports many model types - check out our various Model Types to learn more.&#x20;\n\ndevelopers to create, train, and deploy machine-learning models in the cloud. Monitor and observe models deployed on SageMaker with Arize for data quality issues, performance checks, and drift.&#x20;\n\n{% content-ref url="spell.md" %}\nspell.md\n{% endcontent-ref %}\n\n> Spell is an end-to-end ML platform that provides infrastructure for company to deploy and train models. Visualize your model\'s performance, understand drift & data quality issues, and share insights learned from your models deployed on Spell.\n\n{% content-ref url="ubiops.md" %}\nubiops.md\n{% endcontent-ref %}\n\n> UbiOps is an MLOps platform with APIs to deploy and serve models. The Arize platform can easily integrate with UbiOps to enable model observability, explainability, and monitoring.\n\n{% content-ref url="weights-and-biases.md" %}\nweights-and-biases.md\n{% endcontent-ref %}\n\n> Weights and Biases helps you build better model by logging metrics and visualize your experiments before production. Arize helps you visualize your model performance, understand drift & data quality issues, and share insights learned from your models.',
                                        "query_str": "Does Arize support batch models?",
                                    },
                                },
                                "token_count": {"prompt": 374.0, "total": 382.0, "completion": 8.0},
                            },
                            "output": {"value": "Yes, Arize supports batch models."},
                            "openinference": {"span": {"kind": "LLM"}},
                        },
                        "events": [],
                        "status_code": "OK",
                        "status_message": "",
                        "cumulative_error_count": 0,
                        "cumulative_llm_token_count_prompt": 374,
                        "cumulative_llm_token_count_completion": 8,
                    },
                ],
            )
        ).all()
        await session.execute(
            insert(models.SpanAnnotation),
            [
                {
                    "span_rowid": span_rowids[0],
                    "name": "Hallucination",
                    "label": "hallucinated",
                    "score": 0,
                    "explanation": "The query asks about how to use the SDK to upload a ranking model. The reference text provides information about ranking models and their challenges, and mentions a specific model 'arize-demo-hotel-ranking'. However, it does not provide any information about how to use an SDK to upload a ranking model. The answer talks about following the documentation provided by the SDK to upload the model, which is not mentioned or suggested in the reference text. Therefore, the answer is not based on the reference text.",
                    "metadata_": {},
                    "annotator_kind": "LLM",
                    "created_at": datetime.fromisoformat("2024-05-20T01:42:11+00:00"),
                    "updated_at": datetime.fromisoformat("2024-05-20T01:42:11+00:00"),
                    "identifier": "",
                    "source": "APP",
                    "user_id": None,
                },
                {
                    "span_rowid": span_rowids[5],
                    "name": "Hallucination",
                    "label": "factual",
                    "score": 1,
                    "explanation": "The query asks about the drift metrics supported in Arize. The reference text mentions that Arize calculates drift metrics such as Population Stability Index, KL Divergence, and Wasserstein Distance. The answer states the same information, that Arize supports drift metrics such as Population Stability Index, KL Divergence, and Wasserstein Distance. Therefore, the answer is based on the information provided in the reference text.",
                    "metadata_": {},
                    "annotator_kind": "LLM",
                    "created_at": datetime.fromisoformat("2024-05-20T01:42:11+00:00"),
                    "updated_at": datetime.fromisoformat("2024-05-20T01:42:11+00:00"),
                    "identifier": "",
                    "source": "APP",
                    "user_id": None,
                },
                {
                    "span_rowid": span_rowids[10],
                    "name": "Hallucination",
                    "label": "hallucinated",
                    "score": 0,
                    "explanation": "The query asks if Arize supports batch models. The reference text mentions that Arize supports many model types, but it does not specify if batch models are among those supported. Therefore, the answer assumes information that is not available in the reference text.",
                    "metadata_": {},
                    "annotator_kind": "LLM",
                    "created_at": datetime.fromisoformat("2024-05-20T01:42:11+00:00"),
                    "updated_at": datetime.fromisoformat("2024-05-20T01:42:11+00:00"),
                    "identifier": "",
                    "source": "APP",
                    "user_id": None,
                },
                {
                    "span_rowid": span_rowids[0],
                    "name": "Q&A Correctness",
                    "label": "incorrect",
                    "score": 0,
                    "explanation": "The reference text does not provide any information on how to use the SDK to upload a ranking model. It only mentions the use of a specific model 'arize-demo-hotel-ranking' and some challenges associated with ranking models. The answer, on the other hand, talks about following the documentation provided by the SDK to upload a ranking model. However, since the reference text does not mention anything about an SDK or its documentation, the answer does not correctly answer the question based on the reference text.",
                    "metadata_": {},
                    "annotator_kind": "LLM",
                    "created_at": datetime.fromisoformat("2024-05-20T01:42:11+00:00"),
                    "updated_at": datetime.fromisoformat("2024-05-20T01:42:11+00:00"),
                    "identifier": "",
                    "source": "APP",
                    "user_id": None,
                },
                {
                    "span_rowid": span_rowids[5],
                    "name": "Q&A Correctness",
                    "label": "correct",
                    "score": 1,
                    "explanation": "The reference text clearly states that Arize calculates drift metrics such as Population Stability Index, KL Divergence, and Wasserstein Distance. This directly matches the given answer, which states that Arize supports these same drift metrics. Therefore, the answer is correct.",
                    "metadata_": {},
                    "annotator_kind": "LLM",
                    "created_at": datetime.fromisoformat("2024-05-20T01:42:11+00:00"),
                    "updated_at": datetime.fromisoformat("2024-05-20T01:42:11+00:00"),
                    "identifier": "",
                    "source": "APP",
                    "user_id": None,
                },
                {
                    "span_rowid": span_rowids[10],
                    "name": "Q&A Correctness",
                    "label": "incorrect",
                    "score": 0,
                    "explanation": "The reference text mentions that Arize supports many model types and provides infrastructure for developers to create, train, and deploy machine-learning models in the cloud. However, it does not specifically mention that Arize supports batch models. Therefore, the answer is not supported by the reference text.",
                    "metadata_": {},
                    "annotator_kind": "LLM",
                    "created_at": datetime.fromisoformat("2024-05-20T01:42:11+00:00"),
                    "updated_at": datetime.fromisoformat("2024-05-20T01:42:11+00:00"),
                    "identifier": "",
                    "source": "APP",
                    "user_id": None,
                },
            ],
        )


@dataclass
class _Data:
    spans: list[models.Span] = field(default_factory=list)
    traces: list[models.Trace] = field(default_factory=list)
    project_sessions: list[models.ProjectSession] = field(default_factory=list)
    projects: list[models.Project] = field(default_factory=list)


# The `_data` fixture's sessions 1 and 3 are the two whose root-span input/output carry this
# fragment; the sort tests use it to exercise sorting and cursor paging under a filter.
_IO_CONDITION_ARG = "\"'\\\"\\\\'f' in any_input or '\\\"\\\\'f' in any_output\""


class TestProject:
    @staticmethod
    async def _node(
        field: str,
        project: models.Project,
        httpx_client: httpx.AsyncClient,
    ) -> Any:
        return await _node(field, Project.__name__, project.id, httpx_client)

    async def test_annotation_name_counts(
        self,
        db: DbSessionFactory,
        httpx_client: httpx.AsyncClient,
    ) -> None:
        """Span/trace/session annotation name counts reflect the number of annotations per name."""
        async with db() as session:
            project = await _add_project(session, name="annotation-name-counts-test")
            project_session = await _add_project_session(session, project)
            trace = await _add_trace(session, project, project_session)
            span = await _add_span(session, trace)

            # 2 span annotations named "correctness", 1 named "relevance"
            for name in ("correctness", "correctness", "relevance"):
                session.add(
                    models.SpanAnnotation(
                        span_rowid=span.id,
                        name=name,
                        annotator_kind="HUMAN",
                        source="APP",
                        metadata_={},
                        identifier=token_hex(4),
                    )
                )
            # 3 trace annotations named "quality"
            for _ in range(3):
                session.add(
                    models.TraceAnnotation(
                        trace_rowid=trace.id,
                        name="quality",
                        annotator_kind="HUMAN",
                        source="APP",
                        metadata_={},
                        identifier=token_hex(4),
                    )
                )
            # 1 session annotation named "helpfulness"
            session.add(
                models.ProjectSessionAnnotation(
                    project_session_id=project_session.id,
                    name="helpfulness",
                    annotator_kind="HUMAN",
                    source="APP",
                    metadata_={},
                    identifier=token_hex(4),
                )
            )

        span_counts = await self._node(
            "spanAnnotationNameCounts{name count}", project, httpx_client
        )
        assert span_counts == [
            {"name": "correctness", "count": 2},
            {"name": "relevance", "count": 1},
        ]

        trace_counts = await self._node(
            "traceAnnotationNameCounts{name count}", project, httpx_client
        )
        assert trace_counts == [{"name": "quality", "count": 3}]

        session_counts = await self._node(
            "sessionAnnotationNameCounts{name count}", project, httpx_client
        )
        assert session_counts == [{"name": "helpfulness", "count": 1}]

    async def test_annotation_name_counts_empty(
        self,
        db: DbSessionFactory,
        httpx_client: httpx.AsyncClient,
    ) -> None:
        """Projects without annotations return empty lists for each level."""
        async with db() as session:
            project = await _add_project(session, name="annotation-name-counts-empty")
        assert await self._node("spanAnnotationNameCounts{name count}", project, httpx_client) == []
        assert (
            await self._node("traceAnnotationNameCounts{name count}", project, httpx_client) == []
        )
        assert (
            await self._node("sessionAnnotationNameCounts{name count}", project, httpx_client) == []
        )

    @pytest.fixture
    async def _data(
        self,
        db: DbSessionFactory,
    ) -> _Data:
        projects, project_sessions, traces, spans = [], [], [], []
        async with db() as session:
            projects.append(await _add_project(session))

            project_sessions.append(await _add_project_session(session, projects[-1]))
            traces.append(await _add_trace(session, projects[-1], project_sessions[-1]))
            attributes = {"input": {"value": "a\"'b"}, "output": {"value": "c\"'d"}}
            spans.append(await _add_span(session, traces[-1], attributes=attributes))

            project_sessions.append(await _add_project_session(session, projects[-1]))
            traces.append(await _add_trace(session, projects[-1], project_sessions[-1]))
            attributes = {"input": {"value": "e\"'f"}, "output": {"value": "g\"'h"}}
            spans.append(
                await _add_span(
                    session,
                    traces[-1],
                    attributes=attributes,
                    cumulative_llm_token_count_prompt=1,
                    llm_token_count_prompt=1,
                )
            )
            attributes = {"input": {"value": "i\"'j"}, "output": {"value": "k\"'l"}}
            spans.append(await _add_span(session, parent_span=spans[-1], attributes=attributes))

            project_sessions.append(await _add_project_session(session, projects[-1]))
            traces.append(await _add_trace(session, projects[-1], project_sessions[-1]))
            spans.append(
                await _add_span(
                    session,
                    traces[-1],
                    cumulative_llm_token_count_completion=2,
                    llm_token_count_completion=2,
                )
            )
            traces.append(await _add_trace(session, projects[-1], project_sessions[-1]))
            spans.append(await _add_span(session, traces[-1]))

            project_sessions.append(await _add_project_session(session, projects[-1]))
            traces.append(await _add_trace(session, projects[-1], project_sessions[-1]))
            attributes = {"input": {"value": "g\"'h"}, "output": {"value": "e\"'f"}}
            spans.append(
                await _add_span(
                    session,
                    traces[-1],
                    attributes=attributes,
                    cumulative_llm_token_count_completion=1,
                    llm_token_count_completion=1,
                )
            )
            spans.append(await _add_span(session, traces[-1]))

        return _Data(
            spans=spans,
            traces=traces,
            project_sessions=project_sessions,
            projects=projects,
        )

    @pytest.fixture
    async def _orphan_spans(
        self,
        db: DbSessionFactory,
    ) -> _Data:
        """Creates spans with missing parent references to test orphan span handling.

        This fixture creates a test dataset with spans that have various parent-child relationships:
        - Some spans have no parent (parent_id=None) - these are true root spans
        - Some spans have parent_ids that don't exist - these are orphan spans
        - Some spans have valid parent references

        The fixture is used to test how the API handles orphan spans in different configurations.
        """
        projects, traces = [], []
        spans: list[models.Span] = []
        async with db() as session:
            # Create a test project
            projects.append(models.Project(name=token_hex(8)))
            session.add(projects[-1])
            await session.flush()

            start_time = datetime.fromisoformat("2024-01-01T00:00:00+00:00")
            for i in range(2):
                # Create two traces with different start times
                trace_start_time = start_time + timedelta(seconds=i)
                traces.append(
                    models.Trace(
                        trace_id=token_hex(16),
                        project_rowid=projects[-1].id,
                        start_time=trace_start_time,
                        end_time=trace_start_time + timedelta(seconds=1),
                    )
                )
                session.add(traces[-1])
                await session.flush()

                n = 7
                for j in range(n):
                    # Create spans with different parent relationships:
                    if j % 2:
                        # Odd-indexed spans have non-existent parent IDs (orphan spans)
                        parent_id = token_hex(8)
                    elif j and j == n - 1:
                        # The last span (when j is even and non-zero) references the previous span
                        parent_id = spans[-1].span_id
                    else:
                        # Even-indexed spans (except the last) have no parent (root spans)
                        parent_id = None

                    span_start_time = trace_start_time + timedelta(seconds=j)
                    spans.append(
                        models.Span(
                            trace_rowid=traces[-1].id,
                            span_id=token_hex(8),
                            parent_id=parent_id,
                            name=token_hex(8),
                            span_kind="CHAIN",
                            start_time=span_start_time,
                            end_time=span_start_time + timedelta(seconds=1),
                            # Input value contains sequential digits (0, 01, 012, etc.)
                            # This is used for filtering in tests
                            attributes={"input": {"value": "".join(map(str, range(j)))}},
                            events=[],
                            status_code="OK",
                            status_message="",
                            cumulative_error_count=0,
                            cumulative_llm_token_count_prompt=0,
                            cumulative_llm_token_count_completion=0,
                        )
                    )
            session.add_all(spans)
            await session.flush()

        return _Data(
            spans=spans,
            traces=traces,
            projects=projects,
        )

    async def test_sessions_sort_token_count_total(
        self,
        _data: _Data,
        httpx_client: httpx.AsyncClient,
    ) -> None:
        column = "tokenCountTotal"
        project = _data.projects[0]
        result: list[str] = [
            _gid(_data.project_sessions[2]),
            _gid(_data.project_sessions[3]),
            _gid(_data.project_sessions[1]),
            _gid(_data.project_sessions[0]),
        ]

        for direction, expected in {"desc": result, "asc": result[::-1]}.items():
            field = (
                "sessions(first:50,sort:{col:"
                + column
                + ",dir:"
                + direction
                + "}){edges{node{id}}}"
            )
            res = await self._node(field, project, httpx_client)
            assert [e["node"]["id"] for e in res["edges"]] == expected

        # Test pagination
        first = 2
        cursors = [b"3:INT:2", b"4:INT:1", b"2:INT:1", b"1:INT:0"]
        for direction, (afters, ids) in {
            "desc": ([b""] + cursors, result),
            "asc": ([b""] + cursors[::-1], result[::-1]),
        }.items():
            for i, after in enumerate(afters):
                expected = ids[i : i + first]
                field = (
                    "sessions(sort:{col:"
                    + column
                    + ",dir:"
                    + direction
                    + "},first:"
                    + str(first)
                    + ',after:"'
                    + base64.b64encode(after).decode()
                    + '"){edges{node{id}}}'
                )
                res = await self._node(field, project, httpx_client)
                assert [e["node"]["id"] for e in res["edges"]] == expected

    async def test_sessions_sort_token_count_total_plus_session_filter_condition(
        self,
        _data: _Data,
        httpx_client: httpx.AsyncClient,
    ) -> None:
        column = "tokenCountTotal"
        project = _data.projects[0]
        result: list[str] = [
            _gid(_data.project_sessions[3]),
            _gid(_data.project_sessions[1]),
        ]

        for direction, expected in {"desc": result, "asc": result[::-1]}.items():
            field = (
                "sessions(first:50,sort:{col:"
                + column
                + ",dir:"
                + direction
                + "},sessionFilterCondition:"
                + _IO_CONDITION_ARG
                + "){edges{node{id}}}"
            )
            res = await self._node(field, project, httpx_client)
            assert [e["node"]["id"] for e in res["edges"]] == expected

        # Test pagination
        first = 2
        cursors = [b"4:INT:1", b"2:INT:1"]
        for direction, (afters, ids) in {
            "desc": ([b""] + cursors, result),
            "asc": ([b""] + cursors[::-1], result[::-1]),
        }.items():
            for i, after in enumerate(afters):
                expected = ids[i : i + first]
                field = (
                    "sessions(sort:{col:"
                    + column
                    + ",dir:"
                    + direction
                    + "},sessionFilterCondition:"
                    + _IO_CONDITION_ARG
                    + ",first:"
                    + str(first)
                    + ',after:"'
                    + base64.b64encode(after).decode()
                    + '"){edges{node{id}}}'
                )
                res = await self._node(field, project, httpx_client)
                assert [e["node"]["id"] for e in res["edges"]] == expected

    async def test_sessions_sort_num_traces(
        self,
        _data: _Data,
        httpx_client: httpx.AsyncClient,
    ) -> None:
        column = "numTraces"
        project = _data.projects[0]
        result: list[str] = [
            _gid(_data.project_sessions[2]),
            _gid(_data.project_sessions[3]),
            _gid(_data.project_sessions[1]),
            _gid(_data.project_sessions[0]),
        ]

        for direction, expected in {"desc": result, "asc": result[::-1]}.items():
            field = (
                "sessions(first:50,sort:{col:"
                + column
                + ",dir:"
                + direction
                + "}){edges{node{id}}}"
            )
            res = await self._node(field, project, httpx_client)
            assert [e["node"]["id"] for e in res["edges"]] == expected

        # Test pagination
        first = 2
        cursors = [b"3:INT:2", b"4:INT:1", b"2:INT:1", b"1:INT:1"]
        for direction, (afters, ids) in {
            "desc": ([b""] + cursors, result),
            "asc": ([b""] + cursors[::-1], result[::-1]),
        }.items():
            for i, after in enumerate(afters):
                expected = ids[i : i + first]
                field = (
                    "sessions(sort:{col:"
                    + column
                    + ",dir:"
                    + direction
                    + "},first:"
                    + str(first)
                    + ',after:"'
                    + base64.b64encode(after).decode()
                    + '"){edges{node{id}}}'
                )
                res = await self._node(field, project, httpx_client)
                assert [e["node"]["id"] for e in res["edges"]] == expected

    async def test_sessions_sort_num_traces_plus_session_filter_condition(
        self,
        _data: _Data,
        httpx_client: httpx.AsyncClient,
    ) -> None:
        column = "numTraces"
        project = _data.projects[0]
        result: list[str] = [
            _gid(_data.project_sessions[3]),
            _gid(_data.project_sessions[1]),
        ]

        for direction, expected in {"desc": result, "asc": result[::-1]}.items():
            field = (
                "sessions(first:50,sort:{col:"
                + column
                + ",dir:"
                + direction
                + "},sessionFilterCondition:"
                + _IO_CONDITION_ARG
                + "){edges{node{id}}}"
            )
            res = await self._node(field, project, httpx_client)
            assert [e["node"]["id"] for e in res["edges"]] == expected

        # Test pagination
        first = 2
        cursors = [b"4:INT:1", b"2:INT:1"]
        for direction, (afters, ids) in {
            "desc": ([b""] + cursors, result),
            "asc": ([b""] + cursors[::-1], result[::-1]),
        }.items():
            for i, after in enumerate(afters):
                expected = ids[i : i + first]
                field = (
                    "sessions(sort:{col:"
                    + column
                    + ",dir:"
                    + direction
                    + "},sessionFilterCondition:"
                    + _IO_CONDITION_ARG
                    + ",first:"
                    + str(first)
                    + ',after:"'
                    + base64.b64encode(after).decode()
                    + '"){edges{node{id}}}'
                )
                res = await self._node(field, project, httpx_client)
                assert [e["node"]["id"] for e in res["edges"]] == expected

    async def test_sessions_sort_start_time(
        self,
        db: DbSessionFactory,
        httpx_client: httpx.AsyncClient,
    ) -> None:
        """Test sorting project sessions by start time."""
        async with db() as session:
            project = await _add_project(session, name="start-time-sort-test")

            # Create 5 sessions with different start times
            sessions = []
            base_time = datetime.fromisoformat("2024-01-01T00:00:00+00:00")
            time_deltas = [
                timedelta(hours=3),
                timedelta(hours=1),
                timedelta(hours=5),
                timedelta(hours=2),
                timedelta(hours=4),
            ]

            for i, delta in enumerate(time_deltas):
                start_time = base_time + delta
                end_time = start_time + timedelta(minutes=30)
                ps = await _add_project_session(session, project, start_time=start_time)
                sessions.append(ps)
                trace = await _add_trace(
                    session, project, ps, start_time=start_time, end_time=end_time
                )
                await _add_span(session, trace, start_time=start_time, end_time=end_time)

        column = "startTime"
        # Expected order desc: sessions[2] (5h), sessions[4] (4h), sessions[0] (3h), sessions[3] (2h), sessions[1] (1h)
        result_desc = [
            _gid(sessions[2]),  # 5h
            _gid(sessions[4]),  # 4h
            _gid(sessions[0]),  # 3h
            _gid(sessions[3]),  # 2h
            _gid(sessions[1]),  # 1h
        ]

        # Test descending order
        field = f"sessions(first:50,sort:{{col:{column},dir:desc}}){{edges{{node{{id}}}}}}"
        res = await self._node(field, project, httpx_client)
        assert [e["node"]["id"] for e in res["edges"]] == result_desc

        # Test ascending order
        field = f"sessions(first:50,sort:{{col:{column},dir:asc}}){{edges{{node{{id}}}}}}"
        res = await self._node(field, project, httpx_client)
        assert [e["node"]["id"] for e in res["edges"]] == result_desc[::-1]

        # Test pagination
        first = 2
        field = f"sessions(sort:{{col:{column},dir:desc}},first:{first}){{edges{{node{{id}}}}pageInfo{{hasNextPage}}}}"
        res = await self._node(field, project, httpx_client)
        assert [e["node"]["id"] for e in res["edges"]] == result_desc[:2]
        assert res["pageInfo"]["hasNextPage"] is True

    async def test_sessions_sort_cost_total(
        self,
        db: DbSessionFactory,
        httpx_client: httpx.AsyncClient,
    ) -> None:
        """Test sorting project sessions by total cost."""
        async with db() as session:
            project = await _add_project(session, name="cost-sort-test")
            model = await _add_generative_model(session, name="gpt-4", provider="openai")

            # Create 5 sessions: 4 with costs, 1 without
            sessions = []
            costs = [100.0, 50.0, 200.0, 75.0, None]

            for i, cost in enumerate(costs):
                ps = await _add_project_session(session, project)
                sessions.append(ps)
                trace = await _add_trace(session, project, ps)
                span = await _add_span(session, trace)

                # Only add span cost if cost is not None
                if cost is not None:
                    await _add_span_cost(
                        session,
                        span=span,
                        trace=trace,
                        model=model,
                        total_cost=cost,
                        total_tokens=1000,
                        prompt_cost=cost * 0.75,
                        prompt_tokens=800,
                        completion_cost=cost * 0.25,
                        completion_tokens=200,
                        span_start_time=span.start_time,
                    )

        column = "costTotal"
        # Expected order desc: 200, 100, 75, 50, 0
        result_desc = [
            _gid(sessions[2]),  # 200.0
            _gid(sessions[0]),  # 100.0
            _gid(sessions[3]),  # 75.0
            _gid(sessions[1]),  # 50.0
            _gid(sessions[4]),  # 0.0 (no cost data)
        ]

        # Test descending order
        field = f"sessions(first:50,sort:{{col:{column},dir:desc}}){{edges{{node{{id}}}}}}"
        res = await self._node(field, project, httpx_client)
        assert [e["node"]["id"] for e in res["edges"]] == result_desc

        # Test ascending order: 0, 50, 75, 100, 200
        field = f"sessions(first:50,sort:{{col:{column},dir:asc}}){{edges{{node{{id}}}}}}"
        res = await self._node(field, project, httpx_client)
        assert [e["node"]["id"] for e in res["edges"]] == result_desc[::-1]

        # Test pagination
        first = 2
        field = f"sessions(sort:{{col:{column},dir:desc}},first:{first}){{edges{{node{{id}}}}pageInfo{{hasNextPage}}}}"
        res = await self._node(field, project, httpx_client)
        assert [e["node"]["id"] for e in res["edges"]] == result_desc[:2]
        assert res["pageInfo"]["hasNextPage"] is True

    async def test_sessions_sort_by_annotation_score(
        self,
        db: DbSessionFactory,
        httpx_client: httpx.AsyncClient,
    ) -> None:
        """Test sorting project sessions by annotation score with pagination and multiple NULLs."""
        async with db() as session:
            project = await _add_project(session, name="annotation-sort-test")

            # Create 8 sessions: 5 with scores, 3 with NULL scores (interspersed)
            sessions = []
            # Scores: 0.9, NULL, 0.5, 0.7, NULL, 0.3, 0.8, NULL
            scores = [0.9, None, 0.5, 0.7, None, 0.3, 0.8, None]

            for i, score in enumerate(scores):
                ps = await _add_project_session(session, project)
                sessions.append(ps)
                trace = await _add_trace(session, project, ps)
                await _add_span(session, trace)

                annotation = models.ProjectSessionAnnotation(
                    project_session_id=ps.id,
                    name="Quality",
                    label="good" if score is None else None,
                    score=score,
                    explanation=f"Quality annotation {i}",
                    metadata_={},
                    annotator_kind="LLM",
                    source="APP",
                )
                session.add(annotation)

        # Test descending order: 0.9, 0.8, 0.7, 0.5, 0.3, NULL (desc by ID: 8,5,2)
        field = 'sessions(first:50,sort:{annoResultKey:{name:"Quality",attr:score},dir:desc}){edges{node{id}}}'
        res = await self._node(field, project, httpx_client)
        assert [e["node"]["id"] for e in res["edges"]] == [
            _gid(sessions[0]),  # 0.9
            _gid(sessions[6]),  # 0.8
            _gid(sessions[3]),  # 0.7
            _gid(sessions[2]),  # 0.5
            _gid(sessions[5]),  # 0.3
            _gid(sessions[7]),  # NULL (ID 8)
            _gid(sessions[4]),  # NULL (ID 5)
            _gid(sessions[1]),  # NULL (ID 2)
        ]

        # Test ascending order: 0.3, 0.5, 0.7, 0.8, 0.9, NULL (asc by ID: 2,5,8)
        field = 'sessions(first:50,sort:{annoResultKey:{name:"Quality",attr:score},dir:asc}){edges{node{id}}}'
        res = await self._node(field, project, httpx_client)
        assert [e["node"]["id"] for e in res["edges"]] == [
            _gid(sessions[5]),  # 0.3
            _gid(sessions[2]),  # 0.5
            _gid(sessions[3]),  # 0.7
            _gid(sessions[6]),  # 0.8
            _gid(sessions[0]),  # 0.9
            _gid(sessions[1]),  # NULL (ID 2)
            _gid(sessions[4]),  # NULL (ID 5)
            _gid(sessions[7]),  # NULL (ID 8)
        ]

        # Test pagination with desc order - get first 3, ensuring we paginate through NULLs
        first = 3
        field = f'sessions(sort:{{annoResultKey:{{name:"Quality",attr:score}},dir:desc}},first:{first}){{edges{{node{{id}}}}pageInfo{{hasNextPage}}}}'
        res = await self._node(field, project, httpx_client)
        assert [e["node"]["id"] for e in res["edges"]] == [
            _gid(sessions[0]),  # 0.9
            _gid(sessions[6]),  # 0.8
            _gid(sessions[3]),  # 0.7
        ]
        assert res["pageInfo"]["hasNextPage"] is True

        # Test pagination into NULL territory - page that includes first NULL
        first = 6
        field = f'sessions(sort:{{annoResultKey:{{name:"Quality",attr:score}},dir:asc}},first:{first}){{edges{{node{{id}}cursor}}pageInfo{{endCursor hasNextPage}}}}'
        res = await self._node(field, project, httpx_client)
        assert [e["node"]["id"] for e in res["edges"]] == [
            _gid(sessions[5]),  # 0.3
            _gid(sessions[2]),  # 0.5
            _gid(sessions[3]),  # 0.7
            _gid(sessions[6]),  # 0.8
            _gid(sessions[0]),  # 0.9
            _gid(sessions[1]),  # NULL (ID 2 - first NULL)
        ]
        assert res["pageInfo"]["hasNextPage"] is True

        # Fetch next page AFTER first NULL - should get remaining NULLs
        end_cursor = res["pageInfo"]["endCursor"]
        field = f'sessions(sort:{{annoResultKey:{{name:"Quality",attr:score}},dir:asc}},first:3,after:"{end_cursor}"){{edges{{node{{id}}}}pageInfo{{hasNextPage}}}}'
        res = await self._node(field, project, httpx_client)
        assert [e["node"]["id"] for e in res["edges"]] == [
            _gid(sessions[4]),  # NULL (ID 5)
            _gid(sessions[7]),  # NULL (ID 8)
        ]
        assert res["pageInfo"]["hasNextPage"] is False

    async def test_sessions_sort_by_annotation_label(
        self,
        db: DbSessionFactory,
        httpx_client: httpx.AsyncClient,
    ) -> None:
        """Test sorting project sessions by annotation label with pagination and multiple NULLs."""
        async with db() as session:
            project = await _add_project(session, name="label-sort-test")

            # Create 8 sessions: 5 with labels, 3 with NULL labels (interspersed)
            sessions = []
            # Labels: zebra, NULL, alpha, delta, NULL, beta, charlie, NULL
            labels = ["zebra", None, "alpha", "delta", None, "beta", "charlie", None]

            for i, label in enumerate(labels):
                ps = await _add_project_session(session, project)
                sessions.append(ps)
                trace = await _add_trace(session, project, ps)
                await _add_span(session, trace)

                annotation = models.ProjectSessionAnnotation(
                    project_session_id=ps.id,
                    name="Quality",
                    label=label,
                    score=0.5 if label is None else None,
                    explanation=f"Quality annotation {i}",
                    metadata_={},
                    annotator_kind="LLM",
                    source="APP",
                )
                session.add(annotation)

        # Test descending order: zebra, delta, charlie, beta, alpha, NULL (desc by ID: 8,5,2)
        field = 'sessions(first:50,sort:{annoResultKey:{name:"Quality",attr:label},dir:desc}){edges{node{id}}}'
        res = await self._node(field, project, httpx_client)
        assert [e["node"]["id"] for e in res["edges"]] == [
            _gid(sessions[0]),  # "zebra"
            _gid(sessions[3]),  # "delta"
            _gid(sessions[6]),  # "charlie"
            _gid(sessions[5]),  # "beta"
            _gid(sessions[2]),  # "alpha"
            _gid(sessions[7]),  # NULL (ID 8)
            _gid(sessions[4]),  # NULL (ID 5)
            _gid(sessions[1]),  # NULL (ID 2)
        ]

        # Test ascending order: alpha, beta, charlie, delta, zebra, NULL (asc by ID: 2,5,8)
        field = 'sessions(first:50,sort:{annoResultKey:{name:"Quality",attr:label},dir:asc}){edges{node{id}}}'
        res = await self._node(field, project, httpx_client)
        assert [e["node"]["id"] for e in res["edges"]] == [
            _gid(sessions[2]),  # "alpha"
            _gid(sessions[5]),  # "beta"
            _gid(sessions[6]),  # "charlie"
            _gid(sessions[3]),  # "delta"
            _gid(sessions[0]),  # "zebra"
            _gid(sessions[1]),  # NULL (ID 2)
            _gid(sessions[4]),  # NULL (ID 5)
            _gid(sessions[7]),  # NULL (ID 8)
        ]

        # Test pagination with asc order - get first 3, ensuring we paginate through NULLs
        first = 3
        field = f'sessions(sort:{{annoResultKey:{{name:"Quality",attr:label}},dir:asc}},first:{first}){{edges{{node{{id}}}}pageInfo{{hasNextPage}}}}'
        res = await self._node(field, project, httpx_client)
        assert [e["node"]["id"] for e in res["edges"]] == [
            _gid(sessions[2]),  # "alpha"
            _gid(sessions[5]),  # "beta"
            _gid(sessions[6]),  # "charlie"
        ]
        assert res["pageInfo"]["hasNextPage"] is True

        # Test pagination into NULL territory - page that includes first NULL
        first = 6
        field = f'sessions(sort:{{annoResultKey:{{name:"Quality",attr:label}},dir:desc}},first:{first}){{edges{{node{{id}}cursor}}pageInfo{{endCursor hasNextPage}}}}'
        res = await self._node(field, project, httpx_client)
        assert [e["node"]["id"] for e in res["edges"]] == [
            _gid(sessions[0]),  # "zebra"
            _gid(sessions[3]),  # "delta"
            _gid(sessions[6]),  # "charlie"
            _gid(sessions[5]),  # "beta"
            _gid(sessions[2]),  # "alpha"
            _gid(sessions[7]),  # NULL (ID 8 - first NULL in desc order)
        ]
        assert res["pageInfo"]["hasNextPage"] is True

        # Fetch next page AFTER first NULL - should get remaining NULLs
        end_cursor = res["pageInfo"]["endCursor"]
        field = f'sessions(sort:{{annoResultKey:{{name:"Quality",attr:label}},dir:desc}},first:3,after:"{end_cursor}"){{edges{{node{{id}}}}pageInfo{{hasNextPage}}}}'
        res = await self._node(field, project, httpx_client)
        assert [e["node"]["id"] for e in res["edges"]] == [
            _gid(sessions[4]),  # NULL (ID 5)
            _gid(sessions[1]),  # NULL (ID 2)
        ]
        assert res["pageInfo"]["hasNextPage"] is False

    @pytest.fixture
    async def _case_insensitive_data(
        self,
        db: DbSessionFactory,
    ) -> _Data:
        """Create test data for case-insensitive filtering"""
        projects, project_sessions, traces, spans = [], [], [], []
        async with db() as session:
            projects.append(await _add_project(session))

            # Session 0: matches "\\'\"hello" and "WÖRLD'\""
            project_sessions.append(await _add_project_session(session, projects[-1]))
            traces.append(await _add_trace(session, projects[-1], project_sessions[-1]))
            attributes = {"input": {"value": "\\'\"Hello Wörld'\""}}
            spans.append(await _add_span(session, traces[-1], attributes=attributes))

            # Session 1: matches "\\'\"hello" and "WÖRLD'\""
            project_sessions.append(await _add_project_session(session, projects[-1]))
            traces.append(await _add_trace(session, projects[-1], project_sessions[-1]))
            attributes = {"output": {"value": "\\'\"HELLO wörld'\""}}
            spans.append(await _add_span(session, traces[-1], attributes=attributes))

            # Session 2: matches "%"
            project_sessions.append(await _add_project_session(session, projects[-1]))
            traces.append(await _add_trace(session, projects[-1], project_sessions[-1]))
            attributes = {"input": {"value": "test%data"}}
            spans.append(await _add_span(session, traces[-1], attributes=attributes))

            # Session 3: matches "_"
            project_sessions.append(await _add_project_session(session, projects[-1]))
            traces.append(await _add_trace(session, projects[-1], project_sessions[-1]))
            attributes = {"output": {"value": "query_pattern"}}
            spans.append(await _add_span(session, traces[-1], attributes=attributes))

            # Session 4: matches "\\'\"hello"
            project_sessions.append(await _add_project_session(session, projects[-1]))
            traces.append(await _add_trace(session, projects[-1], project_sessions[-1]))
            attributes = {"input": {"value": "\\'\"Hello 世界"}}
            spans.append(await _add_span(session, traces[-1], attributes=attributes))

        return _Data(
            spans=spans,
            traces=traces,
            project_sessions=project_sessions,
            projects=projects,
        )

    async def test_session_dsl_containment_ignores_case(
        self,
        _case_insensitive_data: _Data,
        httpx_client: httpx.AsyncClient,
    ) -> None:
        """The DSL reaches every session the retired text search did: casing, non-ASCII, and
        SQL wildcard characters carried literally."""
        project = _case_insensitive_data.projects[0]
        sessions = _case_insensitive_data.project_sessions

        async def _matched(condition: str) -> set[str]:
            result = await self._node(
                f"sessions(first:50,sessionFilterCondition:{json.dumps(condition)}){{edges{{node{{id}}}}}}",
                project,
                httpx_client,
            )
            return {e["node"]["id"] for e in result["edges"]}

        # Sessions 0 and 4 spell it "Hello" on the input side; either casing finds both.
        hello_sessions = {_gid(sessions[0]), _gid(sessions[4])}
        assert await _matched("'Hello' in any_input") == hello_sessions
        assert await _matched("'hello' in any_input") == hello_sessions
        # Session 1 spells it "HELLO" on the output side.
        assert await _matched("'hello' in any_output") == {_gid(sessions[1])}
        # Non-ASCII case folding, in both directions.
        assert await _matched("'wörld' in any_input") == {_gid(sessions[0])}
        assert await _matched("'WÖRLD' in any_output") == {_gid(sessions[1])}
        # LIKE wildcards are literal text, not patterns.
        assert await _matched("'%' in any_input") == {_gid(sessions[2])}
        assert await _matched("'_' in any_output") == {_gid(sessions[3])}

    @pytest.mark.parametrize("orphan_span_as_root_span", [False, True])
    async def test_root_spans_only_with_orphan_spans(
        self,
        _orphan_spans: _Data,
        httpx_client: httpx.AsyncClient,
        orphan_span_as_root_span: bool,
    ) -> None:
        """Test pagination of root spans with orphan span handling.

        This test verifies that:
        1. Root spans are correctly identified based on orphan_span_as_root_span setting
        2. Pagination works correctly when fetching spans in chunks
        3. Spans are properly filtered and sorted
        """
        project = _orphan_spans.projects[0]

        # Filter spans containing "2" in their input value
        filtered_spans = [s for s in _orphan_spans.spans if "2" in s.attributes["input"]["value"]]

        # Determine root spans based on configuration
        if orphan_span_as_root_span:
            existing_span_ids = {s.span_id for s in _orphan_spans.spans}
            root_spans = [s for s in filtered_spans if s.parent_id not in existing_span_ids]
        else:
            root_spans = [s for s in filtered_spans if s.parent_id is None]

        # Sort spans by start time and ID
        sorted_spans = sorted(root_spans, key=lambda t: (t.start_time, t.id), reverse=True)

        # Convert to global IDs for comparison
        gids = list(map(_gid, sorted_spans))
        n = len(gids)
        first = n // 2 + 1  # Request half the spans plus one

        # Test pagination
        cursor = ""
        for i in range(n):
            expected = gids[i : i + first]

            # Construct GraphQL query
            field = (
                "spans("
                f"rootSpansOnly:true,"
                f"orphanSpanAsRootSpan:{str(orphan_span_as_root_span).lower()},"
                "sort:{col:startTime,dir:desc},"
                "filterCondition:\"'2' in input.value\","
                f"first:{str(first)},"
                f'after:"{cursor}"'
                "){edges{node{id}cursor}}"
            )

            # Execute query and verify results
            res = await self._node(field, project, httpx_client)
            assert [e["node"]["id"] for e in res["edges"]] == expected
            cursor = res["edges"][0]["cursor"]

    async def test_parent_is_none_matches_orphan_aware_root_spans_only(
        self,
        _orphan_spans: _Data,
        httpx_client: httpx.AsyncClient,
    ) -> None:
        """The DSL predicate ``filterCondition: "parent_span is None"`` selects the same
        spans as ``rootSpansOnly: true, orphanSpanAsRootSpan: true`` — verifying the
        correlated ``NOT EXISTS`` behaves correctly once the resolver's project
        predicate and query structure are applied, not just in direct SpanFilter
        execution.
        """
        project = _orphan_spans.projects[0]

        # Orphan-aware roots computed directly from the fixture: a NULL parent, or a
        # parent_id that references no span in the table.
        existing_span_ids = {s.span_id for s in _orphan_spans.spans}
        expected = {
            _gid(s)
            for s in _orphan_spans.spans
            if s.parent_id is None or s.parent_id not in existing_span_ids
        }
        # The fixture must exercise both kinds of root, or the test proves nothing.
        assert any(s.parent_id is None for s in _orphan_spans.spans)
        assert any(
            s.parent_id is not None and s.parent_id not in existing_span_ids
            for s in _orphan_spans.spans
        )

        dsl_res = await self._node(
            'spans(filterCondition:"parent_span is None",first:100){edges{node{id}}}',
            project,
            httpx_client,
        )
        flag_res = await self._node(
            "spans(rootSpansOnly:true,orphanSpanAsRootSpan:true,first:100){edges{node{id}}}",
            project,
            httpx_client,
        )
        dsl_ids = {e["node"]["id"] for e in dsl_res["edges"]}
        flag_ids = {e["node"]["id"] for e in flag_res["edges"]}
        assert dsl_ids == expected
        assert dsl_ids == flag_ids

    @pytest.mark.parametrize(
        "condition,orphan_span_as_root_span",
        [
            pytest.param("parent_span is None", True, id="orphan-aware-both"),
            pytest.param("parent_id is None", True, id="strict-condition-orphan-aware-flag"),
            pytest.param("parent_id is None", False, id="strict-both"),
            pytest.param("parent_span is None", False, id="orphan-aware-condition-strict-flag"),
            # Disjunction where every branch is strict-scoped: the flag skip now
            # fires off an `or`, so the rows must still match the flag alone.
            pytest.param(
                "(parent_id is None and '1' in input.value)"
                " or (parent_id is None and '3' in input.value)",
                True,
                id="disjunction-all-strict-orphan-flag",
            ),
            pytest.param(
                "(parent_id is None and '1' in input.value)"
                " or (parent_id is None and '3' in input.value)",
                False,
                id="disjunction-all-strict-strict-flag",
            ),
            # Negated root predicate: `not (parent_id is not None)` is strict-scoped,
            # so the skip fires off a `not` and must not change the result.
            pytest.param(
                "not (parent_id is not None) and '1' in input.value",
                True,
                id="negated-predicate-strict-orphan-flag",
            ),
            pytest.param(
                "not (parent_id is not None) and '1' in input.value",
                False,
                id="negated-predicate-strict-strict-flag",
            ),
            # A form the analyzer does *not* recognize (De Morgan over a compound):
            # it under-claims, so the flag's scoping is applied as well. Pinned
            # because an under-claim must stay harmless -- redundant SQL, same rows.
            pytest.param(
                "not (parent_id is not None or '9' in input.value)",
                True,
                id="unrecognized-form-flag-still-applied",
            ),
            # Disjunction whose branches mix strict and orphan-aware: scoped to
            # orphan-aware (the wider), so it may skip the orphan-aware flag but
            # not the strict one.
            pytest.param(
                "(parent_span is None and '1' in input.value)"
                " or (parent_id is None and '3' in input.value)",
                True,
                id="disjunction-mixed-orphan-flag",
            ),
            pytest.param(
                "(parent_span is None and '1' in input.value)"
                " or (parent_id is None and '3' in input.value)",
                False,
                id="disjunction-mixed-strict-flag",
            ),
        ],
    )
    async def test_root_predicate_and_flag_together_match_the_flag_alone(
        self,
        _orphan_spans: _Data,
        httpx_client: httpx.AsyncClient,
        condition: str,
        orphan_span_as_root_span: bool,
    ) -> None:
        """Sending `rootSpansOnly` *and* a root predicate must not change the result.

        The resolver drops the flag's scoping when the condition already implies
        it, so this pins that the optimization is semantics-preserving. The
        strict-flag/orphan-aware-condition case is the one that must NOT be
        optimized away: skipping a strict flag there would admit orphans.
        """
        project = _orphan_spans.projects[0]
        orphan_arg = str(orphan_span_as_root_span).lower()

        async def span_ids(field: str) -> set[str]:
            result = await self._node(field, project, httpx_client)
            return {e["node"]["id"] for e in result["edges"]}

        both = await span_ids(
            f"spans(rootSpansOnly:true,orphanSpanAsRootSpan:{orphan_arg},"
            f'filterCondition:"{condition}",first:100){{edges{{node{{id}}}}}}'
        )
        flag_only = await span_ids(
            f"spans(rootSpansOnly:true,orphanSpanAsRootSpan:{orphan_arg},"
            "first:100){edges{node{id}}}"
        )
        condition_only = await span_ids(
            f'spans(filterCondition:"{condition}",first:100){{edges{{node{{id}}}}}}'
        )
        # Combining them is the intersection, which is what the flag alone gives
        # whenever the condition is at least as strict.
        assert both == flag_only & condition_only
        assert both

    async def test_analyze_span_filter_condition_tracks_actual_query_scope(
        self,
        _orphan_spans: _Data,
        httpx_client: httpx.AsyncClient,
    ) -> None:
        """``analyzeSpanFilterCondition`` tells a client whether a filtered view is
        root-scoped, which is what decides between cumulative and per-span metric
        columns. Its answer has to match what the query actually returns, so this
        checks the verdict against the rows for both a scoped and an unscoped
        condition.
        """
        project = _orphan_spans.projects[0]
        user_condition = "'2' in input.value"
        root_scoped = f"parent_span is None and {user_condition}"

        async def analyze(condition: str) -> bool:
            result = await self._node(
                f'analyzeSpanFilterCondition(condition:"{condition}"){{selectsRootSpansOnly}}',
                project,
                httpx_client,
            )
            return bool(result["selectsRootSpansOnly"])

        async def span_ids(condition: str) -> set[str]:
            result = await self._node(
                f'spans(filterCondition:"{condition}",first:100){{edges{{node{{id}}}}}}',
                project,
                httpx_client,
            )
            return {e["node"]["id"] for e in result["edges"]}

        assert await analyze(user_condition) is False
        assert await analyze(root_scoped) is True

        scoped_ids = await span_ids(root_scoped)
        unscoped_ids = await span_ids(user_condition)
        # The verdict is only meaningful if the predicate actually narrowed the
        # result; a silently dropped predicate would make these equal.
        assert scoped_ids < unscoped_ids
        # ...and it narrows to exactly what the boolean arguments selected.
        flag_res = await self._node(
            "spans(rootSpansOnly:true,orphanSpanAsRootSpan:true,"
            f'filterCondition:"{user_condition}",first:100){{edges{{node{{id}}}}}}',
            project,
            httpx_client,
        )
        assert scoped_ids == {e["node"]["id"] for e in flag_res["edges"]}

    @pytest.fixture
    async def _time_series_data(
        self,
        db: DbSessionFactory,
    ) -> _Data:
        """Creates a dataset for testing span_count_time_series using random timestamps.

        Creates spans with random timestamps distributed across different time periods
        to test time series grouping and filtering functionality. The spans are created
        within the last day to ensure they fall within reasonable time buckets.

        Returns:
            _Data object containing:
            - projects: List of test projects
            - traces: List of test traces
            - spans: List of test spans with random timestamps
        """
        projects: list[models.Project] = []
        traces: list[models.Trace] = []
        spans: list[models.Span] = []

        async with db() as session:
            # Create a test project
            projects.append(models.Project(name=token_hex(8)))
            session.add(projects[-1])
            await session.flush()

            # Create multiple traces with spans at different times
            for _ in range(10):
                # Generate random trace start time within the test time range (2024-01-01)
                trace_start = fake.date_time_between(
                    start_date=datetime.fromisoformat("2024-01-01T00:00:00+00:00"),
                    end_date=datetime.fromisoformat("2024-01-01T23:59:59+00:00"),
                    tzinfo=timezone.utc,
                )
                trace_start = trace_start.replace(microsecond=0)
                trace_end = trace_start + timedelta(seconds=30)

                # Create trace
                trace = models.Trace(
                    trace_id=token_hex(16),
                    project_rowid=projects[-1].id,
                    start_time=trace_start,
                    end_time=trace_end,
                )
                session.add(trace)
                await session.flush()
                traces.append(trace)

                # Create span for this trace
                span = models.Span(
                    trace_rowid=trace.id,
                    span_id=token_hex(8),
                    parent_id=None,
                    name=token_hex(8),
                    span_kind="LLM",
                    start_time=trace_start,
                    end_time=trace_end,
                    attributes={},
                    events=[],
                    status_code="OK",
                    status_message="",
                    cumulative_error_count=0,
                    cumulative_llm_token_count_prompt=0,
                    cumulative_llm_token_count_completion=0,
                )
                spans.append(span)
            session.add_all(spans)

        return _Data(
            spans=spans,
            traces=traces,
            projects=projects,
        )

    @staticmethod
    def _count_rows(
        df: pd.DataFrame,
        field: Literal["minute", "hour", "day", "week", "month", "year"],
        utc_offset_minutes: int,
    ) -> pd.DataFrame:
        offset_tz = timezone(timedelta(minutes=utc_offset_minutes))
        t = df.loc[:, "timestamp"].dt.tz_convert(offset_tz)
        if field == "minute":
            t = t.dt.floor("T")
        elif field == "hour":
            t = t.dt.floor("H")
        elif field == "day":
            t = t.dt.floor("D")
        elif field == "week":
            t = t.dt.to_period("W").dt.start_time.dt.tz_localize(offset_tz)
        elif field == "month":
            t = t.dt.to_period("M").dt.start_time.dt.tz_localize(offset_tz)
        elif field == "year":
            t = t.dt.to_period("Y").dt.start_time.dt.tz_localize(offset_tz)
        else:
            assert_never(field)
        t = t.dt.tz_convert(timezone.utc)
        return df.groupby(t).size().reset_index(name="count")

    @pytest.mark.parametrize(
        "time_range, time_bin_config",
        [
            # === BASIC HOURLY TESTS ===
            pytest.param(
                TimeRange(
                    start=datetime.fromisoformat("2024-01-01T01:00:00+00:00"),
                    end=datetime.fromisoformat("2024-01-02T00:00:00+00:00"),
                ),
                None,
                id="default-hourly-no-offset",
            ),
            pytest.param(
                TimeRange(
                    start=datetime.fromisoformat("2024-01-01T01:00:00+00:00"),
                    end=datetime.fromisoformat("2024-01-01T03:00:00+00:00"),
                ),
                TimeBinConfig(scale=TimeBinScale.HOUR, utc_offset_minutes=0),
                id="hourly-no-offset",
            ),
            # === MINUTE GRANULARITY TESTS ===
            pytest.param(
                TimeRange(
                    start=datetime.fromisoformat("2024-01-01T01:00:00+00:00"),
                    end=datetime.fromisoformat("2024-01-01T01:30:00+00:00"),
                ),
                TimeBinConfig(scale=TimeBinScale.MINUTE, utc_offset_minutes=0),
                id="minute-no-offset",
            ),
            pytest.param(
                TimeRange(
                    start=datetime.fromisoformat("2024-01-01T01:00:00+00:00"),
                    end=datetime.fromisoformat("2024-01-01T02:00:00+00:00"),
                ),
                TimeBinConfig(scale=TimeBinScale.MINUTE, utc_offset_minutes=60),
                id="minute-with-positive-offset",
            ),
            # === DAILY GRANULARITY TESTS ===
            pytest.param(
                TimeRange(
                    start=datetime.fromisoformat("2024-01-01T00:00:00+00:00"),
                    end=datetime.fromisoformat("2024-01-02T00:00:00+00:00"),
                ),
                TimeBinConfig(scale=TimeBinScale.DAY, utc_offset_minutes=0),
                id="daily-no-offset",
            ),
            pytest.param(
                TimeRange(
                    start=datetime.fromisoformat("2024-01-01T00:00:00+00:00"),
                    end=datetime.fromisoformat("2024-01-02T00:00:00+00:00"),
                ),
                TimeBinConfig(scale=TimeBinScale.DAY, utc_offset_minutes=-480),  # PST offset
                id="daily-with-negative-offset",
            ),
            # === WEEKLY GRANULARITY TESTS ===
            pytest.param(
                TimeRange(
                    start=datetime.fromisoformat("2024-01-01T00:00:00+00:00"),
                    end=datetime.fromisoformat("2024-01-08T00:00:00+00:00"),
                ),
                TimeBinConfig(scale=TimeBinScale.WEEK, utc_offset_minutes=0),
                id="weekly-no-offset",
            ),
            pytest.param(
                TimeRange(
                    start=datetime.fromisoformat("2024-01-01T00:00:00+00:00"),
                    end=datetime.fromisoformat("2024-01-08T00:00:00+00:00"),
                ),
                TimeBinConfig(
                    scale=TimeBinScale.WEEK, utc_offset_minutes=330
                ),  # India Standard Time
                id="weekly-with-positive-offset",
            ),
            # === MONTHLY GRANULARITY TESTS ===
            pytest.param(
                TimeRange(
                    start=datetime.fromisoformat("2024-01-01T00:00:00+00:00"),
                    end=datetime.fromisoformat("2024-02-01T00:00:00+00:00"),
                ),
                TimeBinConfig(scale=TimeBinScale.MONTH, utc_offset_minutes=0),
                id="monthly-no-offset",
            ),
            pytest.param(
                TimeRange(
                    start=datetime.fromisoformat("2024-01-01T00:00:00+00:00"),
                    end=datetime.fromisoformat("2024-03-01T00:00:00+00:00"),
                ),
                TimeBinConfig(scale=TimeBinScale.MONTH, utc_offset_minutes=-300),  # EST offset
                id="monthly-with-negative-offset",
            ),
            # === YEARLY GRANULARITY TESTS ===
            pytest.param(
                TimeRange(
                    start=datetime.fromisoformat("2024-01-01T00:00:00+00:00"),
                    end=datetime.fromisoformat("2025-01-01T00:00:00+00:00"),
                ),
                TimeBinConfig(scale=TimeBinScale.YEAR, utc_offset_minutes=0),
                id="yearly-no-offset",
            ),
            pytest.param(
                TimeRange(
                    start=datetime.fromisoformat("2023-01-01T00:00:00+00:00"),
                    end=datetime.fromisoformat("2025-01-01T00:00:00+00:00"),
                ),
                TimeBinConfig(scale=TimeBinScale.YEAR, utc_offset_minutes=540),  # JST offset
                id="yearly-with-positive-offset",
            ),
            # === EDGE CASES ===
            # Empty result set (time range with no data)
            pytest.param(
                TimeRange(
                    start=datetime.fromisoformat("2023-01-01T00:00:00+00:00"),
                    end=datetime.fromisoformat("2023-01-01T01:00:00+00:00"),
                ),
                TimeBinConfig(scale=TimeBinScale.HOUR, utc_offset_minutes=0),
                id="empty-result-set",
            ),
            # Very narrow time range
            pytest.param(
                TimeRange(
                    start=datetime.fromisoformat("2024-01-01T12:00:00+00:00"),
                    end=datetime.fromisoformat("2024-01-01T12:00:01+00:00"),
                ),
                TimeBinConfig(scale=TimeBinScale.MINUTE, utc_offset_minutes=0),
                id="narrow-time-range",
            ),
            # Limited data points
            pytest.param(
                TimeRange(
                    start=datetime.fromisoformat("2024-01-01T00:00:00+00:00"),
                    end=datetime.fromisoformat("2024-01-01T00:30:00+00:00"),
                ),
                TimeBinConfig(scale=TimeBinScale.MINUTE, utc_offset_minutes=0),
                id="limited-data-points",
            ),
            # Boundary condition: Start equals end
            pytest.param(
                TimeRange(
                    start=datetime.fromisoformat("2024-01-01T12:00:00+00:00"),
                    end=datetime.fromisoformat("2024-01-01T12:00:00+00:00"),
                ),
                TimeBinConfig(scale=TimeBinScale.HOUR, utc_offset_minutes=0),
                id="zero-duration-range",
            ),
            # === PARTIAL TIME RANGE TESTS ===
            # Only start time specified
            pytest.param(
                TimeRange(
                    start=datetime.fromisoformat("2024-01-01T12:00:00+00:00"),
                    end=None,
                ),
                TimeBinConfig(scale=TimeBinScale.HOUR, utc_offset_minutes=0),
                id="only-start-time",
            ),
            # === LARGE UTC OFFSET TESTS ===
            pytest.param(
                TimeRange(
                    start=datetime.fromisoformat("2024-01-01T00:00:00+00:00"),
                    end=datetime.fromisoformat("2024-01-02T00:00:00+00:00"),
                ),
                TimeBinConfig(scale=TimeBinScale.HOUR, utc_offset_minutes=780),  # +13 hours
                id="large-positive-offset",
            ),
            pytest.param(
                TimeRange(
                    start=datetime.fromisoformat("2024-01-01T00:00:00+00:00"),
                    end=datetime.fromisoformat("2024-01-02T00:00:00+00:00"),
                ),
                TimeBinConfig(scale=TimeBinScale.HOUR, utc_offset_minutes=-720),  # -12 hours
                id="large-negative-offset",
            ),
            # === BOUNDARY CONDITION TESTS ===
            # Cross-day boundary with offset
            pytest.param(
                TimeRange(
                    start=datetime.fromisoformat("2024-01-01T22:00:00+00:00"),
                    end=datetime.fromisoformat("2024-01-02T02:00:00+00:00"),
                ),
                TimeBinConfig(scale=TimeBinScale.HOUR, utc_offset_minutes=480),  # +8 hours
                id="cross-day-boundary",
            ),
            # Test with fractional hour offset
            pytest.param(
                TimeRange(
                    start=datetime.fromisoformat("2024-01-01T00:00:00+00:00"),
                    end=datetime.fromisoformat("2024-01-01T06:00:00+00:00"),
                ),
                TimeBinConfig(scale=TimeBinScale.HOUR, utc_offset_minutes=90),  # +1.5 hours
                id="fractional-hour-offset",
            ),
            # Test with leap year boundaries
            pytest.param(
                TimeRange(
                    start=datetime.fromisoformat("2024-02-28T00:00:00+00:00"),
                    end=datetime.fromisoformat("2024-03-01T00:00:00+00:00"),
                ),
                TimeBinConfig(scale=TimeBinScale.DAY, utc_offset_minutes=0),
                id="leap-year-boundary",
            ),
            # Test with month boundary at different scales
            pytest.param(
                TimeRange(
                    start=datetime.fromisoformat("2024-01-31T22:00:00+00:00"),
                    end=datetime.fromisoformat("2024-02-01T02:00:00+00:00"),
                ),
                TimeBinConfig(scale=TimeBinScale.HOUR, utc_offset_minutes=0),
                id="month-boundary-hourly",
            ),
            # Test with daylight saving time-like offset changes
            pytest.param(
                TimeRange(
                    start=datetime.fromisoformat("2024-01-01T00:00:00+00:00"),
                    end=datetime.fromisoformat("2024-01-01T06:00:00+00:00"),
                ),
                TimeBinConfig(scale=TimeBinScale.HOUR, utc_offset_minutes=-480),  # PST
                id="dst-like-offset",
            ),
            # Test with extreme small time range at minute level
            pytest.param(
                TimeRange(
                    start=datetime.fromisoformat("2024-01-01T12:30:00+00:00"),
                    end=datetime.fromisoformat("2024-01-01T12:31:00+00:00"),
                ),
                TimeBinConfig(scale=TimeBinScale.MINUTE, utc_offset_minutes=0),
                id="single-minute-range",
            ),
            # Test with year-end boundary
            pytest.param(
                TimeRange(
                    start=datetime.fromisoformat("2024-12-31T22:00:00+00:00"),
                    end=datetime.fromisoformat("2025-01-01T02:00:00+00:00"),
                ),
                TimeBinConfig(scale=TimeBinScale.HOUR, utc_offset_minutes=0),
                id="year-end-boundary",
            ),
            # Test with maximum reasonable offset (UTC+14)
            pytest.param(
                TimeRange(
                    start=datetime.fromisoformat("2024-01-01T00:00:00+00:00"),
                    end=datetime.fromisoformat("2024-01-02T00:00:00+00:00"),
                ),
                TimeBinConfig(
                    scale=TimeBinScale.HOUR, utc_offset_minutes=840
                ),  # +14 hours (Line Islands)
                id="maximum-positive-offset",
            ),
            # Test with minimum reasonable offset (UTC-12)
            pytest.param(
                TimeRange(
                    start=datetime.fromisoformat("2024-01-01T00:00:00+00:00"),
                    end=datetime.fromisoformat("2024-01-02T00:00:00+00:00"),
                ),
                TimeBinConfig(
                    scale=TimeBinScale.HOUR, utc_offset_minutes=-720
                ),  # -12 hours (Baker Island)
                id="minimum-negative-offset",
            ),
        ],
    )
    async def test_span_count_time_series(
        self,
        time_range: TimeRange,
        time_bin_config: TimeBinConfig,
        _time_series_data: _Data,
        gql_client: AsyncGraphQLClient,
    ) -> None:
        """Test the span_count_time_series field using pandas validation.

        This comprehensive test verifies that the SQL-based time series calculation matches
        pandas-based calculations using the same logic. It covers:

        **Time Granularities:**
        - Minute-level aggregation
        - Hourly aggregation (default)
        - Daily aggregation
        - Weekly aggregation
        - Monthly aggregation
        - Yearly aggregation

        **UTC Offset Scenarios:**
        - No offset (UTC+0)
        - Positive offsets (UTC+1, UTC+5.5, UTC+8, UTC+9, UTC+13, UTC+14)
        - Negative offsets (UTC-5, UTC-8, UTC-12)
        - Fractional hour offsets (UTC+1.5)

        **Time Range Edge Cases:**
        - Empty result sets (no data in range)
        - Very narrow time ranges (1 second, 1 minute)
        - Limited data points
        - Boundary conditions (start == end)
        - Partial range specifications (start-only, end-only)
        - No time range specified (all data)

        **Boundary Conditions:**
        - Cross-day boundaries with offsets
        - Month boundaries (including leap year)
        - Year-end boundaries
        - Leap year date boundaries

        **Real-world Timezone Examples:**
        - PST (UTC-8)
        - EST (UTC-5)
        - IST (UTC+5.5)
        - JST (UTC+9)
        - Line Islands (UTC+14)
        - Baker Island (UTC-12)

        Each test case validates both span and trace count time series.
        """
        project = _time_series_data.projects[0]

        # Calculate expected results using pandas (same logic as SQL)
        records = _time_series_data.spans
        data_df = pd.DataFrame([{"timestamp": s.start_time} for s in records]).sort_values(
            "timestamp"
        )

        # Apply time range filtering if specified
        if time_range and time_range.start:
            data_df = data_df[data_df["timestamp"] >= time_range.start]
        if time_range and time_range.end:
            data_df = data_df[data_df["timestamp"] < time_range.end]

        if data_df.empty:
            expected_summary = pd.DataFrame(columns=["timestamp", "count"])
        else:
            expected_summary = self._count_rows(
                data_df,
                field=time_bin_config.scale.value if time_bin_config else "hour",
                utc_offset_minutes=time_bin_config.utc_offset_minutes if time_bin_config else 0,
            )

        # Execute GraphQL query
        project_gid = str(GlobalID(type_name="Project", node_id=str(project.id)))
        variables: dict[str, Any] = {"id": project_gid}

        query = """
            query($id: ID!, $timeRange: TimeRange!, $timeBinConfig: TimeBinConfig) {
                node(id: $id) {
                    ... on Project {
                        spanCountTimeSeries(timeRange: $timeRange, timeBinConfig: $timeBinConfig) {
                            data {
                                timestamp
                                totalCount
                            }
                        }
                    }
                }
            }
        """

        if time_range:
            time_range_vars = {}
            if time_range.start:
                time_range_vars["start"] = time_range.start.isoformat()
            if time_range.end:
                time_range_vars["end"] = time_range.end.isoformat()
            if time_range_vars:
                variables["timeRange"] = time_range_vars
        if time_bin_config:
            variables["timeBinConfig"] = {
                "scale": time_bin_config.scale.value.upper(),
                "utcOffsetMinutes": time_bin_config.utc_offset_minutes,
            }

        # Execute GraphQL query
        response = await gql_client.execute(query=query, variables=variables)
        assert not response.errors
        assert (data := response.data) is not None
        res = data["node"]["spanCountTimeSeries"]

        # Verify the structure of the response
        assert "data" in res
        assert isinstance(res["data"], list)

        # Convert response to DataFrame for comparison
        if not res["data"]:
            actual_summary = pd.DataFrame(columns=["timestamp", "count"])
        else:
            actual_data = []
            for data_point in res["data"]:
                timestamp = datetime.fromisoformat(data_point["timestamp"])
                if (value := data_point["totalCount"]) is not None:
                    actual_data.append({"timestamp": timestamp, "count": value})
            actual_summary = pd.DataFrame(
                actual_data,
                columns=["timestamp", "count"],
            ).sort_values("timestamp")

        # Handle empty results
        if expected_summary.empty:
            assert actual_summary.empty, "Expected empty summary for span"
            return

        actual_summary["timestamp"] = pd.to_datetime(actual_summary["timestamp"])

        # Verify SQL results match pandas calculation
        pd.testing.assert_frame_equal(actual_summary, expected_summary, check_dtype=False)

    @pytest.mark.parametrize(
        "time_range, time_bin_config",
        [
            # === BASIC HOURLY TESTS ===
            pytest.param(
                TimeRange(
                    start=datetime.fromisoformat("2024-01-01T01:00:00+00:00"),
                    end=datetime.fromisoformat("2024-01-02T00:00:00+00:00"),
                ),
                None,
                id="default-hourly-no-offset",
            ),
            pytest.param(
                TimeRange(
                    start=datetime.fromisoformat("2024-01-01T01:00:00+00:00"),
                    end=datetime.fromisoformat("2024-01-01T03:00:00+00:00"),
                ),
                TimeBinConfig(scale=TimeBinScale.HOUR, utc_offset_minutes=0),
                id="hourly-no-offset",
            ),
            # === MINUTE GRANULARITY TESTS ===
            pytest.param(
                TimeRange(
                    start=datetime.fromisoformat("2024-01-01T01:00:00+00:00"),
                    end=datetime.fromisoformat("2024-01-01T01:30:00+00:00"),
                ),
                TimeBinConfig(scale=TimeBinScale.MINUTE, utc_offset_minutes=0),
                id="minute-no-offset",
            ),
            pytest.param(
                TimeRange(
                    start=datetime.fromisoformat("2024-01-01T01:00:00+00:00"),
                    end=datetime.fromisoformat("2024-01-01T02:00:00+00:00"),
                ),
                TimeBinConfig(scale=TimeBinScale.MINUTE, utc_offset_minutes=60),
                id="minute-with-positive-offset",
            ),
            # === DAILY GRANULARITY TESTS ===
            pytest.param(
                TimeRange(
                    start=datetime.fromisoformat("2024-01-01T00:00:00+00:00"),
                    end=datetime.fromisoformat("2024-01-02T00:00:00+00:00"),
                ),
                TimeBinConfig(scale=TimeBinScale.DAY, utc_offset_minutes=0),
                id="daily-no-offset",
            ),
            pytest.param(
                TimeRange(
                    start=datetime.fromisoformat("2024-01-01T00:00:00+00:00"),
                    end=datetime.fromisoformat("2024-01-02T00:00:00+00:00"),
                ),
                TimeBinConfig(scale=TimeBinScale.DAY, utc_offset_minutes=-480),  # PST offset
                id="daily-with-negative-offset",
            ),
            # === WEEKLY GRANULARITY TESTS ===
            pytest.param(
                TimeRange(
                    start=datetime.fromisoformat("2024-01-01T00:00:00+00:00"),
                    end=datetime.fromisoformat("2024-01-08T00:00:00+00:00"),
                ),
                TimeBinConfig(scale=TimeBinScale.WEEK, utc_offset_minutes=0),
                id="weekly-no-offset",
            ),
            pytest.param(
                TimeRange(
                    start=datetime.fromisoformat("2024-01-01T00:00:00+00:00"),
                    end=datetime.fromisoformat("2024-01-08T00:00:00+00:00"),
                ),
                TimeBinConfig(
                    scale=TimeBinScale.WEEK, utc_offset_minutes=330
                ),  # India Standard Time
                id="weekly-with-positive-offset",
            ),
            # === MONTHLY GRANULARITY TESTS ===
            pytest.param(
                TimeRange(
                    start=datetime.fromisoformat("2024-01-01T00:00:00+00:00"),
                    end=datetime.fromisoformat("2024-02-01T00:00:00+00:00"),
                ),
                TimeBinConfig(scale=TimeBinScale.MONTH, utc_offset_minutes=0),
                id="monthly-no-offset",
            ),
            pytest.param(
                TimeRange(
                    start=datetime.fromisoformat("2024-01-01T00:00:00+00:00"),
                    end=datetime.fromisoformat("2024-03-01T00:00:00+00:00"),
                ),
                TimeBinConfig(scale=TimeBinScale.MONTH, utc_offset_minutes=-300),  # EST offset
                id="monthly-with-negative-offset",
            ),
            # === YEARLY GRANULARITY TESTS ===
            pytest.param(
                TimeRange(
                    start=datetime.fromisoformat("2024-01-01T00:00:00+00:00"),
                    end=datetime.fromisoformat("2025-01-01T00:00:00+00:00"),
                ),
                TimeBinConfig(scale=TimeBinScale.YEAR, utc_offset_minutes=0),
                id="yearly-no-offset",
            ),
            pytest.param(
                TimeRange(
                    start=datetime.fromisoformat("2023-01-01T00:00:00+00:00"),
                    end=datetime.fromisoformat("2025-01-01T00:00:00+00:00"),
                ),
                TimeBinConfig(scale=TimeBinScale.YEAR, utc_offset_minutes=540),  # JST offset
                id="yearly-with-positive-offset",
            ),
            # === EDGE CASES ===
            # Empty result set (time range with no data)
            pytest.param(
                TimeRange(
                    start=datetime.fromisoformat("2023-01-01T00:00:00+00:00"),
                    end=datetime.fromisoformat("2023-01-01T01:00:00+00:00"),
                ),
                TimeBinConfig(scale=TimeBinScale.HOUR, utc_offset_minutes=0),
                id="empty-result-set",
            ),
            # Very narrow time range
            pytest.param(
                TimeRange(
                    start=datetime.fromisoformat("2024-01-01T12:00:00+00:00"),
                    end=datetime.fromisoformat("2024-01-01T12:00:01+00:00"),
                ),
                TimeBinConfig(scale=TimeBinScale.MINUTE, utc_offset_minutes=0),
                id="narrow-time-range",
            ),
            # Limited data points
            pytest.param(
                TimeRange(
                    start=datetime.fromisoformat("2024-01-01T00:00:00+00:00"),
                    end=datetime.fromisoformat("2024-01-01T00:30:00+00:00"),
                ),
                TimeBinConfig(scale=TimeBinScale.MINUTE, utc_offset_minutes=0),
                id="limited-data-points",
            ),
            # Boundary condition: Start equals end
            pytest.param(
                TimeRange(
                    start=datetime.fromisoformat("2024-01-01T12:00:00+00:00"),
                    end=datetime.fromisoformat("2024-01-01T12:00:00+00:00"),
                ),
                TimeBinConfig(scale=TimeBinScale.HOUR, utc_offset_minutes=0),
                id="zero-duration-range",
            ),
            # === PARTIAL TIME RANGE TESTS ===
            # Only start time specified
            pytest.param(
                TimeRange(
                    start=datetime.fromisoformat("2024-01-01T12:00:00+00:00"),
                    end=None,
                ),
                TimeBinConfig(scale=TimeBinScale.HOUR, utc_offset_minutes=0),
                id="only-start-time",
            ),
            # === LARGE UTC OFFSET TESTS ===
            pytest.param(
                TimeRange(
                    start=datetime.fromisoformat("2024-01-01T00:00:00+00:00"),
                    end=datetime.fromisoformat("2024-01-02T00:00:00+00:00"),
                ),
                TimeBinConfig(scale=TimeBinScale.HOUR, utc_offset_minutes=780),  # +13 hours
                id="large-positive-offset",
            ),
            pytest.param(
                TimeRange(
                    start=datetime.fromisoformat("2024-01-01T00:00:00+00:00"),
                    end=datetime.fromisoformat("2024-01-02T00:00:00+00:00"),
                ),
                TimeBinConfig(scale=TimeBinScale.HOUR, utc_offset_minutes=-720),  # -12 hours
                id="large-negative-offset",
            ),
            # === BOUNDARY CONDITION TESTS ===
            # Cross-day boundary with offset
            pytest.param(
                TimeRange(
                    start=datetime.fromisoformat("2024-01-01T22:00:00+00:00"),
                    end=datetime.fromisoformat("2024-01-02T02:00:00+00:00"),
                ),
                TimeBinConfig(scale=TimeBinScale.HOUR, utc_offset_minutes=480),  # +8 hours
                id="cross-day-boundary",
            ),
            # Test with fractional hour offset
            pytest.param(
                TimeRange(
                    start=datetime.fromisoformat("2024-01-01T00:00:00+00:00"),
                    end=datetime.fromisoformat("2024-01-01T06:00:00+00:00"),
                ),
                TimeBinConfig(scale=TimeBinScale.HOUR, utc_offset_minutes=90),  # +1.5 hours
                id="fractional-hour-offset",
            ),
            # Test with leap year boundaries
            pytest.param(
                TimeRange(
                    start=datetime.fromisoformat("2024-02-28T00:00:00+00:00"),
                    end=datetime.fromisoformat("2024-03-01T00:00:00+00:00"),
                ),
                TimeBinConfig(scale=TimeBinScale.DAY, utc_offset_minutes=0),
                id="leap-year-boundary",
            ),
            # Test with month boundary at different scales
            pytest.param(
                TimeRange(
                    start=datetime.fromisoformat("2024-01-31T22:00:00+00:00"),
                    end=datetime.fromisoformat("2024-02-01T02:00:00+00:00"),
                ),
                TimeBinConfig(scale=TimeBinScale.HOUR, utc_offset_minutes=0),
                id="month-boundary-hourly",
            ),
            # Test with daylight saving time-like offset changes
            pytest.param(
                TimeRange(
                    start=datetime.fromisoformat("2024-01-01T00:00:00+00:00"),
                    end=datetime.fromisoformat("2024-01-01T06:00:00+00:00"),
                ),
                TimeBinConfig(scale=TimeBinScale.HOUR, utc_offset_minutes=-480),  # PST
                id="dst-like-offset",
            ),
            # Test with extreme small time range at minute level
            pytest.param(
                TimeRange(
                    start=datetime.fromisoformat("2024-01-01T12:30:00+00:00"),
                    end=datetime.fromisoformat("2024-01-01T12:31:00+00:00"),
                ),
                TimeBinConfig(scale=TimeBinScale.MINUTE, utc_offset_minutes=0),
                id="single-minute-range",
            ),
            # Test with year-end boundary
            pytest.param(
                TimeRange(
                    start=datetime.fromisoformat("2024-12-31T22:00:00+00:00"),
                    end=datetime.fromisoformat("2025-01-01T02:00:00+00:00"),
                ),
                TimeBinConfig(scale=TimeBinScale.HOUR, utc_offset_minutes=0),
                id="year-end-boundary",
            ),
            # Test with maximum reasonable offset (UTC+14)
            pytest.param(
                TimeRange(
                    start=datetime.fromisoformat("2024-01-01T00:00:00+00:00"),
                    end=datetime.fromisoformat("2024-01-02T00:00:00+00:00"),
                ),
                TimeBinConfig(
                    scale=TimeBinScale.HOUR, utc_offset_minutes=840
                ),  # +14 hours (Line Islands)
                id="maximum-positive-offset",
            ),
            # Test with minimum reasonable offset (UTC-12)
            pytest.param(
                TimeRange(
                    start=datetime.fromisoformat("2024-01-01T00:00:00+00:00"),
                    end=datetime.fromisoformat("2024-01-02T00:00:00+00:00"),
                ),
                TimeBinConfig(
                    scale=TimeBinScale.HOUR, utc_offset_minutes=-720
                ),  # -12 hours (Baker Island)
                id="minimum-negative-offset",
            ),
        ],
    )
    async def test_trace_count_time_series(
        self,
        time_range: TimeRange,
        time_bin_config: TimeBinConfig,
        _time_series_data: _Data,
        gql_client: AsyncGraphQLClient,
    ) -> None:
        """Test the trace_count_time_series field using pandas validation.

        This comprehensive test verifies that the SQL-based time series calculation matches
        pandas-based calculations using the same logic. It covers:

        **Time Granularities:**
        - Minute-level aggregation
        - Hourly aggregation (default)
        - Daily aggregation
        - Weekly aggregation
        - Monthly aggregation
        - Yearly aggregation

        **UTC Offset Scenarios:**
        - No offset (UTC+0)
        - Positive offsets (UTC+1, UTC+5.5, UTC+8, UTC+9, UTC+13, UTC+14)
        - Negative offsets (UTC-5, UTC-8, UTC-12)
        - Fractional hour offsets (UTC+1.5)

        **Time Range Edge Cases:**
        - Empty result sets (no data in range)
        - Very narrow time ranges (1 second, 1 minute)
        - Limited data points
        - Boundary conditions (start == end)
        - Partial range specifications (start-only, end-only)
        - No time range specified (all data)

        **Boundary Conditions:**
        - Cross-day boundaries with offsets
        - Month boundaries (including leap year)
        - Year-end boundaries
        - Leap year date boundaries

        **Real-world Timezone Examples:**
        - PST (UTC-8)
        - EST (UTC-5)
        - IST (UTC+5.5)
        - JST (UTC+9)
        - Line Islands (UTC+14)
        - Baker Island (UTC-12)

        Each test case validates trace count time series.
        """
        project = _time_series_data.projects[0]

        # Calculate expected results using pandas (same logic as SQL)
        records = _time_series_data.traces
        data_df = pd.DataFrame([{"timestamp": t.start_time} for t in records]).sort_values(
            "timestamp"
        )

        # Apply time range filtering if specified
        if time_range and time_range.start:
            data_df = data_df[data_df["timestamp"] >= time_range.start]
        if time_range and time_range.end:
            data_df = data_df[data_df["timestamp"] < time_range.end]

        if data_df.empty:
            expected_summary = pd.DataFrame(columns=["timestamp", "count"])
        else:
            expected_summary = self._count_rows(
                data_df,
                field=time_bin_config.scale.value if time_bin_config else "hour",
                utc_offset_minutes=time_bin_config.utc_offset_minutes if time_bin_config else 0,
            )

        # Execute GraphQL query
        project_gid = str(GlobalID(type_name="Project", node_id=str(project.id)))
        variables: dict[str, Any] = {"id": project_gid}

        query = """
            query($id: ID!, $timeRange: TimeRange!, $timeBinConfig: TimeBinConfig) {
                node(id: $id) {
                    ... on Project {
                        traceCountTimeSeries(timeRange: $timeRange, timeBinConfig: $timeBinConfig) {
                            data {
                                timestamp
                                value
                            }
                        }
                    }
                }
            }
        """

        if time_range:
            time_range_vars = {}
            if time_range.start:
                time_range_vars["start"] = time_range.start.isoformat()
            if time_range.end:
                time_range_vars["end"] = time_range.end.isoformat()
            if time_range_vars:
                variables["timeRange"] = time_range_vars
        if time_bin_config:
            variables["timeBinConfig"] = {
                "scale": time_bin_config.scale.value.upper(),
                "utcOffsetMinutes": time_bin_config.utc_offset_minutes,
            }

        # Execute GraphQL query
        response = await gql_client.execute(query=query, variables=variables)
        assert not response.errors
        assert (data := response.data) is not None
        res = data["node"]["traceCountTimeSeries"]

        # Verify the structure of the response
        assert "data" in res
        assert isinstance(res["data"], list)

        # Convert response to DataFrame for comparison
        if not res["data"]:
            actual_summary = pd.DataFrame(columns=["timestamp", "count"])
        else:
            actual_data = []
            for data_point in res["data"]:
                timestamp = datetime.fromisoformat(data_point["timestamp"])
                if (value := data_point["value"]) is not None:
                    actual_data.append({"timestamp": timestamp, "count": value})
            actual_summary = pd.DataFrame(
                actual_data,
                columns=["timestamp", "count"],
            ).sort_values("timestamp")

        # Handle empty results
        if expected_summary.empty:
            assert actual_summary.empty, "Expected empty summary for trace count time series"
            return

        actual_summary["timestamp"] = pd.to_datetime(actual_summary["timestamp"])

        # Verify SQL results match pandas calculation
        try:
            pd.testing.assert_frame_equal(actual_summary, expected_summary, check_dtype=False)
        except AssertionError as e:
            raise AssertionError("Test failed for trace count time series") from e

    @pytest.mark.parametrize(
        "expectation,condition,error_message",
        [
            (True, "span_kind == 'LLM'", None),
            (False, "span_kind == 'LLM' and ", None),
            (False, "span_kind == 'LLM' and ''", None),
            (
                False,
                "annotations['rate_limit_type'].label is True",
                "cannot compare string and boolean",
            ),
        ],
    )
    async def test_validate_span_filter_condition(
        self,
        condition: str,
        expectation: bool,
        error_message: Optional[str],
        gql_client: AsyncGraphQLClient,
        db: DbSessionFactory,
    ) -> None:
        async with db() as session:
            project = models.Project(name=token_hex(8))
            session.add(project)
        query = """
            query($id: ID!, $condition: String!) {
              node(id: $id) {
                ... on Project {
                  validateSpanFilterCondition(
                    condition: $condition
                  ) {
                    isValid
                    errorMessage
                  }
                }
              }
            }
        """
        project_gid = str(GlobalID(type_name="Project", node_id=str(project.id)))
        response = await gql_client.execute(
            query=query,
            variables={"id": project_gid, "condition": condition},
        )
        assert not response.errors
        assert (data := response.data) is not None
        result = data["node"]["validateSpanFilterCondition"]
        assert result["isValid"] == expectation
        if error_message is not None:
            assert error_message in result["errorMessage"]

    async def test_invalid_span_filter_returns_actionable_execution_error(
        self,
        gql_client: AsyncGraphQLClient,
        db: DbSessionFactory,
    ) -> None:
        async with db() as session:
            project = models.Project(name=token_hex(8))
            session.add(project)

        query = """
            query($id: ID!, $condition: String!) {
              node(id: $id) {
                ... on Project {
                  spans(first: 1, filterCondition: $condition) {
                    edges { node { id } }
                  }
                }
              }
            }
        """
        response = await gql_client.execute(
            query=query,
            variables={
                "id": str(GlobalID(type_name="Project", node_id=str(project.id))),
                "condition": "annotations['rate_limit_type'].label is True",
            },
        )

        assert response.errors
        assert "cannot compare string and boolean" in response.errors[0].message

    @pytest.mark.parametrize(
        "expectation,condition",
        [
            (True, "num_traces >= 5"),
            (False, "num_traces >= "),
            (False, "nonexistent_field == 1"),
        ],
    )
    async def test_validate_session_filter_condition(
        self,
        condition: str,
        expectation: bool,
        gql_client: AsyncGraphQLClient,
        db: DbSessionFactory,
    ) -> None:
        async with db() as session:
            project = models.Project(name=token_hex(8))
            session.add(project)
        query = """
            query($id: ID!, $condition: String!) {
              node(id: $id) {
                ... on Project {
                  validateSessionFilterCondition(condition: $condition) {
                    isValid
                    errorMessage
                  }
                }
              }
            }
        """
        project_gid = str(GlobalID(type_name="Project", node_id=str(project.id)))
        response = await gql_client.execute(
            query=query,
            variables={"id": project_gid, "condition": condition},
        )
        assert not response.errors
        assert (data := response.data) is not None
        result = data["node"]["validateSessionFilterCondition"]
        assert result["isValid"] == expectation
        if not expectation:
            assert result["errorMessage"]

    async def test_session_filter_vocabulary(
        self,
        gql_client: AsyncGraphQLClient,
        db: DbSessionFactory,
    ) -> None:
        from phoenix.trace.dsl.session_filter import SESSION_BINDINGS

        base_time = datetime.fromisoformat("2024-01-01T00:00:00+00:00")
        async with db() as session:
            project = await _add_project(session)
            project_session = await _add_project_session(session, project, start_time=base_time)
            trace = await _add_trace(
                session,
                project,
                project_session,
                start_time=base_time,
                end_time=base_time + timedelta(seconds=10),
            )
            root_span = await _add_span(
                session,
                trace,
                attributes={
                    "custom": {"leaf": "alpha"},
                    "duplicate": {"leaf": "first"},
                    "input": {"value": "hello"},
                    "metadata": {"tenant": "acme"},
                    "user": {"id": "user-1"},
                },
                start_time=base_time,
                end_time=base_time + timedelta(seconds=10),
            )
            child_tool_span = await _add_span(
                session,
                parent_span=root_span,
                span_kind="TOOL",
                attributes={"child_only": "not cataloged"},
                start_time=base_time + timedelta(seconds=1),
                end_time=base_time + timedelta(seconds=2),
            )
            child_tool_span.name = "search"
            later_trace = await _add_trace(
                session,
                project,
                project_session,
                start_time=base_time + timedelta(minutes=1),
                end_time=base_time + timedelta(minutes=1, seconds=10),
            )
            llm_span = await _add_span(
                session,
                later_trace,
                attributes={"later_only": "not cataloged"},
                start_time=base_time + timedelta(minutes=1),
                end_time=base_time + timedelta(minutes=1, seconds=10),
            )
            llm_span.name = "llm_not_cataloged"
            second_project_session = await _add_project_session(
                session, project, start_time=base_time + timedelta(minutes=2)
            )
            second_trace = await _add_trace(
                session,
                project,
                second_project_session,
                start_time=base_time + timedelta(minutes=2),
                end_time=base_time + timedelta(minutes=2, seconds=10),
            )
            second_tool_span = await _add_span(
                session,
                second_trace,
                span_kind="TOOL",
                attributes={
                    "duplicate": {"leaf": "second"},
                    "flat": "cataloged",
                    # Whole-literal storage of the key the first session stores
                    # nested — the same OTel wire key, so one vocabulary term.
                    "custom.leaf": "whole-literal",
                },
                start_time=base_time + timedelta(minutes=2),
                end_time=base_time + timedelta(minutes=2, seconds=10),
            )
            second_tool_span.name = "lookup"
            other_project = await _add_project(session)
            other_project_session = await _add_project_session(session, other_project)
            other_trace = await _add_trace(session, other_project, other_project_session)
            other_tool_span = await _add_span(
                session,
                other_trace,
                span_kind="TOOL",
                attributes={"other_project": "not cataloged"},
            )
            other_tool_span.name = "other_tool"
            session.add(
                models.ProjectSessionAnnotation(
                    project_session_id=project_session.id,
                    name="quality",
                    label="good",
                    score=1.0,
                    explanation=None,
                    metadata_={},
                    annotator_kind="HUMAN",
                    source="APP",
                )
            )
            await session.flush()
        query = """
            query($id: ID!) {
              node(id: $id) {
                ... on Project {
                  sessionFilterVocabulary {
                    name
                    type
                    description
                    category
                    iterableName
                  }
                }
              }
            }
        """
        project_gid = str(GlobalID(type_name="Project", node_id=str(project.id)))
        response = await gql_client.execute(query=query, variables={"id": project_gid})
        assert not response.errors
        assert (data := response.data) is not None
        terms = data["node"]["sessionFilterVocabulary"]
        # Element-field names repeat across iterables and collide with top-level ones, so the
        # top-level view keys only on unscoped terms; scoped terms are keyed by (iterable, name).
        top_level_terms = [t for t in terms if t["iterableName"] is None]
        term_names = {t["name"] for t in top_level_terms}
        terms_by_name = {term["name"]: term for term in top_level_terms}
        # Drift-proof: static session + aggregate terms are exactly the compiler's binding keys.
        static_names = {
            t["name"] for t in top_level_terms if t["name"] in SESSION_BINDINGS.binding_names
        }
        assert static_names == set(SESSION_BINDINGS.binding_names)
        assert {"any_input", "any_output"}.issubset(static_names)
        # The `in`-only IO terms are typed apart from plain strings so the UI doesn't
        # offer them as equality comparands.
        assert terms_by_name["any_input"]["type"] == "containment"
        assert terms_by_name["any_output"]["type"] == "containment"
        # Every served term carries a non-empty gloss.
        assert all(t["description"] for t in terms)
        # The turn model is an unenforced ingestion convention, so the served gloss leads
        # with the mechanic and labels the metaphor as the approximation it is.
        num_traces_term = next(t for t in terms if t["name"] == "num_traces")
        assert num_traces_term["description"].startswith("Number of traces in the session")
        assert "approximate conversation-turn count" in num_traces_term["description"]
        assert "does not enforce" in num_traces_term["description"]
        assert "earliest root span" in terms_by_name["attributes[...]"]["description"]
        assert 'attributes["user.id"]' in terms_by_name["user.id"]["description"]
        assert 'attributes["metadata.key"]' in terms_by_name['metadata["key"]']["description"]
        # Observed attributes are served in the canonical wire-key spelling,
        # regardless of how ingestion decomposed each key into nested JSON.
        observed_attribute_terms = {
            'attributes["custom.leaf"]',
            'attributes["duplicate.leaf"]',
            'attributes["flat"]',
            'attributes["input.value"]',
            'attributes["metadata.tenant"]',
            'attributes["user.id"]',
        }
        for term_name in observed_attribute_terms:
            assert terms_by_name[term_name]["type"] == "string"
            assert terms_by_name[term_name]["category"] == "attribute"
            assert "earliest root span" in terms_by_name[term_name]["description"]
        # One term per wire key: observed twice nested, and observed nested in
        # one session but whole-literal in another, both collapse.
        served_names = [term["name"] for term in terms]
        assert served_names.count('attributes["duplicate.leaf"]') == 1
        assert served_names.count('attributes["custom.leaf"]') == 1
        # The nested storage-path spelling stays an accepted compiler synonym
        # but is not served as its own term.
        assert 'attributes["custom"]["leaf"]' not in terms_by_name
        absent_attribute_terms = {
            'attributes["child_only"]',
            'attributes["later_only"]',
            'attributes["other_project"]',
        }
        assert absent_attribute_terms.isdisjoint(terms_by_name)
        # Per-project session-annotation names are folded in as fully-typed terms.
        assert 'annotations["quality"].score' in term_names
        assert 'annotations["quality"].label' in term_names
        # Span counts are session totals only — no per-tool-name terms are served.
        assert not any(name.startswith("tool_span_count[") for name in term_names)
        # The comprehension surface is served from the compiler's own catalog, so the iterables
        # and their element fields cannot drift from what actually compiles.
        iterable_terms = {t["name"]: t for t in terms if t["category"] == "iterable"}
        assert iterable_terms.keys() == set(SESSION_BINDINGS.iterables)
        assert all(
            t["type"] == "iterable" and t["iterableName"] is None for t in iterable_terms.values()
        )
        element_fields: dict[str, dict[str, str]] = {}
        for term in terms:
            if term["iterableName"] is not None:
                assert term["category"] == "element"
                element_fields.setdefault(term["iterableName"], {})[term["name"]] = term["type"]
        for iterable_name, grammar in SESSION_BINDINGS.iterables.items():
            element_bindings = grammar.element_bindings
            assert element_fields[iterable_name] == {
                **{n: "string" for n in element_bindings.string_names},
                **{n: "number" for n in element_bindings.float_names},
                **{n: "datetime" for n in element_bindings.datetime_names},
                **{n: "boolean" for n in element_bindings.boolean_names},
            }
        # A scoped name may repeat across iterables and shadow a top-level one; each is its
        # own term, distinguished only by iterableName.
        assert element_fields["spans"]["latency_ms"] == "number"
        assert element_fields["traces"]["latency_ms"] == "number"
        assert element_fields["traces"]["start_time"] == "datetime"
        assert "start_time" in term_names
        assert element_fields["span_cost_details"]["is_prompt"] == "boolean"

    async def test_session_count_shares_dsl_path_with_sessions_and_record_count(
        self,
        _data: _Data,
        httpx_client: httpx.AsyncClient,
    ) -> None:
        project = _data.projects[0]
        condition = "num_traces >= 2"
        # Only session index 2 has two traces; the rest have one.
        expected_session_gid = _gid(_data.project_sessions[2])

        sessions_result = await self._node(
            f'sessions(first:50,sessionFilterCondition:"{condition}"){{edges{{node{{id}}}}}}',
            project,
            httpx_client,
        )
        assert [e["node"]["id"] for e in sessions_result["edges"]] == [expected_session_gid]

        session_count = await self._node(
            f'sessionCount(sessionFilterCondition:"{condition}")', project, httpx_client
        )
        assert session_count == len(sessions_result["edges"]) == 1

        # A non-primary fan-out call site (record_counts) filtered by the same DSL condition:
        # session index 2 owns two spans, versus seven spans project-wide.
        assert await self._node("recordCount", project, httpx_client) == 7
        assert (
            await self._node(
                f'recordCount(sessionFilterCondition:"{condition}")', project, httpx_client
            )
            == 2
        )

    async def test_aside_statistics_follow_the_active_session_filter(
        self,
        db: DbSessionFactory,
        gql_client: AsyncGraphQLClient,
    ) -> None:
        """Every statistic the sessions aside renders is scoped to the filtered table's rows."""
        base_time = datetime.fromisoformat("2024-01-01T00:00:00+00:00")
        async with db() as session:
            project = await _add_project(session, name="aside-filtered-statistics")
            expected_gids = []
            # (traces, duration, label): two traces puts a session past the filter below.
            for trace_count, duration, label in (
                (2, timedelta(seconds=10), "good"),
                (2, timedelta(seconds=20), "bad"),
                (1, timedelta(seconds=100), "good"),
            ):
                project_session = await _add_project_session(
                    session, project, start_time=base_time, end_time=base_time + duration
                )
                for _ in range(trace_count):
                    await _add_trace(session, project, project_session, start_time=base_time)
                session.add(
                    models.ProjectSessionAnnotation(
                        project_session_id=project_session.id,
                        name="quality",
                        label=label,
                        score=1.0,
                        explanation=None,
                        metadata_={},
                        annotator_kind="HUMAN",
                        source="APP",
                    )
                )
                if trace_count >= 2:
                    expected_gids.append(_gid(project_session))
            await session.flush()

        query = """
          query ($id: ID!, $condition: String!) {
            node(id: $id) {
              ... on Project {
                sessions(first: 50, sessionFilterCondition: $condition) { edges { node { id } } }
                sessionCount(sessionFilterCondition: $condition)
                averageSessionDurationMs(sessionFilterCondition: $condition)
                averageTracesPerSession(sessionFilterCondition: $condition)
                sessionDurationMsQuantile(
                  probability: 0.5, sessionFilterCondition: $condition
                )
                sessionAnnotationSummary(
                  annotationName: "quality", sessionFilterCondition: $condition
                ) { count labelFractions { label fraction } }
              }
            }
          }
        """
        response = await gql_client.execute(
            query=query,
            variables={
                "id": str(GlobalID(type_name="Project", node_id=str(project.id))),
                "condition": "num_traces >= 2",
            },
        )
        assert not response.errors
        assert (data := response.data) is not None
        aside = data["node"]
        assert sorted(e["node"]["id"] for e in aside["sessions"]["edges"]) == sorted(expected_gids)
        assert aside["sessionCount"] == 2
        assert aside["averageSessionDurationMs"] == pytest.approx(15_000)
        assert aside["averageTracesPerSession"] == pytest.approx(2.0)
        assert 10_000 <= aside["sessionDurationMsQuantile"] <= 20_000
        summary = aside["sessionAnnotationSummary"]
        assert summary["count"] == 2
        assert {lf["label"]: lf["fraction"] for lf in summary["labelFractions"]} == {
            "good": pytest.approx(0.5),
            "bad": pytest.approx(0.5),
        }

    async def test_sessions_lookup_by_exact_session_id_via_dsl(
        self,
        _data: _Data,
        httpx_client: httpx.AsyncClient,
    ) -> None:
        """`session_id == '...'` is how the DSL expresses a single-session lookup."""
        project = _data.projects[0]
        session = _data.project_sessions[0]

        async def _list_ids(condition: str) -> list[str]:
            result = await self._node(
                f"sessions(first:50,sessionFilterCondition:{json.dumps(condition)}){{edges{{node{{id}}}}}}",
                project,
                httpx_client,
            )
            return [e["node"]["id"] for e in result["edges"]]

        assert await _list_ids(f"session_id == {session.session_id!r}") == [_gid(session)]
        assert await _list_ids(f"session_id == {token_hex(16)!r}") == []

    async def test_sessions_filtered_by_annotation_return_each_session_once(
        self,
        db: DbSessionFactory,
        httpx_client: httpx.AsyncClient,
    ) -> None:
        """A session carrying several annotations under one name is one row on a filtered page.

        Session annotations are unique on `(name, project_session_id, identifier)`, so the join a
        filter on `annotations[...]` makes can match a session more than once.
        """
        async with db() as session:
            project = await _add_project(session, name="annotation-fan-out-page")
            project_session = await _add_project_session(session, project)
            trace = await _add_trace(session, project, project_session)
            await _add_span(session, trace)
            for identifier in ("first", "second"):
                session.add(
                    models.ProjectSessionAnnotation(
                        project_session_id=project_session.id,
                        name="Quality",
                        label="good",
                        score=0.9,
                        metadata_={},
                        annotator_kind="HUMAN",
                        source="APP",
                        identifier=identifier,
                    )
                )

        condition = 'annotations["Quality"].score > 0.5'
        page = await self._node(
            f"sessions(first:50,sessionFilterCondition:{json.dumps(condition)}){{edges{{node{{id}}}}}}",
            project,
            httpx_client,
        )
        assert [e["node"]["id"] for e in page["edges"]] == [_gid(project_session)]

    @pytest.mark.parametrize("sort_column", ("tokenCountTotal", "costTotal"))
    async def test_sessions_sorted_page_and_count_agree_over_comprehension_condition(
        self,
        db: DbSessionFactory,
        httpx_client: httpx.AsyncClient,
        sort_column: str,
    ) -> None:
        """The page and the count select the same sessions, though they compile differently."""
        async with db() as session:
            project = await _add_project(session, name="comprehension-page-count-agreement")
            expected = []
            for span_count in (1, 3, 4):
                project_session = await _add_project_session(session, project)
                trace = await _add_trace(session, project, project_session)
                for _ in range(span_count):
                    await _add_span(session, trace, span_kind="CHAIN")
                if span_count > 2:
                    expected.append(_gid(project_session))

        condition = "len([s for s in spans]) > 2"
        page = await self._node(
            f"sessions(first:50,sort:{{col:{sort_column},dir:desc}},"
            f"sessionFilterCondition:{json.dumps(condition)}){{edges{{node{{id}}}}}}",
            project,
            httpx_client,
        )
        count = await self._node(
            f"sessionCount(sessionFilterCondition:{json.dumps(condition)})", project, httpx_client
        )
        assert sorted(e["node"]["id"] for e in page["edges"]) == sorted(expected)
        assert count == len(expected)

    @staticmethod
    async def _validate_session_filter(
        project: models.Project,
        condition: str,
        gql_client: AsyncGraphQLClient,
    ) -> dict[str, Any]:
        query = """
            query($id: ID!, $condition: String!) {
              node(id: $id) {
                ... on Project {
                  validateSessionFilterCondition(condition: $condition) {
                    isValid
                    errorMessage
                    warnings
                  }
                }
              }
            }
        """
        project_gid = str(GlobalID(type_name="Project", node_id=str(project.id)))
        response = await gql_client.execute(
            query=query,
            variables={"id": project_gid, "condition": condition},
        )
        assert not response.errors
        assert (data := response.data) is not None
        result: dict[str, Any] = data["node"]["validateSessionFilterCondition"]
        return result

    async def test_validate_session_filter_condition_warns_on_unknown_dynamic_names(
        self,
        gql_client: AsyncGraphQLClient,
        db: DbSessionFactory,
    ) -> None:
        """Typo'd names in string-literal subscripts compile but warn against observed names."""
        base_time = datetime.fromisoformat("2024-01-01T00:00:00+00:00")
        async with db() as session:
            project = await _add_project(session)
            project_session = await _add_project_session(session, project, start_time=base_time)
            trace = await _add_trace(
                session,
                project,
                project_session,
                start_time=base_time,
                end_time=base_time + timedelta(seconds=10),
            )
            tool_span = await _add_span(
                session,
                trace,
                span_kind="TOOL",
                start_time=base_time,
                end_time=base_time + timedelta(seconds=1),
            )
            tool_span.name = "search"
            session.add(
                models.ProjectSessionAnnotation(
                    project_session_id=project_session.id,
                    name="quality",
                    label="good",
                    score=1.0,
                    explanation=None,
                    metadata_={},
                    annotator_kind="HUMAN",
                    source="APP",
                )
            )
            await session.flush()

        # Typo'd annotation name: valid (compiles) but warns, naming the unknown + the observed set.
        result = await self._validate_session_filter(
            project, "annotations['typoo'].score > 0.5", gql_client
        )
        assert result["isValid"] is True
        assert result["errorMessage"] is None
        assert any(
            w.startswith("unknown annotation name 'typoo'") and "quality" in w
            for w in result["warnings"]
        )

        # An observed annotation name produces no warnings.
        result = await self._validate_session_filter(
            project,
            "annotations['quality'].score > 0.5 and tool_span_count > 0",
            gql_client,
        )
        assert result["isValid"] is True
        assert result["warnings"] == []

    async def test_session_filter_vocabulary_bounds_attribute_scan_to_recent_sessions(
        self,
        gql_client: AsyncGraphQLClient,
        db: DbSessionFactory,
        monkeypatch: pytest.MonkeyPatch,
    ) -> None:
        """Byte truncation preserves the newest session's root-span attribute names."""
        from phoenix.server.api.types import Project as project_module

        monkeypatch.setattr(project_module, "_VOCABULARY_ATTRIBUTE_SCAN_LIMIT", 2)
        monkeypatch.setattr(project_module, "_VOCABULARY_ATTRIBUTE_SCAN_BYTE_LIMIT", 64)
        base_time = datetime.fromisoformat("2024-01-01T00:00:00+00:00")
        async with db() as session:
            project = await _add_project(session)
            older_session = await _add_project_session(session, project, start_time=base_time)
            older_trace = await _add_trace(
                session,
                project,
                older_session,
                start_time=base_time,
                end_time=base_time + timedelta(seconds=10),
            )
            await _add_span(
                session,
                older_trace,
                attributes={"old_only": "x" * 128},
                start_time=base_time,
                end_time=base_time + timedelta(seconds=10),
            )
            newer_session = await _add_project_session(
                session, project, start_time=base_time + timedelta(minutes=5)
            )
            newer_trace = await _add_trace(
                session,
                project,
                newer_session,
                start_time=base_time + timedelta(minutes=5),
                end_time=base_time + timedelta(minutes=5, seconds=10),
            )
            await _add_span(
                session,
                newer_trace,
                attributes={"new_only": "y" * 128},
                start_time=base_time + timedelta(minutes=5),
                end_time=base_time + timedelta(minutes=5, seconds=10),
            )
            await session.flush()
        query = """
            query($id: ID!) {
              node(id: $id) {
                ... on Project {
                  sessionFilterVocabulary { name }
                }
              }
            }
        """
        project_gid = str(GlobalID(type_name="Project", node_id=str(project.id)))
        response = await gql_client.execute(query=query, variables={"id": project_gid})
        assert not response.errors
        assert (data := response.data) is not None
        term_names = {t["name"] for t in data["node"]["sessionFilterVocabulary"]}
        assert 'attributes["new_only"]' in term_names
        assert 'attributes["old_only"]' not in term_names


@pytest.mark.parametrize(
    "sort_col, sort_dir, expected_order",
    [
        pytest.param(
            "name",
            "asc",
            ["project1", "project2", "project3"],
            id="sort-by-name-asc",
        ),
        pytest.param(
            "name",
            "desc",
            ["project3", "project2", "project1"],
            id="sort-by-name-desc",
        ),
        pytest.param(
            "endTime",
            "asc",
            ["project1", "project2", "project3"],
            id="sort-by-end-time-asc",
        ),
        pytest.param(
            "endTime",
            "desc",
            ["project3", "project2", "project1"],
            id="sort-by-end-time-desc",
        ),
    ],
)
async def test_project_sort(
    sort_col: str,
    sort_dir: str,
    expected_order: list[str],
    gql_client: AsyncGraphQLClient,
    db: DbSessionFactory,
) -> None:
    """Test project sorting capabilities."""
    # Create test projects with controlled timestamps
    base_time = datetime.fromisoformat("2024-01-01T00:00:00+00:00")
    projects: list[models.Project] = []
    async with db() as session:
        # Create projects first
        for i, name in enumerate(["project1", "project2", "project3"]):
            project = models.Project(
                name=name,
                created_at=base_time + timedelta(hours=i),
                updated_at=base_time + timedelta(hours=i),
            )
            session.add(project)
            await session.flush()
            projects.append(project)

        # Now create traces for each project with different end times
        # Each project will have 3 traces with different end times
        # The max end time for each project will be different and match the expected sort order
        for i, project in enumerate(projects):
            for j in range(3):
                trace = models.Trace(
                    trace_id=token_hex(16),
                    project_rowid=project.id,
                    start_time=base_time + timedelta(hours=i),
                    # The max end time for each project will be base_time + (i+1) days
                    # This ensures project1 has max end time of day 1
                    # project2 has max end time of day 2
                    # project3 has max end time of day 3
                    end_time=base_time + timedelta(days=i + 1, hours=j),
                )
                session.add(trace)
        await session.commit()

    query = """
        query ($sort: ProjectSort) {
            projects(sort: $sort) {
                edges {
                    node {
                        name
                    }
                }
            }
        }
    """

    variables = {"sort": {"col": sort_col, "dir": sort_dir}}
    response = await gql_client.execute(query=query, variables=variables)
    assert not response.errors
    assert (data := response.data) is not None
    projects = data["projects"]
    project_names = [edge["node"]["name"] for edge in projects["edges"]]  # type: ignore
    assert project_names == expected_order


@pytest.mark.parametrize(
    "filter_value, expected_names",
    [
        pytest.param(
            "project",
            ["test_project", "project_test"],
            id="filter-matches-all",
        ),
        pytest.param(
            "test",
            ["test_project", "project_test"],
            id="filter-matches-partial",
        ),
        pytest.param(
            "TEST",
            ["test_project", "project_test"],
            id="filter-case-insensitive",
        ),
        pytest.param(
            "nomatch",
            [],
            id="filter-no-matches",
        ),
    ],
)
async def test_project_filter(
    filter_value: str,
    expected_names: list[str],
    gql_client: AsyncGraphQLClient,
    db: DbSessionFactory,
) -> None:
    """Test project filtering capabilities."""
    # Create test projects
    async with db() as session:
        for name in ["test_project", "project_test", "other_name"]:
            project = models.Project(name=name)
            session.add(project)
        await session.commit()

    query = """
        query ($filter: ProjectFilter) {
            projects(filter: $filter) {
                edges {
                    node {
                        name
                    }
                }
            }
        }
    """

    variables = {"filter": {"col": "name", "value": filter_value}}
    response = await gql_client.execute(query=query, variables=variables)
    assert not response.errors
    assert (data := response.data) is not None
    projects = data["projects"]
    project_names = [edge["node"]["name"] for edge in projects["edges"]]
    assert sorted(project_names) == sorted(expected_names)


@pytest.mark.parametrize(
    "sort, filter_value, expected_names",
    [
        pytest.param(
            {"col": "name", "dir": "asc"},
            "test",
            ["project_test", "test_project"],
            id="filter-and-sort-asc",
        ),
        pytest.param(
            {"col": "name", "dir": "desc"},
            "test",
            ["test_project", "project_test"],
            id="filter-and-sort-desc",
        ),
    ],
)
async def test_project_filter_and_sort(
    sort: dict[str, str],
    filter_value: str,
    expected_names: list[str],
    gql_client: AsyncGraphQLClient,
    db: DbSessionFactory,
) -> None:
    """Test combining project filtering and sorting."""
    # Create test projects
    async with db() as session:
        for name in ["test_project", "project_test", "other_name"]:
            project = models.Project(name=name)
            session.add(project)
        await session.commit()

    query = """
        query ($sort: ProjectSort, $filter: ProjectFilter) {
            projects(sort: $sort, filter: $filter) {
                edges {
                    node {
                        name
                    }
                }
            }
        }
    """

    variables = {
        "sort": sort,
        "filter": {"col": "name", "value": filter_value},
    }
    response = await gql_client.execute(query=query, variables=variables)
    assert not response.errors
    assert (data := response.data) is not None
    projects = data["projects"]
    project_names = [edge["node"]["name"] for edge in projects["edges"]]
    assert project_names == expected_names


async def test_paginate_spans_by_trace_start_time(
    db: DbSessionFactory,
    gql_client: AsyncGraphQLClient,
) -> None:
    """Test the _paginate_span_by_trace_start_time optimization function.

    This function is triggered when:
    - rootSpansOnly: true
    - No filter_condition
    - sort.col is SpanColumn.startTime

    Key behaviors tested:
    - Returns one representative span per trace (not all spans)
    - Orders by trace start time (not span start time)
    - Uses cursors based on trace rowids + start times (unusual!)
    - Handles orphan spans based on orphan_span_as_root_span parameter
    - Supports time range filtering on trace start times
    - May return empty edges while has_next_page=True when traces have no matching spans
    - **RETRY LOGIC**: When insufficient edges are found (len(edges) < first) but has_next_page=True,
      the function automatically retries pagination with larger batch sizes (max(first, 1000))
      up to 10 times (retries=10) to collect enough spans. This handles cases where many traces
      exist but lack matching root spans.

    Implementation Details:
    - Uses CTEs (Common Table Expressions) for efficient trace-based pagination
    - PostgreSQL: Uses DISTINCT ON for deduplication
    - SQLite: Uses Python groupby() for deduplication (too complex for SQLite DISTINCT)
    - SQL ordering: trace start_time -> trace id -> span start_time (ASC for earliest) -> span id (DESC)
    - Cursors contain trace rowid + trace start_time, NOT span data
    - Over-fetches by 1 trace to determine has_next_page efficiently

    Test Data Setup:
    ================
    Creates 5 traces with start times at hours 1, 2, 3, 4, 5:

    Trace Index | Hour | Real Root Span | Orphan Span  | Additional Spans | Expected Name
    ------------|------|----------------|--------------|------------------|---------------
    0 (even)    |  1   |      ✓         |      ✗       | +2nd root span   | root-span-1
    1 (odd)     |  2   |      ✗         |      ✓       | +2nd orphan span | orphan-span-2
    2 (even)    |  3   |      ✓         |      ✗       | +2nd root span   | root-span-3
    3 (odd)     |  4   |      ✗         |      ✓       | +2nd orphan span | orphan-span-4
    4 (even)    |  5   |      ✓         |      ✗       | +2nd root span   | root-span-5

    Key Testing Points:
    - ALL traces have multiple candidate spans to test "earliest span per trace" selection
    - Trace 1: 2 root spans → Returns earliest (root-span-1, not second-root-span-1)
    - Trace 2: 2 orphan spans → Returns earliest (orphan-span-2, not second-orphan-span-2)
    - Trace 3: 2 root spans → Returns earliest (root-span-3, not second-root-span-3)
    - Trace 4: 2 orphan spans → Returns earliest (orphan-span-4, not second-orphan-span-4)
    - Trace 5: 2 root spans → Returns earliest (root-span-5, not second-root-span-5)
    - Comprehensive test of SQL ordering: ORDER BY span.start_time ASC, span.id DESC

    With orphan_span_as_root_span=false: Only returns real root spans 1, 3, 5 (3 total)
    With orphan_span_as_root_span=true:  Returns all spans 1, 2, 3, 4, 5 (5 total)
    """
    # ========================================
    # SETUP: Create test data
    # ========================================
    async with db() as session:
        project = models.Project(name=token_hex(8))
        session.add(project)
        await session.flush()

        # Create 5 traces with start times at hours 1, 2, 3, 4, 5
        base_time = datetime.fromisoformat("2024-01-01T00:00:00+00:00")
        traces = []
        spans = []

        for i in range(5):
            # Trace start times: 01:00, 02:00, 03:00, 04:00, 05:00
            trace = models.Trace(
                trace_id=token_hex(16),
                project_rowid=project.id,
                start_time=base_time + timedelta(hours=i + 1),
                end_time=base_time + timedelta(hours=i + 2),
            )
            session.add(trace)
            await session.flush()
            traces.append(trace)

            if i % 2 == 0:
                # EVEN indices (0, 2, 4) → traces at hours 1, 3, 5 → CREATE REAL ROOT SPANS
                # These spans have parent_id=None (true root spans)
                root_span = models.Span(
                    trace_rowid=trace.id,
                    span_id=token_hex(8),
                    parent_id=None,  # ← This makes it a real root span
                    name=f"root-span-{i + 1}",
                    span_kind="CHAIN",
                    start_time=trace.start_time + timedelta(minutes=10),
                    end_time=trace.start_time + timedelta(minutes=20),
                    attributes={},
                    events=[],
                    status_code="OK",
                    status_message="",
                    cumulative_error_count=0,
                    cumulative_llm_token_count_prompt=0,
                    cumulative_llm_token_count_completion=0,
                )
                session.add(root_span)
                spans.append(root_span)

                # Also create a child span to verify only root span is returned per trace
                child_span = models.Span(
                    trace_rowid=trace.id,
                    span_id=token_hex(8),
                    parent_id=root_span.span_id,  # ← Child of the root span
                    name=f"child-span-{i + 1}",
                    span_kind="CHAIN",
                    start_time=trace.start_time + timedelta(minutes=15),
                    end_time=trace.start_time + timedelta(minutes=25),
                    attributes={},
                    events=[],
                    status_code="OK",
                    status_message="",
                    cumulative_error_count=0,
                    cumulative_llm_token_count_prompt=0,
                    cumulative_llm_token_count_completion=0,
                )
                session.add(child_span)

                # Add a SECOND root span with later start time to test "earliest span" selection
                # This span should NOT be returned (only the earliest root span per trace)
                second_root_span = models.Span(
                    trace_rowid=trace.id,
                    span_id=token_hex(8),
                    parent_id=None,  # ← Also a root span
                    name=f"second-root-span-{i + 1}",
                    span_kind="CHAIN",
                    start_time=trace.start_time + timedelta(minutes=30),  # ← Later start time
                    end_time=trace.start_time + timedelta(minutes=40),
                    attributes={},
                    events=[],
                    status_code="OK",
                    status_message="",
                    cumulative_error_count=0,
                    cumulative_llm_token_count_prompt=0,
                    cumulative_llm_token_count_completion=0,
                )
                session.add(second_root_span)
            else:
                # ODD indices (1, 3) → traces at hours 2, 4 → CREATE ORPHAN SPANS
                # These spans have parent_id pointing to non-existent spans (orphans)
                orphan_span = models.Span(
                    trace_rowid=trace.id,
                    span_id=token_hex(8),
                    parent_id=token_hex(8),  # ← Points to non-existent span (orphan)
                    name=f"orphan-span-{i + 1}",
                    span_kind="CHAIN",
                    start_time=trace.start_time + timedelta(minutes=10),
                    end_time=trace.start_time + timedelta(minutes=20),
                    attributes={},
                    events=[],
                    status_code="OK",
                    status_message="",
                    cumulative_error_count=0,
                    cumulative_llm_token_count_prompt=0,
                    cumulative_llm_token_count_completion=0,
                )
                session.add(orphan_span)
                spans.append(orphan_span)

                # Add a SECOND orphan span with later start time to test "earliest span" selection
                # This span should NOT be returned (only the earliest orphan span per trace)
                second_orphan_span = models.Span(
                    trace_rowid=trace.id,
                    span_id=token_hex(8),
                    parent_id=token_hex(8),  # ← Also an orphan span (different parent_id)
                    name=f"second-orphan-span-{i + 1}",
                    span_kind="CHAIN",
                    start_time=trace.start_time + timedelta(minutes=30),  # ← Later start time
                    end_time=trace.start_time + timedelta(minutes=40),
                    attributes={},
                    events=[],
                    status_code="OK",
                    status_message="",
                    cumulative_error_count=0,
                    cumulative_llm_token_count_prompt=0,
                    cumulative_llm_token_count_completion=0,
                )
                session.add(second_orphan_span)

        project_gid = str(GlobalID(type_name="Project", node_id=str(project.id)))

    # ========================================
    # TEST 1: Basic pagination with orphan_span_as_root_span=false
    # Expected: Only real root spans (1, 3, 5) returned, NOT orphan spans (2, 4)
    # ========================================
    query = """
        query ($projectId: ID!, $first: Int!, $after: String) {
            node(id: $projectId) {
                ... on Project {
                    spans(
                        rootSpansOnly: true,
                        orphanSpanAsRootSpan: false,  # ← Exclude orphan spans
                        sort: {col: startTime, dir: desc},
                        first: $first,
                        after: $after
                    ) {
                        edges {
                            node {
                                id
                                name
                            }
                            cursor
                        }
                        pageInfo {
                            hasNextPage
                            hasPreviousPage
                            startCursor
                            endCursor
                        }
                    }
                }
            }
        }
    """

    # Page 1: Request first 2 spans in descending order (by trace start time)
    # Expected: Only root-span-5 (trace 5 is latest, and only that trace has a real root span)
    # Note: trace 4 has an orphan span, but it's excluded by orphanSpanAsRootSpan=false
    response = await gql_client.execute(
        query=query,
        variables={
            "projectId": project_gid,
            "first": 2,
        },
    )

    assert not response.errors
    assert (data := response.data) is not None

    page = data["node"]["spans"]
    edges = page["edges"]
    page_info = page["pageInfo"]

    assert len(edges) == 2
    assert edges[0]["node"]["name"] == "root-span-5"
    assert edges[1]["node"]["name"] == "root-span-3"
    assert page_info["hasNextPage"] is True  # More traces to check
    assert page_info["hasPreviousPage"] is False

    # Verify cursor contains trace rowid (5) and trace start time (05:00:00)
    # This demonstrates the unusual "trace-based cursors" behavior
    assert (
        base64.b64decode(page_info["startCursor"].encode())
        == b"5:DATETIME:2024-01-01T05:00:00+00:00"
    )
    assert (
        base64.b64decode(page_info["endCursor"].encode()) == b"3:DATETIME:2024-01-01T03:00:00+00:00"
    )

    # Page 2: Continue pagination after trace 5
    # Expected: root-span-3 (trace 3 is next latest with real root span)
    # Note: trace 4 is skipped because it only has orphan span (excluded)
    response = await gql_client.execute(
        query=query,
        variables={
            "projectId": project_gid,
            "first": 3,
            "after": base64.b64encode(b"5:DATETIME:2024-01-01T05:00:00+00:00").decode(),
        },
    )

    assert not response.errors
    assert (data := response.data) is not None

    page = data["node"]["spans"]
    edges = page["edges"]
    page_info = page["pageInfo"]

    assert len(edges) == 2
    assert edges[0]["node"]["name"] == "root-span-3"
    assert edges[1]["node"]["name"] == "root-span-1"
    assert page_info["hasNextPage"] is False
    assert page_info["hasPreviousPage"] is False
    assert (
        base64.b64decode(page_info["startCursor"].encode())
        == b"4:DATETIME:2024-01-01T04:00:00+00:00"  # Trace 3 yielded the span
    )
    assert (
        base64.b64decode(page_info["endCursor"].encode()) == b"1:DATETIME:2024-01-01T01:00:00+00:00"
    )

    # Page 3: Continue pagination after trace 3
    # Expected: root-span-1 (trace 1 is oldest with real root span)
    # Note: trace 2 is skipped because it only has orphan span (excluded)
    response = await gql_client.execute(
        query=query,
        variables={
            "projectId": project_gid,
            "first": 4,
            "after": base64.b64encode(b"3:DATETIME:2024-01-01T03:00:00+00:00").decode(),
        },
    )

    assert not response.errors
    assert (data := response.data) is not None

    page = data["node"]["spans"]
    edges = page["edges"]
    page_info = page["pageInfo"]

    # Should return root-span-1 (oldest real root span)
    assert len(edges) == 1
    assert edges[0]["node"]["name"] == "root-span-1"
    assert page_info["hasNextPage"] is False  # No more traces
    assert page_info["hasPreviousPage"] is False
    assert (
        base64.b64decode(page_info["startCursor"].encode())
        == b"2:DATETIME:2024-01-01T02:00:00+00:00"
    )
    assert (
        base64.b64decode(page_info["endCursor"].encode()) == b"1:DATETIME:2024-01-01T01:00:00+00:00"
    )

    # ========================================
    # TEST 2: Ascending order (orphan_span_as_root_span=false)
    # Expected: Same spans but in reverse order: root-span-1, root-span-3, root-span-5
    # ========================================
    response = await gql_client.execute(
        query=query.replace("dir: desc", "dir: asc"),
        variables={
            "projectId": project_gid,
            "first": 2,
        },
    )

    assert not response.errors
    assert (data := response.data) is not None

    asc_page = data["node"]["spans"]
    edges = asc_page["edges"]
    page_info = asc_page["pageInfo"]

    # Should return first span in ascending order (oldest trace with real root span)
    assert len(edges) == 2
    assert edges[0]["node"]["name"] == "root-span-1"
    assert edges[1]["node"]["name"] == "root-span-3"
    assert page_info["hasNextPage"] is True

    # ========================================
    # TEST 3: Bulk query (orphan_span_as_root_span=false)
    # Expected: All 3 real root spans at once
    # ========================================
    response = await gql_client.execute(
        query=query,
        variables={
            "projectId": project_gid,
            "first": 10,
        },
    )

    assert not response.errors
    assert (data := response.data) is not None

    all_spans = data["node"]["spans"]
    edges = all_spans["edges"]
    page_info = all_spans["pageInfo"]
    span_names = [edge["node"]["name"] for edge in edges]

    # Should return all 3 real root spans (excluding orphan spans 2, 4)
    # IMPORTANT: Returns earliest root span per trace (ALL traces have multiple candidates):
    # - Trace 1: root-span-1 (NOT second-root-span-1 which has later start time)
    # - Trace 3: root-span-3 (NOT second-root-span-3 which has later start time)
    # - Trace 5: root-span-5 (NOT second-root-span-5 which has later start time)
    assert len(edges) == 3
    assert span_names == [
        "root-span-5",
        "root-span-3",
        "root-span-1",
    ]
    assert page_info["hasNextPage"] is False

    # ========================================
    # TEST 4: Time range filtering (orphan_span_as_root_span=false)
    # Filter: hours 2-4 (includes traces 2, 3, 4)
    # Expected: Only root-span-3 (trace 3 has real root span, traces 2&4 have orphans)
    # ========================================
    time_range_query = """
        query ($projectId: ID!, $first: Int!, $timeRange: TimeRange) {
            node(id: $projectId) {
                ... on Project {
                    spans(
                        rootSpansOnly: true,
                        orphanSpanAsRootSpan: false,  # ← Exclude orphan spans
                        sort: {col: startTime, dir: desc},
                        first: $first,
                        timeRange: $timeRange
                    ) {
                        edges {
                            node {
                                id
                                name
                            }
                        }
                        pageInfo {
                            hasNextPage
                        }
                    }
                }
            }
        }
    """

    response = await gql_client.execute(
        query=time_range_query,
        variables={
            "projectId": project_gid,
            "first": 10,
            "timeRange": {
                "start": (base_time + timedelta(hours=2)).isoformat(),  # 02:00:00
                "end": (base_time + timedelta(hours=4)).isoformat(),  # 04:00:00
            },
        },
    )

    assert not response.errors
    assert (data := response.data) is not None

    filtered_spans = data["node"]["spans"]
    edges = filtered_spans["edges"]

    # Time range includes traces 2, 3, 4:
    # - Trace 2 (hour 2): has orphan span → excluded by orphanSpanAsRootSpan=false
    # - Trace 3 (hour 3): has real root span → included
    # - Trace 4 (hour 4): has orphan span → excluded by orphanSpanAsRootSpan=false
    assert len(edges) == 1
    assert edges[0]["node"]["name"] == "root-span-3"

    # ========================================
    # TEST 5: Include orphan spans (orphanSpanAsRootSpan=true)
    # Expected: All 5 spans returned (3 real roots + 2 orphans)
    # ========================================
    orphan_query = """
        query ($projectId: ID!, $first: Int!, $after: String) {
            node(id: $projectId) {
                ... on Project {
                    spans(
                        rootSpansOnly: true,
                        orphanSpanAsRootSpan: true,
                        sort: {col: startTime, dir: desc},
                        first: $first,
                        after: $after
                    ) {
                        edges {
                            node {
                                id
                                name
                            }
                            cursor
                        }
                        pageInfo {
                            hasNextPage
                            hasPreviousPage
                            startCursor
                            endCursor
                        }
                    }
                }
            }
        }
    """

    # Test 5a: Basic pagination with orphans included
    # Expected: Now returns 2 spans per page instead of 1 (includes orphan spans)
    response = await gql_client.execute(
        query=orphan_query,
        variables={
            "projectId": project_gid,
            "first": 2,
        },
    )

    assert not response.errors
    assert (data := response.data) is not None

    page = data["node"]["spans"]
    edges = page["edges"]
    page_info = page["pageInfo"]

    # Should return 2 spans: both real root and orphan spans
    assert len(edges) == 2
    assert edges[0]["node"]["name"] == "root-span-5"  # Real root span from trace 5 (latest)
    assert edges[1]["node"]["name"] == "orphan-span-4"  # Orphan span from trace 4 (2nd latest)
    assert page_info["hasNextPage"] is True

    # Test 5b: Bulk query with orphans included
    # Expected: All 5 spans (3 real + 2 orphan) vs 3 spans when orphans excluded
    response = await gql_client.execute(
        query=orphan_query,
        variables={
            "projectId": project_gid,
            "first": 10,
        },
    )

    assert not response.errors
    assert (data := response.data) is not None

    all_spans = data["node"]["spans"]
    edges = all_spans["edges"]
    page_info = all_spans["pageInfo"]
    span_names = [edge["node"]["name"] for edge in edges]

    # Should return ALL 5 spans (3 real root spans + 2 orphan spans) in descending order
    # IMPORTANT: Returns earliest span per trace (ALL traces have multiple candidates):
    # - Trace 1: root-span-1 (NOT second-root-span-1)
    # - Trace 2: orphan-span-2 (NOT second-orphan-span-2)
    # - Trace 3: root-span-3 (NOT second-root-span-3)
    # - Trace 4: orphan-span-4 (NOT second-orphan-span-4)
    # - Trace 5: root-span-5 (NOT second-root-span-5)
    assert len(edges) == 5
    assert span_names == [
        "root-span-5",
        "orphan-span-4",
        "root-span-3",
        "orphan-span-2",
        "root-span-1",
    ]
    assert page_info["hasNextPage"] is False

    # Test 5c: Ascending order with orphans included
    # Expected: Same 5 spans but in reverse order
    response = await gql_client.execute(
        query=orphan_query.replace("dir: desc", "dir: asc"),
        variables={"projectId": project_gid, "first": 3},
    )

    assert not response.errors
    assert (data := response.data) is not None

    asc_page = data["node"]["spans"]
    edges = asc_page["edges"]
    span_names = [edge["node"]["name"] for edge in edges]

    # Should return first 3 spans in ascending order (includes orphan span 2)
    assert len(edges) == 3
    assert span_names == [
        "root-span-1",
        "orphan-span-2",
        "root-span-3",
    ]

    # Test 5d: Time range filtering with orphans included
    orphan_time_range_query = """
        query ($projectId: ID!, $first: Int!, $timeRange: TimeRange) {
            node(id: $projectId) {
                ... on Project {
                    spans(
                        rootSpansOnly: true,
                        orphanSpanAsRootSpan: true,
                        sort: {col: startTime, dir: desc},
                        first: $first,
                        timeRange: $timeRange
                    ) {
                        edges {
                            node {
                                id
                                name
                            }
                        }
                        pageInfo {
                            hasNextPage
                        }
                    }
                }
            }
        }
    """

    # Expected: Now returns 2 spans (includes orphan span 2) vs 1 span when orphans excluded
    response = await gql_client.execute(
        query=orphan_time_range_query,
        variables={
            "projectId": project_gid,
            "first": 10,
            "timeRange": {
                "start": (base_time + timedelta(hours=2)).isoformat(),  # 02:00:00
                "end": (base_time + timedelta(hours=4)).isoformat(),  # 04:00:00
            },
        },
    )

    assert not response.errors
    assert (data := response.data) is not None

    filtered_spans = data["node"]["spans"]
    edges = filtered_spans["edges"]
    span_names = [edge["node"]["name"] for edge in edges]

    # Time range includes traces 2, 3, 4 - with orphans included:
    # - Trace 2 (hour 2): has orphan span → NOW INCLUDED
    # - Trace 3 (hour 3): has real root span → included
    # - Trace 4 (hour 4): has orphan span → NOW INCLUDED
    # But trace 4 is excluded by time range end=04:00:00 (exclusive), so only traces 2 & 3
    assert len(edges) == 2
    assert span_names == [
        "root-span-3",
        "orphan-span-2",
    ]  # Descending order: trace 3, then trace 2


async def test_cost_summary_returns_expected_results(
    gql_client: AsyncGraphQLClient,
    db: DbSessionFactory,
) -> None:
    async with db() as session:
        project = await _add_project(session, name="span-cost-filter-test")
        session1 = await _add_project_session(session, project)
        trace1 = await _add_trace(session, project, session1)
        llm_span_costing_100_dollars = await _add_span(
            session, trace1, span_kind="LLM", attributes={"input": {"value": "llm query"}}
        )
        model = await _add_generative_model(session, "test-model")
        await _add_span_cost(session, llm_span_costing_100_dollars, trace1, model, total_cost=100.0)

        session2 = await _add_project_session(session, project, session_id="cost-session-2")
        trace2 = await _add_trace(session, project, session2)
        chain_span_costing_50_dollars = await _add_span(
            session, trace2, span_kind="CHAIN", attributes={"input": {"value": "chain query"}}
        )
        await _add_span_cost(session, chain_span_costing_50_dollars, trace2, model, total_cost=50.0)

        await session.commit()

    query = """
      query ($projectId: ID!, $filterCondition: String, $sessionFilterCondition: String) {
        node(id: $projectId) {
          ... on Project {
            costSummary(
              filterCondition: $filterCondition
              sessionFilterCondition: $sessionFilterCondition
            ) {
              total {
                cost
              }
              prompt {
                cost
              }
              completion {
                cost
              }
            }
          }
        }
      }
    """

    project_gid = str(GlobalID(type_name="Project", node_id=str(project.id)))

    response_no_filter = await gql_client.execute(query=query, variables={"projectId": project_gid})
    assert not response_no_filter.errors
    assert response_no_filter.data is not None
    assert response_no_filter.data["node"]["costSummary"]["total"]["cost"] == 150.0

    response_llm_span_filter = await gql_client.execute(
        query=query, variables={"projectId": project_gid, "filterCondition": 'span_kind == "LLM"'}
    )
    assert not response_llm_span_filter.errors
    assert response_llm_span_filter.data is not None
    assert response_llm_span_filter.data["node"]["costSummary"]["total"]["cost"] == 100.0

    response_chain_session_filter = await gql_client.execute(
        query=query,
        variables={
            "projectId": project_gid,
            "sessionFilterCondition": "session_id == 'cost-session-2'",
        },
    )
    assert not response_chain_session_filter.errors
    assert response_chain_session_filter.data is not None
    assert response_chain_session_filter.data["node"]["costSummary"]["total"]["cost"] == 50.0


async def test_latency_quantile_with_filters_returns_accurate_percentiles(
    gql_client: AsyncGraphQLClient,
    db: DbSessionFactory,
) -> None:
    async with db() as session:
        project = await _add_project(session, name="latency-filter-test")
        llm_span_session = await _add_project_session(
            session, project, session_id="latency-llm-session"
        )
        llm_span_latencies_ms = [100, 200, 300]
        for latency_ms in llm_span_latencies_ms:
            start_time = datetime.now(timezone.utc)
            trace = await _add_trace(
                session,
                project,
                project_session=llm_span_session,
                start_time=start_time,
                end_time=start_time + timedelta(milliseconds=latency_ms),
            )
            await _add_span(
                session,
                trace,
                span_kind="LLM",
                start_time=trace.start_time,
                end_time=trace.end_time,
                attributes={"input": {"value": "llm span input"}},
            )

        chain_span_session = await _add_project_session(
            session, project, session_id="latency-chain-session"
        )
        chain_span_latencies_ms = [400, 500, 600]
        for latency_ms in chain_span_latencies_ms:
            start_time = datetime.now(timezone.utc)
            trace = await _add_trace(
                session,
                project,
                project_session=chain_span_session,
                start_time=start_time,
                end_time=start_time + timedelta(milliseconds=latency_ms),
            )
            await _add_span(
                session,
                trace,
                span_kind="CHAIN",
                start_time=trace.start_time,
                end_time=trace.end_time,
                attributes={"input": {"value": "chain span input"}},
            )

    query = """
      query ($projectId: ID!, $filterCondition: String, $sessionFilterCondition: String) {
        project: node(id: $projectId) {
          ... on Project {
            latencyMsQuantile(
              probability: 0.5
              filterCondition: $filterCondition
              sessionFilterCondition: $sessionFilterCondition
            )
          }
        }
      }
    """

    project_gid = str(GlobalID(type_name="Project", node_id=str(project.id)))

    response_llm_span_filter = await gql_client.execute(
        query=query, variables={"projectId": project_gid, "filterCondition": "span_kind == 'LLM'"}
    )
    assert not response_llm_span_filter.errors
    assert response_llm_span_filter.data is not None
    assert response_llm_span_filter.data["project"]["latencyMsQuantile"] == 200.0

    response_chain_span_filter = await gql_client.execute(
        query=query, variables={"projectId": project_gid, "filterCondition": "span_kind == 'CHAIN'"}
    )
    assert not response_chain_span_filter.errors
    assert response_chain_span_filter.data is not None
    assert response_chain_span_filter.data["project"]["latencyMsQuantile"] == 500.0

    response_llm_session_filter = await gql_client.execute(
        query=query,
        variables={
            "projectId": project_gid,
            "sessionFilterCondition": "session_id == 'latency-llm-session'",
        },
    )
    assert not response_llm_session_filter.errors
    assert response_llm_session_filter.data is not None
    assert response_llm_session_filter.data["project"]["latencyMsQuantile"] == 200.0

    response_chain_session_filter = await gql_client.execute(
        query=query,
        variables={
            "projectId": project_gid,
            "sessionFilterCondition": "session_id == 'latency-chain-session'",
        },
    )
    assert not response_chain_session_filter.errors
    assert response_chain_session_filter.data is not None
    assert response_chain_session_filter.data["project"]["latencyMsQuantile"] == 500.0


async def test_span_annotation_summary_with_session_filter_returns_expected_results(
    gql_client: AsyncGraphQLClient,
    db: DbSessionFactory,
) -> None:
    async with db() as session:
        project = await _add_project(session, name="annotation-session-test")
        session1 = await _add_project_session(session, project, session_id="priority-session")
        trace1 = await _add_trace(session, project, session1)
        span1 = await _add_span(session, trace1, attributes={"input": {"value": "priority task"}})
        annotation = models.SpanAnnotation(
            span_rowid=span1.id,
            name="test-annotation",
            label="important",
            score=1.0,
            explanation="Test annotation",
            metadata_={},
            annotator_kind="HUMAN",
            source="APP",
        )
        session.add(annotation)

        session2 = await _add_project_session(session, project)
        trace2 = await _add_trace(session, project, session2)
        span2 = await _add_span(session, trace2, attributes={"input": {"value": "normal task"}})
        annotation2 = models.SpanAnnotation(
            span_rowid=span2.id,
            name="test-annotation",
            label="normal",
            score=0.5,
            explanation="Test annotation",
            metadata_={},
            annotator_kind="HUMAN",
            source="APP",
        )
        session.add(annotation2)

    query = """
        query ($projectId: ID!, $sessionFilterCondition: String!) {
            node(id: $projectId) {
                ... on Project {
                    spanAnnotationSummary(annotationName: "test-annotation", sessionFilterCondition: $sessionFilterCondition) {
                        labelFractions {
                            label
                            fraction
                        }
                    }
                }
            }
        }
    """

    project_gid = str(GlobalID(type_name="Project", node_id=str(project.id)))
    response = await gql_client.execute(
        query=query,
        variables={
            "projectId": project_gid,
            "sessionFilterCondition": "session_id == 'priority-session'",
        },
    )

    assert not response.errors
    assert response.data is not None
    summary = response.data["node"]["spanAnnotationSummary"]
    assert summary is not None

    label_fractions = summary["labelFractions"]
    assert len(label_fractions) == 1
    assert label_fractions[0]["label"] == "important"
    assert label_fractions[0]["fraction"] == 1.0


async def test_trace_annotation_summary_with_session_filter_returns_expected_results(
    gql_client: AsyncGraphQLClient,
    db: DbSessionFactory,
) -> None:
    async with db() as session:
        project = await _add_project(session, name="annotation-session-test")
        session1 = await _add_project_session(session, project, session_id="priority-session")
        trace1 = await _add_trace(session, project, session1)
        await _add_span(session, trace1, attributes={"input": {"value": "priority task"}})
        annotation = models.TraceAnnotation(
            trace_rowid=trace1.id,
            name="test-annotation",
            label="important",
            score=1.0,
            explanation="Test annotation",
            metadata_={},
            annotator_kind="HUMAN",
            source="APP",
        )
        session.add(annotation)

        session2 = await _add_project_session(session, project)
        trace2 = await _add_trace(session, project, session2)
        await _add_span(session, trace2, attributes={"input": {"value": "normal task"}})
        annotation2 = models.TraceAnnotation(
            trace_rowid=trace2.id,
            name="test-annotation",
            label="normal",
            score=0.5,
            explanation="Test annotation",
            metadata_={},
            annotator_kind="HUMAN",
            source="APP",
        )
        session.add(annotation2)

    query = """
        query ($projectId: ID!, $sessionFilterCondition: String) {
            node(id: $projectId) {
                ... on Project {
                    traceAnnotationSummary(annotationName: "test-annotation", sessionFilterCondition: $sessionFilterCondition) {
                        labelFractions {
                            label
                            fraction
                        }
                    }
                }
            }
        }
    """

    project_gid = str(GlobalID(type_name="Project", node_id=str(project.id)))
    response = await gql_client.execute(
        query=query,
        variables={
            "projectId": project_gid,
            "sessionFilterCondition": "session_id == 'priority-session'",
        },
    )

    assert not response.errors
    assert response.data is not None
    summary = response.data["node"]["traceAnnotationSummary"]
    assert summary is not None

    label_fractions = summary["labelFractions"]
    assert len(label_fractions) == 1
    assert label_fractions[0]["label"] == "important"
    assert label_fractions[0]["fraction"] == 1.0


async def test_record_count_returns_expected_count(
    gql_client: AsyncGraphQLClient,
    db: DbSessionFactory,
) -> None:
    async with db() as session:
        project = await _add_project(session, name="annotation-session-test")
        session1 = await _add_project_session(session, project, session_id="priority-session")
        trace1 = await _add_trace(session, project, session1)
        await _add_span(session, trace1, attributes={"input": {"value": "priority task"}})
        annotation = models.TraceAnnotation(
            trace_rowid=trace1.id,
            name="test-annotation",
            label="important",
            score=1.0,
            explanation="Test annotation",
            metadata_={},
            annotator_kind="HUMAN",
            source="APP",
        )
        session.add(annotation)

        session2 = await _add_project_session(session, project)
        trace2 = await _add_trace(session, project, session2)
        await _add_span(session, trace2, attributes={"input": {"value": "normal task"}})
        annotation2 = models.TraceAnnotation(
            trace_rowid=trace2.id,
            name="test-annotation",
            label="normal",
            score=0.5,
            explanation="Test annotation",
            metadata_={},
            annotator_kind="HUMAN",
            source="APP",
        )
        session.add(annotation2)

    query = """
      query ($projectId: ID!, $filterCondition: String, $sessionFilterCondition: String) {
        project: node(id: $projectId) {
          ... on Project {
            recordCount(
              filterCondition: $filterCondition
              sessionFilterCondition: $sessionFilterCondition
            )
          }
        }
      }
    """

    project_gid = str(GlobalID(type_name="Project", node_id=str(project.id)))
    response = await gql_client.execute(
        query=query,
        variables={
            "projectId": project_gid,
        },
    )

    assert not response.errors
    assert response.data is not None
    assert response.data["project"]["recordCount"] == 2

    response = await gql_client.execute(
        query=query,
        variables={
            "projectId": project_gid,
            "sessionFilterCondition": "session_id == 'priority-session'",
        },
    )

    assert not response.errors
    assert response.data is not None
    assert response.data["project"]["recordCount"] == 1

    response = await gql_client.execute(
        query=query,
        variables={
            "projectId": project_gid,
            "filterCondition": "input.value == 'normal task'",
        },
    )

    assert not response.errors
    assert response.data is not None
    assert response.data["project"]["recordCount"] == 1


async def test_pagination_spans_by_annotation_score(
    gql_client: AsyncGraphQLClient,
    db: DbSessionFactory,
) -> None:
    """Test pagination of spans sorted by annotation score.

    Creates 3 spans:
    - First span has annotation scores
    - Last two spans have annotation labels but no scores
    Tests that pagination works correctly with score-based sorting.
    """
    spans = []
    traces = []
    score = random()
    base_time = datetime.fromisoformat("2024-01-01T00:00:00+00:00")
    async with db() as session:
        project = await _add_project(session, name="annotation-score-test")

        for i in range(3):
            trace = await _add_trace(
                session,
                project,
                start_time=base_time + timedelta(minutes=i * 10),
                end_time=base_time + timedelta(minutes=i * 10 + 5),
            )
            traces.append(trace)

            span = await _add_span(
                session,
                trace=trace,
                start_time=trace.start_time,
                end_time=trace.end_time,
                span_kind="LLM",
            )
            spans.append(span)

            if i < 1:
                annotation = models.SpanAnnotation(
                    span_rowid=span.id,
                    name="Quality",
                    label=None,
                    score=score,
                    explanation=f"Quality annotation for span {i + 1}",
                    metadata_={},
                    annotator_kind="LLM",
                    source="APP",
                )
            else:
                annotation = models.SpanAnnotation(
                    span_rowid=span.id,
                    name="Quality",
                    label="good",
                    score=None,  # No score
                    explanation=f"Quality annotation for span {i + 1}",
                    metadata_={},
                    annotator_kind="LLM",
                    source="APP",
                )
            session.add(annotation)

    project_gid = str(GlobalID(type_name="Project", node_id=str(project.id)))

    query = """
        query ($projectId: ID!, $first: Int!, $after: String, $sort: SpanSort) {
            node(id: $projectId) {
                ... on Project {
                    spans(
                        first: $first,
                        after: $after,
                        sort: $sort
                    ) {
                        edges {
                            node {
                                id
                                name
                            }
                            cursor
                        }
                        pageInfo {
                            hasNextPage
                            hasPreviousPage
                            startCursor
                            endCursor
                        }
                    }
                }
            }
        }
    """

    response = await gql_client.execute(
        query=query,
        variables={
            "projectId": project_gid,
            "first": 2,
            "sort": {"evalResultKey": {"name": "Quality", "attr": "score"}, "dir": "desc"},
        },
    )

    assert not response.errors
    assert (data := response.data) is not None

    page = data["node"]["spans"]
    edges = page["edges"]
    page_info = page["pageInfo"]

    assert len(edges) == 2
    assert edges[0]["node"]["name"] == spans[0].name
    assert edges[1]["node"]["name"] == spans[2].name
    assert page_info["hasNextPage"] is True
    assert page_info["hasPreviousPage"] is False

    start_cursor = Cursor.from_string(page_info["startCursor"])
    assert start_cursor.sort_column is not None
    assert start_cursor.sort_column.type == CursorSortColumnDataType.FLOAT
    assert start_cursor.sort_column.value == score

    end_cursor = Cursor.from_string(page_info["endCursor"])
    assert end_cursor.sort_column is not None
    assert end_cursor.sort_column.type == CursorSortColumnDataType.NULL
    assert end_cursor.sort_column.value is None

    response = await gql_client.execute(
        query=query,
        variables={
            "projectId": project_gid,
            "first": 2,
            "after": page_info["endCursor"],
            "sort": {"evalResultKey": {"name": "Quality", "attr": "score"}, "dir": "desc"},
        },
    )

    assert not response.errors
    assert (data := response.data) is not None

    page = data["node"]["spans"]
    edges = page["edges"]
    page_info = page["pageInfo"]

    assert len(edges) == 1
    assert edges[0]["node"]["name"] == spans[1].name
    assert page_info["hasNextPage"] is False

    start_cursor = Cursor.from_string(page_info["startCursor"])
    assert start_cursor.sort_column is not None
    assert start_cursor.sort_column.type == CursorSortColumnDataType.NULL
    assert start_cursor.sort_column.value is None


async def test_trace_count_returns_expected_count(
    gql_client: AsyncGraphQLClient,
    db: DbSessionFactory,
) -> None:
    async with db() as session:
        project = await _add_project(session, name="annotation-session-test")
        session1 = await _add_project_session(session, project, session_id="priority-session")
        trace1 = await _add_trace(session, project, session1)
        await _add_span(session, trace1, attributes={"input": {"value": "priority task"}})
        annotation = models.TraceAnnotation(
            trace_rowid=trace1.id,
            name="test-annotation",
            label="important",
            score=1.0,
            explanation="Test annotation",
            metadata_={},
            annotator_kind="HUMAN",
            source="APP",
        )
        session.add(annotation)

        session2 = await _add_project_session(session, project)
        trace2 = await _add_trace(session, project, session2)
        await _add_span(session, trace2, attributes={"input": {"value": "normal task"}})
        await _add_span(session, trace2, attributes={"input": {"value": "normal task"}})
        annotation2 = models.TraceAnnotation(
            trace_rowid=trace2.id,
            name="test-annotation",
            label="normal",
            score=0.5,
            explanation="Test annotation",
            metadata_={},
            annotator_kind="HUMAN",
            source="APP",
        )
        session.add(annotation2)

    query = """
      query ($projectId: ID!, $filterCondition: String, $sessionFilterCondition: String) {
        project: node(id: $projectId) {
          ... on Project {
            traceCount(
              filterCondition: $filterCondition
              sessionFilterCondition: $sessionFilterCondition
            )
          }
        }
      }
    """

    project_gid = str(GlobalID(type_name="Project", node_id=str(project.id)))
    response = await gql_client.execute(
        query=query,
        variables={
            "projectId": project_gid,
        },
    )

    assert not response.errors
    assert response.data is not None
    assert response.data["project"]["traceCount"] == 2

    response = await gql_client.execute(
        query=query,
        variables={
            "projectId": project_gid,
            "sessionFilterCondition": "session_id == 'priority-session'",
        },
    )

    assert not response.errors
    assert response.data is not None
    assert response.data["project"]["traceCount"] == 1

    response = await gql_client.execute(
        query=query,
        variables={
            "projectId": project_gid,
            "filterCondition": "input.value == 'normal task'",
        },
    )

    assert not response.errors
    assert response.data is not None
    assert response.data["project"]["traceCount"] == 1


async def test_session_count_returns_expected_count(
    gql_client: AsyncGraphQLClient,
    db: DbSessionFactory,
) -> None:
    early = datetime.fromisoformat("2021-01-01T00:00:00.000+00:00")
    late = datetime.fromisoformat("2021-06-01T00:00:00.000+00:00")
    async with db() as session:
        project = await _add_project(session, name="session-count-test")
        await _add_project_session(session, project, start_time=early)
        await _add_project_session(session, project, session_id="target-session", start_time=late)

    query = """
      query ($projectId: ID!, $timeRange: TimeRange, $sessionFilterCondition: String) {
        project: node(id: $projectId) {
          ... on Project {
            sessionCount(
              timeRange: $timeRange
              sessionFilterCondition: $sessionFilterCondition
            )
          }
        }
      }
    """

    project_gid = str(GlobalID(type_name="Project", node_id=str(project.id)))

    # Unfiltered: batched through the record_counts dataloader.
    response = await gql_client.execute(
        query=query,
        variables={
            "projectId": project_gid,
        },
    )

    assert not response.errors
    assert response.data is not None
    assert response.data["project"]["sessionCount"] == 2

    # Time-range filter still routes through the dataloader.
    response = await gql_client.execute(
        query=query,
        variables={
            "projectId": project_gid,
            "timeRange": {
                "start": late.isoformat(),
                "end": (late + timedelta(days=1)).isoformat(),
            },
        },
    )

    assert not response.errors
    assert response.data is not None
    assert response.data["project"]["sessionCount"] == 1

    # An expression filter leaves the dataloader path for a direct count query.
    response = await gql_client.execute(
        query=query,
        variables={
            "projectId": project_gid,
            "sessionFilterCondition": "session_id == 'target-session'",
        },
    )

    assert not response.errors
    assert response.data is not None
    assert response.data["project"]["sessionCount"] == 1


class TestAnnotationScoreTimeSeries:
    @pytest.fixture
    async def _annotation_score_data(self, db: DbSessionFactory) -> models.Project:
        """Creates a project with annotated traces and sessions in two hourly buckets.

        Hour one (01:00) has a trace with two "quality" trace annotations (scores 0.2
        and 0.4) and a session with one "satisfaction" annotation (score 0.5). Hour two
        (02:00) has a trace with "quality" (0.8) and "toxicity" (1.0) trace annotations
        and a session with one "satisfaction" annotation (score 0.9). Annotations named
        "incomplete" have null scores and must be excluded from the results.
        """
        hour_one = datetime.fromisoformat("2024-01-01T01:15:00+00:00")
        hour_two = datetime.fromisoformat("2024-01-01T02:20:00+00:00")
        async with db() as session:
            project = await _add_project(session)
            session_one = await _add_project_session(session, project, start_time=hour_one)
            session_two = await _add_project_session(session, project, start_time=hour_two)
            trace_one = await _add_trace(session, project, session_one, start_time=hour_one)
            trace_two = await _add_trace(session, project, session_two, start_time=hour_two)

            def trace_annotation(
                trace: models.Trace,
                name: str,
                score: Optional[float],
                identifier: str = "",
            ) -> models.TraceAnnotation:
                return models.TraceAnnotation(
                    trace_rowid=trace.id,
                    name=name,
                    label=None,
                    score=score,
                    explanation=None,
                    metadata_={},
                    annotator_kind="HUMAN",
                    identifier=identifier,
                    source="APP",
                    user_id=None,
                )

            def session_annotation(
                project_session: models.ProjectSession,
                name: str,
                score: Optional[float],
            ) -> models.ProjectSessionAnnotation:
                return models.ProjectSessionAnnotation(
                    project_session_id=project_session.id,
                    name=name,
                    label=None,
                    score=score,
                    explanation=None,
                    metadata_={},
                    annotator_kind="HUMAN",
                    identifier="",
                    source="APP",
                    user_id=None,
                )

            session.add_all(
                [
                    trace_annotation(trace_one, "quality", 0.2, identifier="a"),
                    trace_annotation(trace_one, "quality", 0.4, identifier="b"),
                    trace_annotation(trace_one, "incomplete", None),
                    trace_annotation(trace_two, "quality", 0.8),
                    trace_annotation(trace_two, "toxicity", 1.0),
                    session_annotation(session_one, "satisfaction", 0.5),
                    session_annotation(session_one, "incomplete", None),
                    session_annotation(session_two, "satisfaction", 0.9),
                ]
            )
        return project

    @staticmethod
    def _variables(project: models.Project) -> dict[str, Any]:
        return {
            "id": str(GlobalID(type_name="Project", node_id=str(project.id))),
            "timeRange": {
                "start": "2024-01-01T01:00:00+00:00",
                "end": "2024-01-01T03:00:00+00:00",
            },
            "timeBinConfig": {"scale": "HOUR", "utcOffsetMinutes": 0},
        }

    @staticmethod
    def _scores_by_timestamp(
        data: list[dict[str, Any]],
    ) -> dict[datetime, dict[str, float]]:
        return {
            datetime.fromisoformat(point["timestamp"]): {
                score_with_label["label"]: score_with_label["score"]
                for score_with_label in point["scoresWithLabels"]
            }
            for point in data
            if point["scoresWithLabels"]
        }

    async def test_trace_annotation_score_time_series(
        self,
        _annotation_score_data: models.Project,
        gql_client: AsyncGraphQLClient,
    ) -> None:
        query = """
            query($id: ID!, $timeRange: TimeRange!, $timeBinConfig: TimeBinConfig) {
                node(id: $id) {
                    ... on Project {
                        traceAnnotationScoreTimeSeries(
                            timeRange: $timeRange
                            timeBinConfig: $timeBinConfig
                        ) {
                            data {
                                timestamp
                                scoresWithLabels {
                                    label
                                    score
                                }
                            }
                            names
                        }
                    }
                }
            }
        """
        response = await gql_client.execute(
            query=query, variables=self._variables(_annotation_score_data)
        )
        assert not response.errors
        assert (data := response.data) is not None
        res = data["node"]["traceAnnotationScoreTimeSeries"]
        assert res["names"] == ["quality", "toxicity"]
        assert self._scores_by_timestamp(res["data"]) == {
            datetime.fromisoformat("2024-01-01T01:00:00+00:00"): {
                "quality": pytest.approx(0.3),
            },
            datetime.fromisoformat("2024-01-01T02:00:00+00:00"): {
                "quality": pytest.approx(0.8),
                "toxicity": pytest.approx(1.0),
            },
        }

    async def test_session_annotation_score_time_series(
        self,
        _annotation_score_data: models.Project,
        gql_client: AsyncGraphQLClient,
    ) -> None:
        query = """
            query($id: ID!, $timeRange: TimeRange!, $timeBinConfig: TimeBinConfig) {
                node(id: $id) {
                    ... on Project {
                        sessionAnnotationScoreTimeSeries(
                            timeRange: $timeRange
                            timeBinConfig: $timeBinConfig
                        ) {
                            data {
                                timestamp
                                scoresWithLabels {
                                    label
                                    score
                                }
                            }
                            names
                        }
                    }
                }
            }
        """
        response = await gql_client.execute(
            query=query, variables=self._variables(_annotation_score_data)
        )
        assert not response.errors
        assert (data := response.data) is not None
        res = data["node"]["sessionAnnotationScoreTimeSeries"]
        assert res["names"] == ["satisfaction"]
        assert self._scores_by_timestamp(res["data"]) == {
            datetime.fromisoformat("2024-01-01T01:00:00+00:00"): {
                "satisfaction": pytest.approx(0.5),
            },
            datetime.fromisoformat("2024-01-01T02:00:00+00:00"): {
                "satisfaction": pytest.approx(0.9),
            },
        }


class TestAnnotationMetricsTimeSeries:
    @pytest.fixture
    async def _annotation_metrics_data(self, db: DbSessionFactory) -> models.Project:
        hour_one = datetime.fromisoformat("2024-01-01T01:15:00+00:00")
        hour_two = datetime.fromisoformat("2024-01-01T02:20:00+00:00")
        async with db() as session:
            project = await _add_project(session)
            session_one = await _add_project_session(session, project, start_time=hour_one)
            session_two = await _add_project_session(session, project, start_time=hour_two)
            trace_one = await _add_trace(session, project, session_one, start_time=hour_one)
            trace_two = await _add_trace(session, project, session_two, start_time=hour_two)
            span_one = await _add_span(session, trace_one, start_time=hour_one)
            span_two = await _add_span(session, trace_two, start_time=hour_two)

            annotation_values = [
                ("mixed", "pass", 0.2, None),
                ("mixed", "fail", 0.4, None),
                ("mixed", None, 0.6, None),
                ("mixed", None, None, "explanation only"),
                ("empty", None, None, "explanation only"),
                ("label-only", "yes", None, None),
                ("score-only", None, 0.5, None),
            ]
            for model, parent_id_field, parent_id in (
                (models.SpanAnnotation, "span_rowid", span_one.id),
                (models.TraceAnnotation, "trace_rowid", trace_one.id),
                (
                    models.ProjectSessionAnnotation,
                    "project_session_id",
                    session_one.id,
                ),
            ):
                session.add_all(
                    [
                        model(
                            **{parent_id_field: parent_id},
                            name=name,
                            label=label,
                            score=score,
                            explanation=explanation,
                            metadata_={},
                            annotator_kind="HUMAN",
                            identifier=str(index),
                            source="APP",
                            user_id=None,
                        )
                        for index, (name, label, score, explanation) in enumerate(annotation_values)
                    ]
                )
            for model, parent_id_field, parent_id in (
                (models.SpanAnnotation, "span_rowid", span_two.id),
                (models.TraceAnnotation, "trace_rowid", trace_two.id),
                (
                    models.ProjectSessionAnnotation,
                    "project_session_id",
                    session_two.id,
                ),
            ):
                session.add(
                    model(
                        **{parent_id_field: parent_id},
                        name="mixed",
                        label="pass",
                        score=0.8,
                        explanation=None,
                        metadata_={},
                        annotator_kind="HUMAN",
                        identifier="hour-two",
                        source="APP",
                        user_id=None,
                    )
                )
        return project

    async def test_annotation_metrics_time_series_at_all_levels(
        self,
        _annotation_metrics_data: models.Project,
        gql_client: AsyncGraphQLClient,
    ) -> None:
        query = """
            query($id: ID!, $timeRange: TimeRange!, $timeBinConfig: TimeBinConfig) {
                node(id: $id) {
                    ... on Project {
                        span: spanAnnotationMetricsTimeSeries(
                            timeRange: $timeRange
                            timeBinConfig: $timeBinConfig
                        ) { ...metrics }
                        trace: traceAnnotationMetricsTimeSeries(
                            timeRange: $timeRange
                            timeBinConfig: $timeBinConfig
                        ) { ...metrics }
                        session: sessionAnnotationMetricsTimeSeries(
                            timeRange: $timeRange
                            timeBinConfig: $timeBinConfig
                        ) { ...metrics }
                    }
                }
            }

            fragment metrics on AnnotationMetricsTimeSeries {
                names
                data {
                    timestamp
                    annotationSummaries {
                        name
                        count
                        scoreCount
                        labelCount
                        meanScore
                        labelFractions { label fraction }
                    }
                }
            }
        """
        response = await gql_client.execute(
            query=query,
            variables={
                "id": str(GlobalID(type_name="Project", node_id=str(_annotation_metrics_data.id))),
                "timeRange": {
                    "start": "2024-01-01T01:00:00+00:00",
                    "end": "2024-01-01T04:00:00+00:00",
                },
                "timeBinConfig": {"scale": "HOUR", "utcOffsetMinutes": 0},
            },
        )
        assert not response.errors
        assert response.data is not None
        for level in ("span", "trace", "session"):
            metrics = response.data["node"][level]
            assert metrics["names"] == ["label-only", "mixed", "score-only"]
            assert [point["timestamp"] for point in metrics["data"]] == [
                "2024-01-01T01:00:00+00:00",
                "2024-01-01T02:00:00+00:00",
                "2024-01-01T03:00:00+00:00",
            ]
            hour_one_summaries = {
                summary["name"]: summary for summary in metrics["data"][0]["annotationSummaries"]
            }
            assert hour_one_summaries["label-only"] == {
                "name": "label-only",
                "count": 1,
                "scoreCount": 0,
                "labelCount": 1,
                "meanScore": None,
                "labelFractions": [{"label": "yes", "fraction": 1.0}],
            }
            assert hour_one_summaries["score-only"] == {
                "name": "score-only",
                "count": 1,
                "scoreCount": 1,
                "labelCount": 0,
                "meanScore": pytest.approx(0.5),
                "labelFractions": [],
            }
            mixed_summary = hour_one_summaries["mixed"]
            assert mixed_summary["count"] == 3
            assert mixed_summary["scoreCount"] == 3
            assert mixed_summary["labelCount"] == 2
            assert mixed_summary["meanScore"] == pytest.approx(0.4)
            assert mixed_summary["labelFractions"] == [
                {"label": "fail", "fraction": pytest.approx(1 / 2)},
                {"label": "pass", "fraction": pytest.approx(1 / 2)},
            ]
            assert metrics["data"][2]["annotationSummaries"] == []

    async def test_annotation_metrics_time_series_filters_by_name_at_all_levels(
        self,
        _annotation_metrics_data: models.Project,
        gql_client: AsyncGraphQLClient,
    ) -> None:
        query = """
            query(
                $id: ID!
                $annotationName: String!
                $timeRange: TimeRange!
                $timeBinConfig: TimeBinConfig
            ) {
                node(id: $id) {
                    ... on Project {
                        span: spanAnnotationMetricsTimeSeries(
                            annotationName: $annotationName
                            timeRange: $timeRange
                            timeBinConfig: $timeBinConfig
                        ) { ...metrics }
                        trace: traceAnnotationMetricsTimeSeries(
                            annotationName: $annotationName
                            timeRange: $timeRange
                            timeBinConfig: $timeBinConfig
                        ) { ...metrics }
                        session: sessionAnnotationMetricsTimeSeries(
                            annotationName: $annotationName
                            timeRange: $timeRange
                            timeBinConfig: $timeBinConfig
                        ) { ...metrics }
                    }
                }
            }

            fragment metrics on AnnotationMetricsTimeSeries {
                names
                data {
                    timestamp
                    annotationSummaries {
                        name
                    }
                }
            }
        """
        response = await gql_client.execute(
            query=query,
            variables={
                "id": str(GlobalID(type_name="Project", node_id=str(_annotation_metrics_data.id))),
                "annotationName": "mixed",
                "timeRange": {
                    "start": "2024-01-01T01:00:00+00:00",
                    "end": "2024-01-01T04:00:00+00:00",
                },
                "timeBinConfig": {"scale": "HOUR", "utcOffsetMinutes": 0},
            },
        )
        assert not response.errors
        assert response.data is not None
        for level in ("span", "trace", "session"):
            metrics = response.data["node"][level]
            assert metrics["names"] == ["mixed"]
            assert [point["timestamp"] for point in metrics["data"]] == [
                "2024-01-01T01:00:00+00:00",
                "2024-01-01T02:00:00+00:00",
                "2024-01-01T03:00:00+00:00",
            ]
            assert [
                [summary["name"] for summary in point["annotationSummaries"]]
                for point in metrics["data"]
            ] == [["mixed"], ["mixed"], []]

    async def test_annotation_metric_names_are_chartable_and_time_range_scoped(
        self,
        db: DbSessionFactory,
        gql_client: AsyncGraphQLClient,
    ) -> None:
        annotations = (
            ("before", "2024-01-01T00:59:00+00:00", 0.1, None, None),
            ("z-start", "2024-01-01T01:00:00+00:00", 0.2, None, None),
            (
                "explanation-only",
                "2024-01-01T01:15:00+00:00",
                None,
                None,
                "not chartable",
            ),
            ("a-within", "2024-01-01T01:30:00+00:00", None, "pass", None),
            ("at-end", "2024-01-01T02:00:00+00:00", 0.9, None, None),
        )
        async with db() as session:
            project = await _add_project(session)
            for name, timestamp_string, score, label, explanation in annotations:
                timestamp = datetime.fromisoformat(timestamp_string)
                project_session = await _add_project_session(session, project, start_time=timestamp)
                trace = await _add_trace(
                    session,
                    project,
                    project_session,
                    start_time=timestamp,
                )
                span = await _add_span(session, trace, start_time=timestamp)
                common_annotation_fields: dict[str, Any] = {
                    "label": label,
                    "score": score,
                    "explanation": explanation,
                    "metadata_": {},
                    "annotator_kind": "HUMAN",
                    "source": "APP",
                    "user_id": None,
                }
                session.add_all(
                    [
                        models.SpanAnnotation(
                            span_rowid=span.id,
                            name=f"span-{name}",
                            identifier=f"span-{name}",
                            **common_annotation_fields,
                        ),
                        models.TraceAnnotation(
                            trace_rowid=trace.id,
                            name=f"trace-{name}",
                            identifier=f"trace-{name}",
                            **common_annotation_fields,
                        ),
                        models.ProjectSessionAnnotation(
                            project_session_id=project_session.id,
                            name=f"session-{name}",
                            identifier=f"session-{name}",
                            **common_annotation_fields,
                        ),
                    ]
                )

            other_project = await _add_project(
                session,
                name="other-annotation-metrics-project",
            )
            other_timestamp = datetime.fromisoformat("2024-01-01T01:45:00+00:00")
            other_project_session = await _add_project_session(
                session,
                other_project,
                start_time=other_timestamp,
            )
            other_trace = await _add_trace(
                session,
                other_project,
                other_project_session,
                start_time=other_timestamp,
            )
            other_span = await _add_span(session, other_trace, start_time=other_timestamp)
            other_annotation_fields: dict[str, Any] = {
                "label": None,
                "score": 1.0,
                "explanation": None,
                "metadata_": {},
                "annotator_kind": "HUMAN",
                "source": "APP",
                "user_id": None,
            }
            session.add_all(
                [
                    models.SpanAnnotation(
                        span_rowid=other_span.id,
                        name="other-project-span",
                        identifier="other-project-span",
                        **other_annotation_fields,
                    ),
                    models.TraceAnnotation(
                        trace_rowid=other_trace.id,
                        name="other-project-trace",
                        identifier="other-project-trace",
                        **other_annotation_fields,
                    ),
                    models.ProjectSessionAnnotation(
                        project_session_id=other_project_session.id,
                        name="other-project-session",
                        identifier="other-project-session",
                        **other_annotation_fields,
                    ),
                ]
            )

        query = """
            query($id: ID!, $timeRange: TimeRange!) {
                node(id: $id) {
                    ... on Project {
                        spanAnnotationMetricNames(timeRange: $timeRange)
                        traceAnnotationMetricNames(timeRange: $timeRange)
                        sessionAnnotationMetricNames(timeRange: $timeRange)
                    }
                }
            }
        """
        response = await gql_client.execute(
            query=query,
            variables={
                "id": str(GlobalID(type_name="Project", node_id=str(project.id))),
                "timeRange": {
                    "start": "2024-01-01T01:00:00+00:00",
                    "end": "2024-01-01T02:00:00+00:00",
                },
            },
        )
        assert not response.errors
        assert response.data is not None
        assert response.data["node"] == {
            "spanAnnotationMetricNames": ["span-a-within", "span-z-start"],
            "traceAnnotationMetricNames": ["trace-a-within", "trace-z-start"],
            "sessionAnnotationMetricNames": ["session-a-within", "session-z-start"],
        }

    async def test_label_fractions_weight_entities_instead_of_annotation_rows(
        self,
        db: DbSessionFactory,
        gql_client: AsyncGraphQLClient,
    ) -> None:
        timestamp = datetime.fromisoformat("2024-01-01T01:15:00+00:00")
        async with db() as session:
            project = await _add_project(session)
            project_sessions = [
                await _add_project_session(session, project, start_time=timestamp) for _ in range(3)
            ]
            traces = [
                await _add_trace(
                    session,
                    project,
                    project_session,
                    start_time=timestamp,
                )
                for project_session in project_sessions
            ]
            spans = [await _add_span(session, trace, start_time=timestamp) for trace in traces]

            for model, parent_id_field, parent_ids in (
                (models.SpanAnnotation, "span_rowid", [span.id for span in spans]),
                (models.TraceAnnotation, "trace_rowid", [trace.id for trace in traces]),
                (
                    models.ProjectSessionAnnotation,
                    "project_session_id",
                    [project_session.id for project_session in project_sessions],
                ),
            ):
                for index, (parent_id, label, score) in enumerate(
                    zip(
                        parent_ids,
                        ("pass", "fail", None),
                        (0.2, 0.4, 0.6),
                    )
                ):
                    session.add(
                        model(
                            **{parent_id_field: parent_id},
                            name="coverage",
                            label=label,
                            score=score,
                            explanation=None,
                            metadata_={},
                            annotator_kind="HUMAN",
                            identifier=f"entity-{index}",
                            source="APP",
                            user_id=None,
                        )
                    )
                session.add_all(
                    [
                        model(
                            **{parent_id_field: parent_ids[0]},
                            name="coverage",
                            label="pass",
                            score=None,
                            explanation=None,
                            metadata_={},
                            annotator_kind="HUMAN",
                            identifier=f"repeat-{index}",
                            source="APP",
                            user_id=None,
                        )
                        for index in range(9)
                    ]
                )

        query = """
            query($id: ID!, $timeRange: TimeRange!, $timeBinConfig: TimeBinConfig) {
                node(id: $id) {
                    ... on Project {
                        span: spanAnnotationMetricsTimeSeries(
                            timeRange: $timeRange
                            timeBinConfig: $timeBinConfig
                        ) { ...metrics }
                        trace: traceAnnotationMetricsTimeSeries(
                            timeRange: $timeRange
                            timeBinConfig: $timeBinConfig
                        ) { ...metrics }
                        session: sessionAnnotationMetricsTimeSeries(
                            timeRange: $timeRange
                            timeBinConfig: $timeBinConfig
                        ) { ...metrics }
                    }
                }
            }

            fragment metrics on AnnotationMetricsTimeSeries {
                data {
                    annotationSummaries {
                        name
                        count
                        scoreCount
                        labelCount
                        meanScore
                        labelFractions { label fraction }
                    }
                }
            }
        """
        response = await gql_client.execute(
            query=query,
            variables={
                "id": str(GlobalID(type_name="Project", node_id=str(project.id))),
                "timeRange": {
                    "start": "2024-01-01T01:00:00+00:00",
                    "end": "2024-01-01T02:00:00+00:00",
                },
                "timeBinConfig": {"scale": "HOUR", "utcOffsetMinutes": 0},
            },
        )

        assert not response.errors
        assert response.data is not None
        for level in ("span", "trace", "session"):
            summary = response.data["node"][level]["data"][0]["annotationSummaries"][0]
            assert summary == {
                "name": "coverage",
                "count": 12,
                "scoreCount": 3,
                "labelCount": 11,
                "meanScore": pytest.approx(0.4),
                "labelFractions": [
                    {"label": "fail", "fraction": pytest.approx(1 / 3)},
                    {"label": "pass", "fraction": pytest.approx(1 / 3)},
                ],
            }

    async def test_each_time_bin_returns_its_observed_labels(
        self,
        db: DbSessionFactory,
        gql_client: AsyncGraphQLClient,
    ) -> None:
        hour_one = datetime.fromisoformat("2024-01-01T01:15:00+00:00")
        hour_two = datetime.fromisoformat("2024-01-01T02:15:00+00:00")
        hour_three = datetime.fromisoformat("2024-01-01T03:15:00+00:00")
        async with db() as session:
            project = await _add_project(session)
            project_session_one = await _add_project_session(session, project, start_time=hour_one)
            project_session_two = await _add_project_session(session, project, start_time=hour_two)
            project_session_three = await _add_project_session(
                session, project, start_time=hour_three
            )
            trace_one = await _add_trace(session, project, project_session_one, start_time=hour_one)
            trace_two = await _add_trace(session, project, project_session_two, start_time=hour_two)
            trace_three = await _add_trace(
                session, project, project_session_three, start_time=hour_three
            )
            labels_by_trace = (
                (trace_one, ["pass", "pass", "pass", "fail"]),
                (trace_two, ["pass", "review"]),
            )
            for trace, labels in labels_by_trace:
                session.add_all(
                    [
                        models.TraceAnnotation(
                            trace_rowid=trace.id,
                            name="quality",
                            label=label,
                            score=None,
                            explanation=None,
                            metadata_={},
                            annotator_kind="CODE",
                            identifier=str(index),
                            source="APP",
                            user_id=None,
                        )
                        for index, label in enumerate(labels)
                    ]
                )
            session.add(
                models.TraceAnnotation(
                    trace_rowid=trace_three.id,
                    name="quality",
                    label=None,
                    score=0.9,
                    explanation=None,
                    metadata_={},
                    annotator_kind="CODE",
                    identifier="score-only-tail",
                    source="APP",
                    user_id=None,
                )
            )

        query = """
          query ($id: ID!, $timeRange: TimeRange!, $timeBinConfig: TimeBinConfig) {
            node(id: $id) {
              ... on Project {
                traceAnnotationMetricsTimeSeries(
                  timeRange: $timeRange
                  timeBinConfig: $timeBinConfig
                ) {
                  data {
                    annotationSummaries {
                      name
                      count
                      scoreCount
                      labelCount
                      labelFractions { label fraction }
                    }
                  }
                }
              }
            }
          }
        """
        response = await gql_client.execute(
            query=query,
            variables={
                "id": str(GlobalID("Project", str(project.id))),
                "timeRange": {
                    "start": "2024-01-01T01:00:00+00:00",
                    "end": "2024-01-01T04:00:00+00:00",
                },
                "timeBinConfig": {"scale": "HOUR", "utcOffsetMinutes": 0},
            },
        )

        assert not response.errors
        assert response.data is not None
        data = response.data["node"]["traceAnnotationMetricsTimeSeries"]["data"]
        first_summary = data[0]["annotationSummaries"][0]
        second_summary = data[1]["annotationSummaries"][0]
        third_summary = data[2]["annotationSummaries"][0]
        assert first_summary["count"] == first_summary["labelCount"] == 4
        assert second_summary["count"] == second_summary["labelCount"] == 2
        assert [item["label"] for item in first_summary["labelFractions"]] == ["fail", "pass"]
        assert [item["label"] for item in second_summary["labelFractions"]] == [
            "pass",
            "review",
        ]
        assert sum(item["fraction"] for item in first_summary["labelFractions"]) == pytest.approx(1)
        assert sum(item["fraction"] for item in second_summary["labelFractions"]) == pytest.approx(
            1
        )
        assert third_summary["scoreCount"] == 1
        assert third_summary["labelFractions"] == []


async def test_trace_resolves_by_otel_id_and_global_node_id(
    db: DbSessionFactory,
    gql_client: AsyncGraphQLClient,
) -> None:
    # The trace resolver accepts either an OpenTelemetry trace id (a hex string)
    # or a Relay global node id, decoding the latter when possible and otherwise
    # falling back to an OpenTelemetry trace id lookup
    # (see https://github.com/Arize-ai/phoenix/issues/13513).
    async with db() as session:
        project = await _add_project(session)
        trace = await _add_trace(session, project)
    project_gid = str(GlobalID("Project", str(project.id)))
    trace_gid = str(GlobalID("Trace", str(trace.id)))
    query = """
    query ($pid: ID!, $tid: ID!) {
      node(id: $pid) {
        ... on Project {
          trace(traceId: $tid) { id }
        }
      }
    }
    """

    # Resolve by OpenTelemetry trace id.
    response = await gql_client.execute(
        query=query,
        variables={"pid": project_gid, "tid": trace.trace_id},
    )
    assert not response.errors
    assert response.data is not None
    assert response.data["node"]["trace"]["id"] == trace_gid

    # Resolve by Relay global node id.
    response = await gql_client.execute(
        query=query,
        variables={"pid": project_gid, "tid": trace_gid},
    )
    assert not response.errors
    assert response.data is not None
    assert response.data["node"]["trace"]["id"] == trace_gid


async def test_trace_with_unmatched_global_node_id_returns_null(
    db: DbSessionFactory,
    gql_client: AsyncGraphQLClient,
) -> None:
    async with db() as session:
        project = await _add_project(session)
    project_gid = str(GlobalID("Project", str(project.id)))
    node_id_as_trace_id = str(GlobalID("Trace", "2003"))
    response = await gql_client.execute(
        query="""
        query ($pid: ID!, $tid: ID!) {
          node(id: $pid) {
            ... on Project {
              trace(traceId: $tid) { id }
            }
          }
        }
        """,
        variables={"pid": project_gid, "tid": node_id_as_trace_id},
    )
    assert not response.errors
    assert response.data is not None
    assert response.data["node"]["trace"] is None


async def test_session_stats(
    gql_client: AsyncGraphQLClient,
    db: DbSessionFactory,
) -> None:
    # session1 has two traces (i.e. two turns) and lasts 10 seconds; session2
    # has one trace and lasts 20 seconds
    base_time = datetime(2024, 1, 1, 12, 0, 0, tzinfo=timezone.utc)
    async with db() as session:
        project = await _add_project(session, name="session-stats-test")
        session1 = await _add_project_session(
            session,
            project,
            session_id="session-1-exact-id",
            start_time=base_time,
            end_time=base_time + timedelta(seconds=10),
        )
        trace1 = await _add_trace(session, project, session1, start_time=base_time)
        await _add_span(session, trace1, attributes={"input": {"value": "alpha task"}})
        trace1b = await _add_trace(
            session, project, session1, start_time=base_time + timedelta(seconds=5)
        )
        await _add_span(session, trace1b, attributes={"input": {"value": "alpha follow-up"}})
        session2 = await _add_project_session(
            session,
            project,
            start_time=base_time + timedelta(minutes=5),
            end_time=base_time + timedelta(minutes=5, seconds=20),
        )
        trace2 = await _add_trace(
            session, project, session2, start_time=base_time + timedelta(minutes=5)
        )
        await _add_span(session, trace2, attributes={"input": {"value": "beta task"}})

    query = """
      query ($projectId: ID!, $timeRange: TimeRange) {
        node(id: $projectId) {
          ... on Project {
            sessionCount(timeRange: $timeRange)
            averageSessionDurationMs(timeRange: $timeRange)
            averageTracesPerSession(timeRange: $timeRange)
            sessionDurationMsP50: sessionDurationMsQuantile(
              probability: 0.5
              timeRange: $timeRange
            )
            sessionDurationMsP99: sessionDurationMsQuantile(
              probability: 0.99
              timeRange: $timeRange
            )
          }
        }
      }
    """
    project_gid = str(GlobalID(type_name="Project", node_id=str(project.id)))

    response = await gql_client.execute(query=query, variables={"projectId": project_gid})
    assert not response.errors
    assert response.data is not None
    assert response.data["node"]["sessionCount"] == 2
    assert response.data["node"]["averageSessionDurationMs"] == 15000.0
    assert response.data["node"]["averageTracesPerSession"] == 1.5
    assert response.data["node"]["sessionDurationMsP50"] == pytest.approx(15000.0)
    assert response.data["node"]["sessionDurationMsP99"] == pytest.approx(19900.0)

    response = await gql_client.execute(
        query=query,
        variables={
            "projectId": project_gid,
            "timeRange": {
                "start": base_time.isoformat(),
                "end": (base_time + timedelta(minutes=1)).isoformat(),
            },
        },
    )
    assert not response.errors
    assert response.data is not None
    assert response.data["node"]["sessionCount"] == 1
    assert response.data["node"]["averageSessionDurationMs"] == 10000.0
    assert response.data["node"]["averageTracesPerSession"] == 2.0
    assert response.data["node"]["sessionDurationMsP50"] == pytest.approx(10000.0)
    assert response.data["node"]["sessionDurationMsP99"] == pytest.approx(10000.0)

    response = await gql_client.execute(
        query=query,
        variables={
            "projectId": project_gid,
            "timeRange": {
                "start": (base_time + timedelta(days=1)).isoformat(),
                "end": (base_time + timedelta(days=2)).isoformat(),
            },
        },
    )
    assert not response.errors
    assert response.data is not None
    assert response.data["node"]["sessionCount"] == 0
    assert response.data["node"]["averageSessionDurationMs"] is None
    assert response.data["node"]["averageTracesPerSession"] is None
    assert response.data["node"]["sessionDurationMsP50"] is None
    assert response.data["node"]["sessionDurationMsP99"] is None


async def test_session_annotation_summary_returns_expected_results(
    gql_client: AsyncGraphQLClient,
    db: DbSessionFactory,
) -> None:
    async with db() as session:
        project = await _add_project(session, name="session-annotation-summary-test")
        session1 = await _add_project_session(
            session, project, session_id="annotated-session-exact-id"
        )
        trace1 = await _add_trace(session, project, session1)
        await _add_span(session, trace1, attributes={"input": {"value": "priority task"}})
        session.add(
            models.ProjectSessionAnnotation(
                project_session_id=session1.id,
                name="test-annotation",
                label="important",
                score=1.0,
                explanation="Test annotation",
                metadata_={},
                annotator_kind="HUMAN",
                source="APP",
            )
        )
        session2 = await _add_project_session(session, project)
        trace2 = await _add_trace(session, project, session2)
        await _add_span(session, trace2, attributes={"input": {"value": "normal task"}})
        session.add(
            models.ProjectSessionAnnotation(
                project_session_id=session2.id,
                name="test-annotation",
                label="normal",
                score=0.5,
                explanation="Test annotation",
                metadata_={},
                annotator_kind="HUMAN",
                source="APP",
            )
        )

    query = """
      query ($projectId: ID!) {
        node(id: $projectId) {
          ... on Project {
            sessionAnnotationSummary(annotationName: "test-annotation") {
              name
              count
              scoreCount
              labelCount
              meanScore
              labelFractions {
                label
                fraction
              }
            }
          }
        }
      }
    """
    project_gid = str(GlobalID(type_name="Project", node_id=str(project.id)))

    response = await gql_client.execute(query=query, variables={"projectId": project_gid})
    assert not response.errors
    assert response.data is not None
    summary = response.data["node"]["sessionAnnotationSummary"]
    assert summary is not None
    assert summary["name"] == "test-annotation"
    assert summary["count"] == 2
    assert summary["scoreCount"] == 2
    assert summary["labelCount"] == 2
    assert summary["meanScore"] == 0.75
    assert summary["labelFractions"] == [
        {"label": "important", "fraction": 0.5},
        {"label": "normal", "fraction": 0.5},
    ]


async def test_session_annotation_summary_time_range_uses_interval_overlap(
    gql_client: AsyncGraphQLClient,
    db: DbSessionFactory,
) -> None:
    # long_running starts before the requested window but is active inside it;
    # before_window ends before the window starts. The summary must use the
    # same interval-overlap semantics as the sessions table and include only
    # long_running's annotation.
    base_time = datetime(2024, 1, 1, 12, 0, 0, tzinfo=timezone.utc)
    async with db() as session:
        project = await _add_project(session, name="session-annotation-overlap-test")
        intervals = {
            "long_running": (timedelta(0), timedelta(minutes=15)),
            "before_window": (timedelta(0), timedelta(minutes=5)),
        }
        for label, (start_offset, end_offset) in intervals.items():
            project_session = await _add_project_session(
                session,
                project,
                start_time=base_time + start_offset,
                end_time=base_time + end_offset,
            )
            session.add(
                models.ProjectSessionAnnotation(
                    project_session_id=project_session.id,
                    name="test-annotation",
                    label=label,
                    score=1.0,
                    explanation="Test annotation",
                    metadata_={},
                    annotator_kind="HUMAN",
                    source="APP",
                )
            )

    query = """
      query ($projectId: ID!, $timeRange: TimeRange) {
        node(id: $projectId) {
          ... on Project {
            sessionAnnotationSummary(
              annotationName: "test-annotation"
              timeRange: $timeRange
            ) {
              count
              labelFractions {
                label
                fraction
              }
            }
          }
        }
      }
    """
    project_gid = str(GlobalID(type_name="Project", node_id=str(project.id)))
    response = await gql_client.execute(
        query=query,
        variables={
            "projectId": project_gid,
            "timeRange": {
                "start": (base_time + timedelta(minutes=10)).isoformat(),
                "end": (base_time + timedelta(minutes=20)).isoformat(),
            },
        },
    )
    assert not response.errors
    assert response.data is not None
    summary = response.data["node"]["sessionAnnotationSummary"]
    assert summary is not None
    assert summary["count"] == 1
    assert summary["labelFractions"] == [{"label": "long_running", "fraction": 1.0}]


@dataclass
class _TimeRangeSessionsData:
    project: models.Project
    sessions_by_name: dict[str, models.ProjectSession]


class TestProjectSessionsTimeRange:
    """The sessions connection uses interval-overlap semantics for timeRange:
    a session is included iff [start_time, end_time] intersects [start, end)."""

    _QUERY = """
        query ($projectId: ID!, $timeRange: TimeRange, $sessionFilterCondition: String) {
          node(id: $projectId) {
            ... on Project {
              sessions(
                first: 50
                timeRange: $timeRange
                sessionFilterCondition: $sessionFilterCondition
              ) {
                edges { node { id } }
              }
            }
          }
        }
    """

    _BASE_TIME = datetime.fromisoformat("2024-01-01T00:00:00+00:00")
    _WINDOW_START = _BASE_TIME + timedelta(minutes=10)
    _WINDOW_END = _BASE_TIME + timedelta(minutes=20)

    @classmethod
    def _minutes(cls, minute: int) -> datetime:
        return cls._BASE_TIME + timedelta(minutes=minute)

    @pytest.fixture
    async def _sessions_data(
        self,
        db: DbSessionFactory,
    ) -> _TimeRangeSessionsData:
        minutes = self._minutes
        intervals = {
            # started before the window, last activity inside it (the long-running case)
            "long_running": (0, 15),
            # entirely before the window
            "before_window": (0, 5),
            # last activity exactly at the window start (closed lower bound)
            "ends_at_window_start": (0, 10),
            # entirely inside the window
            "inside_window": (12, 13),
            # spans the entire window
            "spans_window": (5, 25),
            # starts exactly at the window end (right-exclusive upper bound)
            "starts_at_window_end": (20, 30),
        }
        sessions_by_name = {}
        async with db() as session:
            project = await _add_project(session)
            for name, (start_minute, end_minute) in intervals.items():
                project_session = await _add_project_session(
                    session,
                    project,
                    start_time=minutes(start_minute),
                    end_time=minutes(end_minute),
                )
                sessions_by_name[name] = project_session
                trace = await _add_trace(
                    session,
                    project,
                    project_session,
                    start_time=minutes(start_minute),
                    end_time=minutes(start_minute) + timedelta(seconds=30),
                )
                await _add_span(
                    session,
                    trace,
                    start_time=minutes(start_minute),
                    end_time=minutes(start_minute) + timedelta(seconds=30),
                    attributes={"input": {"value": f"input for {name}"}},
                )
        return _TimeRangeSessionsData(project=project, sessions_by_name=sessions_by_name)

    async def _get_session_ids(
        self,
        gql_client: AsyncGraphQLClient,
        data: _TimeRangeSessionsData,
        time_range: Optional[dict[str, str]],
        session_filter_condition: Optional[str] = None,
    ) -> set[str]:
        response = await gql_client.execute(
            query=self._QUERY,
            variables={
                "projectId": str(GlobalID(Project.__name__, str(data.project.id))),
                "timeRange": time_range,
                "sessionFilterCondition": session_filter_condition,
            },
        )
        assert not response.errors
        assert response.data is not None
        edges = response.data["node"]["sessions"]["edges"]
        return {edge["node"]["id"] for edge in edges}

    def _expected_ids(self, data: _TimeRangeSessionsData, *names: str) -> set[str]:
        return {_gid(data.sessions_by_name[name]) for name in names}

    async def test_includes_sessions_overlapping_the_window(
        self,
        _sessions_data: _TimeRangeSessionsData,
        gql_client: AsyncGraphQLClient,
    ) -> None:
        time_range = {
            "start": self._WINDOW_START.isoformat(),
            "end": self._WINDOW_END.isoformat(),
        }
        actual = await self._get_session_ids(gql_client, _sessions_data, time_range)
        assert actual == self._expected_ids(
            _sessions_data,
            "long_running",
            "ends_at_window_start",
            "inside_window",
            "spans_window",
        )

    async def test_open_ended_lower_bound_filters_by_last_activity(
        self,
        _sessions_data: _TimeRangeSessionsData,
        gql_client: AsyncGraphQLClient,
    ) -> None:
        time_range = {"start": self._WINDOW_START.isoformat()}
        actual = await self._get_session_ids(gql_client, _sessions_data, time_range)
        assert actual == self._expected_ids(
            _sessions_data,
            "long_running",
            "ends_at_window_start",
            "inside_window",
            "spans_window",
            "starts_at_window_end",
        )

    async def test_open_ended_upper_bound_filters_by_start_time(
        self,
        _sessions_data: _TimeRangeSessionsData,
        gql_client: AsyncGraphQLClient,
    ) -> None:
        time_range = {"end": self._WINDOW_END.isoformat()}
        actual = await self._get_session_ids(gql_client, _sessions_data, time_range)
        assert actual == self._expected_ids(
            _sessions_data,
            "long_running",
            "before_window",
            "ends_at_window_start",
            "inside_window",
            "spans_window",
        )

    async def test_content_filter_matches_traces_outside_the_window(
        self,
        _sessions_data: _TimeRangeSessionsData,
        gql_client: AsyncGraphQLClient,
    ) -> None:
        # The long-running session's only trace starts before the window; the session
        # must still be returned because it overlaps the window and its content matches.
        time_range = {
            "start": self._WINDOW_START.isoformat(),
            "end": self._WINDOW_END.isoformat(),
        }
        actual = await self._get_session_ids(
            gql_client,
            _sessions_data,
            time_range,
            session_filter_condition="'input for long_running' in any_input",
        )
        assert actual == self._expected_ids(_sessions_data, "long_running")

    async def test_content_filter_does_not_widen_the_window(
        self,
        _sessions_data: _TimeRangeSessionsData,
        gql_client: AsyncGraphQLClient,
    ) -> None:
        # A session outside the window stays excluded even if its content matches.
        time_range = {
            "start": self._WINDOW_START.isoformat(),
            "end": self._WINDOW_END.isoformat(),
        }
        actual = await self._get_session_ids(
            gql_client,
            _sessions_data,
            time_range,
            session_filter_condition="'input for before_window' in any_input",
        )
        assert actual == set()

    async def test_dsl_start_time_compares_points_where_time_range_tests_overlap(
        self,
        _sessions_data: _TimeRangeSessionsData,
        gql_client: AsyncGraphQLClient,
    ) -> None:
        """`timeRange` selects by overlap; a DSL `start_time` cutoff compares a point.

        `long_running` began before the window and is still active inside it, so overlap
        includes it and a cutoff on its own start does not. Composing the two ANDs, which is
        how the same session can be both in the window and out of the filter.
        """
        time_range = {
            "start": self._WINDOW_START.isoformat(),
            "end": self._WINDOW_END.isoformat(),
        }
        by_overlap = await self._get_session_ids(gql_client, _sessions_data, time_range)
        assert _gid(_sessions_data.sessions_by_name["long_running"]) in by_overlap

        cutoff = self._WINDOW_START.isoformat()
        assert cutoff.endswith("+00:00")
        with_cutoff = await self._get_session_ids(
            gql_client,
            _sessions_data,
            time_range,
            session_filter_condition=f"start_time >= '{cutoff}'",
        )
        assert with_cutoff == self._expected_ids(_sessions_data, "inside_window")
