zai_rs/model/text_rerank/
data.rs1use 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
15pub 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 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 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 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}