convergio_knowledge/
store_pruning.rs1use 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 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 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 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}