From cc36bb075e3b8c8e8bac9b9319ddf54174d00c91 Mon Sep 17 00:00:00 2001 From: thisisjun786 Date: Mon, 10 Aug 2026 10:48:28 +0900 Subject: [PATCH 1/2] perf: bulk load semantic parent windows --- src/archive/parent.rs | 113 ++++++++++++++++++++++++++++++++++++------ 1 file changed, 98 insertions(+), 15 deletions(-) diff --git a/src/archive/parent.rs b/src/archive/parent.rs index 47929f5..9b0ed98 100644 --- a/src/archive/parent.rs +++ b/src/archive/parent.rs @@ -4,6 +4,7 @@ use crate::{ Error, Result, }; use rusqlite::{params, OptionalExtension}; +use std::collections::HashMap; impl Archive { pub fn get_parent_chunk(&self, parent_id: &str) -> Result> { @@ -64,21 +65,52 @@ impl Archive { } pub fn all_parent_chunks(&self) -> Result> { - let mut stmt = self - .conn - .prepare("SELECT parent_id FROM parent_chunks ORDER BY started_at, parent_id") - .map_err(Error::Sql)?; - let ids = stmt - .query_map([], |row| row.get::<_, String>(0)) - .map_err(Error::Sql)? - .collect::, _>>() - .map_err(Error::Sql)?; - ids.into_iter() - .map(|id| { - self.get_parent_chunk(&id) - .and_then(|parent| parent.ok_or(Error::MissingChunk(id))) - }) - .collect() + let tx = self.conn.unchecked_transaction().map_err(Error::Sql)?; + let mut parents = { + let mut stmt = tx + .prepare( + "SELECT parent_id, chat_id, chat_name, started_at, ended_at, + text, message_count + FROM parent_chunks + ORDER BY started_at, parent_id", + ) + .map_err(Error::Sql)?; + let parents = stmt + .query_map([], |row| { + Ok(ParentChunk { + parent_id: row.get(0)?, + chat_id: row.get(1)?, + chat_name: row.get(2)?, + started_at: row.get(3)?, + ended_at: row.get(4)?, + text: row.get(5)?, + message_count: row.get::<_, i64>(6)? as usize, + child_chunk_ids: Vec::new(), + }) + }) + .map_err(Error::Sql)? + .collect::, _>>() + .map_err(Error::Sql)?; + parents + }; + let children = { + let mut stmt = tx + .prepare( + "SELECT parent_id, chunk_id + FROM parent_chunk_children + ORDER BY parent_id, ordinal", + ) + .map_err(Error::Sql)?; + let children = stmt + .query_map([], |row| Ok((row.get(0)?, row.get(1)?))) + .map_err(Error::Sql)? + .collect::, _>>() + .map_err(Error::Sql)?; + children + }; + tx.commit().map_err(Error::Sql)?; + attach_parent_children(&mut parents, children)?; + Ok(parents) } pub(super) fn window_parent_ids(&self, chunk_id: &str) -> Result> { @@ -158,7 +190,58 @@ impl Archive { } } +fn attach_parent_children( + parents: &mut [ParentChunk], + children: impl IntoIterator, +) -> Result<()> { + let parent_indexes = parents + .iter() + .enumerate() + .map(|(index, parent)| (parent.parent_id.clone(), index)) + .collect::>(); + for (parent_id, chunk_id) in children { + let Some(index) = parent_indexes.get(&parent_id).copied() else { + return Err(Error::MissingChunk(parent_id)); + }; + parents[index].child_chunk_ids.push(chunk_id); + } + Ok(()) +} + enum Neighbor { Previous, Next, } + +#[cfg(test)] +mod tests { + use super::*; + + fn parent(parent_id: &str) -> ParentChunk { + ParentChunk { + parent_id: parent_id.to_string(), + chat_id: "chat-fixture".to_string(), + chat_name: "Synthetic room".to_string(), + started_at: "2026-01-01T00:00:00Z".to_string(), + ended_at: "2026-01-01T00:01:00Z".to_string(), + text: "synthetic text".to_string(), + message_count: 1, + child_chunk_ids: Vec::new(), + } + } + + #[test] + fn bulk_parent_assembly_preserves_parent_and_child_order() { + let mut parents = vec![parent("parent-early"), parent("parent-late")]; + let children = vec![ + ("parent-early".to_string(), "child-a".to_string()), + ("parent-early".to_string(), "child-b".to_string()), + ("parent-late".to_string(), "child-c".to_string()), + ]; + + attach_parent_children(&mut parents, children).expect("attach children"); + + assert_eq!(parents[0].child_chunk_ids, ["child-a", "child-b"]); + assert_eq!(parents[1].child_chunk_ids, ["child-c"]); + } +} From 6ee94e78dd48b852faf494b89b97ce93ffc17b92 Mon Sep 17 00:00:00 2001 From: Soohong Kim Date: Wed, 12 Aug 2026 15:41:45 +0900 Subject: [PATCH 2/2] perf: integrate semantic bulk index preparation --- src/archive/parent.rs | 105 +++++++++++--------- src/commands/index_commands.rs | 22 +++-- src/semantic.rs | 15 ++- src/semantic/live.rs | 173 ++++++++++++++++++++++----------- src/semantic/mock.rs | 47 ++++++--- src/semantic/store.rs | 72 +++++++++++--- tests/cli_behaviors.rs | 50 ++++++++++ tests/live_semantic.rs | 14 +++ tests/planned_behaviors.rs | 29 +++++- tests/search_index_recovery.rs | 47 +++++++++ 10 files changed, 430 insertions(+), 144 deletions(-) diff --git a/src/archive/parent.rs b/src/archive/parent.rs index 9b0ed98..d2217e5 100644 --- a/src/archive/parent.rs +++ b/src/archive/parent.rs @@ -3,7 +3,7 @@ use crate::{ types::{Chunk, ChunkContext, ChunkSummary, ParentChunk}, Error, Result, }; -use rusqlite::{params, OptionalExtension}; +use rusqlite::{params, Connection, OptionalExtension}; use std::collections::HashMap; impl Archive { @@ -65,51 +65,12 @@ impl Archive { } pub fn all_parent_chunks(&self) -> Result> { + if !self.conn.is_autocommit() { + return load_all_parent_chunks(&self.conn); + } let tx = self.conn.unchecked_transaction().map_err(Error::Sql)?; - let mut parents = { - let mut stmt = tx - .prepare( - "SELECT parent_id, chat_id, chat_name, started_at, ended_at, - text, message_count - FROM parent_chunks - ORDER BY started_at, parent_id", - ) - .map_err(Error::Sql)?; - let parents = stmt - .query_map([], |row| { - Ok(ParentChunk { - parent_id: row.get(0)?, - chat_id: row.get(1)?, - chat_name: row.get(2)?, - started_at: row.get(3)?, - ended_at: row.get(4)?, - text: row.get(5)?, - message_count: row.get::<_, i64>(6)? as usize, - child_chunk_ids: Vec::new(), - }) - }) - .map_err(Error::Sql)? - .collect::, _>>() - .map_err(Error::Sql)?; - parents - }; - let children = { - let mut stmt = tx - .prepare( - "SELECT parent_id, chunk_id - FROM parent_chunk_children - ORDER BY parent_id, ordinal", - ) - .map_err(Error::Sql)?; - let children = stmt - .query_map([], |row| Ok((row.get(0)?, row.get(1)?))) - .map_err(Error::Sql)? - .collect::, _>>() - .map_err(Error::Sql)?; - children - }; + let parents = load_all_parent_chunks(&tx)?; tx.commit().map_err(Error::Sql)?; - attach_parent_children(&mut parents, children)?; Ok(parents) } @@ -190,6 +151,51 @@ impl Archive { } } +fn load_all_parent_chunks(conn: &Connection) -> Result> { + let mut parents = { + let mut stmt = conn + .prepare( + "SELECT parent_id, chat_id, chat_name, started_at, ended_at, + text, message_count + FROM parent_chunks + ORDER BY started_at, parent_id", + ) + .map_err(Error::Sql)?; + let rows = stmt + .query_map([], |row| { + Ok(ParentChunk { + parent_id: row.get(0)?, + chat_id: row.get(1)?, + chat_name: row.get(2)?, + started_at: row.get(3)?, + ended_at: row.get(4)?, + text: row.get(5)?, + message_count: row.get::<_, i64>(6)? as usize, + child_chunk_ids: Vec::new(), + }) + }) + .map_err(Error::Sql)?; + rows.collect::, _>>() + .map_err(Error::Sql)? + }; + let children = { + let mut stmt = conn + .prepare( + "SELECT parent_id, chunk_id + FROM parent_chunk_children + ORDER BY parent_id, ordinal", + ) + .map_err(Error::Sql)?; + let rows = stmt + .query_map([], |row| Ok((row.get(0)?, row.get(1)?))) + .map_err(Error::Sql)?; + rows.collect::, _>>() + .map_err(Error::Sql)? + }; + attach_parent_children(&mut parents, children)?; + Ok(parents) +} + fn attach_parent_children( parents: &mut [ParentChunk], children: impl IntoIterator, @@ -244,4 +250,15 @@ mod tests { assert_eq!(parents[0].child_chunk_ids, ["child-a", "child-b"]); assert_eq!(parents[1].child_chunk_ids, ["child-c"]); } + + #[test] + fn bulk_parent_load_reuses_caller_transaction() { + let dir = tempfile::tempdir().expect("tempdir"); + let archive = Archive::open(&dir.path().join("archive.sqlite3")).expect("open archive"); + + let parents: Result> = + archive.in_transaction(|| archive.all_parent_chunks()); + + assert!(parents.expect("load parents inside transaction").is_empty()); + } } diff --git a/src/commands/index_commands.rs b/src/commands/index_commands.rs index 7598f0e..2fbb5e3 100644 --- a/src/commands/index_commands.rs +++ b/src/commands/index_commands.rs @@ -3,7 +3,8 @@ use anyhow::{Context, Result}; use katok::{ archive::Archive, config::KatokConfig, - semantic::{index_semantic_live, planned_semantic_documents}, + semantic::{index_semantic_live_for_parents, planned_semantic_documents_for_parents}, + types::ParentChunk, }; use std::path::Path; @@ -20,7 +21,10 @@ pub(crate) fn run( anyhow::bail!("archive is missing; run katok sync before katok index"); } let archive = Archive::open(archive_path).context("open archive")?; - let chunks = archive.all_chunks().context("load chunks")?; + let candidate_chunks = archive.chunk_count().context("count chunks")?; + let parents = archive + .all_parent_chunks() + .context("load semantic parent windows")?; if !dry_run { return run_live_index(LiveIndexInput { full, @@ -28,15 +32,16 @@ pub(crate) fn run( json, config, archive: &archive, + parents: &parents, semantic_dir, - candidate_chunks: chunks.len(), + candidate_chunks, }); } - let documents = planned_semantic_documents(&archive, semantic_dir).context("plan documents")?; + let documents = planned_semantic_documents_for_parents(&parents, semantic_dir); let payload = serde_json::json!({ "full": full, "dry_run": dry_run, - "candidate_chunks": chunks.len(), + "candidate_chunks": candidate_chunks, "written_documents": 0, "embedding_calls": 0, "documents": documents, @@ -52,6 +57,7 @@ struct LiveIndexInput<'a> { json: bool, config: &'a KatokConfig, archive: &'a Archive, + parents: &'a [ParentChunk], semantic_dir: &'a Path, candidate_chunks: usize, } @@ -62,8 +68,9 @@ fn run_live_index(input: LiveIndexInput<'_>) -> Result<()> { .build() .context("create semantic runtime")?; let report = runtime - .block_on(index_semantic_live( + .block_on(index_semantic_live_for_parents( input.archive, + input.parents, input.semantic_dir, input.config, input.full, @@ -71,8 +78,7 @@ fn run_live_index(input: LiveIndexInput<'_>) -> Result<()> { .context("index semantic documents")?; let generation = katok::semantic::current_generation(input.semantic_dir) .context("resolve committed semantic generation")?; - let documents = planned_semantic_documents(input.archive, &generation) - .context("report semantic documents")?; + let documents = planned_semantic_documents_for_parents(input.parents, &generation); let payload = serde_json::json!({ "full": input.full, "dry_run": input.dry_run, diff --git a/src/semantic.rs b/src/semantic.rs index dda4e73..51b2ded 100644 --- a/src/semantic.rs +++ b/src/semantic.rs @@ -8,12 +8,13 @@ use sha2::{Digest, Sha256}; use std::path::{Path, PathBuf}; pub use live::{ - committed_cursor, index_semantic_live, semantic_search_live_with_config, SemanticIndexReport, - STORE_DIR, + committed_cursor, index_semantic_live, index_semantic_live_for_parents, + semantic_search_live_with_config, SemanticIndexReport, STORE_DIR, }; pub use mock::{ - planned_semantic_documents, semantic_search, semantic_search_with_snippet, - write_semantic_documents, SemanticDocument, + planned_semantic_documents, planned_semantic_documents_for_parents, semantic_search, + semantic_search_with_snippet, write_semantic_documents, write_semantic_documents_for_parents, + SemanticDocument, }; pub(crate) const CHUNK_SCHEMA_ID: &str = "katok-kakao-parent-window-v1"; @@ -27,6 +28,10 @@ pub fn semantic_source_dir(root: &std::path::Path) -> std::path::PathBuf { pub fn archive_revision(archive: &Archive) -> Result { let parents = archive.all_parent_chunks()?; + Ok(archive_revision_for_parents(&parents)) +} + +pub(crate) fn archive_revision_for_parents(parents: &[crate::types::ParentChunk]) -> String { let mut material = String::new(); for parent in parents { material.push_str(&parent.parent_id); @@ -34,7 +39,7 @@ pub fn archive_revision(archive: &Archive) -> Result { material.push_str(&content_hash(&parent.text)); material.push('\0'); } - Ok(content_hash(&material)) + content_hash(&material) } pub fn current_generation(root: &Path) -> Result { diff --git a/src/semantic/live.rs b/src/semantic/live.rs index 923b5d8..5f83ac6 100644 --- a/src/semantic/live.rs +++ b/src/semantic/live.rs @@ -1,16 +1,20 @@ use crate::{ - archive::Archive, config::KatokConfig, search::hydrate_parent_hits, types::SearchHit, Error, - Result, + archive::Archive, + config::KatokConfig, + search::hydrate_parent_hits, + types::{ParentChunk, SearchHit}, + Error, Result, }; +use rayon::prelude::*; use rusqlite::Connection; use serde::{Deserialize, Serialize}; -use std::{path::Path, time::Duration}; +use std::{collections::HashMap, path::Path, time::Duration}; use super::{ - archive_revision, content_hash, current_generation, + archive_revision, archive_revision_for_parents, content_hash, current_generation, embedder::create_embedder, - mock::write_semantic_documents_plain, - store::{LocalVectorStore, VectorUpsert}, + mock::write_semantic_documents_plain_for_parents, + store::{LocalVectorStore, StoredVector, VectorUpsert}, CHUNK_SCHEMA_ID, CURRENT_FILE, GENERATIONS_DIR, SOURCE_ID, }; @@ -47,11 +51,22 @@ pub async fn index_semantic_live( root: &Path, config: &KatokConfig, full: bool, +) -> Result { + let parents = archive.all_parent_chunks()?; + index_semantic_live_for_parents(archive, &parents, root, config, full).await +} + +pub async fn index_semantic_live_for_parents( + archive: &Archive, + parents: &[ParentChunk], + root: &Path, + config: &KatokConfig, + full: bool, ) -> Result { crate::paths::ensure_private_dir(root)?; crate::paths::ensure_private_dir(&root.join(GENERATIONS_DIR))?; let _writer = IndexWriterGuard::acquire(root)?; - let revision = archive_revision(archive)?; + let revision = archive_revision_for_parents(parents); let generation_id = format!( "gen-{}-{}-{}", &revision[..16], @@ -63,8 +78,7 @@ pub async fn index_semantic_live( .join(format!(".{generation_id}.staging")); let mut staging_guard = StagingGuard::new(staging.clone()); crate::paths::ensure_private_dir(&staging)?; - let written = write_semantic_documents_plain(archive, &staging)?; - let parents = archive.all_parent_chunks()?; + let written = write_semantic_documents_plain_for_parents(parents, &staging)?; let store = LocalVectorStore::open( &staging.join(STORE_DIR), usize::from(config.vector_dimension), @@ -72,58 +86,47 @@ pub async fn index_semantic_live( let mut embedder = create_embedder(config)?; let prior_generation = current_generation(root); let mut self_healed = false; - let prior_store = if full { - None + let prior_vectors = if full { + HashMap::new() } else { match prior_generation { - Ok(generation) => match validate_generation( - archive, - &generation, - embedder.id(), - config.vector_dimension, - ) { - Ok(()) => Some(LocalVectorStore::open_existing( - &generation.join(STORE_DIR), - usize::from(config.vector_dimension), - )?), - Err(Error::SemanticIndexStale(_) | Error::Embedding(_) | Error::Sql(_)) => { - self_healed = true; - None + Ok(generation) => { + match load_reusable_vectors(&generation, embedder.id(), config.vector_dimension) { + Ok((cursor, vectors)) => { + let expected = expected_content_pairs(parents); + let mut actual = vectors + .values() + .map(|stored| (stored.chunk_id.clone(), stored.content_hash.clone())) + .collect::>(); + actual.sort(); + self_healed = cursor.archive_revision != revision || actual != expected; + vectors + } + Err(Error::SemanticIndexStale(_) | Error::Embedding(_) | Error::Sql(_)) => { + self_healed = true; + HashMap::new() + } + Err(error) => return Err(error), } - Err(error) => return Err(error), - }, - Err(Error::SemanticIndexMissing | Error::SemanticIndexStale(_)) => None, + } + Err(Error::SemanticIndexMissing) => HashMap::new(), + Err(Error::SemanticIndexStale(_)) => { + self_healed = true; + HashMap::new() + } Err(error) => return Err(error), } }; let mut pending = Vec::new(); let mut reused_vectors = 0usize; - for parent in parents { - let hash = content_hash(&parent.text); - let heading_path = format!("{} / parent window", parent.chat_name); - if let Some(stored) = prior_store - .as_ref() - .and_then(|prior| prior.fetch(&parent.parent_id).transpose()) - .transpose()? - .filter(|stored| stored.content_hash == hash) - { - reused_vectors += 1; - store.upsert(&VectorUpsert { - chunk_id: parent.parent_id, - content_hash: hash, - seen_token: revision.clone(), - heading_path, - vector: stored.vector, - })?; - } else { - pending.push(PendingChunk { - chunk_id: parent.parent_id, - content_hash: hash, - seen_token: revision.clone(), - heading_path, - text: parent.text, - }); + for chunk in prepare_parent_chunks(parents, &prior_vectors, &revision) { + match chunk { + PreparedChunk::Reused(item) => { + reused_vectors += 1; + store.upsert(&item)?; + } + PreparedChunk::Pending(item) => pending.push(item), } } @@ -212,12 +215,8 @@ fn validate_generation( current_revision, cursor.archive_revision ))); } - let mut expected = archive - .all_parent_chunks()? - .into_iter() - .map(|parent| (parent.parent_id, content_hash(&parent.text))) - .collect::>(); - expected.sort(); + let parents = archive.all_parent_chunks()?; + let expected = expected_content_pairs(&parents); let actual = LocalVectorStore::open_existing(&generation.join(STORE_DIR), usize::from(dimension))? .content_pairs()?; @@ -238,6 +237,64 @@ struct PendingChunk { text: String, } +enum PreparedChunk { + Reused(VectorUpsert), + Pending(PendingChunk), +} + +fn prepare_parent_chunks( + parents: &[ParentChunk], + stored: &HashMap, + seen_token: &str, +) -> Vec { + parents + .par_iter() + .map(|parent| { + let hash = content_hash(&parent.text); + let heading_path = format!("{} / parent window", parent.chat_name); + match stored.get(&parent.parent_id) { + Some(stored) if stored.content_hash == hash => { + PreparedChunk::Reused(VectorUpsert { + chunk_id: parent.parent_id.clone(), + content_hash: hash, + seen_token: seen_token.to_string(), + heading_path, + vector: stored.vector.clone(), + }) + } + Some(_) | None => PreparedChunk::Pending(PendingChunk { + chunk_id: parent.parent_id.clone(), + content_hash: hash, + seen_token: seen_token.to_string(), + heading_path, + text: parent.text.clone(), + }), + } + }) + .collect() +} + +fn load_reusable_vectors( + generation: &Path, + embedder_id: &str, + dimension: u16, +) -> Result<(SemanticCursor, HashMap)> { + let cursor = load_cursor(generation)?; + validate_cursor(&cursor, embedder_id)?; + let store = + LocalVectorStore::open_existing(&generation.join(STORE_DIR), usize::from(dimension))?; + Ok((cursor, store.stored_vectors()?)) +} + +fn expected_content_pairs(parents: &[ParentChunk]) -> Vec<(String, String)> { + let mut pairs = parents + .par_iter() + .map(|parent| (parent.parent_id.clone(), content_hash(&parent.text))) + .collect::>(); + pairs.sort(); + pairs +} + fn embed_pending( store: &LocalVectorStore, embedder: &mut dyn super::embedder::SemanticEmbedder, diff --git a/src/semantic/mock.rs b/src/semantic/mock.rs index af52672..38fdaea 100644 --- a/src/semantic/mock.rs +++ b/src/semantic/mock.rs @@ -1,4 +1,10 @@ -use crate::{archive::Archive, search::hydrate_parent_hits, types::SearchHit, Error, Result}; +use crate::{ + archive::Archive, + search::hydrate_parent_hits, + types::{ParentChunk, SearchHit}, + Error, Result, +}; +use rayon::prelude::*; use std::{fs, io::Write, path::Path}; use super::{document_path, semantic_source_dir}; @@ -10,35 +16,47 @@ pub struct SemanticDocument { } pub fn planned_semantic_documents(archive: &Archive, dir: &Path) -> Result> { + let parents = archive.all_parent_chunks()?; + Ok(planned_semantic_documents_for_parents(&parents, dir)) +} + +pub fn planned_semantic_documents_for_parents( + parents: &[ParentChunk], + dir: &Path, +) -> Vec { let dir = semantic_source_dir(dir); - archive - .all_parent_chunks()? - .into_iter() - .map(|parent| { - Ok(SemanticDocument { - path: document_path(&dir, &parent.parent_id), - chunk_id: parent.parent_id, - }) + parents + .par_iter() + .map(|parent| SemanticDocument { + path: document_path(&dir, &parent.parent_id), + chunk_id: parent.parent_id.clone(), }) .collect() } pub fn write_semantic_documents(archive: &Archive, dir: &Path) -> Result { - let written = write_semantic_documents_plain(archive, dir)?; + let parents = archive.all_parent_chunks()?; + write_semantic_documents_for_parents(&parents, dir) +} + +pub fn write_semantic_documents_for_parents(parents: &[ParentChunk], dir: &Path) -> Result { + let written = write_semantic_documents_plain_for_parents(parents, dir)?; if let Some(root) = semantic_source_dir(dir).parent().and_then(Path::parent) { fs::write(root.join("INDEXED_WITH_MOCK"), b"mock\n").map_err(Error::Io)?; } Ok(written) } -pub(crate) fn write_semantic_documents_plain(archive: &Archive, dir: &Path) -> Result { +pub(crate) fn write_semantic_documents_plain_for_parents( + parents: &[ParentChunk], + dir: &Path, +) -> Result { let dir = semantic_source_dir(dir); if dir.exists() { fs::remove_dir_all(&dir).map_err(Error::Io)?; } crate::paths::ensure_private_dir(&dir)?; - let parents = archive.all_parent_chunks()?; - for parent in &parents { + parents.par_iter().try_for_each(|parent| -> Result<()> { let path = document_path(&dir, &parent.parent_id); let mut file = fs::File::create(&path).map_err(Error::Io)?; writeln!(file, "parent_id: {}", parent.parent_id).map_err(Error::Io)?; @@ -58,7 +76,8 @@ pub(crate) fn write_semantic_documents_plain(archive: &Archive, dir: &Path) -> R .map_err(Error::Io)?; writeln!(file, "---").map_err(Error::Io)?; writeln!(file, "{}", parent.text).map_err(Error::Io)?; - } + Ok(()) + })?; Ok(parents.len()) } diff --git a/src/semantic/store.rs b/src/semantic/store.rs index ab9d9f1..80e75dc 100644 --- a/src/semantic/store.rs +++ b/src/semantic/store.rs @@ -1,6 +1,6 @@ use crate::{Error, Result}; use rusqlite::{params, Connection, OpenFlags}; -use std::path::Path; +use std::{collections::HashMap, path::Path}; pub(crate) struct LocalVectorStore { conn: Connection, @@ -9,6 +9,7 @@ pub(crate) struct LocalVectorStore { #[derive(Debug, Clone)] pub(crate) struct StoredVector { + pub chunk_id: String, pub content_hash: String, pub vector: Vec, } @@ -51,22 +52,32 @@ impl LocalVectorStore { Ok(Self { conn, dimension }) } - pub(crate) fn fetch(&self, chunk_id: &str) -> Result> { + pub(crate) fn stored_vectors(&self) -> Result> { let mut statement = self .conn - .prepare("SELECT chunk_id, content_hash, vector FROM vectors WHERE chunk_id = ?1") + .prepare("SELECT chunk_id, content_hash, vector FROM vectors") .map_err(Error::Sql)?; - let mut rows = statement.query(params![chunk_id]).map_err(Error::Sql)?; - let Some(row) = rows.next().map_err(Error::Sql)? else { - return Ok(None); - }; - Ok(Some(StoredVector { - content_hash: row.get(1).map_err(Error::Sql)?, - vector: decode_vector( - &row.get::<_, Vec>(2).map_err(Error::Sql)?, - self.dimension, - )?, - })) + let rows = statement + .query_map([], |row| { + Ok(( + row.get::<_, String>(0)?, + row.get::<_, String>(1)?, + row.get::<_, Vec>(2)?, + )) + }) + .map_err(Error::Sql)?; + rows.map(|row| { + let (chunk_id, content_hash, vector) = row.map_err(Error::Sql)?; + Ok(( + chunk_id.clone(), + StoredVector { + chunk_id, + content_hash, + vector: decode_vector(&vector, self.dimension)?, + }, + )) + }) + .collect() } pub(crate) fn content_pairs(&self) -> Result> { @@ -222,3 +233,36 @@ fn dot(left: &[f32], right: &[f32]) -> f32 { .map(|(left, right)| left * right) .sum() } + +#[cfg(test)] +mod tests { + use super::*; + + fn vector(chunk_id: &str, content_hash: &str, values: [f32; 2]) -> VectorUpsert { + VectorUpsert { + chunk_id: chunk_id.to_string(), + content_hash: content_hash.to_string(), + seen_token: "seen".to_string(), + heading_path: "Synthetic room / parent window".to_string(), + vector: values.to_vec(), + } + } + + #[test] + fn loads_vector_metadata_in_one_pass() { + let dir = tempfile::tempdir().expect("tempdir"); + let store = LocalVectorStore::open(dir.path(), 2).expect("open store"); + store + .upsert(&vector("parent-b", "hash-b", [0.0, 1.0])) + .expect("upsert parent b"); + store + .upsert(&vector("parent-a", "hash-a", [1.0, 0.0])) + .expect("upsert parent a"); + + let stored = store.stored_vectors().expect("load stored vectors"); + + assert_eq!(stored.len(), 2); + assert_eq!(stored["parent-a"].content_hash, "hash-a"); + assert_eq!(stored["parent-b"].content_hash, "hash-b"); + } +} diff --git a/tests/cli_behaviors.rs b/tests/cli_behaviors.rs index a098aa7..b372013 100644 --- a/tests/cli_behaviors.rs +++ b/tests/cli_behaviors.rs @@ -1,4 +1,5 @@ use assert_cmd::Command; +use katok::archive::Archive; use predicates::prelude::*; fn fixture_path(name: &str) -> String { @@ -227,6 +228,55 @@ fn cli_reports_semantic_index_states_when_embedder_is_local_test_or_mocked() { .stdout(predicate::str::contains("\"full\": true")); } +#[test] +fn cli_index_counts_candidates_without_loading_chunk_bodies() { + let dir = tempfile::tempdir().expect("create tempdir"); + let data_dir = dir.path(); + let fixture = fixture_path("replies.jsonl"); + + Command::cargo_bin("katok") + .expect("katok binary") + .args([ + "--data-dir", + data_dir.to_str().expect("utf8 path"), + "sync", + "--source", + "fixture", + &fixture, + "--json", + ]) + .assert() + .success(); + + let archive_path = data_dir.join("archive.sqlite3"); + let archive = Archive::open(&archive_path).expect("open archive"); + let expected = archive.chunk_count().expect("count chunks"); + archive + .connection() + .execute( + "UPDATE chunks SET chat_name = X'80' + WHERE rowid = (SELECT rowid FROM chunks ORDER BY rowid LIMIT 1)", + [], + ) + .expect("make chunk body unreadable as utf8"); + drop(archive); + + Command::cargo_bin("katok") + .expect("katok binary") + .args([ + "--data-dir", + data_dir.to_str().expect("utf8 path"), + "index", + "--dry-run", + "--json", + ]) + .assert() + .success() + .stdout(predicate::str::contains(format!( + "\"candidate_chunks\": {expected}" + ))); +} + #[test] fn cli_lists_gap_chunks_and_applies_chunk_output_flags() { let dir = tempfile::tempdir().expect("create tempdir"); diff --git a/tests/live_semantic.rs b/tests/live_semantic.rs index fb1ad82..8bb6c12 100644 --- a/tests/live_semantic.rs +++ b/tests/live_semantic.rs @@ -64,6 +64,20 @@ fn live_semantic_cli_indexes_local_embeddings_and_searches_without_endpoint() { .stdout(predicate::str::contains("\"vectorstore\": \"local\"")) .stdout(predicate::str::contains("\"embedding_calls\": 1")); + Command::cargo_bin("katok") + .expect("katok binary") + .env("KATOK_EMBEDDER", "local-test") + .args([ + "--data-dir", + data_dir.to_str().expect("utf8 data"), + "index", + "--json", + ]) + .assert() + .success() + .stdout(predicate::str::contains("\"embedding_calls\": 0")) + .stdout(predicate::str::contains("\"embedded_texts\": 0")); + Command::cargo_bin("katok") .expect("katok binary") .env("KATOK_EMBEDDER", "local-test") diff --git a/tests/planned_behaviors.rs b/tests/planned_behaviors.rs index 7a77a85..973acf1 100644 --- a/tests/planned_behaviors.rs +++ b/tests/planned_behaviors.rs @@ -3,10 +3,37 @@ use katok::{ chunking::rebuild_chunks, fixture::read_fixture, search::{bm25_search, keyword_search}, - semantic::{semantic_search, write_semantic_documents}, + semantic::{ + planned_semantic_documents_for_parents, semantic_search, write_semantic_documents, + write_semantic_documents_for_parents, + }, types::RawMessage, }; +#[test] +fn semantic_documents_accept_a_loaded_parent_snapshot() { + let dir = tempfile::tempdir().expect("tempdir"); + let archive_path = dir.path().join("archive.sqlite3"); + let archive = Archive::open(&archive_path).expect("open archive"); + let fixture_path = + std::path::Path::new(env!("CARGO_MANIFEST_DIR")).join("tests/fixtures/kakao/replies.jsonl"); + let messages = read_fixture(&fixture_path).expect("read fixture"); + + archive.sync_messages(&messages).expect("sync messages"); + rebuild_chunks(&archive).expect("rebuild chunks"); + let parents = archive.all_parent_chunks().expect("load parents"); + let semantic_dir = dir.path().join("semantic"); + + let planned = planned_semantic_documents_for_parents(&parents, &semantic_dir); + let written = write_semantic_documents_for_parents(&parents, &semantic_dir) + .expect("write semantic documents"); + + assert_eq!(planned.len(), parents.len()); + assert_eq!(written, parents.len()); + assert_eq!(planned[0].chunk_id, parents[0].parent_id); + assert!(planned.iter().all(|document| document.path.exists())); +} + #[test] fn same_sender_reply_and_search_behaviors_when_fixture_is_indexed() { let dir = tempfile::tempdir().expect("create tempdir"); diff --git a/tests/search_index_recovery.rs b/tests/search_index_recovery.rs index af4734b..d4dd86e 100644 --- a/tests/search_index_recovery.rs +++ b/tests/search_index_recovery.rs @@ -117,6 +117,53 @@ fn index_self_heals_an_orphan_generation_and_full_never_reuses_vectors() { assert_eq!(generations.len(), 1); } +#[test] +fn stale_but_structurally_valid_generation_reuses_unchanged_vectors() { + let dir = tempfile::tempdir().expect("tempdir"); + let data = dir.path().join("data"); + run_json( + &data, + &[ + "sync", + "--source", + "fixture", + fixture("replies.jsonl").to_str().unwrap(), + "--json", + ], + ); + let first = run_json(&data, &["index", "--json"]); + assert!(first["embedded_texts"].as_u64().unwrap() > 0); + + let extra = dir.path().join("extra.jsonl"); + std::fs::write( + &extra, + concat!( + "{\"account_hash\":\"acct\",\"chat_id\":\"chat-new\",\"chat_name\":\"Synthetic new room\",", + "\"chat_type\":\"group\",\"message_id\":\"new-1\",\"sender_id\":\"sender-new\",", + "\"sender_nickname\":\"Tester\",\"timestamp\":\"2026-01-02T00:00:00Z\",", + "\"text\":\"brand new semantic material\",\"message_type\":\"text\",\"reply_to_message_id\":null}\n" + ), + ) + .expect("write extra fixture"); + run_json( + &data, + &[ + "sync", + "--source", + "fixture", + extra.to_str().unwrap(), + "--json", + ], + ); + + let second = run_json(&data, &["index", "--json"]); + assert!(second["reused_vectors"].as_u64().unwrap() > 0); + assert!(second["embedded_texts"].as_u64().unwrap() > 0); + assert!( + second["embedded_texts"].as_u64().unwrap() < second["written_documents"].as_u64().unwrap() + ); +} + #[test] fn failed_rebuild_keeps_the_committed_generation_and_returns_json_error() { let dir = tempfile::tempdir().expect("tempdir");