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`). The queue
23//! DB schema in `ingest_claude.rs` still duplicates `open_queue_db` here — a
24//! future pass can unify them.
25
26mod args;
27mod drain_parallel;
28mod drain_serial;
29mod events;
30mod extraction;
31// extraction_providers is a child of extraction.rs (#[path])
32// extraction_ops_* are submodules of extraction.rs (#[path])
33mod postprocess;
34mod predicates;
35mod prompts;
36mod queue;
37mod queue_ops;
38mod run;
39mod scan;
40mod scan_ec;
41mod quality_sample;
42mod scheduler;
43mod schemas;
44mod status;
45
46// Surface schemas + defaults for sibling modules (`super::BINDINGS_SCHEMA`, etc.).
47#[allow(unused_imports)]
48pub(crate) use schemas::*;
49
50pub(crate) const DEFAULT_RATE_LIMIT_WAIT: u64 = 60;
51pub(crate) const DEFAULT_BODY_ENRICH_MIN_CHARS: usize = 500;
52pub(crate) const DEFAULT_BODY_ENRICH_MAX_CHARS: usize = 2000;
53
54pub use args::{EnrichArgs, EnrichMode, EnrichOperation, ReEmbedTarget};
55pub use queue::{cleanup_queue_entry, DeadItem, DeadSummary, EnrichStatus, WaitingItem};
56pub use run::run;
57
58// External types re-exported so historical `use super::*` in child modules keeps working.
59#[allow(unused_imports)]
60pub(crate) use crate::commands::ingest_claude::find_claude_binary;
61#[allow(unused_imports)]
62pub(crate) use crate::constants::MAX_MEMORY_BODY_LEN;
63#[allow(unused_imports)]
64pub(crate) use crate::entity_type::EntityType;
65#[allow(unused_imports)]
66pub(crate) use crate::errors::AppError;
67#[allow(unused_imports)]
68pub(crate) use crate::paths::AppPaths;
69#[allow(unused_imports)]
70pub(crate) use crate::storage::connection::{ensure_db_ready, open_rw};
71#[allow(unused_imports)]
72pub(crate) use crate::storage::entities::{self as entities, NewEntity, NewRelationship};
73#[allow(unused_imports)]
74pub(crate) use crate::storage::memories;
75#[allow(unused_imports)]
76pub(crate) use rusqlite::Connection;
77#[allow(unused_imports)]
78pub(crate) use serde::{Deserialize, Serialize};
79#[allow(unused_imports)]
80pub(crate) use std::io::Write;
81#[allow(unused_imports)]
82pub(crate) use std::path::{Path, PathBuf};
83#[allow(unused_imports)]
84pub(crate) use std::time::Instant;
85#[allow(unused_imports)]
86pub(crate) use prompts::ENTITY_DESCRIPTION_PROMPT_PREFIX;
87
88use queue::{enqueue_candidate_with_priority, open_queue_db, PRIORITY_HOT};
89
90/// GAP-CLI-PRIO-02: enqueue entity-descriptions for a hot set of entity names
91/// into the enrich sidecar queue with elevated priority.
92pub fn enqueue_priority_entity_descriptions(
93    paths: &crate::paths::AppPaths,
94    namespace: &str,
95    entity_names: &[String],
96) -> Result<usize, AppError> {
97    let _ = namespace; // queue keys are entity names; namespace is scoped by DB path
98    let queue_path = crate::paths::sidecar_path(&paths.db, ".enrich-queue.sqlite");
99    let queue = open_queue_db(&queue_path)?;
100    let mut n = 0usize;
101    for name in entity_names {
102        enqueue_candidate_with_priority(
103            &queue,
104            name,
105            "entity",
106            "EntityDescriptions",
107            PRIORITY_HOT,
108        );
109        n += 1;
110    }
111    Ok(n)
112}
113
114
115#[cfg(test)]
116mod tests {
117    use super::*;
118    use super::events::{enrich_operation_cli_name, is_sqlite_interrupt, scan_operation_with_deadline};
119    use rusqlite::ErrorCode;
120    use std::time::{Duration, Instant};
121
122    #[test]
123    fn bindings_schema_is_valid_json() {
124        let _: serde_json::Value =
125            serde_json::from_str(BINDINGS_SCHEMA).expect("BINDINGS_SCHEMA must be valid JSON");
126    }
127
128    #[test]
129    fn entity_description_schema_is_valid_json() {
130        let _: serde_json::Value = serde_json::from_str(ENTITY_DESCRIPTION_SCHEMA)
131            .expect("ENTITY_DESCRIPTION_SCHEMA must be valid JSON");
132    }
133
134    #[test]
135    fn body_enrich_schema_is_valid_json() {
136        let _: serde_json::Value = serde_json::from_str(BODY_ENRICH_SCHEMA)
137            .expect("BODY_ENRICH_SCHEMA must be valid JSON");
138    }
139
140    // v1.1.06 — GAP-ENTITY-CONNECT-SCAN-CARTESIAN observability + interrupt
141
142    #[test]
143    fn enrich_operation_cli_name_pair_ops_are_kebab_case() {
144        assert_eq!(
145            enrich_operation_cli_name(&EnrichOperation::EntityConnect),
146            "entity-connect"
147        );
148        assert_eq!(
149            enrich_operation_cli_name(&EnrichOperation::CrossDomainBridges),
150            "cross-domain-bridges"
151        );
152        assert_eq!(
153            enrich_operation_cli_name(&EnrichOperation::EntityDescriptions),
154            "entity-descriptions"
155        );
156    }
157
158    #[test]
159    fn is_sqlite_interrupt_detects_operation_interrupted() {
160        let ffi_err = rusqlite::ffi::Error {
161            code: ErrorCode::OperationInterrupted,
162            extended_code: 9,
163        };
164        let err = rusqlite::Error::SqliteFailure(ffi_err, Some("interrupted".into()));
165        assert!(is_sqlite_interrupt(&err));
166
167        let busy = rusqlite::ffi::Error {
168            code: ErrorCode::DatabaseBusy,
169            extended_code: 5,
170        };
171        let busy_err = rusqlite::Error::SqliteFailure(busy, None);
172        assert!(!is_sqlite_interrupt(&busy_err));
173    }
174
175    #[test]
176    fn scan_deadline_already_elapsed_returns_timeout() {
177        // Past deadline must fail fast without running SQL (exit path → Timeout).
178        use clap::Parser;
179        let cli = crate::cli::Cli::try_parse_from([
180            "sqlite-graphrag",
181            "enrich",
182            "--operation",
183            "entity-connect",
184            "--mode",
185            "openrouter",
186            "--openrouter-model",
187            "test/model",
188            "--dry-run",
189            "--limit",
190            "1",
191        ])
192        .expect("parse enrich args");
193        let Some(crate::cli::Commands::Enrich(args)) = cli.command else {
194            panic!("expected Commands::Enrich");
195        };
196        let conn = Connection::open_in_memory().unwrap();
197        let past = Instant::now() - Duration::from_secs(1);
198        let err = scan_operation_with_deadline(&conn, "global", &args, Some(past))
199            .expect_err("elapsed deadline must Timeout");
200        match err {
201            AppError::Timeout { .. } => {}
202            other => panic!("expected Timeout, got {other:?}"),
203        }
204    }
205
206    #[test]
207    fn interrupt_handle_maps_long_query_to_sqlite_interrupt() {
208        // Live SQLite: watchdog interrupt aborts a recursive CTE (same mechanism
209        // as scan_operation_with_deadline). Confirms rusqlite InterruptHandle.
210        let conn = Connection::open_in_memory().unwrap();
211        let handle = conn.get_interrupt_handle();
212        let stop = std::sync::Arc::new(std::sync::atomic::AtomicBool::new(false));
213        let stop_w = std::sync::Arc::clone(&stop);
214        let watchdog = std::thread::spawn(move || {
215            std::thread::sleep(Duration::from_millis(30));
216            if !stop_w.load(std::sync::atomic::Ordering::Relaxed) {
217                handle.interrupt();
218            }
219        });
220        let result = conn.query_row(
221            "WITH RECURSIVE t(x) AS (
222                 SELECT 1
223                 UNION ALL
224                 SELECT x + 1 FROM t WHERE x < 500000000
225             )
226             SELECT COUNT(*) FROM t",
227            [],
228            |r| r.get::<_, i64>(0),
229        );
230        stop.store(true, std::sync::atomic::Ordering::Relaxed);
231        let _ = watchdog.join();
232        let err = result.expect_err("recursive CTE must be interrupted");
233        assert!(
234            is_sqlite_interrupt(&err),
235            "expected SQLITE_INTERRUPT, got {err:?}"
236        );
237    }
238}