use super::*;
impl Db {
pub fn attachment_insert(&self, chunk: &AttachmentChunk) -> Result<()> {
let conn = self.conn.lock().unwrap();
ensure_attachment_vec_table(&conn, chunk.embedding.len())?;
conn.execute(
"INSERT INTO attachment_documents(id, chat_id, attachment_id, name, chunk_text, created_at, embed_gen)
VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7)",
params![
chunk.id.to_string(),
chunk.chat_id.to_string(),
chunk.attachment_id.to_string(),
chunk.name,
chunk.text,
chunk.created_at.to_rfc3339(),
current_embed_gen(&conn)?,
],
)?;
let rowid = conn.last_insert_rowid();
conn.execute(
"INSERT INTO attachment_vectors(rowid, chat_id, embedding) VALUES (?1, ?2, ?3)",
params![
rowid,
chunk.chat_id.to_string(),
bytemuck::cast_slice::<f32, u8>(&chunk.embedding),
],
)?;
Ok(())
}
pub fn attachment_search(
&self,
chat_id: Uuid,
query: &[f32],
k: usize,
) -> Result<Vec<AttachmentHit>> {
let conn = self.conn.lock().unwrap();
if !table_exists(&conn, "attachment_vectors")? {
return Ok(Vec::new()); }
let mut stmt = conn.prepare(
"SELECT d.attachment_id, d.name, d.chunk_text, v.distance
FROM attachment_vectors v
JOIN attachment_documents d ON d.rowid = v.rowid
WHERE v.chat_id = ?1 AND v.embedding MATCH ?2 AND k = ?3
AND IFNULL(d.embed_gen, ?4) = ?5
ORDER BY v.distance",
)?;
let hits = stmt
.query_map(
params![
chat_id.to_string(),
bytemuck::cast_slice::<f32, u8>(query),
k as i64,
NULL_EMBED_GEN,
current_embed_gen(&conn)?
],
|r| {
Ok(AttachmentHit {
attachment_id: parse_uuid(r.get::<_, String>(0)?),
name: r.get(1)?,
text: r.get(2)?,
distance: r.get(3)?,
})
},
)?
.collect::<rusqlite::Result<Vec<_>>>()?;
Ok(hits)
}
pub fn attachment_indexed_ids(&self, chat_id: Uuid) -> Result<Vec<Uuid>> {
let conn = self.conn.lock().unwrap();
let mut stmt = conn.prepare(
"SELECT DISTINCT attachment_id FROM attachment_documents
WHERE chat_id = ?1 AND IFNULL(embed_gen, ?2) = ?3",
)?;
let ids = stmt
.query_map(
params![
chat_id.to_string(),
NULL_EMBED_GEN,
current_embed_gen(&conn)?
],
|r| Ok(parse_uuid(r.get::<_, String>(0)?)),
)?
.collect::<rusqlite::Result<Vec<_>>>()?;
Ok(ids)
}
pub fn attachment_known_ids(&self, chat_id: Uuid) -> Result<Vec<Uuid>> {
let conn = self.conn.lock().unwrap();
let mut stmt = conn.prepare(
"SELECT DISTINCT attachment_id FROM attachment_documents WHERE chat_id = ?1",
)?;
let ids = stmt
.query_map(params![chat_id.to_string()], |r| {
Ok(parse_uuid(r.get::<_, String>(0)?))
})?
.collect::<rusqlite::Result<Vec<_>>>()?;
Ok(ids)
}
pub fn attachment_delete(&self, chat_id: Uuid, attachment_id: Uuid) -> Result<usize> {
let conn = self.conn.lock().unwrap();
delete_attachment_rows(&conn, chat_id, |id| id == attachment_id)
}
pub fn attachment_prune(&self, chat_id: Uuid, keep: &[Uuid]) -> Result<usize> {
let conn = self.conn.lock().unwrap();
delete_attachment_rows(&conn, chat_id, |id| !keep.contains(&id))
}
}
fn delete_attachment_rows(
conn: &Connection,
chat_id: Uuid,
pred: impl Fn(Uuid) -> bool,
) -> Result<usize> {
let rows: Vec<(i64, Uuid)> = {
let mut stmt = conn
.prepare("SELECT rowid, attachment_id FROM attachment_documents WHERE chat_id = ?1")?;
stmt.query_map(params![chat_id.to_string()], |r| {
Ok((r.get::<_, i64>(0)?, parse_uuid(r.get::<_, String>(1)?)))
})?
.collect::<rusqlite::Result<Vec<_>>>()?
};
let victims: Vec<i64> = rows
.into_iter()
.filter(|(_, id)| pred(*id))
.map(|(rowid, _)| rowid)
.collect();
if victims.is_empty() {
return Ok(0);
}
let has_vectors = table_exists(conn, "attachment_vectors")?;
for rowid in &victims {
if has_vectors {
conn.execute(
"DELETE FROM attachment_vectors WHERE rowid = ?1",
params![rowid],
)?;
}
conn.execute(
"DELETE FROM attachment_documents WHERE rowid = ?1",
params![rowid],
)?;
}
Ok(victims.len())
}
#[cfg(test)]
mod tests {
use super::*;
fn db() -> Db {
Db::open_in_memory().unwrap()
}
fn chunk(chat: Uuid, att: Uuid, name: &str, text: &str, v: Vec<f32>) -> AttachmentChunk {
AttachmentChunk::new(chat, att, name, text, v)
}
#[test]
fn known_and_indexed_diverge_once_a_generation_is_retired() {
let db = db();
let (chat, att) = (Uuid::new_v4(), Uuid::new_v4());
db.attachment_insert(&chunk(chat, att, "a.txt", "A", vec![1.0, 0.0, 0.0, 0.0]))
.unwrap();
assert_eq!(db.attachment_known_ids(chat).unwrap(), vec![att]);
assert_eq!(db.attachment_indexed_ids(chat).unwrap(), vec![att]);
db.bump_embed_generation().unwrap();
assert_eq!(
db.attachment_known_ids(chat).unwrap(),
vec![att],
"the rows and their text are still there"
);
assert!(
db.attachment_indexed_ids(chat).unwrap().is_empty(),
"but none of them can answer a query"
);
let other = Uuid::new_v4();
assert!(db.attachment_known_ids(other).unwrap().is_empty());
}
#[test]
fn attachment_knn_respects_chat_isolation() {
let db = db();
let (a, b) = (Uuid::new_v4(), Uuid::new_v4());
let (att_a, att_b) = (Uuid::new_v4(), Uuid::new_v4());
db.attachment_insert(&chunk(b, att_b, "b.txt", "B", vec![1.0, 0.0, 0.0, 0.0]))
.unwrap();
db.attachment_insert(&chunk(
a,
att_a,
"a.txt",
"A near",
vec![0.9, 0.1, 0.0, 0.0],
))
.unwrap();
db.attachment_insert(&chunk(a, att_a, "a.txt", "A far", vec![0.0, 0.0, 1.0, 0.0]))
.unwrap();
let hits = db.attachment_search(a, &[1.0, 0.0, 0.0, 0.0], 5).unwrap();
assert_eq!(hits.len(), 2, "only chat A's fragments");
assert_eq!(hits[0].text, "A near");
assert_eq!(hits[0].name, "a.txt");
assert_eq!(hits[0].attachment_id, att_a);
assert_eq!(hits[1].text, "A far");
}
#[test]
fn search_is_empty_before_any_insert() {
let db = db();
assert!(
db.attachment_search(Uuid::new_v4(), &[1.0, 0.0], 5)
.unwrap()
.is_empty()
);
assert!(
db.attachment_indexed_ids(Uuid::new_v4())
.unwrap()
.is_empty()
);
}
#[test]
fn delete_and_prune_scope_to_one_chat() {
let db = db();
let (a, b) = (Uuid::new_v4(), Uuid::new_v4());
let (one, two) = (Uuid::new_v4(), Uuid::new_v4());
db.attachment_insert(&chunk(a, one, "1.txt", "x", vec![1.0, 0.0]))
.unwrap();
db.attachment_insert(&chunk(a, two, "2.txt", "y", vec![0.0, 1.0]))
.unwrap();
db.attachment_insert(&chunk(b, one, "1.txt", "z", vec![1.0, 1.0]))
.unwrap();
assert_eq!(db.attachment_delete(a, one).unwrap(), 1);
assert_eq!(db.attachment_indexed_ids(a).unwrap(), vec![two]);
assert_eq!(db.attachment_indexed_ids(b).unwrap(), vec![one]);
let hits = db.attachment_search(a, &[1.0, 0.0], 5).unwrap();
assert!(hits.iter().all(|h| h.attachment_id == two));
assert_eq!(
db.attachment_prune(a, &[two]).unwrap(),
0,
"nothing to drop"
);
assert_eq!(db.attachment_prune(a, &[]).unwrap(), 1);
assert!(db.attachment_indexed_ids(a).unwrap().is_empty());
assert_eq!(db.attachment_indexed_ids(b).unwrap(), vec![one]);
}
#[test]
fn reindexing_an_attachment_replaces_its_chunks() {
let db = db();
let (chat, att) = (Uuid::new_v4(), Uuid::new_v4());
db.attachment_insert(&chunk(chat, att, "a.txt", "старый текст", vec![1.0, 0.0]))
.unwrap();
db.attachment_delete(chat, att).unwrap();
db.attachment_insert(&chunk(chat, att, "a.txt", "новый текст", vec![1.0, 0.0]))
.unwrap();
let hits = db.attachment_search(chat, &[1.0, 0.0], 5).unwrap();
assert_eq!(hits.len(), 1, "no duplicate from the previous run");
assert_eq!(hits[0].text, "новый текст");
}
#[test]
fn dimension_is_shared_with_rag_and_reset_clears_both() {
let db = db();
let chat = Uuid::new_v4();
let profile = Uuid::new_v4();
db.attachment_insert(&chunk(chat, Uuid::new_v4(), "a.txt", "x", vec![1.0, 0.0]))
.unwrap();
assert_eq!(db.rag_dimension().unwrap(), Some(2));
assert!(
db.rag_insert(&crate::entities::rag::RagDocument::new(
profile,
"s",
"t",
vec![0.0, 1.0, 0.0]
))
.is_err()
);
db.rag_insert(&crate::entities::rag::RagDocument::new(
profile,
"s",
"t",
vec![0.0, 1.0],
))
.unwrap();
assert_eq!(db.reset_vectors().unwrap(), 1, "attachment chunks dropped");
assert_eq!(db.rag_dimension().unwrap(), None);
assert!(db.attachment_indexed_ids(chat).unwrap().is_empty());
assert!(
db.attachment_search(chat, &[1.0, 0.0], 5)
.unwrap()
.is_empty()
);
db.attachment_insert(&chunk(
chat,
Uuid::new_v4(),
"a.txt",
"x",
vec![1.0, 0.0, 0.0],
))
.unwrap();
assert_eq!(db.rag_dimension().unwrap(), Some(3));
}
#[test]
fn rag_search_survives_a_dimension_registered_by_attachments_only() {
let db = db();
db.attachment_insert(&chunk(
Uuid::new_v4(),
Uuid::new_v4(),
"a.txt",
"x",
vec![1.0, 0.0],
))
.unwrap();
assert!(
db.rag_search(Uuid::new_v4(), &[1.0, 0.0], 5)
.unwrap()
.is_empty()
);
}
}