diff --git a/tests/test_llm_detect.py b/tests/test_llm_detect.py new file mode 100644 index 0000000..9ddae5a --- /dev/null +++ b/tests/test_llm_detect.py @@ -0,0 +1,363 @@ +"""Tests for LLM-based redaction detection.""" + +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +from unredact.pipeline.ocr import OcrChar, OcrLine +from unredact.pipeline.llm_detect import ( + LlmRedaction, + _build_prompt, + _find_word_in_chars, + _parse_response, + detect_redactions_llm, +) + + +# --------------------------------------------------------------------------- +# Helpers to build OcrLine fixtures +# --------------------------------------------------------------------------- + +def _make_char(ch: str, x: int, w: int = 10, y: int = 100, h: int = 20) -> OcrChar: + return OcrChar(text=ch, x=x, y=y, w=w, h=h, conf=95.0) + + +def _make_line_from_words( + words: list[str], + start_x: int = 50, + char_w: int = 10, + space_w: int = 8, + y: int = 100, + h: int = 20, +) -> OcrLine: + """Build an OcrLine from a list of words with evenly-spaced characters.""" + chars: list[OcrChar] = [] + x = start_x + for wi, word in enumerate(words): + for ch in word: + chars.append(_make_char(ch, x, w=char_w, y=y, h=h)) + x += char_w + if wi < len(words) - 1: + chars.append(_make_char(" ", x, w=space_w, y=y, h=h)) + x += space_w + line_x = chars[0].x + line_w = (chars[-1].x + chars[-1].w) - line_x + return OcrLine(chars=chars, x=line_x, y=y, w=line_w, h=h) + + +# --------------------------------------------------------------------------- +# _build_prompt +# --------------------------------------------------------------------------- + +class TestBuildPrompt: + def test_includes_all_lines(self): + lines = [ + _make_line_from_words(["Hello", "world"]), + _make_line_from_words(["Second", "line"]), + ] + prompt = _build_prompt(lines) + assert "Hello world" in prompt + assert "Second line" in prompt + + def test_includes_line_indices(self): + lines = [ + _make_line_from_words(["Line", "zero"]), + _make_line_from_words(["Line", "one"]), + ] + prompt = _build_prompt(lines) + # Should have line numbers / indices + assert "0" in prompt + assert "1" in prompt + + def test_empty_lines(self): + prompt = _build_prompt([]) + # Should not crash; should produce valid string + assert isinstance(prompt, str) + + +# --------------------------------------------------------------------------- +# _find_word_in_chars +# --------------------------------------------------------------------------- + +class TestFindWordInChars: + def test_find_first_word(self): + line = _make_line_from_words(["Hello", "world", "foo"]) + result = _find_word_in_chars(line, "Hello", search_from=0, from_right=False) + assert result is not None + start_x, end_x = result + # "Hello" starts at char 0 (x=50) and ends at char 4 (x=90, end=100) + assert start_x == 50 + assert end_x == 50 + 5 * 10 # 5 chars * 10px each + + def test_find_last_word(self): + line = _make_line_from_words(["Hello", "world", "foo"]) + result = _find_word_in_chars(line, "foo", search_from=0, from_right=False) + assert result is not None + + def test_find_from_right(self): + line = _make_line_from_words(["the", "cat", "the", "dog"]) + # from_right should find the last occurrence + result = _find_word_in_chars( + line, "the", search_from=0, from_right=True + ) + assert result is not None + # The second "the" starts further right + start_x, _ = result + # First "the" starts at x=50; second "the" should start later + assert start_x > 50 + + def test_word_not_found(self): + line = _make_line_from_words(["Hello", "world"]) + result = _find_word_in_chars(line, "missing", search_from=0, from_right=False) + assert result is None + + def test_search_from_offset(self): + line = _make_line_from_words(["Hello", "world"]) + # search_from past where "Hello" starts should still find "world" + hello_end_x = 50 + 5 * 10 + 8 # past hello + space + result = _find_word_in_chars( + line, "world", search_from=hello_end_x, from_right=False + ) + assert result is not None + + def test_empty_word_returns_none(self): + line = _make_line_from_words(["Hello", "world"]) + result = _find_word_in_chars(line, "", search_from=0, from_right=False) + assert result is None + + +# --------------------------------------------------------------------------- +# _parse_response +# --------------------------------------------------------------------------- + +class TestParseResponse: + def test_single_redaction_middle_of_line(self): + # Line: "The [REDACTED] cat" + # LLM says redaction between "The" and "cat" on line 0 + line = _make_line_from_words(["The", "[REDACTED]", "cat"]) + tool_input = { + "redactions": [ + {"line_index": 0, "left_word": "The", "right_word": "cat"}, + ] + } + results = _parse_response(tool_input, [line]) + assert len(results) == 1 + r = results[0] + assert isinstance(r, LlmRedaction) + assert r.line_index == 0 + assert r.left_word == "The" + assert r.right_word == "cat" + assert r.line_y == line.y + assert r.line_h == line.h + # left_x should be the right edge of "The" + # right_x should be the left edge of "cat" + assert r.left_x < r.right_x + + def test_multiple_redactions_same_line(self): + line = _make_line_from_words(["A", "XXX", "B", "YYY", "C"]) + tool_input = { + "redactions": [ + {"line_index": 0, "left_word": "A", "right_word": "B"}, + {"line_index": 0, "left_word": "B", "right_word": "C"}, + ] + } + results = _parse_response(tool_input, [line]) + assert len(results) == 2 + # First redaction should be left of second + assert results[0].right_x <= results[1].left_x + + def test_no_redactions(self): + line = _make_line_from_words(["Clean", "text", "here"]) + tool_input = {"redactions": []} + results = _parse_response(tool_input, [line]) + assert results == [] + + def test_redaction_at_line_start(self): + # Redaction at the very beginning — empty left_word + line = _make_line_from_words(["[REDACTED]", "world"]) + tool_input = { + "redactions": [ + {"line_index": 0, "left_word": "", "right_word": "world"}, + ] + } + results = _parse_response(tool_input, [line]) + assert len(results) == 1 + r = results[0] + assert r.left_word == "" + # left_x should be the line start + assert r.left_x == line.x + + def test_redaction_at_line_end(self): + # Redaction at the very end — empty right_word + line = _make_line_from_words(["Hello", "[REDACTED]"]) + tool_input = { + "redactions": [ + {"line_index": 0, "left_word": "Hello", "right_word": ""}, + ] + } + results = _parse_response(tool_input, [line]) + assert len(results) == 1 + r = results[0] + assert r.right_word == "" + # right_x should be the line end + assert r.right_x == line.x + line.w + + def test_invalid_line_index_skipped(self): + line = _make_line_from_words(["Hello", "world"]) + tool_input = { + "redactions": [ + {"line_index": 99, "left_word": "Hello", "right_word": "world"}, + ] + } + results = _parse_response(tool_input, [line]) + assert results == [] + + def test_word_not_found_skipped(self): + line = _make_line_from_words(["Hello", "world"]) + tool_input = { + "redactions": [ + {"line_index": 0, "left_word": "MISSING", "right_word": "world"}, + ] + } + results = _parse_response(tool_input, [line]) + assert results == [] + + def test_multiple_lines(self): + line0 = _make_line_from_words(["A", "XXX", "B"], y=100) + line1 = _make_line_from_words(["C", "YYY", "D"], y=200) + tool_input = { + "redactions": [ + {"line_index": 0, "left_word": "A", "right_word": "B"}, + {"line_index": 1, "left_word": "C", "right_word": "D"}, + ] + } + results = _parse_response(tool_input, [line0, line1]) + assert len(results) == 2 + assert results[0].line_y == 100 + assert results[1].line_y == 200 + + +# --------------------------------------------------------------------------- +# detect_redactions_llm (mocked) +# --------------------------------------------------------------------------- + +class TestDetectRedactionsLlm: + def _make_mock_response(self, tool_input: dict) -> MagicMock: + """Build a mock Anthropic API response with a tool_use content block.""" + tool_block = MagicMock() + tool_block.type = "tool_use" + tool_block.name = "report_redactions" + tool_block.input = tool_input + + response = MagicMock() + response.content = [tool_block] + response.stop_reason = "tool_use" + return response + + @pytest.mark.anyio + async def test_single_redaction(self): + line = _make_line_from_words(["The", "|||", "cat"]) + mock_response = self._make_mock_response({ + "redactions": [ + {"line_index": 0, "left_word": "The", "right_word": "cat"}, + ] + }) + + mock_client = AsyncMock() + mock_client.messages.create = AsyncMock(return_value=mock_response) + + with patch("unredact.pipeline.llm_detect._get_client", return_value=mock_client): + results = await detect_redactions_llm([line]) + + assert len(results) == 1 + assert results[0].left_word == "The" + assert results[0].right_word == "cat" + mock_client.messages.create.assert_called_once() + + @pytest.mark.anyio + async def test_no_redactions(self): + line = _make_line_from_words(["Clean", "text"]) + mock_response = self._make_mock_response({"redactions": []}) + + mock_client = AsyncMock() + mock_client.messages.create = AsyncMock(return_value=mock_response) + + with patch("unredact.pipeline.llm_detect._get_client", return_value=mock_client): + results = await detect_redactions_llm([line]) + + assert results == [] + + @pytest.mark.anyio + async def test_empty_lines_list(self): + """When there are no lines, should return empty without calling LLM.""" + mock_client = AsyncMock() + mock_client.messages.create = AsyncMock() + + with patch("unredact.pipeline.llm_detect._get_client", return_value=mock_client): + results = await detect_redactions_llm([]) + + assert results == [] + mock_client.messages.create.assert_not_called() + + @pytest.mark.anyio + async def test_multiple_redactions_across_lines(self): + line0 = _make_line_from_words(["Foo", "|||", "bar"], y=100) + line1 = _make_line_from_words(["Baz", "XXX", "qux"], y=200) + mock_response = self._make_mock_response({ + "redactions": [ + {"line_index": 0, "left_word": "Foo", "right_word": "bar"}, + {"line_index": 1, "left_word": "Baz", "right_word": "qux"}, + ] + }) + + mock_client = AsyncMock() + mock_client.messages.create = AsyncMock(return_value=mock_response) + + with patch("unredact.pipeline.llm_detect._get_client", return_value=mock_client): + results = await detect_redactions_llm([line0, line1]) + + assert len(results) == 2 + assert results[0].line_index == 0 + assert results[1].line_index == 1 + + @pytest.mark.anyio + async def test_model_configurable_via_env(self): + line = _make_line_from_words(["Hello", "world"]) + mock_response = self._make_mock_response({"redactions": []}) + + mock_client = AsyncMock() + mock_client.messages.create = AsyncMock(return_value=mock_response) + + with ( + patch("unredact.pipeline.llm_detect._get_client", return_value=mock_client), + patch.dict("os.environ", {"UNREDACT_LLM_MODEL": "claude-test-model"}), + ): + await detect_redactions_llm([line]) + + call_kwargs = mock_client.messages.create.call_args + assert call_kwargs.kwargs.get("model") == "claude-test-model" or \ + (call_kwargs.args and call_kwargs.args[0] == "claude-test-model") + + @pytest.mark.anyio + async def test_redaction_at_line_start_and_end(self): + line = _make_line_from_words(["XXX", "middle", "YYY"]) + mock_response = self._make_mock_response({ + "redactions": [ + {"line_index": 0, "left_word": "", "right_word": "middle"}, + {"line_index": 0, "left_word": "middle", "right_word": ""}, + ] + }) + + mock_client = AsyncMock() + mock_client.messages.create = AsyncMock(return_value=mock_response) + + with patch("unredact.pipeline.llm_detect._get_client", return_value=mock_client): + results = await detect_redactions_llm([line]) + + assert len(results) == 2 + # First redaction at line start + assert results[0].left_word == "" + assert results[0].left_x == line.x + # Second redaction at line end + assert results[1].right_word == "" + assert results[1].right_x == line.x + line.w diff --git a/unredact/pipeline/llm_detect.py b/unredact/pipeline/llm_detect.py new file mode 100644 index 0000000..e76e0b3 --- /dev/null +++ b/unredact/pipeline/llm_detect.py @@ -0,0 +1,264 @@ +"""LLM-based redaction detection. + +Receives OCR lines and calls Claude Haiku to identify where redactions are. +The LLM spots broken text patterns (artifacts like ``[``, ``|``, garbled chars) +and returns the clean boundary words around each redaction. +""" + +from __future__ import annotations + +import logging +import os +from dataclasses import dataclass + +import anthropic + +from unredact.pipeline.ocr import OcrLine + +log = logging.getLogger(__name__) + +_DEFAULT_MODEL = "claude-haiku-4-5-20251001" + +_TOOL = { + "name": "report_redactions", + "description": "Report redacted sections found in the OCR text.", + "input_schema": { + "type": "object", + "properties": { + "redactions": { + "type": "array", + "items": { + "type": "object", + "properties": { + "line_index": { + "type": "integer", + "description": "Zero-based line index", + }, + "left_word": { + "type": "string", + "description": ( + "Last clean word before redaction. " + "Empty if at line start." + ), + }, + "right_word": { + "type": "string", + "description": ( + "First clean word after redaction. " + "Empty if at line end." + ), + }, + }, + "required": ["line_index", "left_word", "right_word"], + }, + }, + }, + "required": ["redactions"], + }, +} + +# Module-level cached client +_client: anthropic.AsyncAnthropic | None = None + + +def _get_client() -> anthropic.AsyncAnthropic: + """Return a cached async Anthropic client.""" + global _client + if _client is None: + _client = anthropic.AsyncAnthropic() + return _client + + +@dataclass +class LlmRedaction: + """A redaction detected by the LLM, anchored to OCR pixel positions.""" + + line_index: int + left_word: str + right_word: str + left_x: int + right_x: int + line_y: int + line_h: int + + +def _build_prompt(lines: list[OcrLine]) -> str: + """Format OCR lines into the LLM prompt.""" + header = ( + "You are analyzing OCR output from a scanned legal document. " + "Some parts of the document have been redacted (blacked out). " + "The OCR engine often produces garbled text, brackets, pipes, or " + "random characters where redactions are.\n\n" + "Identify each redacted section. For each one, report:\n" + "- The zero-based line index\n" + "- The last clean word BEFORE the redaction (empty string if " + "redaction is at the start of the line)\n" + "- The first clean word AFTER the redaction (empty string if " + "redaction is at the end of the line)\n\n" + "OCR lines:\n" + ) + + line_texts = [] + for i, line in enumerate(lines): + line_texts.append(f"[{i}] {line.text}") + + return header + "\n".join(line_texts) + + +def _find_word_in_chars( + line: OcrLine, + word: str, + search_from: int, + from_right: bool, +) -> tuple[int, int] | None: + """Find a word in the OCR char list and return (start_x, end_x). + + Args: + line: The OCR line to search in. + word: The word to find. + search_from: Only consider matches starting at or after this x position. + from_right: If True, return the rightmost (last) match instead of the first. + + Returns: + (start_x, end_x) pixel coordinates, or None if not found. + """ + if not word: + return None + + chars = line.chars + text = line.text + best: tuple[int, int] | None = None + + # Find all occurrences of `word` in the line text + start = 0 + while True: + idx = text.find(word, start) + if idx == -1: + break + + # Map character index back to pixel position + char_start = chars[idx] + char_end = chars[idx + len(word) - 1] + start_x = char_start.x + end_x = char_end.x + char_end.w + + if start_x >= search_from: + if from_right: + best = (start_x, end_x) + else: + return (start_x, end_x) + + start = idx + 1 + + return best + + +def _parse_response( + tool_input: dict, + lines: list[OcrLine], +) -> list[LlmRedaction]: + """Map the LLM's tool response back to pixel positions using OCR char data. + + Args: + tool_input: The ``input`` dict from the LLM's tool_use content block. + lines: The OCR lines that were sent to the LLM. + + Returns: + List of LlmRedaction objects with pixel-accurate positions. + """ + redactions: list[LlmRedaction] = [] + + for item in tool_input.get("redactions", []): + line_index = item["line_index"] + left_word = item["left_word"] + right_word = item["right_word"] + + # Validate line index + if line_index < 0 or line_index >= len(lines): + log.warning("LLM returned invalid line_index %d, skipping", line_index) + continue + + line = lines[line_index] + + # Determine left_x: right edge of left_word, or line start + if left_word: + found = _find_word_in_chars(line, left_word, search_from=0, from_right=False) + if found is None: + log.warning( + "Could not find left_word %r in line %d, skipping", + left_word, + line_index, + ) + continue + left_x = found[1] # right edge of left_word + else: + left_x = line.x + + # Determine right_x: left edge of right_word, or line end + if right_word: + search_after = left_x if left_word else 0 + found = _find_word_in_chars( + line, right_word, search_from=search_after, from_right=False + ) + if found is None: + log.warning( + "Could not find right_word %r in line %d, skipping", + right_word, + line_index, + ) + continue + right_x = found[0] # left edge of right_word + else: + right_x = line.x + line.w + + redactions.append( + LlmRedaction( + line_index=line_index, + left_word=left_word, + right_word=right_word, + left_x=left_x, + right_x=right_x, + line_y=line.y, + line_h=line.h, + ) + ) + + return redactions + + +async def detect_redactions_llm( + lines: list[OcrLine], +) -> list[LlmRedaction]: + """Detect redactions by sending OCR text to Claude. + + This is the main entry point. Calls Claude Haiku with tool use for + structured output, then maps the response back to pixel positions. + + Args: + lines: OCR lines from the page. + + Returns: + List of detected redactions with pixel positions. + """ + if not lines: + return [] + + model = os.environ.get("UNREDACT_LLM_MODEL", _DEFAULT_MODEL) + prompt = _build_prompt(lines) + client = _get_client() + + response = await client.messages.create( + model=model, + max_tokens=1024, + tools=[_TOOL], + tool_choice={"type": "tool", "name": "report_redactions"}, + messages=[{"role": "user", "content": prompt}], + ) + + # Extract the tool_use block from the response + for block in response.content: + if block.type == "tool_use" and block.name == "report_redactions": + return _parse_response(block.input, lines) + + log.warning("LLM response did not contain a report_redactions tool call") + return []