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    let action = if was_created {
159        "created".to_string()
160    } else {
161        "already_exists".to_string()
162    };
163
164    let response = LinkResponse {
165        action: action.clone(),
166        from: args.from.clone(),
167        to: args.to.clone(),
168        relation: relation_str.to_string(),
169        weight,
170        namespace: namespace.clone(),
171        elapsed_ms: inicio.elapsed().as_millis() as u64,
172        created_entities,
173    };
174
175    match args.format {
176        OutputFormat::Json => output::emit_json(&response)?,
177        OutputFormat::Text | OutputFormat::Markdown => {
178            output::emit_text(&format!(
179                "{}: {} --[{}]--> {} [{}]",
180                action, response.from, response.relation, response.to, response.namespace
181            ));
182        }
183    }
184
185    Ok(())
186}
187
188#[cfg(test)]
189mod tests {
190    use super::*;
191
192    #[test]
193    fn link_response_without_redundant_aliases() {
194        // P1-O: source/target fields were removed from the JSON response.
195        let resp = LinkResponse {
196            action: "created".to_string(),
197            from: "entity-a".to_string(),
198            to: "entity-b".to_string(),
199            relation: "uses".to_string(),
200            weight: 1.0,
201            namespace: "default".to_string(),
202            elapsed_ms: 0,
203            created_entities: vec![],
204        };
205        let json = serde_json::to_value(&resp).expect("serialization must work");
206        assert_eq!(json["from"], "entity-a");
207        assert_eq!(json["to"], "entity-b");
208        assert!(
209            json.get("source").is_none(),
210            "field 'source' was removed in P1-O"
211        );
212        assert!(
213            json.get("target").is_none(),
214            "field 'target' was removed in P1-O"
215        );
216    }
217
218    #[test]
219    fn link_response_serializes_all_fields() {
220        let resp = LinkResponse {
221            action: "already_exists".to_string(),
222            from: "origin".to_string(),
223            to: "destination".to_string(),
224            relation: "mentions".to_string(),
225            weight: 0.8,
226            namespace: "test".to_string(),
227            elapsed_ms: 5,
228            created_entities: vec![],
229        };
230        let json = serde_json::to_value(&resp).expect("serialization must work");
231        assert!(json.get("action").is_some());
232        assert!(json.get("from").is_some());
233        assert!(json.get("to").is_some());
234        assert!(json.get("relation").is_some());
235        assert!(json.get("weight").is_some());
236        assert!(json.get("namespace").is_some());
237        assert!(json.get("elapsed_ms").is_some());
238    }
239
240    #[test]
241    fn link_response_omits_created_entities_when_empty() {
242        let resp = LinkResponse {
243            action: "created".to_string(),
244            from: "a".to_string(),
245            to: "b".to_string(),
246            relation: "uses".to_string(),
247            weight: 1.0,
248            namespace: "global".to_string(),
249            elapsed_ms: 0,
250            created_entities: vec![],
251        };
252        let json = serde_json::to_value(&resp).expect("serialization");
253        assert!(
254            json.get("created_entities").is_none(),
255            "empty vec must be omitted"
256        );
257    }
258
259    #[test]
260    fn link_response_includes_created_entities_when_present() {
261        let resp = LinkResponse {
262            action: "created".to_string(),
263            from: "new-a".to_string(),
264            to: "new-b".to_string(),
265            relation: "depends-on".to_string(),
266            weight: 0.5,
267            namespace: "test".to_string(),
268            elapsed_ms: 1,
269            created_entities: vec!["new-a".to_string(), "new-b".to_string()],
270        };
271        let json = serde_json::to_value(&resp).expect("serialization");
272        let created = json["created_entities"].as_array().expect("must be array");
273        assert_eq!(created.len(), 2);
274        assert_eq!(created[0], "new-a");
275        assert_eq!(created[1], "new-b");
276    }
277}