sqlite_graphrag/commands/
link.rs

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

sqlite_graphrag/commands/
link.rs
