1use reqwest::{Client, StatusCode};
4use std::sync::{Arc, Mutex};
5use std::time::Duration;
6use tracing::{debug, instrument};
7
8use serde::Deserialize;
9
10use crate::error::{ClientError, Result, ServerErrorCode};
11use crate::types::*;
12
13const DEFAULT_TIMEOUT_SECS: u64 = 30;
15
16#[derive(Debug, Clone)]
18pub struct DakeraClient {
19 pub(crate) client: Client,
21 pub(crate) base_url: String,
23 #[allow(dead_code)]
25 pub(crate) retry_config: RetryConfig,
26 pub(crate) last_rate_limit: Arc<Mutex<Option<RateLimitHeaders>>>,
28}
29
30impl DakeraClient {
31 pub fn new(base_url: impl Into<String>) -> Result<Self> {
41 DakeraClientBuilder::new(base_url).build()
42 }
43
44 pub fn builder(base_url: impl Into<String>) -> DakeraClientBuilder {
46 DakeraClientBuilder::new(base_url)
47 }
48
49 #[instrument(skip(self))]
55 pub async fn health(&self) -> Result<HealthResponse> {
56 let url = format!("{}/health", self.base_url);
57 let response = self.client.get(&url).send().await?;
58
59 if response.status().is_success() {
60 Ok(response.json().await?)
61 } else {
62 Ok(HealthResponse {
64 healthy: true,
65 version: None,
66 uptime_seconds: None,
67 })
68 }
69 }
70
71 #[instrument(skip(self))]
73 pub async fn ready(&self) -> Result<ReadinessResponse> {
74 let url = format!("{}/health/ready", self.base_url);
75 let response = self.client.get(&url).send().await?;
76
77 if response.status().is_success() {
78 Ok(response.json().await?)
79 } else {
80 Ok(ReadinessResponse {
81 ready: false,
82 components: None,
83 })
84 }
85 }
86
87 #[instrument(skip(self))]
89 pub async fn live(&self) -> Result<bool> {
90 let url = format!("{}/health/live", self.base_url);
91 let response = self.client.get(&url).send().await?;
92 Ok(response.status().is_success())
93 }
94
95 #[instrument(skip(self))]
101 pub async fn list_namespaces(&self) -> Result<Vec<String>> {
102 let url = format!("{}/v1/namespaces", self.base_url);
103 let response = self.client.get(&url).send().await?;
104 self.handle_response::<ListNamespacesResponse>(response)
105 .await
106 .map(|r| r.namespaces)
107 }
108
109 #[instrument(skip(self))]
111 pub async fn get_namespace(&self, namespace: &str) -> Result<NamespaceInfo> {
112 let url = format!("{}/v1/namespaces/{}", self.base_url, namespace);
113 let response = self.client.get(&url).send().await?;
114 self.handle_response(response).await
115 }
116
117 #[instrument(skip(self, request))]
119 pub async fn create_namespace(
120 &self,
121 namespace: &str,
122 request: CreateNamespaceRequest,
123 ) -> Result<NamespaceInfo> {
124 let url = format!("{}/v1/namespaces/{}", self.base_url, namespace);
125 let response = self.client.post(&url).json(&request).send().await?;
126 self.handle_response(response).await
127 }
128
129 #[instrument(skip(self, request), fields(namespace = %namespace))]
135 pub async fn configure_namespace(
136 &self,
137 namespace: &str,
138 request: ConfigureNamespaceRequest,
139 ) -> Result<ConfigureNamespaceResponse> {
140 let url = format!("{}/v1/namespaces/{}", self.base_url, namespace);
141 let response = self.client.put(&url).json(&request).send().await?;
142 self.handle_response(response).await
143 }
144
145 #[instrument(skip(self, request), fields(vector_count = request.vectors.len()))]
151 pub async fn upsert(&self, namespace: &str, request: UpsertRequest) -> Result<UpsertResponse> {
152 let url = format!("{}/v1/namespaces/{}/vectors", self.base_url, namespace);
153 debug!(
154 "Upserting {} vectors to {}",
155 request.vectors.len(),
156 namespace
157 );
158
159 let response = self.client.post(&url).json(&request).send().await?;
160 self.handle_response(response).await
161 }
162
163 #[instrument(skip(self, vector))]
165 pub async fn upsert_one(&self, namespace: &str, vector: Vector) -> Result<UpsertResponse> {
166 self.upsert(namespace, UpsertRequest::single(vector)).await
167 }
168
169 #[instrument(skip(self, request), fields(namespace = %namespace, count = request.ids.len()))]
202 pub async fn upsert_columns(
203 &self,
204 namespace: &str,
205 request: ColumnUpsertRequest,
206 ) -> Result<UpsertResponse> {
207 let url = format!(
208 "{}/v1/namespaces/{}/upsert-columns",
209 self.base_url, namespace
210 );
211 debug!(
212 "Upserting {} vectors in column format to {}",
213 request.ids.len(),
214 namespace
215 );
216
217 let response = self.client.post(&url).json(&request).send().await?;
218 self.handle_response(response).await
219 }
220
221 #[instrument(skip(self, request), fields(top_k = request.top_k))]
223 pub async fn query(&self, namespace: &str, request: QueryRequest) -> Result<QueryResponse> {
224 let url = format!("{}/v1/namespaces/{}/query", self.base_url, namespace);
225 debug!(
226 "Querying namespace {} for top {} results",
227 namespace, request.top_k
228 );
229
230 let response = self.client.post(&url).json(&request).send().await?;
231 self.handle_response(response).await
232 }
233
234 #[instrument(skip(self, vector))]
236 pub async fn query_simple(
237 &self,
238 namespace: &str,
239 vector: Vec<f32>,
240 top_k: u32,
241 ) -> Result<QueryResponse> {
242 self.query(namespace, QueryRequest::new(vector, top_k))
243 .await
244 }
245
246 #[instrument(skip(self, request), fields(namespace = %namespace, query_count = request.queries.len()))]
270 pub async fn batch_query(
271 &self,
272 namespace: &str,
273 request: BatchQueryRequest,
274 ) -> Result<BatchQueryResponse> {
275 let url = format!("{}/v1/namespaces/{}/batch-query", self.base_url, namespace);
276 debug!(
277 "Batch querying namespace {} with {} queries",
278 namespace,
279 request.queries.len()
280 );
281
282 let response = self.client.post(&url).json(&request).send().await?;
283 self.handle_response(response).await
284 }
285
286 #[instrument(skip(self, request), fields(id_count = request.ids.len()))]
288 pub async fn delete(&self, namespace: &str, request: DeleteRequest) -> Result<DeleteResponse> {
289 let url = format!(
290 "{}/v1/namespaces/{}/vectors/delete",
291 self.base_url, namespace
292 );
293 debug!("Deleting {} vectors from {}", request.ids.len(), namespace);
294
295 let response = self.client.post(&url).json(&request).send().await?;
296 self.handle_response(response).await
297 }
298
299 #[instrument(skip(self))]
301 pub async fn delete_one(&self, namespace: &str, id: &str) -> Result<DeleteResponse> {
302 self.delete(namespace, DeleteRequest::single(id)).await
303 }
304
305 #[instrument(skip(self, request), fields(doc_count = request.documents.len()))]
311 pub async fn index_documents(
312 &self,
313 namespace: &str,
314 request: IndexDocumentsRequest,
315 ) -> Result<IndexDocumentsResponse> {
316 let url = format!(
317 "{}/v1/namespaces/{}/fulltext/index",
318 self.base_url, namespace
319 );
320 debug!(
321 "Indexing {} documents in {}",
322 request.documents.len(),
323 namespace
324 );
325
326 let response = self.client.post(&url).json(&request).send().await?;
327 self.handle_response(response).await
328 }
329
330 #[instrument(skip(self, document))]
332 pub async fn index_document(
333 &self,
334 namespace: &str,
335 document: Document,
336 ) -> Result<IndexDocumentsResponse> {
337 self.index_documents(
338 namespace,
339 IndexDocumentsRequest {
340 documents: vec![document],
341 },
342 )
343 .await
344 }
345
346 #[instrument(skip(self, request))]
348 pub async fn fulltext_search(
349 &self,
350 namespace: &str,
351 request: FullTextSearchRequest,
352 ) -> Result<FullTextSearchResponse> {
353 let url = format!(
354 "{}/v1/namespaces/{}/fulltext/search",
355 self.base_url, namespace
356 );
357 debug!("Full-text search in {} for: {}", namespace, request.query);
358
359 let response = self.client.post(&url).json(&request).send().await?;
360 self.handle_response(response).await
361 }
362
363 #[instrument(skip(self))]
365 pub async fn search_text(
366 &self,
367 namespace: &str,
368 query: &str,
369 top_k: u32,
370 ) -> Result<FullTextSearchResponse> {
371 self.fulltext_search(namespace, FullTextSearchRequest::new(query, top_k))
372 .await
373 }
374
375 #[instrument(skip(self))]
377 pub async fn fulltext_stats(&self, namespace: &str) -> Result<FullTextStats> {
378 let url = format!(
379 "{}/v1/namespaces/{}/fulltext/stats",
380 self.base_url, namespace
381 );
382 let response = self.client.get(&url).send().await?;
383 self.handle_response(response).await
384 }
385
386 #[instrument(skip(self, request))]
388 pub async fn fulltext_delete(
389 &self,
390 namespace: &str,
391 request: DeleteRequest,
392 ) -> Result<DeleteResponse> {
393 let url = format!(
394 "{}/v1/namespaces/{}/fulltext/delete",
395 self.base_url, namespace
396 );
397 let response = self.client.post(&url).json(&request).send().await?;
398 self.handle_response(response).await
399 }
400
401 #[instrument(skip(self, request), fields(top_k = request.top_k))]
407 pub async fn hybrid_search(
408 &self,
409 namespace: &str,
410 request: HybridSearchRequest,
411 ) -> Result<HybridSearchResponse> {
412 let url = format!("{}/v1/namespaces/{}/hybrid", self.base_url, namespace);
413 debug!(
414 "Hybrid search in {} with vector_weight={}",
415 namespace, request.vector_weight
416 );
417
418 let response = self.client.post(&url).json(&request).send().await?;
419 self.handle_response(response).await
420 }
421
422 #[instrument(skip(self, request), fields(namespace = %namespace))]
459 pub async fn multi_vector_search(
460 &self,
461 namespace: &str,
462 request: MultiVectorSearchRequest,
463 ) -> Result<MultiVectorSearchResponse> {
464 let url = format!("{}/v1/namespaces/{}/multi-vector", self.base_url, namespace);
465 debug!(
466 "Multi-vector search in {} with {} positive vectors",
467 namespace,
468 request.positive_vectors.len()
469 );
470
471 let response = self.client.post(&url).json(&request).send().await?;
472 self.handle_response(response).await
473 }
474
475 #[instrument(skip(self, request), fields(namespace = %namespace))]
509 pub async fn aggregate(
510 &self,
511 namespace: &str,
512 request: AggregationRequest,
513 ) -> Result<AggregationResponse> {
514 let url = format!("{}/v1/namespaces/{}/aggregate", self.base_url, namespace);
515 debug!(
516 "Aggregating in namespace {} with {} aggregations",
517 namespace,
518 request.aggregate_by.len()
519 );
520
521 let response = self.client.post(&url).json(&request).send().await?;
522 self.handle_response(response).await
523 }
524
525 #[instrument(skip(self, request), fields(namespace = %namespace))]
563 pub async fn unified_query(
564 &self,
565 namespace: &str,
566 request: UnifiedQueryRequest,
567 ) -> Result<UnifiedQueryResponse> {
568 let url = format!(
569 "{}/v1/namespaces/{}/unified-query",
570 self.base_url, namespace
571 );
572 debug!(
573 "Unified query in namespace {} with top_k={}",
574 namespace, request.top_k
575 );
576
577 let response = self.client.post(&url).json(&request).send().await?;
578 self.handle_response(response).await
579 }
580
581 #[instrument(skip(self, vector))]
585 pub async fn unified_vector_search(
586 &self,
587 namespace: &str,
588 vector: Vec<f32>,
589 top_k: usize,
590 ) -> Result<UnifiedQueryResponse> {
591 self.unified_query(namespace, UnifiedQueryRequest::vector_search(vector, top_k))
592 .await
593 }
594
595 #[instrument(skip(self))]
599 pub async fn unified_text_search(
600 &self,
601 namespace: &str,
602 field: &str,
603 query: &str,
604 top_k: usize,
605 ) -> Result<UnifiedQueryResponse> {
606 self.unified_query(
607 namespace,
608 UnifiedQueryRequest::fulltext_search(field, query, top_k),
609 )
610 .await
611 }
612
613 #[instrument(skip(self, request), fields(namespace = %namespace))]
650 pub async fn explain_query(
651 &self,
652 namespace: &str,
653 request: QueryExplainRequest,
654 ) -> Result<QueryExplainResponse> {
655 let url = format!("{}/v1/namespaces/{}/explain", self.base_url, namespace);
656 debug!(
657 "Explaining query in namespace {} (query_type={:?}, top_k={})",
658 namespace, request.query_type, request.top_k
659 );
660
661 let response = self.client.post(&url).json(&request).send().await?;
662 self.handle_response(response).await
663 }
664
665 #[instrument(skip(self, request), fields(namespace = %request.namespace, priority = ?request.priority))]
693 pub async fn warm_cache(&self, request: WarmCacheRequest) -> Result<WarmCacheResponse> {
694 let url = format!(
695 "{}/v1/namespaces/{}/cache/warm",
696 self.base_url, request.namespace
697 );
698 debug!(
699 "Warming cache for namespace {} with priority {:?}",
700 request.namespace, request.priority
701 );
702
703 let response = self.client.post(&url).json(&request).send().await?;
704 self.handle_response(response).await
705 }
706
707 #[instrument(skip(self, vector_ids))]
709 pub async fn warm_vectors(
710 &self,
711 namespace: &str,
712 vector_ids: Vec<String>,
713 ) -> Result<WarmCacheResponse> {
714 self.warm_cache(WarmCacheRequest::new(namespace).with_vector_ids(vector_ids))
715 .await
716 }
717
718 #[instrument(skip(self, request), fields(namespace = %namespace))]
751 pub async fn export(&self, namespace: &str, request: ExportRequest) -> Result<ExportResponse> {
752 let url = format!("{}/v1/namespaces/{}/export", self.base_url, namespace);
753 debug!(
754 "Exporting vectors from namespace {} (top_k={}, cursor={:?})",
755 namespace, request.top_k, request.cursor
756 );
757
758 let response = self.client.post(&url).json(&request).send().await?;
759 self.handle_response(response).await
760 }
761
762 #[instrument(skip(self))]
766 pub async fn export_all(&self, namespace: &str) -> Result<ExportResponse> {
767 self.export(namespace, ExportRequest::new()).await
768 }
769
770 #[instrument(skip(self))]
776 pub async fn diagnostics(&self) -> Result<SystemDiagnostics> {
777 let url = format!("{}/ops/diagnostics", self.base_url);
778 let response = self.client.get(&url).send().await?;
779 self.handle_response(response).await
780 }
781
782 #[instrument(skip(self))]
784 pub async fn list_jobs(&self) -> Result<Vec<JobInfo>> {
785 let url = format!("{}/ops/jobs", self.base_url);
786 let response = self.client.get(&url).send().await?;
787 self.handle_response(response).await
788 }
789
790 #[instrument(skip(self))]
792 pub async fn get_job(&self, job_id: &str) -> Result<Option<JobInfo>> {
793 let url = format!("{}/ops/jobs/{}", self.base_url, job_id);
794 let response = self.client.get(&url).send().await?;
795
796 if response.status() == StatusCode::NOT_FOUND {
797 return Ok(None);
798 }
799
800 self.handle_response(response).await.map(Some)
801 }
802
803 #[instrument(skip(self, request))]
805 pub async fn compact(&self, request: CompactionRequest) -> Result<CompactionResponse> {
806 let url = format!("{}/ops/compact", self.base_url);
807 let response = self.client.post(&url).json(&request).send().await?;
808 self.handle_response(response).await
809 }
810
811 #[instrument(skip(self))]
813 pub async fn shutdown(&self) -> Result<()> {
814 let url = format!("{}/ops/shutdown", self.base_url);
815 let response = self.client.post(&url).send().await?;
816
817 if response.status().is_success() {
818 Ok(())
819 } else {
820 let status = response.status().as_u16();
821 let text = response.text().await.unwrap_or_default();
822 Err(ClientError::Server {
823 status,
824 message: text,
825 code: None,
826 })
827 }
828 }
829
830 #[instrument(skip(self, request), fields(id_count = request.ids.len()))]
836 pub async fn fetch(&self, namespace: &str, request: FetchRequest) -> Result<FetchResponse> {
837 let url = format!("{}/v1/namespaces/{}/fetch", self.base_url, namespace);
838 debug!("Fetching {} vectors from {}", request.ids.len(), namespace);
839 let response = self.client.post(&url).json(&request).send().await?;
840 self.handle_response(response).await
841 }
842
843 #[instrument(skip(self))]
845 pub async fn fetch_by_ids(&self, namespace: &str, ids: &[&str]) -> Result<Vec<Vector>> {
846 let request = FetchRequest::new(ids.iter().map(|s| s.to_string()).collect());
847 self.fetch(namespace, request).await.map(|r| r.vectors)
848 }
849
850 #[instrument(skip(self, request), fields(doc_count = request.documents.len()))]
856 pub async fn upsert_text(
857 &self,
858 namespace: &str,
859 request: UpsertTextRequest,
860 ) -> Result<TextUpsertResponse> {
861 let url = format!("{}/v1/namespaces/{}/upsert-text", self.base_url, namespace);
862 debug!(
863 "Upserting {} text documents to {}",
864 request.documents.len(),
865 namespace
866 );
867 let response = self.client.post(&url).json(&request).send().await?;
868 self.handle_response(response).await
869 }
870
871 #[instrument(skip(self, request), fields(top_k = request.top_k))]
873 pub async fn query_text(
874 &self,
875 namespace: &str,
876 request: QueryTextRequest,
877 ) -> Result<TextQueryResponse> {
878 let url = format!("{}/v1/namespaces/{}/query-text", self.base_url, namespace);
879 debug!("Text query in {} for: {}", namespace, request.text);
880 let response = self.client.post(&url).json(&request).send().await?;
881 self.handle_response(response).await
882 }
883
884 #[instrument(skip(self))]
886 pub async fn query_text_simple(
887 &self,
888 namespace: &str,
889 text: &str,
890 top_k: u32,
891 ) -> Result<TextQueryResponse> {
892 self.query_text(namespace, QueryTextRequest::new(text, top_k))
893 .await
894 }
895
896 #[instrument(skip(self, request), fields(query_count = request.queries.len()))]
898 pub async fn batch_query_text(
899 &self,
900 namespace: &str,
901 request: BatchQueryTextRequest,
902 ) -> Result<BatchQueryTextResponse> {
903 let url = format!(
904 "{}/v1/namespaces/{}/batch-query-text",
905 self.base_url, namespace
906 );
907 debug!(
908 "Batch text query in {} with {} queries",
909 namespace,
910 request.queries.len()
911 );
912 let response = self.client.post(&url).json(&request).send().await?;
913 self.handle_response(response).await
914 }
915
916 #[instrument(skip(self, config))]
925 pub async fn configure_namespace_ner(
926 &self,
927 namespace: &str,
928 config: NamespaceNerConfig,
929 ) -> Result<serde_json::Value> {
930 let url = format!("{}/v1/namespaces/{}/config", self.base_url, namespace);
931 let response = self.client.patch(&url).json(&config).send().await?;
932 self.handle_response(response).await
933 }
934
935 #[instrument(skip(self, text, entity_types))]
940 pub async fn extract_entities(
941 &self,
942 text: &str,
943 entity_types: Option<Vec<String>>,
944 ) -> Result<EntityExtractionResponse> {
945 let url = format!("{}/v1/memories/extract", self.base_url);
946 let body = serde_json::json!({
947 "text": text,
948 "entity_types": entity_types,
949 });
950 let response = self.client.post(&url).json(&body).send().await?;
951 self.handle_response(response).await
952 }
953
954 #[instrument(skip(self))]
958 pub async fn memory_entities(&self, memory_id: &str) -> Result<MemoryEntitiesResponse> {
959 let url = format!("{}/v1/memory/entities/{}", self.base_url, memory_id);
960 let response = self.client.get(&url).send().await?;
961 self.handle_response(response).await
962 }
963
964 pub fn last_rate_limit_headers(&self) -> Option<RateLimitHeaders> {
972 self.last_rate_limit.lock().ok()?.clone()
973 }
974
975 pub(crate) async fn handle_response<T: serde::de::DeserializeOwned>(
977 &self,
978 response: reqwest::Response,
979 ) -> Result<T> {
980 let status = response.status();
981
982 if let Ok(mut guard) = self.last_rate_limit.lock() {
984 *guard = Some(RateLimitHeaders::from_response(&response));
985 }
986
987 if status.is_success() {
988 Ok(response.json().await?)
989 } else {
990 let status_code = status.as_u16();
991 let retry_after = response
993 .headers()
994 .get("Retry-After")
995 .and_then(|v| v.to_str().ok())
996 .and_then(|s| s.parse::<u64>().ok());
997 let text = response.text().await.unwrap_or_default();
998
999 if status_code == 429 {
1000 return Err(ClientError::RateLimitExceeded { retry_after });
1001 }
1002
1003 #[derive(Deserialize)]
1004 struct ErrorBody {
1005 error: Option<String>,
1006 code: Option<ServerErrorCode>,
1007 }
1008
1009 let (message, code) = if let Ok(body) = serde_json::from_str::<ErrorBody>(&text) {
1010 (body.error.unwrap_or_else(|| text.clone()), body.code)
1011 } else {
1012 (text, None)
1013 };
1014
1015 match status_code {
1016 401 => Err(ClientError::Server {
1017 status: 401,
1018 message,
1019 code,
1020 }),
1021 403 => Err(ClientError::Authorization {
1022 status: 403,
1023 message,
1024 code,
1025 }),
1026 404 => match &code {
1027 Some(ServerErrorCode::NamespaceNotFound) => {
1028 Err(ClientError::NamespaceNotFound(message))
1029 }
1030 Some(ServerErrorCode::VectorNotFound) => {
1031 Err(ClientError::VectorNotFound(message))
1032 }
1033 _ => Err(ClientError::Server {
1034 status: 404,
1035 message,
1036 code,
1037 }),
1038 },
1039 _ => Err(ClientError::Server {
1040 status: status_code,
1041 message,
1042 code,
1043 }),
1044 }
1045 }
1046 }
1047
1048 pub(crate) async fn handle_text_response(&self, response: reqwest::Response) -> Result<String> {
1050 let status = response.status();
1051
1052 if let Ok(mut guard) = self.last_rate_limit.lock() {
1054 *guard = Some(RateLimitHeaders::from_response(&response));
1055 }
1056
1057 let retry_after = response
1058 .headers()
1059 .get("Retry-After")
1060 .and_then(|v| v.to_str().ok())
1061 .and_then(|s| s.parse::<u64>().ok());
1062 let text = response.text().await.unwrap_or_default();
1063
1064 if status.is_success() {
1065 return Ok(text);
1066 }
1067
1068 let status_code = status.as_u16();
1069
1070 if status_code == 429 {
1071 return Err(ClientError::RateLimitExceeded { retry_after });
1072 }
1073
1074 #[derive(Deserialize)]
1075 struct ErrorBody {
1076 error: Option<String>,
1077 code: Option<ServerErrorCode>,
1078 }
1079
1080 let (message, code) = if let Ok(body) = serde_json::from_str::<ErrorBody>(&text) {
1081 (body.error.unwrap_or_else(|| text.clone()), body.code)
1082 } else {
1083 (text, None)
1084 };
1085
1086 match status_code {
1087 401 => Err(ClientError::Server {
1088 status: 401,
1089 message,
1090 code,
1091 }),
1092 403 => Err(ClientError::Authorization {
1093 status: 403,
1094 message,
1095 code,
1096 }),
1097 _ => Err(ClientError::Server {
1098 status: status_code,
1099 message,
1100 code,
1101 }),
1102 }
1103 }
1104
1105 #[allow(dead_code)]
1113 pub(crate) async fn execute_with_retry<F, Fut, T>(&self, f: F) -> Result<T>
1114 where
1115 F: Fn() -> Fut,
1116 Fut: std::future::Future<Output = Result<T>>,
1117 {
1118 let rc = &self.retry_config;
1119
1120 for attempt in 0..rc.max_retries {
1121 match f().await {
1122 Ok(v) => return Ok(v),
1123 Err(e) => {
1124 let is_last = attempt == rc.max_retries - 1;
1125 if is_last || !e.is_retryable() {
1126 return Err(e);
1127 }
1128
1129 let wait = match &e {
1130 ClientError::RateLimitExceeded {
1131 retry_after: Some(secs),
1132 } => Duration::from_secs(*secs),
1133 _ => {
1134 let base_ms = rc.base_delay.as_millis() as f64;
1135 let backoff_ms = base_ms * 2f64.powi(attempt as i32);
1136 let capped_ms = backoff_ms.min(rc.max_delay.as_millis() as f64);
1137 let final_ms = if rc.jitter {
1138 let seed = (attempt as u64).wrapping_mul(6364136223846793005);
1140 let factor = 0.5 + (seed % 1000) as f64 / 1000.0;
1141 capped_ms * factor
1142 } else {
1143 capped_ms
1144 };
1145 Duration::from_millis(final_ms as u64)
1146 }
1147 };
1148
1149 tokio::time::sleep(wait).await;
1150 }
1151 }
1152 }
1153
1154 Err(ClientError::Config("retry loop exhausted".to_string()))
1156 }
1157}
1158
1159#[derive(Debug)]
1161pub struct DakeraClientBuilder {
1162 base_url: String,
1163 timeout: Duration,
1164 connect_timeout: Option<Duration>,
1165 retry_config: RetryConfig,
1166 user_agent: Option<String>,
1167}
1168
1169impl DakeraClientBuilder {
1170 pub fn new(base_url: impl Into<String>) -> Self {
1172 Self {
1173 base_url: base_url.into(),
1174 timeout: Duration::from_secs(DEFAULT_TIMEOUT_SECS),
1175 connect_timeout: None,
1176 retry_config: RetryConfig::default(),
1177 user_agent: None,
1178 }
1179 }
1180
1181 pub fn timeout(mut self, timeout: Duration) -> Self {
1183 self.timeout = timeout;
1184 self
1185 }
1186
1187 pub fn timeout_secs(mut self, secs: u64) -> Self {
1189 self.timeout = Duration::from_secs(secs);
1190 self
1191 }
1192
1193 pub fn connect_timeout(mut self, timeout: Duration) -> Self {
1195 self.connect_timeout = Some(timeout);
1196 self
1197 }
1198
1199 pub fn retry_config(mut self, config: RetryConfig) -> Self {
1201 self.retry_config = config;
1202 self
1203 }
1204
1205 pub fn max_retries(mut self, max_retries: u32) -> Self {
1207 self.retry_config.max_retries = max_retries;
1208 self
1209 }
1210
1211 pub fn user_agent(mut self, user_agent: impl Into<String>) -> Self {
1213 self.user_agent = Some(user_agent.into());
1214 self
1215 }
1216
1217 pub fn build(self) -> Result<DakeraClient> {
1219 let base_url = self.base_url.trim_end_matches('/').to_string();
1221
1222 if !base_url.starts_with("http://") && !base_url.starts_with("https://") {
1224 return Err(ClientError::InvalidUrl(
1225 "URL must start with http:// or https://".to_string(),
1226 ));
1227 }
1228
1229 let user_agent = self
1230 .user_agent
1231 .unwrap_or_else(|| format!("dakera-client/{}", env!("CARGO_PKG_VERSION")));
1232
1233 let connect_timeout = self.connect_timeout.unwrap_or(self.timeout);
1234
1235 let client = Client::builder()
1236 .timeout(self.timeout)
1237 .connect_timeout(connect_timeout)
1238 .user_agent(user_agent)
1239 .build()
1240 .map_err(|e| ClientError::Config(e.to_string()))?;
1241
1242 Ok(DakeraClient {
1243 client,
1244 base_url,
1245 retry_config: self.retry_config,
1246 last_rate_limit: Arc::new(Mutex::new(None)),
1247 })
1248 }
1249}
1250
1251impl DakeraClient {
1256 pub async fn stream_namespace_events(
1281 &self,
1282 namespace: &str,
1283 ) -> Result<tokio::sync::mpsc::Receiver<Result<crate::events::DakeraEvent>>> {
1284 let url = format!(
1285 "{}/v1/namespaces/{}/events",
1286 self.base_url,
1287 urlencoding::encode(namespace)
1288 );
1289 self.stream_sse(url).await
1290 }
1291
1292 pub async fn stream_global_events(
1299 &self,
1300 ) -> Result<tokio::sync::mpsc::Receiver<Result<crate::events::DakeraEvent>>> {
1301 let url = format!("{}/ops/events", self.base_url);
1302 self.stream_sse(url).await
1303 }
1304
1305 pub async fn stream_memory_events(
1314 &self,
1315 ) -> Result<tokio::sync::mpsc::Receiver<Result<crate::events::MemoryEvent>>> {
1316 let url = format!("{}/v1/events/stream", self.base_url);
1317 self.stream_sse(url).await
1318 }
1319
1320 pub(crate) async fn stream_sse<T>(
1322 &self,
1323 url: String,
1324 ) -> Result<tokio::sync::mpsc::Receiver<Result<T>>>
1325 where
1326 T: serde::de::DeserializeOwned + Send + 'static,
1327 {
1328 use futures_util::StreamExt;
1329
1330 let response = self
1331 .client
1332 .get(&url)
1333 .header("Accept", "text/event-stream")
1334 .header("Cache-Control", "no-cache")
1335 .send()
1336 .await?;
1337
1338 if !response.status().is_success() {
1339 let status = response.status().as_u16();
1340 let body = response.text().await.unwrap_or_default();
1341 return Err(ClientError::Server {
1342 status,
1343 message: body,
1344 code: None,
1345 });
1346 }
1347
1348 let (tx, rx) = tokio::sync::mpsc::channel(64);
1349
1350 tokio::spawn(async move {
1351 let mut byte_stream = response.bytes_stream();
1352 let mut remaining = String::new();
1353 let mut data_lines: Vec<String> = Vec::new();
1354
1355 while let Some(chunk) = byte_stream.next().await {
1356 match chunk {
1357 Ok(bytes) => {
1358 remaining.push_str(&String::from_utf8_lossy(&bytes));
1359 while let Some(pos) = remaining.find('\n') {
1360 let raw = &remaining[..pos];
1361 let line = raw.trim_end_matches('\r').to_string();
1362 remaining = remaining[pos + 1..].to_string();
1363
1364 if line.starts_with(':') {
1365 } else if let Some(data) = line.strip_prefix("data:") {
1367 data_lines.push(data.trim_start().to_string());
1368 } else if line.is_empty() {
1369 if !data_lines.is_empty() {
1370 let payload = data_lines.join("\n");
1371 data_lines.clear();
1372 let result = serde_json::from_str::<T>(&payload)
1373 .map_err(ClientError::Json);
1374 if tx.send(result).await.is_err() {
1375 return; }
1377 }
1378 } else {
1379 }
1381 }
1382 }
1383 Err(e) => {
1384 let _ = tx.send(Err(ClientError::Http(e))).await;
1385 return;
1386 }
1387 }
1388 }
1389 });
1390
1391 Ok(rx)
1392 }
1393}
1394
1395#[cfg(test)]
1396mod tests {
1397 use super::*;
1398
1399 #[test]
1400 fn test_client_builder() {
1401 let client = DakeraClient::new("http://localhost:3000");
1402 assert!(client.is_ok());
1403 }
1404
1405 #[test]
1406 fn test_client_builder_with_options() {
1407 let client = DakeraClient::builder("http://localhost:3000")
1408 .timeout_secs(60)
1409 .user_agent("test-client/1.0")
1410 .build();
1411 assert!(client.is_ok());
1412 }
1413
1414 #[test]
1415 fn test_client_builder_invalid_url() {
1416 let client = DakeraClient::new("invalid-url");
1417 assert!(client.is_err());
1418 }
1419
1420 #[test]
1421 fn test_client_builder_trailing_slash() {
1422 let client = DakeraClient::new("http://localhost:3000/").unwrap();
1423 assert!(!client.base_url.ends_with('/'));
1424 }
1425
1426 #[test]
1427 fn test_vector_creation() {
1428 let v = Vector::new("test", vec![0.1, 0.2, 0.3]);
1429 assert_eq!(v.id, "test");
1430 assert_eq!(v.values.len(), 3);
1431 assert!(v.metadata.is_none());
1432 }
1433
1434 #[test]
1435 fn test_query_request_builder() {
1436 let req = QueryRequest::new(vec![0.1, 0.2], 10)
1437 .with_filter(serde_json::json!({"category": "test"}))
1438 .include_metadata(false);
1439
1440 assert_eq!(req.top_k, 10);
1441 assert!(req.filter.is_some());
1442 assert!(!req.include_metadata);
1443 }
1444
1445 #[test]
1446 fn test_hybrid_search_request() {
1447 let req = HybridSearchRequest::new(vec![0.1], "test query", 5).with_vector_weight(0.7);
1448
1449 assert_eq!(req.vector_weight, 0.7);
1450 assert_eq!(req.text, "test query");
1451 assert!(req.vector.is_some());
1452 }
1453
1454 #[test]
1455 fn test_hybrid_search_weight_clamping() {
1456 let req = HybridSearchRequest::new(vec![0.1], "test", 5).with_vector_weight(1.5); assert_eq!(req.vector_weight, 1.0);
1459 }
1460
1461 #[test]
1462 fn test_hybrid_search_text_only() {
1463 let req = HybridSearchRequest::text_only("bm25 query", 10);
1464
1465 assert!(req.vector.is_none());
1466 assert_eq!(req.text, "bm25 query");
1467 assert_eq!(req.top_k, 10);
1468 let json = serde_json::to_value(&req).unwrap();
1470 assert!(json.get("vector").is_none());
1471 }
1472
1473 #[test]
1474 fn test_text_document_builder() {
1475 let doc = TextDocument::new("doc1", "Hello world").with_ttl(3600);
1476
1477 assert_eq!(doc.id, "doc1");
1478 assert_eq!(doc.text, "Hello world");
1479 assert_eq!(doc.ttl_seconds, Some(3600));
1480 assert!(doc.metadata.is_none());
1481 }
1482
1483 #[test]
1484 fn test_upsert_text_request_builder() {
1485 let docs = vec![
1486 TextDocument::new("doc1", "Hello"),
1487 TextDocument::new("doc2", "World"),
1488 ];
1489 let req = UpsertTextRequest::new(docs).with_model(EmbeddingModel::BgeSmall);
1490
1491 assert_eq!(req.documents.len(), 2);
1492 assert_eq!(req.model, Some(EmbeddingModel::BgeSmall));
1493 }
1494
1495 #[test]
1496 fn test_query_text_request_builder() {
1497 let req = QueryTextRequest::new("semantic search query", 5)
1498 .with_filter(serde_json::json!({"category": "docs"}))
1499 .include_vectors(true)
1500 .with_model(EmbeddingModel::E5Small);
1501
1502 assert_eq!(req.text, "semantic search query");
1503 assert_eq!(req.top_k, 5);
1504 assert!(req.filter.is_some());
1505 assert!(req.include_vectors);
1506 assert_eq!(req.model, Some(EmbeddingModel::E5Small));
1507 }
1508
1509 #[test]
1510 fn test_fetch_request_builder() {
1511 let req = FetchRequest::new(vec!["id1".to_string(), "id2".to_string()]);
1512
1513 assert_eq!(req.ids.len(), 2);
1514 assert!(req.include_values);
1515 assert!(req.include_metadata);
1516 }
1517
1518 #[test]
1519 fn test_create_namespace_request_builder() {
1520 let req = CreateNamespaceRequest::new()
1521 .with_dimensions(384)
1522 .with_index_type("hnsw");
1523
1524 assert_eq!(req.dimensions, Some(384));
1525 assert_eq!(req.index_type.as_deref(), Some("hnsw"));
1526 }
1527
1528 #[test]
1529 fn test_batch_query_text_request() {
1530 let req =
1531 BatchQueryTextRequest::new(vec!["query one".to_string(), "query two".to_string()], 10);
1532
1533 assert_eq!(req.queries.len(), 2);
1534 assert_eq!(req.top_k, 10);
1535 assert!(!req.include_vectors);
1536 assert!(req.model.is_none());
1537 }
1538
1539 #[test]
1544 fn test_retry_config_defaults() {
1545 let rc = RetryConfig::default();
1546 assert_eq!(rc.max_retries, 3);
1547 assert_eq!(rc.base_delay, Duration::from_millis(100));
1548 assert_eq!(rc.max_delay, Duration::from_secs(60));
1549 assert!(rc.jitter);
1550 }
1551
1552 #[test]
1553 fn test_builder_connect_timeout() {
1554 let client = DakeraClient::builder("http://localhost:3000")
1555 .connect_timeout(Duration::from_secs(5))
1556 .timeout_secs(30)
1557 .build()
1558 .unwrap();
1559 assert!(client.base_url.starts_with("http"));
1561 }
1562
1563 #[test]
1564 fn test_builder_max_retries() {
1565 let client = DakeraClient::builder("http://localhost:3000")
1566 .max_retries(5)
1567 .build()
1568 .unwrap();
1569 assert_eq!(client.retry_config.max_retries, 5);
1570 }
1571
1572 #[test]
1573 fn test_builder_retry_config() {
1574 let rc = RetryConfig {
1575 max_retries: 7,
1576 base_delay: Duration::from_millis(200),
1577 max_delay: Duration::from_secs(30),
1578 jitter: false,
1579 };
1580 let client = DakeraClient::builder("http://localhost:3000")
1581 .retry_config(rc)
1582 .build()
1583 .unwrap();
1584 assert_eq!(client.retry_config.max_retries, 7);
1585 assert!(!client.retry_config.jitter);
1586 }
1587
1588 #[test]
1589 fn test_rate_limit_error_retryable() {
1590 let e = ClientError::RateLimitExceeded { retry_after: None };
1591 assert!(e.is_retryable());
1592 }
1593
1594 #[test]
1595 fn test_rate_limit_error_with_retry_after_zero() {
1596 let e = ClientError::RateLimitExceeded {
1598 retry_after: Some(0),
1599 };
1600 assert!(e.is_retryable());
1601 if let ClientError::RateLimitExceeded {
1602 retry_after: Some(secs),
1603 } = &e
1604 {
1605 assert_eq!(*secs, 0u64);
1606 } else {
1607 panic!("unexpected variant");
1608 }
1609 }
1610
1611 #[tokio::test]
1612 async fn test_execute_with_retry_succeeds_immediately() {
1613 let client = DakeraClient::builder("http://localhost:3000")
1614 .max_retries(3)
1615 .build()
1616 .unwrap();
1617
1618 let call_count = std::sync::Arc::new(std::sync::atomic::AtomicU32::new(0));
1619 let cc = call_count.clone();
1620 let result = client
1621 .execute_with_retry(|| {
1622 let cc = cc.clone();
1623 async move {
1624 cc.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
1625 Ok::<u32, ClientError>(42)
1626 }
1627 })
1628 .await;
1629 assert_eq!(result.unwrap(), 42);
1630 assert_eq!(call_count.load(std::sync::atomic::Ordering::SeqCst), 1);
1631 }
1632
1633 #[tokio::test]
1634 async fn test_execute_with_retry_no_retry_on_4xx() {
1635 let client = DakeraClient::builder("http://localhost:3000")
1636 .max_retries(3)
1637 .build()
1638 .unwrap();
1639
1640 let call_count = std::sync::Arc::new(std::sync::atomic::AtomicU32::new(0));
1641 let cc = call_count.clone();
1642 let result = client
1643 .execute_with_retry(|| {
1644 let cc = cc.clone();
1645 async move {
1646 cc.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
1647 Err::<u32, ClientError>(ClientError::Server {
1648 status: 400,
1649 message: "bad request".to_string(),
1650 code: None,
1651 })
1652 }
1653 })
1654 .await;
1655 assert!(result.is_err());
1656 assert_eq!(call_count.load(std::sync::atomic::Ordering::SeqCst), 1);
1658 }
1659
1660 #[tokio::test]
1661 async fn test_execute_with_retry_retries_on_5xx() {
1662 let client = DakeraClient::builder("http://localhost:3000")
1663 .retry_config(RetryConfig {
1664 max_retries: 3,
1665 base_delay: Duration::from_millis(0),
1666 max_delay: Duration::from_millis(0),
1667 jitter: false,
1668 })
1669 .build()
1670 .unwrap();
1671
1672 let call_count = std::sync::Arc::new(std::sync::atomic::AtomicU32::new(0));
1673 let cc = call_count.clone();
1674 let result = client
1675 .execute_with_retry(|| {
1676 let cc = cc.clone();
1677 async move {
1678 let n = cc.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
1679 if n < 2 {
1680 Err::<u32, ClientError>(ClientError::Server {
1681 status: 503,
1682 message: "unavailable".to_string(),
1683 code: None,
1684 })
1685 } else {
1686 Ok(99)
1687 }
1688 }
1689 })
1690 .await;
1691 assert_eq!(result.unwrap(), 99);
1692 assert_eq!(call_count.load(std::sync::atomic::Ordering::SeqCst), 3);
1693 }
1694
1695 #[test]
1700 fn test_batch_recall_request_new() {
1701 use crate::memory::BatchRecallRequest;
1702 let req = BatchRecallRequest::new("agent-1");
1703 assert_eq!(req.agent_id, "agent-1");
1704 assert_eq!(req.limit, 100);
1705 }
1706
1707 #[test]
1708 fn test_batch_recall_request_builder() {
1709 use crate::memory::{BatchMemoryFilter, BatchRecallRequest};
1710 let filter = BatchMemoryFilter::default()
1711 .with_tags(vec!["qa".to_string()])
1712 .with_min_importance(0.7);
1713 let req = BatchRecallRequest::new("agent-1")
1714 .with_filter(filter)
1715 .with_limit(50);
1716 assert_eq!(req.agent_id, "agent-1");
1717 assert_eq!(req.limit, 50);
1718 assert_eq!(
1719 req.filter.tags.as_deref(),
1720 Some(["qa".to_string()].as_slice())
1721 );
1722 assert_eq!(req.filter.min_importance, Some(0.7));
1723 }
1724
1725 #[test]
1726 fn test_batch_recall_request_serialization() {
1727 use crate::memory::{BatchMemoryFilter, BatchRecallRequest};
1728 let filter = BatchMemoryFilter::default().with_min_importance(0.5);
1729 let req = BatchRecallRequest::new("agent-1")
1730 .with_filter(filter)
1731 .with_limit(25);
1732 let json = serde_json::to_value(&req).unwrap();
1733 assert_eq!(json["agent_id"], "agent-1");
1734 assert_eq!(json["limit"], 25);
1735 assert_eq!(json["filter"]["min_importance"], 0.5);
1736 }
1737
1738 #[test]
1739 fn test_batch_forget_request_new() {
1740 use crate::memory::{BatchForgetRequest, BatchMemoryFilter};
1741 let filter = BatchMemoryFilter::default().with_min_importance(0.1);
1742 let req = BatchForgetRequest::new("agent-1", filter);
1743 assert_eq!(req.agent_id, "agent-1");
1744 assert_eq!(req.filter.min_importance, Some(0.1));
1745 }
1746
1747 #[test]
1748 fn test_batch_forget_request_serialization() {
1749 use crate::memory::{BatchForgetRequest, BatchMemoryFilter};
1750 let filter = BatchMemoryFilter {
1751 created_before: Some(1_700_000_000),
1752 ..Default::default()
1753 };
1754 let req = BatchForgetRequest::new("agent-1", filter);
1755 let json = serde_json::to_value(&req).unwrap();
1756 assert_eq!(json["agent_id"], "agent-1");
1757 assert_eq!(json["filter"]["created_before"], 1_700_000_000u64);
1758 }
1759
1760 #[test]
1761 fn test_batch_recall_response_deserialization() {
1762 use crate::memory::BatchRecallResponse;
1763 let json = serde_json::json!({
1764 "memories": [],
1765 "total": 42,
1766 "filtered": 7
1767 });
1768 let resp: BatchRecallResponse = serde_json::from_value(json).unwrap();
1769 assert_eq!(resp.total, 42);
1770 assert_eq!(resp.filtered, 7);
1771 assert!(resp.memories.is_empty());
1772 }
1773
1774 #[test]
1775 fn test_batch_forget_response_deserialization() {
1776 use crate::memory::BatchForgetResponse;
1777 let json = serde_json::json!({ "deleted_count": 13 });
1778 let resp: BatchForgetResponse = serde_json::from_value(json).unwrap();
1779 assert_eq!(resp.deleted_count, 13);
1780 }
1781
1782 #[test]
1787 fn test_rate_limit_headers_default_all_none() {
1788 use crate::types::RateLimitHeaders;
1789 let rl = RateLimitHeaders {
1790 limit: None,
1791 remaining: None,
1792 reset: None,
1793 quota_used: None,
1794 quota_limit: None,
1795 };
1796 assert!(rl.limit.is_none());
1797 assert!(rl.remaining.is_none());
1798 assert!(rl.reset.is_none());
1799 assert!(rl.quota_used.is_none());
1800 assert!(rl.quota_limit.is_none());
1801 }
1802
1803 #[test]
1804 fn test_rate_limit_headers_populated() {
1805 use crate::types::RateLimitHeaders;
1806 let rl = RateLimitHeaders {
1807 limit: Some(1000),
1808 remaining: Some(750),
1809 reset: Some(1_700_000_060),
1810 quota_used: Some(500),
1811 quota_limit: Some(10_000),
1812 };
1813 assert_eq!(rl.limit, Some(1000));
1814 assert_eq!(rl.remaining, Some(750));
1815 assert_eq!(rl.reset, Some(1_700_000_060));
1816 assert_eq!(rl.quota_used, Some(500));
1817 assert_eq!(rl.quota_limit, Some(10_000));
1818 }
1819
1820 #[test]
1821 fn test_last_rate_limit_headers_initially_none() {
1822 let client = DakeraClient::new("http://localhost:3000").unwrap();
1823 assert!(client.last_rate_limit_headers().is_none());
1824 }
1825
1826 #[test]
1831 fn test_namespace_ner_config_default() {
1832 use crate::types::NamespaceNerConfig;
1833 let cfg = NamespaceNerConfig::default();
1834 assert!(!cfg.extract_entities);
1835 assert!(cfg.entity_types.is_none());
1836 }
1837
1838 #[test]
1839 fn test_namespace_ner_config_serialization_skip_none() {
1840 use crate::types::NamespaceNerConfig;
1841 let cfg = NamespaceNerConfig {
1842 extract_entities: true,
1843 entity_types: None,
1844 };
1845 let json = serde_json::to_value(&cfg).unwrap();
1846 assert_eq!(json["extract_entities"], true);
1847 assert!(json.get("entity_types").is_none());
1849 }
1850
1851 #[test]
1852 fn test_namespace_ner_config_serialization_with_types() {
1853 use crate::types::NamespaceNerConfig;
1854 let cfg = NamespaceNerConfig {
1855 extract_entities: true,
1856 entity_types: Some(vec!["PERSON".to_string(), "ORG".to_string()]),
1857 };
1858 let json = serde_json::to_value(&cfg).unwrap();
1859 assert_eq!(json["extract_entities"], true);
1860 assert_eq!(json["entity_types"][0], "PERSON");
1861 assert_eq!(json["entity_types"][1], "ORG");
1862 }
1863
1864 #[test]
1865 fn test_extracted_entity_deserialization() {
1866 use crate::types::ExtractedEntity;
1867 let json = serde_json::json!({
1868 "entity_type": "PERSON",
1869 "value": "Alice",
1870 "score": 0.95
1871 });
1872 let entity: ExtractedEntity = serde_json::from_value(json).unwrap();
1873 assert_eq!(entity.entity_type, "PERSON");
1874 assert_eq!(entity.value, "Alice");
1875 assert!((entity.score - 0.95).abs() < f64::EPSILON);
1876 }
1877
1878 #[test]
1879 fn test_entity_extraction_response_deserialization() {
1880 use crate::types::EntityExtractionResponse;
1881 let json = serde_json::json!({
1882 "entities": [
1883 { "entity_type": "PERSON", "value": "Bob", "score": 0.9 },
1884 { "entity_type": "ORG", "value": "Acme", "score": 0.87 }
1885 ]
1886 });
1887 let resp: EntityExtractionResponse = serde_json::from_value(json).unwrap();
1888 assert_eq!(resp.entities.len(), 2);
1889 assert_eq!(resp.entities[0].entity_type, "PERSON");
1890 assert_eq!(resp.entities[1].value, "Acme");
1891 }
1892
1893 #[test]
1894 fn test_memory_entities_response_deserialization() {
1895 use crate::types::MemoryEntitiesResponse;
1896 let json = serde_json::json!({
1897 "memory_id": "mem-abc-123",
1898 "entities": [
1899 { "entity_type": "LOC", "value": "London", "score": 0.88 }
1900 ]
1901 });
1902 let resp: MemoryEntitiesResponse = serde_json::from_value(json).unwrap();
1903 assert_eq!(resp.memory_id, "mem-abc-123");
1904 assert_eq!(resp.entities.len(), 1);
1905 assert_eq!(resp.entities[0].entity_type, "LOC");
1906 assert_eq!(resp.entities[0].value, "London");
1907 }
1908
1909 #[test]
1910 fn test_configure_namespace_ner_url_pattern() {
1911 let client = DakeraClient::new("http://localhost:3000").unwrap();
1913 let expected = "http://localhost:3000/v1/namespaces/my-ns/config";
1914 let actual = format!("{}/v1/namespaces/{}/config", client.base_url, "my-ns");
1915 assert_eq!(actual, expected);
1916 }
1917
1918 #[test]
1919 fn test_extract_entities_url_pattern() {
1920 let client = DakeraClient::new("http://localhost:3000").unwrap();
1921 let expected = "http://localhost:3000/v1/memories/extract";
1922 let actual = format!("{}/v1/memories/extract", client.base_url);
1923 assert_eq!(actual, expected);
1924 }
1925
1926 #[test]
1927 fn test_memory_entities_url_pattern() {
1928 let client = DakeraClient::new("http://localhost:3000").unwrap();
1929 let memory_id = "mem-xyz-789";
1930 let expected = "http://localhost:3000/v1/memory/entities/mem-xyz-789";
1931 let actual = format!("{}/v1/memory/entities/{}", client.base_url, memory_id);
1932 assert_eq!(actual, expected);
1933 }
1934
1935 #[test]
1940 fn test_feedback_signal_serialization() {
1941 use crate::types::FeedbackSignal;
1942 let upvote = serde_json::to_value(FeedbackSignal::Upvote).unwrap();
1943 assert_eq!(upvote, serde_json::json!("upvote"));
1944 let downvote = serde_json::to_value(FeedbackSignal::Downvote).unwrap();
1945 assert_eq!(downvote, serde_json::json!("downvote"));
1946 let flag = serde_json::to_value(FeedbackSignal::Flag).unwrap();
1947 assert_eq!(flag, serde_json::json!("flag"));
1948 }
1949
1950 #[test]
1951 fn test_feedback_signal_deserialization() {
1952 use crate::types::FeedbackSignal;
1953 let signal: FeedbackSignal = serde_json::from_str("\"upvote\"").unwrap();
1954 assert_eq!(signal, FeedbackSignal::Upvote);
1955 let signal: FeedbackSignal = serde_json::from_str("\"positive\"").unwrap();
1956 assert_eq!(signal, FeedbackSignal::Positive);
1957 }
1958
1959 #[test]
1960 fn test_feedback_response_deserialization() {
1961 use crate::types::{FeedbackResponse, FeedbackSignal};
1962 let json = serde_json::json!({
1963 "memory_id": "mem-abc",
1964 "new_importance": 0.92,
1965 "signal": "upvote"
1966 });
1967 let resp: FeedbackResponse = serde_json::from_value(json).unwrap();
1968 assert_eq!(resp.memory_id, "mem-abc");
1969 assert!((resp.new_importance - 0.92).abs() < f32::EPSILON);
1970 assert_eq!(resp.signal, FeedbackSignal::Upvote);
1971 }
1972
1973 #[test]
1974 fn test_feedback_history_response_deserialization() {
1975 use crate::types::{FeedbackHistoryResponse, FeedbackSignal};
1976 let json = serde_json::json!({
1977 "memory_id": "mem-abc",
1978 "entries": [
1979 {"signal": "upvote", "timestamp": 1774000000_u64, "old_importance": 0.5, "new_importance": 0.575},
1980 {"signal": "downvote", "timestamp": 1774001000_u64, "old_importance": 0.575, "new_importance": 0.489}
1981 ]
1982 });
1983 let resp: FeedbackHistoryResponse = serde_json::from_value(json).unwrap();
1984 assert_eq!(resp.memory_id, "mem-abc");
1985 assert_eq!(resp.entries.len(), 2);
1986 assert_eq!(resp.entries[0].signal, FeedbackSignal::Upvote);
1987 assert_eq!(resp.entries[1].signal, FeedbackSignal::Downvote);
1988 }
1989
1990 #[test]
1991 fn test_agent_feedback_summary_deserialization() {
1992 use crate::types::AgentFeedbackSummary;
1993 let json = serde_json::json!({
1994 "agent_id": "agent-1",
1995 "upvotes": 42_u64,
1996 "downvotes": 7_u64,
1997 "flags": 2_u64,
1998 "total_feedback": 51_u64,
1999 "health_score": 0.78
2000 });
2001 let summary: AgentFeedbackSummary = serde_json::from_value(json).unwrap();
2002 assert_eq!(summary.agent_id, "agent-1");
2003 assert_eq!(summary.upvotes, 42);
2004 assert_eq!(summary.total_feedback, 51);
2005 assert!((summary.health_score - 0.78).abs() < f32::EPSILON);
2006 }
2007
2008 #[test]
2009 fn test_feedback_health_response_deserialization() {
2010 use crate::types::FeedbackHealthResponse;
2011 let json = serde_json::json!({
2012 "agent_id": "agent-1",
2013 "health_score": 0.78,
2014 "memory_count": 120_usize,
2015 "avg_importance": 0.72
2016 });
2017 let health: FeedbackHealthResponse = serde_json::from_value(json).unwrap();
2018 assert_eq!(health.agent_id, "agent-1");
2019 assert!((health.health_score - 0.78).abs() < f32::EPSILON);
2020 assert_eq!(health.memory_count, 120);
2021 }
2022
2023 #[test]
2024 fn test_memory_feedback_body_serialization() {
2025 use crate::types::{FeedbackSignal, MemoryFeedbackBody};
2026 let body = MemoryFeedbackBody {
2027 agent_id: "agent-1".to_string(),
2028 signal: FeedbackSignal::Flag,
2029 };
2030 let json = serde_json::to_value(body).unwrap();
2031 assert_eq!(json["agent_id"], "agent-1");
2032 assert_eq!(json["signal"], "flag");
2033 }
2034
2035 #[test]
2036 fn test_feedback_memory_url_pattern() {
2037 let client = DakeraClient::new("http://localhost:3000").unwrap();
2038 let memory_id = "mem-abc";
2039 let expected_post = "http://localhost:3000/v1/memories/mem-abc/feedback";
2040 let actual_post = format!("{}/v1/memories/{}/feedback", client.base_url, memory_id);
2041 assert_eq!(actual_post, expected_post);
2042
2043 let expected_patch = "http://localhost:3000/v1/memories/mem-abc/importance";
2044 let actual_patch = format!("{}/v1/memories/{}/importance", client.base_url, memory_id);
2045 assert_eq!(actual_patch, expected_patch);
2046 }
2047
2048 #[test]
2049 fn test_feedback_health_url_pattern() {
2050 let client = DakeraClient::new("http://localhost:3000").unwrap();
2051 let agent_id = "agent-1";
2052 let expected = "http://localhost:3000/v1/feedback/health?agent_id=agent-1";
2053 let actual = format!(
2054 "{}/v1/feedback/health?agent_id={}",
2055 client.base_url, agent_id
2056 );
2057 assert_eq!(actual, expected);
2058 }
2059}