From b71ca2ea005e32bd7ef4e1a747a91fb7fd3452f6 Mon Sep 17 00:00:00 2001 From: brianreicher Date: Wed, 11 Oct 2023 10:42:58 -0400 Subject: [PATCH] Pineconde insertion sample --- tokenizer/src/pinecone_utils.rs | 90 ++++----------------------------- 1 file changed, 11 insertions(+), 79 deletions(-) diff --git a/tokenizer/src/pinecone_utils.rs b/tokenizer/src/pinecone_utils.rs index 90168fe..0b50689 100644 --- a/tokenizer/src/pinecone_utils.rs +++ b/tokenizer/src/pinecone_utils.rs @@ -16,85 +16,17 @@ impl PinceoneDriver { PinceoneDriver { client, index } } - pub async fn flush_collection(&self, collection_name: &str) -> Result<(), pinenut::error::Error> { - let collection = self.client.collection(collection_name); - let result = collection.delete_many(Document::new()).await?; - println!("Deleted {} documents from {}", result.deleted_count, collection_name); - Ok(()) - } - - pub async fn remove_collection(&self, collection_name: &str) -> Result<(), pinenut::error::Error> { - let collection = self.client.database(&self.db_name).collection(collection_name); - collection.drop(None).await?; - println!("Dropped {} from {}", collection_name, self.db_name); - Ok(()) - } - - pub async fn create_collection(&self, collection_name: &str) -> Result<(), pinenut::error::Error> { - self.client.create_collection(collection_name, self.index).await?; - println!("Created collection {} in {}", collection_name, self.index); - Ok(()) - } - - pub async fn collection_size(&self, collection_name: &str) -> Result { - let collection = self.client.database(&self.db_name).collection(collection_name); - let count = collection.count_documents(Document::new(), None).await?; - Ok(count) - } - - pub async fn insert_data(&self, collection_name: &str, json_file: &str, clear: bool) -> Result<(), pinenut::error::Error> { - let collection = self.client.database(&self.db_name).collection(collection_name); - - if clear { - self.flush_collection(collection_name).await?; + pub async fn insert_data(&self,) -> Result<(), pinenut::error::Error> { + let vec: Vector = Vector{ + id: "B".to_string(), + values: vec![0.5; 32], + sparse_values: None, + metadata: None + }; + + match self.index.upsert(String::from("odle"), vec![vec]).await { + Ok(_) => assert!(true), + Err(err) => panic!("unable to upsert: {:?}", err) } - - let file_contents = std::fs::read_to_string(json_file)?; - let data: Vec = from_str(&file_contents)?; - - let mut documents = Vec::new(); - for document in data { - if let Value::Object(map) = document { - let bson_doc = Document::try_from(map)?; - documents.push(bson_doc); - } - } - - collection.insert_many(documents, None).await?; - println!("Inserted {} documents into {} in the {}", documents.len(), collection_name, self.db_name); - Ok(()) - } - - pub async fn search_query(&self, collection_name: &str, qu: Document, proj: Document, lim: i64, show: bool) -> Result, mongodb::error::Error> { - let collection = self.client.database(&self.db_name).collection(collection_name); - let cursor = collection.find(qu, None).projection(proj).limit(lim); - let mut result = Vec::new(); - - if show { - while let Some(doc) = cursor.next().await { - match doc { - Ok(document) => { - println!("{:?}", document); - result.push(document); - } - Err(e) => { - eprintln!("Error: {:?}", e); - } - } - } - } else { - while let Some(doc) = cursor.next().await { - match doc { - Ok(document) => { - result.push(document); - } - Err(e) => { - eprintln!("Error: {:?}", e); - } - } - } - } - - Ok(result) } }