Skip to main content

zai_rs/services/tools/
response.rs

1//! Typed responses for document layout parsing and web-page reading.
2
3use std::collections::BTreeMap;
4
5use serde::{Deserialize, Serialize};
6
7use crate::{
8    ZaiResult,
9    client::error::{ZaiError, codes},
10};
11
12fn response_error(message: impl Into<String>) -> ZaiError {
13    ZaiError::ApiError {
14        code: codes::SDK_VALIDATION,
15        message: message.into(),
16    }
17}
18
19/// Element category in a parsed document layout.
20#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
21pub enum LayoutLabel {
22    /// Image element.
23    #[serde(rename = "image")]
24    Image,
25    /// Text element.
26    #[serde(rename = "text")]
27    Text,
28    /// Display-formula element.
29    #[serde(rename = "formula")]
30    Formula,
31    /// Table element.
32    #[serde(rename = "table")]
33    Table,
34}
35
36/// One element in a parsed document layout.
37#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
38pub struct LayoutDetail {
39    /// Element index.
40    pub index: i64,
41    /// Element category.
42    pub label: LayoutLabel,
43    /// Normalized `[x1, y1, x2, y2]` coordinates.
44    #[serde(skip_serializing_if = "Option::is_none")]
45    pub bbox_2d: Option<[f64; 4]>,
46    /// Text, image URL, formula, or table HTML.
47    #[serde(skip_serializing_if = "Option::is_none")]
48    pub content: Option<String>,
49    /// Source page height.
50    #[serde(skip_serializing_if = "Option::is_none")]
51    pub height: Option<i64>,
52    /// Source page width.
53    #[serde(skip_serializing_if = "Option::is_none")]
54    pub width: Option<i64>,
55}
56
57/// Dimensions for one parsed page.
58#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
59pub struct PageInfo {
60    /// Page width.
61    pub width: i64,
62    /// Page height.
63    pub height: i64,
64}
65
66/// Basic information about a parsed document.
67#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
68pub struct DataInfo {
69    /// Number of pages in the document.
70    pub num_pages: i64,
71    /// Per-page dimensions, when returned.
72    #[serde(skip_serializing_if = "Option::is_none")]
73    pub pages: Option<Vec<PageInfo>>,
74}
75
76/// Cached-token details returned for layout parsing.
77#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
78pub struct LayoutPromptTokensDetails {
79    /// Number of cached prompt tokens.
80    #[serde(skip_serializing_if = "Option::is_none")]
81    pub cached_tokens: Option<f64>,
82}
83
84/// Token usage returned by layout parsing.
85#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
86pub struct LayoutUsage {
87    /// Prompt token count.
88    #[serde(skip_serializing_if = "Option::is_none")]
89    pub prompt_tokens: Option<f64>,
90    /// Completion token count.
91    #[serde(skip_serializing_if = "Option::is_none")]
92    pub completion_tokens: Option<f64>,
93    /// Cached-token details.
94    #[serde(skip_serializing_if = "Option::is_none")]
95    pub prompt_tokens_details: Option<LayoutPromptTokensDetails>,
96    /// Total token count.
97    #[serde(skip_serializing_if = "Option::is_none")]
98    pub total_tokens: Option<i64>,
99}
100
101/// Response from the layout-parsing endpoint.
102#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
103pub struct LayoutParsingResponse {
104    /// Task identifier.
105    pub id: String,
106    /// Creation time as a Unix timestamp.
107    pub created: i64,
108    /// Model name returned by the service.
109    pub model: String,
110    /// Markdown parsing result.
111    #[serde(skip_serializing_if = "Option::is_none")]
112    pub md_results: Option<String>,
113    /// Layout elements grouped by page.
114    #[serde(skip_serializing_if = "Option::is_none")]
115    pub layout_details: Option<Vec<Vec<LayoutDetail>>>,
116    /// Layout-visualization image URLs.
117    #[serde(skip_serializing_if = "Option::is_none")]
118    pub layout_visualization: Option<Vec<String>>,
119    /// Document information.
120    #[serde(skip_serializing_if = "Option::is_none")]
121    pub data_info: Option<DataInfo>,
122    /// Token usage.
123    #[serde(skip_serializing_if = "Option::is_none")]
124    pub usage: Option<LayoutUsage>,
125    /// Request identifier.
126    #[serde(skip_serializing_if = "Option::is_none")]
127    pub request_id: Option<String>,
128}
129
130impl LayoutParsingResponse {
131    /// Validate normalized bounding-box coordinates from the response.
132    pub fn validate(&self) -> ZaiResult<()> {
133        let invalid_bbox = self
134            .layout_details
135            .iter()
136            .flatten()
137            .flatten()
138            .filter_map(|detail| detail.bbox_2d)
139            .flatten()
140            .any(|coordinate| !coordinate.is_finite() || !(0.0..=1.0).contains(&coordinate));
141        if invalid_bbox {
142            return Err(response_error(
143                "layout response bbox_2d coordinates must be finite values in 0..=1",
144            ));
145        }
146        Ok(())
147    }
148}
149
150/// One external stylesheet referenced by a reader result.
151#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
152pub struct ReaderStylesheet {
153    /// Stylesheet MIME type.
154    #[serde(rename = "type", skip_serializing_if = "Option::is_none")]
155    pub type_: Option<String>,
156}
157
158/// External resources referenced by a reader result.
159#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
160pub struct ReaderExternalResources {
161    /// Stylesheets keyed by the upstream resource identifier. This is the sole
162    /// `additionalProperties` map in the frozen Reader response.
163    #[serde(skip_serializing_if = "Option::is_none")]
164    pub stylesheet: Option<BTreeMap<String, ReaderStylesheet>>,
165}
166
167/// Page metadata returned by the reader endpoint.
168#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
169pub struct ReaderMetadata {
170    /// Page keywords.
171    #[serde(skip_serializing_if = "Option::is_none")]
172    pub keywords: Option<String>,
173    /// Viewport metadata.
174    #[serde(skip_serializing_if = "Option::is_none")]
175    pub viewport: Option<String>,
176    /// Page description metadata.
177    #[serde(skip_serializing_if = "Option::is_none")]
178    pub description: Option<String>,
179    /// Format-detection metadata.
180    #[serde(rename = "format-detection", skip_serializing_if = "Option::is_none")]
181    pub format_detection: Option<String>,
182}
183
184/// Parsed web-page content returned by the reader endpoint.
185#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
186pub struct ReaderResult {
187    /// Parsed page content.
188    #[serde(skip_serializing_if = "Option::is_none")]
189    pub content: Option<String>,
190    /// Page description.
191    #[serde(skip_serializing_if = "Option::is_none")]
192    pub description: Option<String>,
193    /// Page title.
194    #[serde(skip_serializing_if = "Option::is_none")]
195    pub title: Option<String>,
196    /// Original page URL.
197    #[serde(skip_serializing_if = "Option::is_none")]
198    pub url: Option<String>,
199    /// External page resources.
200    #[serde(skip_serializing_if = "Option::is_none")]
201    pub external: Option<ReaderExternalResources>,
202    /// Page metadata.
203    #[serde(skip_serializing_if = "Option::is_none")]
204    pub metadata: Option<ReaderMetadata>,
205}
206
207/// Response from the reader endpoint.
208#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
209pub struct ReaderResponse {
210    /// Task identifier.
211    #[serde(skip_serializing_if = "Option::is_none")]
212    pub id: Option<String>,
213    /// Creation time as a Unix timestamp.
214    #[serde(skip_serializing_if = "Option::is_none")]
215    pub created: Option<i64>,
216    /// Request identifier.
217    #[serde(skip_serializing_if = "Option::is_none")]
218    pub request_id: Option<String>,
219    /// Model identifier.
220    #[serde(skip_serializing_if = "Option::is_none")]
221    pub model: Option<String>,
222    /// Parsed page result.
223    #[serde(skip_serializing_if = "Option::is_none")]
224    pub reader_result: Option<ReaderResult>,
225}
226
227impl ReaderResponse {
228    /// Enforce the frozen operation's non-empty success-response invariant.
229    pub fn validate(&self) -> ZaiResult<()> {
230        if self.id.is_some()
231            || self.created.is_some()
232            || self.request_id.is_some()
233            || self.model.is_some()
234            || self.reader_result.is_some()
235        {
236            return Ok(());
237        }
238        Err(response_error(
239            "reader response contained no documented fields",
240        ))
241    }
242}
243
244#[cfg(test)]
245mod tests {
246    use super::*;
247
248    #[test]
249    fn layout_required_fields_and_fixed_bbox_length_are_enforced() {
250        assert!(
251            serde_json::from_value::<LayoutParsingResponse>(serde_json::json!({
252                "created": 1,
253                "model": "GLM-OCR"
254            }))
255            .is_err()
256        );
257        assert!(
258            serde_json::from_value::<LayoutParsingResponse>(serde_json::json!({
259                "id": "task-1",
260                "created": 1,
261                "model": "GLM-OCR",
262                "layout_details": [[{"index": 0, "label": "text", "bbox_2d": [0.0, 1.0]}]]
263            }))
264            .is_err()
265        );
266        assert!(serde_json::from_value::<LayoutDetail>(serde_json::json!({"index": 1})).is_err());
267        assert!(serde_json::from_value::<DataInfo>(serde_json::json!({"pages": []})).is_err());
268        assert!(serde_json::from_value::<PageInfo>(serde_json::json!({"width": 1})).is_err());
269    }
270
271    #[test]
272    fn layout_validation_rejects_out_of_range_coordinates() {
273        let response: LayoutParsingResponse = serde_json::from_value(serde_json::json!({
274            "id": "task-1",
275            "created": 1,
276            "model": "GLM-OCR",
277            "layout_details": [[{
278                "index": 0,
279                "label": "text",
280                "bbox_2d": [0.0, 0.0, 1.2, 1.0]
281            }]]
282        }))
283        .unwrap();
284        assert!(response.validate().is_err());
285    }
286
287    #[test]
288    fn reader_response_models_typed_additional_properties() {
289        let response: ReaderResponse = serde_json::from_value(serde_json::json!({
290            "reader_result": {
291                "external": {"stylesheet": {"main": {"type": "text/css"}}},
292                "metadata": {"format-detection": "telephone=no"}
293            }
294        }))
295        .unwrap();
296        assert_eq!(
297            response
298                .reader_result
299                .as_ref()
300                .unwrap()
301                .external
302                .as_ref()
303                .unwrap()
304                .stylesheet
305                .as_ref()
306                .unwrap()["main"]
307                .type_
308                .as_deref(),
309            Some("text/css")
310        );
311        assert!(response.validate().is_ok());
312        assert!(
313            serde_json::from_value::<ReaderResponse>(serde_json::json!({}))
314                .unwrap()
315                .validate()
316                .is_err()
317        );
318    }
319}