Skip to main content

sqlite_graphrag/commands/enrich/
mod.rs

1// v1.0.97: modularised into queue.rs, scan.rs, postprocess.rs, extraction.rs.
2// See ADR-0056 (closes the ADR-0046 "Known Tech Debt (v1.0.89+)" item).
3// v1.1.8 Wave C1: further split into schemas/args/events/run modules (≤800 LOC).
4
5//! Handler for the `enrich` CLI subcommand (GAP-14 + GAP-18).
6//!
7//! Enriches the knowledge graph by running LLM-powered analysis over memories
8//! and entities that are missing key structural data. Operations are:
9//!
10//! - `memory-bindings`: memories without `memory_entities` rows get entity extraction
11//! - `entity-descriptions`: entities with NULL/empty descriptions get LLM descriptions
12//! - `body-enrich`: memories with short bodies get expanded by the LLM (GAP-18)
13//! - `re-embed`: memories without a vector row get re-embedded without rewriting body
14//!
15//! Architecture mirrors `ingest_claude.rs`: SCAN → JUDGE (LLM) → PERSIST, with a
16//! SQLite queue DB derived next to `--db` (GAP-SG-64) for resume/retry support.
17// Workload: Subprocess I/O-bound (claude/codex API calls with network wait)
18//!
19//! # DRY note
20//!
21//! v1.0.97: `claude_runner.rs` now hosts the shared Claude invocation helpers
22//! (`run_claude`, `parse_claude_output`, `spawn_with_memory_limit`).
23//! GAP-SG-121: enrich vs ingest queue table shapes are different products
24//! (item_key/operation vs file_path); only sidecar WAL/busy pragmas are shared
25//! via [`crate::pragmas::apply_sidecar_queue_pragmas`].
26
27mod args;
28mod drain_parallel;
29mod drain_serial;
30mod events;
31mod extraction;
32// extraction_providers is a child of extraction.rs (#[path])
33// extraction_ops_* are submodules of extraction.rs (#[path])
34mod postprocess;
35mod predicates;
36mod prompts;
37mod quality_sample;
38mod queue;
39mod queue_ops;
40mod run;
41mod scan;
42mod scan_ec;
43mod scheduler;
44mod schemas;
45mod status;
46
47pub(crate) const DEFAULT_RATE_LIMIT_WAIT: u64 = 60;
48pub(crate) const DEFAULT_BODY_ENRICH_MIN_CHARS: usize = 500;
49pub(crate) const DEFAULT_BODY_ENRICH_MAX_CHARS: usize = 2000;
50
51pub use args::{EnrichArgs, EnrichMode, EnrichOperation, ReEmbedTarget};
52pub use queue::{cleanup_queue_entry, DeadItem, DeadSummary, EnrichStatus, WaitingItem};
53pub use run::run;
54
55use crate::errors::AppError;
56use queue::{enqueue_candidate_with_priority, open_queue_db, PRIORITY_HOT};
57
58/// GAP-CLI-PRIO-02: enqueue entity-descriptions for a hot set of entity names
59/// into the enrich sidecar queue with elevated priority.
60pub fn enqueue_priority_entity_descriptions(
61    paths: &crate::paths::AppPaths,
62    namespace: &str,
63    entity_names: &[String],
64) -> Result<usize, AppError> {
65    let _ = namespace; // queue keys are entity names; namespace is scoped by DB path
66    let queue_path = crate::paths::sidecar_path(&paths.db, ".enrich-queue.sqlite");
67    let queue = open_queue_db(&queue_path)?;
68    let mut n = 0usize;
69    for name in entity_names {
70        enqueue_candidate_with_priority(&queue, name, "entity", "EntityDescriptions", PRIORITY_HOT);
71        n += 1;
72    }
73    Ok(n)
74}
75
76#[cfg(test)]
77mod tests {
78    use super::events::{
79        enrich_operation_cli_name, is_sqlite_interrupt, scan_operation_with_deadline,
80    };
81    use super::schemas::{BINDINGS_SCHEMA, BODY_ENRICH_SCHEMA, ENTITY_DESCRIPTION_SCHEMA};
82    use super::*;
83    use crate::errors::AppError;
84    use rusqlite::{Connection, ErrorCode};
85    use std::time::{Duration, Instant};
86
87    #[test]
88    fn bindings_schema_is_valid_json() {
89        let _: serde_json::Value =
90            serde_json::from_str(BINDINGS_SCHEMA).expect("BINDINGS_SCHEMA must be valid JSON");
91    }
92
93    #[test]
94    fn entity_description_schema_is_valid_json() {
95        let _: serde_json::Value = serde_json::from_str(ENTITY_DESCRIPTION_SCHEMA)
96            .expect("ENTITY_DESCRIPTION_SCHEMA must be valid JSON");
97    }
98
99    #[test]
100    fn body_enrich_schema_is_valid_json() {
101        let _: serde_json::Value = serde_json::from_str(BODY_ENRICH_SCHEMA)
102            .expect("BODY_ENRICH_SCHEMA must be valid JSON");
103    }
104
105    // v1.1.06 — GAP-ENTITY-CONNECT-SCAN-CARTESIAN observability + interrupt
106
107    #[test]
108    fn enrich_operation_cli_name_pair_ops_are_kebab_case() {
109        assert_eq!(
110            enrich_operation_cli_name(&EnrichOperation::EntityConnect),
111            "entity-connect"
112        );
113        assert_eq!(
114            enrich_operation_cli_name(&EnrichOperation::CrossDomainBridges),
115            "cross-domain-bridges"
116        );
117        assert_eq!(
118            enrich_operation_cli_name(&EnrichOperation::EntityDescriptions),
119            "entity-descriptions"
120        );
121    }
122
123    #[test]
124    fn is_sqlite_interrupt_detects_operation_interrupted() {
125        let ffi_err = rusqlite::ffi::Error {
126            code: ErrorCode::OperationInterrupted,
127            extended_code: 9,
128        };
129        let err = rusqlite::Error::SqliteFailure(ffi_err, Some("interrupted".into()));
130        assert!(is_sqlite_interrupt(&err));
131
132        let busy = rusqlite::ffi::Error {
133            code: ErrorCode::DatabaseBusy,
134            extended_code: 5,
135        };
136        let busy_err = rusqlite::Error::SqliteFailure(busy, None);
137        assert!(!is_sqlite_interrupt(&busy_err));
138    }
139
140    #[test]
141    fn scan_deadline_already_elapsed_returns_timeout() {
142        // Past deadline must fail fast without running SQL (exit path → Timeout).
143        use clap::Parser;
144        let cli = crate::cli::Cli::try_parse_from([
145            "sqlite-graphrag",
146            "enrich",
147            "--operation",
148            "entity-connect",
149            "--mode",
150            "openrouter",
151            "--openrouter-model",
152            "test/model",
153            "--dry-run",
154            "--limit",
155            "1",
156        ])
157        .expect("parse enrich args");
158        let Some(crate::cli::Commands::Enrich(args)) = cli.command else {
159            panic!("expected Commands::Enrich");
160        };
161        let conn = Connection::open_in_memory().unwrap();
162        let past = Instant::now() - Duration::from_secs(1);
163        let err = scan_operation_with_deadline(&conn, "global", &args, Some(past))
164            .expect_err("elapsed deadline must Timeout");
165        match err {
166            AppError::Timeout { .. } => {}
167            other => panic!("expected Timeout, got {other:?}"),
168        }
169    }
170
171    #[test]
172    fn interrupt_handle_maps_long_query_to_sqlite_interrupt() {
173        // Live SQLite: watchdog interrupt aborts a recursive CTE (same mechanism
174        // as scan_operation_with_deadline). Confirms rusqlite InterruptHandle.
175        let conn = Connection::open_in_memory().unwrap();
176        let handle = conn.get_interrupt_handle();
177        let stop = std::sync::Arc::new(std::sync::atomic::AtomicBool::new(false));
178        let stop_w = std::sync::Arc::clone(&stop);
179        let watchdog = std::thread::spawn(move || {
180            std::thread::sleep(Duration::from_millis(30));
181            if !stop_w.load(std::sync::atomic::Ordering::Relaxed) {
182                handle.interrupt();
183            }
184        });
185        let result = conn.query_row(
186            "WITH RECURSIVE t(x) AS (
187                 SELECT 1
188                 UNION ALL
189                 SELECT x + 1 FROM t WHERE x < 500000000
190             )
191             SELECT COUNT(*) FROM t",
192            [],
193            |r| r.get::<_, i64>(0),
194        );
195        stop.store(true, std::sync::atomic::Ordering::Relaxed);
196        let _ = watchdog.join();
197        let err = result.expect_err("recursive CTE must be interrupted");
198        assert!(
199            is_sqlite_interrupt(&err),
200            "expected SQLITE_INTERRUPT, got {err:?}"
201        );
202    }
203}