fix(evaluation): Limit payload size for prompts

backend/app/services/evaluations/prompt_improvement.py

@@ -8,6 +8,7 @@
 import copy
 import json
 import logging
+from collections import Counter
 from uuid import UUID
 
 import anthropic
@@ -25,6 +26,7 @@
 from app.crud.config.version import ConfigVersionCrud
 from app.crud.evaluations.core import get_evaluation_run_by_id
 from app.crud.evaluations.score import (
+    COSINE_SCORE_NAME,
     GROUND_TRUTH_SCORE_NAME,
     KNOWLEDGE_BASE_SCORE_NAME,
     PROMPT_SCORE_NAME,
@@ -50,6 +52,21 @@
 # Headroom for a full prompt rewrite + JSON wrapper; too low truncates into invalid JSON.
 _LLM_MAX_TOKENS = 16384
 
+# Byte budget for the projected trace payload. A byte-level BPE never emits more
+# than one token per UTF-8 byte, so this cannot breach the provider's 1M input
+# window; budgeting bytes rather than characters is also what charges a script
+# that encodes wide (Devanagari at 3 bytes/char) for what it actually costs.
+_TRACE_PAYLOAD_MAX_BYTES = 400_000
+_TRACE_FIELD_MAX_CHARS = 2_000
+_TRUNCATION_MARKER = " …[truncated]"
+
+# Repeats of one question are cut before whole rows: a question the model never
+# sees at all is a bigger loss than the 5th repeat of one it does.
+_MIN_REPEATS_PER_QUESTION = 3
+
+# Sorts unscoreable rows (value = "N/A") last, behind every real score.
+_UNSCOREABLE_SORT_VALUE = float("inf")
+
 # JSON keys expected in the LLM's structured response.
 _LLM_KEY_INSTRUCTIONS = "improved_instructions"
 _LLM_KEY_RATIONALE = "rationale"
@@ -630,6 +647,99 @@ def _call_prompt_drafting_llm(*, user_message_text: str) -> tuple[str, str]:
         ) from exc
 
 
