sqlite_graphrag/commands/enrich/
mod.rs1mod args;
29mod drain_parallel;
30mod drain_serial;
31mod events;
32mod extraction;
33mod 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
53pub(crate) const DEFAULT_ENRICH_MAX_ATTEMPTS: u32 = 8;
61pub(crate) const DEFAULT_ENRICH_STALE_CLAIM_SECS: u64 = 1800;
63pub(crate) const DEFAULT_ENRICH_RATE_LIMIT_BUFFER_SECS: u64 = 300;
65pub(crate) const DEFAULT_ENRICH_CIRCUIT_BREAKER_THRESHOLD: u32 = 5;
67pub(crate) const DEFAULT_ENRICH_PRESERVE_THRESHOLD: f64 = 0.7;
69pub(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
79pub fn enqueue_priority_entity_descriptions(
82 paths: &crate::paths::AppPaths,
83 namespace: &str,
84 entity_names: &[String],
85) -> Result<usize, AppError> {
86 let _ = namespace; 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 #[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 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 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}