Skip to main content

sqlite_graphrag/commands/
memory_entities.rs

1//! Handler for the `memory-entities` CLI subcommand.
2
3use crate::errors::AppError;
4use crate::output;
5use crate::paths::AppPaths;
6use crate::storage::connection::open_ro;
7use rusqlite::params;
8use serde::Serialize;
9
10#[derive(clap::Args)]
11#[command(
12    about = "List entities linked to a memory, or memories linked to an entity",
13    after_long_help = "EXAMPLES:\n  \
14    # List entities connected to a memory (includes description)\n  \
15    sqlite-graphrag memory-entities --name my-memory\n\n  \
16    # Project name + description fields\n  \
17    sqlite-graphrag memory-entities --name my-memory | jaq '.entities[] | {name, description}'\n\n  \
18    # Reverse: list memories bound to an entity\n  \
19    sqlite-graphrag memory-entities --entity rust-lang\n\n  \
20    # With namespace\n  \
21    sqlite-graphrag memory-entities --name my-memory --namespace project"
22)]
23pub struct MemoryEntitiesArgs {
24    #[arg(value_name = "NAME", conflicts_with = "name", help = "Memory name")]
25    pub name_positional: Option<String>,
26    #[arg(long, conflicts_with_all = ["entity"])]
27    pub name: Option<String>,
28    /// Entity name — list memories bound to this entity (reverse lookup).
29    #[arg(long, conflicts_with_all = ["name", "name_positional"])]
30    pub entity: Option<String>,
31    #[arg(long, help = "Namespace (default: global; override via CLI flag)")]
32    pub namespace: Option<String>,
33    #[arg(long, hide = true)]
34    pub json: bool,
35    #[arg(long, help = "Database path (default: XDG data dir)")]
36    pub db: Option<String>,
37}
38
39#[derive(Serialize)]
40struct EntityBinding {
41    entity_id: i64,
42    name: String,
43    entity_type: String,
44    /// Entity description (empty string when NULL/unset in the graph).
45    description: String,
46}
47
48#[derive(Serialize)]
49struct MemoryEntitiesResponse {
50    memory_name: String,
51    entities: Vec<EntityBinding>,
52    count: usize,
53    elapsed_ms: u64,
54}
55
56#[derive(Serialize)]
57struct MemoryBinding {
58    memory_id: i64,
59    name: String,
60    description: String,
61    memory_type: String,
62}
63
64#[derive(Serialize)]
65struct EntityMemoriesResponse {
66    entity_name: String,
67    memories: Vec<MemoryBinding>,
68    count: usize,
69    elapsed_ms: u64,
70}
71
72pub fn run(args: MemoryEntitiesArgs) -> Result<(), AppError> {
73    let start = std::time::Instant::now();
74    let namespace = crate::namespace::resolve_namespace(args.namespace.as_deref())?;
75    let paths = AppPaths::resolve(args.db.as_deref())?;
76    crate::storage::connection::ensure_db_ready(&paths)?;
77    let conn = open_ro(&paths.db)?;
78
79    if let Some(entity_name) = args.entity {
80        let entity_id = crate::storage::entities::find_entity_id(&conn, &namespace, &entity_name)?
81            .ok_or_else(|| {
82                AppError::NotFound(crate::i18n::errors_msg::entity_not_found(
83                    &entity_name,
84                    &namespace,
85                ))
86            })?;
87
88        let mut stmt = conn.prepare_cached(
89            "SELECT m.id, m.name, m.description, m.type
90             FROM memory_entities me
91             JOIN memories m ON m.id = me.memory_id
92             WHERE me.entity_id = ?1 AND m.deleted_at IS NULL
93             ORDER BY m.name",
94        )?;
95
96        let memories: Vec<MemoryBinding> = stmt
97            .query_map(params![entity_id], |r| {
98                Ok(MemoryBinding {
99                    memory_id: r.get(0)?,
100                    name: r.get(1)?,
101                    description: r.get(2)?,
102                    memory_type: r.get(3)?,
103                })
104            })?
105            .collect::<Result<Vec<_>, _>>()?;
106
107        let count = memories.len();
108        output::emit_json(&EntityMemoriesResponse {
109            entity_name,
110            memories,
111            count,
112            elapsed_ms: start.elapsed().as_millis() as u64,
113        })?;
114        return Ok(());
115    }
116
117    let name = args.name_positional.or(args.name).ok_or_else(|| {
118        AppError::Validation(
119            "name required: pass as positional argument, via --name, or use --entity for reverse lookup".to_string(),
120        )
121    })?;
122
123    let memory_id: i64 = conn
124        .query_row(
125            "SELECT id FROM memories WHERE namespace = ?1 AND name = ?2 AND deleted_at IS NULL",
126            params![namespace, name],
127            |r| r.get(0),
128        )
129        .map_err(|_| {
130            AppError::NotFound(crate::i18n::errors_msg::memory_not_found(&name, &namespace))
131        })?;
132
133    let mut stmt = conn.prepare_cached(
134        "SELECT e.id, e.name, e.type AS entity_type,
135                COALESCE(e.description, '') AS description
136         FROM memory_entities me
137         JOIN entities e ON e.id = me.entity_id
138         WHERE me.memory_id = ?1
139         ORDER BY e.name",
140    )?;
141
142    let entities: Vec<EntityBinding> = stmt
143        .query_map(params![memory_id], |r| {
144            Ok(EntityBinding {
145                entity_id: r.get(0)?,
146                name: r.get(1)?,
147                entity_type: r.get(2)?,
148                description: r.get(3)?,
149            })
150        })?
151        .collect::<Result<Vec<_>, _>>()?;
152
153    let count = entities.len();
154
155    output::emit_json(&MemoryEntitiesResponse {
156        memory_name: name,
157        entities,
158        count,
159        elapsed_ms: start.elapsed().as_millis() as u64,
160    })?;
161
162    Ok(())
163}
164
165#[cfg(test)]
166mod tests {
167    use super::*;
168
169    #[test]
170    fn response_serializes_correctly() {
171        let resp = MemoryEntitiesResponse {
172            memory_name: "test-mem".to_string(),
173            entities: vec![EntityBinding {
174                entity_id: 1,
175                name: "rust".to_string(),
176                entity_type: "concept".to_string(),
177                description: "Systems programming language".to_string(),
178            }],
179            count: 1,
180            elapsed_ms: 5,
181        };
182        let json = serde_json::to_value(&resp).unwrap();
183        assert_eq!(json["memory_name"], "test-mem");
184        assert_eq!(json["count"], 1);
185        assert_eq!(json["entities"][0]["name"], "rust");
186        assert_eq!(
187            json["entities"][0]["description"],
188            "Systems programming language"
189        );
190        assert!(
191            json["entities"][0]
192                .as_object()
193                .expect("entity object")
194                .contains_key("description"),
195            "forward lookup must expose description (GAP-CLI-ME-01)"
196        );
197    }
198
199    #[test]
200    fn entity_binding_empty_description_serializes() {
201        let entity = EntityBinding {
202            entity_id: 2,
203            name: "icms-p05".to_string(),
204            entity_type: "concept".to_string(),
205            description: String::new(),
206        };
207        let json = serde_json::to_value(&entity).expect("serialize");
208        assert_eq!(json["description"], "");
209        assert_eq!(json["name"], "icms-p05");
210    }
211
212    #[test]
213    fn entity_memories_response_serializes_correctly() {
214        let resp = EntityMemoriesResponse {
215            entity_name: "rust-lang".to_string(),
216            memories: vec![MemoryBinding {
217                memory_id: 42,
218                name: "design-auth".to_string(),
219                description: "JWT auth design".to_string(),
220                memory_type: "decision".to_string(),
221            }],
222            count: 1,
223            elapsed_ms: 3,
224        };
225        let json = serde_json::to_value(&resp).unwrap();
226        assert_eq!(json["entity_name"], "rust-lang");
227        assert_eq!(json["count"], 1);
228        assert_eq!(json["memories"][0]["name"], "design-auth");
229        assert_eq!(json["memories"][0]["memory_type"], "decision");
230        assert_eq!(json["memories"][0]["memory_id"], 42);
231    }
232
233    #[test]
234    fn entity_memories_response_empty_list() {
235        let resp = EntityMemoriesResponse {
236            entity_name: "orphan-entity".to_string(),
237            memories: vec![],
238            count: 0,
239            elapsed_ms: 1,
240        };
241        let json = serde_json::to_value(&resp).unwrap();
242        assert_eq!(json["entity_name"], "orphan-entity");
243        assert_eq!(json["count"], 0);
244        assert!(json["memories"].as_array().unwrap().is_empty());
245    }
246}