1use 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
23pub 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 {
46 #[arg(long, value_name = "NAME", conflicts_with = "batch")]
48 pub name: Option<String>,
49 #[arg(long, default_value_t = false, conflicts_with = "name")]
51 pub batch: bool,
52 #[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 pub namespace: Option<String>,
58 #[arg(long, default_value_t = false)]
61 pub dry_run: bool,
62 #[arg(long, value_enum, default_value_t = JsonOutputFormat::Json)]
64 pub format: JsonOutputFormat,
65 #[arg(long, hide = true, help = "No-op; JSON is always emitted on stdout")]
67 pub json: bool,
68 #[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 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 elapsed_ms: u64,
100}
101
102pub 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 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 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
179fn 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 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 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 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 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 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 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 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
380fn 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 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 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 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 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 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 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}