Skip to main content

zai_rs/model/text_rerank/
data.rs

1use std::sync::Arc;
2
3use super::{
4    request::{RerankBody, RerankModel},
5    response::RerankResponse,
6};
7use crate::{
8    ZaiResult,
9    client::{
10        endpoints::{ApiBase, EndpointConfig, paths},
11        http::{HttpClient, HttpClientConfig, parse_typed_response},
12    },
13};
14
15/// Text Rerank request client (JSON POST)
16pub struct RerankRequest {
17    pub key: String,
18    url: String,
19    endpoint_config: EndpointConfig,
20    api_base: ApiBase,
21    http_config: Arc<HttpClientConfig>,
22    body: RerankBody,
23}
24
25impl RerankRequest {
26    pub fn new(key: String, query: impl Into<String>, documents: Vec<String>) -> Self {
27        let endpoint_config = EndpointConfig::default();
28        let api_base = ApiBase::PaasV4;
29        let url = endpoint_config.url(&api_base, paths::RERANK);
30        let body = RerankBody::new(RerankModel::Rerank, query, documents);
31        Self {
32            key,
33            url,
34            endpoint_config,
35            api_base,
36            http_config: Arc::new(HttpClientConfig::default()),
37            body,
38        }
39    }
40
41    fn rebuild_url(&mut self) {
42        self.url = self.endpoint_config.url(&self.api_base, paths::RERANK);
43    }
44
45    pub fn with_base_url(mut self, base: impl Into<String>) -> Self {
46        self.api_base = ApiBase::Custom(base.into());
47        self.rebuild_url();
48        self
49    }
50
51    pub fn with_endpoint_config(mut self, endpoint_config: EndpointConfig) -> Self {
52        self.endpoint_config = endpoint_config;
53        self.rebuild_url();
54        self
55    }
56
57    pub fn with_http_config(mut self, config: HttpClientConfig) -> Self {
58        self.http_config = Arc::new(config);
59        self
60    }
61
62    pub fn with_top_n(mut self, n: usize) -> Self {
63        self.body = self.body.with_top_n(n);
64        self
65    }
66    pub fn with_return_documents(mut self, v: bool) -> Self {
67        self.body = self.body.with_return_documents(v);
68        self
69    }
70    pub fn with_return_raw_scores(mut self, v: bool) -> Self {
71        self.body = self.body.with_return_raw_scores(v);
72        self
73    }
74    pub fn with_request_id(mut self, v: impl Into<String>) -> Self {
75        self.body = self.body.with_request_id(v);
76        self
77    }
78    pub fn with_user_id(mut self, v: impl Into<String>) -> Self {
79        self.body = self.body.with_user_id(v);
80        self
81    }
82
83    /// Optional: validate constraints before sending
84    pub fn validate(&self) -> ZaiResult<()> {
85        self.body
86            .validate_constraints()
87            .map_err(|e| crate::client::error::ZaiError::ApiError {
88                code: 1200,
89                message: format!("Validation error: {:?}", e),
90            })
91    }
92
93    /// Send the request and parse typed response.
94    /// Automatically runs `validate()` before sending.
95    pub async fn send(&self) -> ZaiResult<RerankResponse> {
96        self.validate()?;
97        let resp: reqwest::Response = self.post().await?;
98        let parsed = parse_typed_response::<RerankResponse>(resp).await?;
99        Ok(parsed)
100    }
101
102    #[deprecated(note = "Use send() instead")]
103    /// Deprecated: use `send()`.
104    pub async fn execute(&self) -> ZaiResult<RerankResponse> {
105        self.send().await
106    }
107}
108
109impl HttpClient for RerankRequest {
110    type Body = RerankBody;
111    type ApiUrl = String;
112    type ApiKey = String;
113
114    fn api_url(&self) -> &Self::ApiUrl {
115        &self.url
116    }
117    fn api_key(&self) -> &Self::ApiKey {
118        &self.key
119    }
120    fn body(&self) -> &Self::Body {
121        &self.body
122    }
123    fn http_config(&self) -> Arc<HttpClientConfig> {
124        Arc::clone(&self.http_config)
125    }
126}