Skip to main content

zai_rs/model/ocr/
data.rs

1use std::path::{Path, PathBuf};
2
3use validator::Validate;
4
5use super::request::{OcrBody, OcrLanguageType, OcrToolType};
6use crate::client::ZaiClient;
7use crate::client::error::codes;
8
9const MAX_FILE_SIZE: u64 = 8 * 1024 * 1024;
10
11fn file_error(code: u16, message: impl Into<String>) -> crate::ZaiError {
12    crate::ZaiError::FileError {
13        code,
14        message: message.into(),
15    }
16}
17
18fn image_mime_type(path: &Path) -> crate::ZaiResult<&'static str> {
19    match path
20        .extension()
21        .and_then(|extension| extension.to_str())
22        .map(str::to_ascii_lowercase)
23        .as_deref()
24    {
25        Some("png") => Ok("image/png"),
26        Some("jpg" | "jpeg") => Ok("image/jpeg"),
27        Some("bmp") => Ok("image/bmp"),
28        extension => Err(file_error(
29            codes::SDK_FILE_TYPE_UNSUPPORTED,
30            format!(
31                "unsupported OCR image extension {extension:?}; expected png, jpg, jpeg, or bmp"
32            ),
33        )),
34    }
35}
36
37fn validate_file_metadata(path: &Path, metadata: &std::fs::Metadata) -> crate::ZaiResult<()> {
38    if !metadata.is_file() {
39        return Err(file_error(
40            codes::SDK_FILE_NOT_FOUND,
41            format!("file_path is not a regular file: {}", path.display()),
42        ));
43    }
44    if metadata.len() > MAX_FILE_SIZE {
45        return Err(file_error(
46            codes::SDK_FILE_TOO_LARGE,
47            format!(
48                "OCR image exceeds the 8 MiB limit: {} bytes",
49                metadata.len()
50            ),
51        ));
52    }
53    image_mime_type(path).map(|_| ())
54}
55
56/// OCR recognition request encoded as `multipart/form-data`.
57///
58/// Builder for the OCR endpoint. Set the image via
59/// [`OcrRequest::with_file_path`], then call [`OcrRequest::send_via`].
60pub struct OcrRequest {
61    /// Multipart form fields (tool type, language, …).
62    pub body: OcrBody,
63    file_path: Option<PathBuf>,
64}
65
66impl Default for OcrRequest {
67    fn default() -> Self {
68        Self::new()
69    }
70}
71
72impl OcrRequest {
73    /// Create a new OCR request with default options.
74    pub fn new() -> Self {
75        Self {
76            body: OcrBody::new(),
77            file_path: None,
78        }
79    }
80
81    /// Set the local image file path to recognize (required).
82    pub fn with_file_path(mut self, path: impl Into<PathBuf>) -> Self {
83        self.file_path = Some(path.into());
84        self
85    }
86
87    /// Set the OCR tool type (e.g. handwriting).
88    pub fn with_tool_type(mut self, tool_type: OcrToolType) -> Self {
89        self.body = self.body.with_tool_type(tool_type);
90        self
91    }
92
93    /// Set the document language.
94    pub fn with_language_type(mut self, language_type: OcrLanguageType) -> Self {
95        self.body = self.body.with_language_type(language_type);
96        self
97    }
98
99    /// Whether to return per-character recognition probabilities.
100    pub fn with_probability(mut self, probability: bool) -> Self {
101        self.body = self.body.with_probability(probability);
102        self
103    }
104
105    /// Set the client-side request id.
106    pub fn with_request_id(mut self, request_id: impl Into<String>) -> Self {
107        self.body = self.body.with_request_id(request_id);
108        self
109    }
110
111    /// Set the end-user id.
112    pub fn with_user_id(mut self, user_id: impl Into<String>) -> Self {
113        self.body = self.body.with_user_id(user_id);
114        self
115    }
116
117    /// Validate the request: body constraints, file existence, size (≤8MB),
118    /// and supported extension (png/jpg/jpeg/bmp).
119    pub fn validate(&self) -> crate::ZaiResult<()> {
120        self.body
121            .validate()
122            .map_err(crate::client::error::ZaiError::from)?;
123        let path = self.required_file_path()?;
124        let metadata = std::fs::symlink_metadata(path).map_err(|error| {
125            if error.kind() == std::io::ErrorKind::NotFound {
126                file_error(
127                    codes::SDK_FILE_NOT_FOUND,
128                    format!("file_path not found: {}", path.display()),
129                )
130            } else {
131                crate::ZaiError::from(error)
132            }
133        })?;
134        validate_file_metadata(path, &metadata)
135    }
136
137    /// Async counterpart of the file probes in [`validate`](Self::validate).
138    ///
139    /// [`validate`](Self::validate) is a public **sync** fn (kept that way for
140    /// backward compatibility) and uses blocking `std::fs` stats; calling it from
141    /// `async fn send_via` would stall the executor on a slow/networked filesystem.
142    /// This helper performs the same existence/size/extension checks via
143    /// `tokio::fs` so the async path stays non-blocking. Keep the two in sync.
144    async fn validate_file_async(&self) -> crate::ZaiResult<&Path> {
145        let path = self.required_file_path()?;
146        let metadata = match tokio::fs::symlink_metadata(path).await {
147            Ok(m) => m,
148            Err(e) if e.kind() == std::io::ErrorKind::NotFound => {
149                return Err(file_error(
150                    codes::SDK_FILE_NOT_FOUND,
151                    format!("file_path not found: {}", path.display()),
152                ));
153            },
154            Err(e) => return Err(crate::ZaiError::from(e)),
155        };
156        validate_file_metadata(path, &metadata)?;
157        Ok(path)
158    }
159
160    fn required_file_path(&self) -> crate::ZaiResult<&Path> {
161        self.file_path
162            .as_deref()
163            .filter(|path| !path.as_os_str().is_empty())
164            .ok_or_else(|| crate::client::validation::invalid("file_path is required"))
165    }
166
167    /// Submit the multipart request via a [`ZaiClient`] and parse the typed
168    /// OCR response.
169    pub async fn send_via(
170        &self,
171        client: &ZaiClient,
172    ) -> crate::ZaiResult<super::response::OcrResponse> {
173        self.body
174            .validate()
175            .map_err(crate::client::error::ZaiError::from)?;
176        let file_path = self.validate_file_async().await?;
177
178        let route = crate::client::routes::FILES_OCR;
179        let url = client.endpoints().resolve_route(route, &[])?;
180        let mime = image_mime_type(file_path)?;
181        let file_part = crate::client::transport::multipart::FilePart::from_path(file_path)?
182            .with_content_type(mime)?;
183        let tool_type = self
184            .body
185            .tool_type
186            .unwrap_or(OcrToolType::HandWrite)
187            .as_str();
188        let language_type = self.body.language_type.map(OcrLanguageType::as_str);
189        let probability = self.body.probability;
190        let request_id = self.body.request_id.clone();
191        let user_id = self.body.user_id.clone();
192
193        let mut factory = crate::client::transport::multipart::MultipartBodyFactory::new()
194            .field("tool_type", tool_type)?
195            .file_named("file", file_part)?;
196        if let Some(lang) = language_type {
197            factory = factory.field("language_type", lang)?;
198        }
199        if let Some(prob) = probability {
200            factory = factory.field("probability", prob.to_string())?;
201        }
202        if let Some(rid) = request_id {
203            factory = factory.field("request_id", rid)?;
204        }
205        if let Some(uid) = user_id {
206            factory = factory.field("user_id", uid)?;
207        }
208        client
209            .send_multipart::<super::response::OcrResponse>(route.method(), url, &factory)
210            .await
211    }
212}