zai_rs/knowledge/
search.rs

1//! Semantic knowledge search (`POST /llm-application/open/knowledge/retrieve`).
2
3use serde::{Deserialize, Deserializer, Serialize, de::Error as _};
4use validator::Validate;
5
6use crate::{ZaiResult, client::ZaiClient};
7
8/// Retrieval strategy used before optional reranking.
9#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
10#[serde(rename_all = "lowercase")]
11pub enum KnowledgeRecallMethod {
12    /// Vector similarity search.
13    Embedding,
14    /// Keyword search.
15    Keyword,
16    /// Combine vector and keyword search.
17    Mixed,
18}
19
20/// Reranking model applied to recalled chunks.
21#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
22pub enum KnowledgeRerankModel {
23    /// Standard reranker.
24    #[serde(rename = "rerank")]
25    Rerank,
26    /// Higher-quality reranker.
27    #[serde(rename = "rerank-pro")]
28    RerankPro,
29}
30
31/// JSON body for a semantic knowledge search.
32#[derive(Clone, Serialize, Deserialize, Validate)]
33pub struct KnowledgeSearchBody {
34    /// Optional caller-generated request identifier.
35    #[serde(skip_serializing_if = "Option::is_none")]
36    pub request_id: Option<String>,
37    /// Search query text (up to 1000 characters).
38    #[validate(length(min = 1, max = 1000))]
39    pub query: String,
40    /// Knowledge bases to search.
41    #[validate(length(min = 1))]
42    pub knowledge_ids: Vec<String>,
43    /// Restrict retrieval to these document identifiers.
44    #[serde(skip_serializing_if = "Option::is_none")]
45    pub document_ids: Option<Vec<String>>,
46    /// Final result count (`1..=20`, server default `8`).
47    #[serde(skip_serializing_if = "Option::is_none")]
48    #[validate(range(min = 1, max = 20))]
49    pub top_k: Option<u32>,
50    /// Initial recall count (`1..=100`, server default `10`).
51    #[serde(skip_serializing_if = "Option::is_none")]
52    #[validate(range(min = 1, max = 100))]
53    pub top_n: Option<u32>,
54    /// Recall strategy (server default: mixed).
55    #[serde(skip_serializing_if = "Option::is_none")]
56    pub recall_method: Option<KnowledgeRecallMethod>,
57    /// Vector-search weight for mixed retrieval (`1..=99`).
58    #[serde(skip_serializing_if = "Option::is_none")]
59    #[validate(range(min = 1, max = 99))]
60    pub recall_ratio: Option<u8>,
61    /// Whether reranking is enabled (`0` or `1`).
62    #[serde(skip_serializing_if = "Option::is_none")]
63    #[validate(range(max = 1))]
64    pub rerank_status: Option<u8>,
65    /// Reranking model.
66    #[serde(skip_serializing_if = "Option::is_none")]
67    pub rerank_model: Option<KnowledgeRerankModel>,
68    /// Similarity threshold in the open interval `(0, 1)`.
69    ///
70    /// The Rust field keeps its historical name while the wire field follows
71    /// the upstream `fractional_threshold` schema.
72    #[serde(
73        rename = "fractional_threshold",
74        skip_serializing_if = "Option::is_none"
75    )]
76    pub score_threshold: Option<f64>,
77}
78
79impl std::fmt::Debug for KnowledgeSearchBody {
80    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
81        formatter
82            .debug_struct("KnowledgeSearchBody")
83            .field(
84                "request_id",
85                &self.request_id.as_ref().map(|_| "[REDACTED]"),
86            )
87            .field("query", &"[REDACTED]")
88            .field("knowledge_id_count", &self.knowledge_ids.len())
89            .field(
90                "document_id_count",
91                &self.document_ids.as_ref().map(Vec::len),
92            )
93            .field("top_k", &self.top_k)
94            .field("top_n", &self.top_n)
95            .field("recall_method", &self.recall_method)
96            .field("recall_ratio", &self.recall_ratio)
97            .field("rerank_status", &self.rerank_status)
98            .field("rerank_model", &self.rerank_model)
99            .field("score_threshold", &self.score_threshold)
100            .finish()
101    }
102}
103
104/// Request for semantic retrieval across one or more knowledge bases.
105pub struct KnowledgeSearchRequest {
106    /// JSON body sent to the retrieval endpoint.
107    pub body: KnowledgeSearchBody,
108}
109
110impl KnowledgeSearchRequest {
111    /// Create a search request targeting one knowledge base.
112    pub fn new(knowledge_id: impl Into<String>, query: impl Into<String>) -> Self {
113        Self {
114            body: KnowledgeSearchBody {
115                request_id: None,
116                query: query.into(),
117                knowledge_ids: vec![knowledge_id.into()],
118                document_ids: None,
119                top_k: None,
120                top_n: None,
121                recall_method: None,
122                recall_ratio: None,
123                rerank_status: None,
124                rerank_model: None,
125                score_threshold: None,
126            },
127        }
128    }
129
130    /// Replace the knowledge-base selection.
131    pub fn with_knowledge_ids(mut self, knowledge_ids: Vec<String>) -> Self {
132        self.body.knowledge_ids = knowledge_ids;
133        self
134    }
135
136    /// Restrict retrieval to selected documents.
137    pub fn with_document_ids(mut self, document_ids: Vec<String>) -> Self {
138        self.body.document_ids = Some(document_ids);
139        self
140    }
141
142    /// Set a caller-generated request identifier.
143    pub fn with_request_id(mut self, request_id: impl Into<String>) -> Self {
144        self.body.request_id = Some(request_id.into());
145        self
146    }
147
148    /// Set the requested maximum number of final matches.
149    pub fn with_top_k(mut self, top_k: u32) -> Self {
150        self.body.top_k = Some(top_k);
151        self
152    }
153
154    /// Set the number of candidates recalled before reranking.
155    pub fn with_top_n(mut self, top_n: u32) -> Self {
156        self.body.top_n = Some(top_n);
157        self
158    }
159
160    /// Select the retrieval strategy.
161    pub fn with_recall_method(mut self, method: KnowledgeRecallMethod) -> Self {
162        self.body.recall_method = Some(method);
163        self
164    }
165
166    /// Set the vector-search weight for mixed retrieval.
167    pub fn with_recall_ratio(mut self, ratio: u8) -> Self {
168        self.body.recall_ratio = Some(ratio);
169        self
170    }
171
172    /// Configure reranking and its model.
173    pub fn with_reranking(mut self, enabled: bool, model: KnowledgeRerankModel) -> Self {
174        self.body.rerank_status = Some(u8::from(enabled));
175        self.body.rerank_model = Some(model);
176        self
177    }
178
179    /// Set the minimum relevance score sent as `fractional_threshold`.
180    pub fn with_score_threshold(mut self, threshold: f64) -> Self {
181        self.body.score_threshold = Some(threshold);
182        self
183    }
184
185    /// Validate and send the search request.
186    pub async fn send_via(&self, client: &ZaiClient) -> ZaiResult<KnowledgeSearchResponse> {
187        self.body.validate()?;
188        if self.body.query.trim().is_empty() {
189            return Err(crate::client::validation::invalid("query cannot be blank"));
190        }
191        if self
192            .body
193            .knowledge_ids
194            .iter()
195            .any(|knowledge_id| knowledge_id.trim().is_empty())
196        {
197            return Err(crate::client::validation::invalid(
198                "knowledge_ids cannot contain blank values",
199            ));
200        }
201        if self
202            .body
203            .document_ids
204            .as_ref()
205            .is_some_and(|ids| ids.is_empty() || ids.iter().any(|id| id.trim().is_empty()))
206        {
207            return Err(crate::client::validation::invalid(
208                "document_ids must contain at least one non-blank value",
209            ));
210        }
211        if self
212            .body
213            .score_threshold
214            .is_some_and(|threshold| !threshold.is_finite() || threshold <= 0.0 || threshold >= 1.0)
215        {
216            return Err(crate::client::validation::invalid(
217                "score_threshold must be finite and in the range 0.0 < value < 1.0",
218            ));
219        }
220
221        let route = crate::client::routes::KNOWLEDGE_RETRIEVE;
222        let url = client.endpoints().resolve_route(route, &[])?;
223        client
224            .send_json::<_, KnowledgeSearchResponse>(route.method(), url, &self.body)
225            .await
226    }
227}
228
229/// Metadata attached to one retrieved knowledge chunk.
230#[derive(Debug, Clone, Serialize, Deserialize)]
231pub struct KnowledgeSearchMetadata {
232    /// Chunk identifier.
233    #[serde(rename = "_id", skip_serializing_if = "Option::is_none")]
234    pub id: Option<String>,
235    /// Knowledge-base identifier.
236    #[serde(skip_serializing_if = "Option::is_none")]
237    pub knowledge_id: Option<String>,
238    /// Source document identifier.
239    #[serde(skip_serializing_if = "Option::is_none")]
240    pub doc_id: Option<String>,
241    /// Source document name.
242    #[serde(skip_serializing_if = "Option::is_none")]
243    pub doc_name: Option<String>,
244    /// Source document URL.
245    #[serde(skip_serializing_if = "Option::is_none")]
246    pub doc_url: Option<String>,
247    /// Context-enriched text, when contextual retrieval was enabled.
248    #[serde(skip_serializing_if = "Option::is_none")]
249    pub contextual_text: Option<String>,
250}
251
252/// One semantic-retrieval result.
253#[derive(Debug, Clone, Serialize, Deserialize)]
254pub struct KnowledgeSearchResult {
255    /// Retrieved chunk text.
256    #[serde(skip_serializing_if = "Option::is_none")]
257    pub text: Option<String>,
258    /// Similarity score reported by the service.
259    #[serde(skip_serializing_if = "Option::is_none")]
260    pub score: Option<f64>,
261    /// Structured source metadata.
262    #[serde(skip_serializing_if = "Option::is_none")]
263    pub metadata: Option<KnowledgeSearchMetadata>,
264}
265
266/// Semantic search success response.
267#[derive(Debug, Clone, Serialize)]
268pub struct KnowledgeSearchResponse {
269    /// Retrieved chunk objects, when returned.
270    #[serde(skip_serializing_if = "Option::is_none")]
271    pub data: Option<Vec<KnowledgeSearchResult>>,
272    /// Business status code, when returned. The shared transport rejects an
273    /// explicitly non-success code before this type is returned.
274    #[serde(skip_serializing_if = "Option::is_none")]
275    pub code: Option<i64>,
276    /// Human-readable status message.
277    #[serde(default, skip_serializing_if = "Option::is_none")]
278    pub message: Option<String>,
279    /// Server timestamp.
280    #[serde(default, skip_serializing_if = "Option::is_none")]
281    pub timestamp: Option<u64>,
282}
283
284#[derive(Deserialize)]
285struct KnowledgeSearchResponseWire {
286    data: Option<Vec<KnowledgeSearchResult>>,
287    code: Option<i64>,
288    message: Option<String>,
289    timestamp: Option<u64>,
290}
291
292impl<'de> Deserialize<'de> for KnowledgeSearchResponse {
293    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
294    where
295        D: Deserializer<'de>,
296    {
297        let wire = KnowledgeSearchResponseWire::deserialize(deserializer)?;
298        if wire.data.is_none()
299            && wire.code.is_none()
300            && wire.message.is_none()
301            && wire.timestamp.is_none()
302        {
303            return Err(D::Error::custom(
304                "knowledge-search response contained no documented non-null fields",
305            ));
306        }
307        Ok(Self {
308            data: wire.data,
309            code: wire.code,
310            message: wire.message,
311            timestamp: wire.timestamp,
312        })
313    }
314}
315
316impl KnowledgeSearchResponse {
317    /// Borrow the returned results when the service included them.
318    pub fn results(&self) -> Option<&[KnowledgeSearchResult]> {
319        self.data.as_deref()
320    }
321}
322
323#[cfg(test)]
324mod tests {
325    use super::*;
326
327    #[test]
328    fn convenience_builder_matches_official_wire_names() {
329        let request = KnowledgeSearchRequest::new("kb-1", "question")
330            .with_top_k(5)
331            .with_score_threshold(0.25);
332        let value = serde_json::to_value(&request.body).unwrap();
333
334        assert_eq!(value["knowledge_ids"], serde_json::json!(["kb-1"]));
335        assert_eq!(value["fractional_threshold"], 0.25);
336        assert!(value.get("knowledge_id").is_none());
337        assert!(value.get("score_threshold").is_none());
338    }
339
340    #[test]
341    fn request_body_debug_redacts_query_and_identifiers() {
342        let body = KnowledgeSearchRequest::new("private-knowledge", "private query")
343            .with_document_ids(vec!["private-document".to_owned()])
344            .with_request_id("private-request")
345            .body;
346        let debug = format!("{body:?}");
347        for secret in [
348            "private-knowledge",
349            "private query",
350            "private-document",
351            "private-request",
352        ] {
353            assert!(!debug.contains(secret));
354        }
355        assert!(debug.contains("knowledge_id_count: 1"));
356        assert!(debug.contains("document_id_count: Some(1)"));
357    }
358
359    #[test]
360    fn response_uses_typed_metadata_and_follows_optional_frozen_fields() {
361        let response: KnowledgeSearchResponse = serde_json::from_value(serde_json::json!({
362            "code": 200,
363            "data": [{
364                "text": "chunk",
365                "score": 0.9,
366                "metadata": {
367                    "_id": "slice-1",
368                    "knowledge_id": "kb-1",
369                    "doc_id": "doc-1",
370                    "doc_name": "guide.md",
371                    "doc_url": "https://example.com/guide.md",
372                    "contextual_text": "context"
373                }
374            }]
375        }))
376        .unwrap();
377        let result = &response.results().unwrap()[0];
378        assert_eq!(result.text.as_deref(), Some("chunk"));
379        assert_eq!(result.score, Some(0.9));
380        assert_eq!(
381            result
382                .metadata
383                .as_ref()
384                .and_then(|metadata| metadata.id.as_deref()),
385            Some("slice-1")
386        );
387
388        assert!(
389            serde_json::from_value::<KnowledgeSearchResponse>(serde_json::json!({"data": []}))
390                .is_ok()
391        );
392        assert!(
393            serde_json::from_value::<KnowledgeSearchResponse>(serde_json::json!({"code": 200}))
394                .is_ok()
395        );
396        assert!(serde_json::from_value::<KnowledgeSearchResponse>(serde_json::json!({})).is_err());
397    }
398}
zai_rs/knowledge/search.rs

zai_rs/knowledge/
search.rs
