1use std::collections::HashMap;
5use std::sync::Arc;
6use std::time::Duration;
7#[allow(unused_imports)]
8use zeph_db::sql;
9
10use futures::TryStreamExt as _;
11use petgraph::Graph;
12use petgraph::graph::NodeIndex;
13use tokio::sync::Semaphore;
14use tokio::task::JoinSet;
15use zeph_llm::LlmProvider as _;
16use zeph_llm::any::AnyProvider;
17use zeph_llm::provider::{Message, Role};
18
19use crate::error::MemoryError;
20
21use super::store::GraphStore;
22use super::types::{Edge, Entity};
23
24const MAX_LABEL_PROPAGATION_ITERATIONS: usize = 50;
25
26fn scrub_content(s: &str) -> String {
35 s.chars()
36 .filter(|c| {
37 !c.is_control()
38 && !matches!(*c as u32,
39 0x200B..=0x200F | 0x202A..=0x202E | 0x2066..=0x2069 | 0xFEFF
40 )
41 })
42 .collect()
43}
44
45#[derive(Debug, Default)]
47pub struct GraphEvictionStats {
48 pub expired_edges_deleted: usize,
49 pub orphan_entities_deleted: usize,
50 pub capped_entities_deleted: usize,
51}
52
53fn truncate_prompt(prompt: String, max_bytes: usize) -> String {
59 if max_bytes == 0 {
60 return String::new();
61 }
62 if prompt.len() <= max_bytes {
63 return prompt;
64 }
65 let boundary = prompt.floor_char_boundary(max_bytes);
66 format!("{}...", &prompt[..boundary])
67}
68
69fn compute_partition_fingerprint(entity_ids: &[i64], intra_edge_ids: &[i64]) -> String {
75 let mut hasher = blake3::Hasher::new();
76 let mut sorted_entities = entity_ids.to_vec();
77 sorted_entities.sort_unstable();
78 hasher.update(b"entities");
79 for id in &sorted_entities {
80 hasher.update(&id.to_le_bytes());
81 }
82 let mut sorted_edges = intra_edge_ids.to_vec();
83 sorted_edges.sort_unstable();
84 hasher.update(b"edges");
85 for id in &sorted_edges {
86 hasher.update(&id.to_le_bytes());
87 }
88 hasher.finalize().to_hex().to_string()
89}
90
91struct CommunityData {
93 entity_ids: Vec<i64>,
94 entity_names: Vec<String>,
95 intra_facts: Vec<String>,
96 fingerprint: String,
97 name: String,
98}
99
100type UndirectedGraph = Graph<i64, (), petgraph::Undirected>;
101
102fn fold_edge(
107 edge: &Edge,
108 node_map: &HashMap<i64, NodeIndex>,
109 graph: &mut UndirectedGraph,
110 edge_facts_map: &mut HashMap<(i64, i64), Vec<String>>,
111 edge_id_map: &mut HashMap<(i64, i64), Vec<i64>>,
112) {
113 if let (Some(&src_idx), Some(&tgt_idx)) = (
114 node_map.get(&edge.source_entity_id),
115 node_map.get(&edge.target_entity_id),
116 ) {
117 graph.add_edge(src_idx, tgt_idx, ());
118 }
119 let key = (edge.source_entity_id, edge.target_entity_id);
120 edge_facts_map
121 .entry(key)
122 .or_default()
123 .push(edge.fact.clone());
124 edge_id_map.entry(key).or_default().push(edge.id);
125}
126
127async fn build_entity_graph_and_maps(
128 store: &GraphStore,
129 entities: &[Entity],
130 edge_chunk_size: usize,
131) -> Result<
132 (
133 UndirectedGraph,
134 HashMap<(i64, i64), Vec<String>>,
135 HashMap<(i64, i64), Vec<i64>>,
136 ),
137 MemoryError,
138> {
139 let mut graph = UndirectedGraph::new_undirected();
140 let mut node_map: HashMap<i64, NodeIndex> = HashMap::new();
141
142 for entity in entities {
143 let idx = graph.add_node(entity.id.0);
144 node_map.insert(entity.id.0, idx);
145 }
146
147 let mut edge_facts_map: HashMap<(i64, i64), Vec<String>> = HashMap::new();
148 let mut edge_id_map: HashMap<(i64, i64), Vec<i64>> = HashMap::new();
149
150 if edge_chunk_size == 0 {
151 let edges: Vec<_> = store.all_active_edges_stream().try_collect().await?;
152 for edge in &edges {
153 fold_edge(
154 edge,
155 &node_map,
156 &mut graph,
157 &mut edge_facts_map,
158 &mut edge_id_map,
159 );
160 }
161 } else {
162 let limit = i64::try_from(edge_chunk_size).unwrap_or(i64::MAX);
163 let mut last_id: i64 = 0;
164 loop {
165 let chunk = store.edges_after_id(last_id, limit).await?;
166 if chunk.is_empty() {
167 break;
168 }
169 last_id = chunk.last().expect("non-empty chunk has a last element").id;
170 for edge in &chunk {
171 fold_edge(
172 edge,
173 &node_map,
174 &mut graph,
175 &mut edge_facts_map,
176 &mut edge_id_map,
177 );
178 }
179 }
180 }
181
182 Ok((graph, edge_facts_map, edge_id_map))
183}
184
185fn run_label_propagation(graph: &UndirectedGraph) -> HashMap<usize, Vec<i64>> {
186 let mut labels: Vec<usize> = (0..graph.node_count()).collect();
187
188 for _ in 0..MAX_LABEL_PROPAGATION_ITERATIONS {
189 let mut changed = false;
190 for node_idx in graph.node_indices() {
191 let neighbors: Vec<NodeIndex> = graph.neighbors(node_idx).collect();
192 if neighbors.is_empty() {
193 continue;
194 }
195 let mut freq: HashMap<usize, usize> = HashMap::new();
196 for &nbr in &neighbors {
197 *freq.entry(labels[nbr.index()]).or_insert(0) += 1;
198 }
199 let max_count = *freq.values().max().unwrap_or(&0);
200 let best_label = freq
201 .iter()
202 .filter(|&(_, count)| *count == max_count)
203 .map(|(&label, _)| label)
204 .min()
205 .unwrap_or(labels[node_idx.index()]);
206 if labels[node_idx.index()] != best_label {
207 labels[node_idx.index()] = best_label;
208 changed = true;
209 }
210 }
211 if !changed {
212 break;
213 }
214 }
215
216 let mut communities: HashMap<usize, Vec<i64>> = HashMap::new();
217 for node_idx in graph.node_indices() {
218 let entity_id = graph[node_idx];
219 communities
220 .entry(labels[node_idx.index()])
221 .or_default()
222 .push(entity_id);
223 }
224 communities.retain(|_, members| members.len() >= 2);
225 communities
226}
227
228struct ClassifyResult {
229 to_summarize: Vec<CommunityData>,
230 unchanged_count: usize,
231 new_fingerprints: std::collections::HashSet<String>,
232}
233
234fn classify_communities(
235 communities: &HashMap<usize, Vec<i64>>,
236 edge_facts_map: &HashMap<(i64, i64), Vec<String>>,
237 edge_id_map: &HashMap<(i64, i64), Vec<i64>>,
238 entity_name_map: &HashMap<i64, &str>,
239 stored_fingerprints: &HashMap<String, i64>,
240 sorted_labels: &[usize],
241) -> ClassifyResult {
242 let mut to_summarize: Vec<CommunityData> = Vec::new();
243 let mut unchanged_count = 0usize;
244 let mut new_fingerprints: std::collections::HashSet<String> = std::collections::HashSet::new();
245
246 for (label_index, &label) in sorted_labels.iter().enumerate() {
247 let entity_ids = communities[&label].as_slice();
248 let member_set: std::collections::HashSet<i64> = entity_ids.iter().copied().collect();
249
250 let mut intra_facts: Vec<String> = Vec::new();
251 let mut intra_edge_ids: Vec<i64> = Vec::new();
252 for (&(src, tgt), facts) in edge_facts_map {
253 if member_set.contains(&src) && member_set.contains(&tgt) {
254 intra_facts.extend(facts.iter().map(|f| scrub_content(f)));
255 if let Some(ids) = edge_id_map.get(&(src, tgt)) {
256 intra_edge_ids.extend_from_slice(ids);
257 }
258 }
259 }
260
261 let fingerprint = compute_partition_fingerprint(entity_ids, &intra_edge_ids);
262 new_fingerprints.insert(fingerprint.clone());
263
264 if stored_fingerprints.contains_key(&fingerprint) {
265 unchanged_count += 1;
266 continue;
267 }
268
269 let entity_names: Vec<String> = entity_ids
270 .iter()
271 .filter_map(|id| entity_name_map.get(id).map(|&s| scrub_content(s)))
272 .collect();
273
274 let base_name = entity_names
277 .iter()
278 .take(3)
279 .cloned()
280 .collect::<Vec<_>>()
281 .join(", ");
282 let name = format!("{base_name} [{label_index}]");
283
284 to_summarize.push(CommunityData {
285 entity_ids: entity_ids.to_vec(),
286 entity_names,
287 intra_facts,
288 fingerprint,
289 name,
290 });
291 }
292
293 ClassifyResult {
294 to_summarize,
295 unchanged_count,
296 new_fingerprints,
297 }
298}
299
300async fn summarize_and_upsert_communities(
301 store: &GraphStore,
302 provider: &AnyProvider,
303 to_summarize: Vec<CommunityData>,
304 concurrency: usize,
305 community_summary_max_prompt_bytes: usize,
306) -> Result<usize, MemoryError> {
307 let semaphore = Arc::new(Semaphore::new(concurrency.max(1)));
308 let mut join_set: JoinSet<(String, String, Vec<i64>, String)> = JoinSet::new();
309
310 for data in to_summarize {
311 let provider = provider.clone();
312 let sem = Arc::clone(&semaphore);
313 let max_bytes = community_summary_max_prompt_bytes;
314 join_set.spawn(async move {
315 let _permit = sem.acquire().await.expect("semaphore is never closed");
316 let summary = match generate_community_summary(
317 &provider,
318 &data.entity_names,
319 &data.intra_facts,
320 max_bytes,
321 )
322 .await
323 {
324 Ok(text) => text,
325 Err(e) => {
326 tracing::warn!(community = %data.name, "community summary generation failed: {e:#}");
327 String::new()
328 }
329 };
330 (data.name, summary, data.entity_ids, data.fingerprint)
331 });
332 }
333
334 let mut results: Vec<(String, String, Vec<i64>, String)> = Vec::new();
336 while let Some(outcome) = join_set.join_next().await {
337 match outcome {
338 Ok(tuple) => results.push(tuple),
339 Err(e) => {
340 tracing::error!(
341 panicked = e.is_panic(),
342 cancelled = e.is_cancelled(),
343 "community summary task failed"
344 );
345 }
346 }
347 }
348
349 results.sort_unstable_by(|a, b| a.0.cmp(&b.0));
350
351 let mut count = 0usize;
352 for (name, summary, entity_ids, fingerprint) in results {
353 store
354 .upsert_community(&name, &summary, &entity_ids, Some(&fingerprint))
355 .await?;
356 count += 1;
357 }
358
359 Ok(count)
360}
361
362pub async fn detect_communities(
381 store: &GraphStore,
382 provider: &AnyProvider,
383 community_summary_max_prompt_bytes: usize,
384 concurrency: usize,
385 edge_chunk_size: usize,
386) -> Result<usize, MemoryError> {
387 let edge_chunk_size = if edge_chunk_size == 0 {
388 tracing::warn!(
389 "edge_chunk_size is 0, which would load all edges into memory; \
390 using safe default of 10_000"
391 );
392 10_000_usize
393 } else {
394 edge_chunk_size
395 };
396
397 let entities = store.all_entities().await?;
398 if entities.len() < 2 {
399 return Ok(0);
400 }
401
402 let (graph, edge_facts_map, edge_id_map) =
403 build_entity_graph_and_maps(store, &entities, edge_chunk_size).await?;
404
405 let communities = run_label_propagation(&graph);
406
407 let entity_name_map: HashMap<i64, &str> =
408 entities.iter().map(|e| (e.id.0, e.name.as_str())).collect();
409 let stored_fingerprints = store.community_fingerprints().await?;
410
411 let mut sorted_labels: Vec<usize> = communities.keys().copied().collect();
412 sorted_labels.sort_unstable();
413
414 let ClassifyResult {
415 to_summarize,
416 unchanged_count,
417 new_fingerprints,
418 } = classify_communities(
419 &communities,
420 &edge_facts_map,
421 &edge_id_map,
422 &entity_name_map,
423 &stored_fingerprints,
424 &sorted_labels,
425 );
426
427 tracing::debug!(
428 total = sorted_labels.len(),
429 unchanged = unchanged_count,
430 to_summarize = to_summarize.len(),
431 "community detection: partition classification complete"
432 );
433
434 for (stored_fp, community_id) in &stored_fingerprints {
436 if !new_fingerprints.contains(stored_fp.as_str()) {
437 store.delete_community_by_id(*community_id).await?;
438 }
439 }
440
441 let new_count = summarize_and_upsert_communities(
442 store,
443 provider,
444 to_summarize,
445 concurrency,
446 community_summary_max_prompt_bytes,
447 )
448 .await?;
449
450 Ok(unchanged_count + new_count)
451}
452
453pub async fn assign_to_community(
464 store: &GraphStore,
465 entity_id: i64,
466) -> Result<Option<i64>, MemoryError> {
467 let edges = store.edges_for_entity(entity_id).await?;
468 if edges.is_empty() {
469 return Ok(None);
470 }
471
472 let neighbor_ids: Vec<i64> = edges
473 .iter()
474 .map(|e| {
475 if e.source_entity_id == entity_id {
476 e.target_entity_id
477 } else {
478 e.source_entity_id
479 }
480 })
481 .collect();
482
483 let mut community_votes: HashMap<i64, usize> = HashMap::new();
484 for &nbr_id in &neighbor_ids {
485 if let Some(community) = store.community_for_entity(nbr_id).await? {
486 *community_votes.entry(community.id).or_insert(0) += 1;
487 }
488 }
489
490 if community_votes.is_empty() {
491 return Ok(None);
492 }
493
494 let Some((&best_community_id, _)) =
497 community_votes
498 .iter()
499 .max_by(|&(&id_a, &count_a), &(&id_b, &count_b)| {
500 count_a.cmp(&count_b).then(id_b.cmp(&id_a))
501 })
502 else {
503 return Ok(None);
504 };
505
506 if let Some(mut target) = store.find_community_by_id(best_community_id).await? {
507 if !target.entity_ids.iter().any(|eid| eid.0 == entity_id) {
508 target.entity_ids.push(crate::types::EntityId(entity_id));
509 let raw_ids: Vec<i64> = target.entity_ids.iter().map(|eid| eid.0).collect();
510 store
511 .upsert_community(&target.name, &target.summary, &raw_ids, None)
512 .await?;
513 store.clear_community_fingerprint(best_community_id).await?;
515 }
516 return Ok(Some(best_community_id));
517 }
518
519 Ok(None)
520}
521
522pub async fn cleanup_stale_entity_embeddings(
530 store: &GraphStore,
531 embeddings: &crate::embedding_store::EmbeddingStore,
532) -> Result<usize, MemoryError> {
533 const ENTITY_COLLECTION: &str = "zeph_graph_entities";
534
535 let pairs = embeddings.scroll_all_entity_ids(ENTITY_COLLECTION).await?;
539 if pairs.is_empty() {
540 return Ok(0);
541 }
542
543 let qdrant_ids: Vec<i64> = pairs.iter().map(|(_, eid)| *eid).collect();
544 let live: std::collections::HashSet<i64> = store
545 .entity_ids_in(&qdrant_ids)
546 .await?
547 .into_iter()
548 .collect();
549
550 let stale_point_ids: Vec<String> = pairs
551 .into_iter()
552 .filter_map(|(pid, eid)| (!live.contains(&eid)).then_some(pid))
553 .collect();
554
555 if stale_point_ids.is_empty() {
556 return Ok(0);
557 }
558
559 let count = stale_point_ids.len();
560 embeddings
561 .delete_from_collection(ENTITY_COLLECTION, stale_point_ids)
562 .await?;
563 Ok(count)
564}
565
566pub async fn run_graph_eviction(
572 store: &GraphStore,
573 expired_edge_retention_days: u32,
574 max_entities: usize,
575) -> Result<GraphEvictionStats, MemoryError> {
576 let expired_edges_deleted = store
577 .delete_expired_edges(expired_edge_retention_days)
578 .await?;
579 let orphan_entities_deleted = store
580 .delete_orphan_entities(expired_edge_retention_days)
581 .await?;
582 let capped_entities_deleted = if max_entities > 0 {
583 store.cap_entities(max_entities).await?
584 } else {
585 0
586 };
587
588 Ok(GraphEvictionStats {
589 expired_edges_deleted,
590 orphan_entities_deleted,
591 capped_entities_deleted,
592 })
593}
594
595async fn generate_community_summary(
596 provider: &AnyProvider,
597 entity_names: &[String],
598 edge_facts: &[String],
599 max_prompt_bytes: usize,
600) -> Result<String, MemoryError> {
601 let entities_str = entity_names.join(", ");
602 let facts_str = edge_facts
604 .iter()
605 .take(20)
606 .map(|f| format!("- {f}"))
607 .collect::<Vec<_>>()
608 .join("\n");
609
610 let raw_prompt = format!(
611 "Summarize the following group of related entities and their relationships \
612 into a single paragraph (2-3 sentences). Focus on the theme that connects \
613 them and the key relationships.\n\nEntities: {entities_str}\n\
614 Relationships:\n{facts_str}\n\nSummary:"
615 );
616
617 let original_bytes = raw_prompt.len();
618 let truncated = raw_prompt.len() > max_prompt_bytes;
619 let prompt = truncate_prompt(raw_prompt, max_prompt_bytes);
620 if prompt.is_empty() {
621 return Ok(String::new());
622 }
623 if truncated {
624 tracing::warn!(
625 entity_count = entity_names.len(),
626 original_bytes,
627 truncated_bytes = prompt.len(),
628 "community summary prompt truncated"
629 );
630 }
631
632 let messages = [Message::from_legacy(Role::User, prompt)];
633 let response: String = tokio::time::timeout(Duration::from_secs(15), provider.chat(&messages))
634 .await
635 .map_err(|_| {
636 tracing::warn!("community summary generation: LLM call timed out after 15s");
637 MemoryError::Timeout("community summary: LLM call timed out after 15s".into())
638 })?
639 .map_err(MemoryError::Llm)?;
640 Ok(response)
641}
642
643#[cfg(test)]
644mod tests {
645 use std::sync::{Arc, Mutex};
646
647 use super::*;
648 use crate::graph::types::EntityType;
649 use crate::store::SqliteStore;
650
651 async fn setup() -> GraphStore {
652 let store = SqliteStore::new(":memory:").await.unwrap();
653 GraphStore::new(store.pool().clone())
654 }
655
656 fn mock_provider() -> AnyProvider {
657 AnyProvider::Mock(zeph_llm::mock::MockProvider::default())
658 }
659
660 fn recording_provider() -> (
661 AnyProvider,
662 Arc<Mutex<Vec<Vec<zeph_llm::provider::Message>>>>,
663 ) {
664 let (mock, buf) = zeph_llm::mock::MockProvider::default().with_recording();
665 (AnyProvider::Mock(mock), buf)
666 }
667
668 #[tokio::test]
669 async fn test_detect_communities_empty_graph() {
670 let store = setup().await;
671 let provider = mock_provider();
672 let count = detect_communities(&store, &provider, usize::MAX, 4, 0)
673 .await
674 .unwrap();
675 assert_eq!(count, 0);
676 }
677
678 #[tokio::test]
679 async fn test_detect_communities_single_entity() {
680 let store = setup().await;
681 let provider = mock_provider();
682 store
683 .upsert_entity("Solo", "Solo", EntityType::Concept, None, None)
684 .await
685 .unwrap();
686 let count = detect_communities(&store, &provider, usize::MAX, 4, 0)
687 .await
688 .unwrap();
689 assert_eq!(count, 0, "single isolated entity must not form a community");
690 }
691
692 #[tokio::test]
693 async fn test_single_entity_community_filtered() {
694 let store = setup().await;
695 let provider = mock_provider();
696
697 let a = store
699 .upsert_entity("A", "A", EntityType::Concept, None, None)
700 .await
701 .unwrap()
702 .0;
703 let b = store
704 .upsert_entity("B", "B", EntityType::Concept, None, None)
705 .await
706 .unwrap()
707 .0;
708 let c = store
709 .upsert_entity("C", "C", EntityType::Concept, None, None)
710 .await
711 .unwrap()
712 .0;
713 let iso = store
714 .upsert_entity("Isolated", "Isolated", EntityType::Concept, None, None)
715 .await
716 .unwrap()
717 .0;
718
719 store
720 .insert_edge(a, b, "r", "A relates B", 1.0, None, None)
721 .await
722 .unwrap();
723 store
724 .insert_edge(b, c, "r", "B relates C", 1.0, None, None)
725 .await
726 .unwrap();
727
728 let count = detect_communities(&store, &provider, usize::MAX, 4, 0)
729 .await
730 .unwrap();
731 assert_eq!(count, 1, "only the 3-entity cluster should be detected");
733
734 let communities = store.all_communities().await.unwrap();
735 assert_eq!(communities.len(), 1);
736 assert!(
737 !communities[0].entity_ids.iter().any(|eid| eid.0 == iso),
738 "isolated entity must not be in any community"
739 );
740 }
741
742 #[tokio::test]
743 async fn test_label_propagation_basic() {
744 let store = setup().await;
745 let provider = mock_provider();
746
747 let mut cluster_ids: Vec<Vec<i64>> = Vec::new();
749 for cluster in 0..4_i64 {
750 let mut ids = Vec::new();
751 for node in 0..3_i64 {
752 let name = format!("c{cluster}_n{node}");
753 let id = store
754 .upsert_entity(&name, &name, EntityType::Concept, None, None)
755 .await
756 .unwrap()
757 .0;
758 ids.push(id);
759 }
760 store
762 .insert_edge(ids[0], ids[1], "r", "f", 1.0, None, None)
763 .await
764 .unwrap();
765 store
766 .insert_edge(ids[1], ids[2], "r", "f", 1.0, None, None)
767 .await
768 .unwrap();
769 cluster_ids.push(ids);
770 }
771
772 let count = detect_communities(&store, &provider, usize::MAX, 4, 0)
773 .await
774 .unwrap();
775 assert_eq!(count, 4, "expected 4 communities, one per cluster");
776
777 let communities = store.all_communities().await.unwrap();
778 assert_eq!(communities.len(), 4);
779
780 for ids in &cluster_ids {
782 let found = communities
783 .iter()
784 .filter(|c| {
785 ids.iter()
786 .any(|id| c.entity_ids.iter().any(|eid| eid.0 == *id))
787 })
788 .count();
789 assert_eq!(
790 found, 1,
791 "all nodes of a cluster must be in the same community"
792 );
793 }
794 }
795
796 #[tokio::test]
797 async fn test_all_isolated_nodes() {
798 let store = setup().await;
799 let provider = mock_provider();
800
801 for i in 0..5_i64 {
803 store
804 .upsert_entity(
805 &format!("iso_{i}"),
806 &format!("iso_{i}"),
807 EntityType::Concept,
808 None,
809 None,
810 )
811 .await
812 .unwrap();
813 }
814
815 let count = detect_communities(&store, &provider, usize::MAX, 4, 0)
816 .await
817 .unwrap();
818 assert_eq!(count, 0, "zero-edge graph must produce no communities");
819 assert_eq!(store.community_count().await.unwrap(), 0);
820 }
821
822 #[tokio::test]
823 async fn test_eviction_expired_edges() {
824 let store = setup().await;
825
826 let a = store
827 .upsert_entity("EA", "EA", EntityType::Concept, None, None)
828 .await
829 .unwrap()
830 .0;
831 let b = store
832 .upsert_entity("EB", "EB", EntityType::Concept, None, None)
833 .await
834 .unwrap()
835 .0;
836 let edge_id = store
837 .insert_edge(a, b, "r", "f", 1.0, None, None)
838 .await
839 .unwrap();
840 store.invalidate_edge(edge_id).await.unwrap();
841
842 zeph_db::query(sql!(
844 "UPDATE graph_edges SET expired_at = datetime('now', '-200 days') WHERE id = ?1"
845 ))
846 .bind(edge_id)
847 .execute(store.pool())
848 .await
849 .unwrap();
850
851 let stats = run_graph_eviction(&store, 90, 0).await.unwrap();
852 assert_eq!(stats.expired_edges_deleted, 1);
853 }
854
855 #[tokio::test]
856 async fn test_eviction_orphan_entities() {
857 let store = setup().await;
858
859 let iso = store
860 .upsert_entity("Orphan", "Orphan", EntityType::Concept, None, None)
861 .await
862 .unwrap()
863 .0;
864
865 zeph_db::query(sql!(
867 "UPDATE graph_entities SET last_seen_at = datetime('now', '-200 days') WHERE id = ?1"
868 ))
869 .bind(iso)
870 .execute(store.pool())
871 .await
872 .unwrap();
873
874 let stats = run_graph_eviction(&store, 90, 0).await.unwrap();
875 assert_eq!(stats.orphan_entities_deleted, 1);
876 }
877
878 #[tokio::test]
879 async fn test_eviction_entity_cap() {
880 let store = setup().await;
881
882 for i in 0..5_i64 {
884 let name = format!("cap_entity_{i}");
885 store
886 .upsert_entity(&name, &name, EntityType::Concept, None, None)
887 .await
888 .unwrap();
889 }
890
891 let stats = run_graph_eviction(&store, 90, 3).await.unwrap();
892 assert_eq!(
893 stats.capped_entities_deleted, 2,
894 "should delete 5-3=2 entities"
895 );
896 assert_eq!(store.entity_count().await.unwrap(), 3);
897 }
898
899 #[tokio::test]
900 async fn test_assign_to_community_no_neighbors() {
901 let store = setup().await;
902 let entity_id = store
903 .upsert_entity("Loner", "Loner", EntityType::Concept, None, None)
904 .await
905 .unwrap()
906 .0;
907
908 let result = assign_to_community(&store, entity_id).await.unwrap();
909 assert!(result.is_none());
910 }
911
912 #[tokio::test]
913 async fn test_extraction_count_persistence() {
914 use tempfile::NamedTempFile;
915 let tmp = NamedTempFile::new().unwrap();
917 let path = tmp.path().to_str().unwrap().to_owned();
918
919 let store1 = {
920 let s = crate::store::SqliteStore::new(&path).await.unwrap();
921 GraphStore::new(s.pool().clone())
922 };
923
924 store1.set_metadata("extraction_count", "0").await.unwrap();
925 for i in 1..=5_i64 {
926 store1
927 .set_metadata("extraction_count", &i.to_string())
928 .await
929 .unwrap();
930 }
931
932 let store2 = {
934 let s = crate::store::SqliteStore::new(&path).await.unwrap();
935 GraphStore::new(s.pool().clone())
936 };
937 assert_eq!(store2.extraction_count().await.unwrap(), 5);
938 }
939
940 #[test]
941 fn test_scrub_content_ascii_control() {
942 let input = "hello\nworld\r\x00\x01\x09end";
944 assert_eq!(scrub_content(input), "helloworldend");
945 }
946
947 #[test]
948 fn test_scrub_content_bidi_overrides() {
949 let input = "safe\u{202A}inject\u{202E}end\u{2066}iso\u{2069}done".to_string();
952 assert_eq!(scrub_content(&input), "safeinjectendisodone");
953 }
954
955 #[test]
956 fn test_scrub_content_zero_width() {
957 let input = "a\u{200B}b\u{200C}c\u{200D}d\u{200F}e".to_string();
960 assert_eq!(scrub_content(&input), "abcde");
961 }
962
963 #[test]
964 fn test_scrub_content_bom() {
965 let input = "\u{FEFF}hello".to_string();
967 assert_eq!(scrub_content(&input), "hello");
968 }
969
970 #[test]
971 fn test_scrub_content_clean_string_unchanged() {
972 let input = "Hello, World! 123 — normal text.";
973 assert_eq!(scrub_content(input), input);
974 }
975
976 #[test]
977 fn test_truncate_prompt_within_limit() {
978 let result = truncate_prompt("short".into(), 100);
979 assert_eq!(result, "short");
980 }
981
982 #[test]
983 fn test_truncate_prompt_zero_max_bytes() {
984 let result = truncate_prompt("hello".into(), 0);
985 assert_eq!(result, "");
986 }
987
988 #[test]
989 fn test_truncate_prompt_long_facts() {
990 let facts: Vec<String> = (0..20)
991 .map(|i| format!("fact_{i}_{}", "x".repeat(20)))
992 .collect();
993 let prompt = facts.join("\n");
994 let result = truncate_prompt(prompt, 200);
995 assert!(
996 result.ends_with("..."),
997 "truncated prompt must end with '...'"
998 );
999 assert!(result.len() <= 203);
1001 assert!(std::str::from_utf8(result.as_bytes()).is_ok());
1002 }
1003
1004 #[test]
1005 fn test_truncate_prompt_utf8_boundary() {
1006 let prompt = "🔥".repeat(100);
1008 let result = truncate_prompt(prompt, 10);
1009 assert!(
1010 result.ends_with("..."),
1011 "truncated prompt must end with '...'"
1012 );
1013 assert_eq!(result.len(), 8 + 3, "2 emojis (8 bytes) + '...' (3 bytes)");
1015 assert!(std::str::from_utf8(result.as_bytes()).is_ok());
1016 }
1017
1018 #[tokio::test]
1021 async fn test_generate_community_summary_times_out() {
1022 tokio::time::pause();
1023 let mock = zeph_llm::mock::MockProvider::default().with_delay(20_000);
1024 let provider = AnyProvider::Mock(mock);
1025 let fut = async move {
1026 generate_community_summary(
1027 &provider,
1028 &["A".to_owned(), "B".to_owned()],
1029 &["A relates B".to_owned()],
1030 usize::MAX,
1031 )
1032 .await
1033 };
1034 let handle = tokio::spawn(fut); tokio::time::advance(Duration::from_secs(16)).await;
1036 let result = handle.await.expect("task panicked");
1037 match result {
1038 Err(MemoryError::Timeout(msg)) => assert!(
1039 msg.contains("community summary"),
1040 "unexpected timeout message: {msg}"
1041 ),
1042 other => panic!("expected MemoryError::Timeout, got {other:?}"),
1043 }
1044 }
1045
1046 #[tokio::test]
1047 async fn test_assign_to_community_majority_vote() {
1048 let store = setup().await;
1049
1050 let a = store
1052 .upsert_entity("AA", "AA", EntityType::Concept, None, None)
1053 .await
1054 .unwrap()
1055 .0;
1056 let b = store
1057 .upsert_entity("BB", "BB", EntityType::Concept, None, None)
1058 .await
1059 .unwrap()
1060 .0;
1061 let d = store
1062 .upsert_entity("DD", "DD", EntityType::Concept, None, None)
1063 .await
1064 .unwrap()
1065 .0;
1066
1067 store
1068 .upsert_community("test_community", "summary", &[a, b], None)
1069 .await
1070 .unwrap();
1071
1072 store
1073 .insert_edge(d, a, "r", "f", 1.0, None, None)
1074 .await
1075 .unwrap();
1076 store
1077 .insert_edge(d, b, "r", "f", 1.0, None, None)
1078 .await
1079 .unwrap();
1080
1081 let result = assign_to_community(&store, d).await.unwrap();
1082 assert!(result.is_some());
1083
1084 let returned_id = result.unwrap();
1086 let community = store
1087 .find_community_by_id(returned_id)
1088 .await
1089 .unwrap()
1090 .expect("returned community_id must reference an existing row");
1091 assert!(
1092 community.entity_ids.iter().any(|eid| eid.0 == d),
1093 "D should be added to the community"
1094 );
1095 assert!(
1097 community.fingerprint.is_none(),
1098 "fingerprint must be cleared after assign_to_community"
1099 );
1100 }
1101
1102 #[tokio::test]
1104 async fn test_incremental_detection_no_changes_skips_llm() {
1105 let store = setup().await;
1106 let (provider, call_buf) = recording_provider();
1107
1108 let a = store
1109 .upsert_entity("X", "X", EntityType::Concept, None, None)
1110 .await
1111 .unwrap()
1112 .0;
1113 let b = store
1114 .upsert_entity("Y", "Y", EntityType::Concept, None, None)
1115 .await
1116 .unwrap()
1117 .0;
1118 store
1119 .insert_edge(a, b, "r", "X relates Y", 1.0, None, None)
1120 .await
1121 .unwrap();
1122
1123 detect_communities(&store, &provider, usize::MAX, 4, 0)
1125 .await
1126 .unwrap();
1127 let first_calls = call_buf.lock().unwrap().len();
1128 assert_eq!(first_calls, 1, "first run must produce exactly 1 LLM call");
1129
1130 detect_communities(&store, &provider, usize::MAX, 4, 0)
1132 .await
1133 .unwrap();
1134 let second_calls = call_buf.lock().unwrap().len();
1135 assert_eq!(
1136 second_calls, first_calls,
1137 "second run with no graph changes must produce 0 additional LLM calls"
1138 );
1139 }
1140
1141 #[tokio::test]
1143 async fn test_incremental_detection_edge_change_triggers_resummary() {
1144 let store = setup().await;
1145 let (provider, call_buf) = recording_provider();
1146
1147 let a = store
1148 .upsert_entity("P", "P", EntityType::Concept, None, None)
1149 .await
1150 .unwrap()
1151 .0;
1152 let b = store
1153 .upsert_entity("Q", "Q", EntityType::Concept, None, None)
1154 .await
1155 .unwrap()
1156 .0;
1157 store
1158 .insert_edge(a, b, "r", "P relates Q", 1.0, None, None)
1159 .await
1160 .unwrap();
1161
1162 detect_communities(&store, &provider, usize::MAX, 4, 0)
1163 .await
1164 .unwrap();
1165 let after_first = call_buf.lock().unwrap().len();
1166 assert_eq!(after_first, 1);
1167
1168 store
1170 .insert_edge(b, a, "r2", "Q also relates P", 1.0, None, None)
1171 .await
1172 .unwrap();
1173
1174 detect_communities(&store, &provider, usize::MAX, 4, 0)
1175 .await
1176 .unwrap();
1177 let after_second = call_buf.lock().unwrap().len();
1178 assert_eq!(
1179 after_second, 2,
1180 "edge change must trigger one additional LLM call"
1181 );
1182 }
1183
1184 #[tokio::test]
1186 async fn test_incremental_detection_dissolved_community_deleted() {
1187 let store = setup().await;
1188 let provider = mock_provider();
1189
1190 let a = store
1191 .upsert_entity("M1", "M1", EntityType::Concept, None, None)
1192 .await
1193 .unwrap()
1194 .0;
1195 let b = store
1196 .upsert_entity("M2", "M2", EntityType::Concept, None, None)
1197 .await
1198 .unwrap()
1199 .0;
1200 let edge_id = store
1201 .insert_edge(a, b, "r", "M1 relates M2", 1.0, None, None)
1202 .await
1203 .unwrap();
1204
1205 detect_communities(&store, &provider, usize::MAX, 4, 0)
1206 .await
1207 .unwrap();
1208 assert_eq!(store.community_count().await.unwrap(), 1);
1209
1210 store.invalidate_edge(edge_id).await.unwrap();
1212
1213 detect_communities(&store, &provider, usize::MAX, 4, 0)
1214 .await
1215 .unwrap();
1216 assert_eq!(
1217 store.community_count().await.unwrap(),
1218 0,
1219 "dissolved community must be deleted on next refresh"
1220 );
1221 }
1222
1223 #[tokio::test]
1225 async fn test_detect_communities_concurrency_one() {
1226 let store = setup().await;
1227 let provider = mock_provider();
1228
1229 let a = store
1230 .upsert_entity("C1A", "C1A", EntityType::Concept, None, None)
1231 .await
1232 .unwrap()
1233 .0;
1234 let b = store
1235 .upsert_entity("C1B", "C1B", EntityType::Concept, None, None)
1236 .await
1237 .unwrap()
1238 .0;
1239 store
1240 .insert_edge(a, b, "r", "f", 1.0, None, None)
1241 .await
1242 .unwrap();
1243
1244 let count = detect_communities(&store, &provider, usize::MAX, 1, 0)
1245 .await
1246 .unwrap();
1247 assert_eq!(count, 1, "concurrency=1 must still detect the community");
1248 assert_eq!(store.community_count().await.unwrap(), 1);
1249 }
1250
1251 #[test]
1252 fn test_compute_fingerprint_deterministic() {
1253 let fp1 = compute_partition_fingerprint(&[1, 2, 3], &[10, 20]);
1254 let fp2 = compute_partition_fingerprint(&[3, 1, 2], &[20, 10]);
1255 assert_eq!(fp1, fp2, "fingerprint must be order-independent");
1256
1257 let fp3 = compute_partition_fingerprint(&[1, 2, 3], &[10, 30]);
1258 assert_ne!(
1259 fp1, fp3,
1260 "different edge IDs must produce different fingerprint"
1261 );
1262
1263 let fp4 = compute_partition_fingerprint(&[1, 2, 4], &[10, 20]);
1264 assert_ne!(
1265 fp1, fp4,
1266 "different entity IDs must produce different fingerprint"
1267 );
1268 }
1269
1270 #[test]
1275 fn test_compute_fingerprint_domain_separation() {
1276 let fp_a = compute_partition_fingerprint(&[1, 2], &[3]);
1277 let fp_b = compute_partition_fingerprint(&[1], &[2, 3]);
1278 assert_ne!(
1279 fp_a, fp_b,
1280 "entity/edge sequences with same raw bytes must produce different fingerprints"
1281 );
1282 }
1283
1284 #[tokio::test]
1290 async fn test_detect_communities_chunked_correct_membership() {
1291 let store = setup().await;
1292 let provider = mock_provider();
1293
1294 let node_alpha = store
1296 .upsert_entity("CA", "CA", EntityType::Concept, None, None)
1297 .await
1298 .unwrap()
1299 .0;
1300 let node_beta = store
1301 .upsert_entity("CB", "CB", EntityType::Concept, None, None)
1302 .await
1303 .unwrap()
1304 .0;
1305 let node_gamma = store
1306 .upsert_entity("CC", "CC", EntityType::Concept, None, None)
1307 .await
1308 .unwrap()
1309 .0;
1310 let node_delta = store
1311 .upsert_entity("CD", "CD", EntityType::Concept, None, None)
1312 .await
1313 .unwrap()
1314 .0;
1315 let node_epsilon = store
1316 .upsert_entity("CE", "CE", EntityType::Concept, None, None)
1317 .await
1318 .unwrap()
1319 .0;
1320
1321 store
1322 .insert_edge(node_alpha, node_beta, "r", "A-B fact", 1.0, None, None)
1323 .await
1324 .unwrap();
1325 store
1326 .insert_edge(node_beta, node_gamma, "r", "B-C fact", 1.0, None, None)
1327 .await
1328 .unwrap();
1329 store
1330 .insert_edge(node_delta, node_epsilon, "r", "D-E fact", 1.0, None, None)
1331 .await
1332 .unwrap();
1333
1334 let count_chunked = detect_communities(&store, &provider, usize::MAX, 4, 1)
1336 .await
1337 .unwrap();
1338 assert_eq!(
1339 count_chunked, 2,
1340 "chunked loading must detect both communities"
1341 );
1342
1343 let communities = store.all_communities().await.unwrap();
1345 assert_eq!(communities.len(), 2);
1346
1347 let abc_ids = [node_alpha, node_beta, node_gamma];
1348 let de_ids = [node_delta, node_epsilon];
1349 let has_abc = communities.iter().any(|comm| {
1350 abc_ids
1351 .iter()
1352 .all(|id| comm.entity_ids.iter().any(|eid| eid.0 == *id))
1353 });
1354 let has_de = communities.iter().any(|comm| {
1355 de_ids
1356 .iter()
1357 .all(|id| comm.entity_ids.iter().any(|eid| eid.0 == *id))
1358 });
1359 assert!(has_abc, "cluster A-B-C must form a community");
1360 assert!(has_de, "cluster D-E must form a community");
1361 }
1362
1363 #[tokio::test]
1365 async fn test_detect_communities_chunk_size_max() {
1366 let store = setup().await;
1367 let provider = mock_provider();
1368
1369 let x = store
1370 .upsert_entity("MX", "MX", EntityType::Concept, None, None)
1371 .await
1372 .unwrap()
1373 .0;
1374 let y = store
1375 .upsert_entity("MY", "MY", EntityType::Concept, None, None)
1376 .await
1377 .unwrap()
1378 .0;
1379 store
1380 .insert_edge(x, y, "r", "X-Y fact", 1.0, None, None)
1381 .await
1382 .unwrap();
1383
1384 let count = detect_communities(&store, &provider, usize::MAX, 4, usize::MAX)
1385 .await
1386 .unwrap();
1387 assert_eq!(count, 1, "chunk_size=usize::MAX must detect the community");
1388 }
1389
1390 #[tokio::test]
1392 async fn test_detect_communities_chunk_size_zero_fallback() {
1393 let store = setup().await;
1394 let provider = mock_provider();
1395
1396 let p = store
1397 .upsert_entity("ZP", "ZP", EntityType::Concept, None, None)
1398 .await
1399 .unwrap()
1400 .0;
1401 let q = store
1402 .upsert_entity("ZQ", "ZQ", EntityType::Concept, None, None)
1403 .await
1404 .unwrap()
1405 .0;
1406 store
1407 .insert_edge(p, q, "r", "P-Q fact", 1.0, None, None)
1408 .await
1409 .unwrap();
1410
1411 let count = detect_communities(&store, &provider, usize::MAX, 4, 0)
1412 .await
1413 .unwrap();
1414 assert_eq!(
1415 count, 1,
1416 "chunk_size=0 must detect the community via stream fallback"
1417 );
1418 }
1419
1420 #[tokio::test]
1424 async fn test_detect_communities_chunked_edge_map_complete() {
1425 let store = setup().await;
1426 let (provider, call_buf) = recording_provider();
1427
1428 let a = store
1429 .upsert_entity("FA", "FA", EntityType::Concept, None, None)
1430 .await
1431 .unwrap()
1432 .0;
1433 let b = store
1434 .upsert_entity("FB", "FB", EntityType::Concept, None, None)
1435 .await
1436 .unwrap()
1437 .0;
1438 store
1439 .insert_edge(a, b, "r", "edge1 fact", 1.0, None, None)
1440 .await
1441 .unwrap();
1442
1443 detect_communities(&store, &provider, usize::MAX, 4, 1)
1445 .await
1446 .unwrap();
1447 let calls_after_first = call_buf.lock().unwrap().len();
1448 assert_eq!(calls_after_first, 1, "first run must trigger 1 LLM call");
1449
1450 store
1452 .insert_edge(b, a, "r2", "edge2 fact", 1.0, None, None)
1453 .await
1454 .unwrap();
1455
1456 detect_communities(&store, &provider, usize::MAX, 4, 1)
1457 .await
1458 .unwrap();
1459 let calls_after_second = call_buf.lock().unwrap().len();
1460 assert_eq!(
1461 calls_after_second, 2,
1462 "adding an edge must change fingerprint and trigger re-summarization"
1463 );
1464 }
1465
1466 #[tokio::test]
1468 async fn cleanup_stale_empty_collection() {
1469 let store = setup().await;
1470 let sqlite_store = crate::store::SqliteStore::new(":memory:").await.unwrap();
1471 let pool = sqlite_store.pool().clone();
1472 let mem_store = Box::new(crate::in_memory_store::InMemoryVectorStore::new());
1473 let emb_store = crate::embedding_store::EmbeddingStore::with_store(mem_store, pool);
1474 emb_store
1475 .ensure_named_collection("zeph_graph_entities", 4)
1476 .await
1477 .unwrap();
1478
1479 let deleted = cleanup_stale_entity_embeddings(&store, &emb_store)
1480 .await
1481 .unwrap();
1482 assert_eq!(deleted, 0, "nothing to delete from empty collection");
1483 }
1484
1485 #[tokio::test]
1488 async fn cleanup_stale_deletes_orphaned_points() {
1489 use crate::graph::types::EntityType;
1490
1491 let sqlite_store = crate::store::SqliteStore::new(":memory:").await.unwrap();
1492 let pool = sqlite_store.pool().clone();
1493 let graph_store = GraphStore::new(pool.clone());
1494
1495 let mem_store = Box::new(crate::in_memory_store::InMemoryVectorStore::new());
1496 let emb_store = crate::embedding_store::EmbeddingStore::with_store(mem_store, pool.clone());
1497 emb_store
1498 .ensure_named_collection("zeph_graph_entities", 4)
1499 .await
1500 .unwrap();
1501
1502 let live_id = graph_store
1504 .upsert_entity("Live", "live", EntityType::Person, None, None)
1505 .await
1506 .unwrap()
1507 .0;
1508 let stale_id = graph_store
1509 .upsert_entity("Stale", "stale", EntityType::Person, None, None)
1510 .await
1511 .unwrap()
1512 .0;
1513
1514 let live_payload = serde_json::json!({
1516 "entity_id": live_id,
1517 "entity_id_str": live_id.to_string(),
1518 "name": "Live",
1519 });
1520 let stale_payload = serde_json::json!({
1521 "entity_id": stale_id,
1522 "entity_id_str": stale_id.to_string(),
1523 "name": "Stale",
1524 });
1525 emb_store
1526 .store_to_collection(
1527 "zeph_graph_entities",
1528 live_payload,
1529 vec![1.0, 0.0, 0.0, 0.0],
1530 )
1531 .await
1532 .unwrap();
1533 emb_store
1534 .store_to_collection(
1535 "zeph_graph_entities",
1536 stale_payload,
1537 vec![0.0, 1.0, 0.0, 0.0],
1538 )
1539 .await
1540 .unwrap();
1541
1542 zeph_db::query(zeph_db::sql!("DELETE FROM graph_entities WHERE id = ?"))
1544 .bind(stale_id)
1545 .execute(&pool)
1546 .await
1547 .unwrap();
1548
1549 let deleted = cleanup_stale_entity_embeddings(&graph_store, &emb_store)
1550 .await
1551 .unwrap();
1552 assert_eq!(deleted, 1, "exactly one stale point should be removed");
1553
1554 let remaining = emb_store
1556 .scroll_all_entity_ids("zeph_graph_entities")
1557 .await
1558 .unwrap();
1559 assert_eq!(remaining.len(), 1);
1560 assert_eq!(remaining[0].1, live_id);
1561 }
1562}