squigit_storage/threads/
ocr.rs

1// Copyright 2026 a7mddra
2// SPDX-License-Identifier: Apache-2.0
3
4use std::fs;
5
6use crate::error::{Result, StorageError};
7
8use super::paths::ocr_annotations_path;
9use super::{
10    default_ocr_annotations, OcrAnnotationEntry, OcrAnnotations, OcrModelAnnotation, OcrRegion,
11    ThreadStorage, EMPTY_STATE_ASSET_ID,
12};
13
14fn is_supported_ocr_model_id(model_id: &str) -> bool {
15    matches!(
16        model_id,
17        "pp-ocr-v5-en"
18            | "pp-ocr-v5-latin"
19            | "pp-ocr-v5-cyrillic"
20            | "pp-ocr-v5-korean"
21            | "pp-ocr-v5-cjk"
22            | "pp-ocr-v5-devanagari"
23    )
24}
25
26fn canonicalize_ocr_annotations_id(model_id: &str) -> Option<&str> {
27    let trimmed = model_id.trim();
28    if trimmed.is_empty() {
29        return None;
30    }
31    if is_supported_ocr_model_id(trimmed) {
32        return Some(trimmed);
33    }
34    None
35}
36
37pub(super) fn retain_supported_ocr_annotations_ids(annotations: &mut OcrAnnotations) -> bool {
38    let unsupported_keys: Vec<String> = annotations
39        .keys()
40        .filter(|key| key.as_str() != EMPTY_STATE_ASSET_ID && !is_supported_ocr_model_id(key))
41        .cloned()
42        .collect();
43
44    for key in &unsupported_keys {
45        annotations.remove(key);
46    }
47
48    !unsupported_keys.is_empty()
49}
50
51pub(super) fn ensure_empty_state_asset(annotations: &mut OcrAnnotations) -> bool {
52    if matches!(
53        annotations.get(EMPTY_STATE_ASSET_ID),
54        Some(OcrAnnotationEntry::EmptyState(_))
55    ) {
56        return false;
57    }
58
59    annotations.insert(
60        EMPTY_STATE_ASSET_ID.to_string(),
61        OcrAnnotationEntry::EmptyState(Vec::new()),
62    );
63    true
64}
65
66impl ThreadStorage {
67    /// Save OCR data for a specific model into the thread's OCR annotations.
68    pub fn save_ocr_data(
69        &self,
70        thread_id: &str,
71        model_id: &str,
72        ocr_data: &[OcrRegion],
73    ) -> Result<()> {
74        let thread_dir = self.thread_dir(thread_id);
75        fs::create_dir_all(&thread_dir)?;
76        let canonical_model_id = canonicalize_ocr_annotations_id(model_id)
77            .ok_or_else(|| StorageError::InvalidOcrModel(model_id.to_string()))?;
78
79        let ocr_path = ocr_annotations_path(&thread_dir);
80        let mut annotations: OcrAnnotations = if ocr_path.exists() {
81            let json = fs::read_to_string(&ocr_path)?;
82            serde_json::from_str(&json)?
83        } else {
84            default_ocr_annotations()
85        };
86        ensure_empty_state_asset(&mut annotations);
87        retain_supported_ocr_annotations_ids(&mut annotations);
88
89        annotations.insert(
90            canonical_model_id.to_string(),
91            OcrAnnotationEntry::Model(OcrModelAnnotation {
92                scanned_at: Some(chrono::Utc::now()),
93                ocr_data: ocr_data.to_vec(),
94            }),
95        );
96
97        super::atomic_write(
98            &ocr_path,
99            serde_json::to_string_pretty(&annotations)?.as_bytes(),
100        )?;
101        Ok(())
102    }
103
104    /// Get the entire OCR annotations for a thread.
105    pub fn get_ocr_annotations(&self, thread_id: &str) -> Result<OcrAnnotations> {
106        let thread_dir = self.thread_dir(thread_id);
107        let ocr_path = ocr_annotations_path(&thread_dir);
108
109        if !ocr_path.exists() {
110            return Ok(default_ocr_annotations());
111        }
112
113        let json = fs::read_to_string(&ocr_path)?;
114        let mut annotations: OcrAnnotations = serde_json::from_str(&json)?;
115        let mut annotations_changed = ensure_empty_state_asset(&mut annotations);
116        if retain_supported_ocr_annotations_ids(&mut annotations) {
117            annotations_changed = true;
118        }
119        if annotations_changed {
120            super::atomic_write(
121                &ocr_path,
122                serde_json::to_string_pretty(&annotations)?.as_bytes(),
123            )?;
124        }
125        Ok(annotations)
126    }
127}
squigit_storage/threads/ocr.rs

squigit_storage/threads/
ocr.rs
