1use std::path::{Component, Path, PathBuf};
22
23use gaze::Manifest;
24use serde::{Deserialize, Serialize};
25
26use crate::ocr::{
27 detect_image_format, ImageFormat, ImageInput, OcrBackend, OcrError, OcrHints, OcrResult,
28};
29
30#[cfg(feature = "ocr-tesseract")]
31use std::collections::BTreeMap;
32#[cfg(feature = "ocr-tesseract")]
33use std::fs;
34
35#[cfg(any(feature = "ocr-tesseract", feature = "mcp"))]
36use gaze::{
37 Action, ClassRule, CleanDocument, DefaultRule, LocaleTag, Pipeline as GazePipeline,
38 RawDocument, Scope, Session,
39};
40#[cfg(any(feature = "ocr-tesseract", feature = "mcp"))]
41use gaze_recognizers::{
42 AnchoredBoundary, AnchoredMatchRecognizer, CuePosition, NameShape, RegexDetector,
43};
44#[cfg(any(feature = "ocr-tesseract", feature = "mcp"))]
45use gaze_types::{EmittedTokenSpan, PiiClass};
46
47#[cfg(feature = "ocr-tesseract")]
48use crate::extract::InputKind;
49#[cfg(feature = "ocr-tesseract")]
50use crate::{BundleLayoutInvalidReason, DocumentError};
51
52pub const BUNDLE_VERSION: u32 = 2;
54const DEFAULT_LOW_CONFIDENCE_THRESHOLD: f32 = 0.65;
55
56pub const CLEAN_MARKDOWN_FILE: &str = "clean.md";
58pub const MANIFEST_FILE: &str = "manifest.json";
60pub const REPORT_FILE: &str = "report.json";
62
63#[non_exhaustive]
65#[derive(Debug, Clone, PartialEq, Eq)]
66pub struct AgentBundleDir(PathBuf);
67
68impl AgentBundleDir {
69 pub fn new(path: impl Into<PathBuf>) -> Result<Self, DocumentError> {
74 let path = path.into();
75 validate_non_empty_path(&path)?;
76 Ok(Self(path))
77 }
78
79 pub fn as_path(&self) -> &Path {
81 &self.0
82 }
83}
84
85#[non_exhaustive]
87#[derive(Debug, Clone, PartialEq, Eq)]
88pub struct OwnerBundleDir(PathBuf);
89
90impl OwnerBundleDir {
91 pub fn new(path: impl Into<PathBuf>) -> Result<Self, DocumentError> {
96 let path = path.into();
97 validate_non_empty_path(&path)?;
98 Ok(Self(path))
99 }
100
101 pub fn as_path(&self) -> &Path {
103 &self.0
104 }
105}
106
107#[non_exhaustive]
109#[derive(Debug, Clone)]
110pub struct SafeBundle {
111 pub clean_markdown: String,
113 pub manifest: Manifest,
115 pub layout: LayoutSummary,
117 pub preview_png: Option<Vec<u8>>,
119 pub report: BundleReport,
121 pub source_path: PathBuf,
123 pub agent_out_dir: PathBuf,
125 pub owner_out_dir: PathBuf,
127}
128
129impl SafeBundle {
130 #[allow(clippy::too_many_arguments)]
132 pub fn new(
133 clean_markdown: String,
134 manifest: Manifest,
135 layout: LayoutSummary,
136 preview_png: Option<Vec<u8>>,
137 report: BundleReport,
138 source_path: PathBuf,
139 agent_out_dir: PathBuf,
140 owner_out_dir: PathBuf,
141 ) -> Self {
142 Self {
143 clean_markdown,
144 manifest,
145 layout,
146 preview_png,
147 report,
148 source_path,
149 agent_out_dir,
150 owner_out_dir,
151 }
152 }
153}
154
155#[non_exhaustive]
157#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
158pub struct ClassCount {
159 pub class: String,
161 pub count: u32,
164}
165
166impl ClassCount {
167 pub fn new(class: impl Into<String>, count: u32) -> Self {
169 Self {
170 class: class.into(),
171 count,
172 }
173 }
174}
175
176#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
178#[serde(rename_all = "snake_case")]
179pub enum OcrSource {
180 VectorPdf,
182 Ocr,
184}
185
186#[non_exhaustive]
188#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
189pub struct PageReport {
190 pub page_index: i32,
192 pub ocr_source: OcrSource,
194 pub ocr_backend: Option<String>,
196 pub confidence: Option<f32>,
198 pub low_confidence: bool,
200 pub column_count: u32,
202 pub ocr_word_count: usize,
204 pub ocr_mean_confidence: Option<f32>,
206}
207
208impl PageReport {
209 fn new(
210 page_index: i32,
211 ocr_source: OcrSource,
212 ocr_backend: Option<String>,
213 ocr: &OcrResult,
214 column_count: u32,
215 low_confidence_threshold: f32,
216 ) -> Self {
217 let confidence = ocr.mean_confidence_unit();
218 Self {
219 page_index,
220 ocr_source,
221 ocr_backend,
222 confidence,
223 low_confidence: confidence
224 .map(|confidence| confidence < low_confidence_threshold)
225 .unwrap_or(false),
226 column_count,
227 ocr_word_count: ocr.word_count,
228 ocr_mean_confidence: ocr.mean_confidence,
229 }
230 }
231}
232
233#[non_exhaustive]
239#[derive(Debug, Clone, Serialize, Deserialize)]
240pub struct BundleReport {
241 pub bundle_version: u32,
243 pub input_kind: String,
245 pub ocr_mean_confidence: Option<f32>,
247 pub ocr_word_count: usize,
249 pub ocr_lang: String,
251 pub clean_char_count: usize,
253 pub pii_token_count: u32,
260 pub pii_tokens_by_class: Vec<ClassCount>,
265 pub pdf_page_count: Option<i32>,
267 pub pdf_page_index: Option<i32>,
269 #[serde(default)]
271 pub pages: Vec<PageReport>,
272 #[serde(default = "default_low_confidence_threshold")]
274 pub low_confidence_threshold: f32,
275}
276
277impl BundleReport {
278 #[allow(clippy::too_many_arguments)]
280 pub fn new(
281 input_kind: impl Into<String>,
282 ocr: &OcrResult,
283 clean_char_count: usize,
284 pii_token_count: u32,
285 pii_tokens_by_class: Vec<ClassCount>,
286 pdf_page_count: Option<i32>,
287 pdf_page_index: Option<i32>,
288 pages: Vec<PageReport>,
289 low_confidence_threshold: f32,
290 ) -> Self {
291 Self {
292 bundle_version: BUNDLE_VERSION,
293 input_kind: input_kind.into(),
294 ocr_mean_confidence: ocr.mean_confidence,
295 ocr_word_count: ocr.word_count,
296 ocr_lang: ocr.lang.clone(),
297 clean_char_count,
298 pii_token_count,
299 pii_tokens_by_class,
300 pdf_page_count,
301 pdf_page_index,
302 pages,
303 low_confidence_threshold,
304 }
305 }
306}
307
308fn default_low_confidence_threshold() -> f32 {
309 DEFAULT_LOW_CONFIDENCE_THRESHOLD
310}
311
312#[non_exhaustive]
314#[derive(Debug, Clone, Copy)]
315pub struct Pipeline {
316 low_confidence_threshold: f32,
317 column_detection: bool,
318}
319
320impl Pipeline {
321 pub fn new() -> Self {
323 Self {
324 low_confidence_threshold: DEFAULT_LOW_CONFIDENCE_THRESHOLD,
325 column_detection: true,
326 }
327 }
328
329 pub fn with_low_confidence_threshold(mut self, threshold: f32) -> Self {
331 self.low_confidence_threshold = threshold.clamp(0.0, 1.0);
332 self
333 }
334
335 pub fn with_column_detection(mut self, enabled: bool) -> Self {
337 self.column_detection = enabled;
338 self
339 }
340
341 #[cfg(feature = "ocr-tesseract")]
346 #[cfg_attr(docsrs, doc(cfg(feature = "ocr-tesseract")))]
347 pub fn clean_with_ocr_backend(
348 &self,
349 input: &Path,
350 agent_out: AgentBundleDir,
351 owner_out: OwnerBundleDir,
352 ocr_backend: &dyn OcrBackend,
353 ) -> Result<SafeBundle, DocumentError> {
354 clean_with_options(input, agent_out, owner_out, ocr_backend, *self)
355 }
356}
357
358impl Default for Pipeline {
359 fn default() -> Self {
360 Self::new()
361 }
362}
363
364#[non_exhaustive]
369#[derive(Debug, Clone)]
370pub struct LayoutSummary {
371 pub page_count: u32,
373}
374
375impl LayoutSummary {
376 pub fn single_page() -> Self {
378 Self { page_count: 1 }
379 }
380
381 pub fn new(page_count: u32) -> Self {
383 Self { page_count }
384 }
385}
386
387#[cfg(feature = "ocr-tesseract")]
399#[cfg_attr(docsrs, doc(cfg(feature = "ocr-tesseract")))]
400pub fn clean(
401 input: &Path,
402 agent_out: AgentBundleDir,
403 owner_out: OwnerBundleDir,
404) -> Result<SafeBundle, DocumentError> {
405 let backend = crate::ocr::TesseractBackend::new();
406 Pipeline::new().clean_with_ocr_backend(input, agent_out, owner_out, &backend)
407}
408
409#[cfg(feature = "ocr-tesseract")]
420#[cfg_attr(docsrs, doc(cfg(feature = "ocr-tesseract")))]
421pub fn clean_with_ocr_backend(
422 input: &Path,
423 agent_out: AgentBundleDir,
424 owner_out: OwnerBundleDir,
425 ocr_backend: &dyn OcrBackend,
426) -> Result<SafeBundle, DocumentError> {
427 Pipeline::new().clean_with_ocr_backend(input, agent_out, owner_out, ocr_backend)
428}
429
430#[cfg(feature = "ocr-tesseract")]
431fn clean_with_options(
432 input: &Path,
433 agent_out: AgentBundleDir,
434 owner_out: OwnerBundleDir,
435 ocr_backend: &dyn OcrBackend,
436 options: Pipeline,
437) -> Result<SafeBundle, DocumentError> {
438 let kind = InputKind::detect(input)?;
439 let absolute_input = absolutize(input);
440 let (absolute_agent_out, absolute_owner_out) = prepare_bundle_dirs(&agent_out, &owner_out)?;
441
442 let extraction = run_document_extraction(input, kind, ocr_backend, options)?;
443 let normalized_text = crate::ocr::normalize_ocr_artifacts(&extraction.ocr_result.text);
450 let pipeline = build_document_pipeline()?;
451 let session = Session::new(Scope::Ephemeral).map_err(|err| pipeline_err("session", err))?;
452 let locale_chain = [LocaleTag::Global];
453 let (clean_doc, spans, _leak_report) = pipeline
454 .clean_with_safety_net(&session, RawDocument::Text(normalized_text), &locale_chain)
455 .map_err(|err| pipeline_err("redact", err))?;
456
457 let clean_text = match clean_doc {
458 CleanDocument::Text(text) => text,
459 _ => {
460 return Err(DocumentError::Pipeline(
461 "pipeline returned non-text variant for text input".to_string(),
462 ));
463 }
464 };
465
466 let manifest = Manifest::from_spans(spans.clone());
467 let counts = count_pii_by_class(&spans);
468 let pii_token_count: u32 = counts.iter().map(|c| c.count).sum();
469
470 let clean_markdown = format_clean_markdown(&clean_text, kind);
471
472 let report = BundleReport::new(
473 kind_label(kind),
474 &extraction.ocr_result,
475 clean_markdown.chars().count(),
476 pii_token_count,
477 counts,
478 extraction.pdf_page_count,
479 extraction.pdf_page_index,
480 extraction.pages,
481 options.low_confidence_threshold,
482 );
483
484 write_bundle(&agent_out, &owner_out, &clean_markdown, &manifest, &report)?;
485
486 Ok(SafeBundle::new(
487 clean_markdown,
488 manifest,
489 LayoutSummary::new(extraction.page_count),
490 None,
491 report,
492 absolute_input,
493 absolute_agent_out,
494 absolute_owner_out,
495 ))
496}
497
498#[cfg(feature = "ocr-tesseract")]
499struct DocumentExtraction {
500 ocr_result: OcrResult,
501 pdf_page_count: Option<i32>,
502 pdf_page_index: Option<i32>,
503 pages: Vec<PageReport>,
504 page_count: u32,
505}
506
507#[cfg(feature = "ocr-tesseract")]
508#[cfg_attr(not(feature = "mcp"), allow(dead_code))]
509pub(crate) fn run_ocr(
510 input: &Path,
511 kind: InputKind,
512 ocr_backend: &dyn OcrBackend,
513) -> Result<(OcrResult, Option<i32>, Option<i32>), DocumentError> {
514 let extraction = run_document_extraction(input, kind, ocr_backend, Pipeline::new())?;
515 Ok((
516 extraction.ocr_result,
517 extraction.pdf_page_count,
518 extraction.pdf_page_index,
519 ))
520}
521
522#[cfg(feature = "ocr-tesseract")]
523fn run_document_extraction(
524 input: &Path,
525 kind: InputKind,
526 ocr_backend: &dyn OcrBackend,
527 options: Pipeline,
528) -> Result<DocumentExtraction, DocumentError> {
529 match kind {
530 InputKind::Png | InputKind::Jpeg => {
531 let bytes = fs::read(input)?;
532 let format = detect_image_format(&bytes)?;
533 let (result, column_count) = recognize_image(
534 ocr_backend,
535 ImageInput {
536 bytes,
537 format,
538 dpi: None,
539 },
540 options.column_detection,
541 )?;
542 let page_report = PageReport::new(
543 0,
544 OcrSource::Ocr,
545 Some(ocr_backend.name().to_string()),
546 &result,
547 column_count,
548 options.low_confidence_threshold,
549 );
550 Ok(DocumentExtraction {
551 ocr_result: result,
552 pdf_page_count: None,
553 pdf_page_index: None,
554 pages: vec![page_report],
555 page_count: 1,
556 })
557 }
558 InputKind::Pdf => {
559 #[cfg(feature = "pdf-input")]
560 {
561 use crate::extract::pdf::{extract_pages, PdfPagePayload, PdfRasterConfig};
562 let payloads = extract_pages(input, PdfRasterConfig::new())?;
563 let mut page_results = Vec::with_capacity(payloads.len());
564 let mut pages = Vec::with_capacity(payloads.len());
565 let mut pdf_page_count = None;
566 let mut first_page_index = None;
567
568 for payload in payloads {
569 pdf_page_count = Some(payload.page_count());
570 if first_page_index.is_none() {
571 first_page_index = Some(payload.page_index());
572 }
573 match payload {
574 PdfPagePayload::VectorText {
575 text, page_index, ..
576 } => {
577 let result = OcrResult::new(text, None, 0, "vector-pdf".to_string());
578 pages.push(PageReport::new(
579 page_index,
580 OcrSource::VectorPdf,
581 None,
582 &result,
583 1,
584 options.low_confidence_threshold,
585 ));
586 page_results.push(result);
587 }
588 PdfPagePayload::Raster(raster) => {
589 let (result, column_count) = recognize_image(
590 ocr_backend,
591 ImageInput {
592 bytes: raster.png_bytes,
593 format: ImageFormat::Png,
594 dpi: None,
595 },
596 options.column_detection,
597 )?;
598 pages.push(PageReport::new(
599 raster.page_index,
600 OcrSource::Ocr,
601 Some(ocr_backend.name().to_string()),
602 &result,
603 column_count,
604 options.low_confidence_threshold,
605 ));
606 page_results.push(result);
607 }
608 }
609 }
610
611 Ok(DocumentExtraction {
612 ocr_result: merge_page_results(&page_results),
613 pdf_page_count,
614 pdf_page_index: first_page_index,
615 page_count: pages.len() as u32,
616 pages,
617 })
618 }
619 #[cfg(not(feature = "pdf-input"))]
620 {
621 Err(DocumentError::UnsupportedInput {
622 path: input.to_path_buf(),
623 reason: "rebuild gaze-document with `--features pdf-input` for PDF support",
624 })
625 }
626 }
627 }
628}
629
630#[cfg(feature = "ocr-tesseract")]
631fn recognize_image(
632 ocr_backend: &dyn OcrBackend,
633 image: ImageInput,
634 column_detection: bool,
635) -> Result<(OcrResult, u32), DocumentError> {
636 let hints = OcrHints::default();
637 let lang = hints.primary_language().to_string();
638 let image = crate::preprocess::preprocess_image(image);
639 let spans = ocr_backend
640 .recognize(image, hints)
641 .map_err(map_ocr_error_to_document_error)?;
642 Ok(OcrResult::from_spans_with_column_detection(
643 &spans,
644 lang,
645 column_detection,
646 ))
647}
648
649#[cfg(feature = "ocr-tesseract")]
650fn merge_page_results(results: &[OcrResult]) -> OcrResult {
651 let text = results
652 .iter()
653 .map(|result| result.text.as_str())
654 .collect::<Vec<_>>()
655 .join("\n\n");
656 let mut conf_sum = 0.0f64;
657 let mut conf_count = 0usize;
658 for result in results {
659 if let Some(confidence) = result.mean_confidence {
660 conf_sum += confidence as f64 * result.word_count as f64;
661 conf_count += result.word_count;
662 }
663 }
664 let mean_confidence = if conf_count == 0 {
665 None
666 } else {
667 Some((conf_sum / conf_count as f64) as f32)
668 };
669 OcrResult::new(text, mean_confidence, conf_count, "mixed".to_string())
670}
671
672#[cfg(feature = "ocr-tesseract")]
673fn map_ocr_error_to_document_error(err: OcrError) -> DocumentError {
674 match err {
675 OcrError::InitFailed(hint) => DocumentError::TesseractNotFound(hint),
676 OcrError::RecognizeFailed(detail) => DocumentError::TesseractFailed {
677 status: -1,
678 stderr: detail,
679 },
680 OcrError::UnsupportedFormat(format) => DocumentError::UnsupportedInput {
681 path: PathBuf::new(),
682 reason: match format {
683 ImageFormat::Png => "png image format is not supported by the OCR backend",
684 ImageFormat::Jpeg => "jpeg image format is not supported by the OCR backend",
685 ImageFormat::Tiff => "tiff image format is not supported by the OCR backend",
686 },
687 },
688 OcrError::Internal(detail) => DocumentError::Pipeline(format!("ocr: {detail}")),
689 }
690}
691
692#[cfg(feature = "ocr-tesseract")]
693#[cfg(any(feature = "ocr-tesseract", feature = "mcp"))]
694pub(crate) fn build_document_pipeline() -> Result<GazePipeline, DocumentError> {
695 let email = RegexDetector::emails().map_err(|err| pipeline_err("email-regex", err))?;
696 let phone = RegexDetector::new(
700 r"\+?\d{1,3}[-.\s]\(?\d{3}\)?[-.\s]?\d{3,4}[-.\s]?\d{0,4}",
701 PiiClass::custom("phone").expect("valid custom class"),
702 )
703 .map_err(|err| pipeline_err("phone-regex", err))?;
704 let recipient_name = AnchoredMatchRecognizer::new(
712 "gaze_document.name.recipient".to_string(),
713 vec![
714 "Bill to".to_string(),
715 "Invoice to".to_string(),
716 "Ship to".to_string(),
717 "Attention".to_string(),
718 "Attn".to_string(),
719 ],
720 AnchoredBoundary::LineEnd,
721 48,
722 NameShape::PersonName,
723 CuePosition::Before,
724 "invoice_recipient".to_string(),
725 2,
726 0.88,
727 110,
728 );
729 GazePipeline::builder()
730 .detector(email)
731 .detector(phone)
732 .recognizer(recipient_name)
733 .rule(ClassRule::new(PiiClass::Email, Action::Tokenize))
734 .rule(ClassRule::new(
735 PiiClass::custom("phone").expect("valid custom class"),
736 Action::Tokenize,
737 ))
738 .rule(ClassRule::new(PiiClass::Name, Action::Tokenize))
739 .rule(DefaultRule::new(Action::Preserve))
740 .build()
741 .map_err(|err| pipeline_err("build", err))
742}
743
744#[cfg(feature = "ocr-tesseract")]
745fn count_pii_by_class(spans: &[EmittedTokenSpan]) -> Vec<ClassCount> {
746 let mut by_class: BTreeMap<String, u32> = BTreeMap::new();
747 for span in spans {
748 *by_class.entry(span.class.to_canonical_str()).or_insert(0) += 1;
749 }
750 by_class
751 .into_iter()
752 .map(|(class, count)| ClassCount::new(class, count))
753 .collect()
754}
755
756#[cfg(feature = "ocr-tesseract")]
757fn write_bundle(
758 agent_out: &AgentBundleDir,
759 owner_out: &OwnerBundleDir,
760 clean_markdown: &str,
761 manifest: &Manifest,
762 report: &BundleReport,
763) -> Result<(), DocumentError> {
764 fs::write(
765 agent_out.as_path().join(CLEAN_MARKDOWN_FILE),
766 clean_markdown,
767 )?;
768 let manifest_json = serde_json::to_vec_pretty(manifest)?;
769 fs::write(owner_out.as_path().join(MANIFEST_FILE), manifest_json)?;
770 let report_json = serde_json::to_vec_pretty(report)?;
771 fs::write(agent_out.as_path().join(REPORT_FILE), report_json)?;
772 Ok(())
773}
774
775#[cfg(feature = "ocr-tesseract")]
776fn prepare_bundle_dirs(
777 agent_out: &AgentBundleDir,
778 owner_out: &OwnerBundleDir,
779) -> Result<(PathBuf, PathBuf), DocumentError> {
780 let agent = normalize_for_layout(agent_out.as_path());
781 let owner = normalize_for_layout(owner_out.as_path());
782 validate_bundle_layout(&agent, &owner)?;
783
784 fs::create_dir_all(agent_out.as_path())
785 .map_err(|err| DocumentError::OutputDir(agent.clone(), err))?;
786 fs::create_dir_all(owner_out.as_path())
787 .map_err(|err| DocumentError::OutputDir(owner.clone(), err))?;
788
789 let agent = fs::canonicalize(agent_out.as_path()).unwrap_or(agent);
790 let owner = fs::canonicalize(owner_out.as_path()).unwrap_or(owner);
791 validate_bundle_layout(&agent, &owner)?;
792 Ok((agent, owner))
793}
794
795fn validate_non_empty_path(path: &Path) -> Result<(), DocumentError> {
796 if path.as_os_str().is_empty() {
797 return Err(DocumentError::BundleLayoutInvalid {
798 reason: BundleLayoutInvalidReason::EmptyPath,
799 });
800 }
801 Ok(())
802}
803
804fn validate_bundle_layout(agent: &Path, owner: &Path) -> Result<(), DocumentError> {
805 if agent == owner {
806 return Err(DocumentError::BundleLayoutInvalid {
807 reason: BundleLayoutInvalidReason::AgentEqualsOwner,
808 });
809 }
810 if agent.starts_with(owner) {
811 return Err(DocumentError::BundleLayoutInvalid {
812 reason: BundleLayoutInvalidReason::AgentNestedInOwner,
813 });
814 }
815 if owner.starts_with(agent) {
816 return Err(DocumentError::BundleLayoutInvalid {
817 reason: BundleLayoutInvalidReason::OwnerNestedInAgent,
818 });
819 }
820 Ok(())
821}
822
823fn normalize_for_layout(path: &Path) -> PathBuf {
824 let absolute = absolutize(path);
825 let mut normalized = PathBuf::new();
826 for component in absolute.components() {
827 match component {
828 Component::CurDir => {}
829 Component::ParentDir => {
830 normalized.pop();
831 }
832 Component::Prefix(prefix) => normalized.push(prefix.as_os_str()),
833 Component::RootDir => normalized.push(component.as_os_str()),
834 Component::Normal(part) => normalized.push(part),
835 }
836 }
837 normalized
838}
839
840#[cfg(feature = "ocr-tesseract")]
841pub(crate) fn format_clean_markdown(text: &str, kind: InputKind) -> String {
842 let mut out = String::new();
843 out.push_str("# gaze-document safe bundle\n\n");
844 out.push_str(&format!("Source kind: `{}`\n\n", kind_label(kind)));
845 out.push_str("---\n\n");
846 out.push_str(text);
847 if !text.ends_with('\n') {
848 out.push('\n');
849 }
850 out
851}
852
853#[cfg(feature = "ocr-tesseract")]
854pub(crate) fn kind_label(kind: InputKind) -> &'static str {
855 match kind {
856 InputKind::Png => "png",
857 InputKind::Jpeg => "jpeg",
858 InputKind::Pdf => "pdf",
859 }
860}
861
862fn absolutize(path: &Path) -> PathBuf {
863 if path.is_absolute() {
864 path.to_path_buf()
865 } else {
866 std::env::current_dir()
867 .map(|cwd| cwd.join(path))
868 .unwrap_or_else(|_| path.to_path_buf())
869 }
870}
871
872#[cfg(feature = "ocr-tesseract")]
873fn pipeline_err(stage: &'static str, err: impl std::fmt::Display) -> DocumentError {
874 DocumentError::Pipeline(format!("{stage}: {err}"))
875}
876
877#[cfg(all(test, feature = "ocr-tesseract"))]
878mod tests {
879 use super::*;
880 use crate::ocr::{BBox, OcrSpan};
881
882 #[derive(Debug)]
883 struct MockBackend {
884 spans: Vec<OcrSpan>,
885 }
886
887 impl OcrBackend for MockBackend {
888 fn name(&self) -> &str {
889 "mock-ocr"
890 }
891
892 fn recognize(
893 &self,
894 _image: ImageInput,
895 _hints: OcrHints,
896 ) -> Result<Vec<OcrSpan>, OcrError> {
897 Ok(self.spans.clone())
898 }
899 }
900
901 fn span(text: &str, x: u32, y: u32, confidence: f32) -> OcrSpan {
902 OcrSpan {
903 text: text.to_string(),
904 bbox: BBox { x, y, w: 90, h: 16 },
905 confidence: Some(confidence),
906 }
907 }
908
909 fn bundle_dirs(tmp: &tempfile::TempDir) -> (AgentBundleDir, OwnerBundleDir) {
910 (
911 AgentBundleDir::new(tmp.path().join("agent")).expect("agent dir"),
912 OwnerBundleDir::new(tmp.path().join("owner")).expect("owner dir"),
913 )
914 }
915
916 #[test]
917 fn count_pii_by_class_groups_email_and_phone() {
918 let spans = vec![
919 EmittedTokenSpan::new(0..10, 0..10, PiiClass::Email),
920 EmittedTokenSpan::new(20..28, 20..28, PiiClass::Email),
921 EmittedTokenSpan::new(
922 40..50,
923 40..50,
924 PiiClass::custom("phone").expect("valid custom class"),
925 ),
926 ];
927 let counts = count_pii_by_class(&spans);
928 assert_eq!(counts.len(), 2);
929 let by_class: BTreeMap<_, _> = counts.iter().map(|c| (c.class.as_str(), c.count)).collect();
930 assert_eq!(by_class.get("email"), Some(&2));
931 assert_eq!(by_class.get("custom:phone"), Some(&1));
932 }
933
934 #[test]
940 fn class_counts_count_replacements_so_a_fragment_counts_on_its_own() {
941 let spans = vec![
942 EmittedTokenSpan::new(0..10, 0..21, PiiClass::Email),
943 EmittedTokenSpan::residual_fragment(11..17, 21..27, PiiClass::Email),
944 ];
945 let counts = count_pii_by_class(&spans);
946 let by_class: BTreeMap<_, _> = counts.iter().map(|c| (c.class.as_str(), c.count)).collect();
947 assert_eq!(
948 by_class.get("email"),
949 Some(&2),
950 "a residual fragment is its own replacement and is counted"
951 );
952 assert_eq!(
953 spans.iter().filter(|span| span.origin.is_whole()).count(),
954 1,
955 "but only one of them covers a whole recognized value"
956 );
957 }
958
959 #[test]
960 fn report_serializes_with_bundle_version() {
961 let ocr = OcrResult::new("body".into(), Some(91.5), 2, "eng".into());
962 let report = BundleReport::new(
963 "png",
964 &ocr,
965 42,
966 3,
967 vec![
968 ClassCount::new("email", 2),
969 ClassCount::new("custom:phone", 1),
970 ],
971 None,
972 None,
973 vec![PageReport::new(
974 0,
975 OcrSource::Ocr,
976 Some("tesseract".to_string()),
977 &ocr,
978 1,
979 DEFAULT_LOW_CONFIDENCE_THRESHOLD,
980 )],
981 DEFAULT_LOW_CONFIDENCE_THRESHOLD,
982 );
983 let json = serde_json::to_value(&report).expect("serialize");
984 assert_eq!(json["bundle_version"], BUNDLE_VERSION);
985 assert_eq!(json["input_kind"], "png");
986 assert_eq!(json["pii_token_count"], 3);
987 assert_eq!(json["pages"][0]["ocr_source"], "ocr");
988 assert_eq!(
989 json["low_confidence_threshold"],
990 DEFAULT_LOW_CONFIDENCE_THRESHOLD
991 );
992 }
993
994 #[test]
995 fn v1_report_without_page_fields_still_deserializes() {
996 let json = serde_json::json!({
997 "bundle_version": 1,
998 "input_kind": "png",
999 "ocr_mean_confidence": 90.0,
1000 "ocr_word_count": 2,
1001 "ocr_lang": "eng",
1002 "clean_char_count": 12,
1003 "pii_token_count": 1,
1004 "pii_tokens_by_class": [{ "class": "email", "count": 1 }],
1005 "pdf_page_count": null,
1006 "pdf_page_index": null
1007 });
1008
1009 let report: BundleReport = serde_json::from_value(json).expect("v1 parses");
1010
1011 assert_eq!(report.bundle_version, 1);
1012 assert!(report.pages.is_empty());
1013 assert_eq!(
1014 report.low_confidence_threshold,
1015 DEFAULT_LOW_CONFIDENCE_THRESHOLD
1016 );
1017 }
1018
1019 #[test]
1020 fn clean_with_mock_backend_flags_low_confidence_and_columns() {
1021 let backend = MockBackend {
1022 spans: vec![
1023 span("Bill", 20, 10, 0.50),
1024 span("to:", 116, 10, 0.50),
1025 span("Jane", 20, 36, 0.50),
1026 span("Doe", 116, 36, 0.50),
1027 span("Email:", 360, 10, 0.50),
1028 span("alice@example.invalid", 360, 36, 0.50),
1030 ],
1031 };
1032 let tmp = tempfile::tempdir().expect("tempdir");
1033 let input = tmp.path().join("input.png");
1034 fs::write(&input, b"\x89PNG\r\n\x1A\nnot-real-image").expect("write input");
1035 let (agent_out, owner_out) = bundle_dirs(&tmp);
1036
1037 let bundle = Pipeline::new()
1038 .with_low_confidence_threshold(0.65)
1039 .clean_with_ocr_backend(&input, agent_out, owner_out, &backend)
1040 .expect("clean succeeds");
1041
1042 assert_eq!(bundle.report.bundle_version, 2);
1043 assert_eq!(bundle.report.pages.len(), 1);
1044 let page = &bundle.report.pages[0];
1045 assert_eq!(page.ocr_backend.as_deref(), Some("mock-ocr"));
1046 assert_eq!(page.column_count, 2);
1047 assert_eq!(page.confidence, Some(0.5));
1048 assert!(page.low_confidence);
1049 assert!(
1050 bundle.clean_markdown.contains(":Email_"),
1051 "{}",
1052 bundle.clean_markdown
1053 );
1054 assert!(
1055 !bundle.clean_markdown.contains("alice@example.invalid"),
1057 "{}",
1058 bundle.clean_markdown
1059 );
1060 }
1061
1062 #[test]
1063 fn clean_with_mock_backend_preserves_table_cell_context() {
1064 let backend = MockBackend {
1065 spans: vec![
1066 span("Field", 20, 10, 0.92),
1067 span("Value", 160, 10, 0.92),
1068 span("Bill", 20, 40, 0.92),
1069 span("Jane", 160, 40, 0.92),
1070 span("Email", 20, 70, 0.92),
1071 span("alice@example.invalid", 160, 70, 0.92),
1073 ],
1074 };
1075 let tmp = tempfile::tempdir().expect("tempdir");
1076 let input = tmp.path().join("input.png");
1077 fs::write(&input, b"\x89PNG\r\n\x1A\nnot-real-image").expect("write input");
1078 let (agent_out, owner_out) = bundle_dirs(&tmp);
1079
1080 let bundle = Pipeline::new()
1081 .clean_with_ocr_backend(&input, agent_out, owner_out, &backend)
1082 .expect("clean succeeds");
1083
1084 assert_eq!(bundle.report.pages[0].column_count, 1);
1085 assert!(
1086 bundle.clean_markdown.contains("Field\nValue\n\nBill\nJane"),
1087 "{}",
1088 bundle.clean_markdown
1089 );
1090 assert!(
1091 bundle.clean_markdown.contains(":Email_"),
1092 "{}",
1093 bundle.clean_markdown
1094 );
1095 assert!(
1096 !bundle.clean_markdown.contains("alice@example.invalid"),
1098 "{}",
1099 bundle.clean_markdown
1100 );
1101 }
1102
1103 #[test]
1104 fn clean_char_count_matches_persisted_clean_md_length() {
1105 let backend = MockBackend {
1106 spans: vec![
1107 span("Bill", 20, 10, 0.92),
1108 span("to:", 116, 10, 0.92),
1109 span("Jane", 20, 36, 0.92),
1110 span("Doe", 116, 36, 0.92),
1111 span("Email:", 360, 10, 0.92),
1112 span("alice@example.invalid", 360, 36, 0.92),
1114 ],
1115 };
1116 let tmp = tempfile::tempdir().expect("tempdir");
1117 let input = tmp.path().join("input.png");
1118 fs::write(&input, b"\x89PNG\r\n\x1A\nnot-real-image").expect("write input");
1119 let (agent_out, owner_out) = bundle_dirs(&tmp);
1120 let clean_md_path = agent_out.as_path().join(CLEAN_MARKDOWN_FILE);
1121
1122 let bundle = Pipeline::new()
1123 .clean_with_ocr_backend(&input, agent_out, owner_out, &backend)
1124 .expect("clean succeeds");
1125
1126 let on_disk = fs::read_to_string(&clean_md_path).expect("clean.md readable");
1127 assert_eq!(
1128 bundle.report.clean_char_count,
1129 bundle.clean_markdown.chars().count(),
1130 "clean_char_count must match SafeBundle.clean_markdown, got report={} vs markdown={}",
1131 bundle.report.clean_char_count,
1132 bundle.clean_markdown.chars().count(),
1133 );
1134 assert_eq!(
1135 bundle.report.clean_char_count,
1136 on_disk.chars().count(),
1137 "clean_char_count must match the on-disk clean.md length, got report={} vs disk={}",
1138 bundle.report.clean_char_count,
1139 on_disk.chars().count(),
1140 );
1141 assert!(
1142 bundle
1143 .clean_markdown
1144 .starts_with("# gaze-document safe bundle\n"),
1145 "clean_markdown must carry the bundle header: {}",
1146 bundle.clean_markdown,
1147 );
1148 let body = bundle
1149 .clean_markdown
1150 .strip_prefix("# gaze-document safe bundle\n\nSource kind: `png`\n\n---\n\n")
1151 .expect("header prefix present");
1152 let bare = body.trim_end_matches('\n').chars().count();
1153 assert!(
1154 bundle.report.clean_char_count > bare,
1155 "clean_char_count ({}) must exceed the bare redacted body length ({}) \
1156 because the header + appended trailing newline are part of the persisted artifact",
1157 bundle.report.clean_char_count,
1158 bare,
1159 );
1160 }
1161
1162 #[cfg(feature = "pdf-input")]
1163 #[test]
1164 fn clean_preprocesses_rotated_image_before_backend_ocr() {
1165 use image::{GrayImage, ImageFormat as EncodedImageFormat, Luma};
1166
1167 #[derive(Debug)]
1168 struct OrientationSensitiveBackend;
1169
1170 impl OcrBackend for OrientationSensitiveBackend {
1171 fn name(&self) -> &str {
1172 "orientation-sensitive"
1173 }
1174
1175 fn recognize(
1176 &self,
1177 image: ImageInput,
1178 _hints: OcrHints,
1179 ) -> Result<Vec<OcrSpan>, OcrError> {
1180 let decoded = image::load_from_memory(&image.bytes)
1181 .map_err(|err| OcrError::Internal(err.to_string()))?;
1182 if decoded.width() <= decoded.height() {
1183 return Ok(Vec::new());
1184 }
1185 Ok(vec![span("alice@example.invalid", 20, 20, 0.91)])
1187 }
1188 }
1189
1190 let mut image = GrayImage::from_pixel(120, 80, Luma([255]));
1191 for y in 38..42 {
1192 for x in 16..104 {
1193 image.put_pixel(x, y, Luma([0]));
1194 }
1195 }
1196 let sideways = image::imageops::rotate90(&image);
1197 let mut bytes = Vec::new();
1198 sideways
1199 .write_to(
1200 &mut std::io::Cursor::new(&mut bytes),
1201 EncodedImageFormat::Png,
1202 )
1203 .expect("encode png");
1204 let tmp = tempfile::tempdir().expect("tempdir");
1205 let input = tmp.path().join("input.png");
1206 fs::write(&input, bytes).expect("write input");
1207 let (agent_out, owner_out) = bundle_dirs(&tmp);
1208
1209 let bundle = Pipeline::new()
1210 .clean_with_ocr_backend(&input, agent_out, owner_out, &OrientationSensitiveBackend)
1211 .expect("clean succeeds");
1212
1213 assert!(
1214 bundle.clean_markdown.contains(":Email_"),
1215 "{}",
1216 bundle.clean_markdown
1217 );
1218 assert!(
1219 !bundle.clean_markdown.contains("alice@example.invalid"),
1221 "{}",
1222 bundle.clean_markdown
1223 );
1224 }
1225
1226 #[cfg(feature = "pdf-input")]
1227 #[test]
1228 fn clean_deskews_image_before_backend_ocr() {
1229 use image::{GrayImage, ImageFormat as EncodedImageFormat, Luma};
1230 use imageproc::geometric_transformations::{rotate_about_center, Interpolation};
1231
1232 fn horizontal_score(bytes: &[u8]) -> Result<u64, OcrError> {
1233 let decoded = image::load_from_memory(bytes)
1234 .map_err(|err| OcrError::Internal(err.to_string()))?
1235 .to_luma8();
1236 let mut score = 0u64;
1237 for y in 0..decoded.height() {
1238 let mut dark = 0u64;
1239 for x in 0..decoded.width() {
1240 if decoded.get_pixel(x, y).0[0] < 200 {
1241 dark += 1;
1242 }
1243 }
1244 score = score.saturating_add(dark.saturating_mul(dark));
1245 }
1246 Ok(score)
1247 }
1248
1249 #[derive(Debug)]
1250 struct DeskewSensitiveBackend {
1251 minimum_score: u64,
1252 }
1253
1254 impl OcrBackend for DeskewSensitiveBackend {
1255 fn name(&self) -> &str {
1256 "deskew-sensitive"
1257 }
1258
1259 fn recognize(
1260 &self,
1261 image: ImageInput,
1262 _hints: OcrHints,
1263 ) -> Result<Vec<OcrSpan>, OcrError> {
1264 if horizontal_score(&image.bytes)? < self.minimum_score {
1265 return Ok(Vec::new());
1266 }
1267 Ok(vec![span("alice@example.invalid", 20, 20, 0.91)])
1269 }
1270 }
1271
1272 let mut image = GrayImage::from_pixel(120, 80, Luma([255]));
1273 for y in 38..42 {
1274 for x in 16..104 {
1275 image.put_pixel(x, y, Luma([0]));
1276 }
1277 }
1278 let skewed = rotate_about_center(
1279 &image,
1280 4.0_f32.to_radians(),
1281 Interpolation::Nearest,
1282 Luma([255]),
1283 );
1284 let mut bytes = Vec::new();
1285 skewed
1286 .write_to(
1287 &mut std::io::Cursor::new(&mut bytes),
1288 EncodedImageFormat::Png,
1289 )
1290 .expect("encode png");
1291 let raw_score = horizontal_score(&bytes).expect("raw score");
1292 let backend = DeskewSensitiveBackend {
1293 minimum_score: raw_score + 1_000,
1294 };
1295 assert!(
1296 backend
1297 .recognize(
1298 ImageInput {
1299 bytes: bytes.clone(),
1300 format: ImageFormat::Png,
1301 dpi: None
1302 },
1303 OcrHints::default()
1304 )
1305 .expect("raw recognize")
1306 .is_empty(),
1307 "raw skewed payload should miss before preprocessing"
1308 );
1309 let tmp = tempfile::tempdir().expect("tempdir");
1310 let input = tmp.path().join("input.png");
1311 fs::write(&input, bytes).expect("write input");
1312 let (agent_out, owner_out) = bundle_dirs(&tmp);
1313
1314 let bundle = Pipeline::new()
1315 .clean_with_ocr_backend(&input, agent_out, owner_out, &backend)
1316 .expect("clean succeeds");
1317
1318 assert!(
1319 bundle.clean_markdown.contains(":Email_"),
1320 "{}",
1321 bundle.clean_markdown
1322 );
1323 assert!(
1324 !bundle.clean_markdown.contains("alice@example.invalid"),
1326 "{}",
1327 bundle.clean_markdown
1328 );
1329 }
1330
1331 #[test]
1332 fn format_clean_markdown_appends_trailing_newline() {
1333 let md = format_clean_markdown("hello", InputKind::Png);
1334 assert!(md.ends_with('\n'));
1335 assert!(md.contains("Source kind: `png`"));
1336 assert!(md.contains("hello"));
1337 }
1338}