a3s_use_ocr/
client.rs

1use std::path::Path;
2use std::sync::Arc;
3
4use a3s_use_core::{Artifact, UseError, UseResult};
5use sha2::{Digest, Sha256};
6use tokio::io::AsyncReadExt;
7
8use crate::models::{OcrDiagnostic, OcrRequest, OcrResult};
9use crate::provider::{OcrInput, OcrProvider, OcrProviderDescriptor, OcrProviderOutput};
10
11const MAX_INPUT_BYTES: u64 = 32 * 1024 * 1024;
12
13/// Provider-neutral OCR client.
14///
15/// The client owns source validation and evidence. Recognition is delegated to
16/// one injected [`OcrProvider`].
17#[derive(Clone)]
18pub struct OcrClient {
19    provider: Arc<dyn OcrProvider>,
20    descriptor: OcrProviderDescriptor,
21}
22
23impl OcrClient {
24    pub fn with_provider<P>(provider: P) -> UseResult<Self>
25    where
26        P: OcrProvider + 'static,
27    {
28        Self::from_provider(Arc::new(provider))
29    }
30
31    pub fn from_provider(provider: Arc<dyn OcrProvider>) -> UseResult<Self> {
32        let descriptor = provider.descriptor();
33        descriptor.validate()?;
34        Ok(Self {
35            provider,
36            descriptor,
37        })
38    }
39
40    #[cfg(feature = "ppocr-v6")]
41    pub fn from_env() -> UseResult<Self> {
42        Self::with_provider(crate::PpOcrV6Provider::from_env()?)
43    }
44
45    pub fn provider(&self) -> &OcrProviderDescriptor {
46        &self.descriptor
47    }
48
49    pub fn diagnostic(&self) -> OcrDiagnostic {
50        let status = self.provider.diagnostic();
51        OcrDiagnostic {
52            readiness: status.readiness,
53            provider: Some(self.descriptor.id.clone()),
54            engine: Some(self.descriptor.engine.clone()),
55            model: status.model,
56            model_dir: status.model_dir,
57            sends_source_off_device: self.descriptor.sends_source_off_device,
58            message: status.message,
59            suggestions: status.suggestions,
60        }
61    }
62
63    pub async fn extract(&self, request: OcrRequest) -> UseResult<OcrResult> {
64        let input = read_source(&request.path).await?;
65        let source = input.source().clone();
66        let output = self.provider.recognize(input).await?;
67        validate_provider_output(&output)?;
68        Ok(OcrResult {
69            provider: self.descriptor.id.clone(),
70            engine: self.descriptor.engine.clone(),
71            model: output.model,
72            source,
73            text: output.text,
74            blocks: output.blocks,
75            warnings: output.warnings,
76        })
77    }
78}
79
80fn validate_provider_output(output: &OcrProviderOutput) -> UseResult<()> {
81    for block in &output.blocks {
82        if block.page == 0 {
83            return Err(provider_output_error(
84                "OCR providers must number pages starting at 1.",
85            ));
86        }
87        for (name, value) in [
88            ("confidence", block.confidence),
89            ("detection confidence", block.detection_confidence),
90        ] {
91            if value.is_some_and(|value| !value.is_finite() || !(0.0..=1.0).contains(&value)) {
92                return Err(provider_output_error(format!(
93                    "OCR provider {name} must be between 0 and 1."
94                )));
95            }
96        }
97    }
98    Ok(())
99}
100
101fn provider_output_error(message: impl Into<String>) -> UseError {
102    UseError::new("use.ocr.provider_output_invalid", message)
103}
104
105async fn read_source(path: &Path) -> UseResult<OcrInput> {
106    let canonical = tokio::fs::canonicalize(path).await.map_err(|error| {
107        UseError::new(
108            "use.ocr.source_unreadable",
109            format!("Failed to resolve OCR source '{}': {error}", path.display()),
110        )
111    })?;
112    let metadata = tokio::fs::metadata(&canonical).await.map_err(|error| {
113        UseError::new(
114            "use.ocr.source_unreadable",
115            format!(
116                "Failed to inspect OCR source '{}': {error}",
117                canonical.display()
118            ),
119        )
120    })?;
121    if !metadata.is_file() {
122        return Err(UseError::new(
123            "use.ocr.source_invalid",
124            format!(
125                "OCR source '{}' is not a regular file.",
126                canonical.display()
127            ),
128        ));
129    }
130    if metadata.len() == 0 || metadata.len() > MAX_INPUT_BYTES {
131        return Err(UseError::new(
132            "use.ocr.source_too_large",
133            format!(
134                "OCR source '{}' must contain between 1 byte and 32 MiB.",
135                canonical.display()
136            ),
137        )
138        .with_detail("size", metadata.len()));
139    }
140    let file = tokio::fs::File::open(&canonical).await.map_err(|error| {
141        UseError::new(
142            "use.ocr.source_unreadable",
143            format!(
144                "Failed to open OCR source '{}': {error}",
145                canonical.display()
146            ),
147        )
148    })?;
149    let mut bytes = Vec::with_capacity(metadata.len().min(MAX_INPUT_BYTES) as usize);
150    file.take(MAX_INPUT_BYTES + 1)
151        .read_to_end(&mut bytes)
152        .await
153        .map_err(|error| {
154            UseError::new(
155                "use.ocr.source_unreadable",
156                format!(
157                    "Failed to read OCR source '{}': {error}",
158                    canonical.display()
159                ),
160            )
161        })?;
162    if bytes.len() as u64 > MAX_INPUT_BYTES {
163        return Err(UseError::new(
164            "use.ocr.source_too_large",
165            format!(
166                "OCR source '{}' must not exceed 32 MiB.",
167                canonical.display()
168            ),
169        )
170        .with_detail("sizeAtLeast", MAX_INPUT_BYTES + 1));
171    }
172    let media_type = detect_image_type(&bytes).ok_or_else(|| {
173        UseError::new(
174            "use.ocr.source_type_unsupported",
175            "OCR accepts PNG, JPEG, WebP, GIF, BMP, and TIFF image bytes.",
176        )
177    })?;
178    let digest = Sha256::digest(&bytes);
179    Ok(OcrInput::new(
180        Artifact {
181            path: canonical,
182            media_type: media_type.to_string(),
183            size: bytes.len() as u64,
184            sha256: format!("{digest:x}"),
185        },
186        bytes,
187    ))
188}
189
190fn detect_image_type(bytes: &[u8]) -> Option<&'static str> {
191    if bytes.starts_with(b"\x89PNG\r\n\x1a\n") {
192        Some("image/png")
193    } else if bytes.starts_with(b"\xff\xd8\xff") {
194        Some("image/jpeg")
195    } else if bytes.starts_with(b"GIF87a") || bytes.starts_with(b"GIF89a") {
196        Some("image/gif")
197    } else if bytes.starts_with(b"BM") {
198        Some("image/bmp")
199    } else if bytes.starts_with(b"II*\0") || bytes.starts_with(b"MM\0*") {
200        Some("image/tiff")
201    } else if bytes.len() >= 12 && bytes.starts_with(b"RIFF") && &bytes[8..12] == b"WEBP" {
202        Some("image/webp")
203    } else {
204        None
205    }
206}
207
208#[cfg(test)]
209mod tests {
210    use a3s_use_core::Readiness;
211    use async_trait::async_trait;
212
213    use crate::{OcrBlock, OcrProviderStatus};
214
215    use super::*;
216
217    struct FixtureProvider;
218
219    #[async_trait]
220    impl OcrProvider for FixtureProvider {
221        fn descriptor(&self) -> OcrProviderDescriptor {
222            OcrProviderDescriptor::new("fixture-tesseract", "tesseract", true).unwrap()
223        }
224
225        fn diagnostic(&self) -> OcrProviderStatus {
226            OcrProviderStatus {
227                readiness: Readiness::Ready,
228                model: Some("eng".to_string()),
229                model_dir: None,
230                message: "ready".to_string(),
231                suggestions: Vec::new(),
232            }
233        }
234
235        async fn recognize(&self, input: OcrInput) -> UseResult<OcrProviderOutput> {
236            assert_eq!(input.source().media_type, "image/bmp");
237            assert!(input.bytes().starts_with(b"BM"));
238            Ok(OcrProviderOutput {
239                model: Some("eng".to_string()),
240                text: "fixture text".to_string(),
241                blocks: vec![OcrBlock {
242                    page: 1,
243                    text: "fixture text".to_string(),
244                    confidence: None,
245                    detection_confidence: None,
246                    polygon: None,
247                    bounding_box: None,
248                }],
249                warnings: Vec::new(),
250            })
251        }
252    }
253
254    #[test]
255    fn detects_supported_image_signatures() {
256        assert_eq!(
257            detect_image_type(b"\x89PNG\r\n\x1a\nrest"),
258            Some("image/png")
259        );
260        assert_eq!(detect_image_type(b"\xff\xd8\xffrest"), Some("image/jpeg"));
261        assert_eq!(detect_image_type(b"not an image"), None);
262    }
263
264    #[tokio::test]
265    async fn injected_provider_receives_validated_source_and_owns_recognition() {
266        let temp = tempfile::tempdir().unwrap();
267        let source = temp.path().join("scan.bmp");
268        std::fs::write(&source, base64_free_white_bmp()).unwrap();
269        let client = OcrClient::with_provider(FixtureProvider).unwrap();
270
271        let result = client.extract(OcrRequest { path: source }).await.unwrap();
272
273        assert_eq!(result.provider, "fixture-tesseract");
274        assert_eq!(result.engine, "tesseract");
275        assert_eq!(result.model.as_deref(), Some("eng"));
276        assert_eq!(result.text, "fixture text");
277        assert_eq!(result.blocks[0].confidence, None);
278        assert!(client.diagnostic().sends_source_off_device);
279    }
280
281    fn base64_free_white_bmp() -> Vec<u8> {
282        vec![
283            0x42, 0x4d, 0x3a, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x36, 0x00, 0x00, 0x00,
284            0x28, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x01, 0x00,
285            0x18, 0x00, 0x00, 0x00, 0x00, 0x00, 0x04, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
286            0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0xff, 0xff, 0xff, 0x00,
287        ]
288    }
289}
a3s_use_ocr/client.rs

a3s_use_ocr/
client.rs
