sqlite_graphrag/commands/
prune_ner.rs

1//! Handler for the `prune-ner` CLI subcommand.
2//!
3//! Removes NER bindings (rows in `memory_entities`) for a single entity or for
4//! all entities in the namespace. Useful for cleaning up low-quality automatic
5//! extractions without touching the entities or memories themselves.
6
7use crate::errors::AppError;
8use crate::output::{self, OutputFormat};
9use crate::paths::AppPaths;
10use crate::storage::connection::open_rw;
11use serde::Serialize;
12
13#[derive(clap::Args)]
14#[command(after_long_help = "EXAMPLES:\n  \
15    # Preview bindings that would be removed for a single entity\n  \
16    sqlite-graphrag prune-ner --entity jwt-token --dry-run\n\n  \
17    # Remove all NER bindings for a single entity\n  \
18    sqlite-graphrag prune-ner --entity jwt-token --yes\n\n  \
19    # Remove ALL NER bindings in the current namespace\n  \
20    sqlite-graphrag prune-ner --all --yes\n\n  \
21NOTE:\n  \
22    This command deletes rows from memory_entities (the link table between\n  \
23    memories and extracted entities). The entities and memories themselves\n  \
24    are not deleted. Use cleanup-orphans afterwards to remove entity nodes\n  \
25    that have no remaining links.")]
26/// Prune NER args.
27pub struct PruneNerArgs {
28    /// Entity name whose bindings should be removed.
29    /// Mutually exclusive with --all.
30    #[arg(long, conflicts_with = "all", value_name = "NAME")]
31    pub entity: Option<String>,
32
33    /// Remove all NER bindings in the namespace. Mutually exclusive with --entity.
34    #[arg(long, conflicts_with = "entity", default_value_t = false)]
35    pub all: bool,
36
37    /// Namespace scope.
38    #[arg(long)]
39    pub namespace: Option<String>,
40
41    /// Preview count without deleting.
42    #[arg(long, default_value_t = false)]
43    pub dry_run: bool,
44
45    /// Skip confirmation for destructive operation.
46    #[arg(long, default_value_t = false)]
47    pub yes: bool,
48
49    /// Output format.
50    #[arg(long, value_enum, default_value = "json")]
51    pub format: OutputFormat,
52
53    /// Emit machine-readable JSON on stdout.
54    #[arg(long, hide = true, help = "No-op; JSON is always emitted on stdout")]
55    pub json: bool,
56
57    /// Path to the SQLite database file.
58    #[arg(long)]
59    pub db: Option<String>,
60}
61
62#[derive(Serialize)]
63struct PruneNerResponse {
64    action: String,
65    bindings_removed: usize,
66    namespace: String,
67    /// Entity name targeted, when `--entity` was used.
68    #[serde(skip_serializing_if = "Option::is_none")]
69    entity: Option<String>,
70    /// Total execution time in milliseconds from handler start to serialisation.
71    elapsed_ms: u64,
72}
73
74/// Run.
75pub fn run(args: PruneNerArgs) -> Result<(), AppError> {
76    let inicio = std::time::Instant::now();
77
78    if args.entity.is_none() && !args.all {
79        return Err(AppError::Validation(
80            "either --entity <NAME> or --all must be specified".to_string(),
81        ));
82    }
83
84    let namespace = crate::namespace::resolve_namespace(args.namespace.as_deref())?;
85    let paths = AppPaths::resolve(args.db.as_deref())?;
86
87    crate::storage::connection::ensure_db_ready(&paths)?;
88
89    let mut conn = open_rw(&paths.db)?;
90
91    // Count how many rows would be affected.
92    let count: usize = if let Some(ref entity_name) = args.entity {
93        conn.query_row(
94            "SELECT COUNT(*) FROM memory_entities me
95             JOIN entities e ON e.id = me.entity_id
96             WHERE e.name = ?1 AND e.namespace = ?2",
97            rusqlite::params![entity_name, namespace],
98            |r| r.get::<_, i64>(0).map(|v| v as usize),
99        )?
100    } else {
101        conn.query_row(
102            "SELECT COUNT(*) FROM memory_entities me
103             JOIN entities e ON e.id = me.entity_id
104             WHERE e.namespace = ?1",
105            rusqlite::params![namespace],
106            |r| r.get::<_, i64>(0).map(|v| v as usize),
107        )?
108    };
109
110    if args.dry_run {
111        let response = PruneNerResponse {
112            action: "dry_run".to_string(),
113            bindings_removed: count,
114            namespace: namespace.clone(),
115            entity: args.entity.clone(),
116            elapsed_ms: inicio.elapsed().as_millis() as u64,
117        };
118
119        match args.format {
120            OutputFormat::Json => output::emit_json(&response)?,
121            OutputFormat::Text | OutputFormat::Markdown => {
122                output::emit_text(&format!(
123                    "dry_run: {count} NER bindings would be removed [{namespace}]"
124                ));
125            }
126        }
127
128        return Ok(());
129    }
130
131    if !args.yes {
132        let response = PruneNerResponse {
133            action: "aborted".to_string(),
134            bindings_removed: count,
135            namespace: namespace.clone(),
136            entity: args.entity.clone(),
137            elapsed_ms: inicio.elapsed().as_millis() as u64,
138        };
139
140        match args.format {
141            OutputFormat::Json => output::emit_json(&response)?,
142            OutputFormat::Text | OutputFormat::Markdown => {
143                output::emit_text(&format!(
144                    "aborted: {count} NER bindings would be removed; pass --yes to confirm [{namespace}]"
145                ));
146            }
147        }
148
149        return Ok(());
150    }
151
152    // Destructive path: COUNT + DELETE in same transaction for consistency.
153    let removed: usize = if let Some(ref entity_name) = args.entity {
154        // Normalize to match the normalized stored entity names.
155        let entity_name = crate::parsers::normalize_entity_name(entity_name);
156        let tx = conn.transaction_with_behavior(rusqlite::TransactionBehavior::Immediate)?;
157        let n = tx.execute(
158            "DELETE FROM memory_entities WHERE entity_id IN (
159                 SELECT id FROM entities WHERE name = ?1 AND namespace = ?2
160             )",
161            rusqlite::params![entity_name, namespace],
162        )?;
163        tx.commit()?;
164        n
165    } else {
166        let tx = conn.transaction_with_behavior(rusqlite::TransactionBehavior::Immediate)?;
167        let n = tx.execute(
168            "DELETE FROM memory_entities WHERE entity_id IN (
169                 SELECT id FROM entities WHERE namespace = ?1
170             )",
171            rusqlite::params![namespace],
172        )?;
173        tx.commit()?;
174        n
175    };
176
177    conn.execute_batch("PRAGMA wal_checkpoint(TRUNCATE);")?;
178
179    tracing::info!(target: "prune_ner",
180        removed = removed,
181        namespace = %namespace,
182        entity = ?args.entity,
183        "NER bindings pruned"
184    );
185
186    let response = PruneNerResponse {
187        action: "pruned".to_string(),
188        bindings_removed: removed,
189        namespace: namespace.clone(),
190        entity: args.entity.clone(),
191        elapsed_ms: inicio.elapsed().as_millis() as u64,
192    };
193
194    match args.format {
195        OutputFormat::Json => output::emit_json(&response)?,
196        OutputFormat::Text | OutputFormat::Markdown => {
197            output::emit_text(&format!(
198                "pruned: {removed} NER bindings removed [{namespace}]"
199            ));
200        }
201    }
202
203    Ok(())
204}
205
206#[cfg(test)]
207mod tests {
208    use super::*;
209
210    #[test]
211    fn prune_ner_response_dry_run_serializes_correctly() {
212        let resp = PruneNerResponse {
213            action: "dry_run".to_string(),
214            bindings_removed: 42,
215            namespace: "global".to_string(),
216            entity: Some("jwt-token".to_string()),
217            elapsed_ms: 5,
218        };
219        let json = serde_json::to_value(&resp).expect("serialization failed");
220        assert_eq!(json["action"], "dry_run");
221        assert_eq!(json["bindings_removed"], 42);
222        assert_eq!(json["entity"], "jwt-token");
223        assert_eq!(json["namespace"], "global");
224    }
225
226    #[test]
227    fn prune_ner_response_pruned_all_omits_entity() {
228        let resp = PruneNerResponse {
229            action: "pruned".to_string(),
230            bindings_removed: 200,
231            namespace: "project-x".to_string(),
232            entity: None,
233            elapsed_ms: 15,
234        };
235        let json = serde_json::to_value(&resp).expect("serialization failed");
236        assert_eq!(json["action"], "pruned");
237        assert_eq!(json["bindings_removed"], 200);
238        assert!(
239            json.get("entity").is_none(),
240            "entity must be omitted when None"
241        );
242    }
243
244    #[test]
245    fn prune_ner_response_aborted_includes_count() {
246        let resp = PruneNerResponse {
247            action: "aborted".to_string(),
248            bindings_removed: 10,
249            namespace: "global".to_string(),
250            entity: None,
251            elapsed_ms: 1,
252        };
253        let json = serde_json::to_value(&resp).expect("serialization failed");
254        assert_eq!(json["action"], "aborted");
255        assert_eq!(json["bindings_removed"], 10);
256        assert!(json["elapsed_ms"].is_number());
257    }
258
259    #[test]
260    fn prune_ner_response_zero_bindings() {
261        let resp = PruneNerResponse {
262            action: "pruned".to_string(),
263            bindings_removed: 0,
264            namespace: "global".to_string(),
265            entity: Some("nonexistent".to_string()),
266            elapsed_ms: 2,
267        };
268        let json = serde_json::to_value(&resp).expect("serialization failed");
269        assert_eq!(json["bindings_removed"], 0);
270    }
271}
sqlite_graphrag/commands/prune_ner.rs

sqlite_graphrag/commands/
prune_ner.rs
