zai_rs/model/text_rerank/
data.rs1use super::request::{RerankBody, RerankModel};
2use super::response::RerankResponse;
3use crate::client::http::HttpClient;
4
5pub 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 pub fn validate(&self) -> anyhow::Result<()> {
40 self.body.validate_constraints().map_err(Into::into)
41 }
42
43 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 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}