diff --git a/openapi/ga/individual/platform.openapi.yaml b/openapi/ga/individual/platform.openapi.yaml index fc1a6631af..ae5a5bf274 100644 --- a/openapi/ga/individual/platform.openapi.yaml +++ b/openapi/ga/individual/platform.openapi.yaml @@ -3631,7 +3631,7 @@ paths: /apis/intake/v2/workspaces/{workspace}/evaluations/{name}/sessions: get: tags: - - Evaluations + - Experiments summary: List Evaluation Sessions operationId: list_evaluation_sessions_apis_intake_v2_workspaces__workspace__evaluations__name__sessions_get parameters: @@ -3685,6 +3685,22 @@ paths: description: Response mode. summary omits root-span input and output; preview includes both truncated to 300 characters; detailed returns full root-span payloads. + - name: sort + in: query + required: false + schema: + description: "Comma-separated list of fields to sort by, applied in order\ + \ (the first field dominates); prefix a field with '-' for descending\ + \ \u2014 e.g. '-cost_total_usd,latency_ms'. Fields: test_case_id, started_at,\ + \ ended_at, latency_ms, status, cost_total_usd, tokens. When omitted,\ + \ sessions are ordered by started_at ascending." + title: Sort + type: string + description: "Comma-separated list of fields to sort by, applied in order\ + \ (the first field dominates); prefix a field with '-' for descending \u2014\ + \ e.g. '-cost_total_usd,latency_ms'. Fields: test_case_id, started_at, ended_at,\ + \ latency_ms, status, cost_total_usd, tokens. When omitted, sessions are\ + \ ordered by started_at ascending." - in: query name: filter style: deepObject @@ -3701,9 +3717,11 @@ paths: schema: $ref: '#/components/schemas/EvaluationSessionResponsesPage' '400': - description: Invalid filter value + description: Invalid filter value, unsupported sort field, or empty sort '404': description: Evaluation not found + '413': + description: Too many sessions to sort by cost or tokens '503': description: ClickHouse unavailable '422': diff --git a/openapi/ga/openapi.yaml b/openapi/ga/openapi.yaml index fc1a6631af..ae5a5bf274 100644 --- a/openapi/ga/openapi.yaml +++ b/openapi/ga/openapi.yaml @@ -3631,7 +3631,7 @@ paths: /apis/intake/v2/workspaces/{workspace}/evaluations/{name}/sessions: get: tags: - - Evaluations + - Experiments summary: List Evaluation Sessions operationId: list_evaluation_sessions_apis_intake_v2_workspaces__workspace__evaluations__name__sessions_get parameters: @@ -3685,6 +3685,22 @@ paths: description: Response mode. summary omits root-span input and output; preview includes both truncated to 300 characters; detailed returns full root-span payloads. + - name: sort + in: query + required: false + schema: + description: "Comma-separated list of fields to sort by, applied in order\ + \ (the first field dominates); prefix a field with '-' for descending\ + \ \u2014 e.g. '-cost_total_usd,latency_ms'. Fields: test_case_id, started_at,\ + \ ended_at, latency_ms, status, cost_total_usd, tokens. When omitted,\ + \ sessions are ordered by started_at ascending." + title: Sort + type: string + description: "Comma-separated list of fields to sort by, applied in order\ + \ (the first field dominates); prefix a field with '-' for descending \u2014\ + \ e.g. '-cost_total_usd,latency_ms'. Fields: test_case_id, started_at, ended_at,\ + \ latency_ms, status, cost_total_usd, tokens. When omitted, sessions are\ + \ ordered by started_at ascending." - in: query name: filter style: deepObject @@ -3701,9 +3717,11 @@ paths: schema: $ref: '#/components/schemas/EvaluationSessionResponsesPage' '400': - description: Invalid filter value + description: Invalid filter value, unsupported sort field, or empty sort '404': description: Evaluation not found + '413': + description: Too many sessions to sort by cost or tokens '503': description: ClickHouse unavailable '422': diff --git a/openapi/openapi.yaml b/openapi/openapi.yaml index fc1a6631af..ae5a5bf274 100644 --- a/openapi/openapi.yaml +++ b/openapi/openapi.yaml @@ -3631,7 +3631,7 @@ paths: /apis/intake/v2/workspaces/{workspace}/evaluations/{name}/sessions: get: tags: - - Evaluations + - Experiments summary: List Evaluation Sessions operationId: list_evaluation_sessions_apis_intake_v2_workspaces__workspace__evaluations__name__sessions_get parameters: @@ -3685,6 +3685,22 @@ paths: description: Response mode. summary omits root-span input and output; preview includes both truncated to 300 characters; detailed returns full root-span payloads. + - name: sort + in: query + required: false + schema: + description: "Comma-separated list of fields to sort by, applied in order\ + \ (the first field dominates); prefix a field with '-' for descending\ + \ \u2014 e.g. '-cost_total_usd,latency_ms'. Fields: test_case_id, started_at,\ + \ ended_at, latency_ms, status, cost_total_usd, tokens. When omitted,\ + \ sessions are ordered by started_at ascending." + title: Sort + type: string + description: "Comma-separated list of fields to sort by, applied in order\ + \ (the first field dominates); prefix a field with '-' for descending \u2014\ + \ e.g. '-cost_total_usd,latency_ms'. Fields: test_case_id, started_at, ended_at,\ + \ latency_ms, status, cost_total_usd, tokens. When omitted, sessions are\ + \ ordered by started_at ascending." - in: query name: filter style: deepObject @@ -3701,9 +3717,11 @@ paths: schema: $ref: '#/components/schemas/EvaluationSessionResponsesPage' '400': - description: Invalid filter value + description: Invalid filter value, unsupported sort field, or empty sort '404': description: Evaluation not found + '413': + description: Too many sessions to sort by cost or tokens '503': description: ClickHouse unavailable '422': diff --git a/sdk/python/nemo-platform/.github/workflows/ci.yml b/sdk/python/nemo-platform/.github/workflows/ci.yml index 558389ad06..7c53993236 100644 --- a/sdk/python/nemo-platform/.github/workflows/ci.yml +++ b/sdk/python/nemo-platform/.github/workflows/ci.yml @@ -18,7 +18,7 @@ jobs: lint: timeout-minutes: 10 name: lint - runs-on: ${{ github.repository == 'stainless-sdks/nemo-platform-python' && 'depot-ubuntu-24.04' || 'ubuntu-latest' }} + runs-on: ${{ startsWith(github.repository, 'stainless-sdks/') && 'depot-ubuntu-24.04' || 'ubuntu-latest' }} if: (github.event_name == 'push' || github.event.pull_request.head.repo.fork) && (github.event_name != 'push' || github.event.head_commit.message != 'codegen metadata') steps: - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 @@ -41,7 +41,7 @@ jobs: permissions: contents: read id-token: write - runs-on: ${{ github.repository == 'stainless-sdks/nemo-platform-python' && 'depot-ubuntu-24.04' || 'ubuntu-latest' }} + runs-on: ${{ startsWith(github.repository, 'stainless-sdks/') && 'depot-ubuntu-24.04' || 'ubuntu-latest' }} steps: - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 @@ -78,7 +78,7 @@ jobs: test: timeout-minutes: 10 name: test - runs-on: ${{ github.repository == 'stainless-sdks/nemo-platform-python' && 'depot-ubuntu-24.04' || 'ubuntu-latest' }} + runs-on: ${{ startsWith(github.repository, 'stainless-sdks/') && 'depot-ubuntu-24.04' || 'ubuntu-latest' }} if: github.event_name == 'push' || github.event.pull_request.head.repo.fork steps: - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 diff --git a/sdk/python/nemo-platform/.nmpcontext/openapi.yaml b/sdk/python/nemo-platform/.nmpcontext/openapi.yaml index fc1a6631af..ae5a5bf274 100644 --- a/sdk/python/nemo-platform/.nmpcontext/openapi.yaml +++ b/sdk/python/nemo-platform/.nmpcontext/openapi.yaml @@ -3631,7 +3631,7 @@ paths: /apis/intake/v2/workspaces/{workspace}/evaluations/{name}/sessions: get: tags: - - Evaluations + - Experiments summary: List Evaluation Sessions operationId: list_evaluation_sessions_apis_intake_v2_workspaces__workspace__evaluations__name__sessions_get parameters: @@ -3685,6 +3685,22 @@ paths: description: Response mode. summary omits root-span input and output; preview includes both truncated to 300 characters; detailed returns full root-span payloads. + - name: sort + in: query + required: false + schema: + description: "Comma-separated list of fields to sort by, applied in order\ + \ (the first field dominates); prefix a field with '-' for descending\ + \ \u2014 e.g. '-cost_total_usd,latency_ms'. Fields: test_case_id, started_at,\ + \ ended_at, latency_ms, status, cost_total_usd, tokens. When omitted,\ + \ sessions are ordered by started_at ascending." + title: Sort + type: string + description: "Comma-separated list of fields to sort by, applied in order\ + \ (the first field dominates); prefix a field with '-' for descending \u2014\ + \ e.g. '-cost_total_usd,latency_ms'. Fields: test_case_id, started_at, ended_at,\ + \ latency_ms, status, cost_total_usd, tokens. When omitted, sessions are\ + \ ordered by started_at ascending." - in: query name: filter style: deepObject @@ -3701,9 +3717,11 @@ paths: schema: $ref: '#/components/schemas/EvaluationSessionResponsesPage' '400': - description: Invalid filter value + description: Invalid filter value, unsupported sort field, or empty sort '404': description: Evaluation not found + '413': + description: Too many sessions to sort by cost or tokens '503': description: ClickHouse unavailable '422': diff --git a/sdk/python/nemo-platform/src/nemo_platform/resources/evaluations/sessions.py b/sdk/python/nemo-platform/src/nemo_platform/resources/evaluations/sessions.py index 2f0d35701c..25045abb08 100644 --- a/sdk/python/nemo-platform/src/nemo_platform/resources/evaluations/sessions.py +++ b/sdk/python/nemo-platform/src/nemo_platform/resources/evaluations/sessions.py @@ -69,6 +69,7 @@ def list( mode: Literal["summary", "preview", "detailed"] | Omit = omit, page: int | Omit = omit, page_size: int | Omit = omit, + sort: str | Omit = omit, # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs. # The extra values given here take precedence over values defined on the client or passed to this method. extra_headers: Headers | None = None, @@ -89,6 +90,12 @@ def list( page_size: Page size. + sort: Comma-separated list of fields to sort by, applied in order (the first field + dominates); prefix a field with '-' for descending — e.g. + '-cost_total_usd,latency_ms'. Fields: test_case_id, started_at, ended_at, + latency_ms, status, cost_total_usd, tokens. When omitted, sessions are ordered + by started_at ascending. + extra_headers: Send extra headers extra_query: Add additional query parameters to the request @@ -119,6 +126,7 @@ def list( "mode": mode, "page": page, "page_size": page_size, + "sort": sort, }, session_list_params.SessionListParams, ), @@ -156,6 +164,7 @@ def list( mode: Literal["summary", "preview", "detailed"] | Omit = omit, page: int | Omit = omit, page_size: int | Omit = omit, + sort: str | Omit = omit, # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs. # The extra values given here take precedence over values defined on the client or passed to this method. extra_headers: Headers | None = None, @@ -176,6 +185,12 @@ def list( page_size: Page size. + sort: Comma-separated list of fields to sort by, applied in order (the first field + dominates); prefix a field with '-' for descending — e.g. + '-cost_total_usd,latency_ms'. Fields: test_case_id, started_at, ended_at, + latency_ms, status, cost_total_usd, tokens. When omitted, sessions are ordered + by started_at ascending. + extra_headers: Send extra headers extra_query: Add additional query parameters to the request @@ -206,6 +221,7 @@ def list( "mode": mode, "page": page, "page_size": page_size, + "sort": sort, }, session_list_params.SessionListParams, ), diff --git a/sdk/python/nemo-platform/src/nemo_platform/types/evaluations/session_list_params.py b/sdk/python/nemo-platform/src/nemo_platform/types/evaluations/session_list_params.py index 3adef9259a..75ebc5c5e5 100644 --- a/sdk/python/nemo-platform/src/nemo_platform/types/evaluations/session_list_params.py +++ b/sdk/python/nemo-platform/src/nemo_platform/types/evaluations/session_list_params.py @@ -42,3 +42,12 @@ class SessionListParams(TypedDict, total=False): page_size: int """Page size.""" + + sort: str + """ + Comma-separated list of fields to sort by, applied in order (the first field + dominates); prefix a field with '-' for descending — e.g. + '-cost_total_usd,latency_ms'. Fields: test_case_id, started_at, ended_at, + latency_ms, status, cost_total_usd, tokens. When omitted, sessions are ordered + by started_at ascending. + """ diff --git a/sdk/python/nemo-platform/tests/api_resources/evaluations/test_sessions.py b/sdk/python/nemo-platform/tests/api_resources/evaluations/test_sessions.py index dcababa5f1..ca09c4ad42 100644 --- a/sdk/python/nemo-platform/tests/api_resources/evaluations/test_sessions.py +++ b/sdk/python/nemo-platform/tests/api_resources/evaluations/test_sessions.py @@ -55,6 +55,7 @@ def test_method_list_with_all_params(self, client: NeMoPlatform) -> None: mode="summary", page=1, page_size=1, + sort="sort", ) assert_matches_type(SyncDefaultPagination[EvaluationSessionResponse], session, path=["response"]) @@ -129,6 +130,7 @@ async def test_method_list_with_all_params(self, async_client: AsyncNeMoPlatform mode="summary", page=1, page_size=1, + sort="sort", ) assert_matches_type(AsyncDefaultPagination[EvaluationSessionResponse], session, path=["response"]) diff --git a/services/intake/src/nmp/intake/api/v2/experiments/endpoints.py b/services/intake/src/nmp/intake/api/v2/experiments/endpoints.py index 03eadaa640..0d1b086341 100644 --- a/services/intake/src/nmp/intake/api/v2/experiments/endpoints.py +++ b/services/intake/src/nmp/intake/api/v2/experiments/endpoints.py @@ -51,7 +51,7 @@ EvaluationRollupRepository, ScoreRollup, ) -from nmp.intake.spans.evaluation_session_repository import EvaluationSessionRepository +from nmp.intake.spans.evaluation_session_repository import EvaluationSessionRepository, MetricSortTooLargeError from nmp.intake.spans.storage import make_pagination logger = logging.getLogger(__name__) @@ -72,6 +72,20 @@ def _sanitize_for_log(value: str) -> str: # sort by a ClickHouse rollup metric, not just entity columns. `sort` is therefore a free string, # validated against these: an entity column, run_count, or a `.` rollup path. _ENTITY_SORT_FIELDS = frozenset({"name", "created_at", "updated_at", "pinned_at"}) +# Sessions are sorted in ClickHouse (ORDER BY before LIMIT/OFFSET) so sort composes +# correctly with pagination. These are the allowed field names; each maps to an SQL +# expression in the repository - see _list_sql in evaluation_session_repository.py. +_SESSION_SORT_FIELDS = frozenset( + { + "test_case_id", + "started_at", + "ended_at", + "latency_ms", + "cost_total_usd", + "status", + "tokens", + } +) _METRIC_STATS = frozenset({"sum", "mean", "median", "p90", "p95", "p99", "count"}) # Per-group evaluation fetch bound for the in-memory merge. Groups are expected to hold at most # hundreds; a query that selects more than this is rejected rather than sorted on a partial set — the @@ -643,10 +657,10 @@ async def unpin_evaluation( @router.get( "/v2/workspaces/{workspace}/evaluations/{name}/sessions", response_model=Page[EvaluationSessionResponse], - tags=[EVALUATIONS_TAG], responses={ - 400: {"description": "Invalid filter value"}, + 400: {"description": "Invalid filter value, unsupported sort field, or empty sort"}, 404: {"description": "Evaluation not found"}, + 413: {"description": "Too many sessions to sort by cost or tokens"}, 503: {"description": "ClickHouse unavailable"}, }, openapi_extra=generate_openapi_extra_params( @@ -670,8 +684,18 @@ async def list_evaluation_sessions( "300 characters; detailed returns full root-span payloads." ), ), + sort: str | None = Query( + default=None, + description=( + "Comma-separated list of fields to sort by, applied in order (the first field dominates); " + "prefix a field with '-' for descending — e.g. '-cost_total_usd,latency_ms'. Fields: " + "test_case_id, started_at, ended_at, latency_ms, status, cost_total_usd, tokens. When omitted, " + "sessions are ordered by started_at ascending." + ), + ), ) -> Page[EvaluationSessionResponse]: validate_list_query_params(request, additional_params={"mode"}) + sort_keys = _parse_session_sort_keys(sort) if sort is not None else None evaluation = await _get_or_404( entity_client, Evaluation, @@ -703,7 +727,18 @@ async def list_evaluation_sessions( page=page, page_size=page_size, mode=mode, + sort_keys=sort_keys, ) + except MetricSortTooLargeError as exc: + raise HTTPException( + status_code=status.HTTP_413_REQUEST_ENTITY_TOO_LARGE, + detail=( + f"This query selects {exc.total} sessions, exceeding the maximum of " + f"{exc.limit} that can be sorted by cost or tokens in one request. " + "Narrow the result with a filter (e.g. filter[status]=success) or sort by a " + "different field (started_at, latency_ms, status, test_case_id)." + ), + ) from exc except Exception as exc: # Sessions are the response payload (not enrichment), so we can't silently degrade like # _hydrate_rollups does. Convert backend failures (ClickHouse connection drop, query @@ -1024,6 +1059,30 @@ def _parse_sort_keys(sort: str) -> tuple[list[tuple[str, bool]], bool]: return sort_keys, explicit_metric_sort +def _parse_session_sort_keys(sort: str) -> list[tuple[str, bool]]: + """Parse the comma-separated ``sort`` param into an ordered list of ``(field, descending)`` keys. + + Each field may be '-'-prefixed for descending; the keys are applied in order (the first field + dominates). Returns the keys. Raises 400 if a field is unsupported or the list is empty. + """ + sort_keys: list[tuple[str, bool]] = [] + for token in sort.split(","): + field_token = token.strip() + if not field_token: + continue + descending = field_token.startswith("-") + sort_field = field_token[1:] if descending else field_token + if sort_field not in _SESSION_SORT_FIELDS: + raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail=f"Unsupported sort field: {sort_field}") + sort_keys.append((sort_field, descending)) + if not sort_keys: + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail="The 'sort' parameter must contain at least one field.", + ) + return sort_keys + + def _is_metric_field(field: str) -> bool: """True if `field` is *intended* as a rollup metric (by head), valid path or not. diff --git a/services/intake/src/nmp/intake/spans/evaluation_session_repository.py b/services/intake/src/nmp/intake/spans/evaluation_session_repository.py index 0510130e66..b809c82589 100644 --- a/services/intake/src/nmp/intake/spans/evaluation_session_repository.py +++ b/services/intake/src/nmp/intake/spans/evaluation_session_repository.py @@ -28,6 +28,67 @@ ) from nmp.intake.spans.trace_repository import current_spans_sql +# Sort fields that require a pre-pagination spans join to compute. These values live in the +# `spans` table (not `trace_index`), so they don't exist until after the session_metrics join. +# Sorting by them globally requires computing them for ALL sessions before LIMIT/OFFSET runs — +# see `_pre_page_metrics_cte_sql` and the conditional in `_list_sql`. +_PRE_METRICS_SORT_FIELDS = frozenset({"cost_total_usd", "tokens"}) + +# Maps each API sort field to its SQL expression used in _metric_sort_page_ids_sql (the +# pre-metrics sort query). Simple fields come from scoped_sessions (aliased as `s`); +# pre-metrics fields reference the pre_page_metrics CTE (aliased as `pm`). This map is +# NOT used by _list_sql — that path only handles trace_index columns and uses plain names. +_SORT_EXPR_PAGE: dict[str, str] = { + "started_at": "start_time", + "ended_at": "end_time", + "latency_ms": "latency_ms", + "status": "root_span_status", + "test_case_id": "test_case_id", + "cost_total_usd": "pm.cost_total_usd", + "tokens": "pm.total_tokens", +} + +# Same mapping for the final SELECT's ORDER BY. At that point scoped_sessions columns +# are behind the `sessions` alias, and cost/tokens come from the `metrics` CTE +# (session_metrics, computed for just the page). +_SORT_EXPR_FINAL: dict[str, str] = { + "started_at": "sessions.start_time", + "ended_at": "sessions.end_time", + "latency_ms": "sessions.latency_ms", + "status": "sessions.root_span_status", + "test_case_id": "sessions.test_case_id", + "cost_total_usd": "metrics.cost_total_usd", + # Preserve NULL only when both input and output are absent (no span data at all). + # If only one side is NULL (e.g. a failed call with no output tokens but real input tokens), + # coalesce to 0 so the session sorts by its actual partial usage rather than disappearing + # to the bottom via NULLS LAST. + "tokens": ( + "if(metrics.input_tokens IS NULL AND metrics.output_tokens IS NULL, NULL, " + "coalesce(metrics.input_tokens, 0) + coalesce(metrics.output_tokens, 0))" + ), +} + + +# Sessions above this threshold will not be sorted by cost or tokens. Pre-metrics sort joins +# spans for EVERY scoped session before paginating — an unbounded aggregation on the request +# path. Mirror the evaluations list cap (_MAX_GROUP_EVALUATIONS) to prevent runaway queries. +# Raise this if legitimate evaluations routinely exceed it; the right long-term fix is to +# denormalise cost/tokens into trace_index so no pre-pagination join is needed. +_MAX_METRIC_SORT_SESSIONS = 10_000 + + +class MetricSortTooLargeError(Exception): + """Raised when a cost/tokens sort is requested on more sessions than the pre-metrics cap allows. + + This is a domain exception (not HTTPException) so the repository stays HTTP-agnostic. + The endpoint catches it and converts it to 413. + """ + + def __init__(self, total: int, limit: int) -> None: + self.total = total + self.limit = limit + super().__init__(f"Metric sort requested on {total} sessions, limit is {limit}") + @dataclass(frozen=True) class EvaluationSessionRow: @@ -72,6 +133,7 @@ async def list_sessions( page: int, page_size: int, mode: IntakeResponseMode, + sort_keys: list[tuple[str, bool]] | None = None, ) -> EvaluationSessionPage: trace_index_table = self._client.table("trace_index") spans_table = self._client.table("spans") @@ -100,26 +162,86 @@ async def list_sessions( if total == 0: return EvaluationSessionPage(rows=[], total=0) + needs_pre_metrics = sort_keys is not None and any(f in _PRE_METRICS_SORT_FIELDS for f, _ in sort_keys) + if needs_pre_metrics and total > _MAX_METRIC_SORT_SESSIONS: + raise MetricSortTooLargeError(total, _MAX_METRIC_SORT_SESSIONS) + offset = (page - 1) * page_size - list_sql = _list_sql( - trace_index_table=trace_index_table, - spans_table=spans_table, - evaluator_results_table=evaluator_results_table, - scoped_filter_sql=scoped_filter_sql, - mode=mode, - ) - list_parameters = { - **base_parameters, - **scoped_filter_parameters, - **text_query_parameters(mode), - "limit": page_size, - "offset": offset, - } - list_result = await self._client.query( - list_sql, - parameters=list_parameters, - ) - rows = [_row(record) for record in result_rows(list_result)] + + if needs_pre_metrics: + # Two-query path for cost/tokens sorts. + # + # ClickHouse 24.3 inlines CTEs (does not materialise them), so a single query that + # references `page_sessions` from multiple downstream CTEs (current_page_spans, + # session_metrics, session_scores, final SELECT) would re-execute the expensive + # all-session span aggregation once per reference. Splitting into two queries + # ensures the aggregation runs exactly once. + # + # Query 2: compute the sorted page — aggregate cost/tokens once across all sessions, + # apply ORDER BY + LIMIT/OFFSET, return ordered (workspace, session_id) pairs only. + page_ids_sql = _metric_sort_page_ids_sql( + trace_index_table=trace_index_table, + spans_table=spans_table, + scoped_filter_sql=scoped_filter_sql, + sort_keys=sort_keys, # type: ignore[arg-type] # guaranteed non-None here + ) + page_ids_result = await self._client.query( + page_ids_sql, + parameters={ + **base_parameters, + **scoped_filter_parameters, + "limit": page_size, + "offset": offset, + }, + ) + ordered_ids = [record["session_id"] for record in result_rows(page_ids_result)] + + if not ordered_ids: + return EvaluationSessionPage(rows=[], total=total) + + # Query 3: hydrate those session IDs — fetch session fields, page-only span metrics, + # and evaluator scores. Returns rows in unspecified order; we restore Query 2's + # ordering in Python after the fetch. + hydrate_sql = _hydrate_by_ids_sql( + trace_index_table=trace_index_table, + spans_table=spans_table, + evaluator_results_table=evaluator_results_table, + mode=mode, + ) + hydrate_result = await self._client.query( + hydrate_sql, + parameters={ + **base_parameters, + **text_query_parameters(mode), + "session_ids": ordered_ids, + }, + ) + rows_by_id = {record["session_id"]: _row(record) for record in result_rows(hydrate_result)} + # Restore the order from Query 2. `rows_by_id` may be missing a session_id if a + # race caused the trace_index to disagree between queries, so guard with `if sid in`. + rows = [rows_by_id[sid] for sid in ordered_ids if sid in rows_by_id] + else: + # Single-query path for trace_index column sorts (started_at, latency_ms, etc.). + # No pre-metrics join needed; ClickHouse CTE inlining is harmless here because + # page_sessions doesn't depend on an expensive aggregation. + list_sql = _list_sql( + trace_index_table=trace_index_table, + spans_table=spans_table, + evaluator_results_table=evaluator_results_table, + scoped_filter_sql=scoped_filter_sql, + mode=mode, + sort_keys=sort_keys or [], + ) + list_parameters = { + **base_parameters, + **scoped_filter_parameters, + **text_query_parameters(mode), + "limit": page_size, + "offset": offset, + } + list_result = await self._client.query(list_sql, parameters=list_parameters) + rows = [_row(record) for record in result_rows(list_result)] + return EvaluationSessionPage(rows=rows, total=total) @@ -191,6 +313,210 @@ def _count_sql( """ +def _build_order_by(sort_keys: list[tuple[str, bool]], expr_map: dict[str, str], tiebreaker: str) -> str: + """Build a comma-separated ORDER BY clause from an ordered list of (field, descending) pairs. + + Each field is looked up in expr_map to get its SQL expression in the current query scope + (column names differ between page_sessions and the final SELECT — callers pass the right map). + NULLS LAST on every user key so missing values sort at the end rather than the front. + Always appends tiebreaker as the final stable sort key so pages are deterministic when the + user's keys produce ties. + """ + parts = [f"{expr_map[field]} {'DESC' if desc else 'ASC'} NULLS LAST" for field, desc in sort_keys] + parts.append(tiebreaker) + return ", ".join(parts) + + +def _metric_sort_page_ids_sql( + *, + trace_index_table: str, + spans_table: str, + scoped_filter_sql: str, + sort_keys: list[tuple[str, bool]], +) -> str: + """Query 2 of the two-query pre-metrics sort path. + + Aggregates cost/tokens across ALL scoped sessions exactly once, then applies + ORDER BY + LIMIT/OFFSET to return the ordered (workspace, session_id) pairs for + the requested page. Only IDs are returned — row hydration is a separate query. + + Why separate: ClickHouse 24.3 inlines CTEs rather than materialising them, so a + single query that references `page_sessions` from multiple CTEs would re-execute + the expensive all-session span aggregation once per reference. Returning IDs here + and hydrating in _hydrate_by_ids_sql ensures the aggregation runs exactly once. + """ + # summary mode: text payloads are not needed for sorting + scoped_sessions_sql = _scoped_sessions_sql( + trace_index_table, + scoped_filter_sql=scoped_filter_sql, + mode="summary", + ) + all_scoped_spans = current_spans_sql( + spans_table, + extra_where_sql=( + "(span_versions.workspace, span_versions.session_id) IN (SELECT workspace, session_id FROM scoped_sessions)" + ), + ) + order_by = _build_order_by(sort_keys, _SORT_EXPR_PAGE, "s.root_span_id ASC") + return f""" + WITH + scoped_sessions AS ( + {scoped_sessions_sql} + ), + pre_page_metrics AS ( + SELECT + s.workspace AS workspace, + s.session_id AS session_id, + {_guarded_sum_sql("cost_key", scale=COST_SCALE)} AS cost_total_usd, + if( + {_guarded_sum_sql("input_tokens_key")} IS NULL AND {_guarded_sum_sql("output_tokens_key")} IS NULL, + NULL, + coalesce({_guarded_sum_sql("input_tokens_key")}, 0) + coalesce({_guarded_sum_sql("output_tokens_key")}, 0) + ) AS total_tokens + FROM scoped_sessions AS s + LEFT JOIN {all_scoped_spans} AS spans + ON s.workspace = spans.workspace + AND s.session_id = spans.session_id + AND spans.is_deleted = 0 + GROUP BY s.workspace, s.session_id + ) + SELECT s.workspace, s.session_id + FROM scoped_sessions AS s + LEFT JOIN pre_page_metrics AS pm + ON s.workspace = pm.workspace AND s.session_id = pm.session_id + ORDER BY {order_by} + LIMIT %(limit)s OFFSET %(offset)s + """ + + +def _hydrate_by_ids_sql( + *, + trace_index_table: str, + spans_table: str, + evaluator_results_table: str, + mode: IntakeResponseMode, +) -> str: + """Query 3 of the two-query pre-metrics sort path. + + Given a list of session_ids (passed as %(session_ids)s), fetches the full row for + each session: trace_index fields, per-page span metrics, and evaluator scores. + Returns rows in unspecified order — the caller restores the ordering from Query 2 + in Python after the fetch. + + No ORDER BY is emitted; no pre_page_metrics join is needed (the aggregation here + is bounded to the page, which is at most 1000 rows). + """ + select_columns = [ + "workspace", + "evaluation_id", + "session_id", + "test_case_id", + "trace_id", + "root_span_id", + "root_started_at AS start_time", + "root_ended_at AS end_time", + "latency_ms", + "root_status AS root_span_status", + ] + select_columns.extend( + ( + text_select_for_mode("root_input", alias="input", mode=mode), + text_select_for_mode("root_output", alias="output", mode=mode), + ) + ) + select_sql = ",\n ".join(select_columns) + + page_spans = current_spans_sql( + spans_table, + extra_where_sql=( + "(span_versions.workspace, span_versions.session_id) IN (SELECT workspace, session_id FROM page_sessions)" + ), + ) + + return f""" + WITH + page_sessions AS ( + SELECT + {select_sql} + FROM {trace_index_table} FINAL + WHERE workspace = %(workspace)s + AND is_deleted = 0 + AND evaluation_id = %(evaluation_name)s + AND session_id IN %(session_ids)s + LIMIT 1 BY workspace, session_id, evaluation_id + ), + current_page_spans AS ( + {page_spans} + ), + session_metrics AS ( + SELECT + sessions.workspace AS workspace, + sessions.session_id AS session_id, + {_guarded_sum_sql("input_tokens_key")} AS input_tokens, + {_guarded_sum_sql("output_tokens_key")} AS output_tokens, + {_guarded_sum_sql("cached_tokens_key")} AS cached_tokens, + {_guarded_sum_sql("cost_key", scale=COST_SCALE)} AS cost_total_usd + FROM page_sessions AS sessions + LEFT JOIN current_page_spans AS spans + ON sessions.workspace = spans.workspace + AND sessions.session_id = spans.session_id + AND spans.is_deleted = 0 + GROUP BY sessions.workspace, sessions.session_id + ), + session_scores AS ( + SELECT + workspace, + session_id, + mapFromArrays(groupArray(evaluator_name), groupArray(mean_score)) AS evaluator_scores + FROM ( + SELECT + results.workspace AS workspace, + results.session_id AS session_id, + results.name AS evaluator_name, + avg(results.value) AS mean_score + FROM ( + SELECT workspace, session_id, name, value + FROM {evaluator_results_table} FINAL + WHERE workspace = %(workspace)s + AND (workspace, session_id) IN ( + SELECT workspace, session_id + FROM page_sessions + ) + AND data_type IN ('NUMERIC', 'BOOLEAN') + AND value IS NOT NULL + ) AS results + GROUP BY results.workspace, results.session_id, results.name + ) + GROUP BY workspace, session_id + ) + SELECT + sessions.workspace AS workspace, + sessions.evaluation_id AS evaluation_id, + sessions.session_id AS session_id, + sessions.test_case_id AS test_case_id, + sessions.trace_id AS trace_id, + sessions.root_span_id AS root_span_id, + sessions.start_time AS start_time, + sessions.end_time AS end_time, + sessions.latency_ms AS latency_ms, + sessions.root_span_status AS root_span_status, + sessions.input AS input, + sessions.output AS output, + metrics.input_tokens AS input_tokens, + metrics.output_tokens AS output_tokens, + metrics.cached_tokens AS cached_tokens, + metrics.cost_total_usd AS cost_total_usd, + scores.evaluator_scores AS evaluator_scores + FROM page_sessions AS sessions + LEFT JOIN session_metrics AS metrics + ON sessions.workspace = metrics.workspace + AND sessions.session_id = metrics.session_id + LEFT JOIN session_scores AS scores + ON sessions.workspace = scores.workspace + AND sessions.session_id = scores.session_id + """ + + def _list_sql( *, trace_index_table: str, @@ -198,12 +524,38 @@ def _list_sql( evaluator_results_table: str, scoped_filter_sql: str, mode: IntakeResponseMode, + sort_keys: list[tuple[str, bool]], ) -> str: + """Single-query path for trace_index column sorts (started_at, latency_ms, etc.). + + Only called when no sort field requires the pre-metrics spans join (i.e. no cost/tokens + sort). ClickHouse CTE inlining is harmless here because page_sessions doesn't depend on + an expensive all-session aggregation. + """ scoped_sessions_sql = _scoped_sessions_sql( trace_index_table, scoped_filter_sql=scoped_filter_sql, mode=mode, ) + # Empty sort_keys means no sort param — preserve the original default order. + page_order_by = ( + _build_order_by(sort_keys, _SORT_EXPR_PAGE, "root_span_id ASC") + if sort_keys + else "start_time ASC, root_span_id ASC" + ) + # The final SELECT re-orders already-paginated rows as they emerge from the CTE joins. + # ClickHouse does not guarantee CTE output order, so this ORDER BY must match page_order_by. + final_order_by = ( + _build_order_by(sort_keys, _SORT_EXPR_FINAL, "sessions.root_span_id ASC") + if sort_keys + else "sessions.start_time ASC, sessions.root_span_id ASC" + ) + current_page_spans = current_spans_sql( + spans_table, + extra_where_sql=( + "(span_versions.workspace, span_versions.session_id) IN (SELECT workspace, session_id FROM page_sessions)" + ), + ) return f""" WITH scoped_sessions AS ( @@ -224,19 +576,11 @@ def _list_sql( input, output FROM scoped_sessions - ORDER BY start_time ASC, root_span_id ASC + ORDER BY {page_order_by} LIMIT %(limit)s OFFSET %(offset)s ), current_page_spans AS ( - { - current_spans_sql( - spans_table, - extra_where_sql=( - "(span_versions.workspace, span_versions.session_id) IN " - "(SELECT workspace, session_id FROM page_sessions)" - ), - ) - } + {current_page_spans} ), session_metrics AS ( SELECT @@ -304,7 +648,7 @@ def _list_sql( LEFT JOIN session_scores AS scores ON sessions.workspace = scores.workspace AND sessions.session_id = scores.session_id - ORDER BY sessions.start_time ASC, sessions.root_span_id ASC + ORDER BY {final_order_by} """ diff --git a/services/intake/tests/test_evaluation_session_clickhouse_repository.py b/services/intake/tests/test_evaluation_session_clickhouse_repository.py index abc6e087ff..fd5a4f1a3e 100644 --- a/services/intake/tests/test_evaluation_session_clickhouse_repository.py +++ b/services/intake/tests/test_evaluation_session_clickhouse_repository.py @@ -3,7 +3,35 @@ """Evaluation session ClickHouse query tests.""" -from nmp.intake.spans.evaluation_session_repository import _count_sql, _list_sql +from nmp.intake.spans.evaluation_session_repository import ( + _SORT_EXPR_FINAL, + _SORT_EXPR_PAGE, + _build_order_by, + _count_sql, + _hydrate_by_ids_sql, + _list_sql, + _metric_sort_page_ids_sql, +) + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _make_list_sql(**kwargs) -> str: + """Call _list_sql with fixed table names; only pass what the test cares about.""" + return _list_sql( + trace_index_table="trace_index", + spans_table="spans", + evaluator_results_table="evaluator_results", + scoped_filter_sql="", + **kwargs, + ) + + +# --------------------------------------------------------------------------- +# Existing payload-mode tests — updated to pass the now-required sort_keys arg +# --------------------------------------------------------------------------- def test_session_count_query_does_not_read_root_payloads() -> None: @@ -14,13 +42,7 @@ def test_session_count_query_does_not_read_root_payloads() -> None: def test_session_preview_query_truncates_input_and_output_in_clickhouse() -> None: - query = _list_sql( - trace_index_table="trace_index", - spans_table="spans", - evaluator_results_table="evaluator_results", - scoped_filter_sql="", - mode="preview", - ) + query = _make_list_sql(mode="preview", sort_keys=[]) assert "substringUTF8(root_input, 1, %(payload_char_limit)s) AS input" in query assert "substringUTF8(root_output, 1, %(payload_char_limit)s) AS output" in query @@ -29,13 +51,7 @@ def test_session_preview_query_truncates_input_and_output_in_clickhouse() -> Non def test_session_summary_query_omits_input_and_output_columns() -> None: - query = _list_sql( - trace_index_table="trace_index", - spans_table="spans", - evaluator_results_table="evaluator_results", - scoped_filter_sql="", - mode="summary", - ) + query = _make_list_sql(mode="summary", sort_keys=[]) assert "root_input" not in query assert "root_output" not in query @@ -44,14 +60,153 @@ def test_session_summary_query_omits_input_and_output_columns() -> None: def test_session_detailed_query_reads_full_input_and_output() -> None: - query = _list_sql( + query = _make_list_sql(mode="detailed", sort_keys=[]) + + assert "root_input AS input" in query + assert "root_output AS output" in query + assert "substringUTF8(root_input" not in query + + +# --------------------------------------------------------------------------- +# Sort: default (no sort_keys) +# The original ORDER BY must be preserved so existing behaviour is unchanged. +# --------------------------------------------------------------------------- + + +def test_default_sort_preserves_original_order() -> None: + # sort_keys=[] means no sort param was sent — fall back to start_time ASC. + query = _make_list_sql(mode="summary", sort_keys=[]) + + # Both page_sessions and the final SELECT must use the default order. + assert "ORDER BY start_time ASC, root_span_id ASC" in query + assert "ORDER BY sessions.start_time ASC, sessions.root_span_id ASC" in query + # No pre_page_metrics CTE should be injected. + assert "pre_page_metrics" not in query + + +# --------------------------------------------------------------------------- +# Sort: single field (trace_index column — no pre-metrics join needed) +# --------------------------------------------------------------------------- + + +def test_single_field_sort_latency_desc() -> None: + query = _make_list_sql(mode="summary", sort_keys=[("latency_ms", True)]) + + # page_sessions ORDER BY should use the scoped_sessions alias (no prefix). + assert "ORDER BY latency_ms DESC NULLS LAST, root_span_id ASC" in query + # Final SELECT ORDER BY uses the sessions. prefix. + assert "ORDER BY sessions.latency_ms DESC NULLS LAST, sessions.root_span_id ASC" in query + # No pre_page_metrics needed for a trace_index column. + assert "pre_page_metrics" not in query + + +def test_single_field_sort_status_asc() -> None: + query = _make_list_sql(mode="summary", sort_keys=[("status", False)]) + + assert "ORDER BY root_span_status ASC NULLS LAST, root_span_id ASC" in query + assert "ORDER BY sessions.root_span_status ASC NULLS LAST, sessions.root_span_id ASC" in query + assert "pre_page_metrics" not in query + + +# --------------------------------------------------------------------------- +# Sort: multi-field +# --------------------------------------------------------------------------- + + +def test_multi_field_sort_applies_keys_in_order() -> None: + # Primary: cost DESC, tie-break: latency ASC. Cost requires the two-query path. + ids_query = _metric_sort_page_ids_sql( + trace_index_table="trace_index", + spans_table="spans", + scoped_filter_sql="", + sort_keys=[("cost_total_usd", True), ("latency_ms", False)], + ) + # The ids query orders by pm. for cost, plain column for latency, s.root_span_id tiebreaker. + assert "pm.cost_total_usd DESC NULLS LAST, latency_ms ASC NULLS LAST, s.root_span_id ASC" in ids_query + # Hydrate query has no ORDER BY — caller sorts in Python. + hydrate_query = _hydrate_by_ids_sql( trace_index_table="trace_index", spans_table="spans", evaluator_results_table="evaluator_results", + mode="summary", + ) + assert "ORDER BY" not in hydrate_query + + +# --------------------------------------------------------------------------- +# Sort: cost_total_usd and tokens go through the two-query path +# --------------------------------------------------------------------------- + + +def test_cost_sort_uses_two_query_path() -> None: + ids_query = _metric_sort_page_ids_sql( + trace_index_table="trace_index", + spans_table="spans", scoped_filter_sql="", - mode="detailed", + sort_keys=[("cost_total_usd", True)], + ) + # pre_page_metrics must appear before the final SELECT so the ORDER BY can reference pm. + pre_pos = ids_query.index("pre_page_metrics AS (") + select_pos = ids_query.index("SELECT s.workspace, s.session_id") + assert pre_pos < select_pos + assert "LEFT JOIN pre_page_metrics AS pm" in ids_query + # Hydrate uses session_ids IN list, no pre_page_metrics. + hydrate_query = _hydrate_by_ids_sql( + trace_index_table="trace_index", + spans_table="spans", + evaluator_results_table="evaluator_results", + mode="summary", ) + assert "pre_page_metrics" not in hydrate_query + assert "session_id IN %(session_ids)s" in hydrate_query - assert "root_input AS input" in query - assert "root_output AS output" in query - assert "substringUTF8(root_input" not in query + +def test_tokens_sort_uses_two_query_path() -> None: + ids_query = _metric_sort_page_ids_sql( + trace_index_table="trace_index", + spans_table="spans", + scoped_filter_sql="", + sort_keys=[("tokens", False)], + ) + assert "pre_page_metrics AS (" in ids_query + assert "pm.total_tokens ASC NULLS LAST" in ids_query + # NULL-safe coalesce in the pre_page_metrics total_tokens computation. + assert "coalesce" in ids_query + + +# --------------------------------------------------------------------------- +# Sort: stable tiebreaker is always appended +# --------------------------------------------------------------------------- + + +def test_tiebreaker_always_appended() -> None: + # Even a single-field sort should end with root_span_id for determinism. + query = _make_list_sql(mode="summary", sort_keys=[("started_at", False)]) + + assert "root_span_id ASC" in query + + +# --------------------------------------------------------------------------- +# _build_order_by unit tests +# --------------------------------------------------------------------------- + + +def test_build_order_by_single_asc() -> None: + result = _build_order_by([("latency_ms", False)], _SORT_EXPR_PAGE, "root_span_id ASC") + assert result == "latency_ms ASC NULLS LAST, root_span_id ASC" + + +def test_build_order_by_single_desc() -> None: + result = _build_order_by([("cost_total_usd", True)], _SORT_EXPR_PAGE, "s.root_span_id ASC") + assert result == "pm.cost_total_usd DESC NULLS LAST, s.root_span_id ASC" + + +def test_build_order_by_multi() -> None: + result = _build_order_by( + [("cost_total_usd", True), ("latency_ms", False)], + _SORT_EXPR_FINAL, + "sessions.root_span_id ASC", + ) + assert result == ( + "metrics.cost_total_usd DESC NULLS LAST, sessions.latency_ms ASC NULLS LAST, sessions.root_span_id ASC" + )