sqlite_graphrag/commands/enrich/
mod.rs1mod args;
27mod drain_parallel;
28mod drain_serial;
29mod events;
30mod extraction;
31mod 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#[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#[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
90pub fn enqueue_priority_entity_descriptions(
93 paths: &crate::paths::AppPaths,
94 namespace: &str,
95 entity_names: &[String],
96) -> Result<usize, AppError> {
97 let _ = namespace; 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 #[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 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 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}