Skip to main content

sqlite_graphrag/commands/
link.rs

1//! Handler for the `link` CLI subcommand.
2
3use crate::constants::DEFAULT_RELATION_WEIGHT;
4use crate::entity_type::EntityType;
5use crate::errors::AppError;
6use crate::i18n::{errors_msg, validation};
7use crate::output::{self, OutputFormat};
8use crate::paths::AppPaths;
9use crate::storage::connection::open_rw;
10use crate::storage::entities;
11use crate::storage::entities::NewEntity;
12use serde::Serialize;
13
14#[derive(clap::Args)]
15#[command(after_long_help = "EXAMPLES:\n  \
16    # Link two existing graph entities (extracted by GLiNER NER during `remember`)\n  \
17    sqlite-graphrag link --from oauth-flow --to refresh-tokens --relation related\n\n  \
18    # Auto-create entities that don't exist yet\n  \
19    sqlite-graphrag link --from concept-a --to concept-b --relation depends-on --create-missing\n\n  \
20    # Specify entity type for auto-created entities\n  \
21    sqlite-graphrag link --from alice --to acme-corp --relation related --create-missing --entity-type person\n\n  \
22    # Use a custom (non-canonical) relation type\n  \
23    sqlite-graphrag link --from module-a --to module-b --relation implements --create-missing\n\n  \
24    # If the entity does not exist and --create-missing is not set, the command fails with exit 4.\n  \
25    # To list current entity names:\n  \
26    sqlite-graphrag graph entities | jaq '.entities[].name'\n\n  \
27NOTE:\n  \
28    --from and --to expect ENTITY names (graph nodes), not memory names.\n  \
29    Memory names are managed via remember/read/edit/forget; entities are auto-extracted\n  \
30    by GLiNER NER from memory bodies or auto-created via --create-missing.")]
31pub struct LinkArgs {
32    /// Source ENTITY name (graph node, not memory). Entities are extracted by GLiNER NER during
33    /// `remember` or auto-created via `--create-missing`. Use `graph entities` to list
34    /// available entity names. Also accepts the alias `--name`.
35    #[arg(long, alias = "name")]
36    pub from: String,
37    /// Target ENTITY name (graph node, not memory). See `--from` for sourcing entity names.
38    #[arg(long)]
39    pub to: String,
40    /// Relation type between entities. Canonical values: applies-to, uses,
41    /// depends-on, causes, fixes, contradicts, supports, follows, related,
42    /// mentions, replaces, tracked-in. Any kebab-case or snake_case string
43    /// is also accepted as a custom relation.
44    #[arg(long, value_parser = crate::parsers::parse_relation, value_name = "RELATION")]
45    pub relation: String,
46    #[arg(long)]
47    pub weight: Option<f64>,
48    #[arg(long)]
49    pub namespace: Option<String>,
50    #[arg(long, value_enum, default_value = "json")]
51    pub format: OutputFormat,
52    #[arg(long, hide = true, help = "No-op; JSON is always emitted on stdout")]
53    pub json: bool,
54    #[arg(long, env = "SQLITE_GRAPHRAG_DB_PATH")]
55    pub db: Option<String>,
56    /// Auto-create entities when they do not exist. Created entities default to
57    /// type `concept` unless `--entity-type` specifies a different type.
58    #[arg(long, default_value_t = false)]
59    pub create_missing: bool,
60    /// Entity type assigned to auto-created entities (only effective with `--create-missing`).
61    #[arg(long, value_enum, default_value = "concept")]
62    pub entity_type: EntityType,
63}
64
65#[derive(Serialize)]
66struct LinkResponse {
67    action: String,
68    from: String,
69    to: String,
70    relation: String,
71    weight: f64,
72    namespace: String,
73    /// Total execution time in milliseconds from handler start to serialisation.
74    elapsed_ms: u64,
75    /// Entity names that were auto-created by `--create-missing`.
76    #[serde(skip_serializing_if = "Vec::is_empty")]
77    created_entities: Vec<String>,
78}
79
80pub fn run(args: LinkArgs) -> Result<(), AppError> {
81    let inicio = std::time::Instant::now();
82    let namespace = crate::namespace::resolve_namespace(args.namespace.as_deref())?;
83    let paths = AppPaths::resolve(args.db.as_deref())?;
84
85    if args.from == args.to {
86        return Err(AppError::Validation(validation::self_referential_link()));
87    }
88
89    let weight = args.weight.unwrap_or(DEFAULT_RELATION_WEIGHT);
90    if !(0.0..=1.0).contains(&weight) {
91        return Err(AppError::Validation(validation::invalid_link_weight(
92            weight,
93        )));
94    }
95
96    crate::storage::connection::ensure_db_ready(&paths)?;
97
98    crate::parsers::warn_if_non_canonical(&args.relation);
99    let relation_str = &args.relation;
100
101    let mut conn = open_rw(&paths.db)?;
102    let tx = conn.transaction_with_behavior(rusqlite::TransactionBehavior::Immediate)?;
103
104    let mut created_entities: Vec<String> = Vec::with_capacity(2);
105
106    let source_id = match entities::find_entity_id(&tx, &namespace, &args.from)? {
107        Some(id) => id,
108        None if args.create_missing => {
109            let new_entity = NewEntity {
110                name: args.from.clone(),
111                entity_type: args.entity_type,
112                description: None,
113            };
114            created_entities.push(args.from.clone());
115            entities::upsert_entity(&tx, &namespace, &new_entity)?
116        }
117        None => {
118            return Err(AppError::NotFound(errors_msg::entity_not_found(
119                &args.from, &namespace,
120            )));
121        }
122    };
123
124    let target_id = match entities::find_entity_id(&tx, &namespace, &args.to)? {
125        Some(id) => id,
126        None if args.create_missing => {
127            let new_entity = NewEntity {
128                name: args.to.clone(),
129                entity_type: args.entity_type,
130                description: None,
131            };
132            created_entities.push(args.to.clone());
133            entities::upsert_entity(&tx, &namespace, &new_entity)?
134        }
135        None => {
136            return Err(AppError::NotFound(errors_msg::entity_not_found(
137                &args.to, &namespace,
138            )));
139        }
140    };
141
142    let (_rel_id, was_created) = entities::create_or_fetch_relationship(
143        &tx,
144        &namespace,
145        source_id,
146        target_id,
147        relation_str,
148        weight,
149        None,
150    )?;
151
152    if was_created {
153        entities::recalculate_degree(&tx, source_id)?;
154        entities::recalculate_degree(&tx, target_id)?;
155    }
156    tx.commit()?;
157
158    conn.execute_batch("PRAGMA wal_checkpoint(TRUNCATE);")?;
159
160    let action = if was_created {
161        "created".to_string()
162    } else {
163        "already_exists".to_string()
164    };
165
166    let response = LinkResponse {
167        action: action.clone(),
168        from: args.from.clone(),
169        to: args.to.clone(),
170        relation: relation_str.to_string(),
171        weight,
172        namespace: namespace.clone(),
173        elapsed_ms: inicio.elapsed().as_millis() as u64,
174        created_entities,
175    };
176
177    match args.format {
178        OutputFormat::Json => output::emit_json(&response)?,
179        OutputFormat::Text | OutputFormat::Markdown => {
180            output::emit_text(&format!(
181                "{}: {} --[{}]--> {} [{}]",
182                action, response.from, response.relation, response.to, response.namespace
183            ));
184        }
185    }
186
187    Ok(())
188}
189
190#[cfg(test)]
191mod tests {
192    use super::*;
193
194    #[test]
195    fn link_response_without_redundant_aliases() {
196        // P1-O: source/target fields were removed from the JSON response.
197        let resp = LinkResponse {
198            action: "created".to_string(),
199            from: "entity-a".to_string(),
200            to: "entity-b".to_string(),
201            relation: "uses".to_string(),
202            weight: 1.0,
203            namespace: "default".to_string(),
204            elapsed_ms: 0,
205            created_entities: vec![],
206        };
207        let json = serde_json::to_value(&resp).expect("serialization must work");
208        assert_eq!(json["from"], "entity-a");
209        assert_eq!(json["to"], "entity-b");
210        assert!(
211            json.get("source").is_none(),
212            "field 'source' was removed in P1-O"
213        );
214        assert!(
215            json.get("target").is_none(),
216            "field 'target' was removed in P1-O"
217        );
218    }
219
220    #[test]
221    fn link_response_serializes_all_fields() {
222        let resp = LinkResponse {
223            action: "already_exists".to_string(),
224            from: "origin".to_string(),
225            to: "destination".to_string(),
226            relation: "mentions".to_string(),
227            weight: 0.8,
228            namespace: "test".to_string(),
229            elapsed_ms: 5,
230            created_entities: vec![],
231        };
232        let json = serde_json::to_value(&resp).expect("serialization must work");
233        assert!(json.get("action").is_some());
234        assert!(json.get("from").is_some());
235        assert!(json.get("to").is_some());
236        assert!(json.get("relation").is_some());
237        assert!(json.get("weight").is_some());
238        assert!(json.get("namespace").is_some());
239        assert!(json.get("elapsed_ms").is_some());
240    }
241
242    #[test]
243    fn link_response_omits_created_entities_when_empty() {
244        let resp = LinkResponse {
245            action: "created".to_string(),
246            from: "a".to_string(),
247            to: "b".to_string(),
248            relation: "uses".to_string(),
249            weight: 1.0,
250            namespace: "global".to_string(),
251            elapsed_ms: 0,
252            created_entities: vec![],
253        };
254        let json = serde_json::to_value(&resp).expect("serialization");
255        assert!(
256            json.get("created_entities").is_none(),
257            "empty vec must be omitted"
258        );
259    }
260
261    #[test]
262    fn link_response_includes_created_entities_when_present() {
263        let resp = LinkResponse {
264            action: "created".to_string(),
265            from: "new-a".to_string(),
266            to: "new-b".to_string(),
267            relation: "depends-on".to_string(),
268            weight: 0.5,
269            namespace: "test".to_string(),
270            elapsed_ms: 1,
271            created_entities: vec!["new-a".to_string(), "new-b".to_string()],
272        };
273        let json = serde_json::to_value(&resp).expect("serialization");
274        let created = json["created_entities"].as_array().expect("must be array");
275        assert_eq!(created.len(), 2);
276        assert_eq!(created[0], "new-a");
277        assert_eq!(created[1], "new-b");
278    }
279}