Skip to content

Commit c048a64

Browse files
ferhimedamineDakera OpsPaperclip-Paperclip
authored
feat: add upsert_text, query_text, fetch, and create_namespace to Rust SDK (#4)
* feat: add upsert_text, query_text, batch_query_text, fetch, and create_namespace Implements the five operations that were missing from the Rust SDK, bringing it to full parity with the Python, TypeScript, and Go SDKs. New types (types.rs): - EmbeddingModel enum (minilm / bge-small / e5-small) - TextDocument, UpsertTextRequest, TextUpsertResponse - QueryTextRequest, TextQueryResponse, TextSearchResult - BatchQueryTextRequest, BatchQueryTextResponse - FetchRequest, FetchResponse - CreateNamespaceRequest New client methods (client.rs): - upsert_text() — POST /v1/namespaces/{ns}/upsert-text - query_text() — POST /v1/namespaces/{ns}/query-text - query_text_simple() — convenience wrapper - batch_query_text() — POST /v1/namespaces/{ns}/batch-query-text - fetch() — POST /v1/namespaces/{ns}/fetch - fetch_by_ids() — convenience wrapper - create_namespace() — POST /v1/namespaces/{ns} Tests added for all new type builders and request constructors. Tracked in internal issue DAK-5. Co-Authored-By: Paperclip <noreply@paperclip.ing> * fix: apply cargo fmt to resolve CI format check failure Fixes formatting issues in client.rs and types.rs that were causing the Format CI check to fail on PR#4. Co-Authored-By: Paperclip <noreply@paperclip.ing> --------- Co-authored-by: Dakera Ops <ops@dakera.ai> Co-authored-by: Paperclip <noreply@paperclip.ing>
1 parent 0052ad5 commit c048a64

2 files changed

Lines changed: 460 additions & 0 deletions

File tree

src/client.rs

Lines changed: 164 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -106,6 +106,18 @@ impl DakeraClient {
106106
self.handle_response(response).await
107107
}
108108

109+
/// Create a new namespace
110+
#[instrument(skip(self, request))]
111+
pub async fn create_namespace(
112+
&self,
113+
namespace: &str,
114+
request: CreateNamespaceRequest,
115+
) -> Result<NamespaceInfo> {
116+
let url = format!("{}/v1/namespaces/{}", self.base_url, namespace);
117+
let response = self.client.post(&url).json(&request).send().await?;
118+
self.handle_response(response).await
119+
}
120+
109121
// ========================================================================
110122
// Vector Operations
111123
// ========================================================================
@@ -790,6 +802,92 @@ impl DakeraClient {
790802
}
791803
}
792804

