zai_rs/model/text_rerank/
request.rs

1use serde::{Deserialize, Serialize};
2
3/// Rerank model enum
4#[derive(Debug, Clone, Default, Serialize, Deserialize)]
5#[serde(rename_all = "lowercase")]
6pub enum RerankModel {
7    /// The default rerank model.
8    #[default]
9    Rerank,
10}
11
12/// Request body for rerank API
13#[derive(Clone, Serialize)]
14pub struct RerankBody {
15    /// Model identifier; currently always `rerank`.
16    pub model: RerankModel,
17
18    /// Query text (at most 4,096 characters).
19    pub query: String,
20
21    /// Candidate documents (1–128 items, at most 4,096 characters each).
22    pub documents: Vec<String>,
23
24    /// Number of highest-ranked documents to return; `0` or omission returns
25    /// all documents.
26    #[serde(skip_serializing_if = "Option::is_none")]
27    pub top_n: Option<usize>,
28
29    /// Whether to include the original document text; defaults to `false`.
30    #[serde(skip_serializing_if = "Option::is_none")]
31    pub return_documents: Option<bool>,
32
33    /// Whether to include raw relevance scores; defaults to `false`.
34    #[serde(skip_serializing_if = "Option::is_none")]
35    pub return_raw_scores: Option<bool>,
36
37    /// Client-provided request identifier.
38    #[serde(skip_serializing_if = "Option::is_none")]
39    pub request_id: Option<String>,
40
41    /// End-user identifier used for abuse monitoring.
42    #[serde(skip_serializing_if = "Option::is_none")]
43    pub user_id: Option<String>,
44}
45
46impl std::fmt::Debug for RerankBody {
47    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
48        formatter
49            .debug_struct("RerankBody")
50            .field("model", &self.model)
51            .field("query", &"[REDACTED]")
52            .field("documents_len", &self.documents.len())
53            .field("top_n", &self.top_n)
54            .field("return_documents", &self.return_documents)
55            .field("return_raw_scores", &self.return_raw_scores)
56            .field(
57                "request_id",
58                &self.request_id.as_ref().map(|_| "[REDACTED]"),
59            )
60            .field("user_id", &self.user_id.as_ref().map(|_| "[REDACTED]"))
61            .finish()
62    }
63}
64
65impl RerankBody {
66    /// Create a new rerank body from a model, query, and candidate documents.
67    pub fn new(model: RerankModel, query: impl Into<String>, documents: Vec<String>) -> Self {
68        Self {
69            model,
70            query: query.into(),
71            documents,
72            top_n: None,
73            return_documents: None,
74            return_raw_scores: None,
75            request_id: None,
76            user_id: None,
77        }
78    }
79
80    /// Set how many top-ranked documents to return.
81    pub fn with_top_n(mut self, n: usize) -> Self {
82        self.top_n = Some(n);
83        self
84    }
85    /// Whether to include the document text in the response.
86    pub fn with_return_documents(mut self, v: bool) -> Self {
87        self.return_documents = Some(v);
88        self
89    }
90    /// Whether to include raw relevance scores in the response.
91    pub fn with_return_raw_scores(mut self, v: bool) -> Self {
92        self.return_raw_scores = Some(v);
93        self
94    }
95    /// Set the client-side request id.
96    pub fn with_request_id(mut self, v: impl Into<String>) -> Self {
97        self.request_id = Some(v.into());
98        self
99    }
100    /// Set the end-user id.
101    pub fn with_user_id(mut self, v: impl Into<String>) -> Self {
102        self.user_id = Some(v.into());
103        self
104    }
105
106    /// Optional runtime validation for constraints expressed in the docs
107    pub fn validate_constraints(&self) -> crate::ZaiResult<()> {
108        if self.query.trim().is_empty() {
109            return Err(crate::client::validation::invalid(
110                "query must not be empty",
111            ));
112        }
113        if self.query.chars().count() > 4096 {
114            return Err(crate::client::validation::invalid(
115                "query length exceeds 4096 characters",
116            ));
117        }
118        if self.documents.is_empty() {
119            return Err(crate::client::validation::invalid(
120                "documents must not be empty",
121            ));
122        }
123        if self.documents.len() > 128 {
124            return Err(crate::client::validation::invalid(
125                "documents length exceeds 128",
126            ));
127        }
128        for (i, d) in self.documents.iter().enumerate() {
129            if d.trim().is_empty() {
130                return Err(crate::client::validation::invalid(format!(
131                    "document at index {i} must not be empty"
132                )));
133            }
134            if d.chars().count() > 4096 {
135                return Err(crate::client::validation::invalid(format!(
136                    "document at index {i} exceeds 4096 characters"
137                )));
138            }
139        }
140        if let Some(n) = self.top_n
141            && n > self.documents.len()
142        {
143            return Err(crate::client::validation::invalid(
144                "top_n cannot exceed documents length",
145            ));
146        }
147        if let Some(request_id) = self.request_id.as_deref()
148            && !(6..=64).contains(&request_id.chars().count())
149        {
150            return Err(crate::client::validation::invalid(
151                "request_id must contain between 6 and 64 characters",
152            ));
153        }
154        Ok(())
155    }
156}
157
158#[cfg(test)]
159mod tests {
160    use super::*;
161
162    #[test]
163    fn validation_rejects_blank_values_and_accepts_zero_top_n() {
164        assert!(
165            RerankBody::new(RerankModel::Rerank, " ", vec!["doc".into()])
166                .validate_constraints()
167                .is_err()
168        );
169        assert!(
170            RerankBody::new(RerankModel::Rerank, "query", vec![" ".into()])
171                .validate_constraints()
172                .is_err()
173        );
174        assert!(
175            RerankBody::new(RerankModel::Rerank, "query", vec!["doc".into()])
176                .with_top_n(0)
177                .validate_constraints()
178                .is_ok()
179        );
180    }
181
182    #[test]
183    fn debug_redacts_query_documents_and_identifiers() {
184        let body = RerankBody::new(
185            RerankModel::Rerank,
186            "private query",
187            vec!["private document".to_owned()],
188        )
189        .with_request_id("private-request")
190        .with_user_id("private-user");
191        let debug = format!("{body:?}");
192        for secret in [
193            "private query",
194            "private document",
195            "private-request",
196            "private-user",
197        ] {
198            assert!(!debug.contains(secret));
199        }
200    }
201}
zai_rs/model/text_rerank/request.rs

zai_rs/model/text_rerank/
request.rs
