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. */