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