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
56pub struct OcrRequest {
61 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 pub fn new() -> Self {
75 Self {
76 body: OcrBody::new(),
77 file_path: None,
78 }
79 }
80
81 pub fn with_file_path(mut self, path: impl Into<PathBuf>) -> Self {
83 self.file_path = Some(path.into());
84 self
85 }
86
87 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 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 pub fn with_probability(mut self, probability: bool) -> Self {
101 self.body = self.body.with_probability(probability);
102 self
103 }
104
105 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 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 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 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 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}