zai_rs/model/text_rerank/
data.rs

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