1use std::path::Path;
2
3use anyhow::{Context, Result};
4use rusqlite::params;
5
6const SCHEMA_VERSION: i64 = 1;
7
8#[derive(Debug, Clone)]
9pub struct IndexedFile {
10 pub id: i64,
11 pub file_path: String,
12 pub mtime: f64,
13 pub file_type: String,
14}
15
16#[derive(Debug, Clone)]
17pub struct StoredChunk {
18 pub id: i64,
19 pub file_id: i64,
20 pub file_path: String,
21 pub start_line: i64,
22 pub end_line: i64,
23 pub kind: String,
24 pub name: Option<String>,
25 pub content: String,
26 pub file_type: String,
27}
28
29pub struct SearchDb {
30 conn: rusqlite::Connection,
31}
32
33impl SearchDb {
34 pub fn open(db_path: &Path) -> Result<Self> {
35 if let Some(parent) = db_path.parent() {
36 std::fs::create_dir_all(parent).with_context(|| {
37 format!("Creating search index directory: {}", parent.display())
38 })?;
39 }
40
41 let conn = rusqlite::Connection::open(db_path)
42 .with_context(|| format!("Opening search DB at {}", db_path.display()))?;
43
44 let mut db = Self { conn };
45 db.init_schema()?;
46 Ok(db)
47 }
48
49 fn init_schema(&mut self) -> Result<()> {
50 self.conn
51 .execute_batch(
52 "PRAGMA journal_mode=WAL; PRAGMA synchronous=NORMAL; PRAGMA foreign_keys=ON;",
53 )
54 .context("Setting pragmas")?;
55
56 self.conn
57 .execute_batch("CREATE TABLE IF NOT EXISTS meta (key TEXT PRIMARY KEY, value TEXT)")
58 .context("Creating meta table")?;
59
60 let current_version: i64 = self
61 .conn
62 .query_row(
63 "SELECT value FROM meta WHERE key = 'schema_version'",
64 [],
65 |row| row.get::<_, String>(0),
66 )
67 .ok()
68 .and_then(|s| s.parse().ok())
69 .unwrap_or(0);
70
71 if current_version >= SCHEMA_VERSION {
72 return Ok(());
73 }
74
75 self.conn
76 .execute_batch(
77 "CREATE TABLE IF NOT EXISTS files (
78 id INTEGER PRIMARY KEY AUTOINCREMENT,
79 file_path TEXT NOT NULL UNIQUE,
80 mtime REAL NOT NULL,
81 file_type TEXT NOT NULL
82 );
83
84 CREATE TABLE IF NOT EXISTS chunks (
85 id INTEGER PRIMARY KEY AUTOINCREMENT,
86 file_id INTEGER NOT NULL REFERENCES files(id) ON DELETE CASCADE,
87 file_path TEXT NOT NULL,
88 start_line INTEGER NOT NULL,
89 end_line INTEGER NOT NULL,
90 kind TEXT NOT NULL,
91 name TEXT,
92 content TEXT NOT NULL,
93 file_type TEXT NOT NULL
94 );
95
96 CREATE INDEX IF NOT EXISTS idx_chunks_file_id ON chunks(file_id);
97 CREATE INDEX IF NOT EXISTS idx_chunks_file_path ON chunks(file_path);
98
99 CREATE VIRTUAL TABLE IF NOT EXISTS chunks_fts USING fts5(
100 content,
101 name,
102 file_path,
103 content='chunks',
104 content_rowid='id',
105 tokenize='porter unicode61'
106 );
107
108 CREATE TRIGGER IF NOT EXISTS chunks_ai AFTER INSERT ON chunks BEGIN
109 INSERT INTO chunks_fts(rowid, content, name, file_path)
110 VALUES (new.id, new.content, new.name, new.file_path);
111 END;
112
113 CREATE TRIGGER IF NOT EXISTS chunks_ad AFTER DELETE ON chunks BEGIN
114 INSERT INTO chunks_fts(chunks_fts, rowid, content, name, file_path)
115 VALUES ('delete', old.id, old.content, old.name, old.file_path);
116 END;
117
118 CREATE TRIGGER IF NOT EXISTS chunks_au AFTER UPDATE ON chunks BEGIN
119 INSERT INTO chunks_fts(chunks_fts, rowid, content, name, file_path)
120 VALUES ('delete', old.id, old.content, old.name, old.file_path);
121 INSERT INTO chunks_fts(rowid, content, name, file_path)
122 VALUES (new.id, new.content, new.name, new.file_path);
123 END;
124
125 CREATE TABLE IF NOT EXISTS embeddings (
126 chunk_id INTEGER NOT NULL REFERENCES chunks(id) ON DELETE CASCADE,
127 model_name TEXT NOT NULL,
128 vector BLOB NOT NULL,
129 PRIMARY KEY (chunk_id, model_name)
130 );
131
132 CREATE TABLE IF NOT EXISTS imports (
133 source_file_id INTEGER NOT NULL REFERENCES files(id) ON DELETE CASCADE,
134 target_file_path TEXT NOT NULL,
135 PRIMARY KEY (source_file_id, target_file_path)
136 );
137
138 CREATE INDEX IF NOT EXISTS idx_imports_target ON imports(target_file_path);
139
140 CREATE TABLE IF NOT EXISTS symbols (
141 chunk_id INTEGER NOT NULL REFERENCES chunks(id) ON DELETE CASCADE,
142 name TEXT NOT NULL,
143 kind TEXT NOT NULL
144 );
145
146 CREATE INDEX IF NOT EXISTS idx_symbols_name ON symbols(name);
147 ",
148 )
149 .context("Creating schema")?;
150
151 self.conn
152 .execute(
153 "INSERT OR REPLACE INTO meta (key, value) VALUES ('schema_version', ?1)",
154 params![SCHEMA_VERSION.to_string()],
155 )
156 .context("Setting schema version")?;
157
158 Ok(())
159 }
160
161 pub fn upsert_file(&mut self, file_path: &str, mtime: f64, file_type: &str) -> Result<i64> {
162 let existing: Option<i64> = self
163 .conn
164 .query_row(
165 "SELECT id FROM files WHERE file_path = ?1",
166 params![file_path],
167 |row| row.get(0),
168 )
169 .ok();
170
171 if let Some(id) = existing {
172 self.conn.execute(
173 "UPDATE files SET mtime = ?1, file_type = ?2 WHERE id = ?3",
174 params![mtime, file_type, id],
175 )?;
176 return Ok(id);
177 }
178
179 self.conn.execute(
180 "INSERT INTO files (file_path, mtime, file_type) VALUES (?1, ?2, ?3)",
181 params![file_path, mtime, file_type],
182 )?;
183 Ok(self.conn.last_insert_rowid())
184 }
185
186 pub fn get_file(&mut self, file_path: &str) -> Result<Option<IndexedFile>> {
187 let result = self.conn.query_row(
188 "SELECT id, file_path, mtime, file_type FROM files WHERE file_path = ?1",
189 params![file_path],
190 |row| {
191 Ok(IndexedFile {
192 id: row.get(0)?,
193 file_path: row.get(1)?,
194 mtime: row.get(2)?,
195 file_type: row.get(3)?,
196 })
197 },
198 );
199 match result {
200 Ok(f) => Ok(Some(f)),
201 Err(rusqlite::Error::QueryReturnedNoRows) => Ok(None),
202 Err(e) => Err(e.into()),
203 }
204 }
205
206 pub fn get_all_files(&mut self) -> Result<Vec<IndexedFile>> {
207 let mut stmt = self
208 .conn
209 .prepare("SELECT id, file_path, mtime, file_type FROM files")?;
210 let rows = stmt.query_map([], |row| {
211 Ok(IndexedFile {
212 id: row.get(0)?,
213 file_path: row.get(1)?,
214 mtime: row.get(2)?,
215 file_type: row.get(3)?,
216 })
217 })?;
218 rows.collect::<Result<Vec<_>, _>>().map_err(Into::into)
219 }
220
221 pub fn delete_file(&mut self, file_id: i64) -> Result<()> {
222 self.conn
223 .execute("DELETE FROM files WHERE id = ?1", params![file_id])?;
224 Ok(())
225 }
226
227 #[allow(clippy::too_many_arguments)]
228 pub fn insert_chunk(
229 &mut self,
230 file_id: i64,
231 file_path: &str,
232 start_line: i64,
233 end_line: i64,
234 kind: &str,
235 name: Option<&str>,
236 content: &str,
237 file_type: &str,
238 ) -> Result<i64> {
239 self.conn.execute(
240 "INSERT INTO chunks (file_id, file_path, start_line, end_line, kind, name, content, file_type) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8)",
241 params![file_id, file_path, start_line, end_line, kind, name, content, file_type],
242 )?;
243 Ok(self.conn.last_insert_rowid())
244 }
245
246 pub fn delete_chunks_for_file(&mut self, file_id: i64) -> Result<()> {
247 self.conn
248 .execute("DELETE FROM chunks WHERE file_id = ?1", params![file_id])?;
249 Ok(())
250 }
251
252 pub fn get_all_chunks(&mut self) -> Result<Vec<StoredChunk>> {
253 let mut stmt = self.conn.prepare(
254 "SELECT id, file_id, file_path, start_line, end_line, kind, name, content, file_type FROM chunks",
255 )?;
256 let rows = stmt.query_map([], |row| {
257 Ok(StoredChunk {
258 id: row.get(0)?,
259 file_id: row.get(1)?,
260 file_path: row.get(2)?,
261 start_line: row.get(3)?,
262 end_line: row.get(4)?,
263 kind: row.get(5)?,
264 name: row.get(6)?,
265 content: row.get(7)?,
266 file_type: row.get(8)?,
267 })
268 })?;
269 rows.collect::<Result<Vec<_>, _>>().map_err(Into::into)
270 }
271
272 pub fn get_chunks_by_ids(&mut self, chunk_ids: &[i64]) -> Result<Vec<StoredChunk>> {
273 if chunk_ids.is_empty() {
274 return Ok(Vec::new());
275 }
276
277 let batch_size = 500;
278 let mut results = Vec::new();
279
280 for chunk in chunk_ids.chunks(batch_size) {
281 let placeholders: Vec<String> = chunk
282 .iter()
283 .enumerate()
284 .map(|(i, _)| format!("?{}", i + 1))
285 .collect();
286 let sql = format!(
287 "SELECT id, file_id, file_path, start_line, end_line, kind, name, content, file_type FROM chunks WHERE id IN ({})",
288 placeholders.join(",")
289 );
290 let params: Vec<&dyn rusqlite::ToSql> =
291 chunk.iter().map(|id| id as &dyn rusqlite::ToSql).collect();
292 let mut stmt = self.conn.prepare(&sql)?;
293 let rows = stmt.query_map(params.as_slice(), |row| {
294 Ok(StoredChunk {
295 id: row.get(0)?,
296 file_id: row.get(1)?,
297 file_path: row.get(2)?,
298 start_line: row.get(3)?,
299 end_line: row.get(4)?,
300 kind: row.get(5)?,
301 name: row.get(6)?,
302 content: row.get(7)?,
303 file_type: row.get(8)?,
304 })
305 })?;
306 results.extend(rows.collect::<Result<Vec<_>, _>>()?);
307 }
308
309 Ok(results)
310 }
311
312 pub fn upsert_embedding(
313 &mut self,
314 chunk_id: i64,
315 model_name: &str,
316 vector_blob: &[u8],
317 ) -> Result<()> {
318 self.conn.execute(
319 "INSERT OR REPLACE INTO embeddings (chunk_id, model_name, vector) VALUES (?1, ?2, ?3)",
320 params![chunk_id, model_name, vector_blob],
321 )?;
322 Ok(())
323 }
324
325 pub fn batch_upsert_embeddings(&mut self, items: &[(i64, String, Vec<u8>)]) -> Result<()> {
326 let tx = self.conn.transaction()?;
327 for (chunk_id, model_name, vector_blob) in items {
328 tx.execute(
329 "INSERT OR REPLACE INTO embeddings (chunk_id, model_name, vector) VALUES (?1, ?2, ?3)",
330 params![chunk_id, model_name, vector_blob],
331 )
332 .with_context(|| "batch upsert embedding")?;
333 }
334 tx.commit()?;
335 Ok(())
336 }
337
338 pub fn get_all_embeddings(&mut self, model_name: &str) -> Result<Vec<(i64, Vec<u8>)>> {
339 let mut stmt = self
340 .conn
341 .prepare("SELECT chunk_id, vector FROM embeddings WHERE model_name = ?1")?;
342 let rows = stmt.query_map(params![model_name], |row| {
343 let chunk_id: i64 = row.get(0)?;
344 let vector: Vec<u8> = row.get(1)?;
345 Ok((chunk_id, vector))
346 })?;
347 rows.collect::<Result<Vec<_>, _>>().map_err(Into::into)
348 }
349
350 pub fn get_chunk_ids_without_embedding(&mut self, model_name: &str) -> Result<Vec<i64>> {
351 let mut stmt = self.conn.prepare(
352 "SELECT c.id FROM chunks c LEFT JOIN embeddings e ON c.id = e.chunk_id AND e.model_name = ?1 WHERE e.chunk_id IS NULL",
353 )?;
354 let rows = stmt.query_map(params![model_name], |row| row.get::<_, i64>(0))?;
355 rows.collect::<Result<Vec<_>, _>>().map_err(Into::into)
356 }
357
358 pub fn fts_search(&mut self, query: &str, limit: i64) -> Result<Vec<(i64, f64)>> {
359 let result = self
360 .conn
361 .prepare(
362 "SELECT chunks_fts.rowid as chunk_id, -bm25(chunks_fts, 1.0, 10.0, 5.0) as score FROM chunks_fts WHERE chunks_fts MATCH ?1 ORDER BY score DESC LIMIT ?2",
363 );
364 let mut stmt = match result {
365 Ok(s) => s,
366 Err(_) => return Ok(Vec::new()),
367 };
368
369 let rows = stmt.query_map(params![query, limit], |row| {
370 let id: i64 = row.get(0)?;
371 let score: f64 = row.get(1)?;
372 Ok((id, score))
373 });
374 match rows {
375 Ok(mapped) => {
376 let mut results = Vec::new();
377 for row in mapped {
378 match row {
379 Ok(r) => results.push(r),
380 Err(_) => return Ok(Vec::new()),
381 }
382 }
383 Ok(results)
384 }
385 Err(_) => Ok(Vec::new()),
386 }
387 }
388
389 pub fn rebuild_fts(&mut self) -> Result<()> {
390 self.conn
391 .execute_batch("INSERT INTO chunks_fts(chunks_fts) VALUES ('rebuild')")?;
392 Ok(())
393 }
394
395 pub fn insert_import(&mut self, source_file_id: i64, target_file_path: &str) -> Result<()> {
396 self.conn.execute(
397 "INSERT OR IGNORE INTO imports (source_file_id, target_file_path) VALUES (?1, ?2)",
398 params![source_file_id, target_file_path],
399 )?;
400 Ok(())
401 }
402
403 pub fn delete_imports_for_file(&mut self, source_file_id: i64) -> Result<()> {
404 self.conn.execute(
405 "DELETE FROM imports WHERE source_file_id = ?1",
406 params![source_file_id],
407 )?;
408 Ok(())
409 }
410
411 pub fn get_imports_from(&mut self, source_file_id: i64) -> Result<Vec<String>> {
412 let mut stmt = self
413 .conn
414 .prepare("SELECT target_file_path FROM imports WHERE source_file_id = ?1")?;
415 let rows = stmt.query_map(params![source_file_id], |row| row.get::<_, String>(0))?;
416 rows.collect::<Result<Vec<_>, _>>().map_err(Into::into)
417 }
418
419 pub fn get_importers_of(&mut self, target_file_path: &str) -> Result<Vec<i64>> {
420 let mut stmt = self
421 .conn
422 .prepare("SELECT source_file_id FROM imports WHERE target_file_path = ?1")?;
423 let rows = stmt.query_map(params![target_file_path], |row| row.get::<_, i64>(0))?;
424 rows.collect::<Result<Vec<_>, _>>().map_err(Into::into)
425 }
426
427 pub fn insert_symbol(&mut self, chunk_id: i64, name: &str, kind: &str) -> Result<()> {
428 self.conn.execute(
429 "INSERT INTO symbols (chunk_id, name, kind) VALUES (?1, ?2, ?3)",
430 params![chunk_id, name, kind],
431 )?;
432 Ok(())
433 }
434
435 pub fn get_all_symbols(&mut self) -> Result<Vec<(i64, String)>> {
436 let mut stmt = self.conn.prepare("SELECT chunk_id, name FROM symbols")?;
437 let rows = stmt.query_map([], |row| {
438 let chunk_id: i64 = row.get(0)?;
439 let name: String = row.get(1)?;
440 Ok((chunk_id, name))
441 })?;
442 rows.collect::<Result<Vec<_>, _>>().map_err(Into::into)
443 }
444
445 pub fn get_chunk_count(&mut self) -> Result<i64> {
446 self.conn
447 .query_row("SELECT COUNT(*) FROM chunks", [], |row| row.get(0))
448 .map_err(Into::into)
449 }
450
451 pub fn get_file_count(&mut self) -> Result<i64> {
452 self.conn
453 .query_row("SELECT COUNT(*) FROM files", [], |row| row.get(0))
454 .map_err(Into::into)
455 }
456}
457
458#[cfg(test)]
459mod tests {
460 use super::*;
461 use tempfile::TempDir;
462
463 fn test_db() -> SearchDb {
464 let temp = TempDir::new().expect("temp dir");
465 let db_path = temp.path().join("test.db");
466 SearchDb::open(&db_path).expect("should open db")
467 }
468
469 #[test]
470 fn schema_initializes_without_error() {
471 let _db = test_db();
472 }
473
474 #[test]
475 fn upsert_and_get_file() {
476 let mut db = test_db();
477 let id = db
478 .upsert_file("src/main.rs", 1234.5, "rust")
479 .expect("upsert");
480 assert!(id > 0);
481
482 let file = db.get_file("src/main.rs").expect("get");
483 assert!(file.is_some());
484 let f = file.expect("file");
485 assert_eq!(f.file_path, "src/main.rs");
486 assert_eq!(f.file_type, "rust");
487 }
488
489 #[test]
490 fn upsert_file_idempotent() {
491 let mut db = test_db();
492 let id1 = db.upsert_file("test.rs", 1.0, "rust").expect("upsert1");
493 let id2 = db.upsert_file("test.rs", 2.0, "rust").expect("upsert2");
494 assert_eq!(id1, id2);
495 }
496
497 #[test]
498 fn insert_and_get_chunks() {
499 let mut db = test_db();
500 let file_id = db.upsert_file("main.rs", 1.0, "rust").expect("file");
501
502 let chunk_id = db
503 .insert_chunk(
504 file_id,
505 "main.rs",
506 1,
507 5,
508 "function",
509 Some("main"),
510 "fn main() {}",
511 "rust",
512 )
513 .expect("insert chunk");
514
515 assert!(chunk_id > 0);
516
517 let chunks = db.get_all_chunks().expect("get all");
518 assert_eq!(chunks.len(), 1);
519 assert_eq!(chunks[0].name, Some("main".to_string()));
520 }
521
522 #[test]
523 fn delete_file_cascades_to_chunks() {
524 let mut db = test_db();
525 let file_id = db.upsert_file("temp.rs", 1.0, "rust").expect("file");
526 db.insert_chunk(file_id, "temp.rs", 1, 1, "file", None, "code", "rust")
527 .expect("chunk");
528
529 db.delete_file(file_id).expect("delete");
530
531 let chunks = db.get_all_chunks().expect("chunks");
532 assert!(chunks.is_empty());
533 }
534
535 #[test]
536 fn embedding_crud() {
537 let mut db = test_db();
538 let file_id = db.upsert_file("test.rs", 1.0, "rust").expect("file");
539 let chunk_id = db
540 .insert_chunk(file_id, "test.rs", 1, 1, "file", None, "code", "rust")
541 .expect("chunk");
542
543 let vector = vec![0.1_f32, 0.2, 0.3];
544 let blob = crate::vector_store::pack_vector(&vector);
545 db.upsert_embedding(chunk_id, "test-model", &blob)
546 .expect("upsert embedding");
547
548 let embeddings = db.get_all_embeddings("test-model").expect("get embeddings");
549 assert_eq!(embeddings.len(), 1);
550 }
551
552 #[test]
553 fn get_chunk_ids_without_embedding() {
554 let mut db = test_db();
555 let file_id = db.upsert_file("test.rs", 1.0, "rust").expect("file");
556 let c1 = db
557 .insert_chunk(file_id, "test.rs", 1, 1, "file", None, "a", "rust")
558 .expect("chunk");
559 let c2 = db
560 .insert_chunk(file_id, "test.rs", 2, 2, "file", None, "b", "rust")
561 .expect("chunk");
562
563 let vector = vec![0.1_f32];
564 let blob = crate::vector_store::pack_vector(&vector);
565 db.upsert_embedding(c1, "model", &blob).expect("embed");
566
567 let missing = db
568 .get_chunk_ids_without_embedding("model")
569 .expect("missing");
570 assert_eq!(missing, vec![c2]);
571 }
572
573 #[test]
574 fn import_crud() {
575 let mut db = test_db();
576 let f1 = db.upsert_file("main.rs", 1.0, "rust").expect("file");
577 db.upsert_file("lib.rs", 1.0, "rust").expect("file2");
578 db.insert_import(f1, "lib.rs").expect("insert import");
579
580 let imports = db.get_imports_from(f1).expect("imports");
581 assert_eq!(imports, vec!["lib.rs"]);
582
583 let importers = db.get_importers_of("lib.rs").expect("importers");
584 assert_eq!(importers, vec![f1]);
585 }
586
587 #[test]
588 fn symbol_crud() {
589 let mut db = test_db();
590 let file_id = db.upsert_file("test.rs", 1.0, "rust").expect("file");
591 let chunk_id = db
592 .insert_chunk(
593 file_id,
594 "test.rs",
595 1,
596 5,
597 "function",
598 Some("main"),
599 "fn main() {}",
600 "rust",
601 )
602 .expect("chunk");
603
604 db.insert_symbol(chunk_id, "main", "function")
605 .expect("symbol");
606
607 let symbols = db.get_all_symbols().expect("symbols");
608 assert_eq!(symbols.len(), 1);
609 assert_eq!(symbols[0], (chunk_id, "main".to_string()));
610 }
611
612 #[test]
613 fn get_chunks_by_ids() {
614 let mut db = test_db();
615 let file_id = db.upsert_file("test.rs", 1.0, "rust").expect("file");
616 let c1 = db
617 .insert_chunk(file_id, "test.rs", 1, 1, "file", None, "a", "rust")
618 .expect("chunk");
619 let c2 = db
620 .insert_chunk(file_id, "test.rs", 2, 2, "file", None, "b", "rust")
621 .expect("chunk");
622
623 let chunks = db.get_chunks_by_ids(&[c1, c2]).expect("get by ids");
624 assert_eq!(chunks.len(), 2);
625 }
626
627 #[test]
628 fn get_all_files_returns_inserted() {
629 let mut db = test_db();
630 db.upsert_file("a.rs", 1.0, "rust").expect("file");
631 db.upsert_file("b.rs", 2.0, "rust").expect("file");
632
633 let files = db.get_all_files().expect("get all");
634 assert_eq!(files.len(), 2);
635 }
636
637 #[test]
638 fn fts_search_after_insert() {
639 let mut db = test_db();
640 let file_id = db.upsert_file("test.rs", 1.0, "rust").expect("file");
641 db.insert_chunk(
642 file_id,
643 "test.rs",
644 1,
645 5,
646 "function",
647 Some("search_engine"),
648 "fn search_engine() { /* search implementation */ }",
649 "rust",
650 )
651 .expect("chunk");
652
653 let results = db.fts_search("search_engine", 10).expect("fts");
654 assert_eq!(results.len(), 1);
655 assert!(results[0].1 > 0.0);
656 }
657
658 #[test]
659 fn fts_search_malformed_query_returns_empty() {
660 let mut db = test_db();
661 let file_id = db.upsert_file("test.rs", 1.0, "rust").expect("file");
662 db.insert_chunk(file_id, "test.rs", 1, 1, "file", None, "content", "rust")
663 .expect("chunk");
664
665 let results = db.fts_search("??? OR AND NOT", 10).expect("fts");
666 assert!(results.is_empty());
667 }
668
669 #[test]
670 fn batch_upsert_embeddings_transactional() {
671 let mut db = test_db();
672 let file_id = db.upsert_file("test.rs", 1.0, "rust").expect("file");
673 let c1 = db
674 .insert_chunk(file_id, "test.rs", 1, 1, "file", None, "a", "rust")
675 .expect("chunk");
676 let c2 = db
677 .insert_chunk(file_id, "test.rs", 2, 2, "file", None, "b", "rust")
678 .expect("chunk");
679
680 let blob1 = crate::vector_store::pack_vector(&[0.1_f32]);
681 let blob2 = crate::vector_store::pack_vector(&[0.2_f32]);
682
683 db.batch_upsert_embeddings(&[
684 (c1, "model".to_string(), blob1),
685 (c2, "model".to_string(), blob2),
686 ])
687 .expect("batch upsert");
688
689 let embeddings = db.get_all_embeddings("model").expect("get");
690 assert_eq!(embeddings.len(), 2);
691 }
692}