From f1e25f97122a02b45eccdc9b84217c30fc4c7a79 Mon Sep 17 00:00:00 2001 From: Alex-Gilbert Date: Fri, 20 Feb 2026 14:46:45 -0600 Subject: [PATCH] feat: add LLM boundary detection for spot redaction path analyze_spot_redaction is now async and calls identify_boundary_text() to use Claude Haiku for clean left/right text extraction instead of center-point character filtering. Falls back to center-point if LLM is unavailable. Co-Authored-By: Claude Opus 4.6 --- tests/test_analyze_page.py | 60 +++++++++++---- unredact/app.py | 4 +- unredact/pipeline/analyze_page.py | 35 +++++---- unredact/pipeline/llm_detect.py | 121 ++++++++++++++++++++++++++++++ 4 files changed, 185 insertions(+), 35 deletions(-) diff --git a/tests/test_analyze_page.py b/tests/test_analyze_page.py index f17d518..31b615a 100644 --- a/tests/test_analyze_page.py +++ b/tests/test_analyze_page.py @@ -8,7 +8,7 @@ import pytest from PIL import Image from unredact.pipeline.ocr import OcrChar, OcrLine -from unredact.pipeline.llm_detect import LlmRedaction +from unredact.pipeline.llm_detect import LlmRedaction, BoundaryText from unredact.pipeline.detect_redactions import Redaction from unredact.pipeline.font_detect import FontMatch from unredact.pipeline.analyze_page import ( @@ -355,7 +355,8 @@ async def test_analyze_page_uses_precomputed_ocr_lines(): # test_analyze_spot_redaction # --------------------------------------------------------------------------- -def test_analyze_spot_redaction_returns_analysis(): +@pytest.mark.anyio +async def test_analyze_spot_redaction_returns_analysis(): """analyze_spot_redaction should return a RedactionAnalysis for a box that overlaps an OCR line, with correct font, text, and offsets.""" line = _make_line_from_words(["Hello", "|||", "world"], y=100, h=20) @@ -372,30 +373,38 @@ def test_analyze_spot_redaction_returns_analysis(): font_match = _dummy_font_match() page_image = _dummy_page_image() + boundary = BoundaryText(left_text="Hello", right_text="world") + with ( patch( "unredact.pipeline.analyze_page.detect_font_masked", return_value=font_match, ), + patch( + "unredact.pipeline.analyze_page.identify_boundary_text", + new_callable=AsyncMock, + return_value=boundary, + ), patch( "unredact.pipeline.analyze_page.align_text_to_page", return_value=(3, -1), ), ): - result = analyze_spot_redaction(page_image, lines, box) + result = await analyze_spot_redaction(page_image, lines, box) assert result is not None assert isinstance(result, RedactionAnalysis) assert result.box is box assert result.line is line assert result.font is font_match - assert "Hello" in result.left_text - assert "world" in result.right_text + assert result.left_text == "Hello" + assert result.right_text == "world" assert isinstance(result.offset_x, float) assert isinstance(result.offset_y, float) -def test_analyze_spot_redaction_no_matching_line(): +@pytest.mark.anyio +async def test_analyze_spot_redaction_no_matching_line(): """analyze_spot_redaction should return None if no OCR line overlaps.""" # Line at y=100, box at y=500 — no overlap line = _make_line_from_words(["Hello", "world"], y=100, h=20) @@ -404,20 +413,22 @@ def test_analyze_spot_redaction_no_matching_line(): box = Redaction(id="spot2", x=50, y=500, w=100, h=20) page_image = _dummy_page_image() - result = analyze_spot_redaction(page_image, lines, box) + result = await analyze_spot_redaction(page_image, lines, box) assert result is None -def test_analyze_spot_redaction_empty_lines(): +@pytest.mark.anyio +async def test_analyze_spot_redaction_empty_lines(): """analyze_spot_redaction should return None for empty OCR lines list.""" box = Redaction(id="spot3", x=50, y=100, w=100, h=20) page_image = _dummy_page_image() - result = analyze_spot_redaction(page_image, [], box) + result = await analyze_spot_redaction(page_image, [], box) assert result is None -def test_analyze_spot_redaction_no_left_text_skips_alignment(): +@pytest.mark.anyio +async def test_analyze_spot_redaction_no_left_text_skips_alignment(): """When the box covers the start of the line (no left text), alignment should be skipped and offsets should be 0.0.""" # Box starts at x=50 (same as line start), so no chars to the left @@ -429,11 +440,20 @@ def test_analyze_spot_redaction_no_left_text_skips_alignment(): font_match = _dummy_font_match() page_image = _dummy_page_image() - with patch( - "unredact.pipeline.analyze_page.detect_font_masked", - return_value=font_match, + boundary = BoundaryText(left_text="", right_text="world") + + with ( + patch( + "unredact.pipeline.analyze_page.detect_font_masked", + return_value=font_match, + ), + patch( + "unredact.pipeline.analyze_page.identify_boundary_text", + new_callable=AsyncMock, + return_value=boundary, + ), ): - result = analyze_spot_redaction(page_image, lines, box) + result = await analyze_spot_redaction(page_image, lines, box) assert result is not None assert result.left_text == "" @@ -441,7 +461,8 @@ def test_analyze_spot_redaction_no_left_text_skips_alignment(): assert result.offset_y == 0.0 -def test_analyze_spot_redaction_integration(sample_page_image): +@pytest.mark.anyio +async def test_analyze_spot_redaction_integration(sample_page_image): """analyze_spot_redaction should return full analysis for a known bbox.""" from unredact.pipeline.ocr import ocr_page @@ -468,7 +489,14 @@ def test_analyze_spot_redaction_integration(sample_page_image): h=line.h, ) - result = analyze_spot_redaction(sample_page_image, lines, box) + boundary = BoundaryText(left_text="test left", right_text="test right") + with patch( + "unredact.pipeline.analyze_page.identify_boundary_text", + new_callable=AsyncMock, + return_value=boundary, + ): + result = await analyze_spot_redaction(sample_page_image, lines, box) + assert result is not None assert result.box is box assert result.line is line diff --git a/unredact/app.py b/unredact/app.py index 95febbe..926ce78 100644 --- a/unredact/app.py +++ b/unredact/app.py @@ -520,9 +520,7 @@ async def spot(doc_id: str, page: int, data: dict): ocr_lines = pd.get("ocr_lines") analysis_json = None if ocr_lines: - ra = await asyncio.to_thread( - analyze_spot_redaction, page_img, ocr_lines, result, - ) + ra = await analyze_spot_redaction(page_img, ocr_lines, result) if ra is not None: font_id = _make_font_id(ra.font.font_name) segments = [] diff --git a/unredact/pipeline/analyze_page.py b/unredact/pipeline/analyze_page.py index 1edbc95..50c5e66 100644 --- a/unredact/pipeline/analyze_page.py +++ b/unredact/pipeline/analyze_page.py @@ -14,7 +14,7 @@ import numpy as np from PIL import Image from unredact.pipeline.ocr import ocr_page, OcrLine -from unredact.pipeline.llm_detect import detect_redactions_llm, LlmRedaction +from unredact.pipeline.llm_detect import detect_redactions_llm, LlmRedaction, identify_boundary_text from unredact.pipeline.detect_redactions import find_redaction_in_region, Redaction from unredact.pipeline.font_detect import detect_font_masked, align_text_to_page, FontMatch @@ -202,17 +202,17 @@ async def analyze_page( return PageAnalysis(lines=lines, redactions=results) -def analyze_spot_redaction( +async def analyze_spot_redaction( page_image: Image.Image, ocr_lines: list[OcrLine], box: Redaction, ) -> RedactionAnalysis | None: """Run analysis on a single known redaction bounding box. - Uses cached OCR data to: + Uses cached OCR data and an LLM call to: 1. Find the OCR line containing the redaction 2. Detect the font (with redaction masking) - 3. Extract left/right context text + 3. Use LLM to identify clean boundary text (handles garbled OCR) 4. Compute pixel alignment offsets Args: @@ -244,14 +244,15 @@ def analyze_spot_redaction( line = best_line redaction_boxes = [(box.x, box.y, box.w, box.h)] - # Font detection with masking - font = detect_font_masked(line, page_image, redaction_boxes) + # Font detection with masking (blocking — run in thread) + font = await asyncio.to_thread( + detect_font_masked, line, page_image, redaction_boxes, + ) - # Extract left/right text using center-point char filtering - left_chars = [c for c in line.chars if c.x + c.w / 2 < box.x] - right_chars = [c for c in line.chars if c.x + c.w / 2 > box.x + box.w] - left_text = "".join(c.text for c in left_chars).strip() - right_text = "".join(c.text for c in right_chars).strip() + # LLM boundary text identification (async) + boundary = await identify_boundary_text(line, box.x, box.w) + left_text = boundary.left_text + right_text = boundary.right_text # Pixel alignment offset_x = 0.0 @@ -263,11 +264,13 @@ def analyze_spot_redaction( text_region_x2 = min(page_image.width, box.x + 20) text_region_y1 = max(0, line.y - 10) text_region_y2 = min(page_image.height, line.y + line.h + 10) - text_crop = np.array(page_image.convert("L").crop( - (text_region_x1, text_region_y1, text_region_x2, text_region_y2) - )) - align_dx, align_dy = align_text_to_page( - left_text, pil_font, text_crop, + text_crop = await asyncio.to_thread( + lambda: np.array(page_image.convert("L").crop( + (text_region_x1, text_region_y1, text_region_x2, text_region_y2) + )) + ) + align_dx, align_dy = await asyncio.to_thread( + align_text_to_page, left_text, pil_font, text_crop, ) offset_x = float(text_region_x1 + align_dx - line.x) offset_y = float(text_region_y1 + align_dy - line.y) diff --git a/unredact/pipeline/llm_detect.py b/unredact/pipeline/llm_detect.py index d30a151..b463f28 100644 --- a/unredact/pipeline/llm_detect.py +++ b/unredact/pipeline/llm_detect.py @@ -230,6 +230,127 @@ def _parse_response( return redactions +@dataclass +class BoundaryText: + """Clean left/right text around a redaction, as identified by the LLM.""" + + left_text: str + right_text: str + + +_BOUNDARY_TOOL = { + "name": "report_boundary_text", + "description": "Report the clean text on each side of a redaction in an OCR line.", + "input_schema": { + "type": "object", + "properties": { + "left_text": { + "type": "string", + "description": ( + "Clean text before the redaction, with OCR " + "artifacts removed. Empty string if redaction " + "is at line start." + ), + }, + "right_text": { + "type": "string", + "description": ( + "Clean text after the redaction, with OCR " + "artifacts removed. Empty string if redaction " + "is at line end." + ), + }, + }, + "required": ["left_text", "right_text"], + }, +} + + +def _build_boundary_prompt(line: OcrLine, box_x: int, box_w: int) -> str: + """Build prompt for single-line boundary text identification.""" + left_chars = [c for c in line.chars if c.x + c.w / 2 < box_x] + mid_chars = [ + c for c in line.chars + if box_x <= c.x + c.w / 2 <= box_x + box_w + ] + right_chars = [c for c in line.chars if c.x + c.w / 2 > box_x + box_w] + + left_raw = "".join(c.text for c in left_chars) + mid_raw = "".join(c.text for c in mid_chars) + right_raw = "".join(c.text for c in right_chars) + + return ( + "You are analyzing a single line of OCR text from a scanned legal document. " + "Part of this line has been redacted (blacked out), which causes the OCR engine " + "to produce garbled characters near the redaction edges.\n\n" + f'Full OCR line: {line.text}\n' + f'Approximate split: LEFT="{left_raw}" REDACTED="{mid_raw}" RIGHT="{right_raw}"\n\n' + "The characters immediately adjacent to the redaction may be misread or truncated " + "by the OCR engine. For example, 'Se' might actually be 'Sent' with the 'nt' " + "lost to redaction artifacts.\n\n" + "Return the clean, corrected text that should appear to the left and right of " + "the redaction. Fix any obvious OCR errors near the boundary. Do not include " + "any redaction artifacts (brackets, pipes, random characters). " + "If the redaction is at the start or end of the line, use an empty string " + "for that side." + ) + + +async def identify_boundary_text( + line: OcrLine, + box_x: int, + box_w: int, +) -> BoundaryText: + """Use LLM to identify clean boundary text around a single redaction. + + Sends a focused prompt with just the one OCR line and redaction position. + Falls back to center-point character filtering if the LLM call fails. + + Args: + line: The OCR line containing the redaction. + box_x: X pixel position of the redaction box. + box_w: Width of the redaction box in pixels. + + Returns: + BoundaryText with cleaned left and right text. + """ + # Fallback in case LLM is unavailable + def _fallback(): + left = [c for c in line.chars if c.x + c.w / 2 < box_x] + right = [c for c in line.chars if c.x + c.w / 2 > box_x + box_w] + return BoundaryText( + left_text="".join(c.text for c in left).strip(), + right_text="".join(c.text for c in right).strip(), + ) + + try: + model = os.environ.get("UNREDACT_LLM_MODEL", _DEFAULT_MODEL) + prompt = _build_boundary_prompt(line, box_x, box_w) + client = _get_client() + + response = await client.messages.create( + model=model, + max_tokens=512, + tools=[_BOUNDARY_TOOL], + tool_choice={"type": "tool", "name": "report_boundary_text"}, + messages=[{"role": "user", "content": prompt}], + ) + + for block in response.content: + if block.type == "tool_use" and block.name == "report_boundary_text": + return BoundaryText( + left_text=block.input.get("left_text", "").strip(), + right_text=block.input.get("right_text", "").strip(), + ) + + log.warning("LLM boundary response had no tool call, using fallback") + return _fallback() + + except Exception: + log.warning("LLM boundary call failed, using center-point fallback", exc_info=True) + return _fallback() + + async def detect_redactions_llm( lines: list[OcrLine], ) -> list[LlmRedaction]: