Compare commits

..
Author SHA1 Message Date
stumpylogandClaude Sonnet 5 a10f8fc81f chore: remove benchmark-commands spec/plan docs now that work is complete
Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
2026-09-09 07:42:25 -07:00
stumpylogandClaude Sonnet 5 5ac6ae52de Fix: apply final-review fix wave (safety guard, date spread, MariaDB explain, history fields, memory footprint, minors)
Final whole-branch review before this tooling branch settles: add a
required --yes-i-know-this-wipes-the-database confirmation flag ahead of
--reset (it deletes ALL users/groups/documents in the target DB, not just
benchmark-created ones); spread seeded documents' created dates over a
3-year window instead of leaving them all on one default date; fix
capture_explain() to use MariaDB's ANALYZE syntax instead of Postgres-only
EXPLAIN ANALYZE (verified against real MariaDB 12.3 -- the old code was a
silent 1064 syntax error); record db_vendor/document_count in run/profile
history entries; shrink SeededData's memory footprint at scale by
returning counts instead of full ORM instance tuples; and a handful of
minor fixes (storage_path assignment, --explain warning instead of silent
no-op, harness.py repeat<1 guard, type annotations). All changes verified
against real Postgres and MariaDB containers on the VM, not just SQLite.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
2026-09-09 07:42:25 -07:00
stumpylogandClaude Sonnet 5 235be8d15f Fix: correct /api/tags/ endpoint URL in benchmarking skill
The skill documentation stated that run benchmarks /api/tags/ but
src/paperless_benchmark/endpoints.py actually uses
/api/tags/?page_size=100000. Update the skill's Command reference
section to match the actual endpoint definition.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
2026-09-09 07:42:25 -07:00
stumpylog 69af169575 docs(benchmark): add paperless-benchmarking skill 2026-09-09 07:42:25 -07:00
stumpylog 6aecbac2c4 fix(benchmark): fully clear users/groups on reset, grant perf_target model perms 2026-09-09 07:42:25 -07:00
stumpylogandClaude Sonnet 5 0bc7221e12 docs: add Task 10 (quirk fixes) and Task 11 (skill) to benchmark plan
Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
2026-09-09 07:42:25 -07:00
stumpylog 95474bf1ae fix(benchmark): reuse seeded data in profile instead of reseeding
_handle_profile unconditionally called seed_benchmark_dataset every
run, colliding with an IntegrityError against data seeded by a prior
seed/profile call. Look up the existing perf_target user instead,
matching _handle_run's already-correct pattern, and raise a
CommandError pointing at `seed` if no dataset exists.

Scenario.run/queryset_for_explain now take a User directly instead of
a SeededData, since only data.users[0] (perf_target) was ever used.

profile no longer seeds, so it drops the --tier dependency: the
printed line and history entry no longer include a tier/(tier=...)
suffix, and --tier's help text now only mentions `seed`.
2026-09-09 07:42:24 -07:00
stumpylog e9bedd9343 feat(benchmark): add manage.py benchmark command (seed/run/profile/list-scenarios) 2026-09-09 07:42:24 -07:00
stumpylog cee6519a06 feat(benchmark): add pluggable profile-scenario registry 2026-09-09 07:42:24 -07:00
stumpylog c7b94952d2 feat(benchmark): add timed API endpoint scenarios 2026-09-09 07:42:24 -07:00
stumpylogandClaude Sonnet 5 bfae3ef82d docs: add implementation plan for benchmark management commands
Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
2026-09-09 07:42:24 -07:00
stumpylog 7b90eb5069 feat(benchmark): add local JSONL results history log 2026-09-09 07:42:24 -07:00
stumpylog e15dad04bb feat(benchmark): add tiered dataset seeding 2026-09-09 07:42:24 -07:00
stumpylog d0731ce892 feat(benchmark): add per-backend reset and explain helpers 2026-09-09 07:42:24 -07:00
stumpylog dda0c93e80 feat(benchmark): add run_profile timing harness 2026-09-09 07:42:24 -07:00
stumpylogandClaude Sonnet 5 a09ea28636 feat(benchmark): scaffold paperless_benchmark app
Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
2026-09-09 07:42:24 -07:00
stumpylogandClaude Sonnet 5 d5a9c502fc docs: add design spec for benchmark management commands
Merges the standalone run_benchmarks.py/seed_benchmark_data.py scripts
and the tools/profiling branch's Postgres-only harness into one
cross-backend `manage.py benchmark` command suite.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
2026-09-09 07:42:24 -07:00
64 changed files with 1838 additions and 6528 deletions
@@ -0,0 +1,187 @@
---
name: paperless-benchmarking
description: Use when profiling paperless-ngx performance, running `manage.py benchmark`, investigating a slow query or endpoint, or deciding whether a profiling finding should become a permanent registered scenario. Covers command reference, the fork/merge-back branch workflow for perf investigations, and how to read query-plan output.
---
# Paperless-ngx Benchmarking
This repo has a built-in benchmarking/profiling tool: `manage.py benchmark`, in
the `paperless_benchmark` Django app. It replaces ad hoc standalone scripts —
use it instead of writing new one-off seed/timing scripts.
## Command reference
```
manage.py benchmark seed --tier {home,medium,large} [--reset --yes-i-know-this-wipes-the-database] [--seed N]
manage.py benchmark run --repeat 5 [--label baseline]
manage.py benchmark profile <scenario_name> [--repeat 5] [--explain]
manage.py benchmark list-scenarios
```
- **`seed`** builds a realistic dataset at one of three scales: `home` (500
documents — fast, use this for iteration), `medium` (20,000 — the default
when `--tier` is omitted; a multi-minute seed), `large` (360,000 — matches
the scale reported in real large-install bug reports; slow, only use it
when a finding needs confirming at real scale). `--reset` wipes any
previously-seeded benchmark data first — but it is destructive and
irreversible: it deletes **all** users, **all** groups, and **all**
documents/tags/correspondents/document types/storage paths in the target
database, not just benchmark-created rows. Only run it against a disposable
benchmark database, never a real install. Because of that, `--reset` also
requires passing `--yes-i-know-this-wipes-the-database` in the same
invocation, or the command raises an error and does nothing. Omit both
flags if you want to layer more data onto an existing seed instead. `seed`
creates two named users, `perf_target` (mixed owned/shared documents,
realistic guardian permission grants) and `perf_admin` (superuser), plus a
general user/group pool with realistic permission-row ratios.
- **`run`** times the 3 built-in API endpoint benchmarks
(`/api/documents/`, `/api/documents/?page_size=50`, `/api/tags/?page_size=100000`) for both
`perf_target` and `perf_admin`, reporting min/median/max wall-clock and SQL
query count. Requires `seed` to have already run — it reuses that data, it
does not seed its own.
- **`profile`** times one named scenario from the registry (see
`list-scenarios`) via best-of-N repeat timing and SQL query count, against
`perf_target`. `--explain` additionally captures and prints the query plan:
real `EXPLAIN ANALYZE` execution stats on PostgreSQL/MariaDB, or
`EXPLAIN QUERY PLAN` (plan only, no real timing/row counts — clearly labeled
as such) on SQLite. Like `run`, `profile` requires `seed` to have already
run — it does not seed its own data either.
- Every `run`/`profile` invocation appends a JSON line to
`benchmark_results/history.jsonl` at the repo root (local-only, gitignored —
never commit this file). Use it to compare a `before`/`after` pair across
two invocations without hand-copying numbers.
- Full chain example: `seed --reset` once, then `run` and `profile` as many
times as you want against that same seeded data — no need to reseed between
them.
## Adding a new scenario
A "scenario" is a named, registered query/operation that `profile` can time
and explain. To add one, edit `src/paperless_benchmark/scenarios.py`: write a
`_<name>_run(user)` function (returns whatever `run_profile` should time) and
optionally a `_<name>_queryset(user)` function (returns the `QuerySet` for
`--explain` to analyze), then `register(Scenario(name=..., describe=...,
run=..., queryset_for_explain=...))` at module level. Both functions receive
the already-seeded `perf_target` user — they should not seed their own data.
## Branch workflow
`tools/benchmark-management-commands` is a **long-lived tooling branch**, not
a feature branch that gets merged and closed:
1. It is periodically brought up to date with `dev` (merge `dev` into it) so
the tooling doesn't drift from the schema/codebase it profiles. Do this
before starting a new investigation if it's been a while since the last
sync.
2. **Every performance investigation forks its own branch from
`tools/benchmark-management-commands`** (not from `dev`). Do the
investigation there: write throwaway profiling code, try fixes, capture
before/after numbers.
3. **That investigation branch never merges into `dev` or production.** Its
only job is to produce evidence and, optionally, a reusable scenario.
4. If the investigation turns up a scenario worth keeping permanently (see
"When to graduate a scenario" below), open a PR that adds **just that
scenario** back into `tools/benchmark-management-commands` — not the rest
of the investigation branch's throwaway code.
5. Any actual production fix the investigation motivates (e.g. an ORM query
change) goes into its own normal feature branch off `dev`, following the
project's regular contribution process — profiling evidence informs that
PR's description, but the profiling code itself does not travel with it.
## Reading query-plan output
- **PostgreSQL** `EXPLAIN ANALYZE`: look for `Seq Scan` on a large table
(missing index), a large gap between `rows=N` (planner's estimate) and the
actual row count in parentheses (stale statistics or a bad cardinality
estimate), and nested-loop joins driven by an outer relation with many
rows (usually the N+1 pattern this tool exists to catch).
- **MariaDB**: verified against a real MariaDB 12.3 container that MariaDB
does NOT accept MySQL 8.0.18+'s `EXPLAIN ANALYZE` syntax (it's a 1064
syntax error) -- `capture_explain()` instead runs MariaDB's own
`ANALYZE <statement>` form (no `EXPLAIN` keyword), which returns a
tabular plan with real per-row execution columns: `rows` (estimate) vs.
`r_rows` (actual), and `filtered` vs. `r_filtered`. A large gap between
`rows` and `r_rows`, or `type: ALL` (full table scan) on a large table,
are the signals to look for -- the same underlying concerns as Postgres's
`Seq Scan`/estimate-vs-actual gap, just in MariaDB's column-based output
instead of Postgres's nested-tree text format.
- **SQLite** `EXPLAIN QUERY PLAN`: no real timing/row-count data, only the
chosen access path (`SCAN` vs `SEARCH`, which index if any). Useful for
confirming an index is even being considered, not for judging real-world
cost — corroborate any SQLite finding against Postgres/MariaDB before
trusting it, since planner behavior differs meaningfully between them.
- Compare query **count**, not just timing, between before/after: a fix that
keeps the same wall-clock time but drops query count from O(n) to O(1) is
still a real, durable improvement — timing alone is noisy and
environment-dependent, query count is not.
## Cleaning up after an interrupted run
If a `seed`/`run`/`profile` invocation gets killed mid-run (Ctrl-C, `kill -9`,
a timed-out SSH session, etc.), check whether it left anything behind before
trusting the next benchmark's numbers. This was verified for real: a
`benchmark seed --tier large --reset ...` was started against both a fresh
PostgreSQL 18 container and a fresh MariaDB 12.3 container and `kill -9`'d a
few seconds into document seeding. In both cases, the database-side
connection disappeared immediately -- no stuck backend, no lingering query,
no held lock was observed in either backend once the killed process's PID
was confirmed gone. That said, this was one interruption point (mid
bulk-seed, between chunks); a run killed mid-query, or a driver/network
hiccup that doesn't cleanly close the socket, could behave differently, so
still check before trusting a number if any run in the session was
interrupted:
- **PostgreSQL**: look for leftover connections against the benchmark
database:
```sql
SELECT pid, state, query, query_start
FROM pg_stat_activity
WHERE datname = current_database() AND pid <> pg_backend_pid();
```
If a stuck backend shows up, clear it with:
```sql
SELECT pg_terminate_backend(pid)
FROM pg_stat_activity
WHERE datname = current_database() AND pid <> pg_backend_pid();
```
- **MariaDB**: look for leftover connections/queries:
```sql
SHOW FULL PROCESSLIST;
```
If a stuck connection shows up (anything other than your current admin
session), clear it with:
```sql
KILL <id>;
```
A stray connection left running concurrently with a subsequent benchmark run
would add real, contaminating load (extra queries competing for the same
rows, possibly held locks slowing the next run's timings) -- cheap to rule
out, expensive to silently trust a number that was actually measured
alongside a zombie connection.
## When to graduate a one-off finding into a permanent scenario
Register a scenario (rather than leaving it as throwaway code on the
investigation branch) when **both** are true:
- The query pattern is one this codebase is likely to regress on again (e.g.
it involves a permission-check join, a bulk operation, or anything else
with an easy-to-reintroduce N+1) — not a one-time fluke specific to this
investigation.
- Re-running it later, against a fresh seed, would still produce a
meaningful signal (it doesn't depend on investigation-specific throwaway
data or a fix that's already permanently landed and can't regress the same
way).
If a finding doesn't meet both bars, keep it as disposable code on the
investigation branch and let the branch's evidence (captured in the PR
description of whatever production fix it motivates) be the permanent
record instead.
+3
View File
@@ -115,3 +115,6 @@ celerybeat-schedule*
# Git worktree local folder
.worktrees
# Benchmark tooling output (local only, never committed)
/benchmark_results/
-38
View File
@@ -13,8 +13,6 @@ import { DocumentDetailComponent } from './components/document-detail/document-d
import { DocumentListComponent } from './components/document-list/document-list.component'
import { DocumentAttributesComponent } from './components/manage/document-attributes/document-attributes.component'
import { MailComponent } from './components/manage/mail/mail.component'
import { OcrTemplateEditorComponent } from './components/manage/ocr-templates/ocr-template-editor/ocr-template-editor.component'
import { OcrTemplatesComponent } from './components/manage/ocr-templates/ocr-templates.component'
import { SavedViewsComponent } from './components/manage/saved-views/saved-views.component'
import { WorkflowsComponent } from './components/manage/workflows/workflows.component'
import { NotFoundComponent } from './components/not-found/not-found.component'
@@ -276,42 +274,6 @@ export const routes: Routes = [
componentName: 'WorkflowsComponent',
},
},
{
path: 'ocr-templates',
component: OcrTemplatesComponent,
canActivate: [PermissionsGuard],
data: {
requiredPermission: {
action: PermissionAction.View,
type: PermissionType.OcrTemplate,
},
componentName: 'OcrTemplatesComponent',
},
},
{
path: 'ocr-templates/new',
component: OcrTemplateEditorComponent,
canActivate: [PermissionsGuard],
data: {
requiredPermission: {
action: PermissionAction.Add,
type: PermissionType.OcrTemplate,
},
componentName: 'OcrTemplateEditorComponent',
},
},
{
path: 'ocr-templates/:id',
component: OcrTemplateEditorComponent,
canActivate: [PermissionsGuard],
data: {
requiredPermission: {
action: PermissionAction.Change,
type: PermissionType.OcrTemplate,
},
componentName: 'OcrTemplateEditorComponent',
},
},
{
path: 'mail',
component: MailComponent,
@@ -253,14 +253,6 @@
<i-bs class="me-2" name="boxes"></i-bs><span class="nav-link-label"><ng-container i18n>Workflows</ng-container></span>
</a>
</li>
<li class="nav-item app-link"
*pngxIfPermissions="{ action: PermissionAction.View, type: PermissionType.OcrTemplate }">
<a class="nav-link" routerLink="ocr-templates" routerLinkActive="active" (click)="closeMenu()"
ngbPopover="OCR Templates" i18n-ngbPopover [disablePopover]="!slimSidebarEnabled" placement="end"
container="body" triggers="mouseenter:mouseleave" popoverClass="popover-slim">
<i-bs class="me-2" name="file-earmark-break"></i-bs><span><ng-container i18n>OCR Templates</ng-container></span>
</a>
</li>
<li class="nav-item app-link" *pngxIfPermissions="{ action: PermissionAction.View, type: PermissionType.MailAccount }"
tourAnchor="tour.mail">
<a class="nav-link" routerLink="mail" routerLinkActive="active" (click)="closeMenu()" ngbPopover="Mail"
@@ -82,23 +82,6 @@
<i-bs name="pencil" class="me-1"></i-bs><ng-container i18n>PDF Editor</ng-container>
</button>
<button
ngbDropdownItem
(click)="runZoneOcr()"
[disabled]="!userCanEdit || !document?.document_type"
*pngxIfPermissions="{ action: PermissionAction.Change, type: PermissionType.OcrTemplate }"
>
<i-bs width="1em" height="1em" name="file-earmark-ruled" class="me-1"></i-bs><span i18n>Run Zone OCR</span>
</button>
<button
ngbDropdownItem
(click)="createOcrTemplate()"
*pngxIfPermissions="{ action: PermissionAction.Add, type: PermissionType.OcrTemplate }"
>
<i-bs width="1em" height="1em" name="file-earmark-medical" class="me-1"></i-bs><span i18n>Create OCR Template</span>
</button>
@if (userIsOwner && (requiresPassword || password)) {
<button ngbDropdownItem (click)="removePassword()" [disabled]="!password">
<i-bs name="unlock" class="me-1"></i-bs><ng-container i18n>Remove Password</ng-container>
@@ -1449,48 +1449,6 @@ export class DocumentDetailComponent
})
}
runZoneOcr() {
this.documentsService.runZoneOcr(this.document.id).subscribe({
next: (res) => {
const results = res.results ?? []
if (results.length) {
const failed = results.filter(
(r) =>
r.value === null ||
r.value === undefined ||
`${r.value}`.trim() === ''
)
const filled = results.length - failed.length
let msg = $localize`Filled ${filled} of ${results.length} fields`
if (failed.length) {
const names = failed.map((r) => r.zone).join(', ')
msg = `${msg}. ${$localize`Failed to match zones: ${names}`}`
}
this.toastService.showInfo(msg)
} else {
this.toastService.showInfo(
$localize`Zone OCR ran but no results extracted.`
)
}
this.documentsService
.get(this.documentId)
.subscribe((doc) => this.updateComponent(doc))
},
error: (error) => {
this.toastService.showError($localize`Zone OCR failed`, error)
},
})
}
createOcrTemplate() {
this.router.navigate(['/ocr-templates', 'new'], {
queryParams: {
document_type: this.document.document_type,
sample_document: this.document.id,
},
})
}
private getSelectedNonLatestVersionId(): number | null {
const versions = this.document()?.versions ?? []
if (!versions.length || !this.selectedVersionId()) {
@@ -98,9 +98,6 @@
<button ngbDropdownItem (click)="mergeSelectedAsVersions()" [disabled]="!userOwnsAll || !userCanEditAll || !userCanDelete || list.allSelected || list.selectedCount < 2">
<i-bs name="journal-bookmark-fill" class="me-1"></i-bs><ng-container i18n>Merge as versions</ng-container>
</button>
<button ngbDropdownItem (click)="runZoneOcrSelected()" [disabled]="!userCanEditAll || list.allSelected">
<i-bs name="file-earmark-ruled" class="me-1"></i-bs><ng-container i18n>Run Zone OCR</ng-container>
</button>
</div>
</div>
</div>
@@ -19,15 +19,7 @@ import {
} from '@ng-bootstrap/ng-bootstrap'
import { saveAs } from 'file-saver'
import { NgxBootstrapIconsModule } from 'ngx-bootstrap-icons'
import {
first,
forkJoin,
map,
Observable,
Subject,
switchMap,
takeUntil,
} from 'rxjs'
import { first, map, Observable, Subject, switchMap, takeUntil } from 'rxjs'
import { ConfirmDialogComponent } from 'src/app/components/common/confirm-dialog/confirm-dialog.component'
import { CustomField } from 'src/app/data/custom-field'
import { MatchingModel } from 'src/app/data/matching-model'
@@ -947,27 +939,6 @@ export class BulkEditorComponent
})
}
runZoneOcrSelected() {
const ids = Array.from(this.list.selected)
if (!ids.length) return
const modal = this.modalService.open(ConfirmDialogComponent, {
backdrop: 'static',
})
modal.componentInstance.title = $localize`Run Zone OCR`
modal.componentInstance.messageBold = $localize`Run zone OCR on ${this.getSelectionSize()} selected document(s)?`
modal.componentInstance.message = $localize`Each document's type template (if it has one) is applied, overwriting the mapped fields.`
modal.componentInstance.btnCaption = $localize`Proceed`
modal.componentInstance.confirmClicked
.pipe(takeUntil(this.unsubscribeNotifier))
.subscribe(() => {
modal.componentInstance.buttonsEnabled = false
this.executeDocumentAction(
modal,
forkJoin(ids.map((id) => this.documentService.runZoneOcr(id)))
)
})
}
setPermissions() {
let modal = this.modalService.open(PermissionsDialogComponent, {
backdrop: 'static',
@@ -1,34 +0,0 @@
@if (zones.length === 0) {
<p class="text-muted" i18n>
No zones defined. Load a document preview and draw rectangles to add zones.
</p>
}
<div class="list-group">
@for (zone of zones; track $index; let i = $index) {
<div
class="list-group-item list-group-item-action d-flex justify-content-between align-items-center"
[style.box-shadow]="selectedZoneIndex === i ? 'inset 3px 0 0 0 var(--bs-primary)' : null"
>
<div class="flex-grow-1" role="button" style="cursor: pointer;" (click)="zoneSelected.emit(i)">
<div>
<strong [class.text-primary]="selectedZoneIndex === i">
{{ zone.name }}
</strong>
</div>
<div class="small text-muted">
{{ getZoneTargetName(zone) }} - {{ zone.width }}x{{ zone.height }}px
<ng-container i18n>p.</ng-container>{{ zonePage(zone) }}
</div>
</div>
<div class="btn-group">
<button class="btn btn-sm btn-outline-secondary" type="button" (click)="zoneSelected.emit(i)" title="Edit" i18n-title>
<i-bs name="pencil"></i-bs>
</button>
<button class="btn btn-sm btn-outline-danger" type="button" (click)="zoneRemoved.emit(i)" title="Delete" i18n-title>
<i-bs name="trash"></i-bs>
</button>
</div>
</div>
}
</div>
@@ -1,72 +0,0 @@
import { ComponentFixture, TestBed } from '@angular/core/testing'
import { NgxBootstrapIconsModule, allIcons } from 'ngx-bootstrap-icons'
import { CustomField } from 'src/app/data/custom-field'
import { OcrTemplateZone } from 'src/app/data/ocr-template'
import { OcrTemplateEditorZoneListComponent } from './ocr-template-editor-zone-list.component'
function zone(overrides: Partial<OcrTemplateZone> = {}): OcrTemplateZone {
return {
name: 'Zone 1',
target: 'custom_field',
custom_field: 7,
x: 10,
y: 20,
width: 30,
height: 40,
page: 1,
ocr_language: 'eng',
transform: 'strip',
validation_regex: '',
order: 0,
...overrides,
}
}
describe('OcrTemplateEditorZoneListComponent', () => {
let fixture: ComponentFixture<OcrTemplateEditorZoneListComponent>
let component: OcrTemplateEditorZoneListComponent
beforeEach(async () => {
await TestBed.configureTestingModule({
imports: [
OcrTemplateEditorZoneListComponent,
NgxBootstrapIconsModule.pick(allIcons),
],
}).compileComponents()
fixture = TestBed.createComponent(OcrTemplateEditorZoneListComponent)
component = fixture.componentInstance
})
it('shows empty state when no zones are defined', () => {
fixture.detectChanges()
expect(fixture.nativeElement.textContent).toContain('No zones defined')
})
it('renders zone target, size, and page', () => {
component.zones = [zone()]
component.customFields = [{ id: 7, name: 'Invoice Number' } as CustomField]
fixture.detectChanges()
const text = fixture.nativeElement.textContent
expect(text).toContain('Zone 1')
expect(text).toContain('Invoice Number')
expect(text).toContain('30x40px')
expect(text).toContain('p.1')
})
it('emits select and remove events', () => {
component.zones = [zone()]
const selectSpy = jest.spyOn(component.zoneSelected, 'emit')
const removeSpy = jest.spyOn(component.zoneRemoved, 'emit')
fixture.detectChanges()
const buttons = fixture.nativeElement.querySelectorAll('button')
buttons[0].click()
buttons[1].click()
expect(selectSpy).toHaveBeenCalledWith(0)
expect(removeSpy).toHaveBeenCalledWith(0)
})
})
@@ -1,41 +0,0 @@
import { Component, EventEmitter, Input, Output } from '@angular/core'
import { NgxBootstrapIconsModule } from 'ngx-bootstrap-icons'
import { CustomField } from 'src/app/data/custom-field'
import { OCR_BUILTIN_TARGETS, OcrTemplateZone } from 'src/app/data/ocr-template'
import { getZonePage } from '../zone-geometry'
@Component({
selector: 'pngx-ocr-template-zone-list',
imports: [NgxBootstrapIconsModule],
templateUrl: './ocr-template-editor-zone-list.component.html',
})
export class OcrTemplateEditorZoneListComponent {
@Input() zones: OcrTemplateZone[] = []
@Input() selectedZoneIndex: number | null = null
@Input() previewPage = 0
@Input() previewPageCount: number | null = null
@Input() customFields: CustomField[] = []
@Output() zoneSelected = new EventEmitter<number>()
@Output() zoneRemoved = new EventEmitter<number>()
zonePage(zone: OcrTemplateZone): number {
return getZonePage(zone, this.previewPage, this.previewPageCount)
}
getZoneTargetName(zone: OcrTemplateZone): string {
const target = zone.target || 'custom_field'
if (target === 'custom_field') {
return zone.custom_field
? this.getCustomFieldName(zone.custom_field)
: $localize`(no field)`
}
return OCR_BUILTIN_TARGETS.find((t) => t.id === target)?.name ?? target
}
private getCustomFieldName(id: number): string {
return (
this.customFields.find((field) => field.id === id)?.name ?? `Field #${id}`
)
}
}
@@ -1,442 +0,0 @@
<pngx-page-header [title]="pageTitle" [id]="template.id">
<div class="input-group input-group-sm me-5 align-items-center">
<div class="input-group-text">
<i-bs name="file-text"></i-bs>
</div>
<input
type="text"
class="form-control"
[(ngModel)]="previewDocModel"
[ngbTypeahead]="searchDocuments"
[inputFormatter]="documentFormatter"
[resultFormatter]="documentFormatter"
(selectItem)="onPreviewDocSelected($event)"
[editable]="false"
placeholder="Search documents by title..."
i18n-placeholder
/>
</div>
<div class="d-flex align-items-center flex-wrap gap-2">
<div class="input-group input-group-sm ms-2 d-none d-md-flex">
<div class="input-group-text" i18n>Page</div>
<input class="form-control flex-grow-0 w-auto" type="number" min="1" [max]="previewPageCount" [(ngModel)]="previewPageDisplay" />
<div class="input-group-text" i18n>of {{previewPageCount}}</div>
</div>
<button type="button" class="btn btn-sm btn-outline-secondary" i18n-title title="Previous" (click)="prevPage()" [disabled]="!pageImageUrl || previewPage <= 0">
<i-bs width="1.2em" height="1.2em" name="arrow-left"></i-bs>
</button>
<button type="button" class="btn btn-sm btn-outline-secondary" i18n-title title="Next" (click)="nextPage()" [disabled]="!pageImageUrl || previewPage >= (previewPageCount ?? 1) - 1">
<i-bs width="1.2em" height="1.2em" name="arrow-right"></i-bs>
</button>
<div class="input-group input-group-sm">
<button class="btn btn-outline-secondary" (click)="zoomOut()" i18n>-</button>
<span class="input-group-text">{{ zoom * 100 | number: '1.0-0' }}%</span>
<button class="btn btn-outline-secondary" (click)="zoomIn()" i18n>+</button>
</div>
</div>
</pngx-page-header>
<div class="row">
<div class="col-md-4">
<div class="btn-toolbar mb-1 border-bottom">
<div class="btn-group pb-3">
<a routerLink="/ocr-templates" class="btn btn-sm btn-outline-secondary">
<i-bs width="1.2em" height="1.2em" name="x"></i-bs>
<span class="ms-1" i18n>Close</span>
</a>
</div>
<div class="btn-group ms-auto pb-3">
<button class="btn btn-sm btn-primary" (click)="save()" [disabled]="saving">
@if (saving) {
<span class="spinner-border spinner-border-sm me-1"></span>
}
<span i18n>Save</span>
</button>
</div>
</div>
<ul ngbNav #nav="ngbNav" [(activeId)]="activeTab" class="nav-underline flex-nowrap flex-md-wrap overflow-auto">
<li ngbNavItem="settings">
<a ngbNavLink i18n>Settings</a>
<ng-template ngbNavContent>
<div class="row mb-3">
<div class="col-9">
<pngx-input-text [(ngModel)]="template.name" title="Template name" i18n-title></pngx-input-text>
</div>
<div class="col-3">
<pngx-input-switch [(ngModel)]="template.enabled" title="Enabled" i18n-title></pngx-input-switch>
</div>
</div>
<pngx-input-select [(ngModel)]="template.document_type" [items]="documentTypes" bindLabel="name" bindValue="id" title="Document type" i18n-title></pngx-input-select>
<small class="text-muted" i18n>
Draw rectangles on the preview to define extraction zones. Use the
page controls above the preview to add zones on different pages.
</small>
</ng-template>
</li>
<li ngbNavItem="zones">
<a ngbNavLink><ng-container i18n>Zones</ng-container> <span class="badge bg-primary ms-2">{{ template.zones.length }}</span></a>
<ng-template ngbNavContent>
<pngx-ocr-template-zone-list
[zones]="template.zones"
[selectedZoneIndex]="selectedZoneIndex"
[previewPage]="previewPage"
[previewPageCount]="previewPageCount"
[customFields]="customFields"
(zoneSelected)="selectZone($event)"
(zoneRemoved)="removeZone($event)"
></pngx-ocr-template-zone-list>
</ng-template>
</li>
<li ngbNavItem="zone">
<a ngbNavLink i18n>Zone</a>
<ng-template ngbNavContent>
@if (selectedZone; as zone) {
<div class="d-flex justify-content-between align-items-center mb-3">
<strong>{{ zone.name }}</strong>
<div class="d-flex gap-2">
<button class="btn btn-sm btn-primary" (click)="save()" [disabled]="saving">
@if (saving) {
<span class="spinner-border spinner-border-sm me-1"></span>
}
<span i18n>Save</span>
</button>
<button class="btn btn-sm btn-outline-danger" (click)="deleteSelectedZone()">
<i-bs name="trash" class="me-1"></i-bs><ng-container i18n>Delete zone</ng-container>
</button>
</div>
</div>
<div class="mb-3">
<label class="form-label" i18n>Zone Name</label>
<input
type="text"
class="form-control"
[(ngModel)]="zone.name"
/>
</div>
<div class="mb-3">
<label class="form-label" i18n>Page</label>
<input
type="number"
class="form-control"
[(ngModel)]="zone.page"
min="-1"
/>
<small class="text-muted" i18n>Page this zone is on. Use -1 for the last page. Set automatically when you draw it.</small>
</div>
<div class="mb-3">
<label class="form-label" i18n>Field</label>
<div class="input-group">
<select class="form-select" [ngModel]="zoneFieldValue(zone)" (ngModelChange)="setZoneField(zone, $event)">
<optgroup label="Built-in fields" i18n-label>
@for (t of builtinTargets; track t.id) {
<option [ngValue]="t.id">{{ t.name }}</option>
}
</optgroup>
<optgroup label="Custom fields" i18n-label>
@for (cf of customFields; track cf.id) {
<option [ngValue]="cf.id">{{ cf.name }} ({{ cf.data_type }})</option>
}
</optgroup>
</select>
<button
class="btn btn-outline-secondary"
type="button"
(click)="openQuickCreate(selectedZoneIndex)"
title="Create new custom field"
i18n-title
>
<i-bs name="plus"></i-bs>
</button>
</div>
<small class="text-muted" i18n>Write the extracted value to a custom field, or to a built-in field (Title, ASN, Date created).</small>
</div>
@if (isFieldShared(zone)) {
<div class="card mb-3 border-info">
<div class="card-body">
<h6 class="card-title d-flex align-items-center gap-2">
<i-bs name="braces"></i-bs>
<span i18n>Combine zones into this field</span>
</h6>
<p class="small text-muted mb-2" i18n>
More than one zone writes to this field. Build the combined
value below: click a zone to insert its token, and type any
separators or literal text between tokens.
</p>
<div class="d-flex flex-wrap gap-1 mb-2">
@for (z of zonesForField(zone); track $index) {
<button
type="button"
class="btn btn-sm btn-outline-info"
(click)="insertCombineToken(zone, z)"
title="Insert token"
i18n-title
>
+ {{ z.name || 'Zone' }}
</button>
}
</div>
<input
type="text"
class="form-control font-monospace"
[ngModel]="getCombineFormat(zone)"
(ngModelChange)="setCombineFormat(zone, $event)"
placeholder="{Zone 1} - {Zone 2}"
/>
<small class="text-muted" i18n>
Tokens are matched by zone name. An empty zone leaves its
token blank and the stray separator is trimmed. Leave empty
to just join the zones in order with a space.
</small>
</div>
</div>
}
@if (showQuickCreate) {
<div class="card mb-3 border-primary">
<div class="card-body">
<h6 class="card-title" i18n>Create Custom Field</h6>
<div class="mb-2">
<label class="form-label small" i18n>Field Name</label>
<input type="text" class="form-control form-control-sm"
[(ngModel)]="quickCreateName" placeholder="e.g. Invoice Number" />
</div>
<div class="mb-2">
<label class="form-label small" i18n>Field Type</label>
<select class="form-select form-select-sm" [(ngModel)]="quickCreateType">
@for (t of quickCreateTypes; track t.id) {
<option [ngValue]="t.id">{{ t.name }}</option>
}
</select>
</div>
<div class="d-flex gap-2">
<button class="btn btn-primary btn-sm" (click)="submitQuickCreate()"
[disabled]="!quickCreateName.trim()" i18n>
Create & Assign
</button>
<button class="btn btn-outline-secondary btn-sm" (click)="cancelQuickCreate()" i18n>
Cancel
</button>
</div>
</div>
</div>
}
<div class="mb-3">
<label class="form-label" i18n>OCR Language</label>
<ng-select
[items]="ocrLanguageOptions"
bindLabel="name"
bindValue="id"
[multiple]="true"
[closeOnSelect]="false"
[ngModel]="ocrLanguageArray(zone)"
(ngModelChange)="setOcrLanguages(zone, $event)"
placeholder="Select languages"
i18n-placeholder
></ng-select>
</div>
<div class="mb-3">
<label class="form-label" i18n>Transform</label>
<select class="form-select" [(ngModel)]="zone.transform">
@for (opt of transformOptions; track opt.id) {
<option [ngValue]="opt.id">{{ opt.name }}</option>
}
</select>
</div>
@if (zone.transform === dateTransform) {
<div class="mb-3">
<label class="form-label" i18n>Date format</label>
<select class="form-select" [ngModel]="dateFormatChoice(zone)" (ngModelChange)="setDateFormatChoice(zone, $event)">
@for (opt of dateFormatOptions; track opt.id) {
<option [ngValue]="opt.id">{{ opt.name }}</option>
}
<option [ngValue]="customDateFormatChoice" i18n>Custom...</option>
</select>
@if (usesCustomDateFormat(zone)) {
<div class="input-group mt-2">
<input type="text" class="form-control font-monospace" [(ngModel)]="zone.date_format" placeholder="%d.%m.%Y" />
<button class="btn btn-outline-secondary" type="button" [ngbPopover]="dateFmtHelp" [autoClose]="true" title="Date format help" i18n-title>
<i-bs name="question-circle"></i-bs>
</button>
</div>
<ng-template #dateFmtHelp>
<p class="mb-1" i18n>Python date codes:</p>
<ul class="mb-1 ps-3">
<li><code>%d</code> <ng-container i18n>day (01-31)</ng-container></li>
<li><code>%m</code> <ng-container i18n>month (01-12)</ng-container></li>
<li><code>%Y</code> <ng-container i18n>year, 4-digit</ng-container></li>
<li><code>%y</code> <ng-container i18n>year, 2-digit</ng-container></li>
<li><code>%b</code> <ng-container i18n>month name (Jan)</ng-container></li>
</ul>
<span i18n>Example:</span> <code>%d.%m.%Y</code> -> 03.03.2026
</ng-template>
}
</div>
}
<div class="mb-3">
<label class="form-label" i18n>Validation Regex</label>
<input
type="text"
class="form-control font-monospace"
[(ngModel)]="zone.validation_regex"
placeholder="e.g. \d{2}\.\d{2}\.\d{4}"
>
</div>
<div class="text-muted small">
{{ zone.x }}, {{ zone.y }} - {{ zone.width }}x{{ zone.height }}px
</div>
<hr class="my-3" />
<h6 i18n>Test</h6>
@if (!previewDocId) {
<p class="text-muted small mb-0" i18n>
Load a document in the Settings tab to test this zone.
</p>
} @else {
<button class="btn btn-sm btn-outline-secondary" (click)="testZone()" [disabled]="zoneTesting">
@if (zoneTesting) {
<span class="spinner-border spinner-border-sm me-1"></span>
}
<span i18n>Test this zone</span>
</button>
@if (zoneTestResult) {
@if (zoneTestResult.error) {
<div class="alert alert-warning py-2 mt-2 mb-0 small">{{ zoneTestResult.error }}</div>
} @else {
<dl class="row small mt-2 mb-0">
<dt class="col-sm-4" i18n>OCR text</dt>
<dd class="col-sm-8"><code>{{ zoneTestResult.raw_text || '(nothing detected)' }}</code></dd>
<dt class="col-sm-4" i18n>Value</dt>
<dd class="col-sm-8"><code>{{ zoneTestResult.value || '(empty)' }}</code></dd>
@if (zoneTestResult.regex) {
<dt class="col-sm-4" i18n>Validation</dt>
<dd class="col-sm-8">
@if (zoneTestResult.regex_match) {
<span class="badge bg-success" i18n>Regex matches</span>
} @else {
<span class="badge bg-danger" i18n>Regex does not match</span>
}
</dd>
}
</dl>
}
}
}
} @else {
<p class="text-muted" i18n>
Select a zone from the Zones tab, or draw a rectangle on the document to create one.
</p>
}
</ng-template>
</li>
</ul>
<div [ngbNavOutlet]="nav" class="mt-3"></div>
</div>
<!-- Right column: Document preview with zone overlay -->
<div class="col-md-8">
@if (pageImageUrl) {
<div class="zone-preview-scroll border">
<div class="zone-preview-stage" [style.width.%]="zoom * 100">
<img
#pageImage
[src]="pageImageUrl"
(load)="onImageLoad()"
class="zone-preview-image"
[style.visibility]="imageLoaded ? 'visible' : 'hidden'"
crossorigin="use-credentials"
/>
@if (imageLoaded) {
<svg
#zoneOverlay
class="zone-overlay"
[attr.viewBox]="overlayViewBox()"
preserveAspectRatio="none"
[style.cursor]="overlayCursor"
(mousedown)="onOverlayMouseDown($event)"
(mousemove)="onOverlayMouseMove($event)"
(mouseup)="onOverlayMouseUp($event)"
>
@for (zone of template.zones; track $index; let i = $index) {
@if (zoneDisplayRect(i); as rect) {
<g>
<rect
class="zone-rect"
[class.zone-rect-selected]="selectedZoneIndex === i"
[attr.x]="rect.x"
[attr.y]="rect.y"
[attr.width]="rect.w"
[attr.height]="rect.h"
[attr.stroke]="zoneColor(i)"
[attr.fill]="zoneFill(i)"
></rect>
<text
class="zone-label"
[attr.x]="rect.x + overlayUnitSize(6)"
[attr.y]="zoneLabelY(rect)"
[attr.font-size]="overlayFontSize()"
[attr.fill]="zoneColor(i)"
>{{ zoneLabel(zone, i) }}</text>
@if (selectedZoneIndex === i) {
@for (handle of resizeHandles(rect); track handle.handle) {
<rect
class="zone-resize-handle"
[attr.x]="handle.x - overlayHandleSize() / 2"
[attr.y]="handle.y - overlayHandleSize() / 2"
[attr.width]="overlayHandleSize()"
[attr.height]="overlayHandleSize()"
[attr.fill]="zoneColor(i)"
></rect>
}
}
</g>
}
}
@if (drawingRect(); as rect) {
<rect
class="zone-drawing-rect"
[attr.x]="rect.x"
[attr.y]="rect.y"
[attr.width]="rect.w"
[attr.height]="rect.h"
></rect>
}
</svg>
}
@if (!imageLoaded) {
<div class="d-flex justify-content-center p-5">
<div class="spinner-border" role="status">
<span class="visually-hidden" i18n>Loading page...</span>
</div>
</div>
}
</div>
</div>
} @else {
<div class="border rounded p-5 text-center text-muted">
<i-bs name="file-earmark-image" width="48" height="48"></i-bs>
<p class="mt-3" i18n>
Enter a document ID and click "Load" to preview a page and draw extraction zones.
</p>
</div>
}
</div>
</div>
@@ -1,63 +0,0 @@
:host {
display: block;
}
.zone-preview-scroll {
max-height: 78vh;
overflow: auto;
}
.zone-preview-stage {
display: inline-block;
position: relative;
}
.zone-preview-image {
display: block;
width: 100%;
}
.zone-overlay {
height: 100%;
inset: 0;
position: absolute;
touch-action: none;
width: 100%;
}
.zone-rect,
.zone-drawing-rect {
vector-effect: non-scaling-stroke;
}
.zone-rect {
stroke-width: 2;
}
.zone-rect-selected {
stroke-width: 3;
}
.zone-label {
font-family: var(--bs-font-sans-serif);
font-weight: 600;
paint-order: stroke;
pointer-events: none;
stroke: #fff;
stroke-linejoin: round;
stroke-width: 4px;
vector-effect: non-scaling-stroke;
}
.zone-resize-handle {
stroke: #fff;
stroke-width: 1;
vector-effect: non-scaling-stroke;
}
.zone-drawing-rect {
fill: rgba(105, 219, 124, 0.25);
stroke: #69db7c;
stroke-dasharray: 5 5;
stroke-width: 2;
}
@@ -1,962 +0,0 @@
import { CommonModule } from '@angular/common'
import {
Component,
ElementRef,
HostListener,
inject,
OnDestroy,
OnInit,
ViewChild,
} from '@angular/core'
import { FormsModule } from '@angular/forms'
import { ActivatedRoute, Router, RouterModule } from '@angular/router'
import {
NgbNavModule,
NgbPopoverModule,
NgbTypeaheadModule,
NgbTypeaheadSelectItemEvent,
} from '@ng-bootstrap/ng-bootstrap'
import { NgSelectModule } from '@ng-select/ng-select'
import { NgxBootstrapIconsModule } from 'ngx-bootstrap-icons'
import {
catchError,
debounceTime,
distinctUntilChanged,
map,
Observable,
of,
Subject,
switchMap,
takeUntil,
} from 'rxjs'
import { SelectComponent } from 'src/app/components/common/input/select/select.component'
import { SwitchComponent } from 'src/app/components/common/input/switch/switch.component'
import { TextComponent } from 'src/app/components/common/input/text/text.component'
import { PageHeaderComponent } from 'src/app/components/common/page-header/page-header.component'
import { CustomField, CustomFieldDataType } from 'src/app/data/custom-field'
import { Document } from 'src/app/data/document'
import { DocumentType } from 'src/app/data/document-type'
import {
DATE_FORMAT_OPTIONS,
DEFAULT_OCR_ZONE_LANGUAGE,
DEFAULT_OCR_ZONE_TARGET,
DEFAULT_OCR_ZONE_TRANSFORM,
isOcrBuiltinTarget,
OCR_BUILTIN_TARGETS,
OCR_LANGUAGE_OPTIONS,
OCR_ZONE_TARGET,
OCR_ZONE_TRANSFORM,
OcrBuiltinTarget,
OcrTemplate,
OcrTemplateZone,
OcrZoneTestResult,
TRANSFORM_OPTIONS,
ZoneTestRequest,
} from 'src/app/data/ocr-template'
import { CorrespondentService } from 'src/app/services/rest/correspondent.service'
import { CustomFieldsService } from 'src/app/services/rest/custom-fields.service'
import { DocumentTypeService } from 'src/app/services/rest/document-type.service'
import { DocumentService } from 'src/app/services/rest/document.service'
import { OcrTemplateService } from 'src/app/services/rest/ocr-template.service'
import { ToastService } from 'src/app/services/toast.service'
import { OcrTemplateEditorZoneListComponent } from './ocr-template-editor-zone-list/ocr-template-editor-zone-list.component'
import {
DisplayRect,
DrawingRect,
findHandleAt,
findZoneAt,
getZoneDisplayRect,
getZonePage,
HANDLE_SIZE,
isZoneOnPage,
MoveStart,
moveZone,
Point,
ResizeHandle,
resizeZone,
} from './zone-geometry'
type ActiveTab = 'settings' | 'zones' | 'zone'
type ZoneFieldSelection = OcrBuiltinTarget | number | null
type OverlayInteraction =
| { kind: 'idle' }
| { kind: 'drawing'; rect: DrawingRect }
| { kind: 'moving'; zoneIndex: number; start: MoveStart }
| { kind: 'resizing'; zoneIndex: number; handle: ResizeHandle }
interface ResizeHandleMarker extends Point {
handle: ResizeHandle
}
const CUSTOM_DATE_FORMAT_CHOICE = 'custom'
const MIN_DRAWN_ZONE_SIZE = 10
const NO_OVERLAY_INTERACTION: OverlayInteraction = { kind: 'idle' }
const ZONE_COLORS = [
'#4f8ff7',
'#ff6b6b',
'#51cf66',
'#ffd43b',
'#cc5de8',
'#ff922b',
'#20c997',
'#e599f7',
]
const RESIZE_CURSOR: Record<ResizeHandle, string> = {
nw: 'nw-resize',
ne: 'ne-resize',
sw: 'sw-resize',
se: 'se-resize',
n: 'n-resize',
s: 's-resize',
w: 'w-resize',
e: 'e-resize',
}
@Component({
selector: 'pngx-ocr-template-editor',
standalone: true,
imports: [
PageHeaderComponent,
TextComponent,
SelectComponent,
SwitchComponent,
CommonModule,
FormsModule,
RouterModule,
NgbNavModule,
NgbPopoverModule,
NgbTypeaheadModule,
NgSelectModule,
NgxBootstrapIconsModule,
OcrTemplateEditorZoneListComponent,
],
templateUrl: './ocr-template-editor.component.html',
styleUrls: ['./ocr-template-editor.component.scss'],
})
export class OcrTemplateEditorComponent implements OnInit, OnDestroy {
private readonly route = inject(ActivatedRoute)
private readonly router = inject(Router)
private readonly templateService = inject(OcrTemplateService)
private readonly customFieldsService = inject(CustomFieldsService)
private readonly documentTypeService = inject(DocumentTypeService)
private readonly correspondentService = inject(CorrespondentService)
private readonly documentService = inject(DocumentService)
private readonly toastService = inject(ToastService)
private readonly destroy$ = new Subject<void>()
private readonly customDateFormatZones = new WeakSet<OcrTemplateZone>()
@ViewChild('zoneOverlay') overlayRef: ElementRef<SVGSVGElement>
@ViewChild('pageImage') imageRef: ElementRef<HTMLImageElement>
template: OcrTemplate = {
id: null,
name: '',
document_type: null,
sample_document: null,
source_width: 0,
source_height: 0,
enabled: true,
combine_formats: {},
zones: [],
}
customFields: CustomField[] = []
documentTypes: DocumentType[] = []
transformOptions = TRANSFORM_OPTIONS
builtinTargets = OCR_BUILTIN_TARGETS
dateFormatOptions = DATE_FORMAT_OPTIONS
ocrLanguageOptions = OCR_LANGUAGE_OPTIONS
dateTransform = OCR_ZONE_TRANSFORM.Date
customDateFormatChoice = CUSTOM_DATE_FORMAT_CHOICE
isNew = true
saving = false
previewDocId: number | null = null
previewPage = 0
previewPageCount: number | null = null
private pageCountForDoc: number | null = null
pageImageUrl: string | null = null
imageLoaded = false
zoom = 1
previewDocModel: Document | string = ''
private correspondentNames = new Map<number, string>()
public get previewPageDisplay(): number {
return this.previewPage + 1
}
public set previewPageDisplay(value: number) {
this.goToPage(value - 1)
}
activeTab: ActiveTab = 'settings'
selectedZoneIndex: number | null = null
private overlayInteraction: OverlayInteraction = NO_OVERLAY_INTERACTION
overlayCursor = 'crosshair'
zoneTestResult: OcrZoneTestResult | null = null
zoneTesting = false
showQuickCreate = false
quickCreateName = ''
quickCreateType = CustomFieldDataType.String
quickCreateForZoneIndex: number | null = null
quickCreateTypes = [
{ id: CustomFieldDataType.String, name: $localize`String` },
{ id: CustomFieldDataType.Integer, name: $localize`Integer` },
{ id: CustomFieldDataType.Float, name: $localize`Float` },
{ id: CustomFieldDataType.Date, name: $localize`Date` },
{ id: CustomFieldDataType.Monetary, name: $localize`Monetary` },
{ id: CustomFieldDataType.Boolean, name: $localize`Boolean` },
{ id: CustomFieldDataType.Url, name: $localize`URL` },
{ id: CustomFieldDataType.LongText, name: $localize`Long Text` },
]
get selectedZone(): OcrTemplateZone | null {
return this.selectedZoneIndex !== null
? (this.template.zones[this.selectedZoneIndex] ?? null)
: null
}
get pageTitle(): string {
return this.isNew
? $localize`New OCR Template`
: $localize`Edit OCR Template`
}
ngOnInit() {
this.customFieldsService
.listAll()
.pipe(takeUntil(this.destroy$))
.subscribe((r) => (this.customFields = r.results))
this.documentTypeService
.listAll()
.pipe(takeUntil(this.destroy$))
.subscribe((r) => (this.documentTypes = r.results))
this.correspondentService
.listAll()
.pipe(takeUntil(this.destroy$))
.subscribe((r) => {
this.correspondentNames = new Map(r.results.map((c) => [c.id, c.name]))
})
const id = this.route.snapshot.paramMap.get('id')
if (id && id !== 'new') {
this.isNew = false
this.templateService
.get(parseInt(id))
.pipe(takeUntil(this.destroy$))
.subscribe((t) => {
this.template = t
this.template.combine_formats ??= {}
if (t.sample_document) {
this.previewDocId = t.sample_document
this.loadPreview()
}
})
} else {
const qp = this.route.snapshot.queryParams
if (qp['document_type']) {
this.template.document_type = parseInt(qp['document_type'])
}
if (qp['sample_document']) {
const docId = parseInt(qp['sample_document'])
this.template.sample_document = docId
this.previewDocId = docId
this.loadPreview()
}
}
}
searchDocuments = (text$: Observable<string>): Observable<Document[]> =>
text$.pipe(
debounceTime(250),
distinctUntilChanged(),
switchMap((term) => {
if (!term || term.trim().length < 2) return of([])
const params: { title__icontains: string; document_type__id?: number } =
{ title__icontains: term.trim() }
if (this.template.document_type) {
params['document_type__id'] = this.template.document_type
}
return this.documentService.list(1, 10, 'created', true, params).pipe(
map((r) => r.results),
catchError(() => of([]))
)
})
)
documentFormatter = (doc: Document | string): string => {
if (typeof doc === 'string') return doc
const corr = doc.correspondent
? this.correspondentNames.get(doc.correspondent)
: null
return corr
? `#${doc.id} ${doc.title} (${corr})`
: `#${doc.id} ${doc.title}`
}
onPreviewDocSelected(event: NgbTypeaheadSelectItemEvent<Document>) {
event.preventDefault()
const doc: Document = event.item
this.previewDocModel = doc
this.previewDocId = doc.id
if (!this.template.document_type && doc.document_type) {
this.template.document_type = doc.document_type
}
this.previewPage = 0
this.loadPreview()
}
clearPreviewDoc() {
this.previewDocModel = ''
this.previewDocId = null
this.previewPageCount = null
this.pageCountForDoc = null
this.previewPage = 0
this.pageImageUrl = null
this.imageLoaded = false
}
loadPreview() {
if (!this.previewDocId) return
if (this.pageCountForDoc !== this.previewDocId) {
this.pageCountForDoc = this.previewDocId
this.previewPageCount = null
this.documentService
.get(this.previewDocId)
.pipe(takeUntil(this.destroy$))
.subscribe({
next: (doc) => {
this.previewPageCount = doc?.page_count ?? null
if (doc && !this.previewDocModel) this.previewDocModel = doc
},
error: () => (this.previewPageCount = null),
})
}
this.pageImageUrl = this.templateService.getPageImageUrl(
this.previewDocId,
this.previewPage
)
this.imageLoaded = false
}
goToPage(page: number) {
if (!Number.isFinite(page)) return
const max = this.previewPageCount ? this.previewPageCount - 1 : page
const clamped = Math.max(0, Math.min(page, max))
if (clamped === this.previewPage) return
this.previewPage = clamped
this.loadPreview()
}
prevPage() {
this.goToPage(this.previewPage - 1)
}
nextPage() {
this.goToPage(this.previewPage + 1)
}
zoomIn() {
this.zoom = Math.min(4, Math.round((this.zoom + 0.25) * 100) / 100)
}
zoomOut() {
this.zoom = Math.max(0.5, Math.round((this.zoom - 0.25) * 100) / 100)
}
resetZoom() {
this.zoom = 1
}
zonePage(zone: OcrTemplateZone): number {
return getZonePage(zone, this.previewPage, this.previewPageCount)
}
private isOnCurrentPage(zone: OcrTemplateZone): boolean {
return isZoneOnPage(zone, this.previewPage, this.previewPageCount)
}
onImageLoad() {
this.imageLoaded = true
const img = this.imageRef.nativeElement
this.template.source_width = img.naturalWidth
this.template.source_height = img.naturalHeight
}
onOverlayMouseDown(event: MouseEvent) {
const point = this.svgPointFromEvent(event)
if (!point) return
event.preventDefault()
if (this.selectedZoneIndex !== null) {
const handle = this.findHandleAt(point, this.selectedZoneIndex)
if (handle) {
this.overlayInteraction = {
kind: 'resizing',
zoneIndex: this.selectedZoneIndex,
handle,
}
return
}
}
const clickedIdx = this.findZoneAt(point)
if (clickedIdx !== null && !event.shiftKey) {
this.selectZone(clickedIdx)
const zone = this.template.zones[clickedIdx]
this.overlayInteraction = {
kind: 'moving',
zoneIndex: clickedIdx,
start: {
mouseX: point.x,
mouseY: point.y,
zoneX: zone.x,
zoneY: zone.y,
},
}
return
}
// Shift+click or click on empty area starts a new zone.
this.overlayInteraction = {
kind: 'drawing',
rect: {
startX: point.x,
startY: point.y,
endX: point.x,
endY: point.y,
},
}
this.selectedZoneIndex = null
}
onOverlayMouseMove(event: MouseEvent) {
const point = this.svgPointFromEvent(event)
if (!point) return
if (this.overlayInteraction.kind === 'resizing') {
this.applyResize(
this.overlayInteraction.zoneIndex,
this.overlayInteraction.handle,
point
)
return
}
if (this.overlayInteraction.kind === 'moving') {
moveZone(
this.template.zones[this.overlayInteraction.zoneIndex],
point,
this.overlayInteraction.start,
this.imageNaturalSize(),
this.imageNaturalSize()
)
return
}
if (this.overlayInteraction.kind === 'drawing') {
this.overlayInteraction.rect.endX = point.x
this.overlayInteraction.rect.endY = point.y
return
}
this.updateOverlayCursor(point)
}
private updateOverlayCursor(point: Point) {
if (this.selectedZoneIndex !== null) {
const handle = this.findHandleAt(point, this.selectedZoneIndex)
if (handle) {
this.overlayCursor = RESIZE_CURSOR[handle] || 'crosshair'
return
}
}
this.overlayCursor = this.findZoneAt(point) !== null ? 'move' : 'crosshair'
}
onOverlayMouseUp(_event: MouseEvent) {
if (
this.overlayInteraction.kind === 'moving' ||
this.overlayInteraction.kind === 'resizing'
) {
this.stopOverlayInteraction()
return
}
if (this.overlayInteraction.kind !== 'drawing') return
const drawingRect = this.overlayInteraction.rect
this.stopOverlayInteraction()
const rect = this.sourceRectFromDrawing(drawingRect)
// Ignore tiny accidental clicks.
if (rect.w < MIN_DRAWN_ZONE_SIZE || rect.h < MIN_DRAWN_ZONE_SIZE) {
return
}
this.template.zones.push(this.createZoneFromRect(rect))
this.selectZone(this.template.zones.length - 1)
}
private createZoneFromRect(rect: DisplayRect): OcrTemplateZone {
const imageSize = this.imageNaturalSize()
return {
name: `Zone ${this.template.zones.length + 1}`,
target: DEFAULT_OCR_ZONE_TARGET,
custom_field: this.defaultCustomFieldId(),
x: rect.x,
y: rect.y,
width: rect.w,
height: rect.h,
page: this.previewPageDisplay,
ocr_language: DEFAULT_OCR_ZONE_LANGUAGE,
transform: DEFAULT_OCR_ZONE_TRANSFORM,
date_format: '',
validation_regex: '',
order: this.template.zones.length,
zone_source_width: imageSize.width,
zone_source_height: imageSize.height,
}
}
private defaultCustomFieldId(): number | null {
return this.customFields[0]?.id ?? null
}
@HostListener('document:mouseup')
onDocumentMouseUp() {
if (this.overlayInteraction.kind === 'idle') return
this.stopOverlayInteraction()
}
private stopOverlayInteraction() {
this.overlayInteraction = NO_OVERLAY_INTERACTION
this.overlayCursor = 'crosshair'
}
drawingRect(): DisplayRect | null {
return this.overlayInteraction.kind === 'drawing'
? this.displayRectFromDrawing(this.overlayInteraction.rect)
: null
}
zoneDisplayRect(zoneIdx: number): DisplayRect | null {
const img = this.imageRef?.nativeElement
if (!img || !img.naturalWidth) return null
const zone = this.template.zones[zoneIdx]
if (!zone) return null
if (!this.isOnCurrentPage(zone)) return null
return getZoneDisplayRect(
zone,
this.imageNaturalSize(),
this.imageNaturalSize()
)
}
private findHandleAt(point: Point, zoneIdx: number): ResizeHandle | null {
const r = this.zoneDisplayRect(zoneIdx)
if (!r) return null
return findHandleAt(point, r, this.overlayHandleSize())
}
private applyResize(zoneIndex: number, handle: ResizeHandle, point: Point) {
const zone = this.template.zones[zoneIndex]
if (!zone) return
resizeZone(
zone,
handle,
point,
this.imageNaturalSize(),
this.imageNaturalSize()
)
}
private findZoneAt(point: Point): number | null {
const img = this.imageRef.nativeElement
if (!img.naturalWidth) return null
return findZoneAt(
point,
this.template.zones,
this.previewPage,
this.previewPageCount,
this.imageNaturalSize(),
this.imageNaturalSize()
)
}
overlayViewBox(): string {
const imageSize = this.imageNaturalSize()
return `0 0 ${imageSize.width} ${imageSize.height}`
}
zoneColor(index: number): string {
return ZONE_COLORS[index % ZONE_COLORS.length]
}
zoneFill(index: number): string {
return `${this.zoneColor(index)}33`
}
zoneLabel(zone: OcrTemplateZone, index: number): string {
return zone.name || `Zone ${index + 1}`
}
zoneLabelY(rect: DisplayRect): number {
return Math.max(this.overlayUnitSize(14), rect.y - this.overlayUnitSize(4))
}
resizeHandles(rect: DisplayRect): ResizeHandleMarker[] {
return [
{ handle: 'nw', x: rect.x, y: rect.y },
{ handle: 'n', x: rect.x + rect.w / 2, y: rect.y },
{ handle: 'ne', x: rect.x + rect.w, y: rect.y },
{ handle: 'w', x: rect.x, y: rect.y + rect.h / 2 },
{ handle: 'e', x: rect.x + rect.w, y: rect.y + rect.h / 2 },
{ handle: 'sw', x: rect.x, y: rect.y + rect.h },
{ handle: 's', x: rect.x + rect.w / 2, y: rect.y + rect.h },
{ handle: 'se', x: rect.x + rect.w, y: rect.y + rect.h },
]
}
overlayHandleSize(): number {
return this.overlayUnitSize(HANDLE_SIZE)
}
overlayFontSize(): number {
return this.overlayUnitSize(12)
}
overlayUnitSize(screenPixels: number): number {
const img = this.imageRef?.nativeElement
if (!img?.naturalWidth || !img.clientWidth) return screenPixels
return (screenPixels * img.naturalWidth) / img.clientWidth
}
private svgPointFromEvent(event: MouseEvent): Point | null {
const svg = this.overlayRef?.nativeElement
const matrix = svg?.getScreenCTM()
if (!svg || !matrix) return null
const point = svg.createSVGPoint()
point.x = event.clientX
point.y = event.clientY
const svgPoint = point.matrixTransform(matrix.inverse())
return { x: svgPoint.x, y: svgPoint.y }
}
private displayRectFromDrawing(rect: DrawingRect): DisplayRect {
return {
x: Math.min(rect.startX, rect.endX),
y: Math.min(rect.startY, rect.endY),
w: Math.abs(rect.endX - rect.startX),
h: Math.abs(rect.endY - rect.startY),
}
}
private sourceRectFromDrawing(rect: DrawingRect): DisplayRect {
const displayRect = this.displayRectFromDrawing(rect)
return {
x: Math.round(displayRect.x),
y: Math.round(displayRect.y),
w: Math.round(displayRect.w),
h: Math.round(displayRect.h),
}
}
private imageNaturalSize() {
const img = this.imageRef.nativeElement
return { width: img.naturalWidth, height: img.naturalHeight }
}
removeZone(index: number) {
this.template.zones.splice(index, 1)
if (this.selectedZoneIndex === index) {
this.selectedZoneIndex = null
} else if (this.selectedZoneIndex > index) {
this.selectedZoneIndex--
}
}
selectZone(index: number) {
this.selectedZoneIndex = index
this.activeTab = 'zone'
this.zoneTestResult = null
const zone = this.template.zones[index]
if (zone) {
this.seedCombineDefault(zone)
this.goToPage(this.zonePage(zone) - 1)
}
}
testZone() {
const zone = this.selectedZone
if (!zone || !this.previewDocId) return
this.zoneTesting = true
this.zoneTestResult = null
this.templateService
.testZone(this.previewDocId, this.zoneTestRequest(zone))
.pipe(takeUntil(this.destroy$))
.subscribe({
next: (res) => {
this.zoneTestResult = res
this.zoneTesting = false
},
error: (err) => {
this.zoneTestResult = {
error: err.error?.error || $localize`Test failed`,
}
this.zoneTesting = false
},
})
}
private zoneTestRequest(zone: OcrTemplateZone): ZoneTestRequest {
return {
name: zone.name,
x: zone.x,
y: zone.y,
width: zone.width,
height: zone.height,
page: zone.page ?? 1,
ocr_language: zone.ocr_language,
transform: zone.transform,
date_format: zone.date_format,
validation_regex: zone.validation_regex,
zone_source_width: zone.zone_source_width,
zone_source_height: zone.zone_source_height,
}
}
deleteSelectedZone() {
if (this.selectedZoneIndex === null) return
this.removeZone(this.selectedZoneIndex)
this.activeTab = 'zones'
}
save() {
this.saving = true
this.pruneCombineFormats()
this.template.sample_document = this.previewDocId
const obs = this.isNew
? this.templateService.create(this.template)
: this.templateService.update(this.template)
obs.pipe(takeUntil(this.destroy$)).subscribe({
next: (saved) => {
const idx = this.selectedZoneIndex
this.template = saved
this.isNew = false
this.selectedZoneIndex = idx
this.saving = false
this.toastService.showInfo($localize`OCR template saved.`)
},
error: (e) => {
this.saving = false
this.toastService.showError($localize`Error saving OCR template.`, e)
},
})
}
private ocrLangCache = new WeakMap<
OcrTemplateZone,
{ src: string; arr: string[] }
>()
ocrLanguageArray(zone: OcrTemplateZone): string[] {
const src = zone.ocr_language || ''
const cached = this.ocrLangCache.get(zone)
if (cached && cached.src === src) return cached.arr
const arr = src ? src.split('+').filter(Boolean) : []
this.ocrLangCache.set(zone, { src, arr })
return arr
}
setOcrLanguages(zone: OcrTemplateZone, langs: string[]) {
zone.ocr_language = (langs || []).join('+')
this.ocrLangCache.set(zone, {
src: zone.ocr_language,
arr: langs ? [...langs] : [],
})
}
getCustomFieldName(id: number): string {
const cf = this.customFields.find((f) => f.id === id)
return cf ? cf.name : `Field #${id}`
}
/** Value bound to the field select: a built-in id string or a custom-field id. */
zoneFieldValue(zone: OcrTemplateZone): ZoneFieldSelection {
const target = zone.target || DEFAULT_OCR_ZONE_TARGET
return target === OCR_ZONE_TARGET.CustomField ? zone.custom_field : target
}
setZoneField(zone: OcrTemplateZone, value: ZoneFieldSelection) {
if (isOcrBuiltinTarget(value)) {
zone.target = value
zone.custom_field = null
} else {
zone.target = OCR_ZONE_TARGET.CustomField
zone.custom_field = typeof value === 'number' ? value : null
}
this.seedCombineDefault(zone)
}
fieldKeyFor(zone: OcrTemplateZone): string | null {
const v = this.zoneFieldValue(zone)
return v === null || v === undefined ? null : String(v)
}
zonesForField(zone: OcrTemplateZone): OcrTemplateZone[] {
const key = this.fieldKeyFor(zone)
if (!key) return []
return this.template.zones.filter((z) => this.fieldKeyFor(z) === key)
}
isFieldShared(zone: OcrTemplateZone): boolean {
return this.zonesForField(zone).length > 1
}
getCombineFormat(zone: OcrTemplateZone): string {
const key = this.fieldKeyFor(zone)
return (key && this.template.combine_formats?.[key]) || ''
}
setCombineFormat(zone: OcrTemplateZone, value: string) {
const key = this.fieldKeyFor(zone)
if (!key) return
this.template.combine_formats ??= {}
this.template.combine_formats[key] = value
}
insertCombineToken(zone: OcrTemplateZone, tokenZone: OcrTemplateZone) {
const token = `{${tokenZone.name}}`
const current = this.getCombineFormat(zone)
const sep = current && !current.endsWith(' ') ? ' ' : ''
this.setCombineFormat(zone, `${current}${sep}${token}`)
}
private seedCombineDefault(zone: OcrTemplateZone) {
const key = this.fieldKeyFor(zone)
if (!key) return
const shared = this.zonesForField(zone)
if (shared.length <= 1) return
this.template.combine_formats ??= {}
if (!this.template.combine_formats[key]) {
this.template.combine_formats[key] = shared
.map((z) => `{${z.name}}`)
.join(' ')
}
}
private pruneCombineFormats() {
const formats = this.template.combine_formats
if (!formats) return
const counts = new Map<string, number>()
for (const z of this.template.zones) {
const key = this.fieldKeyFor(z)
if (key) counts.set(key, (counts.get(key) ?? 0) + 1)
}
for (const key of Object.keys(formats)) {
if ((counts.get(key) ?? 0) <= 1) delete formats[key]
}
}
/** Value bound to the date-format select: a preset, '' (auto), or 'custom'. */
dateFormatChoice(zone: OcrTemplateZone): string {
return this.usesCustomDateFormat(zone)
? CUSTOM_DATE_FORMAT_CHOICE
: zone.date_format || ''
}
setDateFormatChoice(zone: OcrTemplateZone, value: string) {
if (value === CUSTOM_DATE_FORMAT_CHOICE) {
this.customDateFormatZones.add(zone)
zone.date_format ||= ''
} else {
this.customDateFormatZones.delete(zone)
zone.date_format = value
}
}
usesCustomDateFormat(zone: OcrTemplateZone): boolean {
return (
this.customDateFormatZones.has(zone) ||
(!!zone.date_format &&
!this.dateFormatOptions.some(
(option) => option.id === zone.date_format
))
)
}
getZoneTargetName(zone: OcrTemplateZone): string {
const target = zone.target || DEFAULT_OCR_ZONE_TARGET
if (target === OCR_ZONE_TARGET.CustomField) {
return zone.custom_field
? this.getCustomFieldName(zone.custom_field)
: $localize`(no field)`
}
return this.builtinTargets.find((t) => t.id === target)?.name ?? target
}
getDocumentTypeName(id: number): string {
const dt = this.documentTypes.find((d) => d.id === id)
return dt ? dt.name : `Type #${id}`
}
openQuickCreate(zoneIndex: number | null) {
if (zoneIndex === null) return
this.quickCreateForZoneIndex = zoneIndex
this.quickCreateName = this.template.zones[zoneIndex]?.name || ''
this.quickCreateType = CustomFieldDataType.String
this.showQuickCreate = true
}
cancelQuickCreate() {
this.showQuickCreate = false
this.quickCreateForZoneIndex = null
}
submitQuickCreate() {
if (!this.quickCreateName.trim()) return
this.templateService
.quickCreateField(this.quickCreateName.trim(), this.quickCreateType)
.pipe(takeUntil(this.destroy$))
.subscribe({
next: (result) => {
this.customFieldsService.clearCache()
this.customFieldsService
.listAll()
.pipe(takeUntil(this.destroy$))
.subscribe((r) => {
this.customFields = r.results
if (this.quickCreateForZoneIndex !== null) {
this.template.zones[this.quickCreateForZoneIndex].custom_field =
result.id
this.template.zones[this.quickCreateForZoneIndex].target =
OCR_ZONE_TARGET.CustomField
}
this.showQuickCreate = false
this.quickCreateForZoneIndex = null
})
},
error: (err) => {
this.toastService.showError(
$localize`Failed to create custom field.`,
err
)
},
})
}
ngOnDestroy() {
this.destroy$.next()
this.destroy$.complete()
}
}
@@ -1,140 +0,0 @@
import { OcrTemplateZone } from 'src/app/data/ocr-template'
import {
findHandleAt,
findZoneAt,
getZoneDisplayRect,
getZonePage,
isZoneOnPage,
moveZone,
resizeZone,
sourceRectFromDrawing,
} from './zone-geometry'
function zone(overrides: Partial<OcrTemplateZone> = {}): OcrTemplateZone {
return {
name: 'Zone',
target: 'custom_field',
custom_field: 1,
x: 100,
y: 200,
width: 300,
height: 400,
page: 1,
ocr_language: 'eng',
transform: 'strip',
validation_regex: '',
order: 0,
...overrides,
}
}
describe('OCR template editor geometry', () => {
it('normalizes zone pages', () => {
expect(getZonePage(zone({ page: 2 }), 0, 5)).toBe(2)
expect(getZonePage(zone({ page: -1 }), 0, 5)).toBe(5)
expect(getZonePage(zone({ page: -1 }), 2, null)).toBe(3)
expect(getZonePage(zone({ page: 0 }), 0, 5)).toBe(1)
expect(getZonePage(zone({ page: undefined }), 0, 5)).toBe(1)
})
it('checks whether a zone is on the current preview page', () => {
expect(isZoneOnPage(zone({ page: 2 }), 1, 5)).toBe(true)
expect(isZoneOnPage(zone({ page: 2 }), 0, 5)).toBe(false)
expect(isZoneOnPage(zone({ page: -1 }), 4, 5)).toBe(true)
})
it('scales source coordinates to canvas display coordinates', () => {
expect(
getZoneDisplayRect(
zone({ x: 100, y: 200, width: 300, height: 400 }),
{ width: 500, height: 1000 },
{ width: 1000, height: 2000 }
)
).toEqual({ x: 50, y: 100, w: 150, h: 200 })
})
it('uses per-zone source dimensions when present', () => {
expect(
getZoneDisplayRect(
zone({
x: 100,
y: 100,
width: 100,
height: 100,
zone_source_width: 1000,
zone_source_height: 1000,
}),
{ width: 500, height: 500 },
{ width: 2000, height: 2000 }
)
).toEqual({ x: 50, y: 50, w: 50, h: 50 })
})
it('finds zones from topmost to bottommost on the current page', () => {
const zones = [
zone({ name: 'first', x: 0, y: 0, width: 100, height: 100, page: 1 }),
zone({ name: 'second', x: 0, y: 0, width: 50, height: 50, page: 1 }),
zone({ name: 'third', x: 0, y: 0, width: 50, height: 50, page: 2 }),
]
expect(
findZoneAt(
{ x: 25, y: 25 },
zones,
0,
2,
{ width: 100, height: 100 },
{ width: 100, height: 100 }
)
).toBe(1)
})
it('finds resize handles around a display rect', () => {
const rect = { x: 10, y: 20, w: 100, h: 200 }
expect(findHandleAt({ x: 10, y: 20 }, rect)).toBe('nw')
expect(findHandleAt({ x: 110, y: 220 }, rect)).toBe('se')
expect(findHandleAt({ x: 60, y: 20 }, rect)).toBe('n')
expect(findHandleAt({ x: 90, y: 160 }, rect)).toBeNull()
})
it('moves zones without leaving source image bounds', () => {
const z = zone({ x: 50, y: 50, width: 100, height: 100 })
moveZone(
z,
{ x: 500, y: 500 },
{ mouseX: 50, mouseY: 50, zoneX: 50, zoneY: 50 },
{ width: 500, height: 500 },
{ width: 500, height: 500 }
)
expect(z.x).toBe(400)
expect(z.y).toBe(400)
})
it('resizes zones without leaving source image bounds', () => {
const z = zone({ x: 50, y: 50, width: 100, height: 100 })
resizeZone(
z,
'se',
{ x: 500, y: 500 },
{ width: 500, height: 500 },
{ width: 200, height: 200 }
)
expect(z.width).toBe(150)
expect(z.height).toBe(150)
})
it('converts drawn canvas rectangles to source rectangles', () => {
expect(
sourceRectFromDrawing(
{ startX: 100, startY: 200, endX: 50, endY: 100 },
{ width: 500, height: 1000 },
{ width: 1000, height: 2000 }
)
).toEqual({ x: 100, y: 200, w: 100, h: 200 })
})
})
@@ -1,201 +0,0 @@
import { OcrTemplateZone } from 'src/app/data/ocr-template'
export interface DrawingRect {
startX: number
startY: number
endX: number
endY: number
}
export interface Dimensions {
width: number
height: number
}
export interface Point {
x: number
y: number
}
export interface DisplayRect {
x: number
y: number
w: number
h: number
}
export interface MoveStart {
mouseX: number
mouseY: number
zoneX: number
zoneY: number
}
export type ResizeHandle = 'n' | 's' | 'e' | 'w' | 'ne' | 'nw' | 'se' | 'sw'
export const HANDLE_SIZE = 8
export const MIN_ZONE_SIZE = 10
export function getZonePage(
zone: OcrTemplateZone,
previewPage: number,
previewPageCount: number | null
): number {
const page = zone.page ?? 1
if (page === -1) return previewPageCount ?? previewPage + 1
return page >= 1 ? page : 1
}
export function isZoneOnPage(
zone: OcrTemplateZone,
previewPage: number,
previewPageCount: number | null
): boolean {
return getZonePage(zone, previewPage, previewPageCount) === previewPage + 1
}
export function getZoneSourceSize(
zone: OcrTemplateZone,
imageSize: Dimensions
): Dimensions {
return {
width: zone.zone_source_width || imageSize.width,
height: zone.zone_source_height || imageSize.height,
}
}
export function getZoneDisplayRect(
zone: OcrTemplateZone,
canvasSize: Dimensions,
imageSize: Dimensions
): DisplayRect {
const sourceSize = getZoneSourceSize(zone, imageSize)
const scaleX = canvasSize.width / sourceSize.width
const scaleY = canvasSize.height / sourceSize.height
return {
x: zone.x * scaleX,
y: zone.y * scaleY,
w: zone.width * scaleX,
h: zone.height * scaleY,
}
}
export function findHandleAt(
point: Point,
rect: DisplayRect,
handleSize = HANDLE_SIZE
): ResizeHandle | null {
const handles: [ResizeHandle, number, number][] = [
['nw', rect.x, rect.y],
['n', rect.x + rect.w / 2, rect.y],
['ne', rect.x + rect.w, rect.y],
['w', rect.x, rect.y + rect.h / 2],
['e', rect.x + rect.w, rect.y + rect.h / 2],
['sw', rect.x, rect.y + rect.h],
['s', rect.x + rect.w / 2, rect.y + rect.h],
['se', rect.x + rect.w, rect.y + rect.h],
]
return (
handles.find(
([, x, y]) =>
Math.abs(point.x - x) <= handleSize &&
Math.abs(point.y - y) <= handleSize
)?.[0] ?? null
)
}
export function findZoneAt(
point: Point,
zones: OcrTemplateZone[],
previewPage: number,
previewPageCount: number | null,
canvasSize: Dimensions,
imageSize: Dimensions
): number | null {
for (let i = zones.length - 1; i >= 0; i--) {
const zone = zones[i]
if (!isZoneOnPage(zone, previewPage, previewPageCount)) continue
const rect = getZoneDisplayRect(zone, canvasSize, imageSize)
if (
point.x >= rect.x &&
point.x <= rect.x + rect.w &&
point.y >= rect.y &&
point.y <= rect.y + rect.h
) {
return i
}
}
return null
}
export function moveZone(
zone: OcrTemplateZone,
point: Point,
moveStart: MoveStart,
canvasSize: Dimensions,
imageSize: Dimensions
) {
const sourceSize = getZoneSourceSize(zone, imageSize)
const scaleX = sourceSize.width / canvasSize.width
const scaleY = sourceSize.height / canvasSize.height
const dx = Math.round((point.x - moveStart.mouseX) * scaleX)
const dy = Math.round((point.y - moveStart.mouseY) * scaleY)
zone.x = clamp(moveStart.zoneX + dx, 0, sourceSize.width - zone.width)
zone.y = clamp(moveStart.zoneY + dy, 0, sourceSize.height - zone.height)
}
export function resizeZone(
zone: OcrTemplateZone,
handle: ResizeHandle,
point: Point,
canvasSize: Dimensions,
imageSize: Dimensions
) {
const sourceSize = getZoneSourceSize(zone, imageSize)
const scaleX = sourceSize.width / canvasSize.width
const scaleY = sourceSize.height / canvasSize.height
const imageX = clamp(Math.round(point.x * scaleX), 0, sourceSize.width)
const imageY = clamp(Math.round(point.y * scaleY), 0, sourceSize.height)
if (handle.includes('w')) {
const right = Math.min(zone.x + zone.width, sourceSize.width)
zone.x = clamp(imageX, 0, right - MIN_ZONE_SIZE)
zone.width = right - zone.x
}
if (handle.includes('e')) {
zone.width = Math.max(MIN_ZONE_SIZE, imageX - zone.x)
}
if (handle.includes('n')) {
const bottom = Math.min(zone.y + zone.height, sourceSize.height)
zone.y = clamp(imageY, 0, bottom - MIN_ZONE_SIZE)
zone.height = bottom - zone.y
}
if (handle.includes('s')) {
zone.height = Math.max(MIN_ZONE_SIZE, imageY - zone.y)
}
}
export function sourceRectFromDrawing(
rect: DrawingRect,
canvasSize: Dimensions,
imageSize: Dimensions
): DisplayRect {
const scaleX = imageSize.width / canvasSize.width
const scaleY = imageSize.height / canvasSize.height
return {
x: Math.round(Math.min(rect.startX, rect.endX) * scaleX),
y: Math.round(Math.min(rect.startY, rect.endY) * scaleY),
w: Math.round(Math.abs(rect.endX - rect.startX) * scaleX),
h: Math.round(Math.abs(rect.endY - rect.startY) * scaleY),
}
}
function clamp(value: number, min: number, max: number): number {
return Math.max(min, Math.min(value, max))
}
@@ -1,75 +0,0 @@
<pngx-page-header
title="OCR Templates"
i18n-title
info="Define extraction zones on document types to automatically populate custom fields via OCR."
i18n-info
>
<button type="button" class="btn btn-sm btn-outline-primary" (click)="createTemplate()" *pngxIfPermissions="{ action: PermissionAction.Add, type: PermissionType.OcrTemplate }">
<i-bs name="plus-circle" class="me-1"></i-bs><ng-container i18n>Create Template</ng-container>
</button>
</pngx-page-header>
<ul class="list-group">
<li class="list-group-item">
<div class="row">
<div class="col" i18n>Name</div>
<div class="col d-none d-sm-flex" i18n>Document Type</div>
<div class="col d-none d-sm-flex" i18n>Zones</div>
<div class="col" i18n>Status</div>
<div class="col" i18n>Actions</div>
</div>
</li>
@if (loading && templates.length === 0) {
<li class="list-group-item">
<div class="spinner-border spinner-border-sm me-2" role="status"></div>
<ng-container i18n>Loading...</ng-container>
</li>
}
@for (t of templates; track t.id) {
<li class="list-group-item">
<div class="row fade" [class.show]="show">
<div class="col d-flex align-items-center"><button class="btn btn-link p-0 text-start" type="button" (click)="editTemplate(t)" [disabled]="!permissionsService.currentUserCan(PermissionAction.Change, PermissionType.OcrTemplate)">{{t.name}}</button></div>
<div class="col d-flex align-items-center d-none d-sm-flex">{{getDocumentTypeName(t)}}</div>
<div class="col d-flex align-items-center d-none d-sm-flex"><code>{{t.zones?.length || 0}}</code></div>
<div class="col d-flex align-items-center">
<div class="form-check form-switch mb-0">
<input type="checkbox" class="form-check-input cursor-pointer" [id]="t.id+'_enable'" [(ngModel)]="t.enabled" (change)="toggleTemplate(t)" *pngxIfPermissions="{ action: PermissionAction.Change, type: PermissionType.OcrTemplate }">
<label class="form-check-label cursor-pointer" [for]="t.id+'_enable'">
<code> @if(t.enabled) { <ng-container i18n>Enabled</ng-container> } @else { <span i18n class="text-muted">Disabled</span> }</code>
</label>
</div>
</div>
<div class="col">
<div class="btn-group d-block d-sm-none">
<div ngbDropdown container="body" class="d-inline-block">
<button type="button" class="btn btn-link" id="actionsMenuMobile{{t.id}}" (click)="$event.stopPropagation()" ngbDropdownToggle>
<i-bs name="three-dots-vertical"></i-bs>
</button>
<div ngbDropdownMenu aria-labelledby="actionsMenuMobile{{t.id}}">
<button (click)="editTemplate(t)" *pngxIfPermissions="{ action: PermissionAction.Change, type: PermissionType.OcrTemplate }" ngbDropdownItem i18n>Edit</button>
<button (click)="deleteTemplate(t)" *pngxIfPermissions="{ action: PermissionAction.Delete, type: PermissionType.OcrTemplate }" ngbDropdownItem i18n>Delete</button>
</div>
</div>
</div>
<div class="btn-toolbar d-none d-sm-flex gap-2" role="toolbar">
<div class="btn-group">
<button *pngxIfPermissions="{ action: PermissionAction.Change, type: PermissionType.OcrTemplate }" class="btn btn-sm btn-outline-secondary" type="button" (click)="editTemplate(t)">
<i-bs width="1em" height="1em" name="pencil" class="me-1"></i-bs><ng-container i18n>Edit</ng-container>
</button>
<button *pngxIfPermissions="{ action: PermissionAction.Delete, type: PermissionType.OcrTemplate }" class="btn btn-sm btn-outline-danger" type="button" (click)="deleteTemplate(t)">
<i-bs width="1em" height="1em" name="trash" class="me-1"></i-bs><ng-container i18n>Delete</ng-container>
</button>
</div>
</div>
</div>
</div>
</li>
}
@if (!loading && templates.length === 0) {
<li class="list-group-item" [class.show]="show" i18n>No OCR templates defined.</li>
}
</ul>
@@ -1,109 +0,0 @@
import { Component, OnInit, inject } from '@angular/core'
import { FormsModule } from '@angular/forms'
import { Router } from '@angular/router'
import { NgbDropdownModule, NgbModal } from '@ng-bootstrap/ng-bootstrap'
import { NgxBootstrapIconsModule } from 'ngx-bootstrap-icons'
import { delay, takeUntil, tap } from 'rxjs'
import { OcrTemplate } from 'src/app/data/ocr-template'
import { IfPermissionsDirective } from 'src/app/directives/if-permissions.directive'
import { PermissionsService } from 'src/app/services/permissions.service'
import { DocumentTypeService } from 'src/app/services/rest/document-type.service'
import { OcrTemplateService } from 'src/app/services/rest/ocr-template.service'
import { ToastService } from 'src/app/services/toast.service'
import { ConfirmDialogComponent } from '../../common/confirm-dialog/confirm-dialog.component'
import { PageHeaderComponent } from '../../common/page-header/page-header.component'
import { LoadingComponentWithPermissions } from '../../loading-component/loading.component'
@Component({
selector: 'pngx-ocr-templates',
templateUrl: './ocr-templates.component.html',
imports: [
PageHeaderComponent,
IfPermissionsDirective,
FormsModule,
NgbDropdownModule,
NgxBootstrapIconsModule,
],
})
export class OcrTemplatesComponent
extends LoadingComponentWithPermissions
implements OnInit
{
private readonly service = inject(OcrTemplateService)
private readonly documentTypeService = inject(DocumentTypeService)
private readonly router = inject(Router)
private readonly modalService = inject(NgbModal)
private readonly toastService = inject(ToastService)
permissionsService = inject(PermissionsService)
public templates: OcrTemplate[] = []
private documentTypeNames: Map<number, string> = new Map()
ngOnInit() {
this.documentTypeService
.listAll()
.pipe(takeUntil(this.unsubscribeNotifier))
.subscribe((r) => {
this.documentTypeNames = new Map(
r.results.map((dt) => [dt.id, dt.name])
)
})
this.reload()
}
reload() {
this.loading = true
this.service
.listAll()
.pipe(
takeUntil(this.unsubscribeNotifier),
tap((r) => (this.templates = r.results)),
delay(100)
)
.subscribe(() => {
this.show = true
this.loading = false
})
}
getDocumentTypeName(t: OcrTemplate): string {
return (
this.documentTypeNames.get(t.document_type) ?? `${t.document_type ?? ''}`
)
}
createTemplate() {
this.router.navigate(['/ocr-templates', 'new'])
}
editTemplate(t: OcrTemplate) {
this.router.navigate(['/ocr-templates', t.id])
}
toggleTemplate(t: OcrTemplate) {
// ngModel has already flipped t.enabled; restore it if persistence fails.
const enabled = t.enabled
this.service.patch(t).subscribe({
error: (error) => {
t.enabled = !enabled
this.toastService.showError(
$localize`Error updating OCR template.`,
error
)
},
})
}
deleteTemplate(t: OcrTemplate) {
const modal = this.modalService.open(ConfirmDialogComponent)
modal.componentInstance.title = $localize`Delete OCR Template`
modal.componentInstance.messageBoldPart = t.name
modal.componentInstance.message = $localize`Do you really want to delete this OCR template?`
modal.componentInstance.btnClass = 'btn-danger'
modal.componentInstance.btnCaption = $localize`Delete`
modal.componentInstance.confirmClicked.subscribe(() => {
modal.close()
this.service.delete(t).subscribe(() => this.reload())
})
}
}
-142
View File
@@ -1,142 +0,0 @@
import { ObjectWithId } from './object-with-id'
export type OcrZoneTarget = 'custom_field' | 'title' | 'asn' | 'created'
export type OcrBuiltinTarget = Exclude<OcrZoneTarget, 'custom_field'>
export type OcrZoneTransform =
| 'none'
| 'strip'
| 'uppercase'
| 'lowercase'
| 'numeric'
| 'strip_punctuation'
| 'date'
| 'qr_code'
export const OCR_ZONE_TARGET = {
CustomField: 'custom_field',
Title: 'title',
Asn: 'asn',
Created: 'created',
} as const satisfies Record<string, OcrZoneTarget>
export const OCR_ZONE_TRANSFORM = {
None: 'none',
Strip: 'strip',
Uppercase: 'uppercase',
Lowercase: 'lowercase',
Numeric: 'numeric',
StripPunctuation: 'strip_punctuation',
Date: 'date',
QrCode: 'qr_code',
} as const satisfies Record<string, OcrZoneTransform>
export const DEFAULT_OCR_ZONE_TARGET = OCR_ZONE_TARGET.CustomField
export const DEFAULT_OCR_ZONE_TRANSFORM = OCR_ZONE_TRANSFORM.Strip
export const DEFAULT_OCR_ZONE_LANGUAGE = 'deu+eng'
export function isOcrBuiltinTarget(value: unknown): value is OcrBuiltinTarget {
return (
value === OCR_ZONE_TARGET.Title ||
value === OCR_ZONE_TARGET.Asn ||
value === OCR_ZONE_TARGET.Created
)
}
export const OCR_BUILTIN_TARGETS = [
{ id: OCR_ZONE_TARGET.Title, name: $localize`Title` },
{ id: OCR_ZONE_TARGET.Asn, name: $localize`Archive serial number` },
{ id: OCR_ZONE_TARGET.Created, name: $localize`Date created` },
]
export interface OcrTemplateZone {
id?: number
name: string
target?: OcrZoneTarget
custom_field: number | null
page?: number
x: number
y: number
width: number
height: number
ocr_language: string
transform: OcrZoneTransform
date_format?: string
validation_regex: string
order: number
zone_source_width?: number
zone_source_height?: number
}
export const TRANSFORM_OPTIONS = [
{ id: OCR_ZONE_TRANSFORM.None, name: $localize`None` },
{ id: OCR_ZONE_TRANSFORM.Strip, name: $localize`Strip whitespace` },
{ id: OCR_ZONE_TRANSFORM.Uppercase, name: $localize`Uppercase` },
{ id: OCR_ZONE_TRANSFORM.Lowercase, name: $localize`Lowercase` },
{ id: OCR_ZONE_TRANSFORM.Numeric, name: $localize`Numeric only` },
{
id: OCR_ZONE_TRANSFORM.StripPunctuation,
name: $localize`Remove leading/trailing punctuation`,
},
{ id: OCR_ZONE_TRANSFORM.Date, name: $localize`Parse date` },
{ id: OCR_ZONE_TRANSFORM.QrCode, name: $localize`Read QR/barcode` },
]
export const OCR_LANGUAGE_OPTIONS = [
{ id: 'eng', name: $localize`English` },
{ id: 'deu', name: $localize`German` },
{ id: 'fra', name: $localize`French` },
{ id: 'ita', name: $localize`Italian` },
{ id: 'spa', name: $localize`Spanish` },
{ id: 'por', name: $localize`Portuguese` },
{ id: 'nld', name: $localize`Dutch` },
]
export const DATE_FORMAT_OPTIONS = [
{ id: '', name: $localize`Auto-detect` },
{ id: '%d.%m.%Y', name: 'DD.MM.YYYY' },
{ id: '%Y/%m/%d', name: 'YYYY/MM/DD' },
{ id: '%d/%m/%Y', name: 'DD/MM/YYYY' },
]
export interface OcrTemplate extends ObjectWithId {
name: string
document_type: number
sample_document: number | null
source_width: number
source_height: number
enabled: boolean
combine_formats?: Record<string, string>
created?: string
updated?: string
zones: OcrTemplateZone[]
}
export interface ZoneTestRequest {
name: string
x: number
y: number
width: number
height: number
page: number
ocr_language: string
transform: OcrZoneTransform
date_format?: string
validation_regex: string
zone_source_width?: number
zone_source_height?: number
}
export interface OcrZoneTestResult {
raw_text?: string | null
value?: string | null
regex?: string
regex_match?: boolean | null
error?: string
}
export interface OcrZoneRunResult {
template: string
zone: string
custom_field: string
value: string | number | null
}
@@ -29,7 +29,6 @@ export enum PermissionType {
ShareLinkBundle = '%s_sharelinkbundle',
CustomField = '%s_customfield',
Workflow = '%s_workflow',
OcrTemplate = '%s_ocrtemplate',
ProcessedMail = '%s_processedmail',
GlobalStatistics = '%s_global_statistics',
SystemMonitoring = '%s_system_monitoring',
@@ -12,7 +12,6 @@ import {
import { DocumentMetadata } from 'src/app/data/document-metadata'
import { DocumentSuggestions } from 'src/app/data/document-suggestions'
import { FilterRule } from 'src/app/data/filter-rule'
import { OcrZoneRunResult } from 'src/app/data/ocr-template'
import { Results, SelectionData } from 'src/app/data/results'
import { SETTINGS_KEYS } from 'src/app/data/ui-settings'
import { queryParamsFromFilterRules } from '../../utils/query-params'
@@ -360,13 +359,6 @@ export class DocumentService extends AbstractPaperlessService<Document> {
})
}
runZoneOcr(id: number): Observable<{ results: OcrZoneRunResult[] }> {
return this.http.post<{ results: OcrZoneRunResult[] }>(
this.getResourceUrl(id, 'run-zone-ocr'),
{}
)
}
rotateDocuments(
selection: DocumentSelectionQuery,
degrees: number,
@@ -1,47 +0,0 @@
import { Injectable } from '@angular/core'
import { Observable } from 'rxjs'
import {
OcrTemplate,
OcrZoneTestResult,
ZoneTestRequest,
} from '../../data/ocr-template'
import { AbstractPaperlessService } from './abstract-paperless-service'
export interface QuickCreateFieldResult {
id: number
name: string
data_type: string
created: boolean
}
@Injectable({ providedIn: 'root' })
export class OcrTemplateService extends AbstractPaperlessService<OcrTemplate> {
constructor() {
super()
this.resourceName = 'ocr_templates'
}
getPageImageUrl(docId: number, page: number): string {
return `${this.baseUrl}${this.resourceName}/document-page-image/${docId}/${page}/`
}
testZone(
docId: number,
zone: ZoneTestRequest
): Observable<OcrZoneTestResult> {
return this.http.post<OcrZoneTestResult>(
`${this.baseUrl}${this.resourceName}/test-zone/`,
{ document: docId, zone }
)
}
quickCreateField(
name: string,
dataType: string
): Observable<QuickCreateFieldResult> {
return this.http.post<QuickCreateFieldResult>(
`${this.baseUrl}${this.resourceName}/quick-create-field/`,
{ name, data_type: dataType }
)
}
}
-4
View File
@@ -89,7 +89,6 @@ import {
exclamationTriangleFill,
eye,
fileEarmark,
fileEarmarkBreak,
fileEarmarkCheck,
fileEarmarkDiff,
fileEarmarkFill,
@@ -100,7 +99,6 @@ import {
fileEarmarkPlus,
fileEarmarkRichtext,
fileEarmarkSpreadsheet,
fileEarmarkRuled,
fileText,
files,
filter,
@@ -340,7 +338,6 @@ const icons = {
exclamationTriangleFill,
eye,
fileEarmark,
fileEarmarkBreak,
fileEarmarkCheck,
fileEarmarkDiff,
fileEarmarkFill,
@@ -351,7 +348,6 @@ const icons = {
fileEarmarkPlus,
fileEarmarkRichtext,
fileEarmarkSpreadsheet,
fileEarmarkRuled,
files,
fileText,
filter,
-13
View File
@@ -13,11 +13,8 @@ class DocumentsConfig(AppConfig):
from documents.signals.handlers import add_inbox_tags
from documents.signals.handlers import add_or_update_document_in_llm_index
from documents.signals.handlers import add_to_index
from documents.signals.handlers import capture_old_document_type
from documents.signals.handlers import run_workflows_added
from documents.signals.handlers import run_workflows_updated
from documents.signals.handlers import run_zone_ocr_extraction
from documents.signals.handlers import run_zone_ocr_on_type_change
from documents.signals.handlers import send_websocket_document_updated
from documents.signals.handlers import set_correspondent
from documents.signals.handlers import set_document_type
@@ -32,16 +29,6 @@ class DocumentsConfig(AppConfig):
document_consumption_finished.connect(add_to_index)
document_consumption_finished.connect(run_workflows_added)
document_consumption_finished.connect(add_or_update_document_in_llm_index)
document_consumption_finished.connect(run_zone_ocr_extraction)
from django.db.models.signals import post_save
from django.db.models.signals import pre_save
from documents.models import Document
pre_save.connect(capture_old_document_type, sender=Document)
post_save.connect(run_zone_ocr_on_type_change, sender=Document)
document_updated.connect(run_workflows_updated)
document_updated.connect(send_websocket_document_updated)
document_updated.connect(add_or_update_document_in_llm_index)
+1 -2
View File
@@ -2,7 +2,6 @@ from __future__ import annotations
import logging
import tempfile
import uuid
from pathlib import Path
from typing import TYPE_CHECKING
from typing import Literal
@@ -380,7 +379,7 @@ def delete(doc_ids: list[int]) -> Literal["OK"]:
)
delete_ids = list({*doc_ids, *version_ids})
Document.objects.filter(id__in=delete_ids).delete(transaction_id=uuid.uuid4())
Document.objects.filter(id__in=delete_ids).delete()
from documents.search import get_backend
@@ -156,15 +156,6 @@ class FileStabilityTracker:
logger.debug(f"File disappeared during stability check: {path}")
continue
# Stable, but empty: some scanners create a zero byte placeholder
# and only write the page some time later. Consuming it now can
# only fail so drop it and let the writer's next event
# (or the periodic rescan) bring it back once it has content
if not tracked.last_size:
to_remove.append(path)
logger.debug("Ignoring stable but empty file: %s", path)
continue
# File is stable, we can return it
to_yield.append(path)
logger.info(f"File is stable: {path}")
@@ -1,267 +0,0 @@
# Generated by Django 5.2.14 on 2026-06-16 17:36
import django.core.validators
import django.db.models.deletion
import django.utils.timezone
from django.db import migrations
from django.db import models
class Migration(migrations.Migration):
dependencies = [
("documents", "0021_widen_workflow_integer_fields"),
]
operations = [
migrations.CreateModel(
name="OcrTemplate",
fields=[
(
"id",
models.AutoField(
auto_created=True,
primary_key=True,
serialize=False,
verbose_name="ID",
),
),
("name", models.CharField(max_length=128, verbose_name="name")),
(
"source_width",
models.PositiveIntegerField(
help_text="Width of the image the zones were drawn on (px)",
validators=[django.core.validators.MinValueValidator(1)],
verbose_name="source width",
),
),
(
"source_height",
models.PositiveIntegerField(
help_text="Height of the image the zones were drawn on (px)",
validators=[django.core.validators.MinValueValidator(1)],
verbose_name="source height",
),
),
("enabled", models.BooleanField(default=True, verbose_name="enabled")),
(
"combine_formats",
models.JSONField(
blank=True,
default=dict,
help_text="Per-target format strings for combining several zones into one field, keyed by target (custom field id, or 'title'/'asn'/'created'). Tokens like {Zone Name} are replaced with that zone's value.",
verbose_name="combine formats",
),
),
(
"created",
models.DateTimeField(
db_index=True,
default=django.utils.timezone.now,
editable=False,
verbose_name="created",
),
),
(
"updated",
models.DateTimeField(auto_now=True, verbose_name="updated"),
),
(
"document_type",
models.ForeignKey(
on_delete=django.db.models.deletion.CASCADE,
related_name="ocr_templates",
to="documents.documenttype",
verbose_name="document type",
),
),
(
"sample_document",
models.ForeignKey(
blank=True,
help_text="Document used for previewing zones in the editor",
null=True,
on_delete=django.db.models.deletion.SET_NULL,
related_name="+",
to="documents.document",
verbose_name="sample document",
),
),
],
options={
"verbose_name": "OCR template",
"verbose_name_plural": "OCR templates",
"ordering": ("name",),
},
),
migrations.CreateModel(
name="OcrTemplateZone",
fields=[
(
"id",
models.AutoField(
auto_created=True,
primary_key=True,
serialize=False,
verbose_name="ID",
),
),
(
"name",
models.CharField(
help_text="Descriptive name for this zone (e.g. 'Invoice Number')",
max_length=128,
verbose_name="zone name",
),
),
(
"target",
models.CharField(
choices=[
("custom_field", "Custom field"),
("title", "Title"),
("asn", "Archive serial number"),
("created", "Date created"),
],
default="custom_field",
help_text="Where the extracted value is written: a custom field, or a built-in document field (title, ASN, created date)",
max_length=20,
verbose_name="target",
),
),
(
"page",
models.IntegerField(
blank=True,
help_text="Page (1 = first, -1 = last; blank uses the template default)",
null=True,
verbose_name="page",
),
),
(
"x",
models.PositiveIntegerField(
help_text="Left edge (px)",
verbose_name="x",
),
),
(
"y",
models.PositiveIntegerField(
help_text="Top edge (px)",
verbose_name="y",
),
),
(
"width",
models.PositiveIntegerField(
help_text="Zone width (px)",
validators=[django.core.validators.MinValueValidator(1)],
verbose_name="width",
),
),
(
"height",
models.PositiveIntegerField(
help_text="Zone height (px)",
validators=[django.core.validators.MinValueValidator(1)],
verbose_name="height",
),
),
(
"zone_source_width",
models.PositiveIntegerField(
blank=True,
help_text="Width of the page image this zone was drawn on (px). Falls back to template source_width if unset.",
null=True,
verbose_name="zone source width",
),
),
(
"zone_source_height",
models.PositiveIntegerField(
blank=True,
help_text="Height of the page image this zone was drawn on (px). Falls back to template source_height if unset.",
null=True,
verbose_name="zone source height",
),
),
(
"ocr_language",
models.CharField(
default="deu+eng",
help_text="Tesseract language code(s), e.g. 'deu+eng'",
max_length=20,
verbose_name="OCR language",
),
),
(
"transform",
models.CharField(
choices=[
("none", "None"),
("strip", "Strip whitespace"),
("uppercase", "Uppercase"),
("lowercase", "Lowercase"),
("numeric", "Numeric only"),
(
"strip_punctuation",
"Remove leading/trailing punctuation",
),
("date", "Parse date"),
("qr_code", "Read QR/barcode"),
],
default="strip",
max_length=20,
verbose_name="transform",
),
),
(
"date_format",
models.CharField(
blank=True,
default="",
help_text="Python strptime format for the 'Parse date' transform (e.g. %d.%m.%Y). Blank = auto-detect.",
max_length=64,
verbose_name="date format",
),
),
(
"validation_regex",
models.CharField(
blank=True,
default="",
help_text="Optional regex pattern — extracted text is only accepted if it matches",
max_length=256,
verbose_name="validation regex",
),
),
("order", models.PositiveIntegerField(default=0, verbose_name="order")),
(
"custom_field",
models.ForeignKey(
blank=True,
help_text="Target custom field (only used when target is 'custom_field')",
null=True,
on_delete=django.db.models.deletion.CASCADE,
related_name="ocr_zones",
to="documents.customfield",
verbose_name="custom field",
),
),
(
"template",
models.ForeignKey(
on_delete=django.db.models.deletion.CASCADE,
related_name="zones",
to="documents.ocrtemplate",
verbose_name="template",
),
),
],
options={
"verbose_name": "OCR template zone",
"verbose_name_plural": "OCR template zones",
"ordering": ("template", "order"),
},
),
]
+2 -255
View File
@@ -1,5 +1,4 @@
import datetime
import uuid
from pathlib import Path
from typing import Final
@@ -515,20 +514,13 @@ class Document(SoftDeleteModel, ModelWithOwner): # type: ignore[django-manager-
def delete(
self,
*args,
transaction_id=None,
**kwargs,
):
# Versions must share the root's transaction ID so they are restored
# together by django-softdelete.
if transaction_id is None:
transaction_id = uuid.uuid4()
# If deleting a root document, move all its versions to trash as well.
if self.root_document_id is None:
Document.objects.filter(root_document=self).delete(
transaction_id=transaction_id,
)
Document.objects.filter(root_document=self).delete()
return super().delete(
*args,
transaction_id=transaction_id,
**kwargs,
)
@@ -2033,248 +2025,3 @@ class WorkflowRun(SoftDeleteModel):
def __str__(self) -> str:
return f"WorkflowRun of {self.workflow} at {self.run_at} on {self.document}"
class OcrTemplate(models.Model):
"""
Defines a set of OCR extraction zones for a specific document type.
When a document of that type is consumed, each zone in the template is
cropped from the document image and OCR'd separately. The extracted text
is written to the configured custom field or built-in document field.
"""
name = models.CharField(
_("name"),
max_length=128,
)
document_type = models.ForeignKey(
"documents.DocumentType",
on_delete=models.CASCADE,
related_name="ocr_templates",
verbose_name=_("document type"),
db_index=True,
)
source_width = models.PositiveIntegerField(
_("source width"),
validators=[MinValueValidator(1)],
help_text=_("Width of the image the zones were drawn on (px)"),
)
source_height = models.PositiveIntegerField(
_("source height"),
validators=[MinValueValidator(1)],
help_text=_("Height of the image the zones were drawn on (px)"),
)
sample_document = models.ForeignKey(
"documents.Document",
on_delete=models.SET_NULL,
null=True,
blank=True,
related_name="+",
verbose_name=_("sample document"),
help_text=_("Document used for previewing zones in the editor"),
)
enabled = models.BooleanField(_("enabled"), default=True)
combine_formats = models.JSONField(
_("combine formats"),
default=dict,
blank=True,
help_text=_(
"Per-target format strings for combining several zones into one "
"field, keyed by target (custom field id, or 'title'/'asn'/'created'). "
"Tokens like {Zone Name} are replaced with that zone's value.",
),
)
created = models.DateTimeField(
_("created"),
default=timezone.now,
db_index=True,
editable=False,
)
updated = models.DateTimeField(
_("updated"),
auto_now=True,
)
class Meta:
ordering = ("name",)
verbose_name = _("OCR template")
verbose_name_plural = _("OCR templates")
def __str__(self) -> str:
return f"{self.name} ({self.document_type})"
class OcrTemplateZone(models.Model):
"""
A rectangular region within a document page to OCR and extract into a custom
field or built-in document field. Coordinates are relative to the source
image dimensions stored on the template.
"""
template = models.ForeignKey(
OcrTemplate,
on_delete=models.CASCADE,
related_name="zones",
verbose_name=_("template"),
)
name = models.CharField(
_("zone name"),
max_length=128,
help_text=_("Descriptive name for this zone (e.g. 'Invoice Number')"),
)
class TargetType(models.TextChoices):
CUSTOM_FIELD = ("custom_field", _("Custom field"))
TITLE = ("title", _("Title"))
ASN = ("asn", _("Archive serial number"))
CREATED = ("created", _("Date created"))
target = models.CharField(
_("target"),
max_length=20,
choices=TargetType.choices,
default=TargetType.CUSTOM_FIELD,
help_text=_(
"Where the extracted value is written: a custom field, or a "
"built-in document field (title, ASN, created date)",
),
)
custom_field = models.ForeignKey(
"documents.CustomField",
on_delete=models.CASCADE,
related_name="ocr_zones",
verbose_name=_("custom field"),
null=True,
blank=True,
help_text=_("Target custom field (only used when target is 'custom_field')"),
)
page = models.IntegerField(
_("page"),
null=True,
blank=True,
help_text=_("Page (1 = first, -1 = last; blank uses the template default)"),
)
x = models.PositiveIntegerField(_("x"), help_text=_("Left edge (px)"))
y = models.PositiveIntegerField(_("y"), help_text=_("Top edge (px)"))
width = models.PositiveIntegerField(
_("width"),
validators=[MinValueValidator(1)],
help_text=_("Zone width (px)"),
)
height = models.PositiveIntegerField(
_("height"),
validators=[MinValueValidator(1)],
help_text=_("Zone height (px)"),
)
# Per-zone source dimensions for coordinate scaling.
# Stored from the page image the zone was drawn on.
# If null, falls back to the template's source_width/source_height.
# This handles PDFs with mixed page sizes (e.g. landscape + portrait,
# or different paper formats across pages).
zone_source_width = models.PositiveIntegerField(
_("zone source width"),
null=True,
blank=True,
help_text=_(
"Width of the page image this zone was drawn on (px). "
"Falls back to template source_width if unset.",
),
)
zone_source_height = models.PositiveIntegerField(
_("zone source height"),
null=True,
blank=True,
help_text=_(
"Height of the page image this zone was drawn on (px). "
"Falls back to template source_height if unset.",
),
)
ocr_language = models.CharField(
_("OCR language"),
max_length=20,
default="deu+eng",
help_text=_("Tesseract language code(s), e.g. 'deu+eng'"),
)
class TransformType(models.TextChoices):
NONE = ("none", _("None"))
STRIP = ("strip", _("Strip whitespace"))
UPPERCASE = ("uppercase", _("Uppercase"))
LOWERCASE = ("lowercase", _("Lowercase"))
NUMERIC = ("numeric", _("Numeric only"))
STRIP_PUNCTUATION = (
"strip_punctuation",
_("Remove leading/trailing punctuation"),
)
DATE = ("date", _("Parse date"))
QR_CODE = ("qr_code", _("Read QR/barcode"))
transform = models.CharField(
_("transform"),
max_length=20,
choices=TransformType.choices,
default=TransformType.STRIP,
)
date_format = models.CharField(
_("date format"),
max_length=64,
blank=True,
default="",
help_text=_(
"Python strptime format for the 'Parse date' transform "
"(e.g. %d.%m.%Y). Blank = auto-detect.",
),
)
validation_regex = models.CharField(
_("validation regex"),
max_length=256,
blank=True,
default="",
help_text=_(
"Optional regex pattern — extracted text is only accepted if it matches",
),
)
order = models.PositiveIntegerField(_("order"), default=0)
class Meta:
ordering = ("template", "order")
verbose_name = _("OCR template zone")
verbose_name_plural = _("OCR template zones")
def __str__(self) -> str:
return f"{self.template.name} -> {self.name}"
# Custom field data types that zone OCR can extract into. DOCUMENTLINK and
# SELECT are excluded (they reference other objects, not free text). Single
# source of truth for the serializer, the quick-create endpoint and the engine.
OCR_SUPPORTED_FIELD_TYPES = frozenset(
{
CustomField.FieldDataType.STRING,
CustomField.FieldDataType.URL,
CustomField.FieldDataType.DATE,
CustomField.FieldDataType.INT,
CustomField.FieldDataType.FLOAT,
CustomField.FieldDataType.MONETARY,
CustomField.FieldDataType.LONG_TEXT,
CustomField.FieldDataType.BOOL,
},
)
-129
View File
@@ -56,7 +56,6 @@ if settings.AUDIT_LOG_ENABLED:
from documents import bulk_edit
from documents.data_models import DocumentSource
from documents.filters import CustomFieldQueryParser
from documents.models import OCR_SUPPORTED_FIELD_TYPES
from documents.models import Correspondent
from documents.models import CustomField
from documents.models import CustomFieldInstance
@@ -64,8 +63,6 @@ from documents.models import Document
from documents.models import DocumentType
from documents.models import MatchingModel
from documents.models import Note
from documents.models import OcrTemplate
from documents.models import OcrTemplateZone
from documents.models import PaperlessTask
from documents.models import SavedView
from documents.models import SavedViewFilterRule
@@ -3665,129 +3662,3 @@ class StoragePathTestSerializer(SerializerWithPerms):
document_field.queryset = Document.objects.filter(
id__in=permitted_document_ids(user),
)
class OcrTemplateZoneSerializer(serializers.ModelSerializer):
class Meta:
model = OcrTemplateZone
fields = [
"id",
"name",
"target",
"custom_field",
"page",
"x",
"y",
"width",
"height",
"ocr_language",
"transform",
"date_format",
"order",
"zone_source_width",
"zone_source_height",
"validation_regex",
]
def validate_width(self, value):
if value < 1:
raise serializers.ValidationError("Width must be at least 1.")
return value
def validate_height(self, value):
if value < 1:
raise serializers.ValidationError("Height must be at least 1.")
return value
def validate_custom_field(self, value):
if value is None:
# Built-in target (title/asn/created) — no custom field required.
return value
if value.data_type not in OCR_SUPPORTED_FIELD_TYPES:
raise serializers.ValidationError(
f"Custom field type '{value.data_type}' is not supported for OCR extraction. "
f"Use string, integer, float, date, monetary, boolean, URL, or long text.",
)
return value
class OcrTemplateSerializer(serializers.ModelSerializer):
zones = OcrTemplateZoneSerializer(many=True, required=False)
class Meta:
model = OcrTemplate
fields = [
"id",
"name",
"document_type",
"source_width",
"source_height",
"sample_document",
"enabled",
"combine_formats",
"created",
"updated",
"zones",
]
read_only_fields = ["created", "updated"]
def validate_source_width(self, value):
if value < 1:
raise serializers.ValidationError("Source width must be at least 1.")
return value
def validate_source_height(self, value):
if value < 1:
raise serializers.ValidationError("Source height must be at least 1.")
return value
def validate_zones(self, zones_data):
"""Validate zone coordinates are within the source dimensions."""
# source_width/height may not be in initial_data during partial updates
source_width = self.initial_data.get("source_width") or (
self.instance.source_width if self.instance else None
)
source_height = self.initial_data.get("source_height") or (
self.instance.source_height if self.instance else None
)
if source_width and source_height:
for zone in zones_data:
x = zone.get("x", 0)
y = zone.get("y", 0)
w = zone.get("width", 0)
h = zone.get("height", 0)
if x + w > int(source_width):
raise serializers.ValidationError(
f"Zone '{zone.get('name', '?')}' extends beyond source width "
f"({x + w} > {source_width}).",
)
if y + h > int(source_height):
raise serializers.ValidationError(
f"Zone '{zone.get('name', '?')}' extends beyond source height "
f"({y + h} > {source_height}).",
)
return zones_data
def create(self, validated_data):
zones_data = validated_data.pop("zones", [])
template = OcrTemplate.objects.create(**validated_data)
for zone_data in zones_data:
OcrTemplateZone.objects.create(template=template, **zone_data)
return template
def update(self, instance, validated_data):
zones_data = validated_data.pop("zones", None)
for attr, value in validated_data.items():
setattr(instance, attr, value)
instance.save()
if zones_data is not None:
# Replace all zones with the new set
instance.zones.all().delete()
for zone_data in zones_data:
OcrTemplateZone.objects.create(template=instance, **zone_data)
return instance
-70
View File
@@ -1398,76 +1398,6 @@ def close_connection_pool_on_worker_init(**kwargs) -> None:
conn.close_pool()
def run_zone_ocr_extraction(sender, document, original_file=None, **kwargs):
"""
Run zone-based OCR extraction if the document's type has an active template.
"""
try:
from documents.zone_ocr import run_zone_extraction
run_zone_extraction(document, Path(original_file) if original_file else None)
except Exception:
logger.exception(
"Zone OCR extraction failed for document %s",
document.pk,
)
def capture_old_document_type(sender, instance, **kwargs):
"""pre_save: remember the document's previous type so the post_save handler
can tell whether the type actually changed (vs. every other save)."""
if instance.pk:
instance._old_document_type_id = (
Document.objects.filter(pk=instance.pk)
.values_list("document_type_id", flat=True)
.first()
)
else:
instance._old_document_type_id = None
def run_zone_ocr_on_type_change(sender, instance, *, created=False, **kwargs):
"""
Run zone OCR only when a document's TYPE actually changes (and the new type
has an enabled template). NOT on every save zone OCR overwrites fields, so
re-running it on each edit would clobber the user's changes. Newly created
documents are handled by the consumption signal, and the user can always
trigger extraction manually via the run-zone-ocr action.
"""
if created or not instance.pk or not instance.document_type_id:
return
# Only proceed if the type changed compared to what was in the DB before.
old_type = getattr(instance, "_old_document_type_id", None)
if old_type == instance.document_type_id:
return
from documents.models import OcrTemplate
if not OcrTemplate.objects.filter(
document_type_id=instance.document_type_id,
enabled=True,
).exists():
return
try:
from documents.zone_ocr import run_zone_extraction
doc_path = instance.archive_path or instance.source_path
if doc_path and Path(doc_path).is_file():
logger.info(
"Zone OCR: running extraction for document %d (type %d)",
instance.pk,
instance.document_type_id,
)
run_zone_extraction(instance, None)
except Exception:
logger.exception(
"Zone OCR extraction failed for document %s",
instance.pk,
)
@worker_process_shutdown.connect
def close_connection_pool_on_worker_shutdown(**kwargs) -> None: # pragma: no cover
"""
@@ -1,449 +0,0 @@
"""Tests for the OCR Template API."""
import json
from django.contrib.auth.models import User
from rest_framework import status
from rest_framework.test import APITestCase
from documents.models import CustomField
from documents.models import DocumentType
from documents.models import OcrTemplate
from documents.models import OcrTemplateZone
from documents.tests.utils import DirectoriesMixin
class TestOcrTemplatesAPI(DirectoriesMixin, APITestCase):
ENDPOINT = "/api/ocr_templates/"
def setUp(self) -> None:
self.user = User.objects.create_superuser(username="temp_admin")
self.client.force_authenticate(user=self.user)
self.doc_type = DocumentType.objects.create(name="Invoice")
self.custom_field_text = CustomField.objects.create(
name="Invoice Number",
data_type=CustomField.FieldDataType.STRING,
)
self.custom_field_date = CustomField.objects.create(
name="Invoice Date",
data_type=CustomField.FieldDataType.DATE,
)
self.custom_field_int = CustomField.objects.create(
name="Amount",
data_type=CustomField.FieldDataType.INT,
)
self.custom_field_doclink = CustomField.objects.create(
name="Related Docs",
data_type=CustomField.FieldDataType.DOCUMENTLINK,
)
return super().setUp()
def _make_template_data(self, **overrides):
data = {
"name": "Invoice Template",
"document_type": self.doc_type.pk,
"default_page": 0,
"source_width": 2480,
"source_height": 3508,
"enabled": True,
"zones": [],
}
data.update(overrides)
return data
def _make_zone_data(self, **overrides):
data = {
"name": "Zone 1",
"custom_field": self.custom_field_text.pk,
"x": 100,
"y": 100,
"width": 200,
"height": 50,
"ocr_language": "deu+eng",
"transform": "strip",
"order": 0,
}
data.update(overrides)
return data
# --- Create ---
def test_create_template(self):
"""
GIVEN:
- A document type and custom fields exist
WHEN:
- API request to create an OCR template with one zone
THEN:
- The template and zone are created
"""
data = self._make_template_data(
zones=[
self._make_zone_data(
name="Invoice Number",
x=1500,
y=200,
width=800,
height=100,
),
],
)
resp = self.client.post(
self.ENDPOINT,
data=json.dumps(data),
content_type="application/json",
)
self.assertEqual(resp.status_code, status.HTTP_201_CREATED)
result = resp.json()
self.assertEqual(result["name"], "Invoice Template")
self.assertEqual(result["document_type"], self.doc_type.pk)
self.assertEqual(len(result["zones"]), 1)
self.assertEqual(result["zones"][0]["name"], "Invoice Number")
self.assertEqual(OcrTemplate.objects.count(), 1)
self.assertEqual(OcrTemplateZone.objects.count(), 1)
def test_create_template_multiple_zones(self):
"""
GIVEN:
- Multiple custom fields exist
WHEN:
- A template with multiple zones is created
THEN:
- All zones are created
"""
data = self._make_template_data(
zones=[
self._make_zone_data(
name="Invoice Number",
custom_field=self.custom_field_text.pk,
),
self._make_zone_data(
name="Invoice Date",
custom_field=self.custom_field_date.pk,
order=1,
),
],
)
resp = self.client.post(
self.ENDPOINT,
data=json.dumps(data),
content_type="application/json",
)
self.assertEqual(resp.status_code, status.HTTP_201_CREATED)
self.assertEqual(len(resp.json()["zones"]), 2)
self.assertEqual(OcrTemplateZone.objects.count(), 2)
def test_create_template_no_zones(self):
"""
GIVEN:
- Valid template data without zones
WHEN:
- Template is created
THEN:
- Template is created with no zones
"""
data = self._make_template_data()
resp = self.client.post(
self.ENDPOINT,
data=json.dumps(data),
content_type="application/json",
)
self.assertEqual(resp.status_code, status.HTTP_201_CREATED)
self.assertEqual(len(resp.json()["zones"]), 0)
# --- Validation ---
def test_create_template_zero_source_width_rejected(self):
"""
GIVEN:
- Template data with source_width=0
WHEN:
- Create is attempted
THEN:
- 400 error is returned
"""
data = self._make_template_data(source_width=0)
resp = self.client.post(
self.ENDPOINT,
data=json.dumps(data),
content_type="application/json",
)
self.assertEqual(resp.status_code, status.HTTP_400_BAD_REQUEST)
def test_create_template_zero_source_height_rejected(self):
data = self._make_template_data(source_height=0)
resp = self.client.post(
self.ENDPOINT,
data=json.dumps(data),
content_type="application/json",
)
self.assertEqual(resp.status_code, status.HTTP_400_BAD_REQUEST)
def test_create_zone_zero_width_rejected(self):
data = self._make_template_data(
zones=[self._make_zone_data(width=0)],
)
resp = self.client.post(
self.ENDPOINT,
data=json.dumps(data),
content_type="application/json",
)
self.assertEqual(resp.status_code, status.HTTP_400_BAD_REQUEST)
def test_create_zone_zero_height_rejected(self):
data = self._make_template_data(
zones=[self._make_zone_data(height=0)],
)
resp = self.client.post(
self.ENDPOINT,
data=json.dumps(data),
content_type="application/json",
)
self.assertEqual(resp.status_code, status.HTTP_400_BAD_REQUEST)
def test_create_zone_exceeds_source_width_rejected(self):
"""Zone that extends beyond the source image width should be rejected."""
data = self._make_template_data(
source_width=1000,
zones=[self._make_zone_data(x=800, width=300)], # 800+300 > 1000
)
resp = self.client.post(
self.ENDPOINT,
data=json.dumps(data),
content_type="application/json",
)
self.assertEqual(resp.status_code, status.HTTP_400_BAD_REQUEST)
def test_create_zone_exceeds_source_height_rejected(self):
data = self._make_template_data(
source_height=1000,
zones=[self._make_zone_data(y=900, height=200)], # 900+200 > 1000
)
resp = self.client.post(
self.ENDPOINT,
data=json.dumps(data),
content_type="application/json",
)
self.assertEqual(resp.status_code, status.HTTP_400_BAD_REQUEST)
def test_create_zone_unsupported_custom_field_type_rejected(self):
"""DOCUMENTLINK and SELECT fields can't be populated via OCR."""
data = self._make_template_data(
zones=[self._make_zone_data(custom_field=self.custom_field_doclink.pk)],
)
resp = self.client.post(
self.ENDPOINT,
data=json.dumps(data),
content_type="application/json",
)
self.assertEqual(resp.status_code, status.HTTP_400_BAD_REQUEST)
# --- List ---
def test_list_templates(self):
template = OcrTemplate.objects.create(
name="Test Template",
document_type=self.doc_type,
source_width=2480,
source_height=3508,
)
OcrTemplateZone.objects.create(
template=template,
name="Zone 1",
custom_field=self.custom_field_text,
x=100,
y=100,
width=200,
height=50,
)
resp = self.client.get(self.ENDPOINT)
self.assertEqual(resp.status_code, status.HTTP_200_OK)
data = resp.json()
self.assertEqual(data["count"], 1)
self.assertEqual(len(data["results"][0]["zones"]), 1)
def test_list_empty(self):
resp = self.client.get(self.ENDPOINT)
self.assertEqual(resp.status_code, status.HTTP_200_OK)
self.assertEqual(resp.json()["count"], 0)
# --- Update ---
def test_update_template_replaces_zones(self):
"""PUT should replace all zones with the new set."""
template = OcrTemplate.objects.create(
name="Old Name",
document_type=self.doc_type,
source_width=2480,
source_height=3508,
)
OcrTemplateZone.objects.create(
template=template,
name="Old Zone",
custom_field=self.custom_field_text,
x=0,
y=0,
width=100,
height=100,
)
data = self._make_template_data(
name="New Name",
zones=[
self._make_zone_data(
name="New Zone",
custom_field=self.custom_field_date.pk,
),
],
)
resp = self.client.put(
f"{self.ENDPOINT}{template.pk}/",
data=json.dumps(data),
content_type="application/json",
)
self.assertEqual(resp.status_code, status.HTTP_200_OK)
template.refresh_from_db()
self.assertEqual(template.name, "New Name")
self.assertEqual(OcrTemplateZone.objects.count(), 1)
self.assertEqual(OcrTemplateZone.objects.first().name, "New Zone")
# --- Delete ---
def test_delete_template_cascades_zones(self):
template = OcrTemplate.objects.create(
name="To Delete",
document_type=self.doc_type,
source_width=2480,
source_height=3508,
)
OcrTemplateZone.objects.create(
template=template,
name="Zone",
custom_field=self.custom_field_text,
x=0,
y=0,
width=100,
height=100,
)
resp = self.client.delete(f"{self.ENDPOINT}{template.pk}/")
self.assertEqual(resp.status_code, status.HTTP_204_NO_CONTENT)
self.assertEqual(OcrTemplate.objects.count(), 0)
self.assertEqual(OcrTemplateZone.objects.count(), 0)
def test_delete_nonexistent_returns_404(self):
resp = self.client.delete(f"{self.ENDPOINT}99999/")
self.assertEqual(resp.status_code, status.HTTP_404_NOT_FOUND)
# --- Patch ---
def test_patch_toggle_enabled(self):
template = OcrTemplate.objects.create(
name="Toggle Test",
document_type=self.doc_type,
source_width=2480,
source_height=3508,
enabled=True,
)
resp = self.client.patch(
f"{self.ENDPOINT}{template.pk}/",
data=json.dumps({"enabled": False}),
content_type="application/json",
)
self.assertEqual(resp.status_code, status.HTTP_200_OK)
template.refresh_from_db()
self.assertFalse(template.enabled)
def test_patch_preserves_zones(self):
"""PATCH without zones field should not delete existing zones."""
template = OcrTemplate.objects.create(
name="Patch Test",
document_type=self.doc_type,
source_width=2480,
source_height=3508,
)
OcrTemplateZone.objects.create(
template=template,
name="Existing Zone",
custom_field=self.custom_field_text,
x=0,
y=0,
width=100,
height=100,
)
resp = self.client.patch(
f"{self.ENDPOINT}{template.pk}/",
data=json.dumps({"name": "Updated Name"}),
content_type="application/json",
)
self.assertEqual(resp.status_code, status.HTTP_200_OK)
self.assertEqual(OcrTemplateZone.objects.count(), 1)
# --- Auth ---
def test_unauthenticated_rejected(self):
self.client.logout()
resp = self.client.get(self.ENDPOINT)
self.assertIn(
resp.status_code,
(status.HTTP_401_UNAUTHORIZED, status.HTTP_403_FORBIDDEN),
)
# --- Quick create field ---
def test_quick_create_field(self):
"""Creating a custom field inline from the template editor."""
resp = self.client.post(
f"{self.ENDPOINT}quick-create-field/",
data=json.dumps({"name": "New Field", "data_type": "string"}),
content_type="application/json",
)
self.assertEqual(resp.status_code, status.HTTP_201_CREATED)
data = resp.json()
self.assertEqual(data["name"], "New Field")
self.assertEqual(data["data_type"], "string")
self.assertTrue(data["created"])
self.assertTrue(CustomField.objects.filter(name="New Field").exists())
def test_quick_create_field_existing(self):
"""If a field with the same name exists, return it without creating."""
resp = self.client.post(
f"{self.ENDPOINT}quick-create-field/",
data=json.dumps({"name": "Invoice Number", "data_type": "string"}),
content_type="application/json",
)
self.assertEqual(resp.status_code, status.HTTP_200_OK)
data = resp.json()
self.assertEqual(data["id"], self.custom_field_text.pk)
self.assertFalse(data["created"])
def test_quick_create_field_empty_name_rejected(self):
resp = self.client.post(
f"{self.ENDPOINT}quick-create-field/",
data=json.dumps({"name": "", "data_type": "string"}),
content_type="application/json",
)
self.assertEqual(resp.status_code, status.HTTP_400_BAD_REQUEST)
def test_quick_create_field_unsupported_type_rejected(self):
resp = self.client.post(
f"{self.ENDPOINT}quick-create-field/",
data=json.dumps({"name": "Bad Field", "data_type": "documentlink"}),
content_type="application/json",
)
self.assertEqual(resp.status_code, status.HTTP_400_BAD_REQUEST)
def test_quick_create_field_select_type_rejected(self):
resp = self.client.post(
f"{self.ENDPOINT}quick-create-field/",
data=json.dumps({"name": "Bad Field", "data_type": "select"}),
content_type="application/json",
)
self.assertEqual(resp.status_code, status.HTTP_400_BAD_REQUEST)
-62
View File
@@ -207,65 +207,3 @@ class TestTrashAPI(DirectoriesMixin, APITestCase):
)
self.assertEqual(resp.status_code, status.HTTP_400_BAD_REQUEST)
self.assertIn("have not yet been deleted", resp.data["documents"][0])
def _make_versioned_document(self) -> tuple[Document, list[Document]]:
root = Document.objects.create(
title="root",
content="root-content",
checksum="root",
mime_type="application/pdf",
)
versions = [
Document.objects.create(
title=f"v{index}",
content=f"v{index}-content",
checksum=f"v{index}",
mime_type="application/pdf",
root_document=root,
version_index=index,
)
for index in range(1, 3)
]
return root, versions
def test_api_trash_restore_document_restores_its_versions(self) -> None:
"""
GIVEN:
- Existing document with two versions
WHEN:
- API request to delete the document
- API request to restore it from the trash
THEN:
- Only the document itself is listed in the trash
- A version cannot be restored without its root
- The document is restored together with all of its versions
"""
root, versions = self._make_versioned_document()
self.client.force_login(user=self.user)
self.client.delete(f"/api/documents/{root.pk}/")
self.assertEqual(Document.deleted_objects.count(), 3)
resp = self.client.get("/api/trash/")
self.assertEqual(resp.status_code, status.HTTP_200_OK)
self.assertEqual(resp.data["count"], 1)
self.assertEqual(resp.data["results"][0]["id"], root.pk)
# A version cannot be restored while its root remains in the trash.
resp = self.client.post(
"/api/trash/",
{"action": "restore", "documents": [versions[0].pk]},
)
self.assertEqual(resp.status_code, status.HTTP_400_BAD_REQUEST)
self.assertIn("Restore the root document", resp.data["documents"][0])
resp = self.client.post(
"/api/trash/",
{"action": "restore", "documents": [root.pk]},
)
self.assertEqual(resp.status_code, status.HTTP_200_OK)
self.assertEqual(Document.deleted_objects.count(), 0)
self.assertCountEqual(
Document.objects.filter(root_document=root).values_list("id", flat=True),
[version.pk for version in versions],
)
-5
View File
@@ -392,11 +392,6 @@ class TestBulkEdit(DirectoriesMixin, TestCase):
self.assertFalse(Document.objects.filter(id=self.doc1.id).exists())
self.assertFalse(Document.objects.filter(id=version.id).exists())
Document.deleted_objects.get(id=self.doc1.id).restore(strict=False)
self.assertTrue(Document.objects.filter(id=self.doc1.id).exists())
self.assertTrue(Document.objects.filter(id=version.id).exists())
def test_delete_version_document_keeps_root(self) -> None:
version = Document.objects.create(
checksum="A-v1",
+1 -5
View File
@@ -110,7 +110,7 @@ class TestDocument(TestCase):
checksum="checksum",
mime_type="application/pdf",
)
version = Document.objects.create(
Document.objects.create(
root_document=root,
correspondent=root.correspondent,
title="Version",
@@ -124,10 +124,6 @@ class TestDocument(TestCase):
self.assertEqual(Document.objects.count(), 0)
self.assertEqual(Document.deleted_objects.count(), 2)
root.restore(strict=False)
self.assertTrue(Document.objects.filter(pk=version.pk).exists())
def test_file_name(self) -> None:
doc = Document(
mime_type="application/pdf",
@@ -136,23 +136,6 @@ def wait_for_mock_call(
return False
def sleep_past_stability(
owner: FileStabilityTracker | ConsumerThread,
*,
windows: float = 1.5,
) -> None:
"""
Block until a tracked file's stability window has certainly elapsed.
Args:
owner: The tracker, or the consumer thread running one, whose
configured stability delay sets the wait.
windows: How many stability windows to wait, giving slop for a slow
or loaded test runner.
"""
sleep(owner.stability_delay * windows)
class TestTrackedFile:
"""Tests for the TrackedFile dataclass."""
@@ -278,56 +261,6 @@ class TestFileStabilityTracker:
assert len(stable) == 0
assert stability_tracker.pending_count == 1
def test_get_stable_files_skips_empty_file(
self,
stability_tracker: FileStabilityTracker,
tmp_path: Path,
) -> None:
"""
GIVEN:
- A zero byte file, tracked and past its stability delay
WHEN:
- Stable files are collected
THEN:
- The file is not yielded for consumption
- The file is dropped from tracking rather than held, so an
abandoned placeholder does not keep the watch loop awake
"""
empty = tmp_path / "scan.pdf"
empty.write_bytes(b"")
stability_tracker.track(empty, Change.added)
sleep_past_stability(stability_tracker)
stable = list(stability_tracker.get_stable_files())
assert stable == []
assert stability_tracker.pending_count == 0
def test_empty_file_is_yielded_once_content_arrives(
self,
stability_tracker: FileStabilityTracker,
tmp_path: Path,
) -> None:
"""
GIVEN:
- A zero byte file which was dropped from tracking while empty
WHEN:
- The writer fills the file and a new event re-tracks it
THEN:
- The file is yielded for consumption once it is stable
"""
target = tmp_path / "scan.pdf"
target.write_bytes(b"")
stability_tracker.track(target, Change.added)
sleep_past_stability(stability_tracker)
assert list(stability_tracker.get_stable_files()) == []
target.write_bytes(b"%PDF-1.4 content")
stability_tracker.track(target, Change.modified)
sleep_past_stability(stability_tracker)
assert list(stability_tracker.get_stable_files()) == [target]
def test_get_stable_files_deleted_during_check(self, temp_file: Path) -> None:
"""Test deleted file is not returned during stability check."""
tracker = FileStabilityTracker(stability_delay=0.1)
@@ -946,51 +879,6 @@ class TestCommandWatch:
mock_consume_file_delay.apply_async.assert_called()
def test_scanner_placeholder_is_not_consumed_while_empty(
self,
consumption_dir: Path,
sample_pdf: Path,
mock_consume_file_delay: MagicMock,
start_consumer: Callable[..., ConsumerThread],
) -> None:
"""
GIVEN:
- A scanner which creates a zero byte placeholder and only writes
the page some time later (GH discussion #13969)
WHEN:
- The placeholder sits untouched well past the stability delay
- The scanner then writes the real content
THEN:
- The empty placeholder is never queued, as it could only fail
with "Unsupported mime type inode/x-empty"
- The file is queued exactly once, when the content lands
"""
thread = start_consumer(stability_delay=0.2)
target = consumption_dir / "scan.pdf"
target.write_bytes(b"") # the scanner's placeholder
# Well past the stability delay: the old behaviour queued it here.
sleep_past_stability(thread, windows=5)
if thread.exception:
raise thread.exception
assert mock_consume_file_delay.apply_async.call_count == 0
shutil.copy(sample_pdf, target) # the scanner finishes the page
assert wait_for_mock_call(
mock_consume_file_delay.apply_async,
timeout_s=5.0,
)
if thread.exception:
raise thread.exception
assert mock_consume_file_delay.apply_async.call_count == 1
queued_doc = mock_consume_file_delay.apply_async.call_args.kwargs["kwargs"][
"input_doc"
]
assert queued_doc.original_file.name == "scan.pdf"
def test_ignores_macos_files(
self,
consumption_dir: Path,
-33
View File
@@ -32,7 +32,6 @@ from documents.signals.handlers import update_llm_suggestions_cache
from documents.tests.utils import DirectoriesMixin
from documents.tests.utils import read_streaming_response
from paperless.models import ApplicationConfiguration
from paperless_ai.exceptions import LLMProviderError
from paperless_ai.exceptions import LLMTimeoutError
@@ -738,38 +737,6 @@ class TestAISuggestions(DirectoriesMixin, TestCase):
get_llm_suggestion_cache(self.document.pk, backend="openai-like"),
)
@patch("documents.views.get_ai_document_classification")
@override_settings(
AI_ENABLED=True,
LLM_BACKEND="openai-like",
)
def test_ai_suggestions_with_llm_provider_error(
self,
mock_get_ai_classification,
) -> None:
mock_get_ai_classification.side_effect = LLMProviderError(
"confidential provider response",
)
self.client.force_login(user=self.user)
response = self.client.get(
f"/api/documents/{self.document.pk}/ai_suggestions/",
)
self.assertEqual(response.status_code, status.HTTP_502_BAD_GATEWAY)
self.assertEqual(
response.json(),
{
"ai": [
"AI backend rejected the request. Check logs for details.",
],
},
)
self.assertNotIn("confidential provider response", response.content.decode())
self.assertIsNone(
get_llm_suggestion_cache(self.document.pk, backend="openai-like"),
)
@patch("documents.views.get_ai_document_classification")
@override_settings(
AI_ENABLED=True,
-454
View File
@@ -1,454 +0,0 @@
"""Tests for the zone-based OCR extraction engine."""
import tempfile
from pathlib import Path
from unittest.mock import MagicMock
from unittest.mock import patch
from django.test import TestCase
from documents.models import CustomField
from documents.models import CustomFieldInstance
from documents.models import Document
from documents.models import DocumentType
from documents.models import OcrTemplate
from documents.models import OcrTemplateZone
from documents.zone_ocr import _apply_transform
from documents.zone_ocr import _convert_value
from documents.zone_ocr import _detect_mime
from documents.zone_ocr import _resolve_doc_path
from documents.zone_ocr import run_zone_extraction
class TestApplyTransform(TestCase):
"""Tests for the _apply_transform function."""
def test_strip(self):
self.assertEqual(_apply_transform(" hello ", "strip"), "hello")
def test_none_transform(self):
self.assertEqual(_apply_transform(" hello ", "none"), "hello")
def test_uppercase(self):
self.assertEqual(_apply_transform("hello world", "uppercase"), "HELLO WORLD")
def test_lowercase(self):
self.assertEqual(_apply_transform("HELLO WORLD", "lowercase"), "hello world")
def test_numeric_basic(self):
self.assertEqual(_apply_transform("INV-2026-001", "numeric"), "2026-001")
def test_numeric_with_currency(self):
self.assertEqual(_apply_transform("€1,234.56", "numeric"), "1,234.56")
def test_numeric_empty_result_falls_back(self):
self.assertEqual(_apply_transform("abc", "numeric"), "abc")
def test_date_dmy_dots(self):
self.assertEqual(_apply_transform("13.04.2026", "date_dmy"), "2026-04-13")
def test_date_dmy_slashes(self):
self.assertEqual(_apply_transform("01/12/2025", "date_dmy"), "2025-12-01")
def test_date_dmy_two_digit_year(self):
self.assertEqual(_apply_transform("13.04.26", "date_dmy"), "2026-04-13")
def test_date_dmy_with_prefix(self):
self.assertEqual(_apply_transform("Date: 01/12/2025", "date_dmy"), "2025-12-01")
def test_date_dmy_invalid_falls_back(self):
self.assertEqual(_apply_transform("32.13.2026", "date_dmy"), "32.13.2026")
def test_date_dmy_no_match_falls_back(self):
self.assertEqual(_apply_transform("not a date", "date_dmy"), "not a date")
def test_date_ymd_dashes(self):
self.assertEqual(_apply_transform("2026-04-13", "date_ymd"), "2026-04-13")
def test_date_ymd_slashes(self):
self.assertEqual(_apply_transform("2026/04/13", "date_ymd"), "2026-04-13")
def test_date_ymd_invalid_falls_back(self):
self.assertEqual(_apply_transform("2026-13-32", "date_ymd"), "2026-13-32")
def test_empty_string(self):
self.assertEqual(_apply_transform("", "strip"), "")
def test_whitespace_only(self):
self.assertEqual(_apply_transform(" ", "strip"), "")
def test_unknown_transform_strips(self):
self.assertEqual(_apply_transform(" hello ", "unknown"), "hello")
class TestConvertValue(TestCase):
"""Tests for the _convert_value function."""
def test_string(self):
self.assertEqual(
_convert_value("Hello", CustomField.FieldDataType.STRING),
"Hello",
)
def test_string_truncation(self):
result = _convert_value("x" * 200, CustomField.FieldDataType.STRING)
self.assertEqual(len(result), 128)
def test_url(self):
self.assertEqual(
_convert_value("https://example.com", CustomField.FieldDataType.URL),
"https://example.com",
)
def test_long_text(self):
long = "x" * 500
self.assertEqual(
_convert_value(long, CustomField.FieldDataType.LONG_TEXT),
long,
)
def test_int_simple(self):
self.assertEqual(_convert_value("42", CustomField.FieldDataType.INT), 42)
def test_int_with_noise(self):
self.assertEqual(_convert_value("INV-123", CustomField.FieldDataType.INT), 123)
def test_int_negative(self):
self.assertEqual(_convert_value("-42", CustomField.FieldDataType.INT), -42)
def test_int_empty_returns_none(self):
self.assertIsNone(_convert_value("abc", CustomField.FieldDataType.INT))
def test_int_only_dash_returns_none(self):
self.assertIsNone(_convert_value("-", CustomField.FieldDataType.INT))
def test_float_simple(self):
self.assertAlmostEqual(
_convert_value("1234.56", CustomField.FieldDataType.FLOAT),
1234.56,
)
def test_float_european_format(self):
self.assertAlmostEqual(
_convert_value("1.234,56", CustomField.FieldDataType.FLOAT),
1234.56,
)
def test_float_us_format(self):
self.assertAlmostEqual(
_convert_value("1,234.56", CustomField.FieldDataType.FLOAT),
1234.56,
)
def test_float_comma_only(self):
self.assertAlmostEqual(
_convert_value("1234,56", CustomField.FieldDataType.FLOAT),
1234.56,
)
def test_float_empty_returns_none(self):
self.assertIsNone(_convert_value("abc", CustomField.FieldDataType.FLOAT))
def test_float_only_separator_returns_none(self):
self.assertIsNone(_convert_value(",", CustomField.FieldDataType.FLOAT))
def test_date_iso(self):
self.assertEqual(
_convert_value("2026-04-13", CustomField.FieldDataType.DATE),
"2026-04-13",
)
def test_date_invalid_returns_none(self):
self.assertIsNone(_convert_value("not a date", CustomField.FieldDataType.DATE))
def test_date_invalid_values_returns_none(self):
self.assertIsNone(_convert_value("2026-13-32", CustomField.FieldDataType.DATE))
def test_monetary_simple(self):
self.assertEqual(
_convert_value("123.45", CustomField.FieldDataType.MONETARY),
"123.45",
)
def test_monetary_european(self):
self.assertEqual(
_convert_value("1.234,56", CustomField.FieldDataType.MONETARY),
"1234.56",
)
def test_monetary_with_currency_symbol(self):
self.assertEqual(
_convert_value("€1,234.56", CustomField.FieldDataType.MONETARY),
"1234.56",
)
def test_monetary_empty_returns_none(self):
self.assertIsNone(_convert_value("CHF", CustomField.FieldDataType.MONETARY))
def test_bool_true(self):
for val in ("true", "True", "yes", "1", "ja", "x", "X"):
self.assertTrue(
_convert_value(val, CustomField.FieldDataType.BOOL),
f"Expected True for {val!r}",
)
def test_bool_false(self):
for val in ("false", "False", "no", "0", "nein"):
self.assertFalse(
_convert_value(val, CustomField.FieldDataType.BOOL),
f"Expected False for {val!r}",
)
def test_bool_unknown_returns_none(self):
self.assertIsNone(_convert_value("maybe", CustomField.FieldDataType.BOOL))
def test_unsupported_type_returns_none(self):
self.assertIsNone(
_convert_value("test", CustomField.FieldDataType.DOCUMENTLINK),
)
self.assertIsNone(
_convert_value("test", CustomField.FieldDataType.SELECT),
)
def test_empty_string_returns_none(self):
self.assertIsNone(_convert_value("", CustomField.FieldDataType.STRING))
class TestDetectMime(TestCase):
"""Tests for _detect_mime."""
def test_pdf_extension(self):
self.assertEqual(_detect_mime(Path("test.pdf")), "application/pdf")
def test_png_extension(self):
self.assertEqual(_detect_mime(Path("test.png")), "image/png")
def test_jpg_extension(self):
self.assertEqual(_detect_mime(Path("test.jpg")), "image/jpeg")
def test_unknown_extension(self):
self.assertIsNone(_detect_mime(Path("test.xyz")))
def test_webp_extension(self):
self.assertEqual(_detect_mime(Path("test.webp")), "image/webp")
class TestResolveDocPath(TestCase):
"""Tests for _resolve_doc_path."""
def test_returns_none_when_no_files_exist(self):
doc = MagicMock()
doc.has_archive_version = False
doc.source_path = Path("/nonexistent/source.pdf")
result = _resolve_doc_path(doc, None)
self.assertIsNone(result)
def test_returns_original_file_as_fallback(self):
doc = MagicMock()
doc.has_archive_version = False
doc.source_path = Path("/nonexistent/source.pdf")
with tempfile.NamedTemporaryFile(suffix=".pdf") as f:
result = _resolve_doc_path(doc, Path(f.name))
self.assertEqual(result, Path(f.name))
def test_returns_none_for_none_original_file(self):
doc = MagicMock()
doc.has_archive_version = False
doc.source_path = Path("/nonexistent/source.pdf")
result = _resolve_doc_path(doc, None)
self.assertIsNone(result)
class TestRunZoneExtraction(TestCase):
"""Tests for the full extraction pipeline."""
def setUp(self):
self.doc_type = DocumentType.objects.create(name="Invoice")
self.custom_field = CustomField.objects.create(
name="Invoice Number",
data_type=CustomField.FieldDataType.STRING,
)
def test_skips_document_without_type(self):
doc = Document.objects.create(
title="No Type",
content="test",
mime_type="application/pdf",
)
run_zone_extraction(doc, Path("/nonexistent"))
self.assertEqual(CustomFieldInstance.objects.count(), 0)
def test_skips_document_without_matching_template(self):
other_type = DocumentType.objects.create(name="Other")
doc = Document.objects.create(
title="No Template",
content="test",
mime_type="application/pdf",
document_type=other_type,
)
run_zone_extraction(doc, Path("/nonexistent"))
self.assertEqual(CustomFieldInstance.objects.count(), 0)
def test_skips_disabled_template(self):
template = OcrTemplate.objects.create(
name="Disabled",
document_type=self.doc_type,
source_width=2480,
source_height=3508,
enabled=False,
)
OcrTemplateZone.objects.create(
template=template,
name="Zone",
custom_field=self.custom_field,
x=0,
y=0,
width=100,
height=50,
)
doc = Document.objects.create(
title="Test",
content="test",
mime_type="application/pdf",
document_type=self.doc_type,
)
run_zone_extraction(doc, Path("/nonexistent"))
self.assertEqual(CustomFieldInstance.objects.count(), 0)
def test_skips_template_with_no_zones(self):
OcrTemplate.objects.create(
name="Empty",
document_type=self.doc_type,
source_width=2480,
source_height=3508,
enabled=True,
)
doc = Document.objects.create(
title="Test",
content="test",
mime_type="application/pdf",
document_type=self.doc_type,
)
with tempfile.NamedTemporaryFile(suffix=".pdf") as f:
f.write(b"%PDF-1.4 fake")
f.flush()
run_zone_extraction(doc, Path(f.name))
self.assertEqual(CustomFieldInstance.objects.count(), 0)
@patch("documents.zone_ocr._process_template")
def test_calls_process_for_enabled_template(self, mock_process):
template = OcrTemplate.objects.create(
name="Active",
document_type=self.doc_type,
source_width=2480,
source_height=3508,
enabled=True,
)
OcrTemplateZone.objects.create(
template=template,
name="Zone",
custom_field=self.custom_field,
x=0,
y=0,
width=100,
height=50,
)
doc = Document.objects.create(
title="Test",
content="test",
mime_type="application/pdf",
document_type=self.doc_type,
)
with tempfile.NamedTemporaryFile(suffix=".pdf") as f:
f.write(b"%PDF-1.4 fake")
f.flush()
run_zone_extraction(doc, Path(f.name))
self.assertTrue(mock_process.called)
@patch("documents.zone_ocr._process_template")
def test_handles_process_exception_gracefully(self, mock_process):
"""A failing template should not prevent other templates from running."""
mock_process.side_effect = RuntimeError("test error")
template = OcrTemplate.objects.create(
name="Failing",
document_type=self.doc_type,
source_width=2480,
source_height=3508,
enabled=True,
)
OcrTemplateZone.objects.create(
template=template,
name="Zone",
custom_field=self.custom_field,
x=0,
y=0,
width=100,
height=50,
)
doc = Document.objects.create(
title="Test",
content="test",
mime_type="application/pdf",
document_type=self.doc_type,
)
with tempfile.NamedTemporaryFile(suffix=".pdf") as f:
f.write(b"%PDF-1.4 fake")
f.flush()
# Should not raise
run_zone_extraction(doc, Path(f.name))
def test_handles_none_original_file(self):
"""Should not crash when original_file is None."""
doc = Document.objects.create(
title="Test",
content="test",
mime_type="application/pdf",
document_type=self.doc_type,
)
# No template, so it exits early — but shouldn't crash on None
run_zone_extraction(doc, None)
@patch("documents.zone_ocr._process_template")
def test_multiple_templates_all_process(self, mock_process):
"""Multiple enabled templates for the same type should all run."""
for i in range(3):
template = OcrTemplate.objects.create(
name=f"Template {i}",
document_type=self.doc_type,
source_width=2480,
source_height=3508,
enabled=True,
)
OcrTemplateZone.objects.create(
template=template,
name=f"Zone {i}",
custom_field=self.custom_field,
x=0,
y=0,
width=100,
height=50,
)
doc = Document.objects.create(
title="Test",
content="test",
mime_type="application/pdf",
document_type=self.doc_type,
)
with tempfile.NamedTemporaryFile(suffix=".pdf") as f:
f.write(b"%PDF-1.4 fake")
f.flush()
run_zone_extraction(doc, Path(f.name))
self.assertEqual(mock_process.call_count, 3)
+2 -322
View File
@@ -3,7 +3,6 @@ import logging
import os
import platform
import re
import subprocess
import tempfile
import zipfile
from collections import defaultdict
@@ -150,14 +149,12 @@ from documents.matching import match_correspondents
from documents.matching import match_document_types
from documents.matching import match_storage_paths
from documents.matching import match_tags
from documents.models import OCR_SUPPORTED_FIELD_TYPES
from documents.models import Correspondent
from documents.models import CustomField
from documents.models import CustomFieldInstance
from documents.models import Document
from documents.models import DocumentType
from documents.models import Note
from documents.models import OcrTemplate
from documents.models import PaperlessTask
from documents.models import SavedView
from documents.models import ShareLink
@@ -204,7 +201,6 @@ from documents.serialisers import EmailSerializer
from documents.serialisers import MergeDocumentsAsVersionsSerializer
from documents.serialisers import MergeDocumentsSerializer
from documents.serialisers import NotesSerializer
from documents.serialisers import OcrTemplateSerializer
from documents.serialisers import PostDocumentSerializer
from documents.serialisers import RemovePasswordDocumentsSerializer
from documents.serialisers import ReprocessDocumentsSerializer
@@ -256,7 +252,6 @@ from paperless.views import StandardPagination
from paperless_ai.ai_classifier import get_ai_document_classification
from paperless_ai.ai_classifier import get_llm_output_language
from paperless_ai.chat import stream_chat_with_documents
from paperless_ai.exceptions import LLMProviderError
from paperless_ai.exceptions import LLMTimeoutError
from paperless_ai.matching import extract_unmatched_names
from paperless_ai.matching import match_correspondents_by_name
@@ -1608,22 +1603,6 @@ class DocumentViewSet(
{"ai": [_("AI backend request timed out.")]},
status=status.HTTP_503_SERVICE_UNAVAILABLE,
)
except LLMProviderError:
logger.exception(
"AI backend rejected the request for document %s",
doc.pk,
)
return Response(
{
"ai": [
_(
"AI backend rejected the request. "
"Check logs for details.",
),
],
},
status=status.HTTP_502_BAD_GATEWAY,
)
set_llm_suggestions_cache(
doc.pk,
llm_suggestions,
@@ -2172,73 +2151,6 @@ class DocumentViewSet(
},
),
)
@action(methods=["post"], detail=True, url_path="run-zone-ocr")
def run_zone_ocr(self, request, pk=None):
"""Run zone-based OCR extraction on this document."""
try:
document = Document.objects.get(pk=pk)
except Document.DoesNotExist:
raise Http404
if not document.document_type_id:
return Response(
{"error": "Document has no type assigned"},
status=status.HTTP_400_BAD_REQUEST,
)
templates = OcrTemplate.objects.filter(
document_type_id=document.document_type_id,
enabled=True,
)
if not templates.exists():
return Response(
{"error": "No OCR templates found for this document type"},
status=status.HTTP_404_NOT_FOUND,
)
doc_path = document.archive_path or document.source_path
if not doc_path or not Path(doc_path).is_file():
return Response(
{"error": "Document file not found"},
status=status.HTTP_404_NOT_FOUND,
)
from documents.zone_ocr import run_zone_extraction
run_zone_extraction(document, None)
# Collect results
results = []
builtin_labels = {"title": "Title", "asn": "ASN", "created": "Created"}
for template in templates.prefetch_related("zones", "zones__custom_field"):
for zone in template.zones.all():
target = getattr(zone, "target", None) or "custom_field"
if target == "custom_field" and zone.custom_field_id:
cf_instance = document.custom_fields.filter(
field=zone.custom_field,
).first()
field_name = zone.custom_field.name
value = cf_instance.value if cf_instance else None
else:
field_name = builtin_labels.get(target, target)
value = {
"title": document.title,
"asn": document.archive_serial_number,
"created": document.created.isoformat()
if document.created
else None,
}.get(target)
results.append(
{
"template": template.name,
"zone": zone.name,
"custom_field": field_name,
"value": value,
},
)
return Response({"results": results})
@action(
methods=["delete"],
detail=True,
@@ -5525,10 +5437,7 @@ class TrashView(ListModelMixin, PassUserMixin):
model = Document
# A version is listed separately only when its root is not in the trash.
queryset = Document.deleted_objects.exclude(
root_document_id__in=Document.deleted_objects.values("id"),
)
queryset = Document.deleted_objects.all()
def get(self, request: Request, format: str | None = None) -> Response:
self.serializer_class = DocumentSerializer
@@ -5559,15 +5468,7 @@ class TrashView(ListModelMixin, PassUserMixin):
return HttpResponseForbidden("Insufficient permissions")
action = serializer.validated_data.get("action")
if action == "restore":
restored = list(self.get_queryset().filter(id__in=doc_ids))
if len(restored) != len(doc_ids):
raise ValidationError(
{
"documents": [
"Restore the root document instead of one of its versions.",
],
},
)
restored = list(Document.deleted_objects.filter(id__in=doc_ids))
for doc in restored:
doc.restore(strict=False)
if restored:
@@ -5613,224 +5514,3 @@ def serve_logo(request: HttpRequest, filename: str | None = None) -> FileRespons
filename=logo_name,
as_attachment=True,
)
class OcrTemplateViewSet(ModelViewSet):
"""CRUD for OCR templates with zone definitions."""
queryset = (
OcrTemplate.objects.all()
.prefetch_related(
"zones",
"zones__custom_field",
)
.order_by("name")
)
serializer_class = OcrTemplateSerializer
permission_classes = (IsAuthenticated, PaperlessObjectPermissions)
pagination_class = StandardPagination
@action(
detail=False,
methods=["get"],
url_path=r"document-page-image/(?P<doc_id>[0-9]+)/(?P<page>[0-9]+)",
)
def document_page_image(self, request, doc_id=None, page=None):
"""Render a specific page of a document as a PNG image.
Used by the frontend template editor to display document pages
as images that users can draw zones on.
"""
try:
document = Document.objects.get(pk=doc_id)
except Document.DoesNotExist:
raise Http404("Document not found")
page_num = int(page)
# Validate page number
if document.page_count and page_num >= document.page_count:
raise Http404(
f"Page {page_num} out of range (document has {document.page_count} pages)",
)
doc_path = document.archive_path or document.source_path
if not doc_path or not Path(doc_path).is_file():
raise Http404("Document file not found")
# Check if document is an image (single page, no PDF rendering needed)
if document.mime_type and document.mime_type.startswith("image/"):
content = Path(doc_path).read_bytes()
return HttpResponse(content, content_type=document.mime_type)
with tempfile.TemporaryDirectory(dir=settings.SCRATCH_DIR) as tmp_dir:
output_prefix = Path(tmp_dir) / "page"
try:
subprocess.run(
[
"pdftoppm",
"-png",
"-r",
"150", # Lower DPI for preview
"-f",
str(page_num + 1),
"-l",
str(page_num + 1),
str(doc_path),
str(output_prefix),
],
check=True,
capture_output=True,
timeout=30,
)
except subprocess.CalledProcessError as e:
raise Http404(
f"Failed to render page: {e.stderr.decode(errors='replace')[:200]}",
)
except FileNotFoundError:
raise Http404("pdftoppm not available - is poppler-utils installed?")
rendered = sorted(Path(tmp_dir).glob("page-*.png"))
if not rendered:
raise Http404("No rendered page found")
content = rendered[0].read_bytes()
return HttpResponse(content, content_type="image/png")
@action(detail=False, methods=["post"], url_path="test-zone")
def test_zone(self, request):
"""Run OCR on a single ad-hoc zone of a document and return what it
yields: the raw OCR text, the transformed value, and whether the
validation regex matches. Non-destructive - writes nothing. Used by the
editor's per-zone test so a user can tune the zone/regex before saving.
Accepts: {"document": <id>, "zone": {x, y, width, height, page,
ocr_language, transform, validation_regex, zone_source_width,
zone_source_height}}.
"""
from documents.models import OcrTemplateZone
from documents.zone_ocr import extract_zone_preview
zone_data = request.data.get("zone") or {}
try:
document = Document.objects.get(pk=request.data.get("document"))
except (Document.DoesNotExist, ValueError, TypeError):
return Response(
{"error": "Document not found"},
status=status.HTTP_404_NOT_FOUND,
)
doc_path = document.archive_path or document.source_path
if not doc_path or not Path(doc_path).is_file():
return Response(
{"error": "Document file not found"},
status=status.HTTP_404_NOT_FOUND,
)
try:
zone = OcrTemplateZone(
name=zone_data.get("name") or "test",
x=int(zone_data.get("x", 0)),
y=int(zone_data.get("y", 0)),
width=int(zone_data.get("width", 0)),
height=int(zone_data.get("height", 0)),
page=zone_data.get("page"),
ocr_language=zone_data.get("ocr_language") or "eng",
transform=zone_data.get("transform") or "strip",
date_format=zone_data.get("date_format") or "",
validation_regex=zone_data.get("validation_regex") or "",
)
except (ValueError, TypeError):
return Response(
{"error": "Invalid zone definition"},
status=status.HTTP_400_BAD_REQUEST,
)
if zone.width < 2 or zone.height < 2:
return Response(
{"error": "Zone is too small to test"},
status=status.HTTP_400_BAD_REQUEST,
)
result = extract_zone_preview(
Path(doc_path),
zone,
int(zone_data.get("zone_source_width") or 0),
int(zone_data.get("zone_source_height") or 0),
document.page_count,
)
regex_match = None
if zone.validation_regex and result.get("value") is not None:
try:
regex_match = (
re.fullmatch(zone.validation_regex, result["value"]) is not None
)
except re.error:
regex_match = None
return Response(
{
"raw_text": result.get("raw_text"),
"value": result.get("value"),
"regex": zone.validation_regex,
"regex_match": regex_match,
},
)
@action(detail=False, methods=["post"], url_path="quick-create-field")
def quick_create_field(self, request):
"""Create a custom field inline from the template editor.
Accepts: {"name": "Invoice Number", "data_type": "string"}
Returns the created field so the frontend can immediately use it.
"""
name = request.data.get("name", "").strip()
data_type = request.data.get("data_type", "").strip()
if not name:
return Response(
{"error": "Field name is required"},
status=status.HTTP_400_BAD_REQUEST,
)
if data_type not in OCR_SUPPORTED_FIELD_TYPES:
return Response(
{
"error": f"Unsupported data type '{data_type}'. "
f"Supported: {', '.join(sorted(OCR_SUPPORTED_FIELD_TYPES))}",
},
status=status.HTTP_400_BAD_REQUEST,
)
# Check if field already exists
existing = CustomField.objects.filter(name=name).first()
if existing:
return Response(
{
"id": existing.pk,
"name": existing.name,
"data_type": existing.data_type,
"created": False,
},
)
# Check user has permission to create custom fields
if not request.user.has_perm("documents.add_customfield"):
return Response(
{"error": "You don't have permission to create custom fields"},
status=status.HTTP_403_FORBIDDEN,
)
field = CustomField.objects.create(name=name, data_type=data_type)
return Response(
{
"id": field.pk,
"name": field.name,
"data_type": field.data_type,
"created": True,
},
status=status.HTTP_201_CREATED,
)
-757
View File
@@ -1,757 +0,0 @@
"""
Zone-based OCR extraction engine.
After a document is consumed, this module checks if the document's type has
an active OCR template. If so, it renders the relevant pages as images,
crops each zone, runs Tesseract OCR on the crop, applies transforms,
and writes the results to the mapped custom fields.
"""
from __future__ import annotations
import logging
import re
import string
import subprocess
import tempfile
from datetime import date
from datetime import datetime
from pathlib import Path
from django.conf import settings
from PIL import Image
from documents.models import CustomField
from documents.models import CustomFieldInstance
from documents.models import Document
from documents.models import OcrTemplate
from documents.models import OcrTemplateZone
logger = logging.getLogger("paperless.zone_ocr")
def run_zone_extraction(
document: Document,
original_file: Path | None,
) -> None:
"""
Run zone-based OCR extraction for a document if its type has an active template.
Called from the document_consumption_finished signal handler.
"""
if not document.document_type_id:
return
templates = OcrTemplate.objects.filter(
document_type_id=document.document_type_id,
enabled=True,
).prefetch_related("zones", "zones__custom_field")
if not templates.exists():
return
# Resolve the document file: prefer archive (PDF/A), then source, then signal arg
doc_path = _resolve_doc_path(document, original_file)
if doc_path is None:
logger.warning(
"Zone OCR: no accessible file for document %d",
document.pk,
)
return
for template in templates:
zones = list(template.zones.all())
if not zones:
continue
logger.info(
"Zone OCR: processing template '%s' for document %d (%d zones)",
template.name,
document.pk,
len(zones),
)
try:
_process_template(document, doc_path, template, zones)
except Exception:
logger.exception(
"Zone OCR: error processing template '%s' for document %d",
template.name,
document.pk,
)
def _resolve_doc_path(
document: Document,
original_file: Path | None,
) -> Path | None:
"""Find an accessible file for the document."""
candidates = []
if document.has_archive_version:
candidates.append(document.archive_path)
candidates.append(document.source_path)
if original_file is not None:
candidates.append(original_file)
for path in candidates:
if path is not None and Path(path).is_file():
return Path(path)
return None
def _resolve_page_idx(page_value, page_count) -> int:
"""Resolve a 1-indexed page (1 = first, -1 = last) to a 0-indexed image
index. A blank page_value defaults to the first page."""
if page_value is None:
return 0
if page_value == -1:
return (page_count - 1) if page_count else 0
if page_value >= 1:
return page_value - 1
return 0
def _process_template(
document: Document,
doc_path: Path,
template: OcrTemplate,
zones: list[OcrTemplateZone],
) -> None:
"""Process all zones in a template against a document.
Each zone is OCR'd independently, then zones are grouped by their target
field and each field is written exactly once. When several zones share a
field, their values are combined via the template's per-field format string
(or joined in order if none is set) this avoids the zones overwriting each
other's value.
"""
pages_needed: set[int] = {
_resolve_page_idx(zone.page, document.page_count) for zone in zones
}
with tempfile.TemporaryDirectory(dir=settings.SCRATCH_DIR) as tmp_dir:
tmp_path = Path(tmp_dir)
page_images = _render_pages(
doc_path,
pages_needed,
tmp_path,
document.page_count,
)
# Pass 1: OCR every zone into a value (or None if it failed/was rejected).
zone_values: dict[int, str | None] = {}
for zone in zones:
page_idx = _resolve_page_idx(zone.page, document.page_count)
if page_idx not in page_images:
logger.warning(
"Zone OCR: page %d not available for zone '%s'",
page_idx,
zone.name,
)
continue
src_w = zone.zone_source_width or template.source_width
src_h = zone.zone_source_height or template.source_height
extracted = _extract_zone(
page_images[page_idx],
zone,
src_w,
src_h,
tmp_path,
)
if (
extracted is not None
and zone.validation_regex
and not re.fullmatch(zone.validation_regex, extracted)
):
logger.info(
"Zone OCR: '%s' value %r rejected by regex '%s'",
zone.name,
extracted[:100],
zone.validation_regex,
)
extracted = None
zone_values[id(zone)] = extracted
# Pass 2: group zones by target field and write each field once.
grouped: dict[str, list[OcrTemplateZone]] = {}
for zone in zones:
grouped.setdefault(_field_key(zone), []).append(zone)
combine_formats = template.combine_formats or {}
for key, field_zones in grouped.items():
value = _combine_field_value(
combine_formats.get(key, ""),
field_zones,
zone_values,
)
if not value:
continue
target_zone = field_zones[0]
_write_zone_value(document, target_zone, value)
logger.info(
"Zone OCR: %s = %r (from %d zone(s))",
_zone_target_label(target_zone),
value[:100] if len(value) > 100 else value,
len(field_zones),
)
def _field_key(zone: OcrTemplateZone) -> str:
"""Identify a zone's target field. Custom fields key by id, built-in targets
by their name. Matches the key used in OcrTemplate.combine_formats and on the
frontend field select."""
target = getattr(zone, "target", None) or "custom_field"
if target == "custom_field" and zone.custom_field_id:
return str(zone.custom_field_id)
return target
def _combine_field_value(
fmt: str,
field_zones: list[OcrTemplateZone],
zone_values: dict[int, str | None],
) -> str:
"""Combine the OCR values of all zones targeting one field.
With a format string, `{Zone Name}` tokens are replaced by that zone's value
and literal text is kept; separators left dangling by an empty token are
cleaned up. Without a format, the zone values are joined in order by a space.
"""
values = {z.name: (zone_values.get(id(z)) or "") for z in field_zones}
if not fmt:
parts = [zone_values.get(id(z)) or "" for z in field_zones]
return " ".join(p for p in parts if p).strip()
def _replace(match: re.Match) -> str:
return values.get(match.group(1).strip(), "")
combined = re.sub(r"\{([^{}]+)\}", _replace, fmt)
# Tidy up separators an empty token may have left behind.
combined = re.sub(r"\s{2,}", " ", combined)
combined = re.sub(r"([^\w\s])\s*\1+", r"\1", combined)
return combined.strip().strip("-/.,;:| \t")
def _render_pages(
doc_path: Path,
pages: set[int],
tmp_dir: Path,
page_count: int | None,
) -> dict[int, Path]:
"""Render specific PDF pages as PNG images using pdftoppm (poppler-utils)."""
result: dict[int, Path] = {}
mime = _detect_mime(doc_path)
if mime and mime.startswith("image/"):
# Single-image document — use it directly as page 0.
result[0] = doc_path
return result
# Callers pass already-resolved 0-indexed page numbers (see _resolve_page_idx).
for actual_page in pages:
if actual_page < 0:
logger.warning("Zone OCR: invalid page index %d", actual_page)
continue
output_prefix = tmp_dir / f"page_{actual_page}"
try:
subprocess.run(
[
"pdftoppm",
"-png",
"-r",
"300",
"-f",
str(actual_page + 1), # pdftoppm is 1-indexed
"-l",
str(actual_page + 1),
str(doc_path),
str(output_prefix),
],
check=True,
capture_output=True,
timeout=60,
)
except subprocess.TimeoutExpired:
logger.error("Zone OCR: pdftoppm timed out for page %d", actual_page)
continue
except subprocess.CalledProcessError as e:
logger.error(
"Zone OCR: pdftoppm failed for page %d: %s",
actual_page,
e.stderr.decode(errors="replace") if e.stderr else str(e),
)
continue
except FileNotFoundError:
logger.error("Zone OCR: pdftoppm not found — is poppler-utils installed?")
return result # No point trying other pages
# pdftoppm names output as prefix-NNNN.png
rendered = sorted(tmp_dir.glob(f"page_{actual_page}-*.png"))
if rendered:
result[actual_page] = rendered[0]
return result
def _crop_zone(
page_img: Path,
zone: OcrTemplateZone,
source_width: int,
source_height: int,
tmp_dir: Path,
) -> Image.Image | None:
"""Crop a zone from the page image and return the PIL Image."""
try:
with Image.open(page_img) as img:
img_width, img_height = img.size
scale_x = img_width / source_width
scale_y = img_height / source_height
crop_left = int(zone.x * scale_x)
crop_top = int(zone.y * scale_y)
crop_right = int((zone.x + zone.width) * scale_x)
crop_bottom = int((zone.y + zone.height) * scale_y)
# Clamp to the image so an oversized zone can't crop out of bounds.
crop_left = max(0, min(crop_left, img_width))
crop_top = max(0, min(crop_top, img_height))
crop_right = max(crop_left + 1, min(crop_right, img_width))
crop_bottom = max(crop_top + 1, min(crop_bottom, img_height))
if crop_right - crop_left < 2 or crop_bottom - crop_top < 2:
logger.warning("Zone OCR: crop too small for zone '%s'", zone.name)
return None
return img.crop((crop_left, crop_top, crop_right, crop_bottom)).copy()
except Exception:
logger.exception("Zone OCR: crop failed for zone '%s'", zone.name)
return None
def _read_barcode(cropped: Image.Image, zone_name: str) -> str | None:
"""Read QR/barcode from a cropped image using zxingcpp."""
try:
import zxingcpp
results = zxingcpp.read_barcodes(cropped)
if results:
text = results[0].text
logger.debug(
"Zone OCR: barcode found in zone '%s': %s",
zone_name,
text[:100],
)
return text
logger.debug("Zone OCR: no barcode found in zone '%s'", zone_name)
return None
except ImportError:
logger.error("Zone OCR: zxingcpp not available — install zxing-cpp")
return None
except Exception:
logger.exception("Zone OCR: barcode read failed for zone '%s'", zone_name)
return None
def _ocr_text(cropped: Image.Image, zone: OcrTemplateZone, tmp_dir: Path) -> str | None:
"""OCR a cropped image with Tesseract."""
crop_path = tmp_dir / f"zone_{zone.pk}.png"
cropped.save(crop_path)
try:
proc = subprocess.run(
[
"tesseract",
str(crop_path),
"stdout",
"-l",
zone.ocr_language,
"--psm",
"6", # Assume uniform block of text
],
capture_output=True,
text=True,
timeout=30,
check=True,
)
return proc.stdout.strip() or None
except subprocess.TimeoutExpired:
logger.error("Zone OCR: Tesseract timed out for zone '%s'", zone.name)
return None
except subprocess.CalledProcessError as e:
logger.error(
"Zone OCR: Tesseract failed for zone '%s': %s",
zone.name,
e.stderr[:200] if e.stderr else str(e),
)
return None
except FileNotFoundError:
logger.error("Zone OCR: Tesseract not found — is tesseract-ocr installed?")
return None
def _extract_zone(
page_img: Path,
zone: OcrTemplateZone,
source_width: int,
source_height: int,
tmp_dir: Path,
) -> str | None:
"""Crop a zone from the page image and extract text via OCR or barcode reader."""
cropped = _crop_zone(page_img, zone, source_width, source_height, tmp_dir)
if cropped is None:
return None
# QR/barcode zones skip Tesseract entirely
if zone.transform == "qr_code":
text = _read_barcode(cropped, zone.name)
if not text:
return None
return _apply_transform(
text,
zone.transform,
getattr(zone, "date_format", "") or "",
)
text = _ocr_text(cropped, zone, tmp_dir)
if not text:
return None
return _apply_transform(
text,
zone.transform,
getattr(zone, "date_format", "") or "",
)
def extract_zone_preview(
doc_path: Path,
zone: OcrTemplateZone,
source_width: int,
source_height: int,
page_count: int | None,
) -> dict:
"""Non-destructive single-zone extraction for the editor's per-zone test.
Renders the zone's page, crops it, runs OCR (or the barcode reader) and
applies the transform WITHOUT writing any custom field. Returns the raw
OCR text and the transformed value so the user can see what the zone yields
(and tune the validation regex) before saving.
"""
# zone.page is 1-indexed (1 = first, -1 = last); resolve to a 0-indexed
# image index exactly like the production extraction path does.
page_idx = _resolve_page_idx(zone.page, page_count)
with tempfile.TemporaryDirectory(dir=settings.SCRATCH_DIR) as tmp_dir:
tmp_path = Path(tmp_dir)
page_images = _render_pages(doc_path, {page_idx}, tmp_path, page_count)
if page_idx not in page_images:
return {"raw_text": None, "value": None}
if not source_width or not source_height:
with Image.open(page_images[page_idx]) as im:
source_width, source_height = im.size
cropped = _crop_zone(
page_images[page_idx],
zone,
source_width,
source_height,
tmp_path,
)
if cropped is None:
return {"raw_text": None, "value": None}
if zone.transform == "qr_code":
raw_text = _read_barcode(cropped, zone.name)
else:
raw_text = _ocr_text(cropped, zone, tmp_path)
value = (
_apply_transform(
raw_text,
zone.transform,
getattr(zone, "date_format", "") or "",
)
if raw_text
else None
)
return {"raw_text": raw_text, "value": value}
def _parse_date(text: str, fmt: str) -> str:
"""Parse a date from OCR text. With a Python strptime `fmt`, try that first;
otherwise (or on failure) fall back to dateparser auto-detection. Returns an
ISO date string, or the original text if nothing parses."""
text = text.strip()
if not text:
return text
if fmt:
try:
return datetime.strptime(text, fmt).date().isoformat()
except ValueError:
pass
try:
import dateparser
parsed = dateparser.parse(
text,
settings={
"PREFER_DAY_OF_MONTH": "first",
"RETURN_AS_TIMEZONE_AWARE": False,
},
)
if parsed:
return parsed.date().isoformat()
except Exception:
logger.debug("Zone OCR: dateparser failed for %r", text[:50])
return text
def _apply_transform(text: str, transform: str, date_format: str = "") -> str:
"""Apply post-processing transform to extracted text."""
text = text.strip()
if not text:
return text
if transform in ("strip", "none"):
return text
elif transform == "date":
return _parse_date(text, date_format)
elif transform == "uppercase":
return text.upper()
elif transform == "lowercase":
return text.lower()
elif transform == "numeric":
result = re.sub(r"[^\d.,\-]", "", text)
return result if result else text
elif transform == "strip_punctuation":
return text.strip(string.punctuation + " \t\r\n")
elif transform == "qr_code":
# Barcode/QR content as read by _read_barcode.
return text
return text
def _zone_target_label(zone: OcrTemplateZone) -> str:
"""Human label of a zone's write target (for logging)."""
target = getattr(zone, "target", None) or "custom_field"
if target == "custom_field":
return zone.custom_field.name if zone.custom_field_id else "(no field)"
return {"title": "Title", "asn": "ASN", "created": "Created"}.get(target, target)
def _parse_created_datetime(value: str):
"""Parse an extracted value into a tz-aware datetime for document.created.
Prefers an ISO date (the zone should use a date transform); falls back to
dateparser. Returns None if no date can be parsed.
"""
from django.utils import timezone as djtz
m = re.search(r"(\d{4})-(\d{2})-(\d{2})", value)
if m:
try:
dt = datetime(int(m[1]), int(m[2]), int(m[3]))
return djtz.make_aware(dt) if djtz.is_naive(dt) else dt
except ValueError:
pass
try:
import dateparser
parsed = dateparser.parse(
value,
settings={"RETURN_AS_TIMEZONE_AWARE": False},
)
if parsed:
return djtz.make_aware(parsed) if djtz.is_naive(parsed) else parsed
except Exception:
logger.debug("Zone OCR: dateparser failed for created value %r", value[:50])
return None
def _write_zone_value(
document: Document,
zone: OcrTemplateZone,
value: str,
) -> None:
"""Write an extracted value to the zone's target — a custom field, or a
built-in document field (title / archive_serial_number / created)."""
target = getattr(zone, "target", None) or "custom_field"
if target == "custom_field":
if zone.custom_field_id:
_write_custom_field(document, zone.custom_field, value)
else:
logger.debug("Zone OCR: zone '%s' has no custom field set", zone.name)
return
if target == "title":
document.title = value[:128]
document.save(update_fields=["title"])
elif target == "asn":
digits = re.sub(r"[^\d]", "", value)
if not digits:
logger.debug(
"Zone OCR: ASN zone '%s' produced no digits (%r)",
zone.name,
value[:50],
)
return
document.archive_serial_number = int(digits)
document.save(update_fields=["archive_serial_number"])
elif target == "created":
parsed = _parse_created_datetime(value)
if parsed is None:
logger.debug(
"Zone OCR: created zone '%s' could not parse a date (%r)",
zone.name,
value[:50],
)
return
document.created = parsed
document.save(update_fields=["created"])
def _write_custom_field(
document: Document,
custom_field: CustomField,
value: str,
) -> None:
"""Write an extracted value to a document's custom field."""
typed_value = _convert_value(value, custom_field.data_type)
if typed_value is None:
logger.debug(
"Zone OCR: skipping custom field '%s' — value conversion returned None",
custom_field.name,
)
return
value_field_name = CustomFieldInstance.get_value_field_name(custom_field.data_type)
CustomFieldInstance.objects.update_or_create(
document=document,
field=custom_field,
defaults={value_field_name: typed_value},
)
def _convert_value(value: str, data_type: str) -> object | None:
"""Convert an extracted OCR string to the appropriate type for the custom field."""
if not value:
return None
try:
if data_type in (
CustomField.FieldDataType.STRING,
CustomField.FieldDataType.URL,
):
return value[:128]
elif data_type == CustomField.FieldDataType.LONG_TEXT:
return value
elif data_type == CustomField.FieldDataType.INT:
digits = re.sub(r"[^\d\-]", "", value)
# Handle edge case: only dashes or empty
digits = digits.lstrip("-") or ""
if not digits:
return None
# Restore leading minus if original had one
if value.strip().startswith("-"):
digits = "-" + digits
return int(digits)
elif data_type == CustomField.FieldDataType.FLOAT:
# Handle European format: 1.234,56 → 1234.56
cleaned = re.sub(r"[^\d.,\-]", "", value)
if not cleaned or cleaned in (".", ",", "-"):
return None
# If both . and , present, the last one is the decimal separator
if "," in cleaned and "." in cleaned:
if cleaned.rindex(",") > cleaned.rindex("."):
# European: 1.234,56
cleaned = cleaned.replace(".", "").replace(",", ".")
else:
# US: 1,234.56
cleaned = cleaned.replace(",", "")
elif "," in cleaned:
# Only comma — treat as decimal separator
cleaned = cleaned.replace(",", ".")
return float(cleaned)
elif data_type == CustomField.FieldDataType.DATE:
match = re.search(r"(\d{4})-(\d{2})-(\d{2})", value)
if match:
y, m, d = match.groups()
# Validate the date
date(int(y), int(m), int(d))
return f"{y}-{m}-{d}"
return None
elif data_type == CustomField.FieldDataType.MONETARY:
cleaned = re.sub(r"[^\d.,\-]", "", value)
if not cleaned or cleaned in (".", ",", "-"):
return None
if "," in cleaned and "." in cleaned:
if cleaned.rindex(",") > cleaned.rindex("."):
cleaned = cleaned.replace(".", "").replace(",", ".")
else:
cleaned = cleaned.replace(",", "")
elif "," in cleaned:
cleaned = cleaned.replace(",", ".")
# Validate it parses as a number
float(cleaned)
return cleaned
elif data_type == CustomField.FieldDataType.BOOL:
lower = value.lower().strip()
if lower in ("true", "yes", "1", "ja", "oui", "si", "x"):
return True
elif lower in ("false", "no", "0", "nein", "non"):
return False
return None
else:
# Unsupported types (DOCUMENTLINK, SELECT) — can't OCR into these
logger.debug(
"Zone OCR: unsupported custom field type %s for OCR extraction",
data_type,
)
return None
except (ValueError, TypeError) as e:
logger.warning("Zone OCR: could not convert %r to %s: %s", value, data_type, e)
return None
def _detect_mime(path: Path) -> str | None:
"""Detect MIME type of a file."""
try:
import magic
return magic.from_file(str(path), mime=True)
except ImportError:
pass
except Exception:
logger.debug("Zone OCR: magic failed for %s, falling back to extension", path)
suffix = path.suffix.lower()
return {
".pdf": "application/pdf",
".png": "image/png",
".jpg": "image/jpeg",
".jpeg": "image/jpeg",
".tiff": "image/tiff",
".tif": "image/tiff",
".webp": "image/webp",
".bmp": "image/bmp",
".gif": "image/gif",
}.get(suffix)
File diff suppressed because it is too large Load Diff
+1
View File
@@ -150,6 +150,7 @@ INSTALLED_APPS = [
"drf_spectacular",
"drf_spectacular_sidecar",
"treenode",
"paperless_benchmark.apps.PaperlessBenchmarkConfig",
*env_apps,
]
-2
View File
@@ -29,7 +29,6 @@ from documents.views import IndexView
from documents.views import LogViewSet
from documents.views import MergeDocumentsAsVersionsView
from documents.views import MergeDocumentsView
from documents.views import OcrTemplateViewSet
from documents.views import PostDocumentView
from documents.views import RemoteVersionView
from documents.views import RemovePasswordDocumentsView
@@ -88,7 +87,6 @@ api_router.register(r"workflow_triggers", WorkflowTriggerViewSet)
api_router.register(r"workflow_actions", WorkflowActionViewSet)
api_router.register(r"workflows", WorkflowViewSet)
api_router.register(r"custom_fields", CustomFieldViewSet)
api_router.register(r"ocr_templates", OcrTemplateViewSet)
api_router.register(r"config", ApplicationConfigurationViewSet)
api_router.register(r"processed_mail", ProcessedMailViewSet)
+47 -97
View File
@@ -4,24 +4,21 @@ from django.conf import settings
from django.contrib.auth.models import User
from documents.models import Document
from documents.permissions import permitted_object_ids
from documents.permissions import restrict_queryset_to_visible
from documents.permissions import user_is_unrestricted
from documents.permissions import get_objects_for_user_owner_aware
from paperless.config import AIConfig
from paperless_ai.base_model import ClassificationSuggestions
from paperless_ai.base_model import TaxonomyChoiceDict
from paperless_ai.base_model import classification_suggestions_to_model
from paperless_ai.client import AIClient
from paperless_ai.db import db_connection_released
from paperless_ai.indexing import _node_document_ids
from paperless_ai.indexing import retrieve_similar_nodes
from paperless_ai.indexing import truncate_content
from paperless_ai.prompts.context import ClassificationPromptContext
from paperless_ai.prompts.context import LocalizationPromptContext
from paperless_ai.prompts.context import RagContextPromptContext
from paperless_ai.prompts.render import render_prompt
from paperless_ai.taxonomy import SimilarDocument
from paperless_ai.taxonomy import TaxonomyCandidates
from paperless_ai.taxonomy import _node_document_weights
from paperless_ai.taxonomy import build_taxonomy_candidates
from paperless_ai.taxonomy import empty_taxonomy_candidates
from paperless_ai.taxonomy import format_taxonomy_for_prompt
@@ -40,48 +37,6 @@ logger = logging.getLogger("paperless_ai.rag_classifier")
TAXONOMY_CANDIDATE_TOP_K = 15
def _fulltext_similar_documents(
document: Document,
user: User | None,
top_k: int,
) -> list[SimilarDocument]:
"""Rank-based fallback when no embedding backend is configured. Uses
Tantivy's "More Like This" (term-overlap similarity) instead of vector
similarity - cruder, but far better than no candidates at all.
more_like_this_ids returns only a ranked ID list, no scores, so weight is
synthesized from rank (descending from top_k) rather than claiming a
similarity magnitude that doesn't exist. An unrestricted user (none, or an
active superuser - see user_is_unrestricted) is normalized to ``None``
before calling, since the backend's permission filter has no superuser
short-circuit of its own. Results are re-checked with
restrict_queryset_to_visible() since Tantivy's indexed permission fields
lag the DB via async reindexing.
"""
from documents.search import get_backend
unrestricted = user_is_unrestricted(user)
search_user = None if unrestricted else user
backend = get_backend()
similar_ids = backend.more_like_this_ids(
document.pk,
user=search_user,
limit=top_k,
)
if not unrestricted:
allowed_ids = set(
restrict_queryset_to_visible(
Document.objects.filter(pk__in=similar_ids),
user,
"view_document",
).values_list("pk", flat=True),
)
similar_ids = [doc_id for doc_id in similar_ids if doc_id in allowed_ids]
return [
SimilarDocument(document_id=doc_id, weight=float(top_k - rank))
for rank, doc_id in enumerate(similar_ids)
]
def get_language_name(language_code: str) -> str:
normalized_language_code = language_code.lower()
for code, name in settings.LANGUAGES:
@@ -181,52 +136,43 @@ def get_taxonomy_context(
user: User | None = None,
max_docs: int = 5,
) -> tuple[TaxonomyCandidates, str]:
"""One retrieval feeds both taxonomy candidates and RAG text context. Uses
vector similarity when an embedding backend is configured, otherwise
falls back to Tantivy full-text "More Like This" similarity - see
_fulltext_similar_documents. On any retrieval failure, degrades to empty
candidates/context rather than propagating the exception - neither a
vector-store outage nor a search-index issue should block classification,
only its context-assisted enrichment.
"""One retrieval feeds both taxonomy candidates and RAG text context.
On any retrieval failure, degrades to empty candidates/context rather than
propagating the exception - a vector-store outage should not block
classification, only its RAG-assisted enrichment.
"""
ai_config = AIConfig()
try:
if ai_config.llm_embedding_backend:
# None means "no restriction" to retrieve_similar_nodes. An
# unrestricted user (no user at all, or an active superuser -- see
# user_is_unrestricted) can see every document, so skip
# materializing every visible pk into a Python list and passing it
# through as an IN filter: for a large library that is a wasted
# quadratic scan in the vector store at best, and past ~32,763
# documents a hard sqlite3.OperationalError (SQLite's
# bound-parameter limit) at worst.
# permitted_object_ids() has its own superuser shortcut that would
# return every Document's id anyway, so this changes nothing about
# which documents are considered -- only how we get there.
visible_document_ids = (
None
if user_is_unrestricted(user)
else list(permitted_object_ids(user, Document, "view_document"))
)
nodes = retrieve_similar_nodes(
document,
top_k=TAXONOMY_CANDIDATE_TOP_K,
document_ids=visible_document_ids,
)
similar_documents = _node_document_weights(nodes)
else:
# See _fulltext_similar_documents: it applies its own permission
# filter via `user`, so no visible-document-id list is needed here.
similar_documents = _fulltext_similar_documents(
document,
user,
top_k=TAXONOMY_CANDIDATE_TOP_K,
# None means "no restriction" to retrieve_similar_nodes. A superuser
# (like no user at all) can see every document, so skip materializing
# every visible pk into a Python list and passing it through as an IN
# filter: for a large library that is a wasted quadratic scan in the
# vector store at best, and past ~32,763 documents a hard
# sqlite3.OperationalError (SQLite's bound-parameter limit) at worst.
# get_objects_for_user_owner_aware() would return every Document for a
# superuser anyway (guardian's own with_superuser shortcut), so this
# changes nothing about which documents are considered -- only how we
# get there.
visible_document_ids = (
None
if user is None or user.is_superuser
else list(
get_objects_for_user_owner_aware(
user,
"view_document",
Document,
).values_list("pk", flat=True),
)
)
nodes = retrieve_similar_nodes(
document,
top_k=TAXONOMY_CANDIDATE_TOP_K,
document_ids=visible_document_ids,
)
candidates = build_taxonomy_candidates(similar_documents, user)
candidates = build_taxonomy_candidates(nodes, user)
# similar_documents is already ordered by descending weight; don't lose it.
similar_document_ids = [s["document_id"] for s in similar_documents]
# ``nodes`` are already ordered by descending vector similarity; don't lose it.
similar_document_ids = list(dict.fromkeys(_node_document_ids(nodes)))
similar_documents_by_id = Document.objects.in_bulk(similar_document_ids)
similar_docs = [
similar_documents_by_id[document_id]
@@ -240,8 +186,8 @@ def get_taxonomy_context(
context_blocks.append(f"TITLE: {title}\n{text}")
except Exception:
logger.exception(
"Failed to retrieve similar-document context for document %s; "
"continuing without taxonomy candidates or similar-document context.",
"Failed to retrieve RAG neighbours for document %s; continuing "
"without taxonomy candidates or similar-document context.",
document.pk,
)
return empty_taxonomy_candidates(), ""
@@ -295,13 +241,17 @@ def get_ai_document_classification(
) -> ClassificationSuggestions:
ai_config = AIConfig()
candidates, context = get_taxonomy_context(document, user)
prompt = build_prompt_with_rag(
document,
ai_config,
candidates=candidates,
context=context,
)
if ai_config.llm_embedding_backend:
candidates, context = get_taxonomy_context(document, user)
prompt = build_prompt_with_rag(
document,
ai_config,
candidates=candidates,
context=context,
)
else:
candidates = empty_taxonomy_candidates()
prompt = build_prompt_without_rag(document, ai_config, candidates=candidates)
client = AIClient()
# Hand the pooled DB connection back while the (slow) LLM query runs so it
+3 -19
View File
@@ -22,7 +22,6 @@ from paperless.network import validate_outbound_http_url
from paperless_ai.base_model import ClassificationSuggestions
from paperless_ai.base_model import DocumentClassifierSchema
from paperless_ai.base_model import model_to_classification_suggestions
from paperless_ai.exceptions import LLMProviderError
from paperless_ai.exceptions import LLMTimeoutError
logger = logging.getLogger("paperless_ai.client")
@@ -133,7 +132,7 @@ class AIClient:
from llama_index.core.llms import ChatMessage
if self.settings.llm_backend == LLMBackend.OLLAMA:
with self._normalize_errors():
with self._normalize_timeouts():
result = self.llm.chat(
[ChatMessage(role="user", content=prompt)],
format=DocumentClassifierSchema.model_json_schema(),
@@ -154,7 +153,7 @@ class AIClient:
content=f"{prompt}\n\n"
f"Answer by calling the {tool.metadata.name} tool. Do not write the answer as text.",
)
with self._normalize_errors():
with self._normalize_timeouts():
result = self.llm.chat_with_tools(
tools=[tool],
user_msg=user_msg,
@@ -174,7 +173,7 @@ class AIClient:
)
@contextmanager
def _normalize_errors(self) -> Iterator[None]:
def _normalize_timeouts(self) -> Iterator[None]:
try:
yield
except httpx.TimeoutException as exc:
@@ -182,23 +181,8 @@ class AIClient:
except Exception as exc:
if self._is_openai_timeout(exc):
raise LLMTimeoutError from exc
if self._is_provider_error(exc):
raise LLMProviderError from exc
raise
def _is_provider_error(self, exc: Exception) -> bool:
if self.settings.llm_backend == LLMBackend.OLLAMA:
from ollama import ResponseError
return isinstance(exc, ResponseError)
if self.settings.llm_backend == LLMBackend.OPENAI_LIKE:
from openai import APIStatusError
return isinstance(exc, APIStatusError)
return False
def _is_openai_timeout(self, exc: Exception) -> bool:
if self.settings.llm_backend != LLMBackend.OPENAI_LIKE:
return False
-4
View File
@@ -1,6 +1,2 @@
class LLMTimeoutError(Exception):
pass
class LLMProviderError(Exception):
"""The LLM backend rejected the request."""
+17
View File
@@ -721,3 +721,20 @@ def retrieve_similar_nodes(
continue
filtered.append(node)
return filtered
def _node_document_ids(nodes: list["NodeWithScore"]) -> list[int]:
document_ids: list[int] = []
for node in nodes:
document_id = node.metadata.get("document_id")
if document_id is None: # pragma: no cover
# See the matching guard in retrieve_similar_nodes() above.
continue
try:
document_ids.append(int(document_id))
except ValueError: # pragma: no cover
logger.warning(
"Skipping LLM index result with invalid document_id %r.",
document_id,
)
return document_ids
+14 -31
View File
@@ -31,11 +31,6 @@ class TaxonomyCandidate(TypedDict):
weight: float
class SimilarDocument(TypedDict):
document_id: int
weight: float
class TaxonomyCandidates(TypedDict):
tags: list[TaxonomyCandidate]
document_types: list[TaxonomyCandidate]
@@ -54,10 +49,10 @@ def empty_taxonomy_candidates() -> TaxonomyCandidates:
)
def _node_document_weights(nodes: list["NodeWithScore"]) -> list[SimilarDocument]:
"""Sum each node's similarity score into its document_id (a document can
appear via multiple chunks/nodes) and return one SimilarDocument per
distinct document_id."""
def _node_document_weights(nodes: list["NodeWithScore"]) -> dict[int, float]:
"""document_id -> that node's similarity score, summed if a document_id
appears more than once across the retrieved nodes (e.g. multiple chunks
of the same source document)."""
weights: dict[int, float] = defaultdict(float)
for node in nodes:
document_id = node.metadata.get("document_id")
@@ -70,14 +65,7 @@ def _node_document_weights(nodes: list["NodeWithScore"]) -> list[SimilarDocument
weights[int(document_id)] += float(node.score or 0.0)
except (TypeError, ValueError): # pragma: no cover
continue
return sorted(
(
SimilarDocument(document_id=document_id, weight=weight)
for document_id, weight in weights.items()
),
key=lambda similar: similar["weight"],
reverse=True,
)
return weights
def _visible_ranked_candidates(
@@ -113,25 +101,20 @@ def _visible_ranked_candidates(
def build_taxonomy_candidates(
similar_documents: list[SimilarDocument],
nodes: list["NodeWithScore"],
user: User | None,
) -> TaxonomyCandidates:
"""Resolve each similar document's id to a live Document, read its
*current* tags/type/correspondent/storage_path via the ORM (never any
possibly-stale names an adapter's source might have cached), weight each
distinct taxonomy object by aggregate similarity weight, permission-filter
"""Resolve each neighbour node's document_id to a live Document, read its
*current* tags/type/correspondent/storage_path via the ORM (never the
possibly-stale names cached in vector-index node metadata), weight each
distinct taxonomy object by aggregate neighbour similarity, permission-filter
against what ``user`` can see, and return each category ranked by weight
and capped. ``similar_documents`` may come from either the vector-RAG
adapter or the full-text fallback adapter - both produce this same shape.
and capped.
"""
if not similar_documents:
return empty_taxonomy_candidates()
# Both adapters guarantee at most one SimilarDocument per document_id, so
# this never silently drops a duplicate's weight.
document_weights: dict[int, float] = {
s["document_id"]: s["weight"] for s in similar_documents
}
document_weights = _node_document_weights(nodes)
if not document_weights:
return empty_taxonomy_candidates()
# Only .tags.all() needs prefetching (a reverse M2M, one extra query for
# the whole batch). document_type/correspondent/storage_path are read
+22 -306
View File
@@ -1,5 +1,4 @@
import datetime
from collections.abc import Generator
from types import SimpleNamespace
from unittest.mock import MagicMock
from unittest.mock import patch
@@ -7,24 +6,18 @@ from unittest.mock import patch
import pytest
import pytest_mock
from django.test import override_settings
from guardian.shortcuts import assign_perm
from guardian.shortcuts import remove_perm
from documents.models import Document
from documents.search import TantivyBackend
from documents.tests.factories import DocumentFactory
from documents.tests.factories import TagFactory
from documents.tests.factories import UserFactory
from paperless.config import AIConfig
from paperless_ai.ai_classifier import TAXONOMY_CANDIDATE_TOP_K
from paperless_ai.ai_classifier import _fulltext_similar_documents
from paperless_ai.ai_classifier import build_localization_prompt
from paperless_ai.ai_classifier import build_prompt_with_rag
from paperless_ai.ai_classifier import build_prompt_without_rag
from paperless_ai.ai_classifier import get_ai_document_classification
from paperless_ai.ai_classifier import get_language_name
from paperless_ai.ai_classifier import get_taxonomy_context
from paperless_ai.taxonomy import SimilarDocument
from paperless_ai.taxonomy import TaxonomyCandidate
from paperless_ai.taxonomy import TaxonomyCandidates
@@ -227,10 +220,12 @@ def test_use_rag_if_configured(
@pytest.mark.django_db
@patch("paperless_ai.client.AIClient.run_llm_query")
@patch("paperless_ai.ai_classifier.build_prompt_with_rag")
@patch("paperless_ai.ai_classifier.build_prompt_without_rag")
@patch("paperless_ai.ai_classifier.AIConfig")
@override_settings(LLM_BACKEND="ollama", LLM_MODEL="some_model")
def test_use_rag_prompt_even_without_embedding_backend(
mock_build_prompt_with_rag,
def test_use_without_rag_if_not_configured(
mock_ai_config,
mock_build_prompt_without_rag,
mock_run_llm_query,
mock_document,
):
@@ -240,13 +235,13 @@ def test_use_rag_prompt_even_without_embedding_backend(
WHEN:
- get_ai_document_classification() is called
THEN:
- The RAG-context prompt builder is still used (fed by the full-text
fallback's context/candidates instead of the vector store's)
- The non-RAG prompt builder is used
"""
mock_build_prompt_with_rag.return_value = "Prompt with RAG"
mock_ai_config.return_value.llm_embedding_backend = None
mock_build_prompt_without_rag.return_value = "Prompt without RAG"
mock_run_llm_query.return_value = NESTED_SUGGESTIONS
get_ai_document_classification(mock_document)
mock_build_prompt_with_rag.assert_called_once()
mock_build_prompt_without_rag.assert_called_once()
@pytest.mark.django_db
@@ -325,7 +320,6 @@ def test_build_localization_prompt_preserves_unicode_characters():
@pytest.mark.django_db
@override_settings(LLM_EMBEDDING_BACKEND="huggingface")
def test_get_taxonomy_context_assembles_rag_text_and_candidates():
"""
GIVEN:
@@ -360,7 +354,6 @@ def test_get_taxonomy_context_assembles_rag_text_and_candidates():
@pytest.mark.django_db
@override_settings(LLM_EMBEDDING_BACKEND="huggingface")
def test_get_taxonomy_context_preserves_similarity_order_and_distinct_documents():
"""
GIVEN:
@@ -431,7 +424,6 @@ def test_get_taxonomy_context_preserves_similarity_order_and_distinct_documents(
@pytest.mark.django_db
@override_settings(LLM_EMBEDDING_BACKEND="huggingface")
def test_get_taxonomy_context_no_similar_docs():
"""
GIVEN:
@@ -455,67 +447,6 @@ def test_get_taxonomy_context_no_similar_docs():
}
@pytest.mark.django_db
def test_get_taxonomy_context_uses_fulltext_fallback_when_no_embedding_backend(
mocker: pytest_mock.MockerFixture,
) -> None:
"""
GIVEN:
- No LLM embedding backend is configured (the default test settings)
WHEN:
- get_taxonomy_context() is called
THEN:
- _fulltext_similar_documents() is called with the document, the user
and TAXONOMY_CANDIDATE_TOP_K
- retrieve_similar_nodes() (the vector path) is never called
"""
document = DocumentFactory.create(content="Some content")
mock_fulltext = mocker.patch(
"paperless_ai.ai_classifier._fulltext_similar_documents",
return_value=[],
)
mock_retrieve = mocker.patch("paperless_ai.ai_classifier.retrieve_similar_nodes")
get_taxonomy_context(document, user=None)
mock_fulltext.assert_called_once_with(
document,
None,
top_k=TAXONOMY_CANDIDATE_TOP_K,
)
mock_retrieve.assert_not_called()
@pytest.mark.django_db
@override_settings(LLM_EMBEDDING_BACKEND="huggingface")
def test_get_taxonomy_context_uses_vector_path_when_embedding_backend_configured(
mocker: pytest_mock.MockerFixture,
) -> None:
"""
GIVEN:
- An LLM embedding backend is configured
WHEN:
- get_taxonomy_context() is called
THEN:
- retrieve_similar_nodes() (the vector path) is called
- _fulltext_similar_documents() (the no-embedding-backend fallback)
is never called
"""
document = DocumentFactory.create(content="Some content")
mock_retrieve = mocker.patch(
"paperless_ai.ai_classifier.retrieve_similar_nodes",
return_value=[],
)
mock_fulltext = mocker.patch(
"paperless_ai.ai_classifier._fulltext_similar_documents",
)
get_taxonomy_context(document, user=None)
mock_retrieve.assert_called_once()
mock_fulltext.assert_not_called()
class TestGetTaxonomyContextVisibility:
"""get_taxonomy_context must not materialize every visible document id
for a user who can already see the whole library: a superuser (like no
@@ -528,7 +459,6 @@ class TestGetTaxonomyContextVisibility:
"""
@pytest.mark.django_db
@override_settings(LLM_EMBEDDING_BACKEND="huggingface")
def test_skips_permission_lookup_for_superuser(
self,
mocker: pytest_mock.MockerFixture,
@@ -547,18 +477,17 @@ class TestGetTaxonomyContextVisibility:
"paperless_ai.ai_classifier.retrieve_similar_nodes",
return_value=[],
)
mock_permitted = mocker.patch(
"paperless_ai.ai_classifier.permitted_object_ids",
mock_get_objects = mocker.patch(
"paperless_ai.ai_classifier.get_objects_for_user_owner_aware",
)
user = UserFactory.create(is_superuser=True)
get_taxonomy_context(document, user)
mock_permitted.assert_not_called()
mock_get_objects.assert_not_called()
assert mock_retrieve.call_args.kwargs["document_ids"] is None
@pytest.mark.django_db
@override_settings(LLM_EMBEDDING_BACKEND="huggingface")
def test_skips_permission_lookup_when_no_user(
self,
mocker: pytest_mock.MockerFixture,
@@ -577,17 +506,16 @@ class TestGetTaxonomyContextVisibility:
"paperless_ai.ai_classifier.retrieve_similar_nodes",
return_value=[],
)
mock_permitted = mocker.patch(
"paperless_ai.ai_classifier.permitted_object_ids",
mock_get_objects = mocker.patch(
"paperless_ai.ai_classifier.get_objects_for_user_owner_aware",
)
get_taxonomy_context(document, None)
mock_permitted.assert_not_called()
mock_get_objects.assert_not_called()
assert mock_retrieve.call_args.kwargs["document_ids"] is None
@pytest.mark.django_db
@override_settings(LLM_EMBEDDING_BACKEND="huggingface")
def test_restricts_to_visible_documents_for_non_superuser(
self,
mocker: pytest_mock.MockerFixture,
@@ -598,7 +526,7 @@ class TestGetTaxonomyContextVisibility:
WHEN:
- get_taxonomy_context() is called
THEN:
- The user's permitted document ids are looked up and passed to
- The user's visible document ids are looked up and passed to
retrieve_similar_nodes() as a restriction
"""
document = DocumentFactory.create(content="Some content")
@@ -606,232 +534,21 @@ class TestGetTaxonomyContextVisibility:
"paperless_ai.ai_classifier.retrieve_similar_nodes",
return_value=[],
)
mock_permitted = mocker.patch(
"paperless_ai.ai_classifier.permitted_object_ids",
return_value=[1, 2, 3],
mock_queryset = mocker.MagicMock()
mock_queryset.values_list.return_value = [1, 2, 3]
mock_get_objects = mocker.patch(
"paperless_ai.ai_classifier.get_objects_for_user_owner_aware",
return_value=mock_queryset,
)
user = UserFactory.create(is_superuser=False)
get_taxonomy_context(document, user)
mock_permitted.assert_called_once_with(user, Document, "view_document")
mock_get_objects.assert_called_once_with(user, "view_document", Document)
assert mock_retrieve.call_args.kwargs["document_ids"] == [1, 2, 3]
@pytest.mark.django_db
class TestFulltextSimilarDocuments:
"""_fulltext_similar_documents is the no-embedding-backend fallback: it
asks the Tantivy full-text index for "More Like This" neighbours instead
of the vector store, and synthesizes a rank-based weight since Tantivy's
more_like_this_ids returns only an ordered id list, no scores.
"""
@pytest.fixture
def fulltext_backend(
self,
mocker: pytest_mock.MockerFixture,
) -> Generator[TantivyBackend, None, None]:
"""An in-memory Tantivy backend, wired up as the module-level
singleton _fulltext_similar_documents resolves via get_backend()."""
backend = TantivyBackend(path=None)
backend.open()
mocker.patch("documents.search.get_backend", return_value=backend)
try:
yield backend
finally:
backend.close()
def test_ranks_by_rank_based_weight_descending(
self,
fulltext_backend: TantivyBackend,
) -> None:
"""
GIVEN:
- A source document and two similar documents indexed in Tantivy
WHEN:
- _fulltext_similar_documents() is called
THEN:
- Each result's weight reflects its rank (first result weighted
higher than the second), not a raw similarity score
"""
source = DocumentFactory.create(content="quarterly financial report details")
first = DocumentFactory.create(content="quarterly financial report details")
second = DocumentFactory.create(content="financial report")
for doc in (source, first, second):
fulltext_backend.add_or_update(doc)
result = _fulltext_similar_documents(source, user=None, top_k=5)
assert len(result) == 2
weight_by_id = {s["document_id"]: s["weight"] for s in result}
assert weight_by_id[first.pk] > weight_by_id[second.pk]
def test_excludes_source_document(
self,
fulltext_backend: TantivyBackend,
) -> None:
"""
GIVEN:
- A source document indexed in Tantivy with no other documents
WHEN:
- _fulltext_similar_documents() is called
THEN:
- An empty list is returned - the source document is never its
own similar document
"""
source = DocumentFactory.create(content="unique unrelated content")
fulltext_backend.add_or_update(source)
result = _fulltext_similar_documents(source, user=None, top_k=5)
assert result == []
def test_empty_index_returns_empty_list(
self,
fulltext_backend: TantivyBackend,
) -> None:
"""
GIVEN:
- A document that has never been indexed (fresh/empty Tantivy index)
WHEN:
- _fulltext_similar_documents() is called
THEN:
- An empty list is returned rather than raising
"""
source = DocumentFactory.create(content="never indexed")
result = _fulltext_similar_documents(source, user=None, top_k=5)
assert result == []
def test_respects_top_k_limit(
self,
fulltext_backend: TantivyBackend,
) -> None:
"""
GIVEN:
- A source document and four similar documents indexed
WHEN:
- _fulltext_similar_documents() is called with top_k=2
THEN:
- At most 2 results are returned
"""
source = DocumentFactory.create(content="shared overlapping keyword text")
fulltext_backend.add_or_update(source)
for _ in range(4):
fulltext_backend.add_or_update(
DocumentFactory.create(content="shared overlapping keyword text"),
)
result = _fulltext_similar_documents(source, user=None, top_k=2)
assert len(result) == 2
def test_result_shape_is_similar_document(
self,
fulltext_backend: TantivyBackend,
) -> None:
"""
GIVEN:
- A source document and one similar document indexed
WHEN:
- _fulltext_similar_documents() is called
THEN:
- Each result is a SimilarDocument (document_id + weight only)
"""
source = DocumentFactory.create(content="shared content phrase")
other = DocumentFactory.create(content="shared content phrase")
fulltext_backend.add_or_update(source)
fulltext_backend.add_or_update(other)
result = _fulltext_similar_documents(source, user=None, top_k=5)
# rank 0 (the only/best result) with top_k=5 -> weight = top_k - rank = 5.0,
# per the "first result gets top_k, the last gets 1" formula.
assert result == [SimilarDocument(document_id=other.pk, weight=5.0)]
def test_superuser_sees_other_users_documents(
self,
fulltext_backend: TantivyBackend,
) -> None:
"""
GIVEN:
- A source document owned by one user and a similar document
owned by a different user, with no sharing between them
WHEN:
- _fulltext_similar_documents() is called with a superuser
THEN:
- The other user's document is still returned as a similar
document - a superuser must not be narrowed by the backend's
owner-based permission filter
"""
owner = UserFactory.create()
other_owner = UserFactory.create()
superuser = UserFactory.create(is_superuser=True)
source = DocumentFactory.create(
content="shared content phrase",
owner=owner,
)
other = DocumentFactory.create(
content="shared content phrase",
owner=other_owner,
)
fulltext_backend.add_or_update(source)
fulltext_backend.add_or_update(other)
result = _fulltext_similar_documents(source, user=superuser, top_k=5)
assert [s["document_id"] for s in result] == [other.pk]
def test_excludes_stale_permitted_document_for_regular_user(
self,
fulltext_backend: TantivyBackend,
) -> None:
"""
GIVEN:
- A regular (non-superuser) user
- A similar document the user is permitted to view, and another
similar document indexed while the user still had view
permission but which has since had that permission revoked in
the database, i.e. the Tantivy index has stale permission data
WHEN:
- _fulltext_similar_documents() is called with that user
THEN:
- Only the still-permitted document is returned - the DB
re-check via restrict_queryset_to_visible() must catch the
document Tantivy's stale index still thinks is visible
"""
owner = UserFactory.create()
viewer = UserFactory.create(is_superuser=False)
source = DocumentFactory.create(
content="shared content phrase",
owner=owner,
)
permitted = DocumentFactory.create(
content="shared content phrase",
owner=owner,
)
now_private = DocumentFactory.create(
content="shared content phrase",
owner=owner,
)
assign_perm("view_document", viewer, permitted)
assign_perm("view_document", viewer, now_private)
fulltext_backend.add_or_update(source)
fulltext_backend.add_or_update(permitted)
fulltext_backend.add_or_update(now_private)
# Revoke access after indexing, without reindexing: the index still
# carries viewer as a permitted viewer for `now_private`.
remove_perm("view_document", viewer, now_private)
result = _fulltext_similar_documents(source, user=viewer, top_k=5)
assert [s["document_id"] for s in result] == [permitted.pk]
@pytest.mark.django_db
@override_settings(LLM_EMBEDDING_BACKEND="huggingface")
@patch("paperless_ai.ai_classifier.retrieve_similar_nodes")
def test_get_taxonomy_context_retrieval_failure_degrades_to_no_hints(mock_retrieve):
"""
@@ -858,7 +575,6 @@ def test_get_taxonomy_context_retrieval_failure_degrades_to_no_hints(mock_retrie
@pytest.mark.django_db
@override_settings(LLM_EMBEDDING_BACKEND="huggingface")
@patch("paperless_ai.ai_classifier.build_taxonomy_candidates")
@patch("paperless_ai.ai_classifier.retrieve_similar_nodes")
def test_get_taxonomy_context_candidate_building_failure_degrades_to_no_hints(
+5 -3
View File
@@ -1188,7 +1188,9 @@ class TestRetrieveSimilarNodesAgainstRealIndex:
nodes = indexing.retrieve_similar_nodes(a, document_ids=[b.id])
assert all(int(node.metadata["document_id"]) == b.id for node in nodes)
assert all(
document_id == b.id for document_id in indexing._node_document_ids(nodes)
)
def test_excludes_self(
self,
@@ -1210,7 +1212,7 @@ class TestRetrieveSimilarNodesAgainstRealIndex:
nodes = indexing.retrieve_similar_nodes(a, top_k=5)
assert {int(node.metadata["document_id"]) for node in nodes} == {b.id}
assert set(indexing._node_document_ids(nodes)) == {b.id}
def test_excludes_self_with_multiple_chunks(
self,
@@ -1233,4 +1235,4 @@ class TestRetrieveSimilarNodesAgainstRealIndex:
nodes = indexing.retrieve_similar_nodes(a, top_k=3)
assert {int(node.metadata["document_id"]) for node in nodes} == {b.id}
assert set(indexing._node_document_ids(nodes)) == {b.id}
-48
View File
@@ -4,7 +4,6 @@ from unittest.mock import MagicMock
from unittest.mock import patch
import httpx
import ollama
import openai
import pytest
from llama_index.core.llms.llm import ToolSelection
@@ -12,7 +11,6 @@ from llama_index.core.llms.llm import ToolSelection
from paperless_ai.client import LLM_SYSTEM_PROMPT
from paperless_ai.client import PLACEHOLDER_API_KEY
from paperless_ai.client import AIClient
from paperless_ai.exceptions import LLMProviderError
from paperless_ai.exceptions import LLMTimeoutError
@@ -216,52 +214,6 @@ def test_run_llm_query_openai_timeout_raises_local_error(
client.run_llm_query("test_prompt")
def test_run_llm_query_openai_status_error_raises_provider_error(
mock_ai_config,
mock_openai_llm,
):
mock_ai_config.llm_backend = "openai-like"
mock_ai_config.llm_model = "test_model"
mock_ai_config.llm_endpoint = "http://test-url"
request = httpx.Request("POST", "http://test-url/v1/chat/completions")
body = {"error": {"message": "Thinking mode does not support this tool_choice"}}
mock_openai_llm.return_value.chat_with_tools.side_effect = openai.BadRequestError(
"Error code: 400",
response=httpx.Response(400, request=request, json=body),
body=body,
)
client = AIClient()
with pytest.raises(LLMProviderError) as exc_info:
client.run_llm_query("test_prompt")
assert str(exc_info.value) == ""
assert isinstance(exc_info.value.__cause__, openai.BadRequestError)
def test_run_llm_query_ollama_response_error_raises_provider_error(
mock_ai_config,
mock_ollama_llm,
):
mock_ai_config.llm_backend = "ollama"
mock_ai_config.llm_model = "test_model"
mock_ai_config.llm_endpoint = "http://test-url"
response_error = ollama.ResponseError(
"confidential provider response",
status_code=400,
)
mock_ollama_llm.return_value.chat.side_effect = response_error
client = AIClient()
with pytest.raises(LLMProviderError) as exc_info:
client.run_llm_query("test_prompt")
assert str(exc_info.value) == ""
assert exc_info.value.__cause__ is response_error
def test_run_llm_query_httpx_timeout_raises_local_error(
mock_ai_config,
mock_ollama_llm,
+31 -33
View File
@@ -1,4 +1,5 @@
import json
from types import SimpleNamespace
import pytest
import pytest_mock
@@ -9,14 +10,14 @@ from documents.tests.factories import DocumentTypeFactory
from documents.tests.factories import StoragePathFactory
from documents.tests.factories import TagFactory
from documents.tests.factories import UserFactory
from paperless_ai.taxonomy import SimilarDocument
from paperless_ai.taxonomy import TaxonomyCandidates
from paperless_ai.taxonomy import build_taxonomy_candidates
from paperless_ai.taxonomy import format_taxonomy_for_prompt
def make_similar(document_id: int, weight: float) -> SimilarDocument:
return SimilarDocument(document_id=document_id, weight=weight)
def make_node(document_id: int, score: float) -> SimpleNamespace:
"""A stand-in for NodeWithScore: only ``.metadata``/``.score`` are read."""
return SimpleNamespace(metadata={"document_id": str(document_id)}, score=score)
@pytest.mark.django_db
@@ -52,9 +53,9 @@ class TestBuildTaxonomyCandidates:
doc_a.tags.add(tag)
doc_b = DocumentFactory.create()
doc_b.tags.add(tag)
similar_documents = [make_similar(doc_a.pk, 0.9), make_similar(doc_b.pk, 0.4)]
nodes = [make_node(doc_a.pk, 0.9), make_node(doc_b.pk, 0.4)]
result = build_taxonomy_candidates(similar_documents, user=None)
result = build_taxonomy_candidates(nodes, user=None)
assert len(result["tags"]) == 1
assert result["tags"][0]["id"] == tag.pk
@@ -79,9 +80,9 @@ class TestBuildTaxonomyCandidates:
document.tags.add(tag)
tag.name = "New Name"
tag.save()
similar_documents = [make_similar(document.pk, 0.5)]
nodes = [make_node(document.pk, 0.5)]
result = build_taxonomy_candidates(similar_documents, user=None)
result = build_taxonomy_candidates(nodes, user=None)
assert result["tags"][0]["name"] == "New Name"
@@ -101,9 +102,9 @@ class TestBuildTaxonomyCandidates:
document = DocumentFactory.create()
document.tags.add(tag)
tag.delete()
similar_documents = [make_similar(document.pk, 0.5)]
nodes = [make_node(document.pk, 0.5)]
result = build_taxonomy_candidates(similar_documents, user=None)
result = build_taxonomy_candidates(nodes, user=None)
assert result["tags"] == []
@@ -122,12 +123,9 @@ class TestBuildTaxonomyCandidates:
strong_doc.tags.add(strong_tag)
weak_doc = DocumentFactory.create()
weak_doc.tags.add(weak_tag)
similar_documents = [
make_similar(strong_doc.pk, 0.9),
make_similar(weak_doc.pk, 0.1),
]
nodes = [make_node(strong_doc.pk, 0.9), make_node(weak_doc.pk, 0.1)]
result = build_taxonomy_candidates(similar_documents, user=None)
result = build_taxonomy_candidates(nodes, user=None)
assert [c["name"] for c in result["tags"]] == ["Strong", "Weak"]
@@ -143,9 +141,9 @@ class TestBuildTaxonomyCandidates:
document = DocumentFactory.create()
for i in range(15):
document.tags.add(TagFactory.create(name=f"Tag{i}"))
similar_documents = [make_similar(document.pk, 0.5)]
nodes = [make_node(document.pk, 0.5)]
result = build_taxonomy_candidates(similar_documents, user=None)
result = build_taxonomy_candidates(nodes, user=None)
assert len(result["tags"]) == 10
@@ -159,12 +157,12 @@ class TestBuildTaxonomyCandidates:
- Only 5 correspondents are returned
"""
correspondents = CorrespondentFactory.create_batch(7)
similar_documents = [
make_similar(DocumentFactory.create(correspondent=c).pk, 0.5)
nodes = [
make_node(DocumentFactory.create(correspondent=c).pk, 0.5)
for c in correspondents
]
result = build_taxonomy_candidates(similar_documents, user=None)
result = build_taxonomy_candidates(nodes, user=None)
assert len(result["correspondents"]) == 5
@@ -179,9 +177,9 @@ class TestBuildTaxonomyCandidates:
"""
document_type = DocumentTypeFactory.create(name="Invoice")
document = DocumentFactory.create(document_type=document_type)
similar_documents = [make_similar(document.pk, 0.5)]
nodes = [make_node(document.pk, 0.5)]
result = build_taxonomy_candidates(similar_documents, user=None)
result = build_taxonomy_candidates(nodes, user=None)
assert len(result["document_types"]) == 1
assert result["document_types"][0]["id"] == document_type.pk
@@ -197,12 +195,12 @@ class TestBuildTaxonomyCandidates:
- Only 5 document_types are returned
"""
document_types = DocumentTypeFactory.create_batch(7)
similar_documents = [
make_similar(DocumentFactory.create(document_type=dt).pk, 0.5)
nodes = [
make_node(DocumentFactory.create(document_type=dt).pk, 0.5)
for dt in document_types
]
result = build_taxonomy_candidates(similar_documents, user=None)
result = build_taxonomy_candidates(nodes, user=None)
assert len(result["document_types"]) == 5
@@ -217,9 +215,9 @@ class TestBuildTaxonomyCandidates:
"""
storage_path = StoragePathFactory.create(name="Invoices")
document = DocumentFactory.create(storage_path=storage_path)
similar_documents = [make_similar(document.pk, 0.5)]
nodes = [make_node(document.pk, 0.5)]
result = build_taxonomy_candidates(similar_documents, user=None)
result = build_taxonomy_candidates(nodes, user=None)
assert len(result["storage_paths"]) == 1
assert result["storage_paths"][0]["id"] == storage_path.pk
@@ -235,12 +233,12 @@ class TestBuildTaxonomyCandidates:
- Only 5 storage_paths are returned
"""
storage_paths = StoragePathFactory.create_batch(7)
similar_documents = [
make_similar(DocumentFactory.create(storage_path=sp).pk, 0.5)
nodes = [
make_node(DocumentFactory.create(storage_path=sp).pk, 0.5)
for sp in storage_paths
]
result = build_taxonomy_candidates(similar_documents, user=None)
result = build_taxonomy_candidates(nodes, user=None)
assert len(result["storage_paths"]) == 5
@@ -260,14 +258,14 @@ class TestBuildTaxonomyCandidates:
tag = TagFactory.create(name="Restricted")
document = DocumentFactory.create()
document.tags.add(tag)
similar_documents = [make_similar(document.pk, 0.5)]
nodes = [make_node(document.pk, 0.5)]
user = UserFactory.create()
mocker.patch(
"documents.permissions.permitted_object_ids",
return_value=[], # user cannot see this tag
)
result = build_taxonomy_candidates(similar_documents, user=user)
result = build_taxonomy_candidates(nodes, user=user)
assert result["tags"] == []
@@ -297,10 +295,10 @@ class TestBuildTaxonomyCandidates:
tag.save()
document = DocumentFactory.create()
document.tags.add(tag)
similar_documents = [make_similar(document.pk, 0.5)]
nodes = [make_node(document.pk, 0.5)]
spy = mocker.patch("documents.permissions.permitted_object_ids")
result = build_taxonomy_candidates(similar_documents, user=None)
result = build_taxonomy_candidates(nodes, user=None)
assert result["tags"][0]["name"] == "Owned"
spy.assert_not_called()
View File
+8
View File
@@ -0,0 +1,8 @@
from django.apps import AppConfig
from django.utils.translation import gettext_lazy as _
class PaperlessBenchmarkConfig(AppConfig):
name = "paperless_benchmark"
verbose_name = _("Paperless benchmark")
+136
View File
@@ -0,0 +1,136 @@
from __future__ import annotations
from typing import TYPE_CHECKING
from django.db import connection
if TYPE_CHECKING:
from django.db.models import QuerySet
def _reset_table_names() -> list[str]:
from guardian.models import GroupObjectPermission
from guardian.models import UserObjectPermission
from documents.models import Correspondent
from documents.models import Document
from documents.models import DocumentType
from documents.models import StoragePath
from documents.models import Tag
return [
Document.tags.through._meta.db_table,
Document._meta.db_table,
Tag._meta.db_table,
Correspondent._meta.db_table,
DocumentType._meta.db_table,
StoragePath._meta.db_table,
UserObjectPermission._meta.db_table,
GroupObjectPermission._meta.db_table,
]
def _delete_all_users_and_groups() -> None:
# ASSUMPTION: this tool assumes a disposable benchmark database, never
# point it at a real install. This deletes EVERY user and group in the
# database (not just benchmark-created ones) -- there is no way to
# distinguish "real" users from seeded ones, so this is only safe against
# a database that exists solely to run this benchmarking tool. The
# `benchmark seed --reset` CLI path requires an explicit
# `--yes-i-know-this-wipes-the-database` flag before reaching here; do
# not remove that guard.
from django.contrib.auth import get_user_model
from django.contrib.auth.models import Group
get_user_model().objects.all().delete()
Group.objects.all().delete()
def _reset_postgresql() -> None:
tables = _reset_table_names()
with connection.cursor() as cursor:
cursor.execute(f"TRUNCATE TABLE {', '.join(tables)} RESTART IDENTITY CASCADE;")
_delete_all_users_and_groups()
def _reset_mariadb() -> None:
# MariaDB's TRUNCATE has no CASCADE clause and refuses to truncate a
# table referenced by a foreign key while checks are enabled, so
# checks are disabled for the duration of the reset.
tables = _reset_table_names()
with connection.cursor() as cursor:
cursor.execute("SET FOREIGN_KEY_CHECKS = 0;")
try:
for table in tables:
cursor.execute(f"TRUNCATE TABLE {table};")
finally:
cursor.execute("SET FOREIGN_KEY_CHECKS = 1;")
_delete_all_users_and_groups()
def _reset_sqlite() -> None:
from documents.models import Correspondent
from documents.models import Document
from documents.models import DocumentType
from documents.models import StoragePath
from documents.models import Tag
Document.global_objects.all().delete()
Tag.objects.all().delete()
Correspondent.objects.all().delete()
DocumentType.objects.all().delete()
StoragePath.objects.all().delete()
_delete_all_users_and_groups()
def reset_benchmark_data() -> None:
"""
Remove all previously-seeded benchmark data (documents, tags,
correspondents, document types, storage paths, guardian permission
rows, users, and groups) so a fresh `benchmark seed` run starts
from an empty slate. Dispatches per-backend because TRUNCATE syntax
and cascade behavior differ across the 3 supported databases.
"""
if connection.vendor == "postgresql":
_reset_postgresql()
elif connection.vendor == "mysql":
# MariaDB also reports vendor == "mysql" under Django's mysql backend.
_reset_mariadb()
else:
_reset_sqlite()
def capture_explain(queryset: QuerySet) -> str:
"""
Return the query plan for `queryset` using the current backend's
explain facility. PostgreSQL supports `EXPLAIN ANALYZE {sql}` (real
execution stats). MariaDB does NOT accept that syntax -- verified
against a real MariaDB 12.3 container: `EXPLAIN ANALYZE {sql}` raises a
1064 syntax error, while MariaDB's own `ANALYZE {sql}` form (no
`EXPLAIN` keyword) works and returns real per-row execution stats
(`r_rows`, `r_filtered`, etc. columns) -- this is MariaDB's
EXPLAIN-ANALYZE-equivalent, distinct from MySQL 8.0.18+'s
`EXPLAIN ANALYZE` syntax, which MariaDB does not implement. SQLite only
supports EXPLAIN QUERY PLAN (the chosen plan, not real timing/row
counts) -- that output is clearly labeled rather than silently looking
equivalent to the other two backends' output.
"""
sql, params = queryset.query.sql_with_params()
with connection.cursor() as cursor:
if connection.vendor == "postgresql":
cursor.execute(f"EXPLAIN ANALYZE {sql}", params)
return "\n".join(str(row[0]) for row in cursor.fetchall())
if connection.vendor == "mysql":
# MariaDB also reports vendor == "mysql" under Django's mysql
# backend. Unlike MySQL 8.0.18+, MariaDB has no `EXPLAIN
# ANALYZE` syntax -- its equivalent is `ANALYZE <statement>`.
cursor.execute(f"ANALYZE {sql}", params)
columns = [c[0] for c in cursor.description]
header = " | ".join(columns)
rows = "\n".join(
" | ".join(str(c) for c in row) for row in cursor.fetchall()
)
return f"{header}\n{rows}"
cursor.execute(f"EXPLAIN QUERY PLAN {sql}", params)
rows = "\n".join(" | ".join(str(c) for c in row) for row in cursor.fetchall())
return f"(plan only -- no execution stats on SQLite)\n{rows}"
+81
View File
@@ -0,0 +1,81 @@
# src/paperless_benchmark/endpoints.py
from __future__ import annotations
import statistics
import time
from dataclasses import dataclass
from typing import TYPE_CHECKING
if TYPE_CHECKING:
from django.contrib.auth.models import User
from rest_framework.test import APIClient
ENDPOINTS: tuple[tuple[str, str], ...] = (
("documents_default", "/api/documents/"),
("documents_page50", "/api/documents/?page_size=50"),
("tags_all", "/api/tags/?page_size=100000"),
)
@dataclass(frozen=True, slots=True)
class EndpointTiming:
user_label: str
endpoint_name: str
query_count: int
min_ms: float
median_ms: float
max_ms: float
def _timed_requests(client: APIClient, url: str, n: int) -> list[float]:
times = []
for _ in range(n):
t0 = time.perf_counter()
resp = client.get(url)
t1 = time.perf_counter()
if resp.status_code != 200:
raise RuntimeError(
f"GET {url} -> {resp.status_code}: {resp.content[:300]!r}",
)
times.append(t1 - t0)
return times
def _query_count(client: APIClient, url: str) -> int:
from django.db import connection
from django.test.utils import CaptureQueriesContext
with CaptureQueriesContext(connection) as ctx:
resp = client.get(url)
if resp.status_code != 200:
raise RuntimeError(f"GET {url} -> {resp.status_code}: {resp.content[:300]!r}")
return len(ctx.captured_queries)
def run_endpoint_benchmarks(
*,
perf_target: User,
perf_admin: User,
repeat: int,
) -> list[EndpointTiming]:
from rest_framework.test import APIClient
results: list[EndpointTiming] = []
for user_label, user in (("target", perf_target), ("admin", perf_admin)):
client = APIClient()
client.force_authenticate(user=user)
for name, url in ENDPOINTS:
client.get(url) # warm-up request, not counted
qcount = _query_count(client, url)
times_ms = [t * 1000 for t in _timed_requests(client, url, repeat)]
results.append(
EndpointTiming(
user_label=user_label,
endpoint_name=name,
query_count=qcount,
min_ms=min(times_ms),
median_ms=statistics.median(times_ms),
max_ms=max(times_ms),
),
)
return results
+56
View File
@@ -0,0 +1,56 @@
# src/paperless_benchmark/harness.py
from __future__ import annotations
import time
from dataclasses import dataclass
from typing import TYPE_CHECKING
from typing import Generic
from typing import TypeVar
from django.db import connection
from django.test.utils import CaptureQueriesContext
if TYPE_CHECKING:
from collections.abc import Callable
T = TypeVar("T")
@dataclass(frozen=True, slots=True)
class ProfileResult(Generic[T]):
best_seconds: float
all_seconds: tuple[float, ...]
query_count: int
result: T
def run_profile(fn: Callable[[], T], *, repeat: int = 5) -> ProfileResult[T]:
"""
Call `fn` `repeat` times, capturing wall-clock time for every call and
the SQL query count for the final call. Returns the best (minimum)
time across all repeats, since the first call(s) can be skewed by
connection warm-up or cold caches.
"""
if repeat < 1:
raise ValueError("repeat must be >= 1")
all_seconds: list[float] = []
result: T | None = None
query_count = 0
for i in range(repeat):
with CaptureQueriesContext(connection) as ctx:
start = time.perf_counter()
result = fn()
all_seconds.append(time.perf_counter() - start)
if i == repeat - 1:
query_count = len(ctx.captured_queries)
# Purely a type-narrowing aid for the type checker: the `repeat < 1`
# guard above already turns the one case that could leave `result`
# unset into a clear ValueError, so this is unreachable in practice.
assert result is not None
return ProfileResult(
best_seconds=min(all_seconds),
all_seconds=tuple(all_seconds),
query_count=query_count,
result=result,
)
@@ -0,0 +1,222 @@
# src/paperless_benchmark/management/commands/benchmark.py
from __future__ import annotations
from typing import Any
from django.core.management.base import BaseCommand
from django.core.management.base import CommandError
from django.core.management.base import CommandParser
class Command(BaseCommand):
help = "Seed, run, and profile paperless-ngx performance benchmarks."
def add_arguments(self, parser: CommandParser) -> None:
parser.add_argument(
"action",
choices=["seed", "run", "profile", "list-scenarios"],
help="Which benchmark action to perform.",
)
parser.add_argument(
"scenario",
nargs="?",
default=None,
help="Scenario name (required for `profile`; see `list-scenarios`).",
)
parser.add_argument(
"--tier",
choices=["home", "medium", "large"],
default="medium",
help="Dataset scale tier for `seed` (default: medium).",
)
parser.add_argument(
"--reset",
action="store_true",
default=False,
help="For `seed`: truncate existing benchmark data first.",
)
parser.add_argument(
"--yes-i-know-this-wipes-the-database",
action="store_true",
default=False,
help=(
"Required alongside --reset: confirms you understand `seed "
"--reset` deletes ALL users, ALL groups, and ALL documents/"
"tags/correspondents/document types/storage paths in this "
"database, not just benchmark-created ones."
),
)
parser.add_argument(
"--seed",
type=int,
default=42,
help="RNG seed for reproducible datasets (default: 42).",
)
parser.add_argument(
"--repeat",
type=int,
default=5,
help="Number of timed repetitions for `run`/`profile` (default: 5).",
)
parser.add_argument(
"--label",
default="baseline",
help="Free-text tag for a `run`, printed and recorded in history only.",
)
parser.add_argument(
"--explain",
action="store_true",
default=False,
help="For `profile`: also capture and print the query plan.",
)
def handle(self, *args: Any, **options: Any) -> None:
action = options["action"]
if action == "seed":
self._handle_seed(options)
elif action == "run":
self._handle_run(options)
elif action == "profile":
self._handle_profile(options)
else:
self._handle_list_scenarios()
def _handle_seed(self, options: dict[str, Any]) -> None:
from paperless_benchmark.db import reset_benchmark_data
from paperless_benchmark.seeding import seed_benchmark_dataset
if options["reset"]:
if not options["yes_i_know_this_wipes_the_database"]:
raise CommandError(
"--reset requires --yes-i-know-this-wipes-the-database. "
"This deletes ALL users, ALL groups, and ALL documents, "
"tags, correspondents, document types, and storage paths "
"in this database -- not just benchmark-created ones. "
"Only run this against a disposable benchmark database, "
"never a real install. Re-run with "
"--reset --yes-i-know-this-wipes-the-database to proceed.",
)
self.stdout.write("Resetting existing benchmark data...")
reset_benchmark_data()
data = seed_benchmark_dataset(options["tier"], seed=options["seed"])
self.stdout.write(
self.style.SUCCESS(
f"Seeded tier={options['tier']!r}: {data.documents} documents, "
f"{len(data.users)} users, {len(data.groups)} groups.",
),
)
def _handle_run(self, options: dict[str, Any]) -> None:
from django.contrib.auth import get_user_model
from django.db import connection
from documents.models import Document
from paperless_benchmark.endpoints import run_endpoint_benchmarks
from paperless_benchmark.results import append_history
user_model = get_user_model()
try:
perf_target = user_model.objects.get(username="perf_target")
perf_admin = user_model.objects.get(username="perf_admin")
except user_model.DoesNotExist as e:
raise CommandError(
"No benchmark dataset found. Run `manage.py benchmark seed` first.",
) from e
db_vendor = connection.vendor
document_count = Document.objects.count()
results = run_endpoint_benchmarks(
perf_target=perf_target,
perf_admin=perf_admin,
repeat=options["repeat"],
)
self.stdout.write(f"# label={options['label']} repeat={options['repeat']}")
self.stdout.write(
f"{'user':7s} {'endpoint':20s} {'queries':>8s} "
f"{'min_ms':>9s} {'median_ms':>10s} {'max_ms':>9s}",
)
for r in results:
self.stdout.write(
f"{r.user_label:7s} {r.endpoint_name:20s} {r.query_count:8d} "
f"{r.min_ms:9.1f} {r.median_ms:10.1f} {r.max_ms:9.1f}",
)
append_history(
{
"mode": "run",
"label": options["label"],
"user": r.user_label,
"endpoint": r.endpoint_name,
"query_count": r.query_count,
"min_ms": r.min_ms,
"median_ms": r.median_ms,
"max_ms": r.max_ms,
"db_vendor": db_vendor,
"document_count": document_count,
},
)
def _handle_profile(self, options: dict[str, Any]) -> None:
from django.contrib.auth import get_user_model
from django.db import connection
from documents.models import Document
from paperless_benchmark.db import capture_explain
from paperless_benchmark.harness import run_profile
from paperless_benchmark.results import append_history
from paperless_benchmark.scenarios import get as get_scenario
if not options["scenario"]:
raise CommandError(
"`profile` requires a scenario name; see `list-scenarios`.",
)
scenario = get_scenario(options["scenario"])
user_model = get_user_model()
try:
perf_target = user_model.objects.get(username="perf_target")
except user_model.DoesNotExist as e:
raise CommandError(
"No benchmark dataset found. Run `manage.py benchmark seed` first.",
) from e
profile = run_profile(
lambda: scenario.run(perf_target),
repeat=options["repeat"],
)
self.stdout.write(
f"{scenario.name}: best={profile.best_seconds:.4f}s "
f"queries={profile.query_count}",
)
if options["explain"]:
if scenario.queryset_for_explain is not None:
plan = capture_explain(scenario.queryset_for_explain(perf_target))
self.stdout.write(plan)
else:
self.stdout.write(
self.style.WARNING(
f"--explain was requested but scenario {scenario.name!r} "
"does not support it (no queryset_for_explain); skipping.",
),
)
append_history(
{
"mode": "profile",
"scenario": scenario.name,
"best_seconds": profile.best_seconds,
"query_count": profile.query_count,
"db_vendor": connection.vendor,
"document_count": Document.objects.count(),
},
)
def _handle_list_scenarios(self) -> None:
from paperless_benchmark.scenarios import all_scenarios
for scenario in all_scenarios():
self.stdout.write(f"{scenario.name}: {scenario.describe}")
+39
View File
@@ -0,0 +1,39 @@
# src/paperless_benchmark/results.py
from __future__ import annotations
import json
import subprocess
from datetime import UTC
from datetime import datetime
from pathlib import Path
from typing import Any
RESULTS_DIR = Path(__file__).resolve().parent.parent.parent / "benchmark_results"
def _current_git_ref() -> str:
result = subprocess.run(
["git", "rev-parse", "--short", "HEAD"],
capture_output=True,
text=True,
check=False,
)
return result.stdout.strip() or "unknown"
def append_history(entry: dict[str, Any], *, code_ref: str | None = None) -> None:
"""
Append one line to benchmark_results/history.jsonl -- a local-only,
append-only, cross-session record of every `benchmark run`/`profile`
invocation. Unlike a single overwritten snapshot file, this survives
across sessions so a benchmarking effort picked back up days later has
a full timeline instead of only the most recent result.
"""
RESULTS_DIR.mkdir(exist_ok=True)
record = {
"timestamp": datetime.now(UTC).isoformat(),
"code_ref": code_ref or _current_git_ref(),
**entry,
}
with (RESULTS_DIR / "history.jsonl").open("a") as f:
f.write(json.dumps(record) + "\n")
+110
View File
@@ -0,0 +1,110 @@
# src/paperless_benchmark/scenarios.py
from __future__ import annotations
from dataclasses import dataclass
from typing import TYPE_CHECKING
from typing import Any
if TYPE_CHECKING:
from collections.abc import Callable
from django.contrib.auth.models import User
from django.db.models import QuerySet
@dataclass(frozen=True, slots=True)
class Scenario:
name: str
describe: str
run: Callable[[User], Any]
queryset_for_explain: Callable[[User], QuerySet] | None = None
_SCENARIOS: dict[str, Scenario] = {}
def register(scenario: Scenario) -> None:
_SCENARIOS[scenario.name] = scenario
def get(name: str) -> Scenario:
from django.core.management.base import CommandError
try:
return _SCENARIOS[name]
except KeyError:
available = ", ".join(sorted(_SCENARIOS)) or "(none registered)"
raise CommandError(
f"Unknown benchmark scenario {name!r}. Available: {available}",
) from None
def all_scenarios() -> tuple[Scenario, ...]:
return tuple(_SCENARIOS.values())
def _guardian_visibility_query_run(user: User) -> list[int]:
from documents.models import Document
from documents.permissions import get_objects_for_user_owner_aware
return list(
get_objects_for_user_owner_aware(
user,
"documents.view_document",
Document,
).values_list("id", flat=True),
)
def _guardian_visibility_query_queryset(user: User) -> QuerySet:
from documents.models import Document
from documents.permissions import get_objects_for_user_owner_aware
return get_objects_for_user_owner_aware(user, "documents.view_document", Document)
register(
Scenario(
name="guardian_visibility_query",
describe=(
"Document-visibility queryset for a user with mixed owned/shared "
"documents -- exercises documents.permissions."
"get_objects_for_user_owner_aware's guardian permission join."
),
run=_guardian_visibility_query_run,
queryset_for_explain=_guardian_visibility_query_queryset,
),
)
def _permitted_document_ids_run(user: User) -> list[int]:
from documents.models import Document
from documents.permissions import permitted_document_ids
return list(
Document.objects.filter(id__in=permitted_document_ids(user)).values_list(
"id",
flat=True,
),
)
def _permitted_document_ids_queryset(user: User) -> QuerySet:
from documents.models import Document
from documents.permissions import permitted_document_ids
return Document.objects.filter(id__in=permitted_document_ids(user))
register(
Scenario(
name="permitted_document_ids",
describe=(
"Document-visibility query built from documents.permissions."
"permitted_document_ids -- the resolved-ID-set alternative to "
"guardian_visibility_query, for side-by-side comparison."
),
run=_permitted_document_ids_run,
queryset_for_explain=_permitted_document_ids_queryset,
),
)
+445
View File
@@ -0,0 +1,445 @@
# src/paperless_benchmark/seeding.py
from __future__ import annotations
import datetime
import random
import time
from dataclasses import dataclass
from typing import TYPE_CHECKING
from typing import Literal
if TYPE_CHECKING:
from django.contrib.auth.models import Group
from django.contrib.auth.models import User
Tier = Literal["home", "medium", "large"]
CHUNK_SIZE = 5_000
MIME_TYPES = (
"application/pdf",
"image/png",
"image/jpeg",
"text/plain",
)
# Ownership split for documents, mirroring the shape used in the #11950
# perf-benchmark dataset (owned-by-target / owned-by-other / unowned).
OWNED_BY_TARGET_FRACTION = 0.60
OWNED_BY_OTHER_FRACTION = 0.30
# remainder (0.10) is unowned
# Of documents owned by "other" users, the fraction explicitly shared
# (view, or view+change) with perf_target via guardian permissions --
# this is what exercises the get_user_can_change() per-row N+1 that the
# `run` endpoint benchmarks measure.
SHARED_WITH_TARGET_FRACTION = 0.5
SHARED_WITH_CHANGE_FRACTION = 0.5
# Guardian permission-row ratios measured from a real install (discussion
# #13276): 1,414 user-perm rows / 27,232 group-perm rows over 12,000
# documents. Layered across ALL owned documents for the general user/group
# pool (not just perf_target's shares), so `profile` scenarios exercise a
# realistic permission-join shape for arbitrary users, not only perf_target.
USER_PERM_ROWS_PER_DOC = 1_414 / 12_000
GROUP_PERM_ROWS_PER_DOC = 27_232 / 12_000
@dataclass(frozen=True, slots=True)
class _TierCounts:
documents: int
tags: int
correspondents: int
document_types: int
storage_paths: int
other_users: int
groups: int
tags_per_doc: tuple[int, int]
TIERS: dict[Tier, _TierCounts] = {
"home": _TierCounts(
documents=500,
tags=20,
correspondents=10,
document_types=8,
storage_paths=5,
other_users=3,
groups=2,
tags_per_doc=(1, 3),
),
"medium": _TierCounts(
documents=20_000,
tags=100,
correspondents=300,
document_types=50,
storage_paths=20,
other_users=10,
groups=5,
tags_per_doc=(2, 6),
),
"large": _TierCounts(
documents=360_000,
tags=1_000,
correspondents=5_000,
document_types=300,
storage_paths=50,
other_users=25,
groups=10,
tags_per_doc=(3, 7),
),
}
def log(msg: str) -> None:
print(f"[{time.strftime('%H:%M:%S')}] {msg}", flush=True) # noqa: T201
@dataclass(frozen=True, slots=True)
class SeededData:
"""
Summary of a completed seed run. `documents`/`tags`/`correspondents`/
`document_types`/`storage_paths` are counts, not the seeded ORM
instances: at the `large` tier (360,000 documents) holding every
instance in memory simultaneously is a real risk for zero benefit, since
no caller consumes anything but the counts. `perf_target`/`perf_admin`/
`users`/`groups` stay as real objects -- at most ~26 users/12 groups even
at `large` tier, and small enough to be useful to a future caller.
"""
perf_target: User
perf_admin: User
users: tuple[User, ...]
groups: tuple[Group, ...]
documents: int
tags: int
correspondents: int
document_types: int
storage_paths: int
def _grant_model_level_permissions(user: User) -> None:
"""
Grant perf_target Django model-level view/add/change permissions on
Document and Tag, on top of the per-object guardian grants seeding
creates elsewhere. DRF's PaperlessObjectPermissions checks model-level
permissions before guardian's object-level ones are ever consulted, so
without this perf_target gets a blanket 403 on /api/documents/ and
/api/tags/ regardless of which documents guardian says it can see.
"""
from django.contrib.auth.models import Permission
from django.contrib.contenttypes.models import ContentType
from documents.models import Document
from documents.models import Tag
for model in (Document, Tag):
content_type = ContentType.objects.get_for_model(model)
codenames = [
f"{action}_{model._meta.model_name}" for action in ("view", "add", "change")
]
perms = Permission.objects.filter(
content_type=content_type,
codename__in=codenames,
)
user.user_permissions.add(*perms)
def _create_users_and_groups(
counts: _TierCounts,
) -> tuple[User, User, tuple[User, ...], tuple[Group, ...]]:
from django.contrib.auth.models import Group
from documents.tests.factories import UserFactory
perf_target = UserFactory.create(username="perf_target")
_grant_model_level_permissions(perf_target)
perf_admin = UserFactory.create(username="perf_admin", superuser=True)
other_users = tuple(UserFactory.create_batch(counts.other_users))
groups = tuple(
Group.objects.create(name=f"benchmark_group_{i}") for i in range(counts.groups)
)
log(
f"Created users: 1 target, 1 superuser, {len(other_users)} other, "
f"{len(groups)} groups.",
)
return perf_target, perf_admin, other_users, groups
def _create_lookup_tables(counts: _TierCounts):
from documents.models import Correspondent
from documents.models import DocumentType
from documents.models import StoragePath
from documents.models import Tag
from documents.tests.factories import CorrespondentFactory
from documents.tests.factories import DocumentTypeFactory
from documents.tests.factories import StoragePathFactory
from documents.tests.factories import TagFactory
tags = tuple(Tag.objects.bulk_create(TagFactory.build_batch(counts.tags)))
correspondents = tuple(
Correspondent.objects.bulk_create(
CorrespondentFactory.build_batch(counts.correspondents),
),
)
document_types = tuple(
DocumentType.objects.bulk_create(
DocumentTypeFactory.build_batch(counts.document_types),
),
)
storage_paths = tuple(
StoragePath.objects.bulk_create(
StoragePathFactory.build_batch(counts.storage_paths),
),
)
log(
f"Created {len(tags)} tags, {len(correspondents)} correspondents, "
f"{len(document_types)} document types, {len(storage_paths)} storage paths.",
)
return tags, correspondents, document_types, storage_paths
def _assign_owner(rng: random.Random, perf_target: User, other_users: tuple[User, ...]):
roll = rng.random()
if roll < OWNED_BY_TARGET_FRACTION:
return perf_target, "target"
if roll < OWNED_BY_TARGET_FRACTION + OWNED_BY_OTHER_FRACTION:
return rng.choice(other_users), "other"
return None, "unowned"
def _seed_documents(
rng: random.Random,
counts: _TierCounts,
tags,
correspondents,
document_types,
storage_paths,
perf_target,
other_users,
) -> tuple[int, list[int]]:
"""
Bulk-create `counts.documents` documents in chunks. Returns the total
document count plus a lightweight list of pks for documents that ended
up with an owner (target or other) -- that's all
`_grant_general_permissions` needs to sample from, so full `Document`
instances aren't accumulated across chunks (a real memory concern at the
`large` tier's 360,000 documents).
"""
from django.contrib.auth.models import Permission
from django.contrib.contenttypes.models import ContentType
from guardian.models import UserObjectPermission
from documents.models import Document
from documents.tests.factories import DocumentFactory
tag_ids = [t.pk for t in tags]
correspondent_ids = [c.pk for c in correspondents]
document_type_ids = [d.pk for d in document_types]
storage_path_ids = [s.pk for s in storage_paths]
doc_content_type = ContentType.objects.get_for_model(Document)
view_perm = Permission.objects.get(
codename="view_document",
content_type=doc_content_type,
)
change_perm = Permission.objects.get(
codename="change_document",
content_type=doc_content_type,
)
through_model = Document.tags.through
document_count = 0
owned_document_pks: list[int] = []
remaining = counts.documents
while remaining > 0:
chunk_n = min(CHUNK_SIZE, remaining)
remaining -= chunk_n
batch = []
owner_buckets = []
for _ in range(chunk_n):
doc = DocumentFactory.build(
mime_type=rng.choice(MIME_TYPES),
page_count=rng.randint(1, 30),
correspondent_id=(
rng.choice(correspondent_ids)
if correspondent_ids and rng.random() < 0.8
else None
),
document_type_id=(
rng.choice(document_type_ids)
if document_type_ids and rng.random() < 0.6
else None
),
storage_path_id=(
rng.choice(storage_path_ids)
if storage_path_ids and rng.random() < 0.8
else None
),
# Document.created is a plain DateField (default: today).
# Leaving it unset would give every seeded document the same
# date, collapsing Document's ("-created",) ordering index
# into a single-valued sort key -- spread it over a
# realistic multi-year window instead.
created=datetime.date.today()
- datetime.timedelta(days=rng.randint(0, 365 * 3)),
)
owner, bucket = _assign_owner(rng, perf_target, other_users)
doc.owner_id = owner.pk if owner else None
batch.append(doc)
owner_buckets.append(bucket)
created = Document.objects.bulk_create(batch, batch_size=CHUNK_SIZE)
through_rows = []
for doc in created:
k = rng.randint(*counts.tags_per_doc)
for tag_id in rng.sample(tag_ids, min(k, len(tag_ids))):
through_rows.append(through_model(document_id=doc.pk, tag_id=tag_id))
if through_rows:
through_model.objects.bulk_create(through_rows, batch_size=CHUNK_SIZE)
perm_rows = []
for doc, bucket in zip(created, owner_buckets, strict=True):
if bucket == "unowned":
continue
owned_document_pks.append(doc.pk)
if bucket != "other":
continue
if rng.random() >= SHARED_WITH_TARGET_FRACTION:
continue
perm_rows.append(
UserObjectPermission(
permission=view_perm,
content_type=doc_content_type,
object_pk=str(doc.pk),
user=perf_target,
),
)
if rng.random() < SHARED_WITH_CHANGE_FRACTION:
perm_rows.append(
UserObjectPermission(
permission=change_perm,
content_type=doc_content_type,
object_pk=str(doc.pk),
user=perf_target,
),
)
if perm_rows:
UserObjectPermission.objects.bulk_create(perm_rows, batch_size=CHUNK_SIZE)
document_count += len(created)
log(f" {document_count}/{counts.documents} documents seeded")
return document_count, owned_document_pks
def _grant_general_permissions(
rng: random.Random,
owned_document_pks: list[int],
users,
groups,
) -> None:
"""
Layer realistic (issue #13276-derived) guardian permission-row ratios
across owned documents for the general user/group pool, so `profile`
scenarios exercise the same permission-join shape regardless of which
user they check visibility for (not just perf_target).
"""
from django.contrib.auth.models import Permission
from django.contrib.contenttypes.models import ContentType
from guardian.models import GroupObjectPermission
from guardian.models import UserObjectPermission
from documents.models import Document
if not owned_document_pks or not users:
return
doc_content_type = ContentType.objects.get_for_model(Document)
view_perm = Permission.objects.get(
codename="view_document",
content_type=doc_content_type,
)
n_user_perms = round(len(owned_document_pks) * USER_PERM_ROWS_PER_DOC)
n_group_perms = (
round(len(owned_document_pks) * GROUP_PERM_ROWS_PER_DOC) if groups else 0
)
user_rows = [
UserObjectPermission(
permission=view_perm,
content_type=doc_content_type,
object_pk=str(rng.choice(owned_document_pks)),
user=rng.choice(users),
)
for _ in range(n_user_perms)
]
if user_rows:
UserObjectPermission.objects.bulk_create(
user_rows,
batch_size=CHUNK_SIZE,
ignore_conflicts=True,
)
group_rows = [
GroupObjectPermission(
permission=view_perm,
content_type=doc_content_type,
object_pk=str(rng.choice(owned_document_pks)),
group=rng.choice(groups),
)
for _ in range(n_group_perms)
]
if group_rows:
GroupObjectPermission.objects.bulk_create(
group_rows,
batch_size=CHUNK_SIZE,
ignore_conflicts=True,
)
log(f" Granted {len(user_rows)} user perms, {len(group_rows)} group perms.")
def seed_benchmark_dataset(tier: Tier, *, seed: int = 42) -> SeededData:
"""
Build a benchmark dataset at the given scale tier: a named perf_target
(mixed owned/shared documents) and perf_admin (superuser) for endpoint
benchmarking, plus a general user/group pool with realistic guardian
permission-row ratios for profile scenarios.
"""
counts = TIERS[tier]
rng = random.Random(seed)
log(f"Seeding tier={tier!r}")
perf_target, perf_admin, other_users, groups = _create_users_and_groups(counts)
tags, correspondents, document_types, storage_paths = _create_lookup_tables(counts)
document_count, owned_document_pks = _seed_documents(
rng,
counts,
tags,
correspondents,
document_types,
storage_paths,
perf_target,
other_users,
)
all_users = (perf_target, *other_users)
_grant_general_permissions(rng, owned_document_pks, all_users, groups)
log(f"Done. {document_count} documents seeded for tier={tier!r}.")
return SeededData(
perf_target=perf_target,
perf_admin=perf_admin,
users=all_users,
groups=groups,
documents=document_count,
tags=len(tags),
correspondents=len(correspondents),
document_types=len(document_types),
storage_paths=len(storage_paths),
)
+2 -22
View File
@@ -1,9 +1,7 @@
import logging
from celery import Task
from celery import shared_task
from documents.models import PaperlessTask
from paperless_mail.mail import MailAccountHandler
from paperless_mail.mail import MailError
from paperless_mail.models import MailAccount
@@ -12,26 +10,8 @@ from paperless_mail.models import MailRule
logger = logging.getLogger("paperless.mail.tasks")
@shared_task(bind=True)
def process_mail_accounts(self: Task, account_ids: list[int] | None = None) -> str:
# A scheduled check can still be running (or queued) when the next one
# ProcessedMail dedup only records a message once its
# handling has finished, so an overlapping run can still pick up the same
# not-yet-recorded message. Skip outright rather than race it.
other_mail_fetch_running = (
PaperlessTask.objects.filter(
task_type=PaperlessTask.TaskType.MAIL_FETCH,
status__in=[PaperlessTask.Status.PENDING, PaperlessTask.Status.STARTED],
)
.exclude(task_id=self.request.id)
.exists()
)
if other_mail_fetch_running:
logger.info(
"Mail account processing is already running; skipping this run.",
)
return "Skipped: mail account processing already in progress."
@shared_task
def process_mail_accounts(account_ids: list[int] | None = None) -> str:
total_new_documents = 0
accounts = (
MailAccount.objects.filter(pk__in=account_ids)
@@ -1,134 +0,0 @@
from typing import Final
import pytest
import pytest_mock
from documents.models import PaperlessTask
from documents.tests.factories import PaperlessTaskFactory
from paperless_mail import tasks
from paperless_mail.tests.factories import MailAccountFactory
from paperless_mail.tests.factories import MailRuleFactory
NO_DOCUMENTS_ADDED: Final = "No new documents were added."
SKIPPED: Final = "Skipped: mail account processing already in progress."
@pytest.mark.django_db
@pytest.mark.usefixtures("account_with_rule")
class TestProcessMailAccountsOverlap:
@pytest.fixture
def account_with_rule(self) -> None:
"""An enabled mail account with a single enabled rule."""
account = MailAccountFactory.create()
MailRuleFactory.create(account=account, enabled=True)
@pytest.mark.parametrize(
("status", "expected_result", "expected_call_count"),
[
pytest.param(
PaperlessTask.Status.PENDING,
SKIPPED,
0,
id="pending-task-blocks",
),
pytest.param(
PaperlessTask.Status.STARTED,
SKIPPED,
0,
id="started-task-blocks",
),
pytest.param(
PaperlessTask.Status.SUCCESS,
NO_DOCUMENTS_ADDED,
1,
id="finished-task-does-not-block",
),
],
)
def test_skips_only_while_another_mail_fetch_task_runs(
self,
mocker: pytest_mock.MockerFixture,
status: PaperlessTask.Status,
expected_result: str,
expected_call_count: int,
) -> None:
"""
GIVEN:
- An enabled mail account with a rule
- Another mail fetch task row in the given status
WHEN:
- Mail accounts are processed
THEN:
- Processing is skipped only if that other task is pending or running
"""
PaperlessTaskFactory.create(
task_type=PaperlessTask.TaskType.MAIL_FETCH,
trigger_source=PaperlessTask.TriggerSource.SCHEDULED,
status=status,
)
mocked_handle = mocker.patch.object(
tasks.MailAccountHandler,
"handle_mail_account",
return_value=0,
)
result = tasks.process_mail_accounts()
assert mocked_handle.call_count == expected_call_count
assert result == expected_result
def test_runs_when_no_other_mail_fetch_task_exists(
self,
mocker: pytest_mock.MockerFixture,
) -> None:
"""
GIVEN:
- An enabled mail account with a rule
- No other mail fetch task rows
WHEN:
- Mail accounts are processed
THEN:
- The account is handled
"""
mocked_handle = mocker.patch.object(
tasks.MailAccountHandler,
"handle_mail_account",
return_value=0,
)
result = tasks.process_mail_accounts()
mocked_handle.assert_called_once()
assert result == NO_DOCUMENTS_ADDED
def test_does_not_skip_due_to_its_own_task_row(
self,
mocker: pytest_mock.MockerFixture,
) -> None:
"""
GIVEN:
- An enabled mail account with a rule
- A running mail fetch task row belonging to this very task
WHEN:
- Mail accounts are processed under that task id
THEN:
- The task does not skip itself and handles the account
"""
PaperlessTaskFactory.create(
task_id="self-task-id",
task_type=PaperlessTask.TaskType.MAIL_FETCH,
trigger_source=PaperlessTask.TriggerSource.SCHEDULED,
status=PaperlessTask.Status.STARTED,
)
mocked_handle = mocker.patch.object(
tasks.MailAccountHandler,
"handle_mail_account",
return_value=0,
)
result = tasks.process_mail_accounts.apply(task_id="self-task-id").result
mocked_handle.assert_called_once()
assert result == NO_DOCUMENTS_ADDED