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
12pub 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 self.body
89 .validate()
90 .map_err(crate::client::error::ZaiError::from)?;
91
92 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 let metadata = std::fs::metadata(p)?;
110 let file_size = metadata.len();
111 const MAX_SIZE: u64 = 8 * 1024 * 1024; 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 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 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}