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 is SCAN → JUDGE (LLM) → PERSIST, with a SQLite queue DB derived
16//! next to `--db` (GAP-SG-64) for resume/retry support. The shape was inherited
17//! from the retired `ingest_claude.rs`, which v1.2.0 deleted along with every
18//! headless-subprocess frontend.
19// Workload: network-bound. JUDGE issues OpenRouter chat-completions requests
20// over HTTP in-process; `--rest-concurrency` is the only fan-out knob.
21//!
22//! # DRY note
23//!
24//! GAP-SG-121: enrich vs ingest queue table shapes are different products
25//! (item_key/operation vs file_path); only sidecar WAL/busy pragmas are shared
26//! via [`crate::pragmas::apply_sidecar_queue_pragmas`].
27
28mod args;
29mod drain_parallel;
30mod drain_serial;
31mod events;
32mod extraction;
33// extraction_providers is a child of extraction.rs (#[path])
34// extraction_ops_* are submodules of extraction.rs (#[path])
35mod postprocess;
36mod predicates;
37mod prompts;
38mod quality_sample;
39mod queue;
40mod queue_ops;
41mod reembed;
42mod run;
43mod scan;
44mod scan_ec;
45mod scheduler;
46mod schemas;
47mod status;
48
49pub(crate) const DEFAULT_RATE_LIMIT_WAIT: u64 = 60;
50pub(crate) const DEFAULT_BODY_ENRICH_MIN_CHARS: usize = 500;
51pub(crate) const DEFAULT_BODY_ENRICH_MAX_CHARS: usize = 2000;
52
53// GAP-SG-149: single source of truth for the enrich knobs that both the clap
54// surface (`default_value_t`) and the synthesized `EnrichArgs` of
55// `ingest --enrich-after` must agree on. Before this, `enrich_after.rs`
56// carried its own literals and two of them DIVERGED from the documented
57// default, so the auto-pass silently ran under a different contract than the
58// one `enrich --help` advertises.
59/// Attempts before a queue row is dead-lettered.
60pub(crate) const DEFAULT_ENRICH_MAX_ATTEMPTS: u32 = 8;
61/// Seconds after which a `processing` claim is considered abandoned.
62pub(crate) const DEFAULT_ENRICH_STALE_CLAIM_SECS: u64 = 1800;
63/// Seconds of headroom kept below a provider rate-limit reset.
64pub(crate) const DEFAULT_ENRICH_RATE_LIMIT_BUFFER_SECS: u64 = 300;
65/// Consecutive hard failures that trip the circuit breaker.
66pub(crate) const DEFAULT_ENRICH_CIRCUIT_BREAKER_THRESHOLD: u32 = 5;
67/// Minimum similarity for an enriched body to be accepted as preserving.
68pub(crate) const DEFAULT_ENRICH_PRESERVE_THRESHOLD: f64 = 0.7;
69/// Minimum grounding score for a generated entity description.
70pub(crate) const DEFAULT_ENRICH_GROUNDING_THRESHOLD: f64 = 0.12;
71
72pub use args::{EnrichArgs, EnrichMode, EnrichOperation, ReEmbedTarget};
73pub use queue::{cleanup_queue_entry, DeadItem, DeadSummary, EnrichStatus, WaitingItem};
74pub use run::run;
75
76use crate::errors::AppError;
77use queue::{enqueue_candidate_with_priority, open_queue_db, PRIORITY_HOT};
78
79/// GAP-CLI-PRIO-02: enqueue entity-descriptions for a hot set of entity names
80/// into the enrich sidecar queue with elevated priority.
81pub fn enqueue_priority_entity_descriptions(
82    paths: &crate::paths::AppPaths,
83    namespace: &str,
84    entity_names: &[String],
85) -> Result<usize, AppError> {
86    let _ = namespace; // queue keys are entity names; namespace is scoped by DB path
87    let queue_path = crate::paths::sidecar_path(&paths.db, ".enrich-queue.sqlite");
88    let queue = open_queue_db(&queue_path)?;
89    let mut n = 0usize;
90    for name in entity_names {
91        enqueue_candidate_with_priority(&queue, name, "entity", "EntityDescriptions", PRIORITY_HOT);
92        n += 1;
93    }
94    Ok(n)
95}
96
97#[cfg(test)]
98mod tests {
99    use super::events::{
100        enrich_operation_cli_name, is_sqlite_interrupt, scan_operation_with_deadline,
101    };
102    use super::schemas::{BINDINGS_SCHEMA, BODY_ENRICH_SCHEMA, ENTITY_DESCRIPTION_SCHEMA};
103    use super::*;
104    use crate::errors::AppError;
105    use rusqlite::{Connection, ErrorCode};
106    use std::time::{Duration, Instant};
107
108    #[test]
109    fn bindings_schema_is_valid_json() {
110        let _: serde_json::Value =
111            serde_json::from_str(BINDINGS_SCHEMA).expect("BINDINGS_SCHEMA must be valid JSON");
112    }
113
114    #[test]
115    fn entity_description_schema_is_valid_json() {
116        let _: serde_json::Value = serde_json::from_str(ENTITY_DESCRIPTION_SCHEMA)
117            .expect("ENTITY_DESCRIPTION_SCHEMA must be valid JSON");
118    }
119
120    #[test]
121    fn body_enrich_schema_is_valid_json() {
122        let _: serde_json::Value = serde_json::from_str(BODY_ENRICH_SCHEMA)
123            .expect("BODY_ENRICH_SCHEMA must be valid JSON");
124    }
125
126    // v1.1.06 — GAP-ENTITY-CONNECT-SCAN-CARTESIAN observability + interrupt
127
128    #[test]
129    fn enrich_operation_cli_name_pair_ops_are_kebab_case() {
130        assert_eq!(
131            enrich_operation_cli_name(&EnrichOperation::EntityConnect),
132            "entity-connect"
133        );
134        assert_eq!(
135            enrich_operation_cli_name(&EnrichOperation::CrossDomainBridges),
136            "cross-domain-bridges"
137        );
138        assert_eq!(
139            enrich_operation_cli_name(&EnrichOperation::EntityDescriptions),
140            "entity-descriptions"
141        );
142    }
143
144    #[test]
145    fn is_sqlite_interrupt_detects_operation_interrupted() {
146        let ffi_err = rusqlite::ffi::Error {
147            code: ErrorCode::OperationInterrupted,
148            extended_code: 9,
149        };
150        let err = rusqlite::Error::SqliteFailure(ffi_err, Some("interrupted".into()));
151        assert!(is_sqlite_interrupt(&err));
152
153        let busy = rusqlite::ffi::Error {
154            code: ErrorCode::DatabaseBusy,
155            extended_code: 5,
156        };
157        let busy_err = rusqlite::Error::SqliteFailure(busy, None);
158        assert!(!is_sqlite_interrupt(&busy_err));
159    }
160
161    #[test]
162    fn scan_deadline_already_elapsed_returns_timeout() {
163        // Past deadline must fail fast without running SQL (exit path → Timeout).
164        use clap::Parser;
165        let cli = crate::cli::Cli::try_parse_from([
166            "sqlite-graphrag",
167            "enrich",
168            "--operation",
169            "entity-connect",
170            "--mode",
171            "openrouter",
172            "--openrouter-model",
173            "test/model",
174            "--dry-run",
175            "--limit",
176            "1",
177        ])
178        .expect("parse enrich args");
179        let Some(crate::cli::Commands::Enrich(args)) = cli.command else {
180            panic!("expected Commands::Enrich");
181        };
182        let conn = Connection::open_in_memory().unwrap();
183        let past = Instant::now() - Duration::from_secs(1);
184        let err = scan_operation_with_deadline(&conn, "global", &args, Some(past))
185            .expect_err("elapsed deadline must Timeout");
186        match err {
187            AppError::Timeout { .. } => {}
188            other => panic!("expected Timeout, got {other:?}"),
189        }
190    }
191
192    #[test]
193    fn interrupt_handle_maps_long_query_to_sqlite_interrupt() {
194        // Live SQLite: watchdog interrupt aborts a recursive CTE (same mechanism
195        // as scan_operation_with_deadline). Confirms rusqlite InterruptHandle.
196        let conn = Connection::open_in_memory().unwrap();
197        let handle = conn.get_interrupt_handle();
198        let stop = std::sync::Arc::new(std::sync::atomic::AtomicBool::new(false));
199        let stop_w = std::sync::Arc::clone(&stop);
200        let watchdog = std::thread::spawn(move || {
201            std::thread::sleep(Duration::from_millis(30));
202            if !stop_w.load(std::sync::atomic::Ordering::Relaxed) {
203                handle.interrupt();
204            }
205        });
206        let result = conn.query_row(
207            "WITH RECURSIVE t(x) AS (
208                 SELECT 1
209                 UNION ALL
210                 SELECT x + 1 FROM t WHERE x < 500000000
211             )
212             SELECT COUNT(*) FROM t",
213            [],
214            |r| r.get::<_, i64>(0),
215        );
216        stop.store(true, std::sync::atomic::Ordering::Relaxed);
217        let _ = watchdog.join();
218        let err = result.expect_err("recursive CTE must be interrupted");
219        assert!(
220            is_sqlite_interrupt(&err),
221            "expected SQLITE_INTERRUPT, got {err:?}"
222        );
223    }
224}