sqlite_graphrag/commands/
split_body.rs

1//! Handler for the `split-body` CLI subcommand (v1.1.03, GAP-V8).
2//!
3//! Splits memories whose body exceeds a character threshold into N child
4//! memories, reusing the lossless section-based chunker
5//! [`crate::chunking::split_body_by_sections`]. The original memory is NOT
6//! soft-deleted: its `metadata` is annotated with `superseded_by_split` so
7//! history and searchability are preserved. Each child memory gets a
8//! canonical `replaces` graph relationship pointing back to the original.
9
10use crate::chunking::split_body_by_sections;
11use crate::constants::{DEFAULT_RELATION_WEIGHT, MAX_MEMORY_NAME_LEN};
12use crate::entity_type::EntityType;
13use crate::errors::AppError;
14use crate::i18n::errors_msg;
15use crate::output::{self, JsonOutputFormat};
16use crate::paths::AppPaths;
17use crate::storage::connection::open_rw;
18use crate::storage::entities::{self, NewEntity};
19use crate::storage::{memories, versions};
20use rusqlite::params;
21use serde::Serialize;
22
23/// Character threshold below which a memory is left untouched.
24pub const DEFAULT_SPLIT_THRESHOLD: usize = 25_000;
25
26#[derive(clap::Args)]
27#[command(
28    about = "Split oversized memories into N child memories at Markdown section boundaries",
29    after_long_help = "EXAMPLES:\n  \
30        # Split a single memory by name using the default 25k char threshold\n  \
31        sqlite-graphrag split-body --name big-doc\n\n  \
32        # Preview the split without writing\n  \
33        sqlite-graphrag split-body --name big-doc --dry-run\n\n  \
34        # Batch-split every oversized memory in the current namespace\n  \
35        sqlite-graphrag split-body --batch\n\n  \
36        # Batch-split with a custom threshold and namespace\n  \
37        sqlite-graphrag split-body --batch --threshold 50000 --namespace my-project\n\n  \
38        NOTE:\n  \
39            The original memory is NEVER soft-deleted. It is annotated with\n  \
40            metadata.superseded_by_split=true and metadata.split_into=[child names]\n  \
41            so its history and searchability are preserved. Each child memory\n  \
42            gets a canonical `replaces` graph relationship to the original."
43)]
44pub struct SplitBodyArgs {
45    /// Memory name to split (single mode). Mutually exclusive with `--batch`.
46    #[arg(long, value_name = "NAME", conflicts_with = "batch")]
47    pub name: Option<String>,
48    /// Batch mode: split every active memory whose `LENGTH(body) > threshold`.
49    #[arg(long, default_value_t = false, conflicts_with = "name")]
50    pub batch: bool,
51    /// Only memories with `LENGTH(body) > threshold` are split. Default 25000 chars.
52    #[arg(long, value_name = "N", default_value_t = DEFAULT_SPLIT_THRESHOLD)]
53    pub threshold: usize,
54    #[arg(
55        long,
56        help = "Namespace (env: SQLITE_GRAPHRAG_NAMESPACE, default: global)"
57    )]
58    pub namespace: Option<String>,
59    /// Preview the split(s) without writing. Emits the planned child names and
60    /// body lengths.
61    #[arg(long, default_value_t = false)]
62    pub dry_run: bool,
63    /// Output format.
64    #[arg(long, value_enum, default_value_t = JsonOutputFormat::Json)]
65    pub format: JsonOutputFormat,
66    #[arg(long, hide = true, help = "No-op; JSON is always emitted on stdout")]
67    pub json: bool,
68    #[arg(long, env = "SQLITE_GRAPHRAG_DB_PATH")]
69    pub db: Option<String>,
70}
71
72#[derive(Serialize)]
73struct ChildPlan {
74    name: String,
75    body_len: usize,
76}
77
78#[derive(Serialize)]
79struct SplitResult {
80    original: String,
81    #[serde(skip_serializing_if = "Option::is_none")]
82    original_id: Option<i64>,
83    children: Vec<ChildPlan>,
84    threshold: usize,
85    dry_run: bool,
86    /// Total execution time in milliseconds.
87    elapsed_ms: u64,
88}
89
90#[derive(Serialize)]
91struct BatchResult {
92    namespace: String,
93    threshold: usize,
94    dry_run: bool,
95    split: Vec<SplitResult>,
96    skipped: usize,
97    /// Total execution time in milliseconds.
98    elapsed_ms: u64,
99}
100
101/// Validates the parsed args and dispatches to single or batch mode.
102pub fn run(args: SplitBodyArgs) -> Result<(), AppError> {
103    let inicio = std::time::Instant::now();
104    let _ = args.format;
105
106    if !args.batch && args.name.is_none() {
107        return Err(AppError::Validation(
108            "either --name <NAME> or --batch is required".to_string(),
109        ));
110    }
111    if args.threshold == 0 {
112        return Err(AppError::Validation(
113            "--threshold must be greater than zero".to_string(),
114        ));
115    }
116
117    let namespace = crate::namespace::resolve_namespace(args.namespace.as_deref())?;
118    let paths = AppPaths::resolve(args.db.as_deref())?;
119    crate::storage::connection::ensure_db_ready(&paths)?;
120
121    if args.batch {
122        run_batch(&paths.db, &namespace, args.threshold, args.dry_run, inicio)
123    } else if let Some(name) = args.name.as_deref() {
124        let result = split_single(&paths.db, &namespace, name, args.threshold, args.dry_run)?;
125        output::emit_json(&result)?;
126        Ok(())
127    } else {
128        // Unreachable: validated at the top of run().
129        Err(AppError::Validation(
130            "either --name <NAME> or --batch is required".to_string(),
131        ))
132    }
133}
134
135fn run_batch(
136    db_path: &std::path::Path,
137    namespace: &str,
138    threshold: usize,
139    dry_run: bool,
140    inicio: std::time::Instant,
141) -> Result<(), AppError> {
142    let conn = open_rw(db_path)?;
143    let mut stmt =
144        conn.prepare("SELECT name FROM memories WHERE namespace = ?1 AND deleted_at IS NULL AND LENGTH(body) > ?2")?;
145    let mut rows = stmt.query(params![namespace, threshold as i64])?;
146    let mut names: Vec<String> = Vec::new();
147    while let Some(row) = rows.next()? {
148        names.push(row.get::<_, String>(0)?);
149    }
150    drop(rows);
151    drop(stmt);
152    drop(conn);
153
154    let mut split: Vec<SplitResult> = Vec::new();
155    let mut skipped = 0usize;
156    for name in &names {
157        match split_single(db_path, namespace, name, threshold, dry_run) {
158            Ok(r) => split.push(r),
159            Err(AppError::NotFound(_)) => {
160                // Race: memory was removed between scan and split. Skip quietly.
161                skipped += 1;
162            }
163            Err(e) => return Err(e),
164        }
165    }
166
167    output::emit_json(&BatchResult {
168        namespace: namespace.to_string(),
169        threshold,
170        dry_run,
171        split,
172        skipped,
173        elapsed_ms: inicio.elapsed().as_millis() as u64,
174    })?;
175    Ok(())
176}
177
178/// Splits a single memory by name into N child memories.
179///
180/// Reuses the lossless [`split_body_by_sections`] chunker so concatenating
181/// every child body reproduces the original body. The original memory is kept
182/// active and its `metadata` is annotated with `superseded_by_split` and the
183/// list of child names. Each child is inserted as a new memory (auto-embedded
184/// downstream by the remember path or `enrich --operation re-embed`) and linked
185/// to the original entity via a canonical `replaces` relationship.
186fn split_single(
187    db_path: &std::path::Path,
188    namespace: &str,
189    name: &str,
190    threshold: usize,
191    dry_run: bool,
192) -> Result<SplitResult, AppError> {
193    let start = std::time::Instant::now();
194
195    let mut conn = open_rw(db_path)?;
196    let row = memories::read_by_name(&conn, namespace, name)?
197        .ok_or_else(|| AppError::NotFound(errors_msg::memory_not_found(name, namespace)))?;
198
199    if row.body.len() <= threshold {
200        return Ok(SplitResult {
201            original: name.to_string(),
202            original_id: Some(row.id),
203            children: Vec::new(),
204            threshold,
205            dry_run,
206            elapsed_ms: start.elapsed().as_millis() as u64,
207        });
208    }
209
210    let partitions = split_body_by_sections(&row.body);
211    if partitions.len() < 2 {
212        // Body crosses the threshold but the chunker kept it as a single
213        // partition (e.g. under byte/chunk/token budgets). Nothing to split.
214        return Ok(SplitResult {
215            original: name.to_string(),
216            original_id: Some(row.id),
217            children: Vec::new(),
218            threshold,
219            dry_run,
220            elapsed_ms: start.elapsed().as_millis() as u64,
221        });
222    }
223
224    // Plan child names first so we can surface them in dry-run and metadata.
225    let child_names: Vec<String> = (0..partitions.len())
226        .map(|i| format!("{name}-part-{i}"))
227        .collect();
228
229    for child in &child_names {
230        validate_child_name(child, name)?;
231    }
232
233    let children_plan: Vec<ChildPlan> = partitions
234        .iter()
235        .zip(child_names.iter())
236        .map(|(body, n)| ChildPlan {
237            name: n.clone(),
238            body_len: body.len(),
239        })
240        .collect();
241
242    if dry_run {
243        return Ok(SplitResult {
244            original: name.to_string(),
245            original_id: Some(row.id),
246            children: children_plan,
247            threshold,
248            dry_run: true,
249            elapsed_ms: start.elapsed().as_millis() as u64,
250        });
251    }
252
253    // Parse existing metadata (stored as TEXT JSON) and merge our markers.
254    let mut metadata: serde_json::Value =
255        serde_json::from_str(&row.metadata).unwrap_or_else(|_| serde_json::json!({}));
256    let map = match metadata.as_object_mut() {
257        Some(m) => m,
258        None => {
259            return Err(AppError::Internal(anyhow::anyhow!(
260                "metadata for memory '{name}' is not a JSON object"
261            )));
262        }
263    };
264    map.insert("superseded_by_split".to_string(), serde_json::json!(true));
265    map.insert(
266        "split_into".to_string(),
267        serde_json::to_value(&child_names)?,
268    );
269    let metadata_str = serde_json::to_string(&metadata)?;
270
271    let tx = conn.transaction_with_behavior(rusqlite::TransactionBehavior::Immediate)?;
272
273    // Insert each child as a new memory of type `document` (auto-embedded by
274    // downstream paths). We do NOT call the embedding pipeline here to keep
275    // the split atomic and side-effect-free; re-embed can be run separately.
276    for (body_part, child_name) in partitions.iter().zip(child_names.iter()) {
277        let child_hash = blake3_hash(body_part);
278        let new_memory = memories::NewMemory {
279            namespace: namespace.to_string(),
280            name: child_name.clone(),
281            memory_type: "document".to_string(),
282            description: format!("Partição de '{name}' (split-body)"),
283            body: body_part.clone(),
284            body_hash: child_hash,
285            session_id: row.session_id.clone(),
286            source: "system".to_string(),
287            metadata: serde_json::json!({ "split_from": name }),
288        };
289        memories::insert(&tx, &new_memory)?;
290    }
291
292    // Annotate the original memory metadata (no soft-delete: history preserved).
293    let affected = tx.execute(
294        "UPDATE memories SET metadata = ?2 WHERE id = ?1 AND deleted_at IS NULL",
295        params![row.id, metadata_str],
296    )?;
297    if affected == 0 {
298        return Err(AppError::Conflict(format!(
299            "memory '{name}' was modified concurrently; retry"
300        )));
301    }
302
303    // Record a version snapshot of the metadata annotation.
304    let next_v = versions::next_version(&tx, row.id)?;
305    versions::insert_version(
306        &tx,
307        row.id,
308        next_v,
309        &row.name,
310        &row.memory_type,
311        &row.description,
312        &row.body,
313        &metadata_str,
314        None,
315        "edit",
316    )?;
317
318    // Graph: each child entity `replaces` the original. Auto-create both
319    // endpoints so split-body is self-contained (no dependency on prior NER).
320    let original_entity_name = crate::parsers::normalize_entity_name(name);
321    let original_entity = NewEntity {
322        name: original_entity_name.clone(),
323        entity_type: EntityType::Memory,
324        description: None,
325    };
326    let original_entity_id = entities::upsert_entity(&tx, namespace, &original_entity)?;
327
328    for child_name in &child_names {
329        let child_entity_name = crate::parsers::normalize_entity_name(child_name);
330        let child_entity = NewEntity {
331            name: child_entity_name.clone(),
332            entity_type: EntityType::Memory,
333            description: None,
334        };
335        let child_entity_id = entities::upsert_entity(&tx, namespace, &child_entity)?;
336        let (_, was_created) = entities::create_or_fetch_relationship(
337            &tx,
338            namespace,
339            child_entity_id,
340            original_entity_id,
341            "replaces",
342            DEFAULT_RELATION_WEIGHT,
343            None,
344        )?;
345        if was_created {
346            entities::recalculate_degree(&tx, child_entity_id)?;
347            entities::recalculate_degree(&tx, original_entity_id)?;
348        }
349    }
350
351    tx.commit()?;
352    conn.execute_batch("PRAGMA wal_checkpoint(TRUNCATE);")?;
353
354    Ok(SplitResult {
355        original: name.to_string(),
356        original_id: Some(row.id),
357        children: children_plan,
358        threshold,
359        dry_run: false,
360        elapsed_ms: start.elapsed().as_millis() as u64,
361    })
362}
363
364fn validate_child_name(child: &str, parent: &str) -> Result<(), AppError> {
365    if child.is_empty() || child.len() > MAX_MEMORY_NAME_LEN {
366        return Err(AppError::Validation(format!(
367            "child name '{child}' derived from '{parent}' exceeds MAX_MEMORY_NAME_LEN ({MAX_MEMORY_NAME_LEN})"
368        )));
369    }
370    let slug_re = crate::constants::name_slug_regex();
371    if !slug_re.is_match(child) {
372        return Err(AppError::Validation(format!(
373            "child name '{child}' is not kebab-case ASCII; rename the parent memory"
374        )));
375    }
376    Ok(())
377}
378
379/// Computes the BLAKE3 hex digest of `body`, matching the project convention.
380fn blake3_hash(body: &str) -> String {
381    use blake3::Hasher;
382    let mut hasher = Hasher::new();
383    hasher.update(body.as_bytes());
384    hasher.finalize().to_hex().to_string()
385}
386
387#[cfg(test)]
388mod tests {
389    use super::*;
390    use crate::storage::memories::{insert, NewMemory};
391    use rusqlite::Connection;
392    use tempfile::TempDir;
393
394    /// `split_single` needs a `db_path` on disk so it can reopen a connection
395    /// for the transaction. This helper inserts the oversized memory and returns
396    /// the temp dir + db path.
397    fn setup_with_memory(name: &str, body_len: usize) -> (TempDir, std::path::PathBuf) {
398        crate::storage::connection::register_vec_extension();
399        let dir = TempDir::new().unwrap();
400        let db_path = dir.path().join("test.db");
401        let mut conn = Connection::open(&db_path).unwrap();
402        crate::migrations::runner().run(&mut conn).unwrap();
403        let body = body_with_headers(body_len);
404        insert(
405            &conn,
406            &NewMemory {
407                namespace: "global".to_string(),
408                name: name.to_string(),
409                memory_type: "document".to_string(),
410                description: "desc".to_string(),
411                body,
412                body_hash: format!("hash-{name}"),
413                session_id: None,
414                source: "agent".to_string(),
415                metadata: serde_json::json!({}),
416            },
417        )
418        .unwrap();
419        drop(conn);
420        (dir, db_path)
421    }
422
423    /// Builds a Markdown body with many ATX sections so the section splitter
424    /// produces multiple partitions well above the 80 KiB byte budget.
425    fn body_with_headers(total_bytes: usize) -> String {
426        let mut body = String::new();
427        let mut i = 0;
428        while body.len() < total_bytes {
429            body.push_str(&format!(
430                "# Section {i}\n\n{}\n\n",
431                "body content text here. ".repeat(200)
432            ));
433            i += 1;
434        }
435        body
436    }
437
438    #[test]
439    fn split_body_divides_long_memory_into_parts() {
440        let (_dir, db_path) = setup_with_memory("big-doc", 100_000);
441        let result = split_single(&db_path, "global", "big-doc", 25_000, false).unwrap();
442        assert!(
443            result.children.len() >= 2,
444            "expected 2+ children, got {}",
445            result.children.len()
446        );
447        // Verify each child memory was actually persisted.
448        let conn = Connection::open(&db_path).unwrap();
449        for child in &result.children {
450            let row = crate::storage::memories::read_by_name(&conn, "global", &child.name)
451                .unwrap()
452                .expect("child memory should exist");
453            assert_eq!(row.memory_type, "document");
454            assert!(
455                row.body.len() <= crate::constants::AUTOSPLIT_PARTITION_MAX_BYTES,
456                "child body {} exceeds partition budget",
457                row.body.len()
458            );
459        }
460    }
461
462    #[test]
463    fn split_body_marks_original_as_superseded() {
464        let (_dir, db_path) = setup_with_memory("super-test", 100_000);
465        split_single(&db_path, "global", "super-test", 25_000, false).unwrap();
466
467        let conn = Connection::open(&db_path).unwrap();
468        let row = crate::storage::memories::read_by_name(&conn, "global", "super-test")
469            .unwrap()
470            .expect("original memory must remain");
471        let metadata: serde_json::Value = serde_json::from_str(&row.metadata).unwrap();
472        assert_eq!(
473            metadata
474                .get("superseded_by_split")
475                .and_then(|v| v.as_bool()),
476            Some(true),
477            "metadata.superseded_by_split must be true after split"
478        );
479        let split_into = metadata
480            .get("split_into")
481            .and_then(|v| v.as_array())
482            .expect("metadata.split_into must be an array");
483        assert!(
484            split_into.len() >= 2,
485            "metadata.split_into must list 2+ children"
486        );
487    }
488
489    #[test]
490    fn split_body_creates_replaces_relations() {
491        let (_dir, db_path) = setup_with_memory("rel-source", 100_000);
492        let result = split_single(&db_path, "global", "rel-source", 25_000, false).unwrap();
493
494        let conn = Connection::open(&db_path).unwrap();
495        let original_entity = crate::parsers::normalize_entity_name("rel-source");
496        let original_id =
497            crate::storage::entities::find_entity_id(&conn, "global", &original_entity)
498                .unwrap()
499                .expect("original entity must exist");
500
501        let mut count = 0;
502        for child in &result.children {
503            let child_entity = crate::parsers::normalize_entity_name(&child.name);
504            let child_id = crate::storage::entities::find_entity_id(&conn, "global", &child_entity)
505                .unwrap()
506                .expect("child entity must exist");
507            let exists: bool = conn
508                .query_row(
509                    "SELECT EXISTS(SELECT 1 FROM relationships
510                     WHERE source_id = ?1 AND target_id = ?2 AND relation = 'replaces')",
511                    rusqlite::params![child_id, original_id],
512                    |r| r.get(0),
513                )
514                .unwrap();
515            assert!(
516                exists,
517                "child '{}' must `replaces` the original",
518                child.name
519            );
520            count += 1;
521        }
522        assert!(count >= 2, "expected 2+ replaces relations, got {count}");
523    }
524
525    #[test]
526    fn split_body_preserves_history() {
527        let (_dir, db_path) = setup_with_memory("hist-keep", 100_000);
528        split_single(&db_path, "global", "hist-keep", 25_000, false).unwrap();
529
530        let conn = Connection::open(&db_path).unwrap();
531        // read_by_name filters deleted_at IS NULL: original must remain visible.
532        let row = crate::storage::memories::read_by_name(&conn, "global", "hist-keep")
533            .unwrap()
534            .expect("original must remain active (not soft-deleted)");
535        assert!(row.deleted_at.is_none(), "deleted_at must stay NULL");
536    }
537
538    #[test]
539    fn dry_run_writes_nothing() {
540        let (_dir, db_path) = setup_with_memory("dry-only", 100_000);
541        let result = split_single(&db_path, "global", "dry-only", 25_000, true).unwrap();
542        assert!(result.dry_run);
543        assert!(result.children.len() >= 2);
544
545        let conn = Connection::open(&db_path).unwrap();
546        // No child should have been persisted.
547        for child in &result.children {
548            assert!(
549                crate::storage::memories::read_by_name(&conn, "global", &child.name)
550                    .unwrap()
551                    .is_none(),
552                "dry-run must not persist child '{}'",
553                child.name
554            );
555        }
556        let row = crate::storage::memories::read_by_name(&conn, "global", "dry-only")
557            .unwrap()
558            .expect("original must exist");
559        let metadata: serde_json::Value = serde_json::from_str(&row.metadata).unwrap();
560        assert!(
561            metadata.get("superseded_by_split").is_none(),
562            "dry-run must not annotate metadata"
563        );
564    }
565
566    #[test]
567    fn below_threshold_returns_empty_children() {
568        let (_dir, db_path) = setup_with_memory("small-doc", 1_000);
569        let result = split_single(&db_path, "global", "small-doc", 25_000, false).unwrap();
570        assert!(result.children.is_empty());
571    }
572
573    #[test]
574    fn validate_child_name_rejects_too_long_parent() {
575        // A parent name so long that even "{name}-part-0" exceeds MAX_MEMORY_NAME_LEN.
576        let long_parent = "a".repeat(MAX_MEMORY_NAME_LEN);
577        let child = format!("{long_parent}-part-0");
578        let err = validate_child_name(&child, &long_parent).unwrap_err();
579        assert!(matches!(err, AppError::Validation(_)));
580    }
581}
sqlite_graphrag/commands/split_body.rs

sqlite_graphrag/commands/
split_body.rs
