fix: batch LLM validation, stream results, improve scoring prompt

- Batch candidates in groups of 50 (was 200) to avoid output token truncation that caused all-zero scores - Stream scored results via SSE as each batch completes instead of waiting for all batches to finish - Include solve_id in done event so validate works when results fit in a single page - Rewrite prompt to judge sentence coherence ("does it make sense") rather than word predictability ("is it expected") - Frontend inserts results in score-sorted order as they arrive Co-Authored-By: Claude Opus 4.6 <[email protected]>
2026-10-27 08:40:29 +00:00 · 2026-02-23 21:17:00 -06:00
1 parent fa6709cca4
commit 672bdc3223
@@ -446,7 +446,7 @@ async def test_cancel_clears_buffer():

@pytest.mark.anyio
 async def test_validate_endpoint(monkeypatch):
-    """POST /api/solve/{id}/validate should return LLM-scored results."""
+    """POST /api/solve/{id}/validate should stream SSE with LLM-scored results."""
    from unredact.pipeline import llm_validate

    fake_id = "validate_test"
@@ -456,10 +456,15 @@ async def test_validate_endpoint(monkeypatch):
        {"text": "running", "width_px": 50.0, "error_px": 0.3, "source": "words"},
    ]

-    async def mock_validate(left, right, candidates):
-        return [95, 10, 5]
+    async def mock_score_batch(client, left, right, candidates):
+        scores = {"Smith": 95, "house": 10, "running": 5}
+        return [scores.get(c, 0) for c in candidates]

-    monkeypatch.setattr(llm_validate, "validate_candidates", mock_validate)
+    monkeypatch.setattr(llm_validate, "_score_batch", mock_score_batch)
+
+    # Mock _get_client since the endpoint imports it directly
+    from unredact.pipeline import llm_detect
+    monkeypatch.setattr(llm_detect, "_get_client", lambda: None)

    transport = ASGITransport(app=app)
    async with AsyncClient(transport=transport, base_url="http://test") as client:
@@ -468,14 +473,26 @@ async def test_validate_endpoint(monkeypatch):
            "right_context": ", we are writing",
        })
        assert resp.status_code == 200
-        data = resp.json()
+
+        # Parse SSE events from response
+        events = []
+        for line in resp.text.split("\n"):
+            if line.startswith("data: "):
+                events.append(json.loads(line[6:]))
+
+        batch_events = [e for e in events if e.get("status") == "batch_done"]
+        assert len(batch_events) == 1  # 3 items fits in one batch
+        data = batch_events[0]
        assert len(data["results"]) == 3
-        assert data["total"] == 3
-        # Sorted by llm_score descending
-        assert data["results"][0]["text"] == "Smith"
-        assert data["results"][0]["llm_score"] == 95
-        assert data["results"][1]["llm_score"] == 10
-        assert data["results"][2]["llm_score"] == 5
+        # Results within a batch are unsorted (frontend sorts on insert)
+        scores = {r["text"]: r["llm_score"] for r in data["results"]}
+        assert scores["Smith"] == 95
+        assert scores["house"] == 10
+        assert scores["running"] == 5
+
+        done_event = [e for e in events if e.get("status") == "done"]
+        assert len(done_event) == 1
+        assert done_event[0]["total"] == 3

    _solve_results.pop(fake_id, None)

@@ -18,7 +18,7 @@ class TestBuildValidationPrompt:
        assert "1. Smith" in prompt
        assert "2. house" in prompt
        assert "3. running" in prompt
