Add PII scan and auto-annotate agent tool - #1624
Conversation
Adds `scan_and_annotate_pii`, an agent tool that scans documents for sensitive information via the privacy-filter microservice and creates one labeled annotation per detection (TOKEN_LABEL via PlasmaPDF for PDFs, SPAN_LABEL for plain text). - New httpx async client with chunking, offset remap, and dedup - 8 PII categories auto-mapped to colored/iconed AnnotationLabels - Agent knobs: min_score, entity_groups, dry_run, start_char/end_char - New privacy_filter Docker service (local.yml, production.yml) - Settings: PRIVACY_FILTER_URL, PRIVACY_FILTER_API_KEY, _TIMEOUT_SECONDS - 18 unit tests; privacy-filter client mocked so CI needs no image
Code Review: Add PII scan and auto-annotate agent toolOverviewThis PR introduces a 🔴 Critical
# production.yml — no `required: false`
privacy_filter:
condition: service_started
The setting 🟠 HighUnhandled The # Current — network errors bubble up uncaught
resp = await client.post(detect_url, json={"text": chunk}, headers=headers)
# Suggested — wrap in try/except httpx.HTTPError
try:
resp = await client.post(detect_url, json={"text": chunk}, headers=headers)
except httpx.HTTPError as exc:
raise RuntimeError(f"privacy-filter request failed: {exc}") from exc
The diff doesn't touch any 🟡 Medium
If an agent passes File type handling is undocumented in the tool description The tool only supports No retry / backoff for transient HTTP failures A single 503 from the privacy-filter service immediately fails the entire scan. A simple retry (e.g., 2–3 attempts with exponential backoff) would make this more robust in practice, especially since large documents generate multiple sequential chunk requests. ⚪ Low / NitsFragile test assertion in assert idx >= 0, "Test fixture must contain the string 'Agreement'."The test is tightly coupled to the content of Duplicate
Irregular import order in The new import ✅ What's done well
|
- Mark privacy_filter as required: false in production.yml so existing prod deployments that don't run the PII service still start. - Wrap the httpx.post() call in adetect_pii in try/except httpx.HTTPError and re-raise as RuntimeError so transport-level failures land in the agent fault-tolerance layer instead of bubbling out raw. - Validate the entity_groups allowlist up-front in ascan_and_annotate_pii and raise ValueError for unknown names instead of silently filtering every detection. - Mention supported file types (PDF + plain text) in the scan_and_annotate_pii tool description. - Add regression tests for the new transport-failure path and the allowlist validation.
Code Review: PR #1624 — PII Scan & Auto-Annotate Agent ToolOverall this is a well-structured addition. The chunking/dedup logic is sound, the fault-tolerance contract with the agent framework is correctly respected, and the test coverage is solid. A few items worth addressing before merge: Bugs / Correctness1. Non-standard MIME type in 2. Import order in 3. Bare type annotation without assignment in allowlist: set[str] | None
if entity_groups:
...
allowlist = set(entity_groups)
else:
allowlist = NoneThe bare annotation on its own line is valid Python but unusual and could confuse static analysis tools that expect assignments before use. Prefer Security4. 5. Error messages could expose partial API key context Production Behaviour Inconsistency6. PR description says "fail-closed" but
But Code Quality / Minor Issues7. Return type annotation is too loose 8. 9. No Test CoverageThe test suite is thorough and covers:
One gap: there is no test for Summary
The core logic (chunking, dedup, PlasmaPDF integration, label auto-creation) is well-implemented and the fault-tolerance contract with |
- isort reordered core_tools/__init__.py to alphabetize the new pii import ahead of notes/page_images per project convention. - Removed bare-annotation-then-assign pattern in ascan_and_annotate_pii; the variable is now initialised at the annotation site so static analysers (and humans) see one clear binding. - Tightened the return type to dict[str, Any] so callers and IDEs get a more useful contract than bare ``dict``. - Added a docs comment to the local.yml privacy_filter service noting it is internal-bridge-only with no host port mapping (smoke-test from inside django container). - New regression test covers the inverted start_char/end_char case: the tool must short-circuit and must not issue a request to the privacy-filter service.
Code Review — PR #1624: PII Scan & Auto-Annotate Agent ToolOverall this is a solid, well-structured addition. The async patterns are correct, the test suite is thorough, and the failure-mode design (RuntimeError → agent error string) fits the existing fault-tolerance contract. A few items worth addressing: Potential Bug: Production Docker Compose Variable Expansion
environment:
API_KEYS: ${PRIVACY_FILTER_API_KEY:?PRIVACY_FILTER_API_KEY must be set}Docker Compose evaluates Options:
Minor IssuesMagic string file types in Per CLAUDE.md ("No magic numbers"), bare string literals like Non-standard MIME type retained
Sequential chunk processing
import asyncio
async def _fetch_chunk(client, detect_url, headers, chunk_start, text):
chunk = text[chunk_start : chunk_start + CHUNK_SIZE]
resp = await client.post(detect_url, json={"text": chunk}, headers=headers)
...
return chunk_start, detections
results = await asyncio.gather(*[_fetch_chunk(..., s, text) for s in starts])Redundant bounds check
The Test Coverage — PraiseThe 18-test suite is comprehensive:
The one gap: there's no test verifying that Nit:
|
| Area | Status |
|---|---|
| Async correctness | ✅ All tools are async, DB calls go through _db_sync_to_async |
| Tool registration | ✅ Registered with requires_approval=True, requires_write_permission=True |
| Error contract | ✅ RuntimeError → agent error string path works correctly |
| Test coverage | ✅ 18 tests, comprehensive |
httpx dependency |
✅ Already in requirements/base.txt |
| Production deploy safety | ${...:?} expansion may break opt-out deployments |
| Constants convention | |
| Performance | ℹ️ Sequential chunks; parallel would be a nice follow-up |
Code Review: PII Scan & Auto-Annotate Agent ToolOverviewThis PR adds a Bugs / Correctness Issues🔴 High:
But the service definition itself uses: API_KEYS: ${PRIVACY_FILTER_API_KEY:?PRIVACY_FILTER_API_KEY must be set}Docker Compose performs variable substitution before trying to start any container. If Fix options (pick one):
🟡 Medium: Overlap dedup discards higher-confidence scores In This is not a crash bug, but it means annotations can carry the lower confidence score rather than the best one. Consider keeping the detection with the maximum # Instead of `if key in seen: continue`, track scores:
if key in seen:
if float(det["score"]) > results[seen[key]]["score"]:
results[seen[key]] = Detection(...) # replace with better score
continue
seen[key] = len(results)
results.append(...)Security ConsiderationsNo issues with the API key handling — it goes over the internal Docker bridge only, which is the correct pattern. Permission enforcement: The Code Quality✅ Good patterns throughout:
Minor code style notes:
Infrastructure🟡 No healthcheck on Consider adding: healthcheck:
test: ["CMD", "curl", "-sf", "http://localhost:8000/health"]
interval: 10s
retries: 5and changing the dependency to No resource limits: The privacy_filter service runs an ML model CPU-only. Without memory/CPU limits it can starve other containers on resource-constrained hosts. Not blocking, but worth a follow-up. Test CoverageThe 18 tests are well-structured and cover the important cases. A few gaps worth noting:
Summary
The production.yml env var issue is the only change I'd consider blocking — everything else is polished and well-thought-out. Great addition overall. |
- production.yml: switch privacy_filter API_KEYS to empty default
(`${PRIVACY_FILTER_API_KEY:-}`). The previous `:?required` form failed
Compose interpolation at parse time, so the whole production stack
refused to start unless the env var was set — contradicting the
`required: false` opt-in declared on the depends_on entry.
- production.yml + local.yml: add HTTP /health healthcheck and a 20s
start_period so dependents can wait on `service_healthy` once the
model weights are loaded, instead of `service_started` which only
proves the container process spawned.
- _privacy_filter_client.py: when chunk overlap produces a duplicate
(start, end, entity_group) detection, keep the higher-confidence
score rather than first-seen. Switches the dedup tracker from
`set` to `dict[..., int]` indexing into `results`.
Code Review: PII Scan & Auto-Annotate Agent ToolOverall this is a well-structured addition. The async HTTP client, chunking/overlap/dedup logic, tool registration, and test coverage are all solid. A few issues worth addressing before merge. Medium Issues1.
# Current (pii.py line ~75)
if file_type in {"application/txt", "text/plain"}:
# Should be
from opencontractserver.constants.document_processing import TEXT_MIMETYPES
if file_type in TEXT_MIMETYPES:This also silently excludes Markdown files from PII scanning, which may be intentional, but worth a comment if so. 2.
3. Production API key silently defaults to empty The CHANGELOG test plan says "stack will fail fast at startup" if 4. No retry logic in the HTTP client A single transient timeout or TCP reset permanently fails the entire tool invocation — the LLM sees an error string and the whole annotation run is lost. Given that large documents trigger many sequential HTTP requests (every 39.5 KB of text), even a 1% per-request failure rate compounds quickly. A simple retry-with-backoff (e.g. Minor Issues5. Hardcoded chunk constants should be in the constants file Per project conventions ( 6. The unnamed 5-tuple @dataclass
class _DocTextResult:
doc: Document
corpus: Corpus
doc_text: str
file_type: str
pdf_layer: Any # None for non-PDFNits
What's Working Well
|
Codecov Report❌ Patch coverage is
📢 Thoughts on this report? Let us know! |
CI linter (pre-existing on main):
- Widen requests.post json= payloads in users/tasks.py, docxodus_parser,
and the two embedder microservices from invariant dict[str, str]
literals to dict[str, Any] so they satisfy JsonType.
PII scan tool review feedback:
- Use TEXT_MIMETYPES from constants instead of inline {"text/plain",
"application/txt"} — picks up markdown alongside plain-text and keeps
the supported-types list in lockstep with the ingestion pipeline.
- Move CHUNK_SIZE / CHUNK_OVERLAP into constants/document_processing.py
(PRIVACY_FILTER_CHUNK_SIZE / _OVERLAP) so the magic numbers are
discoverable per CLAUDE.md conventions.
- Convert _load_doc_text_sync's positional 5-tuple to _DocTextResult
NamedTuple — unpacking sites still work positionally, and named fields
make the contract self-documenting and grep-friendly.
- Make _persist_annotations_sync return (annotation_id, detection) pairs
and derive detection_count / by_entity_group from those persisted pairs
rather than the input detections list. Previously, if any detection
was filtered out by the OOB / unknown-group guards inside persist,
detection_count would overcount relative to annotation_ids — a
confusing failure mode for LLM callers debugging a slice-offset bug.
Code Review: Add PII scan and auto-annotate agent toolOverall assessment: This is a well-designed, well-tested addition. The architecture is sound, the error-handling contract with the tool fault-tolerance layer is well-documented, and the test coverage is solid. A few items worth addressing before merge. IssuesSecurity: Production service accepts unauthenticated requests when
|
| Priority | Item |
|---|---|
| Medium | Production service starts with empty API_KEYS if env var is unset |
| Low | Dead code in _label_for (unreachable ValueError) |
| Low | Race condition in ensure_label_and_labelset when called inside atomic() |
| Nit | Unrelated type-annotation changes in 4 files should be a separate commit |
| Nit | Consider parallelising chunk requests with asyncio.gather |
ensure_label_and_labelset uses a check-then-create pattern that races under concurrent scans. When the losing race fires INSERT inside the outer transaction.atomic() block, the resulting IntegrityError poisons the savepoint and rolls back every annotation created by the current scan. Pull the label get_or_create work up before the atomic block — the inner block now only inserts Annotation rows, which carry no cross-row uniqueness constraint. Also drops the now-dead ValueError inside _label_for: the outer loop already filters on ENTITY_GROUP_LABELS membership, so a label fetch for an unknown group is unreachable. Removing the parallel guard keeps the validation in one place.
Code Review: Add PII scan and auto-annotate agent toolOverall this is a well-structured addition. The chunking logic, dedup strategy, async plumbing, and test coverage are all solid. A few things worth addressing before merge: Bugs / CorrectnessMixed-document-type corpora will create duplicate label records
If SecurityProduction service starts with empty API key when
The service has no
Design / Maintainability
The function name implies lazy iteration but it materialises the whole list. Minor, but for very large documents the whole list is trivially small so this isn't a performance concern — it's just a naming mismatch. Chunks are sent sequentially; no backoff or concurrency
Code Quality
The Comment in )(
document_id, corpus_id
) # NamedTuple — fields stay named for IDE/grepThe TestsTest coverage is genuinely comprehensive — chunking dedup, offset remapping, OOB guard, all eight entity categories, One style note: Summary
|
|
Test review comment. |
|
Code Review: PR 1624 - PII Scan and Auto-Annotate Agent Tool Overview Well-structured PR overall: clean async/sync separation, solid overlap-dedup logic, thoughtful edge-case handling, good test coverage. A few things worth addressing before merge. Issue 1 - Performance: N+1 saves inside the atomic block In pii.py _persist_annotations_sync, ann.save() is called once per detection inside transaction.atomic(). For a document with hundreds of PII hits this becomes N individual INSERTs. Prefer Annotation.objects.bulk_create() with a pre-built list. bulk_create returns created objects with PKs populated in Django >= 4.1, so the (ann.pk, det) pairs callers rely on are still available. Issue 2 - Security: production service can start with an empty API key production.yml now uses PRIVACY_FILTER_API_KEY with a fail-open fallback. If an operator forgets to set the variable, the container starts with no key - any caller on the docker bridge can POST to it unauthenticated. There is no ports mapping so external exposure is low, but worth documenting. Consider a Django system-check or startup log warning when PRIVACY_FILTER_URL is set but PRIVACY_FILTER_API_KEY is empty. Issue 3 - Label pre-creation race: comment could point to the guarantee The comment in _persist_annotations_sync correctly identifies the race and mentions that ensure_label_and_labelset uses a check-then-create pattern, but does not state what prevents a duplicate-key error when two callers collide simultaneously. A one-liner pointing to the DB-level uniqueness constraint or the get_or_create call inside that method would save future readers from tracing the call chain. Issue 4 - Minor: env() vs env.str() inconsistency New settings use bare env(), but every other string setting in base.py uses env.str(). Prefer env.str() for consistency. Issue 5 - Missing test: corpus_action_id propagation None of the 18 tests pass a non-None corpus_action_id and verify it appears on the created annotation. A single round-trip test would close this gap. Non-blocking observations
Summary Two items worth addressing before merge: the N+1 save() loop (performance) and the unauthenticated-service risk (security). The env.str() inconsistency and missing corpus_action_id test are minor. The async architecture, error-to-RuntimeError wrapping, dry-run mode, and the pre-creation-outside-atomic label pattern are all solid. |
…ace comment, cover corpus_action_id propagation - _privacy_filter_client emits a one-time warning when PRIVACY_FILTER_URL is set but PRIVACY_FILTER_API_KEY is empty so misconfigurations show up in logs instead of silently shipping with an unauthenticated service. - Rewrite the label pre-creation race comment in _persist_annotations_sync to explain *why* doing get-or-create outside the atomic block is the correct trade-off (no DB-level uniqueness on AnnotationLabel; the pre-creation isolates a rare duplicate to the labels rather than rolling back the whole annotation batch). - New test_corpus_action_id_propagates_to_created_annotations covers the round-trip from CorpusAction.id → kwarg → Annotation.corpus_action_id so future regressions in the propagation path don't slip through.
![github-code-quality[bot]](https://avatars.githubusercontent.com/u/9919?s=60&v=4){kind=small}
| "requests will be sent unauthenticated. Set PRIVACY_FILTER_API_KEY " | ||
| "to gate access to the privacy-filter service." | ||
| ) | ||
| _warned_about_missing_api_key = True |
Code Review: Add PII scan and auto-annotate agent toolOverall, this is a well-structured PR with solid test coverage, clear separation of concerns, and good documentation. The chunking/dedup logic in the HTTP client is particularly clean. A few issues worth addressing before merge: Bugs / Correctness1. # Instead of:
ann.save()
persisted.append((ann.pk, det))
# Consider:
anns_to_create = [...]
created = Annotation.objects.bulk_create(anns_to_create)
persisted = list(zip([a.pk for a in created], filtered_dets))2. Security3. Production 4. No explicit permission guard in Code Quality5. Module-level mutable global 6. 7. Noisy type-annotation-only diffs in unrelated files Test Coverage Gaps8. No test for the empty-API-key warning path 9. No test for the concurrent-label-creation race 10. Minor Nits
SummaryThe core logic (chunked HTTP client, offset remapping, overlap dedup, PDF vs text annotation dispatch) is sound and well-tested. The main items to address before merge:
|
…ote, empty-API-key warning test - _iter_chunk_starts: document the case where total_len lands just past a chunk boundary and the final chunk overlaps by more than CHUNK_OVERLAP (de-dup on (start, end, group) keeps the over-coverage harmless). - _persist_annotations_sync: clarify that ensure_label_and_labelset already wraps its check-then-create in transaction.atomic(), but READ COMMITTED isolation does not prevent the duplicate-label race because there is no DB-level uniqueness on AnnotationLabel(text, label_type). Pre-creating outside the inner atomic block isolates the rare duplicate to the label table so the larger Annotation insert batch still commits. - _load_doc_text_sync: add docstring note explaining the helper deliberately has no in-line authorization check; the agent tool framework gates writes via requires_write_permission=True and the function is module-private precisely because direct callers would bypass that gate. - Add test_empty_api_key_logs_warning_once covering the suppression-after- first-emit path on _warned_about_missing_api_key, with explicit reset in try/finally so the test is order-independent.
Code Review — PR #1624: PII Scan & Auto-Annotate Agent ToolOverviewThis PR adds a well-architected Issues1.
|
* Address PR #1624 review follow-ups (closes #1638) Six items raised in #1638 against the PII scan agent tool that landed in PR #1624: - _persist_annotations_sync now uses Annotation.objects.bulk_create with manual transaction.on_commit embedding-task queueing — mirroring the established pattern in annotations/models.py — so a document with 200+ PII hits no longer pays 200+ round-trips inside the atomic block. PDF TOKEN_LABEL JSON is explicitly compact_annotation_json()'d before bulk_create because Annotation.save()'s auto-compaction is skipped on the bulk path. - plasmapdf.models.types imports hoisted to module scope in pii.py so static analyzers stop flagging SpanAnnotation / TextSpan as possibly-undefined inside the for-loop's PDF branch. - _privacy_filter_client.py no longer re-exports CHUNK_SIZE / CHUNK_OVERLAP as aliases. Tests import the canonical PRIVACY_FILTER_CHUNK_SIZE / PRIVACY_FILTER_CHUNK_OVERLAP names directly from opencontractserver.constants.document_processing. - New agents.W001 Django system check warns on startup when PRIVACY_FILTER_URL is set but PRIVACY_FILTER_API_KEY is empty. Production env template grows a Privacy Filter section calling out the same risk. The existing one-shot logger.warning in _privacy_filter_client.py on first request stays in place. - New test_persist_annotations_sync_duplicate_label_race regression test documents the accepted-duplicate behavior of Corpus.ensure_label_and_labelset under PostgreSQL READ COMMITTED (no DB-level uniqueness on AnnotationLabel(text, label_type)). Simulates the post-race outcome via patch.object on ensure_label_and_labelset and asserts duplicate labels exist while both annotation inserts still succeed. - httpx is already declared in requirements/base.txt (line 13) — no change needed, item closed via verification. * Address PR #1642 review: fix mypy lambda + pin unknown-group filter test * Address PR review: fix pytest failures, add agents.W001 tests, refactor - ScanAndAnnotateKnobsTests (3 failing tests) and sibling persist-side PII tool test classes pin serialized_rollback=True. _persist_annotations_sync registers transaction.on_commit Celery enqueues against calculate_embedding_for_annotation_text, which under CELERY_TASK_ALWAYS_EAGER=True runs the eager task body once the atomic block commits. The eager task reads the default embedder path from the PipelineSettings singleton seeded by migration 0031. TransactionTestCase truncates *all* tables including documents_pipelinesettings between tests by default, so a prior class on the same pytest-xdist worker can leave us with no default embedder and the eager retry chain raises out of the on_commit callback. serialized_rollback=True restores the migration-seeded row after each test's truncation. Same fix applied to every PII test class that drives _persist_annotations_sync. - agents.W001 check moved to Tags.security so manage.py check --tag security surfaces it alongside the rest of the security family (PR review feedback). - _queue_embed closure factory lifted to module scope (PR review feedback) so it isn't redefined on every _persist_annotations_sync call. Added defensive assert that bulk_create populated ann.pk per PostgreSQL RETURNING — the test suite runs against PostgreSQL so the assert never fires in practice, but surfaces a clearer failure mode than a None-keyed task if the backend ever changes. - New test_agents_system_checks.py covers the four PRIVACY_FILTER_URL × PRIVACY_FILTER_API_KEY truth-table cases plus a whitespace-treated-as- unset case and a registered-with-security-tag invariant. * Mock _queue_embed in PII tests instead of serialized_rollback Replaces the serialized_rollback=True approach (which produced django_content_type UniqueViolation errors at TransactionTestCase setup when multiple sibling classes both restored the migration-seeded fixture) with a surgical patch of the module-level _queue_embed factory to a no-op callback. The patch lives in _PiiPersistEmbeddingNoopMixin mixed into every persist-side test class; the two later persist-side classes (PersistAnnotationsLabelRaceTests, PersistAnnotationsUnknownGroupTests) inline the patch directly so the patch lifecycle stays obvious next to the rest of their setUp. Rationale spelled out in the mixin's docstring. Tests in this module don't exercise embedding behaviour at all, so stubbing the on-commit callback is the cleanest way to avoid the eager-celery + missing PipelineSettings cross-test ordering bug. * Fix mypy: broaden requests headers types to dict[str, str | bytes] types-requests upgrade made requests.post/get expect MutableMapping[str, str | bytes], which is invariant in values, so dict[str, str] no longer matches. Broaden header annotations and maybe_add_cloud_run_auth() signature accordingly. * Address PR review: replace assert in production path, apply mixin to race + unknown-group tests - pii._persist_annotations_sync: raise RuntimeError instead of assert when bulk_create returns rows without pk, so the invariant survives interpreters launched with -O. - pii._queue_embed docstring: drop PR-number reference per CLAUDE.md. - test_pii_scan_tool: apply _PiiPersistEmbeddingNoopMixin to PersistAnnotationsLabelRaceTests and PersistAnnotationsUnknownGroupTests (DRY); replace fragile slice-based new_labels query with snapshot-and-exclude idiom. * Apply black formatting to test_pii_scan_tool.py * Address review: hoist bulk batch size to constant, narrow JSON compaction exception, document race_ensure mock signature coupling --------- Co-authored-by: Claude <noreply@anthropic.com>
1 participant
Summary
scan_and_annotate_piiagent tool that scans documents for sensitive information via the privacy-filter microservice and creates one labeled annotation per detection — TOKEN_LABEL via PlasmaPDF for PDFs, SPAN_LABEL for plain text.AnnotationLabels viaCorpus.ensure_label_and_labelset.min_score,entity_groupsallowlist,dry_runpreview mode,start_char/end_charto scope the scan to a slice. Newprivacy_filterDocker service inlocal.yml(withrequired: falsefor dev opt-out) andproduction.yml(strict${PRIVACY_FILTER_API_KEY:?...}). Three new settings:PRIVACY_FILTER_URL,PRIVACY_FILTER_API_KEY,PRIVACY_FILTER_TIMEOUT_SECONDS.Test plan
docker compose -f test.yml run --rm django pytest opencontractserver/tests/test_pii_scan_tool.py opencontractserver/tests/test_privacy_filter_client.py -n 0 -v(expect 18 passed: 12 tool + 6 client)docker compose -f test.yml run --rm django python -c "from opencontractserver.llms.tools.tool_registry import ToolFunctionRegistry; print(ToolFunctionRegistry.get().to_core_tool('scan_and_annotate_pii') is not None)"printsTruedocker compose -f test.yml run --rm django python -c "from opencontractserver.llms.tools.core_tools import ascan_and_annotate_pii, ENTITY_GROUP_LABELS; print(len(ENTITY_GROUP_LABELS))"prints8docker compose -f local.yml up -d privacy_filter && docker compose -f local.yml exec django curl -sf -X POST http://privacy_filter:8000/v1/detect -H "X-API-Key: dev-only-not-secret" -H "Content-Type: application/json" -d '{"text":"Email me at alice@example.com tomorrow."}'PRIVACY_FILTER_API_KEYis set in the env or the stack will fail fast at startup (intentional — production fail-closed)