mirror of
https://github.com/tinyhumansai/openhuman.git
synced 2026-07-27 21:08:00 +00:00
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) <noreply@anthropic.com> * style: fix Prettier formatting in MemoryWorkspace and tauriCommands Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * 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) <noreply@anthropic.com> * style: fix Prettier and cargo fmt formatting in test files Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * style: fix Prettier formatting in tauriCommandsMemory test Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * fix: merge duplicate tauriCommands import in MemoryWorkspace Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * style: fix Prettier formatting in MemoryWorkspace Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * 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) <noreply@anthropic.com> --------- Co-authored-by: sanil jain <jainsanil18@gmail.com> Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
sanil-23
co-authored by
sanil jain
Claude Opus 4.6
sanil jain
Claude Opus 4.6
parent
906b55b63e
commit
77fd5f9edd
@@ -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"
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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<String> {
|
||||
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<String, Option<String>> = 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();
|
||||
|
||||
Reference in New Issue
Block a user