Skip to main content

sqlite_graphrag/commands/
link.rs

1use crate::cli::RelationKind;
2use crate::constants::DEFAULT_RELATION_WEIGHT;
3use crate::errors::AppError;
4use crate::i18n::{erros, validacao};
5use crate::output::{self, OutputFormat};
6use crate::paths::AppPaths;
7use crate::storage::connection::open_rw;
8use crate::storage::entities;
9use serde::Serialize;
10
11#[derive(clap::Args)]
12pub struct LinkArgs {
13    /// Entidade de origem. Aceita alias `--source` para compatibilidade com doc bilíngue.
14    #[arg(long, alias = "source")]
15    pub from: String,
16    /// Entidade de destino. Aceita alias `--target` para compatibilidade com doc bilíngue.
17    #[arg(long, alias = "target")]
18    pub to: String,
19    #[arg(long, value_enum)]
20    pub relation: RelationKind,
21    #[arg(long)]
22    pub weight: Option<f64>,
23    #[arg(long)]
24    pub namespace: Option<String>,
25    #[arg(long, value_enum, default_value = "json")]
26    pub format: OutputFormat,
27    #[arg(long, hide = true, help = "No-op; JSON is always emitted on stdout")]
28    pub json: bool,
29    #[arg(long, env = "SQLITE_GRAPHRAG_DB_PATH")]
30    pub db: Option<String>,
31}
32
33#[derive(Serialize)]
34struct LinkResponse {
35    action: String,
36    from: String,
37    /// Duplicata de `from` para compatibilidade com docs que usam `source`.
38    source: String,
39    to: String,
40    /// Duplicata de `to` para compatibilidade com docs que usam `target`.
41    target: String,
42    relation: String,
43    weight: f64,
44    namespace: String,
45    /// Tempo total de execução em milissegundos desde início do handler até serialização.
46    elapsed_ms: u64,
47}
48
49pub fn run(args: LinkArgs) -> Result<(), AppError> {
50    let inicio = std::time::Instant::now();
51    let namespace = crate::namespace::resolve_namespace(args.namespace.as_deref())?;
52    let paths = AppPaths::resolve(args.db.as_deref())?;
53
54    if args.from == args.to {
55        return Err(AppError::Validation(validacao::link_auto_referencial()));
56    }
57
58    let weight = args.weight.unwrap_or(DEFAULT_RELATION_WEIGHT);
59    if !(0.0..=1.0).contains(&weight) {
60        return Err(AppError::Validation(validacao::link_peso_invalido(weight)));
61    }
62
63    if !paths.db.exists() {
64        return Err(AppError::NotFound(erros::banco_nao_encontrado(
65            &paths.db.display().to_string(),
66        )));
67    }
68
69    let relation_str = args.relation.as_str();
70
71    let mut conn = open_rw(&paths.db)?;
72
73    let source_id = entities::find_entity_id(&conn, &namespace, &args.from)?.ok_or_else(|| {
74        AppError::NotFound(erros::entidade_nao_encontrada(&args.from, &namespace))
75    })?;
76    let target_id = entities::find_entity_id(&conn, &namespace, &args.to)?
77        .ok_or_else(|| AppError::NotFound(erros::entidade_nao_encontrada(&args.to, &namespace)))?;
78
79    let tx = conn.transaction_with_behavior(rusqlite::TransactionBehavior::Immediate)?;
80    let (_rel_id, was_created) = entities::create_or_fetch_relationship(
81        &tx,
82        &namespace,
83        source_id,
84        target_id,
85        relation_str,
86        weight,
87        None,
88    )?;
89
90    if was_created {
91        entities::recalculate_degree(&tx, source_id)?;
92        entities::recalculate_degree(&tx, target_id)?;
93    }
94    tx.commit()?;
95
96    let action = if was_created {
97        "created".to_string()
98    } else {
99        "already_exists".to_string()
100    };
101
102    let response = LinkResponse {
103        action: action.clone(),
104        from: args.from.clone(),
105        source: args.from.clone(),
106        to: args.to.clone(),
107        target: args.to.clone(),
108        relation: relation_str.to_string(),
109        weight,
110        namespace: namespace.clone(),
111        elapsed_ms: inicio.elapsed().as_millis() as u64,
112    };
113
114    match args.format {
115        OutputFormat::Json => output::emit_json(&response)?,
116        OutputFormat::Text | OutputFormat::Markdown => {
117            output::emit_text(&format!(
118                "{}: {} --[{}]--> {} [{}]",
119                action, response.from, response.relation, response.to, response.namespace
120            ));
121        }
122    }
123
124    Ok(())
125}
126
127#[cfg(test)]
128mod testes {
129    use super::*;
130
131    #[test]
132    fn link_response_source_duplica_from() {
133        let resp = LinkResponse {
134            action: "created".to_string(),
135            from: "entidade-a".to_string(),
136            source: "entidade-a".to_string(),
137            to: "entidade-b".to_string(),
138            target: "entidade-b".to_string(),
139            relation: "uses".to_string(),
140            weight: 1.0,
141            namespace: "default".to_string(),
142            elapsed_ms: 0,
143        };
144        let json = serde_json::to_value(&resp).expect("serialização deve funcionar");
145        assert_eq!(json["source"], json["from"]);
146        assert_eq!(json["target"], json["to"]);
147        assert_eq!(json["source"], "entidade-a");
148        assert_eq!(json["target"], "entidade-b");
149    }
150
151    #[test]
152    fn link_response_serializa_todos_campos() {
153        let resp = LinkResponse {
154            action: "already_exists".to_string(),
155            from: "origem".to_string(),
156            source: "origem".to_string(),
157            to: "destino".to_string(),
158            target: "destino".to_string(),
159            relation: "mentions".to_string(),
160            weight: 0.8,
161            namespace: "teste".to_string(),
162            elapsed_ms: 5,
163        };
164        let json = serde_json::to_value(&resp).expect("serialização deve funcionar");
165        assert!(json.get("action").is_some());
166        assert!(json.get("from").is_some());
167        assert!(json.get("source").is_some());
168        assert!(json.get("to").is_some());
169        assert!(json.get("target").is_some());
170        assert!(json.get("relation").is_some());
171        assert!(json.get("weight").is_some());
172        assert!(json.get("namespace").is_some());
173        assert!(json.get("elapsed_ms").is_some());
174    }
175}