diff --git a/tests/test_app.py b/tests/test_app.py
index a4bc7c6..a4d88b5 100644
--- a/tests/test_app.py
+++ b/tests/test_app.py
@@ -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)
diff --git a/tests/test_llm_validate.py b/tests/test_llm_validate.py
index 3b5d766..4a424fe 100644
--- a/tests/test_llm_validate.py
+++ b/tests/test_llm_validate.py
@@ -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(
diff --git a/unredact/app.py b/unredact/app.py
index 439ee27..6a51031 100644
--- a/unredact/app.py
+++ b/unredact/app.py
@@ -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}")
diff --git a/unredact/pipeline/llm_validate.py b/unredact/pipeline/llm_validate.py
index c6a6aaf..a5d8857 100644
--- a/unredact/pipeline/llm_validate.py
+++ b/unredact/pipeline/llm_validate.py
@@ -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
diff --git a/unredact/static/solver.js b/unredact/static/solver.js
index de75cf8..f33e8ba 100644
--- a/unredact/static/solver.js
+++ b/unredact/static/solver.js
@@ -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 = `
+ ${escapeHtml(item.text)}
+ ${item.error_px.toFixed(1)}px ${item.source || ""}
+ `;
+ 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. */