Skip to main content

zai_rs/model/text_rerank/
data.rs

1use super::{
2    request::{RerankBody, RerankModel},
3    response::RerankResponse,
4};
5use crate::{ZaiResult, client::http::HttpClient};
6
7/// Text Rerank request client (JSON POST)
8pub struct RerankRequest {
9    pub key: String,
10    body: RerankBody,
11}
12
13impl RerankRequest {
14    pub fn new(key: String, query: impl Into<String>, documents: Vec<String>) -> Self {
15        let body = RerankBody::new(RerankModel::Rerank, query, documents);
16        Self { key, body }
17    }
18
19    pub fn with_top_n(mut self, n: usize) -> Self {
20        self.body = self.body.with_top_n(n);
21        self
22    }
23    pub fn with_return_documents(mut self, v: bool) -> Self {
24        self.body = self.body.with_return_documents(v);
25        self
26    }
27    pub fn with_return_raw_scores(mut self, v: bool) -> Self {
28        self.body = self.body.with_return_raw_scores(v);
29        self
30    }
31    pub fn with_request_id(mut self, v: impl Into<String>) -> Self {
32        self.body = self.body.with_request_id(v);
33        self
34    }
35    pub fn with_user_id(mut self, v: impl Into<String>) -> Self {
36        self.body = self.body.with_user_id(v);
37        self
38    }
39
40    /// Optional: validate constraints before sending
41    pub fn validate(&self) -> ZaiResult<()> {
42        self.body
43            .validate_constraints()
44            .map_err(|e| crate::client::error::ZaiError::ApiError {
45                code: 1200,
46                message: format!("Validation error: {:?}", e),
47            })
48    }
49
50    /// Send the request and parse typed response.
51    /// Automatically runs `validate()` before sending.
52    pub async fn send(&self) -> ZaiResult<RerankResponse> {
53        self.validate()?;
54        let resp: reqwest::Response = self.post().await?;
55        let parsed = resp.json::<RerankResponse>().await?;
56        Ok(parsed)
57    }
58
59    #[deprecated(note = "Use send() instead")]
60    /// Deprecated: use `send()`.
61    pub async fn execute(&self) -> ZaiResult<RerankResponse> {
62        self.send().await
63    }
64}
65
66impl HttpClient for RerankRequest {
67    type Body = RerankBody;
68    type ApiUrl = &'static str;
69    type ApiKey = String;
70
71    fn api_url(&self) -> &Self::ApiUrl {
72        &"https://open.bigmodel.cn/api/paas/v4/rerank"
73    }
74    fn api_key(&self) -> &Self::ApiKey {
75        &self.key
76    }
77    fn body(&self) -> &Self::Body {
78        &self.body
79    }
80}