Skip to main content

zai_rs/knowledge/
search.rs

1//! Knowledge search (POST /llm-application/open/knowledge/retrieve) — plan P06.
2
3use crate::ZaiResult;
4use crate::client::ZaiClient;
5use serde::Serialize;
6
7/// POST /knowledge/retrieve search request (P06 knowledge.retrieve).
8#[derive(Debug, Clone, Serialize)]
9pub struct KnowledgeSearchBody {
10    /// Knowledge base id.
11    pub knowledge_id: String,
12    /// Search query text.
13    pub query: String,
14    /// Top-k results (optional).
15    #[serde(skip_serializing_if = "Option::is_none")]
16    pub top_k: Option<u32>,
17    /// Score threshold (optional).
18    #[serde(skip_serializing_if = "Option::is_none")]
19    pub score_threshold: Option<f64>,
20}
21
22pub struct KnowledgeSearchRequest {
23    pub body: KnowledgeSearchBody,
24}
25
26impl KnowledgeSearchRequest {
27    pub fn new(knowledge_id: impl Into<String>, query: impl Into<String>) -> Self {
28        Self {
29            body: KnowledgeSearchBody {
30                knowledge_id: knowledge_id.into(),
31                query: query.into(),
32                top_k: None,
33                score_threshold: None,
34            },
35        }
36    }
37
38    pub fn with_top_k(mut self, top_k: u32) -> Self {
39        self.body.top_k = Some(top_k);
40        self
41    }
42
43    pub fn with_score_threshold(mut self, threshold: f64) -> Self {
44        self.body.score_threshold = Some(threshold);
45        self
46    }
47
48    pub async fn send_via(&self, client: &ZaiClient) -> ZaiResult<KnowledgeSearchResponse> {
49        let route = crate::client::routes::KNOWLEDGE_RETRIEVE;
50        let url = client.endpoints().resolve_route(route, &[])?;
51        client
52            .send_json::<_, KnowledgeSearchResponse>(route.method(), url, &self.body)
53            .await
54    }
55}
56
57#[derive(Debug, Clone, serde::Deserialize)]
58pub struct KnowledgeSearchResponse {
59    #[serde(default)]
60    pub data: Vec<serde_json::Value>,
61}