Skip to main content

zai_rs/model/ocr/
data.rs

1use std::{path::Path, sync::Arc};
2
3use validator::Validate;
4
5use super::request::{OcrBody, OcrLanguageType, OcrToolType};
6use crate::client::{
7    endpoints::{ApiBase, EndpointConfig, paths},
8    error::codes,
9    http::{HttpClient, HttpClientConfig, parse_typed_response, send_multipart_request},
10};
11
12/// OCR recognition request (multipart/form-data)
13pub struct OcrRequest {
14    pub key: String,
15    url: String,
16    endpoint_config: EndpointConfig,
17    api_base: ApiBase,
18    http_config: Arc<HttpClientConfig>,
19    pub body: OcrBody,
20    file_path: Option<String>,
21}
22
23impl OcrRequest {
24    pub fn new(key: String) -> Self {
25        let endpoint_config = EndpointConfig::default();
26        let api_base = ApiBase::PaasV4;
27        let url = endpoint_config.url(&api_base, paths::FILES_OCR);
28        Self {
29            key,
30            url,
31            endpoint_config,
32            api_base,
33            http_config: Arc::new(HttpClientConfig::default()),
34            body: OcrBody::new(),
35            file_path: None,
36        }
37    }
38
39    pub fn with_file_path(mut self, path: impl Into<String>) -> Self {
40        self.file_path = Some(path.into());
41        self
42    }
43
44    pub fn with_tool_type(mut self, tool_type: OcrToolType) -> Self {
45        self.body = self.body.with_tool_type(tool_type);
46        self
47    }
48
49    pub fn with_language_type(mut self, language_type: OcrLanguageType) -> Self {
50        self.body = self.body.with_language_type(language_type);
51        self
52    }
53
54    pub fn with_probability(mut self, probability: bool) -> Self {
55        self.body = self.body.with_probability(probability);
56        self
57    }
58
59    pub fn with_request_id(mut self, request_id: impl Into<String>) -> Self {
60        self.body = self.body.with_request_id(request_id);
61        self
62    }
63
64    pub fn with_user_id(mut self, user_id: impl Into<String>) -> Self {
65        self.body = self.body.with_user_id(user_id);
66        self
67    }
68
69    pub fn with_base_url(mut self, base_url: impl Into<String>) -> Self {
70        self.api_base = ApiBase::Custom(base_url.into());
71        self.url = self.endpoint_config.url(&self.api_base, paths::FILES_OCR);
72        self
73    }
74
75    pub fn with_endpoint_config(mut self, endpoint_config: EndpointConfig) -> Self {
76        self.endpoint_config = endpoint_config;
77        self.url = self.endpoint_config.url(&self.api_base, paths::FILES_OCR);
78        self
79    }
80
81    pub fn with_http_config(mut self, config: HttpClientConfig) -> Self {
82        self.http_config = Arc::new(config);
83        self
84    }
85
86    pub fn validate(&self) -> crate::ZaiResult<()> {
87        // Check body constraints
88        self.body
89            .validate()
90            .map_err(crate::client::error::ZaiError::from)?;
91
92        // Ensure file path exists
93        let p =
94            self.file_path
95                .as_ref()
96                .ok_or_else(|| crate::client::error::ZaiError::ApiError {
97                    code: 1200,
98                    message: "file_path is required".to_string(),
99                })?;
100
101        if !Path::new(p).exists() {
102            return Err(crate::client::error::ZaiError::FileError {
103                code: codes::SDK_FILE_NOT_FOUND,
104                message: format!("file_path not found: {}", p),
105            });
106        }
107
108        // Validate file size (max 8MB)
109        let metadata = std::fs::metadata(p)?;
110        let file_size = metadata.len();
111        const MAX_SIZE: u64 = 8 * 1024 * 1024; // 8MB
112        if file_size > MAX_SIZE {
113            return Err(crate::client::error::ZaiError::FileError {
114                code: codes::SDK_FILE_TOO_LARGE,
115                message: format!("file_size exceeds 8MB limit: {} bytes", file_size),
116            });
117        }
118
119        // Validate file extension
120        let ext = Path::new(p)
121            .extension()
122            .and_then(|s| s.to_str())
123            .map(|s| s.to_ascii_lowercase());
124        let valid_ext = matches!(
125            ext.as_deref(),
126            Some("png") | Some("jpg") | Some("jpeg") | Some("bmp")
127        );
128        if !valid_ext {
129            return Err(crate::client::error::ZaiError::FileError {
130                code: codes::SDK_FILE_TYPE_UNSUPPORTED,
131                message: format!(
132                    "invalid file format: {:?}. Only PNG, JPG, JPEG, BMP are supported",
133                    ext
134                ),
135            });
136        }
137
138        Ok(())
139    }
140
141    pub async fn send(&self) -> crate::ZaiResult<super::response::OcrResponse> {
142        self.validate()?;
143
144        let resp = self.post().await?;
145
146        parse_typed_response::<super::response::OcrResponse>(resp).await
147    }
148}
149
150impl HttpClient for OcrRequest {
151    type Body = OcrBody;
152    type ApiUrl = String;
153    type ApiKey = String;
154
155    fn api_url(&self) -> &Self::ApiUrl {
156        &self.url
157    }
158
159    fn api_key(&self) -> &Self::ApiKey {
160        &self.key
161    }
162
163    fn body(&self) -> &Self::Body {
164        &self.body
165    }
166
167    fn post(
168        &self,
169    ) -> impl std::future::Future<Output = crate::ZaiResult<reqwest::Response>> + Send {
170        let key = self.key.clone();
171        let url = self.url.clone();
172        let config = self.http_config.clone();
173        let body = self.body.clone();
174        let file_path_opt = self.file_path.clone();
175
176        async move {
177            let file_path =
178                file_path_opt.ok_or_else(|| crate::client::error::ZaiError::ApiError {
179                    code: 1200,
180                    message: "file_path is required".to_string(),
181                })?;
182
183            let file_name = Path::new(&file_path)
184                .file_name()
185                .and_then(|s| s.to_str())
186                .unwrap_or("image.png")
187                .to_string();
188            let file_bytes = tokio::fs::read(&file_path).await?;
189
190            // Determine MIME type by extension
191            let ext = Path::new(&file_path)
192                .extension()
193                .and_then(|s| s.to_str())
194                .map(|s| s.to_ascii_lowercase());
195            let mime = match ext.as_deref() {
196                Some("png") => "image/png",
197                Some("jpg") | Some("jpeg") => "image/jpeg",
198                Some("bmp") => "image/bmp",
199                _ => "image/png",
200            };
201
202            let tool_type_str = match &body.tool_type {
203                Some(OcrToolType::HandWrite) => "hand_write",
204                None => "hand_write",
205            }
206            .to_string();
207
208            let language_type = body.language_type.as_ref().map(|lang| {
209                serde_json::to_string(lang)
210                    .unwrap_or_default()
211                    .trim_matches('"')
212                    .to_string()
213            });
214            let probability = body.probability;
215            let request_id = body.request_id.clone();
216            let user_id = body.user_id.clone();
217
218            send_multipart_request(reqwest::Method::POST, url, key, config, move || {
219                let part = reqwest::multipart::Part::bytes(file_bytes.clone())
220                    .file_name(file_name.clone())
221                    .mime_str(mime)?;
222                let mut form = reqwest::multipart::Form::new()
223                    .part("file", part)
224                    .text("tool_type", tool_type_str.clone());
225                if let Some(lang) = language_type.as_ref() {
226                    form = form.text("language_type", lang.clone());
227                }
228                if let Some(prob) = probability {
229                    form = form.text("probability", prob.to_string());
230                }
231                if let Some(rid) = request_id.as_ref() {
232                    form = form.text("request_id", rid.clone());
233                }
234                if let Some(uid) = user_id.as_ref() {
235                    form = form.text("user_id", uid.clone());
236                }
237                Ok(form)
238            })
239            .await
240        }
241    }
242
243    fn http_config(&self) -> Arc<HttpClientConfig> {
244        self.http_config.clone()
245    }
246}