cascade_agent/knowledge/
vectordb.rs

1//! LanceDB-backed vector store for knowledge retrieval.
2
3use std::sync::Arc;
4
5use arrow_array::{
6    types::Float32Type, Array, FixedSizeListArray, Float32Array, RecordBatch, StringArray,
7    UInt64Array,
8};
9use arrow_schema::{DataType, Field, Schema};
10use futures::TryStreamExt;
11use lancedb::connect;
12use lancedb::database::CreateTableMode;
13use lancedb::query::{ExecutableQuery, QueryBase};
14use serde::{Deserialize, Serialize};
15
16use super::embeddings::Embedder;
17use crate::error::{AgentError, Result};
18
19// ---------------------------------------------------------------------------
20// Data types
21// ---------------------------------------------------------------------------
22
23/// A single entry to store in the vector database.
24#[derive(Debug, Clone, Serialize, Deserialize)]
25pub struct KnowledgeEntry {
26    pub text: String,
27    pub source: String,
28    #[serde(default)]
29    pub metadata: serde_json::Value,
30    /// Unix timestamp (seconds since epoch). Default: now.
31    #[serde(default = "default_timestamp")]
32    pub timestamp: i64,
33}
34
35fn default_timestamp() -> i64 {
36    chrono::Utc::now().timestamp()
37}
38
39/// A single result returned from a vector search.
40#[derive(Debug, Clone, Serialize, Deserialize)]
41pub struct SearchResult {
42    pub text: String,
43    pub source: String,
44    pub score: f32,
45    pub metadata: serde_json::Value,
46    pub timestamp: i64,
47}
48
49// ---------------------------------------------------------------------------
50// VectorStore
51// ---------------------------------------------------------------------------
52
53/// LanceDB-backed vector store.
54pub struct VectorStore {
55    conn: lancedb::Connection,
56    embedder: Arc<Embedder>,
57    default_collection: String,
58    dim: usize,
59}
60
61impl std::fmt::Debug for VectorStore {
62    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
63        f.debug_struct("VectorStore")
64            .field("default_collection", &self.default_collection)
65            .field("dim", &self.dim)
66            .field("embedder", &self.embedder)
67            .finish_non_exhaustive()
68    }
69}
70
71impl VectorStore {
72    /// Connect to (or create) a LanceDB database at `db_path`.
73    pub async fn new(
74        db_path: &str,
75        embedder: Arc<Embedder>,
76        default_collection: &str,
77    ) -> Result<Self> {
78        let conn = connect(db_path).execute().await.map_err(|e| {
79            AgentError::KnowledgeError(format!(
80                "Failed to connect to LanceDB at '{}': {}",
81                db_path, e
82            ))
83        })?;
84        let dim = embedder.dimension();
85        Ok(Self {
86            conn,
87            embedder,
88            default_collection: default_collection.to_owned(),
89            dim,
90        })
91    }
92
93    /// Build the Arrow schema used for knowledge tables.
94    fn schema(dim: i32) -> Arc<Schema> {
95        Arc::new(Schema::new(vec![
96            Field::new("id", DataType::Utf8, false),
97            Field::new(
98                "vector",
99                DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Float32, true)), dim),
100                true,
101            ),
102            Field::new("text", DataType::Utf8, false),
103            Field::new("source", DataType::Utf8, false),
104            Field::new("metadata_json", DataType::Utf8, true),
105            Field::new("timestamp", DataType::UInt64, false),
106        ]))
107    }
108
109    /// Create a new collection (table). Uses `exist_ok` mode.
110    pub async fn create_collection(&self, name: &str) -> Result<()> {
111        let schema = Self::schema(self.dim as i32);
112        self.conn
113            .create_empty_table(name, schema)
114            .mode(CreateTableMode::exist_ok(|req| req))
115            .execute()
116            .await
117            .map_err(|e| {
118                AgentError::KnowledgeError(format!("Failed to create collection '{}': {}", name, e))
119            })?;
120        Ok(())
121    }
122
123    /// Insert entries into a collection.
124    pub async fn insert(&self, collection: &str, entries: Vec<KnowledgeEntry>) -> Result<()> {
125        if entries.is_empty() {
126            return Ok(());
127        }
128
129        // Ensure the collection exists.
130        self.create_collection(collection).await?;
131
132        // Embed passages (blocking, offloaded to blocking thread pool).
133        let texts: Vec<String> = entries.iter().map(|e| e.text.clone()).collect();
134        let embedder = self.embedder.clone();
135        let embeddings = tokio::task::spawn_blocking(move || embedder.embed_batch_passages(&texts))
136            .await
137            .map_err(|e| AgentError::KnowledgeError(format!("Embedding task panicked: {}", e)))?
138            .map_err(|e| AgentError::KnowledgeError(format!("Embedding failed: {}", e)))?;
139
140        // Build Arrow columns.
141        let n = entries.len();
142        let dim = self.dim as i32;
143
144        let ids: Vec<String> = (0..n).map(|_| uuid::Uuid::new_v4().to_string()).collect();
145        let id_array = StringArray::from(ids);
146
147        let vector_array = FixedSizeListArray::from_iter_primitive::<Float32Type, _, _>(
148            embeddings
149                .iter()
150                .map(|emb| Some(emb.iter().map(|&v| Some(v)))),
151            dim,
152        );
153
154        let text_array =
155            StringArray::from(entries.iter().map(|e| e.text.as_str()).collect::<Vec<_>>());
156        let source_array = StringArray::from(
157            entries
158                .iter()
159                .map(|e| e.source.as_str())
160                .collect::<Vec<_>>(),
161        );
162        let metadata_json_array = StringArray::from(
163            entries
164                .iter()
165                .map(|e| {
166                    if e.metadata.is_null() {
167                        None
168                    } else {
169                        Some(serde_json::to_string(&e.metadata).unwrap_or_default())
170                    }
171                })
172                .collect::<Vec<_>>(),
173        );
174        let timestamp_array = UInt64Array::from(
175            entries
176                .iter()
177                .map(|e| e.timestamp as u64)
178                .collect::<Vec<_>>(),
179        );
180
181        let schema = Self::schema(dim);
182        let batch = RecordBatch::try_new(
183            schema.clone(),
184            vec![
185                Arc::new(id_array),
186                Arc::new(vector_array),
187                Arc::new(text_array),
188                Arc::new(source_array),
189                Arc::new(metadata_json_array),
190                Arc::new(timestamp_array),
191            ],
192        )
193        .map_err(|e| AgentError::KnowledgeError(format!("Failed to build RecordBatch: {}", e)))?;
194
195        // Open table and add.
196        let table = self
197            .conn
198            .open_table(collection)
199            .execute()
200            .await
201            .map_err(|e| {
202                AgentError::KnowledgeError(format!(
203                    "Failed to open collection '{}': {}",
204                    collection, e
205                ))
206            })?;
207
208        table.add(batch).execute().await.map_err(|e| {
209            AgentError::KnowledgeError(format!(
210                "Failed to insert into collection '{}': {}",
211                collection, e
212            ))
213        })?;
214
215        Ok(())
216    }
217
218    /// Search a collection for entries similar to `query`.
219    pub async fn search(
220        &self,
221        collection: &str,
222        query: &str,
223        limit: usize,
224    ) -> Result<Vec<SearchResult>> {
225        // Embed the query (blocking, offloaded).
226        let embedder = self.embedder.clone();
227        let query_owned = query.to_owned();
228        let query_vec = tokio::task::spawn_blocking(move || embedder.embed_query(&query_owned))
229            .await
230            .map_err(|e| AgentError::KnowledgeError(format!("Query embed panicked: {}", e)))?
231            .map_err(|e| AgentError::KnowledgeError(format!("Query embed failed: {}", e)))?;
232
233        // Open table (graceful if missing).
234        let table = match self.conn.open_table(collection).execute().await {
235            Ok(t) => t,
236            Err(_) => return Ok(Vec::new()),
237        };
238
239        // Execute the vector search.
240        let results: Vec<RecordBatch> = table
241            .query()
242            .nearest_to(query_vec.as_slice())
243            .map_err(|e| AgentError::KnowledgeError(format!("Failed to build query: {}", e)))?
244            .limit(limit)
245            .execute()
246            .await
247            .map_err(|e| AgentError::KnowledgeError(format!("Search failed: {}", e)))?
248            .try_collect()
249            .await
250            .map_err(|e| AgentError::KnowledgeError(format!("Failed to collect results: {}", e)))?;
251
252        // Parse results.
253        let mut search_results = Vec::new();
254        for batch in &results {
255            let schema = batch.schema();
256            let text_idx = schema.index_of("text").unwrap_or(0);
257            let source_idx = schema.index_of("source").unwrap_or(1);
258            let metadata_idx = schema.index_of("metadata_json").ok();
259            let timestamp_idx = schema.index_of("timestamp").unwrap_or(5);
260
261            // The distance column is _distance.
262            let distance_idx = schema.index_of("_distance").ok();
263
264            for row in 0..batch.num_rows() {
265                let text_col = batch
266                    .column(text_idx)
267                    .as_any()
268                    .downcast_ref::<StringArray>()
269                    .unwrap();
270                let source_col = batch
271                    .column(source_idx)
272                    .as_any()
273                    .downcast_ref::<StringArray>()
274                    .unwrap();
275                let ts_col = batch
276                    .column(timestamp_idx)
277                    .as_any()
278                    .downcast_ref::<UInt64Array>()
279                    .unwrap();
280
281                let text = text_col.value(row).to_owned();
282                let source = source_col.value(row).to_owned();
283                let timestamp = ts_col.value(row) as i64;
284
285                let metadata = if let Some(mi) = metadata_idx {
286                    let meta_col = batch
287                        .column(mi)
288                        .as_any()
289                        .downcast_ref::<StringArray>()
290                        .unwrap();
291                    if meta_col.is_null(row) {
292                        serde_json::Value::Null
293                    } else {
294                        serde_json::from_str(meta_col.value(row)).unwrap_or(serde_json::Value::Null)
295                    }
296                } else {
297                    serde_json::Value::Null
298                };
299
300                let score = if let Some(di) = distance_idx {
301                    let dist_col = batch
302                        .column(di)
303                        .as_any()
304                        .downcast_ref::<Float32Array>()
305                        .unwrap();
306                    // Convert distance to a similarity-like score (lower distance = higher score).
307                    // Use 1.0 / (1.0 + distance) for a simple mapping.
308                    let dist = dist_col.value(row);
309                    1.0 / (1.0 + dist)
310                } else {
311                    1.0
312                };
313
314                search_results.push(SearchResult {
315                    text,
316                    source,
317                    score,
318                    metadata,
319                    timestamp,
320                });
321            }
322        }
323
324        Ok(search_results)
325    }
326
327    /// List all collection (table) names.
328    pub async fn list_collections(&self) -> Result<Vec<String>> {
329        let names = self.conn.table_names().execute().await.map_err(|e| {
330            AgentError::KnowledgeError(format!("Failed to list collections: {}", e))
331        })?;
332        Ok(names)
333    }
334
335    /// Returns the default collection name.
336    pub fn default_collection(&self) -> &str {
337        &self.default_collection
338    }
339}
cascade_agent/knowledge/vectordb.rs

cascade_agent/knowledge/
vectordb.rs
