From 660cb78d486c4646593209c0dbdcd198638f8cf5 Mon Sep 17 00:00:00 2001 From: Sopaco Date: Tue, 14 Apr 2026 19:03:04 +0800 Subject: [PATCH] Refactor vector operations to use gRPC instead of REST Replace direct Qdrant REST API calls with gRPC methods in the CLI vector commands (reindex, prune, status). Add new admin methods to QdrantVectorStore for collection management, point scrolling, counting, and deletion. This provides cleaner abstraction and consistent use of the Qdrant gRPC client throughout the codebase. --- Cargo.lock | 105 +++------- Cargo.toml | 2 +- cortex-mem-cli/src/commands/vector.rs | 210 ++++---------------- cortex-mem-core/src/vector_store/qdrant.rs | 220 ++++++++++++++++++++- cortex-mem-tools/src/operations.rs | 5 + 5 files changed, 287 insertions(+), 255 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 9e83791..9f41679 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -600,16 +600,16 @@ dependencies = [ [[package]] name = "cortex-mem-cli" -version = "2.7.0" +version = "2.7.1" dependencies = [ "anyhow", "assert_cmd", "chrono", "clap", "colored", - "cortex-mem-config 2.7.0", - "cortex-mem-core 2.7.0", - "cortex-mem-tools 2.7.0", + "cortex-mem-config", + "cortex-mem-core", + "cortex-mem-tools", "predicates", "reqwest 0.12.24", "serde", @@ -622,20 +622,7 @@ dependencies = [ [[package]] name = "cortex-mem-config" -version = "2.7.0" -dependencies = [ - "anyhow", - "directories 5.0.1", - "serde", - "thiserror 2.0.17", - "toml", -] - -[[package]] -name = "cortex-mem-config" -version = "2.7.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "fb96c8de89ec4ee0d4b634ce86ebed4c14bc990b604a3cf86eaa0fc64d49298d" +version = "2.7.1" dependencies = [ "anyhow", "directories 5.0.1", @@ -646,12 +633,12 @@ dependencies = [ [[package]] name = "cortex-mem-core" -version = "2.7.0" +version = "2.7.1" dependencies = [ "anyhow", "async-trait", "chrono", - "cortex-mem-config 2.7.0", + "cortex-mem-config", "dyn-clone", "futures", "qdrant-client", @@ -672,45 +659,17 @@ dependencies = [ "walkdir", ] -[[package]] -name = "cortex-mem-core" -version = "2.7.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ccc9c7a65901b61be1b2b73bc1d0e588dc8a40aff974df44b79beb4bcf94644c" -dependencies = [ - "anyhow", - "async-trait", - "chrono", - "cortex-mem-config 2.7.0 (registry+https://github.com/rust-lang/crates.io-index)", - "dyn-clone", - "futures", - "qdrant-client", - "regex", - "reqwest 0.12.24", - "rig-core", - "schemars 0.8.22", - "serde", - "serde_json", - "sha2", - "thiserror 2.0.17", - "tokio", - "tracing", - "tracing-subscriber", - "uuid", - "walkdir", -] - [[package]] name = "cortex-mem-mcp" -version = "2.7.0" +version = "2.7.1" dependencies = [ "anyhow", "async-trait", "chrono", "clap", - "cortex-mem-config 2.7.0 (registry+https://github.com/rust-lang/crates.io-index)", - "cortex-mem-core 2.7.0 (registry+https://github.com/rust-lang/crates.io-index)", - "cortex-mem-tools 2.7.0 (registry+https://github.com/rust-lang/crates.io-index)", + "cortex-mem-config", + "cortex-mem-core", + "cortex-mem-tools", "dirs", "rmcp", "schemars 1.2.1", @@ -725,12 +684,12 @@ dependencies = [ [[package]] name = "cortex-mem-rig" -version = "2.7.0" +version = "2.7.1" dependencies = [ "anyhow", "async-trait", - "cortex-mem-core 2.7.0", - "cortex-mem-tools 2.7.0", + "cortex-mem-core", + "cortex-mem-tools", "rig-core", "serde", "serde_json", @@ -741,14 +700,14 @@ dependencies = [ [[package]] name = "cortex-mem-service" -version = "2.7.0" +version = "2.7.1" dependencies = [ "anyhow", "axum", "chrono", "clap", - "cortex-mem-config 2.7.0", - "cortex-mem-core 2.7.0", + "cortex-mem-config", + "cortex-mem-core", "futures", "http-body-util", "serde", @@ -764,7 +723,7 @@ dependencies = [ [[package]] name = "cortex-mem-tars" -version = "2.7.0" +version = "2.7.1" dependencies = [ "anyhow", "async-stream", @@ -772,10 +731,10 @@ dependencies = [ "chrono", "clap", "clipboard", - "cortex-mem-config 2.7.0", - "cortex-mem-core 2.7.0", + "cortex-mem-config", + "cortex-mem-core", "cortex-mem-rig", - "cortex-mem-tools 2.7.0", + "cortex-mem-tools", "cpal", "crossterm", "directories 6.0.0", @@ -804,12 +763,12 @@ dependencies = [ [[package]] name = "cortex-mem-tools" -version = "2.7.0" +version = "2.7.1" dependencies = [ "anyhow", "async-trait", "chrono", - "cortex-mem-core 2.7.0", + "cortex-mem-core", "futures", "serde", "serde_json", @@ -820,24 +779,6 @@ dependencies = [ "uuid", ] -[[package]] -name = "cortex-mem-tools" -version = "2.7.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "545813b8e08ac2f9f25c209d8d2a65b559b9e519d829d655378d401f6da1bf62" -dependencies = [ - "anyhow", - "async-trait", - "chrono", - "cortex-mem-core 2.7.0 (registry+https://github.com/rust-lang/crates.io-index)", - "serde", - "serde_json", - "thiserror 2.0.17", - "tokio", - "tracing", - "uuid", -] - [[package]] name = "cpal" version = "0.17.3" diff --git a/Cargo.toml b/Cargo.toml index 252c9fd..76b3128 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -12,7 +12,7 @@ members = [ ] [workspace.package] -version = "2.7.0" +version = "2.7.1" edition = "2024" rust-version = "1.86" authors = ["Sopaco"] diff --git a/cortex-mem-cli/src/commands/vector.rs b/cortex-mem-cli/src/commands/vector.rs index bf80cfc..0332e41 100644 --- a/cortex-mem-cli/src/commands/vector.rs +++ b/cortex-mem-cli/src/commands/vector.rs @@ -3,16 +3,15 @@ use colored::Colorize; use cortex_mem_tools::MemoryOperations; use std::sync::Arc; -fn qdrant_url() -> String { - std::env::var("QDRANT_URL").unwrap_or_else(|_| "http://localhost:6334".to_string()) -} - /// Reindex: clean up no-URI vectors, then full sync pub async fn reindex(operations: Arc) -> Result<()> { println!("{} Starting vector reindex...\n", "๐Ÿ”„".bold()); - // Step 1: delete stale (no-URI) vectors - match clean_no_uri_vectors().await { + // Step 1: delete stale (no-URI) vectors using gRPC + let vector_store = operations.vector_store(); + let collection_name = vector_store.collection_name().to_string(); + + match vector_store.delete_no_uri_points(&collection_name).await { Ok(n) => println!(" {} Removed {} stale vectors (no URI metadata)", "โœ…".green(), n), Err(e) => println!(" {} Failed to clean stale vectors: {} (continuing...)", "โš ๏ธ".yellow(), e), } @@ -45,67 +44,32 @@ pub async fn prune(operations: Arc, dry_run: bool) -> Result<( println!("{} Scanning for dangling vectors (files deleted from disk)...\n", "๐Ÿงน".bold()); } - let url = qdrant_url(); - let client = reqwest::Client::new(); - - let resp: serde_json::Value = client.get(format!("{}/collections", url)) - .send().await?.json().await?; - let collections: Vec = resp["result"]["collections"] - .as_array().unwrap_or(&vec![]) - .iter() - .filter_map(|c| c["name"].as_str().map(|s| s.to_string())) - .collect(); + let vector_store = operations.vector_store(); + let collection_name = vector_store.collection_name().to_string(); + + // Get all points with URI using gRPC scroll + let points = vector_store + .scroll_points_with_uri(&collection_name, 200) + .await?; let mut total_checked = 0u64; - let mut dangling_ids: Vec<(String, serde_json::Value)> = vec![]; // (collection, point_id) - - for coll in &collections { - let mut offset: Option = None; - - loop { - let mut body = serde_json::json!({ - "limit": 200, - "with_payload": ["uri"], - "with_vector": false, - "filter": { - "must_not": [{"is_empty": {"key": "uri"}}] - } - }); - if let Some(ref off) = offset { - body["offset"] = off.clone(); - } - - let scroll: serde_json::Value = client - .post(format!("{}/collections/{}/points/scroll", url, coll)) - .json(&body) - .send().await?.json().await?; - - let points = match scroll["result"]["points"].as_array() { - Some(p) if !p.is_empty() => p.clone(), - _ => break, - }; - - for pt in &points { - total_checked += 1; - let uri = match pt["payload"]["uri"].as_str() { - Some(u) if !u.is_empty() => u, - _ => continue, - }; - - // Check if the file still exists in cortex filesystem - let exists = operations.exists(uri).await.unwrap_or(true); // assume exists on error - if !exists { - if dry_run { - println!(" {} would delete: {}", "โ†’".dimmed(), uri); - } - dangling_ids.push((coll.clone(), pt["id"].clone())); - } - } - - offset = scroll["result"]["next_page_offset"].clone().into(); - if offset.as_ref().map(|v| v.is_null()).unwrap_or(true) { - break; + let mut dangling_ids: Vec = vec![]; + + for (point_id, uri_opt) in &points { + total_checked += 1; + + let uri = match uri_opt { + Some(u) if !u.is_empty() => u, + _ => continue, + }; + + // Check if the file still exists in cortex filesystem + let exists = operations.exists(uri).await.unwrap_or(true); // assume exists on error + if !exists { + if dry_run { + println!(" {} would delete: {}", "โ†’".dimmed(), uri); } + dangling_ids.push(point_id.clone()); } } @@ -125,25 +89,10 @@ pub async fn prune(operations: Arc, dry_run: bool) -> Result<( return Ok(()); } - // Group by collection and batch-delete - let mut by_coll: std::collections::HashMap> = - std::collections::HashMap::new(); - for (coll, id) in dangling_ids { - by_coll.entry(coll).or_default().push(id); - } - - let mut total_deleted = 0usize; - for (coll, ids) in &by_coll { - let del: serde_json::Value = client - .post(format!("{}/collections/{}/points/delete", url, coll)) - .json(&serde_json::json!({"points": ids})) - .send().await?.json().await?; - if del["status"].as_str() == Some("ok") { - total_deleted += ids.len(); - } - } + // Delete dangling points using gRPC + vector_store.delete_points_by_ids(&collection_name, &dangling_ids).await?; - println!("\n{} Pruned {} dangling vectors.", "โœ…".green(), total_deleted); + println!("\n{} Pruned {} dangling vectors.", "โœ…".green(), dangling_ids.len()); Ok(()) } @@ -160,7 +109,7 @@ pub async fn status(operations: Arc) -> Result<()> { } println!(" Total tracked files: ~{}", total_files); - match fetch_collection_stats().await { + match fetch_collection_stats(operations).await { Ok((total_pts, no_uri_pts)) => { println!(" Vectors in Qdrant: {}", total_pts); if no_uri_pts > 0 { @@ -183,90 +132,15 @@ pub async fn status(operations: Arc) -> Result<()> { // โ”€โ”€ Helpers โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ -async fn fetch_collection_stats() -> Result<(u64, u64)> { - let url = qdrant_url(); - let client = reqwest::Client::new(); - - let resp: serde_json::Value = client.get(format!("{}/collections", url)) - .send().await?.json().await?; - - let collections: Vec = resp["result"]["collections"] - .as_array().unwrap_or(&vec![]) - .iter() - .filter_map(|c| c["name"].as_str().map(|s| s.to_string())) - .collect(); - - if collections.is_empty() { - anyhow::bail!("No Qdrant collections found"); - } - - let mut total_pts = 0u64; - let mut no_uri_pts = 0u64; - - for coll in &collections { - let info: serde_json::Value = client - .get(format!("{}/collections/{}", url, coll)) - .send().await?.json().await?; - total_pts += info["result"]["points_count"].as_u64().unwrap_or(0); - - let count_resp: serde_json::Value = client - .post(format!("{}/collections/{}/points/count", url, coll)) - .json(&serde_json::json!({ - "filter": { - "must": [{"is_empty": {"key": "uri"}}] - } - })) - .send().await?.json().await?; - - no_uri_pts += count_resp["result"]["count"].as_u64().unwrap_or(0); - } +async fn fetch_collection_stats(operations: Arc) -> Result<(u64, u64)> { + let vector_store = operations.vector_store(); + let collection_name = vector_store.collection_name().to_string(); + + // Get total points count + let total_pts = vector_store.get_collection_points_count(&collection_name).await?; + + // Get no-URI points count + let no_uri_pts = vector_store.count_no_uri_points(&collection_name).await?; Ok((total_pts, no_uri_pts)) -} - -async fn clean_no_uri_vectors() -> Result { - let url = qdrant_url(); - let client = reqwest::Client::new(); - - let resp: serde_json::Value = client.get(format!("{}/collections", url)) - .send().await?.json().await?; - - let collections: Vec = resp["result"]["collections"] - .as_array().unwrap_or(&vec![]) - .iter() - .filter_map(|c| c["name"].as_str().map(|s| s.to_string())) - .collect(); - - let mut total_deleted = 0usize; - - for coll in &collections { - let count_resp: serde_json::Value = client - .post(format!("{}/collections/{}/points/count", url, coll)) - .json(&serde_json::json!({ - "filter": { - "must": [{"is_empty": {"key": "uri"}}] - } - })) - .send().await?.json().await?; - - let count = count_resp["result"]["count"].as_u64().unwrap_or(0); - if count == 0 { - continue; - } - - let del_resp: serde_json::Value = client - .post(format!("{}/collections/{}/points/delete", url, coll)) - .json(&serde_json::json!({ - "filter": { - "must": [{"is_empty": {"key": "uri"}}] - } - })) - .send().await?.json().await?; - - if del_resp["status"].as_str() == Some("ok") { - total_deleted += count as usize; - } - } - - Ok(total_deleted) -} +} \ No newline at end of file diff --git a/cortex-mem-core/src/vector_store/qdrant.rs b/cortex-mem-core/src/vector_store/qdrant.rs index 74c8b07..6527dd2 100644 --- a/cortex-mem-core/src/vector_store/qdrant.rs +++ b/cortex-mem-core/src/vector_store/qdrant.rs @@ -2,10 +2,11 @@ use async_trait::async_trait; use qdrant_client::{ Qdrant, qdrant::{ - Condition, CreateCollection, DeletePoints, Distance, FieldCondition, Filter, GetPoints, - Match, PointId, PointStruct, PointsIdsList, PointsSelector, Range, ScoredPoint, - ScrollPoints, SearchPoints, UpsertPoints, VectorParams, VectorsConfig, condition, r#match, - point_id, points_selector, vector_output, vectors_config, vectors_output, + Condition, CountPoints, CreateCollection, DeletePoints, Distance, FieldCondition, Filter, + GetPoints, IsEmptyCondition, Match, PointId, PointStruct, PointsIdsList, PointsSelector, + Range, ScoredPoint, ScrollPoints, SearchPoints, UpsertPoints, VectorParams, VectorsConfig, + condition, point_id, points_selector, r#match, vector_output, vectors_config, + vectors_output, }, }; use std::collections::HashMap; @@ -834,6 +835,217 @@ impl QdrantVectorStore { pub fn set_embedding_dim(&mut self, dim: usize) { self.embedding_dim = Some(dim); } + + // โ”€โ”€ Admin/Management Methods โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ + + /// List all collection names in Qdrant + pub async fn list_collections(&self) -> Result> { + let response = self + .client + .list_collections() + .await + .map_err(|e| Error::VectorStore(e))?; + + let names: Vec = response.collections.into_iter().map(|c| c.name).collect(); + Ok(names) + } + + /// Get collection info (points count, etc.) + pub async fn get_collection_points_count(&self, collection_name: &str) -> Result { + let response = self + .client + .collection_info(collection_name) + .await + .map_err(|e| Error::VectorStore(e))?; + + let count = response + .result + .and_then(|info| info.points_count) + .unwrap_or(0); + + Ok(count) + } + + /// Scroll through all points with URI payload in a collection + /// + /// Returns a stream of (point_id, uri) pairs where uri is the file URI + pub async fn scroll_points_with_uri( + &self, + collection_name: &str, + batch_size: u32, + ) -> Result)>> { + let mut results = Vec::new(); + let mut offset: Option = None; + + loop { + let scroll_points = ScrollPoints { + collection_name: collection_name.to_string(), + limit: Some(batch_size), + offset: offset.clone(), + with_payload: Some(true.into()), + with_vectors: Some(false.into()), + filter: None, + ..Default::default() + }; + + let response = self + .client + .scroll(scroll_points) + .await + .map_err(|e| Error::VectorStore(e))?; + + if response.result.is_empty() { + break; + } + + for point in &response.result { + let point_id = match &point.id { + Some(PointId { + point_id_options: Some(pid), + }) => match pid { + point_id::PointIdOptions::Uuid(uuid) => uuid.clone(), + point_id::PointIdOptions::Num(num) => num.to_string(), + }, + _ => continue, + }; + + let uri = point.payload.get("uri").and_then(|v| match v { + qdrant_client::qdrant::Value { + kind: Some(qdrant_client::qdrant::value::Kind::StringValue(s)), + } => Some(s.clone()), + _ => None, + }); + + results.push((point_id, uri)); + } + + // Get next page offset + offset = response.next_page_offset; + if offset.is_none() { + break; + } + } + + Ok(results) + } + + /// Count points in collection, optionally with a filter + pub async fn count_points(&self, collection_name: &str, filter: Option) -> Result { + let count_request = CountPoints { + collection_name: collection_name.to_string(), + filter, + exact: Some(true), + ..Default::default() + }; + + let response = self + .client + .count(count_request) + .await + .map_err(|e| Error::VectorStore(e))?; + + Ok(response.result.map(|r| r.count).unwrap_or(0)) + } + + /// Count points with empty URI (stale vectors) + pub async fn count_no_uri_points(&self, collection_name: &str) -> Result { + let filter = Filter { + must: vec![Condition { + condition_one_of: Some(condition::ConditionOneOf::IsEmpty(IsEmptyCondition { + key: "uri".to_string(), + })), + }], + ..Default::default() + }; + + self.count_points(collection_name, Some(filter)).await + } + + /// Delete points by filter (e.g., points without URI) + pub async fn delete_points_by_filter( + &self, + collection_name: &str, + filter: Filter, + ) -> Result<()> { + let points_selector = PointsSelector { + points_selector_one_of: Some(points_selector::PointsSelectorOneOf::Filter(filter)), + }; + + let delete_request = DeletePoints { + collection_name: collection_name.to_string(), + points: Some(points_selector), + ..Default::default() + }; + + self.client + .delete_points(delete_request) + .await + .map_err(|e| Error::VectorStore(e))?; + + debug!("Deleted points by filter from collection: {}", collection_name); + Ok(()) + } + + /// Delete points without URI (stale vectors) + pub async fn delete_no_uri_points(&self, collection_name: &str) -> Result { + // First count them + let count = self.count_no_uri_points(collection_name).await?; + + if count == 0 { + return Ok(0); + } + + let filter = Filter { + must: vec![Condition { + condition_one_of: Some(condition::ConditionOneOf::IsEmpty(IsEmptyCondition { + key: "uri".to_string(), + })), + }], + ..Default::default() + }; + + self.delete_points_by_filter(collection_name, filter).await?; + Ok(count) + } + + /// Delete specific points by IDs + pub async fn delete_points_by_ids(&self, collection_name: &str, ids: &[String]) -> Result<()> { + if ids.is_empty() { + return Ok(()); + } + + let point_ids: Vec = ids + .iter() + .map(|id| PointId { + point_id_options: Some(point_id::PointIdOptions::Uuid(id.clone())), + }) + .collect(); + + let points_selector = PointsSelector { + points_selector_one_of: Some(points_selector::PointsSelectorOneOf::Points( + PointsIdsList { ids: point_ids }, + )), + }; + + let delete_request = DeletePoints { + collection_name: collection_name.to_string(), + points: Some(points_selector), + ..Default::default() + }; + + self.client + .delete_points(delete_request) + .await + .map_err(|e| Error::VectorStore(e))?; + + debug!("Deleted {} points from collection: {}", ids.len(), collection_name); + Ok(()) + } + + /// Get the collection name (for admin operations that need to know the collection) + pub fn collection_name(&self) -> &str { + &self.collection_name + } } #[async_trait] diff --git a/cortex-mem-tools/src/operations.rs b/cortex-mem-tools/src/operations.rs index 9188932..fa6a6a6 100644 --- a/cortex-mem-tools/src/operations.rs +++ b/cortex-mem-tools/src/operations.rs @@ -94,6 +94,11 @@ impl MemoryOperations { self.memory_event_tx.as_ref() } + /// Get the vector store (for admin operations like prune, reindex) + pub fn vector_store(&self) -> &Arc { + &self.vector_store + } + /// Create from data directory with tenant isolation, LLM support, and vector search /// /// This is the primary constructor that requires all dependencies.