1use std::collections::{BTreeMap, HashMap};
2use std::sync::Arc;
3
4use anyhow::Result;
5use chrono::{DateTime, Utc};
6use locus_core_rs::domain::contracts::{NodeStore, SemanticIndexStore};
7use locus_core_rs::domain::models::{NodeQuery, SttpNode};
8use locus_core_rs::storage::derive_tenant_id_from_session;
9use serde_json::json;
10
11use crate::application::memory_filters::{
12 build_session_filter, node_matches_common_filters, resolve_indexed_sync_keys,
13};
14use crate::domain::graph::{MemoryGraphRequest, MemoryGraphResult};
15use crate::domain::memory::clamp_limit;
16
17const TENANT_SCAN_LIMIT: usize = 5000;
18
19pub struct MemoryGraphService {
20 store: Arc<dyn NodeStore>,
21 semantic_index: Option<Arc<dyn SemanticIndexStore>>,
22}
23
24impl MemoryGraphService {
25 pub fn new(store: Arc<dyn NodeStore>) -> Self {
26 Self {
27 store,
28 semantic_index: None,
29 }
30 }
31
32 pub fn with_semantic_index(
33 mut self,
34 semantic_index: Arc<dyn SemanticIndexStore>,
35 ) -> Self {
36 self.semantic_index = Some(semantic_index);
37 self
38 }
39
40 pub async fn execute(&self, request: &MemoryGraphRequest) -> Result<MemoryGraphResult> {
41 let capped_limit = clamp_limit(if request.limit == 0 {
42 1000
43 } else {
44 request.limit
45 })
46 .clamp(1, 5000);
47
48 let single_session = request
49 .scope
50 .session_ids
51 .as_deref()
52 .filter(|sessions| sessions.len() == 1)
53 .and_then(|sessions| sessions.first().cloned());
54
55 let backend_limit = if single_session.is_some() {
56 capped_limit
57 } else {
58 TENANT_SCAN_LIMIT
59 };
60
61 let mut nodes = self
62 .store
63 .query_nodes_async(NodeQuery {
64 limit: backend_limit,
65 session_id: single_session.clone(),
66 from_utc: request.scope.from_utc,
67 to_utc: request.scope.to_utc,
68 tiers: request.scope.tiers.clone(),
69 })
70 .await?;
71
72 let session_filter = build_session_filter(&request.scope);
73 let tenant_id = request
74 .scope
75 .tenant_id
76 .clone()
77 .or_else(|| {
78 single_session
79 .as_deref()
80 .map(derive_tenant_id_from_session)
81 })
82 .unwrap_or_else(|| "default".to_string());
83
84 let indexed_sync_keys = if let Some(index) = self.semantic_index.as_ref() {
85 resolve_indexed_sync_keys(
86 index.as_ref(),
87 &tenant_id,
88 &request.filter,
89 single_session.as_deref(),
90 backend_limit,
91 )
92 .await?
93 } else {
94 None
95 };
96
97 nodes.retain(|node| {
98 if let Some(keys) = &indexed_sync_keys
99 && !keys.contains(&node.sync_key)
100 {
101 return false;
102 }
103
104 node_matches_common_filters(
105 node,
106 &request.scope,
107 &request.filter,
108 session_filter.as_ref(),
109 )
110 });
111
112 nodes.sort_by(|left, right| right.timestamp.cmp(&left.timestamp));
113 nodes.truncate(capped_limit);
114
115 let include_topology = request.include_session_topology;
116 let include_lineage = request.include_lineage;
117 let include_semantic = request.include_semantic;
118 let rel_filter = request
119 .rel
120 .as_deref()
121 .map(|value| value.trim().to_ascii_lowercase())
122 .filter(|value| !value.is_empty());
123 let target_prefix = request
124 .target_prefix
125 .as_deref()
126 .map(|value| value.trim().to_ascii_lowercase())
127 .filter(|value| !value.is_empty());
128
129 #[derive(Clone)]
130 struct SessionGroup {
131 id: String,
132 label: String,
133 nodes: Vec<SttpNode>,
134 node_count: usize,
135 avg_psi: f32,
136 last_modified: DateTime<Utc>,
137 size: usize,
138 }
139
140 let mut grouped_map: BTreeMap<String, Vec<SttpNode>> = BTreeMap::new();
141 for node in &nodes {
142 grouped_map
143 .entry(node.session_id.clone())
144 .or_default()
145 .push(node.clone());
146 }
147
148 let mut grouped = grouped_map
149 .into_iter()
150 .map(|(id, mut session_nodes)| {
151 session_nodes.sort_by(|left, right| right.timestamp.cmp(&left.timestamp));
152 let node_count = session_nodes.len();
153 let avg_psi = if node_count == 0 {
154 0.0
155 } else {
156 session_nodes.iter().map(|node| node.psi).sum::<f32>() / node_count as f32
157 };
158 let last_modified = session_nodes
159 .first()
160 .map(|node| node.timestamp)
161 .unwrap_or_else(Utc::now);
162 let size = 16 + std::cmp::min(28, node_count * 2);
163
164 SessionGroup {
165 label: id.clone(),
166 id,
167 nodes: session_nodes,
168 node_count,
169 avg_psi,
170 last_modified,
171 size,
172 }
173 })
174 .collect::<Vec<_>>();
175
176 grouped.sort_by(|left, right| right.last_modified.cmp(&left.last_modified));
177
178 let node_by_id = nodes
179 .iter()
180 .map(|node| (graph_node_id(node), node.clone()))
181 .collect::<HashMap<_, _>>();
182
183 let sessions = if include_topology {
184 grouped
185 .iter()
186 .map(|session| {
187 json!({
188 "id": format!("s:{}", session.id),
189 "label": session.label,
190 "nodeCount": session.node_count,
191 "avgPsi": session.avg_psi,
192 "lastModified": session.last_modified.to_rfc3339(),
193 "size": session.size
194 })
195 })
196 .collect::<Vec<_>>()
197 } else {
198 Vec::new()
199 };
200
201 let graph_nodes = nodes
202 .iter()
203 .map(|node| {
204 json!({
205 "id": graph_node_id(node),
206 "sessionId": node.session_id,
207 "label": format!("{} {}", node.tier, node.timestamp.format("%m-%d %H:%M")),
208 "tier": node.tier,
209 "timestamp": node.timestamp.to_rfc3339(),
210 "psi": node.psi,
211 "parentNodeId": node.parent_node_id,
212 "semanticTags": node.semantic_tags,
213 "size": 9
214 })
215 })
216 .collect::<Vec<_>>();
217
218 let mut edges = Vec::new();
219
220 if include_topology {
221 for index in 0..grouped.len().saturating_sub(1) {
222 edges.push(json!({
223 "id": format!("t-{index}"),
224 "source": format!("s:{}", grouped[index].id),
225 "target": format!("s:{}", grouped[index + 1].id),
226 "kind": "timeline"
227 }));
228 }
229
230 for index in 0..grouped.len() {
231 let from = &grouped[index];
232 let mut nearest: Option<usize> = None;
233 let mut nearest_distance = f32::MAX;
234
235 for (other_index, other) in grouped.iter().enumerate() {
236 if index == other_index {
237 continue;
238 }
239 let distance = (from.avg_psi - other.avg_psi).abs();
240 if distance < nearest_distance {
241 nearest_distance = distance;
242 nearest = Some(other_index);
243 }
244 }
245
246 if let Some(nearest_index) = nearest
247 && index < nearest_index
248 {
249 edges.push(json!({
250 "id": format!("s-{index}-{nearest_index}"),
251 "source": format!("s:{}", from.id),
252 "target": format!("s:{}", grouped[nearest_index].id),
253 "kind": "similarity"
254 }));
255 }
256 }
257 }
258
259 for session in &grouped {
260 for index in 0..session.nodes.len() {
261 let current = &session.nodes[index];
262 let current_id = graph_node_id(current);
263
264 if include_topology {
265 edges.push(json!({
266 "id": format!("m-{}-{index}", session.id),
267 "source": format!("s:{}", session.id),
268 "target": current_id,
269 "kind": "membership"
270 }));
271
272 if index + 1 < session.nodes.len() {
273 let older = &session.nodes[index + 1];
274 edges.push(json!({
275 "id": format!("nt-{}-{index}", session.id),
276 "source": current_id,
277 "target": graph_node_id(older),
278 "kind": "node_timeline"
279 }));
280 }
281 }
282
283 if include_lineage
284 && let Some(parent) = current.parent_node_id.as_ref()
285 && node_by_id.contains_key(parent)
286 {
287 edges.push(json!({
288 "id": format!("l-{}-{index}", session.id),
289 "source": current_id,
290 "target": parent,
291 "kind": "lineage"
292 }));
293 }
294
295 if include_semantic
296 && let Some(links) = current.semantic_links.as_ref()
297 {
298 for (link_index, link) in links.iter().enumerate() {
299 if let Some(expected_rel) = rel_filter.as_deref()
300 && link.rel.trim().to_ascii_lowercase() != expected_rel
301 {
302 continue;
303 }
304
305 if let Some(prefix) = target_prefix.as_deref()
306 && !link.target.trim().to_ascii_lowercase().starts_with(prefix)
307 {
308 continue;
309 }
310
311 let target = if let Some(reference) = link.target.strip_prefix("ref:") {
312 let reference = reference.trim();
313 if node_by_id.contains_key(reference) {
314 reference.to_string()
315 } else {
316 link.target.clone()
317 }
318 } else {
319 link.target.clone()
320 };
321
322 let mut edge = json!({
323 "id": format!("sl-{}-{index}-{link_index}", session.id),
324 "source": current_id,
325 "target": target,
326 "kind": "semantic",
327 "rel": link.rel,
328 });
329 if let Some(confidence) = link.confidence {
330 edge.as_object_mut().map(|object| {
331 object.insert("confidence".to_string(), json!(confidence));
332 });
333 }
334 edges.push(edge);
335 }
336 }
337 }
338 }
339
340 Ok(MemoryGraphResult {
341 sessions,
342 nodes: graph_nodes,
343 edges,
344 retrieved: nodes.len(),
345 })
346 }
347}
348
349pub fn graph_node_id(node: &SttpNode) -> String {
350 format!(
351 "n:{}|{}|{}|{:.4}",
352 node.session_id,
353 node.timestamp.to_rfc3339(),
354 node.compression_depth,
355 node.psi
356 )
357}
358
359#[cfg(test)]
360mod tests {
361 use chrono::Utc;
362 use locus_core_rs::domain::models::{AvecState, SemanticLink, SttpNode};
363
364 use super::graph_node_id;
365
366 #[test]
367 fn graph_node_id_is_stable() {
368 let node = SttpNode {
369 raw: "raw".to_string(),
370 session_id: "demo".to_string(),
371 tier: "raw".to_string(),
372 timestamp: Utc::now(),
373 compression_depth: 1,
374 parent_node_id: None,
375 sync_key: "sync".to_string(),
376 updated_at: Utc::now(),
377 source_metadata: None,
378 context_summary: None,
379 semantic_tags: None,
380 semantic_links: Some(vec![SemanticLink {
381 rel: "evaluates".to_string(),
382 target: "ref:child".to_string(),
383 confidence: Some(0.9),
384 }]),
385 embedding: None,
386 embedding_model: None,
387 embedding_dimensions: None,
388 embedded_at: None,
389 user_avec: AvecState::zero(),
390 model_avec: AvecState::zero(),
391 compression_avec: None,
392 rho: 0.9,
393 kappa: 0.9,
394 psi: 2.5,
395 };
396
397 assert!(graph_node_id(&node).starts_with("n:demo|"));
398 }
399}