1use std::path::Path;
3use std::sync::Arc;
4
5use arrow_array::{
6 types::Float32Type, Array, FixedSizeListArray, RecordBatch, RecordBatchIterator,
7 RecordBatchReader, StringArray,
8};
9use arrow_schema::{DataType, Field, Schema};
10use futures::StreamExt;
11use lancedb::connection::Connection;
12use lancedb::query::{ExecutableQuery, QueryBase};
13use tokio::sync::Mutex;
14use tracing::{info, warn};
15
16use crate::types::{KnowledgeEntry, SearchResult, StoreStats};
17
18pub const EMBEDDING_DIM: usize = 384;
19
20pub struct LanceVectorStore {
21 pub(crate) conn: Mutex<Connection>,
22 pub(crate) table_name: String,
23}
24
25impl LanceVectorStore {
26 pub async fn open(path: &Path) -> Result<Self, String> {
27 let conn = lancedb::connect(path.to_str().unwrap_or("knowledge.lance"))
28 .execute()
29 .await
30 .map_err(|e| format!("lance connect: {e}"))?;
31 let store = Self {
32 conn: Mutex::new(conn),
33 table_name: "knowledge".into(),
34 };
35 store.ensure_table().await?;
36 Ok(store)
37 }
38
39 fn schema() -> Arc<Schema> {
40 Arc::new(Schema::new(vec![
41 Field::new("id", DataType::Utf8, false),
42 Field::new("content", DataType::Utf8, false),
43 Field::new("source_type", DataType::Utf8, false),
44 Field::new("source_id", DataType::Utf8, false),
45 Field::new("org_id", DataType::Utf8, true),
46 Field::new("agent_id", DataType::Utf8, true),
47 Field::new("project_id", DataType::Utf8, true),
48 Field::new("visibility", DataType::Utf8, false),
49 Field::new("created_at", DataType::Utf8, false),
50 Field::new(
51 "vector",
52 DataType::FixedSizeList(
53 Arc::new(Field::new("item", DataType::Float32, true)),
54 EMBEDDING_DIM as i32,
55 ),
56 false,
57 ),
58 ]))
59 }
60
61 async fn ensure_table(&self) -> Result<(), String> {
62 let conn = self.conn.lock().await;
63 let names = conn
64 .table_names()
65 .execute()
66 .await
67 .map_err(|e| format!("table_names: {e}"))?;
68 if !names.contains(&self.table_name) {
69 Self::create_empty_table(&conn, &self.table_name).await?;
70 } else {
71 let table = conn
73 .open_table(&self.table_name)
74 .execute()
75 .await
76 .map_err(|e| format!("open: {e}"))?;
77 let has_visibility = table
78 .schema()
79 .await
80 .map(|s| s.field_with_name("visibility").is_ok())
81 .unwrap_or(false);
82 if !has_visibility {
83 info!("knowledge table missing visibility column, recreating");
84 conn.drop_table(&self.table_name, &[])
85 .await
86 .map_err(|e| format!("drop table: {e}"))?;
87 Self::create_empty_table(&conn, &self.table_name).await?;
88 }
89 }
90 Ok(())
91 }
92 async fn create_empty_table(conn: &Connection, name: &str) -> Result<(), String> {
93 let schema = Self::schema();
94 let batch = RecordBatch::new_empty(schema.clone());
95 let batches = RecordBatchIterator::new(vec![Ok(batch)], schema);
96 let reader: Box<dyn RecordBatchReader + Send> = Box::new(batches);
97 conn.create_table(name, reader)
98 .execute()
99 .await
100 .map_err(|e| format!("create table: {e}"))?;
101 info!("created knowledge lance table");
102 Ok(())
103 }
104
105 #[allow(clippy::too_many_arguments)]
106 pub async fn insert(
107 &self,
108 content: &str,
109 source_type: &str,
110 source_id: &str,
111 org_id: Option<&str>,
112 agent_id: Option<&str>,
113 project_id: Option<&str>,
114 visibility: &str,
115 embedding: &[f32],
116 ) -> Result<String, String> {
117 if embedding.len() != EMBEDDING_DIM {
118 return Err(format!(
119 "embedding dim {}, expected {EMBEDDING_DIM}",
120 embedding.len()
121 ));
122 }
123 let id = format!("ke-{}", uuid::Uuid::new_v4());
124 let now = chrono::Utc::now().to_rfc3339();
125
126 let schema = Self::schema();
127 let vectors = FixedSizeListArray::from_iter_primitive::<Float32Type, _, _>(
128 vec![Some(embedding.iter().map(|v| Some(*v)).collect::<Vec<_>>())],
129 EMBEDDING_DIM as i32,
130 );
131 let batch = RecordBatch::try_new(
132 schema.clone(),
133 vec![
134 Arc::new(StringArray::from(vec![id.as_str()])),
135 Arc::new(StringArray::from(vec![content])),
136 Arc::new(StringArray::from(vec![source_type])),
137 Arc::new(StringArray::from(vec![source_id])),
138 Arc::new(StringArray::from(vec![org_id.unwrap_or("")])),
139 Arc::new(StringArray::from(vec![agent_id.unwrap_or("")])),
140 Arc::new(StringArray::from(vec![project_id.unwrap_or("")])),
141 Arc::new(StringArray::from(vec![visibility])),
142 Arc::new(StringArray::from(vec![now.as_str()])),
143 Arc::new(vectors),
144 ],
145 )
146 .map_err(|e| format!("build batch: {e}"))?;
147
148 let conn = self.conn.lock().await;
149 let table = conn
150 .open_table(&self.table_name)
151 .execute()
152 .await
153 .map_err(|e| format!("open table: {e}"))?;
154 let batches = RecordBatchIterator::new(vec![Ok(batch)], schema);
155 let reader: Box<dyn RecordBatchReader + Send> = Box::new(batches);
156 table
157 .add(reader)
158 .execute()
159 .await
160 .map_err(|e| format!("insert: {e}"))?;
161
162 info!(id = %id, source_type, "knowledge entry stored via LanceDB");
163 Ok(id)
164 }
165
166 pub async fn search(
167 &self,
168 query_embedding: &[f32],
169 limit: usize,
170 ) -> Result<Vec<SearchResult>, String> {
171 let conn = self.conn.lock().await;
172 let table = conn
173 .open_table(&self.table_name)
174 .execute()
175 .await
176 .map_err(|e| format!("open: {e}"))?;
177
178 let count = table.count_rows(None).await.unwrap_or(0);
179 if count == 0 {
180 return Ok(vec![]);
181 }
182
183 let stream = table
184 .query()
185 .nearest_to(query_embedding)
186 .map_err(|e| format!("nearest_to: {e}"))?
187 .limit(limit)
188 .execute()
189 .await
190 .map_err(|e| format!("execute: {e}"))?;
191
192 let mut out = Vec::new();
193 let mut stream = stream;
194 while let Some(batch_result) = stream.next().await {
195 let batch = match batch_result {
196 Ok(b) => b,
197 Err(e) => {
198 warn!(error = %e, "lance stream error");
199 continue;
200 }
201 };
202 let col = |name: &str| -> Option<&StringArray> {
203 batch.column_by_name(name)?.as_any().downcast_ref()
204 };
205 let (Some(ids), Some(contents), Some(src_types), Some(src_ids)) = (
206 col("id"),
207 col("content"),
208 col("source_type"),
209 col("source_id"),
210 ) else {
211 warn!("missing columns in lance result");
212 continue;
213 };
214 let org_ids = col("org_id");
215 let agent_ids = col("agent_id");
216 let project_ids = col("project_id");
217 let visibilities = col("visibility");
218 let created = col("created_at");
219 let distances = batch
220 .column_by_name("_distance")
221 .and_then(|c| c.as_any().downcast_ref::<arrow_array::Float32Array>());
222
223 for i in 0..batch.num_rows() {
224 let dist = distances.map(|d| d.value(i)).unwrap_or(0.0);
225 let score = 1.0 / (1.0 + dist as f64);
226 out.push(SearchResult {
227 entry: KnowledgeEntry {
228 id: ids.value(i).to_string(),
229 content: contents.value(i).to_string(),
230 source_type: src_types.value(i).to_string(),
231 source_id: src_ids.value(i).to_string(),
232 org_id: org_ids.map(|a| a.value(i).to_string()),
233 agent_id: agent_ids.map(|a| a.value(i).to_string()),
234 project_id: project_ids.map(|a| a.value(i).to_string()),
235 visibility: visibilities
236 .map(|a| a.value(i).to_string())
237 .unwrap_or_else(|| "org".to_string()),
238 created_at: created.map(|a| a.value(i).to_string()).unwrap_or_default(),
239 },
240 score,
241 });
242 }
243 }
244 Ok(out)
245 }
246 pub async fn stats(&self) -> Result<StoreStats, String> {
247 let conn = self.conn.lock().await;
248 let table = conn
249 .open_table(&self.table_name)
250 .execute()
251 .await
252 .map_err(|e| format!("open: {e}"))?;
253 let count = table
254 .count_rows(None)
255 .await
256 .map_err(|e| format!("count: {e}"))?;
257 Ok(StoreStats {
258 total_entries: count as i64,
259 total_by_source: vec![],
260 embedding_dimensions: EMBEDDING_DIM,
261 })
262 }
263
264 pub async fn delete(&self, id: &str) -> Result<bool, String> {
265 let conn = self.conn.lock().await;
266 let table = conn
267 .open_table(&self.table_name)
268 .execute()
269 .await
270 .map_err(|e| format!("open: {e}"))?;
271 let id_esc = id.replace('\'', "''");
272 table
273 .delete(&format!("id = '{id_esc}'"))
274 .await
275 .map_err(|e| format!("delete: {e}"))?;
276 Ok(true)
277 }
278
279 pub async fn count_by_source(
280 &self,
281 source_type: &str,
282 source_id: &str,
283 ) -> Result<usize, String> {
284 let conn = self.conn.lock().await;
285 let table = conn
286 .open_table(&self.table_name)
287 .execute()
288 .await
289 .map_err(|e| format!("open: {e}"))?;
290 let (st, si) = (
291 source_type.replace('\'', "''"),
292 source_id.replace('\'', "''"),
293 );
294 let count = table
295 .count_rows(Some(format!("source_type = '{st}' AND source_id = '{si}'")))
296 .await
297 .map_err(|e| format!("count_by_source: {e}"))?;
298 Ok(count)
299 }
300}