Skip to main content

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