sqlite_graphrag/commands/enrich/
mod.rs1mod args;
28mod drain_parallel;
29mod drain_serial;
30mod events;
31mod extraction;
32mod 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
58pub fn enqueue_priority_entity_descriptions(
61 paths: &crate::paths::AppPaths,
62 namespace: &str,
63 entity_names: &[String],
64) -> Result<usize, AppError> {
65 let _ = namespace; 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 #[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 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 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}