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    /// Source entity.
14    #[arg(long)]
15    pub from: String,
16    /// Target entity.
17    #[arg(long)]
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, 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    to: String,
38    relation: String,
39    weight: f64,
40    namespace: String,
41    /// Tempo total de execução em milissegundos desde início do handler até serialização.
42    elapsed_ms: u64,
43}
44
45pub fn run(args: LinkArgs) -> Result<(), AppError> {
46    let inicio = std::time::Instant::now();
47    let namespace = crate::namespace::resolve_namespace(args.namespace.as_deref())?;
48    let paths = AppPaths::resolve(args.db.as_deref())?;
49
50    if args.from == args.to {
51        return Err(AppError::Validation(validacao::link_auto_referencial()));
52    }
53
54    let weight = args.weight.unwrap_or(DEFAULT_RELATION_WEIGHT);
55    if !(0.0..=1.0).contains(&weight) {
56        return Err(AppError::Validation(validacao::link_peso_invalido(weight)));
57    }
58
59    if !paths.db.exists() {
60        return Err(AppError::NotFound(erros::banco_nao_encontrado(
61            &paths.db.display().to_string(),
62        )));
63    }
64
65    let relation_str = args.relation.as_str();
66
67    let mut conn = open_rw(&paths.db)?;
68
69    let source_id = entities::find_entity_id(&conn, &namespace, &args.from)?.ok_or_else(|| {
70        AppError::NotFound(erros::entidade_nao_encontrada(&args.from, &namespace))
71    })?;
72    let target_id = entities::find_entity_id(&conn, &namespace, &args.to)?
73        .ok_or_else(|| AppError::NotFound(erros::entidade_nao_encontrada(&args.to, &namespace)))?;
74
75    let tx = conn.transaction_with_behavior(rusqlite::TransactionBehavior::Immediate)?;
76    let (_rel_id, was_created) = entities::create_or_fetch_relationship(
77        &tx,
78        &namespace,
79        source_id,
80        target_id,
81        relation_str,
82        weight,
83        None,
84    )?;
85
86    if was_created {
87        entities::recalculate_degree(&tx, source_id)?;
88        entities::recalculate_degree(&tx, target_id)?;
89    }
90    tx.commit()?;
91
92    let action = if was_created {
93        "created".to_string()
94    } else {
95        "already_exists".to_string()
96    };
97
98    let response = LinkResponse {
99        action: action.clone(),
100        from: args.from.clone(),
101        to: args.to.clone(),
102        relation: relation_str.to_string(),
103        weight,
104        namespace: namespace.clone(),
105        elapsed_ms: inicio.elapsed().as_millis() as u64,
106    };
107
108    match args.format {
109        OutputFormat::Json => output::emit_json(&response)?,
110        OutputFormat::Text | OutputFormat::Markdown => {
111            output::emit_text(&format!(
112                "{}: {} --[{}]--> {} [{}]",
113                action, response.from, response.relation, response.to, response.namespace
114            ));
115        }
116    }
117
118    Ok(())
119}
120
121#[cfg(test)]
122mod testes {
123    use super::*;
124
125    #[test]
126    fn link_response_sem_aliases_redundantes() {
127        // P1-O: campos source/target foram removidos do JSON de resposta.
128        let resp = LinkResponse {
129            action: "created".to_string(),
130            from: "entidade-a".to_string(),
131            to: "entidade-b".to_string(),
132            relation: "uses".to_string(),
133            weight: 1.0,
134            namespace: "default".to_string(),
135            elapsed_ms: 0,
136        };
137        let json = serde_json::to_value(&resp).expect("serialização deve funcionar");
138        assert_eq!(json["from"], "entidade-a");
139        assert_eq!(json["to"], "entidade-b");
140        assert!(
141            json.get("source").is_none(),
142            "campo 'source' foi removido em P1-O"
143        );
144        assert!(
145            json.get("target").is_none(),
146            "campo 'target' foi removido em P1-O"
147        );
148    }
149
150    #[test]
151    fn link_response_serializa_todos_campos() {
152        let resp = LinkResponse {
153            action: "already_exists".to_string(),
154            from: "origem".to_string(),
155            to: "destino".to_string(),
156            relation: "mentions".to_string(),
157            weight: 0.8,
158            namespace: "teste".to_string(),
159            elapsed_ms: 5,
160        };
161        let json = serde_json::to_value(&resp).expect("serialização deve funcionar");
162        assert!(json.get("action").is_some());
163        assert!(json.get("from").is_some());
164        assert!(json.get("to").is_some());
165        assert!(json.get("relation").is_some());
166        assert!(json.get("weight").is_some());
167        assert!(json.get("namespace").is_some());
168        assert!(json.get("elapsed_ms").is_some());
169    }
170}