kimun_notes/server_client/
mod.rs1use std::collections::HashMap;
16use std::time::Duration;
17
18pub mod dto;
19pub mod observer;
20pub mod reconcile;
21
22pub mod sync;
23
24use async_trait::async_trait;
25use dto::{
26 AnswerResult, DeleteRequest, EmbeddingsResponse, Health, HistoryTurn, IndexDocsRequest,
27 JobAccepted, JobStatus, QueryRequest, WireDoc,
28};
29
30pub use dto::{ChunkResult, WireSection};
31pub use observer::{DirtyOp, DirtySet, RagObserver};
32pub use reconcile::{ReconcilePlan, diff as reconcile_diff};
33
34#[derive(Debug, thiserror::Error)]
35pub enum RagError {
36 #[error("http error: {0}")]
37 Http(#[from] reqwest::Error),
38 #[error("server returned {status}: {body}")]
39 Status { status: u16, body: String },
40 #[error("{0}")]
43 Protocol(String),
44}
45
46impl RagError {
47 pub fn is_auth(&self) -> bool {
51 matches!(
52 self,
53 RagError::Status {
54 status: 401 | 403,
55 ..
56 }
57 )
58 }
59}
60
61#[async_trait]
65pub trait RagTransport: Send + Sync {
66 async fn push_docs(&self, docs: Vec<WireDoc>) -> Result<(), RagError>;
67 async fn delete_paths(&self, paths: Vec<String>) -> Result<(), RagError>;
68 async fn server_hashes(&self) -> Result<HashMap<String, String>, RagError>;
69}
70
71#[async_trait]
72impl RagTransport for RagClient {
73 async fn push_docs(&self, docs: Vec<WireDoc>) -> Result<(), RagError> {
74 RagClient::push_docs(self, docs).await.map(|_job_id| ())
77 }
78 async fn delete_paths(&self, paths: Vec<String>) -> Result<(), RagError> {
79 RagClient::delete_paths(self, paths).await
80 }
81 async fn server_hashes(&self) -> Result<HashMap<String, String>, RagError> {
82 RagClient::server_hashes(self).await
83 }
84}
85
86pub fn hash_string(hash: u64) -> String {
91 hash.to_string()
92}
93
94#[derive(Debug, Clone, Copy)]
97pub enum ContextSize {
98 Small,
99 Medium,
100 Large,
101}
102
103impl ContextSize {
104 fn as_str(self) -> &'static str {
105 match self {
106 ContextSize::Small => "small",
107 ContextSize::Medium => "medium",
108 ContextSize::Large => "large",
109 }
110 }
111}
112
113const CONNECT_TIMEOUT: Duration = Duration::from_secs(5);
117
118const REQUEST_TIMEOUT: Duration = Duration::from_secs(30);
121
122const PUSH_TIMEOUT: Duration = Duration::from_secs(120);
126
127fn shared_http() -> reqwest::Client {
131 static HTTP: std::sync::OnceLock<reqwest::Client> = std::sync::OnceLock::new();
132 HTTP.get_or_init(|| {
133 reqwest::Client::builder()
134 .connect_timeout(CONNECT_TIMEOUT)
135 .timeout(REQUEST_TIMEOUT)
136 .build()
137 .expect("build HTTP client")
140 })
141 .clone()
142}
143
144#[derive(Clone)]
146pub struct RagClient {
147 http: reqwest::Client,
148 base_url: String,
149 token: Option<String>,
150 vault_id: String,
151}
152
153impl RagClient {
154 pub fn new(
157 base_url: impl Into<String>,
158 token: Option<String>,
159 vault_id: impl Into<String>,
160 ) -> Self {
161 let base_url = base_url.into().trim_end_matches('/').to_string();
162 Self {
163 http: shared_http(),
164 base_url,
165 token,
166 vault_id: vault_id.into(),
167 }
168 }
169
170 fn url(&self, path: &str) -> String {
171 format!("{}{}", self.base_url, path)
172 }
173
174 fn auth(&self, req: reqwest::RequestBuilder) -> reqwest::RequestBuilder {
176 match &self.token {
177 Some(token) => req.bearer_auth(token),
178 None => req,
179 }
180 }
181
182 async fn ok(resp: reqwest::Response) -> Result<reqwest::Response, RagError> {
185 if resp.status().is_success() {
186 Ok(resp)
187 } else {
188 let status = resp.status().as_u16();
189 let body = resp.text().await.unwrap_or_default();
190 Err(RagError::Status { status, body })
191 }
192 }
193
194 pub async fn health(&self) -> Result<Health, RagError> {
196 let resp = self.auth(self.http.get(self.url("/health"))).send().await?;
197 Ok(Self::ok(resp).await?.json::<Health>().await?)
198 }
199
200 pub async fn push_docs(&self, docs: Vec<WireDoc>) -> Result<String, RagError> {
202 let body = IndexDocsRequest {
203 vault_id: self.vault_id.clone(),
204 docs,
205 };
206 let resp = self
207 .auth(self.http.post(self.url("/api/index/docs")).json(&body))
208 .timeout(PUSH_TIMEOUT)
209 .send()
210 .await?;
211 Ok(Self::ok(resp).await?.json::<JobAccepted>().await?.job_id)
212 }
213
214 pub async fn delete_paths(&self, paths: Vec<String>) -> Result<(), RagError> {
216 let body = DeleteRequest {
217 vault_id: self.vault_id.clone(),
218 paths,
219 };
220 let resp = self
221 .auth(self.http.post(self.url("/api/index/delete")).json(&body))
222 .send()
223 .await?;
224 Self::ok(resp).await?;
225 Ok(())
226 }
227
228 pub async fn server_hashes(&self) -> Result<HashMap<String, String>, RagError> {
234 let path = format!("/api/collections/{}/hashes", self.vault_id);
235 let resp = self.auth(self.http.get(self.url(&path))).send().await?;
236 Ok(Self::ok(resp)
237 .await?
238 .json::<HashMap<String, String>>()
239 .await?)
240 }
241
242 pub async fn search(
245 &self,
246 query: &str,
247 context_size: Option<ContextSize>,
248 ) -> Result<Vec<ChunkResult>, RagError> {
249 let body = QueryRequest {
250 vault_id: self.vault_id.clone(),
251 query: query.to_string(),
252 context_size: context_size.map(|c| c.as_str().to_string()),
253 history: vec![],
254 };
255 let resp = self
256 .auth(self.http.post(self.url("/api/embeddings")).json(&body))
257 .send()
258 .await?;
259 Ok(Self::ok(resp)
260 .await?
261 .json::<EmbeddingsResponse>()
262 .await?
263 .chunks)
264 }
265
266 pub async fn ask(
269 &self,
270 query: &str,
271 history: &[(String, String)],
272 context_size: Option<ContextSize>,
273 ) -> Result<AnswerResult, RagError> {
274 let body = QueryRequest {
275 vault_id: self.vault_id.clone(),
276 query: query.to_string(),
277 context_size: context_size.map(|c| c.as_str().to_string()),
278 history: history
279 .iter()
280 .map(|(q, a)| HistoryTurn {
281 question: q.clone(),
282 answer: a.clone(),
283 })
284 .collect(),
285 };
286 let resp = self
287 .auth(self.http.post(self.url("/api/answer")).json(&body))
288 .send()
289 .await?;
290 let job_id = Self::ok(resp).await?.json::<JobAccepted>().await?.job_id;
291 self.poll_answer(&job_id).await
292 }
293
294 const ANSWER_POLL_ATTEMPTS: u32 = 720;
301
302 async fn poll_answer(&self, job_id: &str) -> Result<AnswerResult, RagError> {
304 let path = format!("/api/job/{job_id}");
305 let mut consecutive_errors = 0u32;
306 for _ in 0..Self::ANSWER_POLL_ATTEMPTS {
307 let status = match self.poll_once(&path).await {
311 Ok(s) => {
312 consecutive_errors = 0;
313 s
314 }
315 Err(e) => {
316 consecutive_errors += 1;
317 if consecutive_errors >= 15 {
318 return Err(e);
319 }
320 tokio::time::sleep(std::time::Duration::from_secs(1)).await;
321 continue;
322 }
323 };
324 match status.status.as_str() {
325 "completed" => {
326 let result = status
327 .result
328 .ok_or_else(|| RagError::Protocol("completed job had no result".into()))?;
329 return serde_json::from_value::<AnswerResult>(result)
330 .map_err(|e| RagError::Protocol(format!("bad answer result: {e}")));
331 }
332 "failed" => {
333 return Err(RagError::Protocol(
334 status.error.unwrap_or_else(|| "answer job failed".into()),
335 ));
336 }
337 _ => tokio::time::sleep(std::time::Duration::from_secs(1)).await,
338 }
339 }
340 Err(RagError::Protocol("answer job timed out".into()))
341 }
342
343 async fn poll_once(&self, path: &str) -> Result<JobStatus, RagError> {
345 let resp = self.auth(self.http.get(self.url(path))).send().await?;
346 Ok(Self::ok(resp).await?.json::<JobStatus>().await?)
347 }
348}
349
350#[cfg(test)]
351mod tests {
352 use super::*;
353
354 #[test]
355 fn is_auth_matches_only_credential_rejections() {
356 let status = |status| RagError::Status {
357 status,
358 body: String::new(),
359 };
360 assert!(status(401).is_auth());
361 assert!(status(403).is_auth());
362 assert!(!status(500).is_auth());
363 assert!(!status(404).is_auth());
364 assert!(!RagError::Protocol("boom".into()).is_auth());
365 }
366}