+def _truncate(text: object) -> str:
+    """Clamp one trace field; long judge rationales are the main source of bloat."""
+    value = "" if text is None else str(text)
+    if len(value) <= _TRACE_FIELD_MAX_CHARS:
+        return value
+    return value[:_TRACE_FIELD_MAX_CHARS] + _TRUNCATION_MARKER
+
+
+def _project_trace(trace: dict, *, is_judge_run: bool) -> dict:
+    """Keep only the fields the brief tells the model to read."""
+    projected: dict = {
+        "question": _truncate(trace.get("question")),
+        "ground_truth_answer": _truncate(trace.get("ground_truth_answer")),
+        "llm_answer": _truncate(trace.get("llm_answer")),
+        "scores": [
+            {
+                "name": score.get("name"),
+                "value": score.get("value"),
+                **({"unscoreable": True} if score.get("unscoreable") else {}),
+                **(
+                    {"comment": _truncate(score.get("comment"))} if is_judge_run else {}
+                ),
+            }
+            for score in trace.get("scores") or []
+        ],
+    }
+    if not is_judge_run and trace.get("category"):
+        projected["category"] = trace["category"]
+    return projected
+
+
+def _limit_repeats(rows: list[dict]) -> list[dict]:
+    """Keep at most `_MIN_REPEATS_PER_QUESTION` traces per question, in file order.
+
+    A trace with no `question_id` is keyed uniquely, so datasets uploaded before
+    ids existed are not collapsed into one group and cut down to a handful of rows.
+    """
+    kept: list[dict] = []
+    seen: Counter[object] = Counter()
+    for index, row in enumerate(rows):
+        question_id = row.get("question_id")
+        key = question_id if question_id or question_id == 0 else f"__row_{index}"
+        if seen[key] >= _MIN_REPEATS_PER_QUESTION:
+            continue
+        seen[key] += 1
+        kept.append(row)
+    return kept
+
+
+def _primary_score(trace: dict, *, score_name: str) -> float:
+    """The metric the brief calls the primary signal; +inf when it is unusable."""
+    for score in trace.get("scores") or []:
+        if score.get("name") != score_name or score.get("unscoreable"):
+            continue
+        value = score.get("value")
+        if isinstance(value, bool) or not isinstance(value, (int, float)):
+            continue
+        return float(value)
+    return _UNSCOREABLE_SORT_VALUE
+
+
+def _pack_traces(rows: list[dict], *, is_judge_run: bool) -> list[dict]:
+    """Project rows into the byte budget, worst-scoring first."""
+    score_name = GROUND_TRUTH_SCORE_NAME if is_judge_run else COSINE_SCORE_NAME
+    kept: list[dict] = []
+    used = 0
+    for row in sorted(rows, key=lambda r: _primary_score(r, score_name=score_name)):
+        projected = _project_trace(row, is_judge_run=is_judge_run)
+        size = len(json.dumps(projected, ensure_ascii=False).encode())
+        if kept and used + size > _TRACE_PAYLOAD_MAX_BYTES:
+            break
+        kept.append(projected)
+        used += size
+    return kept
+
+
+def _select_traces(rows: list[dict], *, is_judge_run: bool) -> tuple[list[dict], int]:
+    """Fit the run's traces into the byte budget; returns (kept, total).
+
+    Worst-first both selects and orders — the brief tells the model to focus on the
+    failing rows, so those are the ones a cap has to keep. Repeats are only capped
+    when the whole run does not fit.
+    """
+    kept = _pack_traces(rows, is_judge_run=is_judge_run)
+    if len(kept) < len(rows):
+        kept = _pack_traces(_limit_repeats(rows), is_judge_run=is_judge_run)
+        logger.warning(
+            f"[_select_traces] Trimmed traces to fit the prompt budget | "
+            f"kept={len(kept)} total={len(rows)} is_judge_run={is_judge_run}"
+        )
+    return kept, len(rows)
+
+
 def _draft_improved_prompt(
     *,
     current_instructions: str,
@@ -714,11 +824,23 @@ def _draft_improved_prompt(
         f"sentence (≤ {_RATIONALE_MAX_LENGTH} characters): what you changed and why."
     )
     target_config = _target_config_from_params(config_params)
+    raw_rows = traces if isinstance(traces, list) else traces.get("traces") or []
+    rows = [row for row in raw_rows if isinstance(row, dict)]
+    selected_traces, total_traces = _select_traces(rows, is_judge_run=is_judge_run)
+    trace_scope = (
+        ""
+        if len(selected_traces) == total_traces
+        else (
+            f" You are seeing {len(selected_traces)} of {total_traces} traces — the "
+            "lowest-scoring rows; the rest are omitted."
+        )
+    )
 
     user_message_text = (
         "You are a prompt engineer. Below is a JSON array of evaluation traces"
-        f"{trace_description}\n\n"
-        f"## Evaluation traces\n```\n{json.dumps(traces)}\n```\n\n"
+        f"{trace_description}{trace_scope}\n\n"
+        "## Evaluation traces\n```\n"
+        f"{json.dumps(selected_traces, ensure_ascii=False)}\n```\n\n"
         f"## Current system prompt\n```\n{current_instructions}\n```\n\n"
         "## Target configuration (read-only — do NOT change any of these)\n"
         f"```\n{json.dumps(target_config)}\n```\n\n"

backend/app/tests/api/routes/test_improve_prompt.py

@@ -42,16 +42,21 @@
 
 from app.core.config import settings
 from app.crud.config.version import ConfigVersionCrud
+from app.crud.evaluations.score import GROUND_TRUTH_SCORE_NAME
 from app.crud.jobs import JobCrud
 from app.models import EvaluationDataset, EvaluationRun
 from app.models.config.config import ConfigTag
 from app.models.job import Job, JobStatus, JobType
 from app.services.evaluations.prompt_improvement import (
+    _MIN_REPEATS_PER_QUESTION,
+    _UNSCOREABLE_SORT_VALUE,
     _UPSTREAM_ERROR_DETAIL_MAX_LENGTH,
     AI_GENERATED_MARKER,
     COMMIT_MESSAGE_MAX_LENGTH,
     _anthropic_error_detail,
     _call_prompt_drafting_llm,
+    _draft_improved_prompt,
+    _primary_score,
     execute_prompt_improvement,
     start_prompt_improvement_job,
     validate_improve_prompt,
@@ -1039,3 +1044,133 @@ def test_missing_platform_key_is_flagged_as_misconfiguration(
 
         assert "ANTHROPIC_API_KEY" in str(raised.value)
         assert "anthropic_response" not in str(raised.value)
+
+
+class TestTraceBudget:
+    """Oversized runs must be trimmed to fit the model's input window, keeping the
+    rows the brief actually asks the model to fix."""
+
+    @staticmethod
+    def _judge_trace(question_id: int, ground_truth_value: float | str) -> dict:
+        unscoreable = ground_truth_value == "N/A"
+        return {
+            "trace_id": f"t{question_id}",
+            "question_id": question_id,
+            "question": f"Q{question_id} " + "क" * 3000,
+            "ground_truth_answer": "A" * 5000,
+            "llm_answer": "B" * 5000,
+            "scores": [
+                {
+                    "name": GROUND_TRUTH_SCORE_NAME,
+                    "value": ground_truth_value,
+                    "data_type": "NUMERIC",
+                    "comment": "C" * 8000,
+                    **({"unscoreable": True} if unscoreable else {}),
+                }
+            ],
+        }
+
+    @staticmethod
+    def _repeated_run(questions: int, repeats: int) -> list[dict]:
+        """A judged run of `questions` questions scored `repeats` times each."""
+        return [
+            TestTraceBudget._judge_trace(question_id, 2.0)
+            for question_id in range(questions)
+            for _ in range(repeats)
+        ]
+
+    @staticmethod
+    def _capture_brief(traces: list) -> str:
+        """Return the brief the drafting call would have received."""
+        captured: dict[str, str] = {}
+
+        def fake_call(*, user_message_text: str) -> tuple[str, str]:
+            captured["text"] = user_message_text
+            return "improved", "why"
+
+        with patch(f"{_SERVICE}._call_prompt_drafting_llm", side_effect=fake_call):
+            assert _draft_improved_prompt(
+                current_instructions="be helpful",
+                config_params={"model": "gpt-4o"},
+                traces=traces,
+                is_judge_run=True,
+            ) == ("improved", "why")
+        return captured["text"]
+
+    def test_oversized_traces_are_capped_worst_first(self) -> None:
+        # The uniquely BEST row is first in input order and the uniquely WORST is
+        # last, so an unsorted (or reversed) selection fails both assertions.
+        # Untrimmed this payload is ~1.3M chars.
+        traces = [
+            self._judge_trace(1000, 5.0),
+            *(self._judge_trace(i, 3.0) for i in range(1, 60)),
+            self._judge_trace(999, "N/A"),
+            self._judge_trace(0, 0.0),
+        ]
+
+        text = self._capture_brief(traces)
+
+        assert len(text) < 450_000
+        assert "Q0 " in text
+        assert "Q1000 " not in text
+        # Fields the brief never names are not shipped.
+        assert "trace_id" not in text
+        assert f"of {len(traces)} traces" in text
+
+    def test_small_run_is_sent_whole_without_a_scope_note(self) -> None:
+        text = self._capture_brief([self._judge_trace(1, 5.0)])
+
+        assert " traces —" not in text
+
+    def test_repeats_are_capped_before_any_question_is_dropped(self) -> None:
+        # Over budget at 5 repeats per question, under it at the floor of 3.
+        questions = 8
+        traces = self._repeated_run(questions=questions, repeats=5)
+
+        text = self._capture_brief(traces)
+
+        # Every question is still represented, three times each — no question was
+        # dropped to make room.
+        for question_id in range(questions):
+            assert text.count(f"Q{question_id} ") == _MIN_REPEATS_PER_QUESTION
+
+    def test_rows_are_dropped_when_the_repeat_floor_is_not_enough(self) -> None:
+        questions = 15
+        traces = self._repeated_run(questions=questions, repeats=5)
+
+        text = self._capture_brief(traces)
+
+        kept = sum(text.count(f"Q{question_id} ") for question_id in range(questions))
+        assert 0 < kept < questions * _MIN_REPEATS_PER_QUESTION
+        assert f"of {len(traces)} traces" in text
+
+    def test_traces_without_question_ids_still_get_trimmed(self) -> None:
+        """v1 traces from older datasets can carry no `question_id`.
+
+        Those must never be grouped as repeats of one another, so capping repeats is
+        a no-op for them and the byte budget is the only thing that trims the brief.
+        """
+        questions = 40
+        traces = [self._judge_trace(i, 2.0) for i in range(questions)]
+        for trace in traces:
+            del trace["question_id"]
+
+        text = self._capture_brief(traces)
+
+        kept = sum(text.count(f"Q{i} ") for i in range(questions))
+        assert 0 < kept < questions
+
+    def test_unscoreable_score_sorts_behind_every_real_score(self) -> None:
+        """`value` is the string "N/A" on unscoreable rows; mixed str/float must not
+        raise, and those rows must rank last."""
+        unscoreable = self._judge_trace(1, "N/A")
+        best = self._judge_trace(2, 5.0)
+
+        assert _primary_score(
+            unscoreable, score_name=GROUND_TRUTH_SCORE_NAME
+        ) > _primary_score(best, score_name=GROUND_TRUTH_SCORE_NAME)
+        # A missing metric is treated the same way.
+        assert (
+            _primary_score({"scores": []}, score_name=GROUND_TRUTH_SCORE_NAME)
+            == _UNSCOREABLE_SORT_VALUE
+        )

docs/wiki/modules/evaluations.md

@@ -40,6 +40,7 @@ v2 judge field on `EvaluationRun`: `is_judge_run` (bool marker gating native jud
 - Fast text runs (v1 cosine + v2 judge) fan out `ceil(total_items / EVAL_FAST_CHUNK_SIZE)` `run_evaluation_fast_chunk` tasks (responses only), then a cron barrier (`dispatch_fast_evaluation_barriers`) enqueues one `run_evaluation_fast_aggregate` once every chunk has a `raw_output_url`. The aggregate merges chunks, then (v2) judges **every** row in that single task — so the judge pool is sized by its own `EVAL_JUDGE_CONCURRENCY` (not the response stage's `EVAL_FAST_API_CONCURRENCY`) to clear the max dataset (`EVAL_FAST_MAX_UNIQUE_ROWS` × `duplication_factor`) under the aggregate's `CELERY_TASK_SOFT_TIME_LIMIT`. No judge fan-out / second barrier.
 - Prompt improvement is job-based with callback delivery: `POST /evaluations/{id}/improve-prompt` validates preconditions, enqueues a `Job` (`JobType.PROMPT_IMPROVEMENT`, `models/job.py`) run by Celery task `run_prompt_improvement` (`celery/tasks/job_execution.py`), and returns `202` with an `LLMJobImmediatePublic` handle. On finish the worker POSTs a single best-effort callback to the caller-supplied `callback_url` (SSRF-guarded via `validate_callback_url`): an `APIResponse[PromptImprovementJobPublic]` (`models/evaluation.py`) carrying the new `ConfigVersion` on success or `error_message` on failure (for an Anthropic fault `error_message` carries the provider's own response body, appended by `_anthropic_error_detail` in `services/evaluations/prompt_improvement.py`). The `ConfigVersion` is persisted regardless of callback outcome. Celery redelivery of a `SUCCESS` job re-sends the callback without re-running the LLM.
 - v2 prompt iteration reuses the same job/Celery/config-version machinery (`services/evaluations/prompt_improvement.py`) and branches on `run.is_judge_run`: the v2 route calls `start_prompt_improvement_job(..., require_judge_run=True)` (a non-judge run → `422 not_a_judge_run`), and the worker drafts from the three-metric judge trace (`_draft_improved_prompt(is_judge_run=True)`, using each metric's score + reasoning) and delivers a `PromptRecommendationJobPublic` callback carrying `recommendation_type` (`Literal["prompt"]`, `models/evaluation.py`; widens to a union when knowledge-base / model recommendations land). Non-judge (v1) runs keep the default `_draft_improved_prompt` brief + `PromptImprovementJobPublic` path unchanged.
+- Both prompt-improvement briefs are size-bounded before the Anthropic call (`services/evaluations/prompt_improvement.py`) — a large run used to exceed the provider's 1M-token input limit and fail the job with `prompt_generation_failed: ... HTTP 400`. Every trace is projected down to the fields the brief names and its long text fields (notably the judge `comment`) are clamped; `_select_traces` then fills `_TRACE_PAYLOAD_MAX_BYTES` worst-scoring-first (the brief tells the model to focus on the failing rows, so those are what a cap keeps), and only when the whole run does not fit does `_limit_repeats` cap repeats of a question at `_MIN_REPEATS_PER_QUESTION` (3) — a question the model never sees is a bigger loss than its 5th repeat, and a trace with no `question_id` is keyed uniquely so pre-id datasets are never collapsed into one group. The budget is bytes rather than characters because a byte-level BPE never emits more than one token per UTF-8 byte: the bound needs no `count_tokens` round trip, and a wide-encoding script (Devanagari at 3 bytes/char) is charged for what it actually costs. The brief states how many traces it is not seeing, and trimming is logged. `crud/evaluations/summary.py` still sends every trace whole (its own `ponytail:` note) and has the same latent ceiling.
 - Evaluation iteration loop (`POST /evaluations/iterations`) chains fast-eval + v2 prompt improvement into a self-driving cycle via LangGraph (`services/evaluations/iteration_graph.py`): `start_eval_node` → `wait_eval_node` → (conditional) `start_improve_node` → `wait_improve_node` → loops back to `start_eval_node`, or → `finalize_node`. Stop-score = mean(`Adherence to Ground Truth`, `Adherence to Prompt`) from `EvaluationRun.score["summary_scores"]` (`compute_round_scores`); `Adherence to Knowledge Base` is recorded per round for visibility only, never gates stopping. Stops on 3 consecutive rounds with `<EVAL_ITERATION_CEILING_DELTA_THRESHOLD` (0.05) improvement (`ceiling_reached`), a `max_rounds` cap (`max_rounds_reached`), or a hard round failure (`round_failed`). No orchestrator polls a provider directly: each `wait_*_node` reads `EvaluationRun.status`/`Job.status` (both already maintained by the existing fast-eval and prompt-improvement machinery) and calls LangGraph's `interrupt()` if not yet terminal — a Postgres checkpointer (`get_evaluation_iteration_checkpointer`) persists state across the pause. Resumption is driven by the existing `/cron/evaluations` tick, not a new scheduler: `dispatch_pending_evaluation_iteration_resumes` (`crud/evaluations/cron.py`) re-dispatches `run_evaluation_iteration_graph_step` (Celery, `celery/tasks/job_execution.py`, priority 6) for every `EvaluationIterationRun` with `status=processing`. Every node opens/closes its own DB session — no session is held open across an `interrupt()`. `finalize_node` persists the thin row's terminal `status`/`stop_reason` and delivers the round history + best round as an `EvaluationIterationReportPublic` callback, same best-effort delivery convention as prompt improvement above.
 
 - Completion webhook (`callback_url`): `update_evaluation_run` (`crud/evaluations/core.py`) fires on the same terminal-transition guard as the completion email — when the run enters `completed`/`failed` and `callback_url` is set, `_enqueue_eval_completion_callback` runs Celery task `send_eval_completion_callback` → `services/notifications/eval_completion.py::execute_eval_completion_callback`, which POSTs an `APIResponse` (success on completed, `error` set to the run's `error_message` on failed) via `send_callback` (HMAC-signed with `get_webhook_secret`); `data` is a slim run snapshot (id/run_name/dataset_name/status/run_mode/timestamps only — no score/cost). Best-effort, at-least-once; only v2 runs set `callback_url`.