-        assert "[REDACTED]" in prompt
+        assert "_____" in prompt

    def test_empty_context(self):
        prompt = build_validation_prompt(
@@ -562,6 +562,7 @@ async def solve(req: SolveRequest):
            yield json.dumps({
                "status": "done",
                "total_found": len(found_texts),
+                "solve_id": solve_id,
            })
        finally:
            _active_solves.pop(solve_id, None)
@@ -595,24 +596,58 @@ async def validate_solve(solve_id: str, req: ValidateRequest):
    if solve_id not in _solve_results:
        return JSONResponse({"error": "solve not found"}, status_code=404)

-    from unredact.pipeline.llm_validate import validate_candidates
+    from unredact.pipeline.llm_validate import _score_batch, _BATCH_SIZE

-    buf = _solve_results[solve_id]
+    buf = list(_solve_results[solve_id])  # snapshot
    candidates = [r["text"] for r in buf]
+    total = len(candidates)
+    total_batches = (total + _BATCH_SIZE - 1) // _BATCH_SIZE

-    try:
-        scores = await validate_candidates(
-            req.left_context, req.right_context, candidates,
-        )
-    except Exception as e:
-        return JSONResponse({"error": str(e)}, status_code=500)
+    async def event_generator():
+        yield json.dumps({"status": "started", "total": total, "batches": total_batches})

-    scored = []
-    for i, result in enumerate(buf):
-        scored.append({**result, "llm_score": scores[i] if i < len(scores) else 0})
-    scored.sort(key=lambda x: x["llm_score"], reverse=True)
+        scored_so_far = 0
+        try:
+            from unredact.pipeline.llm_detect import _get_client
+            client = _get_client()

-    return {"results": scored, "total": len(scored)}
+            for batch_idx in range(total_batches):
+                start = batch_idx * _BATCH_SIZE
+                end = min(start + _BATCH_SIZE, total)
+                batch = candidates[start:end]
+
+                yield json.dumps({
+                    "status": "scoring",
+                    "batch": batch_idx + 1,
+                    "batches": total_batches,
+                    "scored": scored_so_far,
+                    "total": total,
+                })
+
+                batch_scores = await _score_batch(
+                    client, req.left_context, req.right_context, batch,
+                )
+
+                # Stream this batch's scored results immediately
+                batch_results = []
+                for i, s in enumerate(batch_scores):
+                    result = {**buf[start + i], "llm_score": s}
+                    batch_results.append(result)
+
+                scored_so_far += len(batch_results)
+                yield json.dumps({
+                    "status": "batch_done",
+                    "results": batch_results,
+                    "scored": scored_so_far,
+                    "total": total,
+                })
+        except Exception as e:
+            yield json.dumps({"status": "error", "error": str(e)})
+            return
+
+        yield json.dumps({"status": "done", "scored": scored_so_far, "total": total})
+
+    return EventSourceResponse(event_generator())


@app.delete("/api/solve/{solve_id}")
@@ -7,10 +7,12 @@ Returns a contextual fit score (0-100) for each candidate.
 from __future__ import annotations

 import logging
+from typing import Callable

 log = logging.getLogger(__name__)

 _VALIDATION_MODEL = "claude-sonnet-4-6"
+_BATCH_SIZE = 50

 SCORE_TOOL = {
    "name": "score_candidates",
@@ -51,24 +53,18 @@ def build_validation_prompt(
    """Build the prompt for LLM validation of solve candidates."""
    numbered = "\n".join(f"{i+1}. {c}" for i, c in enumerate(candidates))
    return (
-        "You are analyzing a redacted document. A section of text has been "
-        "blacked out. The text surrounding the redaction reads:\n\n"
-        f'Left context: "{left_context}"\n'
-        "[REDACTED]\n"
-        f'Right context: "{right_context}"\n\n'
-        "Below is a list of candidate words/phrases that fit the redacted "
-        "space by pixel width. Score each from 0-100 on how well it fits "
-        "contextually:\n\n"
-        "- 90-100: Near-certain fit (grammatically correct, semantically "
-        "meaningful, contextually expected)\n"
-        "- 60-89: Plausible (makes sense but not the most likely)\n"
-        "- 30-59: Unlikely (grammatically possible but doesn't make much sense)\n"
-        "- 0-29: Very poor fit (nonsensical, wrong part of speech, doesn't "
-        "work in context)\n\n"
-        'Example: If left context is "Dear Mr." and right is ", we are '
-        'writing to inform you":\n'
-        '- "Smith" -> 95 (common surname, perfect fit)\n'
-        '- "house" -> 5 (not a surname, makes no sense after "Mr.")\n\n'
+        "You are analyzing a redacted document. A word or phrase has been "
+        "blacked out. The original sentence reads:\n\n"
+        f'"{left_context} _____ {right_context}"\n\n'
+        "For each candidate below, mentally substitute it into the blank and "
+        "evaluate: Does the resulting sentence read naturally? Is it "
+        "grammatically correct? Does it make sense?\n\n"
+        "Score each candidate 0-100 based on whether it produces a "
+        "coherent sentence:\n"
+        "- 80-100: The sentence reads naturally and makes sense\n"
+        "- 50-79: Grammatically acceptable but slightly awkward or unusual\n"
+        "- 20-49: Grammatically questionable or semantically odd\n"
+        "- 0-19: Ungrammatical, nonsensical, or wrong part of speech\n\n"
        f"Candidates:\n{numbered}"
    )

@@ -76,20 +72,14 @@ def build_validation_prompt(
 from unredact.pipeline.llm_detect import _get_client


-async def validate_candidates(
+async def _score_batch(
+    client,
    left_context: str,
    right_context: str,
    candidates: list[str],
 ) -> list[int]:
-    """Score candidates using LLM. Returns list of scores in same order as candidates.
-
-    Raises on API error — caller should handle.
-    """
-    if not candidates:
-        return []
-
+    """Score a single batch of candidates. Returns scores in same order."""
    prompt = build_validation_prompt(left_context, right_context, candidates)
-    client = _get_client()

    response = await client.messages.create(
        model=_VALIDATION_MODEL,
@@ -110,3 +100,40 @@ async def validate_candidates(
            break

    return scores
+
+
+async def validate_candidates(
+    left_context: str,
+    right_context: str,
+    candidates: list[str],
+    on_progress: Callable[[int, int], None] | None = None,
+) -> list[int]:
+    """Score candidates using LLM. Returns list of scores in same order as candidates.
+
+    Batches candidates into groups to stay within output token limits.
+    Calls on_progress(scored_so_far, total) after each batch if provided.
+
+    Raises on API error — caller should handle.
+    """
+    if not candidates:
+        return []
+
+    client = _get_client()
+    scores = [0] * len(candidates)
+    total_batches = (len(candidates) + _BATCH_SIZE - 1) // _BATCH_SIZE
+
+    for batch_idx in range(total_batches):
+        start = batch_idx * _BATCH_SIZE
+        end = min(start + _BATCH_SIZE, len(candidates))
+        batch = candidates[start:end]
+
+        log.info("Scoring batch %d/%d (%d candidates)", batch_idx + 1, total_batches, len(batch))
+        batch_scores = await _score_batch(client, left_context, right_context, batch)
+
+        for i, s in enumerate(batch_scores):
+            scores[start + i] = s
+
+        if on_progress:
+            on_progress(end, len(candidates))
+
+    return scores
@@ -206,6 +206,7 @@ function handleSolveEvent(data, redactionId) {
    solveStatus.textContent = `Found ${displayedCount} matches, searching for more...`;
  } else if (data.status === "done") {
    totalFound = data.total_found;
+    if (data.solve_id) currentSolveId = data.solve_id;
    if (totalFound > displayedCount) {
      solveLoadMore.textContent = `Load more (showing ${displayedCount} of ${totalFound})`;
      solveLoadMore.hidden = false;
@@ -305,7 +306,14 @@ async function runValidation() {
  if (!currentSolveId) return;
  validateRun.disabled = true;
  validateRun.textContent = "Validating...";
-  solveStatus.textContent = `Validating ${totalFound} results...`;
+  solveStatus.textContent = "Starting validation...";
+
+  // Clear existing results for fresh scored list
+  solveResults.innerHTML = "";
+  displayedCount = 0;
+  solveLoadMore.hidden = true;
+
+  const redactionId = state.activeRedaction;

  try {
    const resp = await fetch(`/api/solve/${currentSolveId}/validate`, {
@@ -316,21 +324,44 @@ async function runValidation() {
        right_context: validateRight.value,
      }),
    });
-    if (!resp.ok) throw new Error("Validation failed");
-    const data = await resp.json();
-
-    // Replace results with scored list
-    solveResults.innerHTML = "";
-    displayedCount = 0;
-    solveLoadMore.hidden = true;
-
-    const redactionId = state.activeRedaction;
-    for (const item of data.results) {
-      renderScoredResult(item, redactionId);
+    if (!resp.ok) {
+      const errBody = await resp.json().catch(() => ({}));
+      throw new Error(errBody.error || `Validation failed (${resp.status})`);
    }

-    solveStatus.textContent = `Validated. ${data.total} results scored.`;
-    validatePanel.hidden = true;
+    const reader = resp.body.getReader();
+    const decoder = new TextDecoder();
+    let buffer = "";
+
+    while (true) {
+      const { done, value } = await reader.read();
+      if (done) break;
+      buffer += decoder.decode(value, { stream: true });
+      const lines = buffer.split("\n");
+      buffer = lines.pop();
+
+      for (const line of lines) {
+        if (!line.startsWith("data: ")) continue;
+        let data;
+        try { data = JSON.parse(line.slice(6)); } catch { continue; }
+
+        if (data.status === "started") {
+          solveStatus.textContent = `Validating ${data.total} results (${data.batches} batches)...`;
+        } else if (data.status === "scoring") {
+          solveStatus.textContent = `Scoring batch ${data.batch}/${data.batches} (${data.scored}/${data.total} scored)...`;
+        } else if (data.status === "batch_done") {
+          for (const item of data.results) {
+            renderScoredResult(item, redactionId);
+          }
+          solveStatus.textContent = `Scored ${data.scored}/${data.total}...`;
+        } else if (data.status === "error") {
+          throw new Error(data.error);
+        } else if (data.status === "done") {
+          solveStatus.textContent = `Validated. ${data.total} results scored.`;
+          validatePanel.hidden = true;
+        }
+      }
+    }
  } catch (err) {
    solveStatus.textContent = "Validation error: " + err.message;
  } finally {
@@ -339,27 +370,69 @@ async function runValidation() {
  }
 }

-function renderScoredResult(data, redactionId) {
-  handleSolveEvent({ status: "match", ...data }, redactionId);
+/**
+ * Render a single scored result, inserting in descending score order.
+ */
+function renderScoredResult(item, redactionId) {
+  const score = item.llm_score ?? 0;

-  if (data.llm_score !== undefined) {
-    const results = solveResults.children;
-    for (const el of results) {
-      const textEl = el.querySelector(".result-text");
-      if (textEl && textEl.textContent === data.text && !el.querySelector(".llm-score")) {
-        const badge = document.createElement("span");
-        badge.className = "llm-score";
-        const score = data.llm_score;
-        if (score >= 70) badge.classList.add("score-high");
-        else if (score >= 30) badge.classList.add("score-mid");
-        else badge.classList.add("score-low");
-        badge.textContent = String(score);
-        badge.title = "LLM contextual fit score";
-        el.prepend(badge);
-        break;
-      }
+  const div = document.createElement("div");
+  div.className = "solve-result";
+  div.dataset.llmScore = String(score);
+
+  // Score badge
+  const badge = document.createElement("span");
+  badge.className = "llm-score";
+  if (score >= 70) badge.classList.add("score-high");
+  else if (score >= 30) badge.classList.add("score-mid");
+  else badge.classList.add("score-low");
+  badge.textContent = String(score);
+  badge.title = "LLM contextual fit score";
+
+  div.innerHTML = `
+    <span class="result-text">${escapeHtml(item.text)}</span>
+    <span class="result-error">${item.error_px.toFixed(1)}px ${item.source || ""}</span>
+  `;
+  div.prepend(badge);
+
+  // Click to preview
+  div.addEventListener("click", () => {
+    const r = state.redactions[redactionId];
+    if (!r) return;
+
+    let previewText = item.text;
+    const ks = solveKnownStart.value;
+    const ke = solveKnownEnd.value;
+    if (ks && previewText.toLowerCase().startsWith(ks.toLowerCase())) {
+      previewText = previewText.slice(ks.length);
+    }
+    if (ke && previewText.toLowerCase().endsWith(ke.toLowerCase())) {
+      previewText = previewText.slice(0, -ke.length);
+    }
+
+    r.preview = previewText;
+    r.solveFullText = item.text;
+    redactionMarker.value = previewText;
+    redactionMarker.className = "redaction-marker preview";
+    renderCanvas();
+    solveResults.querySelectorAll(".solve-result").forEach(el => el.classList.remove("active"));
+    div.classList.add("active");
+    solveAccept.hidden = false;
+  });
+
+  // Insert in sorted position (descending by score)
+  let inserted = false;
+  for (const existing of solveResults.children) {
+    const exScore = parseInt(/** @type {HTMLElement} */ (existing).dataset.llmScore || "0");
+    if (score > exScore) {
+      solveResults.insertBefore(div, existing);
+      inserted = true;
+      break;
    }
  }
+  if (!inserted) solveResults.appendChild(div);
+
+  displayedCount++;
 }

 /** Set up solver button listeners. Call once from main.js. */