Skip to main content

zai_rs/services/tools/
request.rs

1//! Typed requests for document layout parsing and web-page reading.
2
3use serde::Serialize;
4
5use crate::{
6    ZaiResult,
7    client::{
8        ZaiClient,
9        error::{ZaiError, codes},
10    },
11};
12
13use super::response::{LayoutParsingResponse, ReaderResponse};
14
15fn validation_error(message: impl Into<String>) -> ZaiError {
16    ZaiError::ApiError {
17        code: codes::SDK_VALIDATION,
18        message: message.into(),
19    }
20}
21
22fn invalid_optional_length(value: Option<&str>, min: usize, max: usize) -> bool {
23    value.is_some_and(|value| !(min..=max).contains(&value.chars().count()))
24}
25
26#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)]
27enum LayoutParsingModel {
28    #[serde(rename = "glm-ocr")]
29    GlmOcr,
30}
31
32/// Request for `POST /layout_parsing`.
33#[derive(Clone, Serialize)]
34pub struct LayoutParsingRequest {
35    model: LayoutParsingModel,
36    file: String,
37    #[serde(skip_serializing_if = "Option::is_none")]
38    return_crop_images: Option<bool>,
39    #[serde(skip_serializing_if = "Option::is_none")]
40    need_layout_visualization: Option<bool>,
41    #[serde(skip_serializing_if = "Option::is_none")]
42    start_page_id: Option<u32>,
43    #[serde(skip_serializing_if = "Option::is_none")]
44    end_page_id: Option<u32>,
45    #[serde(skip_serializing_if = "Option::is_none")]
46    request_id: Option<String>,
47    #[serde(skip_serializing_if = "Option::is_none")]
48    user_id: Option<String>,
49}
50
51impl std::fmt::Debug for LayoutParsingRequest {
52    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
53        formatter
54            .debug_struct("LayoutParsingRequest")
55            .field("model", &self.model())
56            .field("file", &"[REDACTED]")
57            .field("return_crop_images", &self.return_crop_images)
58            .field("need_layout_visualization", &self.need_layout_visualization)
59            .field("start_page_id", &self.start_page_id)
60            .field("end_page_id", &self.end_page_id)
61            .field(
62                "request_id",
63                &self.request_id.as_ref().map(|_| "[REDACTED]"),
64            )
65            .field("user_id", &self.user_id.as_ref().map(|_| "[REDACTED]"))
66            .finish()
67    }
68}
69
70impl LayoutParsingRequest {
71    /// Create a request for a URL or base64-encoded PDF/JPG/PNG input.
72    pub fn new(file: impl Into<String>) -> Self {
73        Self {
74            model: LayoutParsingModel::GlmOcr,
75            file: file.into(),
76            return_crop_images: None,
77            need_layout_visualization: None,
78            start_page_id: None,
79            end_page_id: None,
80            request_id: None,
81            user_id: None,
82        }
83    }
84
85    /// Configure cropped-image output.
86    pub fn with_crop_images(mut self, enabled: bool) -> Self {
87        self.return_crop_images = Some(enabled);
88        self
89    }
90
91    /// Configure layout-visualization output.
92    pub fn with_layout_visualization(mut self, enabled: bool) -> Self {
93        self.need_layout_visualization = Some(enabled);
94        self
95    }
96
97    /// Select the inclusive PDF page range.
98    pub fn with_page_range(mut self, start: u32, end: u32) -> Self {
99        self.start_page_id = Some(start);
100        self.end_page_id = Some(end);
101        self
102    }
103
104    /// Set only the first PDF page to parse.
105    pub fn with_start_page_id(mut self, start: u32) -> Self {
106        self.start_page_id = Some(start);
107        self
108    }
109
110    /// Set only the last PDF page to parse.
111    pub fn with_end_page_id(mut self, end: u32) -> Self {
112        self.end_page_id = Some(end);
113        self
114    }
115
116    /// Set the caller-provided request identifier.
117    pub fn with_request_id(mut self, request_id: impl Into<String>) -> Self {
118        self.request_id = Some(request_id.into());
119        self
120    }
121
122    /// Set the end-user identifier used for abuse monitoring.
123    pub fn with_user_id(mut self, user_id: impl Into<String>) -> Self {
124        self.user_id = Some(user_id.into());
125        self
126    }
127
128    /// Return the fixed model identifier sent on the wire.
129    pub const fn model(&self) -> &'static str {
130        "glm-ocr"
131    }
132
133    /// Borrow the URL or base64 input.
134    pub fn file(&self) -> &str {
135        &self.file
136    }
137
138    /// Return the selected PDF page range.
139    pub const fn page_range(&self) -> Option<(u32, u32)> {
140        match (self.start_page_id, self.end_page_id) {
141            (Some(start), Some(end)) => Some((start, end)),
142            _ => None,
143        }
144    }
145
146    /// Return the optional first PDF page.
147    pub const fn start_page_id(&self) -> Option<u32> {
148        self.start_page_id
149    }
150
151    /// Return the optional last PDF page.
152    pub const fn end_page_id(&self) -> Option<u32> {
153        self.end_page_id
154    }
155
156    /// Return the crop-image flag.
157    pub const fn crop_images_enabled(&self) -> Option<bool> {
158        self.return_crop_images
159    }
160
161    /// Return the layout-visualization flag.
162    pub const fn layout_visualization_enabled(&self) -> Option<bool> {
163        self.need_layout_visualization
164    }
165
166    /// Borrow the caller-provided request identifier.
167    pub fn request_id(&self) -> Option<&str> {
168        self.request_id.as_deref()
169    }
170
171    /// Borrow the end-user identifier.
172    pub fn user_id(&self) -> Option<&str> {
173        self.user_id.as_deref()
174    }
175
176    /// Validate frozen OpenAPI numeric and identifier constraints.
177    pub fn validate(&self) -> ZaiResult<()> {
178        if self.file.trim().is_empty() {
179            return Err(validation_error("layout parsing file must not be blank"));
180        }
181        if self.start_page_id == Some(0) || self.end_page_id == Some(0) {
182            return Err(validation_error("layout page ids must be at least 1"));
183        }
184        if matches!(
185            (self.start_page_id, self.end_page_id),
186            (Some(start), Some(end)) if start > end
187        ) {
188            return Err(validation_error(
189                "layout start_page_id must not exceed end_page_id",
190            ));
191        }
192        if invalid_optional_length(self.request_id.as_deref(), 6, 64) {
193            return Err(validation_error(
194                "layout request_id must contain between 6 and 64 characters",
195            ));
196        }
197        if invalid_optional_length(self.user_id.as_deref(), 6, 128) {
198            return Err(validation_error(
199                "layout user_id must contain between 6 and 128 characters",
200            ));
201        }
202        Ok(())
203    }
204
205    /// Validate, send, and decode a layout-parsing response.
206    pub async fn send_via(&self, client: &ZaiClient) -> ZaiResult<LayoutParsingResponse> {
207        self.validate()?;
208        let route = crate::client::routes::TOOLS_LAYOUT;
209        let url = client.endpoints().resolve_route(route, &[])?;
210        let response = client
211            .send_json::<_, LayoutParsingResponse>(route.method(), url, self)
212            .await?;
213        response.validate()?;
214        Ok(response)
215    }
216}
217
218/// Request for `POST /reader`.
219#[derive(Clone, Serialize)]
220pub struct ReaderRequest {
221    url: String,
222    #[serde(skip_serializing_if = "Option::is_none")]
223    timeout: Option<u64>,
224    #[serde(skip_serializing_if = "Option::is_none")]
225    no_cache: Option<bool>,
226    #[serde(skip_serializing_if = "Option::is_none")]
227    return_format: Option<String>,
228    #[serde(skip_serializing_if = "Option::is_none")]
229    retain_images: Option<bool>,
230    #[serde(skip_serializing_if = "Option::is_none")]
231    no_gfm: Option<bool>,
232    #[serde(skip_serializing_if = "Option::is_none")]
233    keep_img_data_url: Option<bool>,
234    #[serde(skip_serializing_if = "Option::is_none")]
235    with_images_summary: Option<bool>,
236    #[serde(skip_serializing_if = "Option::is_none")]
237    with_links_summary: Option<bool>,
238}
239
240impl std::fmt::Debug for ReaderRequest {
241    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
242        formatter
243            .debug_struct("ReaderRequest")
244            .field("url", &"[REDACTED]")
245            .field("timeout", &self.timeout)
246            .field("no_cache", &self.no_cache)
247            .field("return_format", &self.return_format)
248            .field("retain_images", &self.retain_images)
249            .field("no_gfm", &self.no_gfm)
250            .field("keep_img_data_url", &self.keep_img_data_url)
251            .field("with_images_summary", &self.with_images_summary)
252            .field("with_links_summary", &self.with_links_summary)
253            .finish()
254    }
255}
256
257impl ReaderRequest {
258    /// Create a reader request for the required URL.
259    pub fn new(url: impl Into<String>) -> Self {
260        Self {
261            url: url.into(),
262            timeout: None,
263            no_cache: None,
264            return_format: None,
265            retain_images: None,
266            no_gfm: None,
267            keep_img_data_url: None,
268            with_images_summary: None,
269            with_links_summary: None,
270        }
271    }
272
273    /// Set the upstream request timeout in seconds.
274    pub fn with_timeout_seconds(mut self, timeout: u64) -> Self {
275        self.timeout = Some(timeout);
276        self
277    }
278
279    /// Configure cache bypass.
280    pub fn with_cache_disabled(mut self, disabled: bool) -> Self {
281        self.no_cache = Some(disabled);
282        self
283    }
284
285    /// Set the requested output format, such as `markdown` or `text`.
286    pub fn with_return_format(mut self, format: impl Into<String>) -> Self {
287        self.return_format = Some(format.into());
288        self
289    }
290
291    /// Configure whether images remain in the parsed content.
292    pub fn with_retained_images(mut self, retain: bool) -> Self {
293        self.retain_images = Some(retain);
294        self
295    }
296
297    /// Configure GitHub-Flavored Markdown conversion.
298    pub fn with_gfm_disabled(mut self, disabled: bool) -> Self {
299        self.no_gfm = Some(disabled);
300        self
301    }
302
303    /// Configure whether image data URLs are preserved.
304    pub fn with_image_data_urls(mut self, keep: bool) -> Self {
305        self.keep_img_data_url = Some(keep);
306        self
307    }
308
309    /// Configure image summaries.
310    pub fn with_image_summaries(mut self, enabled: bool) -> Self {
311        self.with_images_summary = Some(enabled);
312        self
313    }
314
315    /// Configure link summaries.
316    pub fn with_link_summaries(mut self, enabled: bool) -> Self {
317        self.with_links_summary = Some(enabled);
318        self
319    }
320
321    /// Borrow the source URL.
322    pub fn url(&self) -> &str {
323        &self.url
324    }
325
326    /// Return the configured timeout in seconds.
327    pub const fn timeout_seconds(&self) -> Option<u64> {
328        self.timeout
329    }
330
331    /// Borrow the requested return format.
332    pub fn return_format(&self) -> Option<&str> {
333        self.return_format.as_deref()
334    }
335
336    /// Return whether cache bypass was configured.
337    pub const fn cache_disabled(&self) -> Option<bool> {
338        self.no_cache
339    }
340
341    /// Return whether parsed images are retained.
342    pub const fn retained_images(&self) -> Option<bool> {
343        self.retain_images
344    }
345
346    /// Return whether GitHub-Flavored Markdown is disabled.
347    pub const fn gfm_disabled(&self) -> Option<bool> {
348        self.no_gfm
349    }
350
351    /// Return whether image data URLs are preserved.
352    pub const fn keeps_image_data_urls(&self) -> Option<bool> {
353        self.keep_img_data_url
354    }
355
356    /// Return whether image summaries are requested.
357    pub const fn image_summaries_enabled(&self) -> Option<bool> {
358        self.with_images_summary
359    }
360
361    /// Return whether link summaries are requested.
362    pub const fn link_summaries_enabled(&self) -> Option<bool> {
363        self.with_links_summary
364    }
365
366    /// Validate the required URL string.
367    pub fn validate(&self) -> ZaiResult<()> {
368        if self.url.trim().is_empty() {
369            return Err(validation_error("reader url must not be blank"));
370        }
371        Ok(())
372    }
373
374    /// Validate, send, and decode a reader response.
375    pub async fn send_via(&self, client: &ZaiClient) -> ZaiResult<ReaderResponse> {
376        self.validate()?;
377        let route = crate::client::routes::TOOLS_READER;
378        let url = client.endpoints().resolve_route(route, &[])?;
379        let response = client
380            .send_json::<_, ReaderResponse>(route.method(), url, self)
381            .await?;
382        response.validate()?;
383        Ok(response)
384    }
385}
386
387#[cfg(test)]
388mod tests {
389    use super::*;
390
391    #[test]
392    fn layout_request_serializes_fixed_model_and_exact_wire_names() {
393        let request = LayoutParsingRequest::new("https://example.test/document.pdf")
394            .with_crop_images(true)
395            .with_layout_visualization(false)
396            .with_page_range(2, 4)
397            .with_request_id("request-1")
398            .with_user_id("user-1");
399        assert_eq!(
400            serde_json::to_value(request).unwrap(),
401            serde_json::json!({
402                "model": "glm-ocr",
403                "file": "https://example.test/document.pdf",
404                "return_crop_images": true,
405                "need_layout_visualization": false,
406                "start_page_id": 2,
407                "end_page_id": 4,
408                "request_id": "request-1",
409                "user_id": "user-1"
410            })
411        );
412        let debug = format!(
413            "{:?}",
414            LayoutParsingRequest::new("private-file-data")
415                .with_request_id("private-request")
416                .with_user_id("private-user")
417        );
418        assert!(!debug.contains("private-file-data"));
419        assert!(!debug.contains("private-request"));
420        assert!(!debug.contains("private-user"));
421    }
422
423    #[test]
424    fn request_validation_enforces_documented_constraints() {
425        assert!(LayoutParsingRequest::new(" ").validate().is_err());
426        assert!(
427            LayoutParsingRequest::new("data")
428                .with_page_range(2, 1)
429                .validate()
430                .is_err()
431        );
432        assert!(
433            LayoutParsingRequest::new("data")
434                .with_request_id("short")
435                .validate()
436                .is_err()
437        );
438        assert!(ReaderRequest::new(" ").validate().is_err());
439    }
440
441    #[test]
442    fn reader_request_serializes_only_selected_closed_fields() {
443        let request = ReaderRequest::new("https://example.test/page")
444            .with_timeout_seconds(30)
445            .with_cache_disabled(true)
446            .with_return_format("markdown")
447            .with_retained_images(true)
448            .with_image_summaries(true);
449        assert_eq!(
450            serde_json::to_value(request).unwrap(),
451            serde_json::json!({
452                "url": "https://example.test/page",
453                "timeout": 30,
454                "no_cache": true,
455                "return_format": "markdown",
456                "retain_images": true,
457                "with_images_summary": true
458            })
459        );
460        assert!(
461            !format!("{:?}", ReaderRequest::new("https://secret.test/token"))
462                .contains("secret.test")
463        );
464    }
465}