805+
// ========================================================================
806+
// Fetch by ID
807+
// ========================================================================
808+
809+
/// Fetch vectors by their IDs
810+
#[instrument(skip(self, request), fields(id_count = request.ids.len()))]
811+
pub async fn fetch(&self, namespace: &str, request: FetchRequest) -> Result<FetchResponse> {
812+
let url = format!("{}/v1/namespaces/{}/fetch", self.base_url, namespace);
813+
debug!("Fetching {} vectors from {}", request.ids.len(), namespace);
814+
let response = self.client.post(&url).json(&request).send().await?;
815+
self.handle_response(response).await
816+
}
817+
818+
/// Fetch vectors by IDs (convenience method)
819+
#[instrument(skip(self))]
820+
pub async fn fetch_by_ids(&self, namespace: &str, ids: &[&str]) -> Result<Vec<Vector>> {
821+
let request = FetchRequest::new(ids.iter().map(|s| s.to_string()).collect());
822+
self.fetch(namespace, request).await.map(|r| r.vectors)
823+
}
824+
825+
// ========================================================================
826+
// Text Auto-Embedding Operations
827+
// ========================================================================
828+
829+
/// Upsert text documents with automatic server-side embedding generation
830+
#[instrument(skip(self, request), fields(doc_count = request.documents.len()))]
831+
pub async fn upsert_text(
832+
&self,
833+
namespace: &str,
834+
request: UpsertTextRequest,
835+
) -> Result<TextUpsertResponse> {
836+
let url = format!("{}/v1/namespaces/{}/upsert-text", self.base_url, namespace);
837+
debug!(
838+
"Upserting {} text documents to {}",
839+
request.documents.len(),
840+
namespace
841+
);
842+
let response = self.client.post(&url).json(&request).send().await?;
843+
self.handle_response(response).await
844+
}
845+
846+
/// Query using natural language text with automatic server-side embedding
847+
#[instrument(skip(self, request), fields(top_k = request.top_k))]
848+
pub async fn query_text(
849+
&self,
850+
namespace: &str,
851+
request: QueryTextRequest,
852+
) -> Result<TextQueryResponse> {
853+
let url = format!("{}/v1/namespaces/{}/query-text", self.base_url, namespace);
854+
debug!("Text query in {} for: {}", namespace, request.text);
855+
let response = self.client.post(&url).json(&request).send().await?;
856+
self.handle_response(response).await
857+
}
858+
859+
/// Query text (convenience method)
860+
#[instrument(skip(self))]
861+
pub async fn query_text_simple(
862+
&self,
863+
namespace: &str,
864+
text: &str,
865+
top_k: u32,
866+
) -> Result<TextQueryResponse> {
867+
self.query_text(namespace, QueryTextRequest::new(text, top_k))
868+
.await
869+
}
870+
871+
/// Execute multiple text queries with automatic embedding in a single request
872+
#[instrument(skip(self, request), fields(query_count = request.queries.len()))]
873+
pub async fn batch_query_text(
874+
&self,
875+
namespace: &str,
876+
request: BatchQueryTextRequest,
877+
) -> Result<BatchQueryTextResponse> {
878+
let url = format!(
879+
"{}/v1/namespaces/{}/batch-query-text",
880+
self.base_url, namespace
881+
);
882+
debug!(
883+
"Batch text query in {} with {} queries",
884+
namespace,
885+
request.queries.len()
886+
);
887+
let response = self.client.post(&url).json(&request).send().await?;
888+
self.handle_response(response).await
889+
}
890+
793891
// ========================================================================
794892
// Private Helpers
795893
// ========================================================================
@@ -951,4 +1049,70 @@ mod tests {
9511049

9521050
assert_eq!(req.vector_weight, 1.0);
9531051
}
1052+
1053+
#[test]
1054+
fn test_text_document_builder() {
1055+
let doc = TextDocument::new("doc1", "Hello world").with_ttl(3600);
1056+
1057+
assert_eq!(doc.id, "doc1");
1058+
assert_eq!(doc.text, "Hello world");
1059+
assert_eq!(doc.ttl_seconds, Some(3600));
1060+
assert!(doc.metadata.is_none());
1061+
}
1062+
1063+
#[test]
1064+
fn test_upsert_text_request_builder() {
1065+
let docs = vec![
1066+
TextDocument::new("doc1", "Hello"),
1067+
TextDocument::new("doc2", "World"),
1068+
];
1069+
let req = UpsertTextRequest::new(docs).with_model(EmbeddingModel::BgeSmall);
1070+
1071+
assert_eq!(req.documents.len(), 2);
1072+
assert_eq!(req.model, Some(EmbeddingModel::BgeSmall));
1073+
}
1074+
1075+
#[test]
1076+
fn test_query_text_request_builder() {
1077+
let req = QueryTextRequest::new("semantic search query", 5)
1078+
.with_filter(serde_json::json!({"category": "docs"}))
1079+
.include_vectors(true)
1080+
.with_model(EmbeddingModel::E5Small);
1081+
1082+
assert_eq!(req.text, "semantic search query");
1083+
assert_eq!(req.top_k, 5);
1084+
assert!(req.filter.is_some());
1085+
assert!(req.include_vectors);
1086+
assert_eq!(req.model, Some(EmbeddingModel::E5Small));
1087+
}
1088+
1089+
#[test]
1090+
fn test_fetch_request_builder() {
1091+
let req = FetchRequest::new(vec!["id1".to_string(), "id2".to_string()]);
1092+
1093+
assert_eq!(req.ids.len(), 2);
1094+
assert!(req.include_values);
1095+
assert!(req.include_metadata);
1096+
}
1097+
1098+
#[test]
1099+
fn test_create_namespace_request_builder() {
1100+
let req = CreateNamespaceRequest::new()
1101+
.with_dimensions(384)
1102+
.with_index_type("hnsw");
1103+
1104+
assert_eq!(req.dimensions, Some(384));
1105+
assert_eq!(req.index_type.as_deref(), Some("hnsw"));
1106+
}
1107+
1108+
#[test]
1109+
fn test_batch_query_text_request() {
1110+
let req =
1111+
BatchQueryTextRequest::new(vec!["query one".to_string(), "query two".to_string()], 10);
1112+
1113+
assert_eq!(req.queries.len(), 2);
1114+
assert_eq!(req.top_k, 10);
1115+
assert!(!req.include_vectors);
1116+
assert!(req.model.is_none());
1117+
}
9541118
}

0 commit comments

Comments
 (0)