From 77fd5f9eddef6b0c389cad8eb1eb78811c5ec332 Mon Sep 17 00:00:00 2001 From: sanil-23 Date: Wed, 1 Apr 2026 01:07:03 +0530 Subject: [PATCH] fix: propagate entity_type into retrieval context for graph visualization (#135) * feat(memory): connect graph query and doc ingest APIs to frontend Wire up memory.graph.query and memory.doc.ingest RPC endpoints and integrate graph relations into the MemoryWorkspace UI, replacing backend-only entity counts with local graph store data. Co-Authored-By: Claude Opus 4.6 (1M context) * style: fix Prettier formatting in MemoryWorkspace and tauriCommands Co-Authored-By: Claude Opus 4.6 (1M context) * test: add unit tests for memory graph query, doc ingest, and dispatch routing Cover the new graph query and doc ingest APIs added in this PR: - Rust dispatch tests: routing, param validation, unknown method fallthrough - Frontend tauriCommands tests: Tauri guard, RPC forwarding for memoryGraphQuery/memoryDocIngest - MemoryWorkspace component tests: graph relations rendering, evidence badges, empty states Co-Authored-By: Claude Opus 4.6 (1M context) * style: fix Prettier and cargo fmt formatting in test files Co-Authored-By: Claude Opus 4.6 (1M context) * style: fix Prettier formatting in tauriCommandsMemory test Co-Authored-By: Claude Opus 4.6 (1M context) * fix: merge duplicate tauriCommands import in MemoryWorkspace Co-Authored-By: Claude Opus 4.6 (1M context) * style: fix Prettier formatting in MemoryWorkspace Co-Authored-By: Claude Opus 4.6 (1M context) * fix: propagate entity_type from graph relation attrs into retrieval context build_retrieval_context was discarding entity types that were already present in relation attrs.entity_types (populated during ingestion by GLiNER relex). Entities in MemoryRetrievalEntity now carry their type (e.g. PERSON, PROJECT, WORK_ITEM) instead of always being None. Adds unit tests for both typed and untyped paths, plus an ignored GLiNER smoke test (gline_rs_smoke) that verifies the full pipeline from Notion fixture ingestion through graph storage to retrieval context with the real ONNX model. Co-Authored-By: Claude Opus 4.6 (1M context) --------- Co-authored-by: sanil jain Co-authored-by: Claude Opus 4.6 (1M context) --- src/openhuman/memory/ingestion.rs | 139 ++++++++++++++++++++++++++++++ src/openhuman/memory/ops.rs | 96 +++++++++++++++++++-- 2 files changed, 228 insertions(+), 7 deletions(-) diff --git a/src/openhuman/memory/ingestion.rs b/src/openhuman/memory/ingestion.rs index d5e26caf7..e1b3cda6b 100644 --- a/src/openhuman/memory/ingestion.rs +++ b/src/openhuman/memory/ingestion.rs @@ -1784,4 +1784,143 @@ mod tests { .iter() .any(|hit| !hit.supporting_relations.is_empty())); } + + /// Smoke test using the real GLiNER relex ONNX model with the Notion fixture. + /// Verifies that entity types extracted by the model flow through ingestion + /// into graph relations (attrs.entity_types) and into retrieval context + /// (MemoryRetrievalEntity.entity_type) via build_retrieval_context. + /// + /// Run: cargo test -p openhuman --lib gline_rs_smoke -- --ignored --nocapture + #[tokio::test] + #[ignore] // requires GLiNER ONNX model on disk + async fn gline_rs_smoke_notion_entity_types_flow_through() { + use crate::openhuman::memory::ops::build_retrieval_context; + + let tmp = TempDir::new().unwrap(); + let memory = UnifiedMemory::new(tmp.path(), Arc::new(NoopEmbedding), None).unwrap(); + + // Use default config so the real GLiNER model runs + let result = memory + .ingest_document(MemoryIngestionRequest { + document: NamespaceDocumentInput { + namespace: "skill-notion".to_string(), + key: "notion-roadmap".to_string(), + title: "OpenHuman Memory Layer Roadmap".to_string(), + content: fixture("notion_page_example.txt"), + source_type: "notion".to_string(), + priority: "high".to_string(), + tags: Vec::new(), + metadata: json!({}), + category: "core".to_string(), + session_id: None, + document_id: None, + }, + config: MemoryIngestionConfig::default(), + }) + .await + .unwrap(); + + // 1. Verify GLiNER extracted entities with types + println!("--- Extracted entities ({}) ---", result.entities.len()); + for entity in &result.entities { + println!(" {} [{}]", entity.name, entity.entity_type); + } + assert!( + !result.entities.is_empty(), + "GLiNER should extract at least some entities" + ); + assert!( + result.entities.iter().any(|e| !e.entity_type.is_empty()), + "at least one entity should have a non-empty entity_type" + ); + + // 2. Verify relations carry subject_type / object_type + println!("--- Extracted relations ({}) ---", result.relations.len()); + for rel in &result.relations { + println!( + " {} [{}] -[{}]-> {} [{}] (conf={:.2})", + rel.subject, + rel.subject_type, + rel.predicate, + rel.object, + rel.object_type, + rel.confidence + ); + } + let typed_relations = result + .relations + .iter() + .filter(|r| !r.subject_type.is_empty() && !r.object_type.is_empty()) + .count(); + assert!( + typed_relations > 0, + "at least one relation should have typed subject and object" + ); + + // 3. Verify graph relations have entity_types in attrs + let graph_rows = memory + .graph_query_namespace("skill-notion", None, None) + .await + .unwrap(); + println!("--- Graph relations ({}) ---", graph_rows.len()); + let mut graph_has_entity_types = false; + for row in &graph_rows { + let et = row.get("attrs").and_then(|a| a.get("entity_types")); + if let Some(et) = et { + println!( + " {} -> {} -> {} entity_types={}", + row["subject"], row["predicate"], row["object"], et + ); + graph_has_entity_types = true; + } + } + assert!( + graph_has_entity_types, + "at least one graph relation should have attrs.entity_types" + ); + + // 4. Verify build_retrieval_context propagates entity_type from query hits + let context_data = memory + .query_namespace_context_data("skill-notion", "who owns what", 10) + .await + .unwrap(); + let retrieval = build_retrieval_context(&context_data.hits); + println!("--- Retrieval entities ({}) ---", retrieval.entities.len()); + for entity in &retrieval.entities { + println!( + " {} [type={:?}]", + entity.name, + entity.entity_type.as_deref().unwrap_or("None") + ); + } + let typed_entities = retrieval + .entities + .iter() + .filter(|e| e.entity_type.is_some()) + .count(); + println!( + "Typed entities: {}/{}", + typed_entities, + retrieval.entities.len() + ); + // If there are any supporting relations with entity_types, then + // build_retrieval_context should have picked them up. + let has_typed_supporting_relations = context_data + .hits + .iter() + .flat_map(|hit| hit.supporting_relations.iter()) + .any(|rel| { + rel.attrs + .get("entity_types") + .and_then(|et| et.get("subject")) + .and_then(|v| v.as_str()) + .is_some_and(|s| !s.is_empty()) + }); + if has_typed_supporting_relations { + assert!( + typed_entities > 0, + "build_retrieval_context should propagate entity_type from supporting relations" + ); + } + } } diff --git a/src/openhuman/memory/ops.rs b/src/openhuman/memory/ops.rs index 13970ada2..b652a6381 100644 --- a/src/openhuman/memory/ops.rs +++ b/src/openhuman/memory/ops.rs @@ -148,18 +148,33 @@ fn chunk_metadata(hit: &NamespaceMemoryHit) -> Value { }) } -fn build_retrieval_context(hits: &[NamespaceMemoryHit]) -> MemoryRetrievalContext { - let mut entities = BTreeSet::new(); +fn extract_entity_type(attrs: &Value, role: &str) -> Option { + attrs + .get("entity_types") + .and_then(|et| et.get(role)) + .and_then(|v| v.as_str()) + .filter(|s| !s.is_empty()) + .map(|s| s.to_string()) +} + +pub(crate) fn build_retrieval_context(hits: &[NamespaceMemoryHit]) -> MemoryRetrievalContext { + let mut entity_types: BTreeMap> = BTreeMap::new(); let mut relations = BTreeMap::new(); let chunks = hits .iter() .map(|hit| { for relation in &hit.supporting_relations { if !relation.subject.trim().is_empty() { - entities.insert(relation.subject.clone()); + let entry = entity_types.entry(relation.subject.clone()).or_insert(None); + if entry.is_none() { + *entry = extract_entity_type(&relation.attrs, "subject"); + } } if !relation.object.trim().is_empty() { - entities.insert(relation.object.clone()); + let entry = entity_types.entry(relation.object.clone()).or_insert(None); + if entry.is_none() { + *entry = extract_entity_type(&relation.attrs, "object"); + } } relations .entry(relation_identity(relation)) @@ -186,12 +201,12 @@ fn build_retrieval_context(hits: &[NamespaceMemoryHit]) -> MemoryRetrievalContex .collect(); MemoryRetrievalContext { - entities: entities + entities: entity_types .into_iter() - .map(|name| MemoryRetrievalEntity { + .map(|(name, entity_type)| MemoryRetrievalEntity { id: None, name, - entity_type: None, + entity_type, score: None, metadata: json!({}), }) @@ -1028,6 +1043,73 @@ mod tests { assert_eq!(context.relations[0].predicate, "OWNS"); } + fn sample_hit_with_entity_types() -> NamespaceMemoryHit { + NamespaceMemoryHit { + id: "doc-2".to_string(), + kind: MemoryItemKind::Document, + namespace: "team".to_string(), + key: "atlas-status".to_string(), + title: Some("Atlas status".to_string()), + content: "Project Atlas is owned by Alice.".to_string(), + category: "core".to_string(), + source_type: Some("doc".to_string()), + updated_at: 1_700_000_000.0, + score: 0.92, + score_breakdown: RetrievalScoreBreakdown { + keyword_relevance: 0.3, + vector_similarity: 0.4, + graph_relevance: 0.9, + freshness: 0.0, + final_score: 0.92, + }, + document_id: Some("doc-2".to_string()), + chunk_id: Some("doc-2#chunk-1".to_string()), + supporting_relations: vec![GraphRelationRecord { + namespace: Some("team".to_string()), + subject: "Alice".to_string(), + predicate: "OWNS".to_string(), + object: "Atlas".to_string(), + attrs: json!({ + "source": "ingestion", + "entity_types": { + "subject": "PERSON", + "object": "PROJECT" + } + }), + updated_at: 1_700_000_000.0, + evidence_count: 2, + order_index: Some(1), + document_ids: vec!["doc-2".to_string()], + chunk_ids: vec!["doc-2#chunk-1".to_string()], + }], + } + } + + #[test] + fn build_retrieval_context_extracts_entity_types_from_attrs() { + let context = build_retrieval_context(&[sample_hit_with_entity_types()]); + assert_eq!(context.entities.len(), 2); + + let alice = context.entities.iter().find(|e| e.name == "Alice").unwrap(); + assert_eq!(alice.entity_type.as_deref(), Some("PERSON")); + + let atlas = context.entities.iter().find(|e| e.name == "Atlas").unwrap(); + assert_eq!(atlas.entity_type.as_deref(), Some("PROJECT")); + } + + #[test] + fn build_retrieval_context_entity_type_none_when_attrs_missing() { + let context = build_retrieval_context(&[sample_hit()]); + assert_eq!(context.entities.len(), 2); + + for entity in &context.entities { + assert_eq!( + entity.entity_type, None, + "entity_type should be None when attrs has no entity_types" + ); + } + } + #[test] fn helpers_filter_document_ids_and_format_context_message() { let hit = sample_hit();