sqlite_graphrag/commands/
reclassify.rs

1//! Handler for the `reclassify` CLI subcommand (GAP-18).
2//!
3//! Reclassifies one entity (single mode) or a whole group of entities (batch
4//! mode) by updating the `type` column in the `entities` table.
5//!
6//! Single mode: `--name <entity>` changes the type of one entity.
7//! Batch mode: `--from-type <old> --to-type <new> --batch` changes every
8//! entity in the namespace that currently has `<old>` as its type.
9
10use crate::entity_type::EntityType;
11use crate::errors::AppError;
12use crate::i18n::errors_msg;
13use crate::output::{self, OutputFormat};
14use crate::paths::AppPaths;
15use crate::storage::connection::open_rw;
16use crate::storage::entities;
17use rusqlite::params;
18use serde::Serialize;
19
20#[derive(clap::Args)]
21#[command(after_long_help = "EXAMPLES:\n  \
22    # Reclassify a single entity from its current type to 'tool'\n  \
23    sqlite-graphrag reclassify --name tokio-runtime --new-type tool\n\n  \
24    # Reclassify all 'concept' entities to 'tool' in one shot (batch)\n  \
25    sqlite-graphrag reclassify --from-type concept --to-type tool --batch\n\n  \
26    # Reclassify in a specific namespace\n  \
27    sqlite-graphrag reclassify --name alice --new-type person --namespace my-project\n\n\
28NOTE:\n  \
29    Single mode requires --name and --new-type.\n  \
30    Batch mode requires --from-type, --to-type and --batch.\n  \
31    Providing --name together with --batch is an error.")]
32pub struct ReclassifyArgs {
33    /// Entity name to reclassify (single mode). Mutually exclusive with --from-type + --batch.
34    #[arg(long, conflicts_with_all = ["from_type", "batch"])]
35    pub name: Option<String>,
36    /// New entity type for single mode.
37    #[arg(long, value_enum, value_name = "TYPE")]
38    pub new_type: Option<EntityType>,
39    /// Current entity type to match in batch mode. Requires --to-type and --batch.
40    #[arg(
41        long,
42        value_enum,
43        value_name = "TYPE",
44        requires = "to_type",
45        requires = "batch"
46    )]
47    pub from_type: Option<EntityType>,
48    /// New entity type to assign in batch mode. Requires --from-type and --batch.
49    #[arg(long, value_enum, value_name = "TYPE", requires = "from_type")]
50    pub to_type: Option<EntityType>,
51    /// Enable batch reclassification (--from-type to --to-type). Requires --from-type and --to-type.
52    #[arg(long, default_value_t = false, requires = "from_type")]
53    pub batch: bool,
54    #[arg(long)]
55    pub namespace: Option<String>,
56    #[arg(long, value_enum, default_value = "json")]
57    pub format: OutputFormat,
58    #[arg(long, hide = true, help = "No-op; JSON is always emitted on stdout")]
59    pub json: bool,
60    #[arg(long, env = "SQLITE_GRAPHRAG_DB_PATH")]
61    pub db: Option<String>,
62}
63
64#[derive(Serialize)]
65struct ReclassifyResponse {
66    action: String,
67    count: usize,
68    namespace: String,
69    /// Total execution time in milliseconds from handler start to serialisation.
70    elapsed_ms: u64,
71}
72
73pub fn run(args: ReclassifyArgs) -> Result<(), AppError> {
74    let inicio = std::time::Instant::now();
75    let namespace = crate::namespace::resolve_namespace(args.namespace.as_deref())?;
76    let paths = AppPaths::resolve(args.db.as_deref())?;
77
78    crate::storage::connection::ensure_db_ready(&paths)?;
79
80    let mut conn = open_rw(&paths.db)?;
81
82    let count = if args.batch {
83        // Batch mode: --from-type + --to-type + --batch
84        let from_type = args.from_type.ok_or_else(|| {
85            AppError::Validation("--from-type is required in batch mode".to_string())
86        })?;
87        let to_type = args.to_type.ok_or_else(|| {
88            AppError::Validation("--to-type is required in batch mode".to_string())
89        })?;
90
91        let tx = conn.transaction_with_behavior(rusqlite::TransactionBehavior::Immediate)?;
92        let affected = tx.execute(
93            "UPDATE entities SET type = ?1, updated_at = unixepoch()
94             WHERE type = ?2 AND namespace = ?3",
95            params![to_type.as_str(), from_type.as_str(), namespace],
96        )?;
97        tx.commit()?;
98        if affected == 0 {
99            tracing::warn!(
100                from_type = from_type.as_str(),
101                namespace = %namespace,
102                "reclassify batch matched zero entities — verify --from-type value exists"
103            );
104        }
105        affected
106    } else {
107        // Single mode: --name + --new-type
108        let entity_name = args
109            .name
110            .as_deref()
111            .ok_or_else(|| AppError::Validation("--name is required in single mode".to_string()))?;
112        let new_type = args.new_type.ok_or_else(|| {
113            AppError::Validation("--new-type is required in single mode".to_string())
114        })?;
115
116        // Verify entity exists.
117        entities::find_entity_id(&conn, &namespace, entity_name)?.ok_or_else(|| {
118            AppError::NotFound(errors_msg::entity_not_found(entity_name, &namespace))
119        })?;
120
121        let tx = conn.transaction_with_behavior(rusqlite::TransactionBehavior::Immediate)?;
122        let affected = tx.execute(
123            "UPDATE entities SET type = ?1, updated_at = unixepoch()
124             WHERE name = ?2 AND namespace = ?3",
125            params![new_type.as_str(), entity_name, namespace],
126        )?;
127        tx.commit()?;
128        affected
129    };
130
131    conn.execute_batch("PRAGMA wal_checkpoint(TRUNCATE);")?;
132
133    let response = ReclassifyResponse {
134        action: "reclassified".to_string(),
135        count,
136        namespace: namespace.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                "reclassified: {} entities [{}]",
145                response.count, response.namespace
146            ));
147        }
148    }
149
150    Ok(())
151}
152
153#[cfg(test)]
154mod tests {
155    use super::*;
156
157    #[test]
158    fn reclassify_response_serializes_all_fields() {
159        let resp = ReclassifyResponse {
160            action: "reclassified".to_string(),
161            count: 5,
162            namespace: "global".to_string(),
163            elapsed_ms: 12,
164        };
165        let json = serde_json::to_value(&resp).expect("serialization failed");
166        assert_eq!(json["action"], "reclassified");
167        assert_eq!(json["count"], 5);
168        assert_eq!(json["namespace"], "global");
169        assert!(json["elapsed_ms"].is_number());
170    }
171
172    #[test]
173    fn reclassify_response_count_zero_is_valid() {
174        let resp = ReclassifyResponse {
175            action: "reclassified".to_string(),
176            count: 0,
177            namespace: "my-project".to_string(),
178            elapsed_ms: 3,
179        };
180        let json = serde_json::to_value(&resp).expect("serialization failed");
181        assert_eq!(json["count"], 0);
182        assert_eq!(json["action"], "reclassified");
183    }
184
185    #[test]
186    fn reclassify_response_action_is_reclassified() {
187        let resp = ReclassifyResponse {
188            action: "reclassified".to_string(),
189            count: 1,
190            namespace: "ns".to_string(),
191            elapsed_ms: 1,
192        };
193        assert_eq!(resp.action, "reclassified");
194    }
195}
sqlite_graphrag/commands/reclassify.rs

sqlite_graphrag/commands/
reclassify.rs
