Skip to main content

convergio_knowledge/
store_pruning.rs

1//! Pruning helpers for the vector store — list_all_metadata + delete_batch.
2//!
3//! Extracted from store.rs to stay under the 300-line limit.
4
5use std::collections::HashSet;
6
7use arrow_array::StringArray;
8use futures::StreamExt;
9use lancedb::query::ExecutableQuery;
10
11use crate::store::LanceVectorStore;
12use crate::types::KnowledgeEntry;
13
14impl LanceVectorStore {
15    /// Collect all source_ids in the store (used by seed dedup).
16    pub async fn all_source_ids(&self) -> Result<HashSet<String>, String> {
17        let conn = self.conn.lock().await;
18        let tbl = conn
19            .open_table(&self.table_name)
20            .execute()
21            .await
22            .map_err(|e| format!("open: {e}"))?;
23        if tbl.count_rows(None).await.unwrap_or(0) == 0 {
24            return Ok(HashSet::new());
25        }
26        let mut stream = tbl
27            .query()
28            .execute()
29            .await
30            .map_err(|e| format!("query source_ids: {e}"))?;
31        let mut ids = HashSet::new();
32        while let Some(Ok(batch)) = stream.next().await {
33            if let Some(col) = batch
34                .column_by_name("source_id")
35                .and_then(|c| c.as_any().downcast_ref::<StringArray>())
36            {
37                for i in 0..batch.num_rows() {
38                    ids.insert(col.value(i).to_string());
39                }
40            }
41        }
42        Ok(ids)
43    }
44
45    /// List all entries (lightweight metadata only — no vectors).
46    /// Used by the pruning system to scan for stale/duplicate entries.
47    pub async fn list_all_metadata(&self) -> Result<Vec<KnowledgeEntry>, String> {
48        let conn = self.conn.lock().await;
49        let tbl = conn
50            .open_table(&self.table_name)
51            .execute()
52            .await
53            .map_err(|e| format!("open: {e}"))?;
54        if tbl.count_rows(None).await.unwrap_or(0) == 0 {
55            return Ok(vec![]);
56        }
57        let mut stream = tbl
58            .query()
59            .execute()
60            .await
61            .map_err(|e| format!("query all: {e}"))?;
62
63        let mut out = Vec::new();
64        while let Some(Ok(batch)) = stream.next().await {
65            let col = |name: &str| -> Option<&StringArray> {
66                batch.column_by_name(name)?.as_any().downcast_ref()
67            };
68            let (Some(ids), Some(contents), Some(src_types), Some(src_ids)) = (
69                col("id"),
70                col("content"),
71                col("source_type"),
72                col("source_id"),
73            ) else {
74                continue;
75            };
76            let org_ids = col("org_id");
77            let agent_ids = col("agent_id");
78            let project_ids = col("project_id");
79            let visibilities = col("visibility");
80            let created = col("created_at");
81
82            for i in 0..batch.num_rows() {
83                out.push(KnowledgeEntry {
84                    id: ids.value(i).to_string(),
85                    content: contents.value(i).to_string(),
86                    source_type: src_types.value(i).to_string(),
87                    source_id: src_ids.value(i).to_string(),
88                    org_id: org_ids.map(|a| a.value(i).to_string()),
89                    agent_id: agent_ids.map(|a| a.value(i).to_string()),
90                    project_id: project_ids.map(|a| a.value(i).to_string()),
91                    visibility: visibilities
92                        .map(|a| a.value(i).to_string())
93                        .unwrap_or_else(|| "org".to_string()),
94                    created_at: created.map(|a| a.value(i).to_string()).unwrap_or_default(),
95                });
96            }
97        }
98        Ok(out)
99    }
100
101    /// Delete multiple entries by ID in a single filter expression.
102    pub async fn delete_batch(&self, ids: &[String]) -> Result<usize, String> {
103        if ids.is_empty() {
104            return Ok(0);
105        }
106        let conn = self.conn.lock().await;
107        let table = conn
108            .open_table(&self.table_name)
109            .execute()
110            .await
111            .map_err(|e| format!("open: {e}"))?;
112        let escaped: Vec<String> = ids
113            .iter()
114            .map(|id| format!("'{}'", id.replace('\'', "''")))
115            .collect();
116        let filter = format!("id IN ({})", escaped.join(", "));
117        table
118            .delete(&filter)
119            .await
120            .map_err(|e| format!("delete_batch: {e}"))?;
121        Ok(ids.len())
122    }
123}