diff --git a/Cargo.lock b/Cargo.lock index f4cccd48b..263cb2c55 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1888,6 +1888,9 @@ name = "esaxx-rs" version = "0.1.10" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d817e038c30374a4bcb22f94d0a8a0e216958d4c3dcde369b1439fec4bdda6e6" +dependencies = [ + "cc", +] [[package]] name = "esp-idf-part" @@ -2074,7 +2077,7 @@ dependencies = [ "safetensors", "serde", "serde_json", - "tokenizers", + "tokenizers 0.22.2", ] [[package]] @@ -3446,6 +3449,16 @@ dependencies = [ "cc", ] +[[package]] +name = "libloading" +version = "0.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "754ca22de805bb5744484a5b151a9e1a8e837d5dc232c2d7d8c2e3492edc8b60" +dependencies = [ + "cfg-if", + "windows-link", +] + [[package]] name = "libm" version = "0.2.16" @@ -4590,6 +4603,7 @@ dependencies = [ "log", "mail-parser", "matrix-sdk", + "ndarray", "nu-ansi-term 0.46.0", "nusb 0.2.3", "once_cell", @@ -4598,6 +4612,7 @@ dependencies = [ "opentelemetry", "opentelemetry-otlp", "opentelemetry_sdk", + "ort", "parking_lot", "pdf-extract", "postgres", @@ -4623,6 +4638,7 @@ dependencies = [ "socketioxide", "tempfile", "thiserror 2.0.18", + "tokenizers 0.21.4", "tokio", "tokio-rustls", "tokio-serial", @@ -4782,6 +4798,7 @@ version = "2.0.0-rc.11" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "4a5df903c0d2c07b56950f1058104ab0c8557159f2741782223704de9be73c3c" dependencies = [ + "libloading", "ndarray", "ort-sys", "smallvec", @@ -7133,6 +7150,40 @@ version = "0.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1f3ccbac311fea05f86f61904b462b55fb3df8837a366dfc601a0161d0532f20" +[[package]] +name = "tokenizers" +version = "0.21.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a620b996116a59e184c2fa2dfd8251ea34a36d0a514758c6f966386bd2e03476" +dependencies = [ + "ahash", + "aho-corasick", + "compact_str", + "dary_heap", + "derive_builder", + "esaxx-rs", + "getrandom 0.3.4", + "indicatif", + "itertools 0.14.0", + "log", + "macro_rules_attribute", + "monostate", + "onig", + "paste", + "rand 0.9.2", + "rayon", + "rayon-cond", + "regex", + "regex-syntax", + "serde", + "serde_json", + "spm_precompiled", + "thiserror 2.0.18", + "unicode-normalization-alignments", + "unicode-segmentation", + "unicode_categories", +] + [[package]] name = "tokenizers" version = "0.22.2" diff --git a/Cargo.toml b/Cargo.toml index 0a6b7dd28..86a817817 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -37,6 +37,7 @@ async-trait = "0.1" chacha20poly1305 = "0.10" hex = "0.4" fastembed = "5.13" +ort = { version = "=2.0.0-rc.11", features = ["load-dynamic"] } tokio-util = { version = "0.7", features = ["rt"] } tokio-tungstenite = { version = "0.24", features = ["rustls-tls-webpki-roots"] } futures = "0.3" @@ -62,6 +63,8 @@ dialoguer = { version = "0.12", features = ["fuzzy-select"] } console = "0.16" glob = "0.3" regex = "1.10" +ndarray = "0.17.2" +tokenizers = "0.21.0" hostname = "0.4.2" rustyline = { version = "15", features = ["with-file-history"] } rustls = { version = "0.23", features = ["ring"] } diff --git a/src/openhuman/memory/ingestion.rs b/src/openhuman/memory/ingestion.rs new file mode 100644 index 000000000..23461835c --- /dev/null +++ b/src/openhuman/memory/ingestion.rs @@ -0,0 +1,1777 @@ +use std::collections::{BTreeMap, BTreeSet, HashMap}; +use std::sync::OnceLock; + +use regex::Regex; +use serde::{Deserialize, Serialize}; +use serde_json::{json, Map, Value}; + +use super::relex; +use crate::openhuman::memory::store::types::NamespaceDocumentInput; +use crate::openhuman::memory::UnifiedMemory; + +pub const DEFAULT_GLINER_RELEX_MODEL: &str = "knowledgator/gliner-relex-large-v0.5"; +const DEFAULT_CHUNK_TOKENS: usize = 225; + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum ExtractionMode { + Sentence, + Chunk, +} + +impl Default for ExtractionMode { + fn default() -> Self { + Self::Sentence + } +} + +impl ExtractionMode { + fn as_str(self) -> &'static str { + match self { + Self::Sentence => "sentence", + Self::Chunk => "chunk", + } + } +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct MemoryIngestionConfig { + pub model_name: String, + #[serde(default)] + pub extraction_mode: ExtractionMode, + #[serde(default = "default_entity_threshold")] + pub entity_threshold: f32, + #[serde(default = "default_relation_threshold")] + pub relation_threshold: f32, + #[serde(default = "default_adjacency_threshold")] + pub adjacency_threshold: f32, + #[serde(default = "default_batch_size")] + pub batch_size: usize, +} + +fn default_entity_threshold() -> f32 { + 0.45 +} + +fn default_relation_threshold() -> f32 { + 0.30 +} + +fn default_adjacency_threshold() -> f32 { + 0.50 +} + +fn default_batch_size() -> usize { + 16 +} + +impl Default for MemoryIngestionConfig { + fn default() -> Self { + Self { + model_name: DEFAULT_GLINER_RELEX_MODEL.to_string(), + extraction_mode: ExtractionMode::Sentence, + entity_threshold: default_entity_threshold(), + relation_threshold: default_relation_threshold(), + adjacency_threshold: default_adjacency_threshold(), + batch_size: default_batch_size(), + } + } +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct MemoryIngestionRequest { + pub document: NamespaceDocumentInput, + #[serde(default)] + pub config: MemoryIngestionConfig, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct ExtractedEntity { + pub name: String, + pub entity_type: String, + #[serde(default)] + pub aliases: Vec, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct ExtractedRelation { + pub subject: String, + pub subject_type: String, + pub predicate: String, + pub object: String, + pub object_type: String, + pub confidence: f32, + pub evidence_count: u32, + pub chunk_ids: Vec, + pub order_index: Option, + pub metadata: Value, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct MemoryIngestionResult { + pub document_id: String, + pub namespace: String, + pub model_name: String, + pub extraction_mode: String, + pub chunk_count: usize, + pub entity_count: usize, + pub relation_count: usize, + pub preference_count: usize, + pub decision_count: usize, + #[serde(default)] + pub tags: Vec, + #[serde(default)] + pub entities: Vec, + #[serde(default)] + pub relations: Vec, +} + +#[derive(Debug, Clone)] +struct RawEntity { + name: String, + entity_type: String, + confidence: f32, +} + +#[derive(Debug, Clone)] +struct RawRelation { + subject: String, + subject_type: String, + predicate: String, + object: String, + object_type: String, + confidence: f32, + chunk_indexes: BTreeSet, + order_index: i64, + metadata: Map, +} + +#[derive(Debug, Clone)] +struct ExtractionUnit { + text: String, + chunk_index: usize, + order_index: i64, +} + +#[derive(Debug, Default)] +struct ExtractionAccumulator { + entities: HashMap, + relations: Vec, + tags: BTreeSet, + decisions: BTreeSet, + preferences: BTreeSet, + doc_kind: Option, + primary_subject: Option, + document_title: Option, + current_subject: Option, + current_sender: Option, + known_people: HashMap, +} + +#[derive(Debug)] +struct ParsedIngestion { + tags: Vec, + metadata: Value, + entities: Vec, + relations: Vec, + chunk_count: usize, + preference_count: usize, + decision_count: usize, +} + +#[derive(Debug)] +struct RelationRule { + canonical: &'static str, + allowed_head: &'static [&'static str], + allowed_tail: &'static [&'static str], +} + +const PERSON_TYPES: &[&str] = &["PERSON"]; +const ORG_TYPES: &[&str] = &[ + "ORGANIZATION", + "PROJECT", + "PRODUCT", + "TOOL", + "TOPIC", + "WORK_ITEM", +]; +const PLACE_TYPES: &[&str] = &["PLACE", "LOCATION", "ROOM"]; +const DATE_TYPES: &[&str] = &["DATE"]; + +fn relation_rule(predicate: &str) -> Option { + let normalized = UnifiedMemory::normalize_graph_predicate(predicate); + let rule = match normalized.as_str() { + "OWNS" | "WORKS_ON" | "RESPONSIBLE_FOR" | "REVIEWS" => RelationRule { + canonical: "OWNS", + allowed_head: PERSON_TYPES, + allowed_tail: ORG_TYPES, + }, + "USES" | "KEEPS" | "ADOPTS" => RelationRule { + canonical: "USES", + allowed_head: ORG_TYPES, + allowed_tail: ORG_TYPES, + }, + "WORKS_FOR" => RelationRule { + canonical: "WORKS_FOR", + allowed_head: PERSON_TYPES, + allowed_tail: &["ORGANIZATION", "PROJECT", "PRODUCT"], + }, + "DEPENDS_ON" => RelationRule { + canonical: "DEPENDS_ON", + allowed_head: ORG_TYPES, + allowed_tail: ORG_TYPES, + }, + "PREFERS" => RelationRule { + canonical: "PREFERS", + allowed_head: PERSON_TYPES, + allowed_tail: &["TOPIC", "WORK_ITEM", "MODE", "PRODUCT", "TOOL"], + }, + "HAS_DEADLINE" | "DUE_ON" => RelationRule { + canonical: "HAS_DEADLINE", + allowed_head: ORG_TYPES, + allowed_tail: DATE_TYPES, + }, + "COMMUNICATES_WITH" => RelationRule { + canonical: "COMMUNICATES_WITH", + allowed_head: PERSON_TYPES, + allowed_tail: PERSON_TYPES, + }, + "INVESTIGATES" | "EVALUATES" => RelationRule { + canonical: "INVESTIGATES", + allowed_head: PERSON_TYPES, + allowed_tail: ORG_TYPES, + }, + "NORTH_OF" => RelationRule { + canonical: "NORTH_OF", + allowed_head: PLACE_TYPES, + allowed_tail: PLACE_TYPES, + }, + "SOUTH_OF" => RelationRule { + canonical: "SOUTH_OF", + allowed_head: PLACE_TYPES, + allowed_tail: PLACE_TYPES, + }, + "EAST_OF" => RelationRule { + canonical: "EAST_OF", + allowed_head: PLACE_TYPES, + allowed_tail: PLACE_TYPES, + }, + "WEST_OF" => RelationRule { + canonical: "WEST_OF", + allowed_head: PLACE_TYPES, + allowed_tail: PLACE_TYPES, + }, + "AVOIDS" => RelationRule { + canonical: "AVOIDS", + allowed_head: ORG_TYPES, + allowed_tail: ORG_TYPES, + }, + _ => return None, + }; + Some(rule) +} + +fn type_allowed(actual: &str, allowed: &[&str]) -> bool { + allowed.is_empty() || allowed.iter().any(|candidate| candidate == &actual) +} + +fn resolve_person_alias(name: &str, known_people: &HashMap) -> String { + let upper = name.to_uppercase(); + known_people.get(&upper).cloned().unwrap_or_else(|| upper) +} + +impl ExtractionAccumulator { + fn remember_person_aliases(&mut self, canonical_name: &str) { + let parts = canonical_name.split_whitespace().collect::>(); + if let Some(first_name) = parts.first() { + self.known_people + .entry(first_name.to_uppercase()) + .or_insert_with(|| canonical_name.to_string()); + } + } + + fn add_entity(&mut self, name: &str, entity_type: &str, confidence: f32) -> Option { + let cleaned = sanitize_entity_name(name); + if cleaned.is_empty() { + return None; + } + let resolved_name = if entity_type == "PERSON" { + resolve_person_alias(&cleaned, &self.known_people) + } else { + cleaned.clone() + }; + let entry = self + .entities + .entry(resolved_name.clone()) + .or_insert_with(|| RawEntity { + name: resolved_name.clone(), + entity_type: entity_type.to_string(), + confidence, + }); + if confidence > entry.confidence { + entry.confidence = confidence; + } + if entity_type == "PERSON" { + self.remember_person_aliases(&resolved_name); + } + Some(resolved_name) + } + + fn add_relation( + &mut self, + subject: &str, + subject_type: &str, + predicate: &str, + object: &str, + object_type: &str, + confidence: f32, + chunk_index: usize, + order_index: i64, + metadata: Map, + ) { + let Some(rule) = relation_rule(predicate) else { + return; + }; + let Some(subject_name) = self.add_entity(subject, subject_type, confidence) else { + return; + }; + let Some(object_name) = self.add_entity(object, object_type, confidence) else { + return; + }; + if subject_name == object_name { + return; + } + let actual_subject_type = self + .entities + .get(&subject_name) + .map(|value| value.entity_type.as_str()) + .unwrap_or(subject_type); + let actual_object_type = self + .entities + .get(&object_name) + .map(|value| value.entity_type.as_str()) + .unwrap_or(object_type); + if !type_allowed(actual_subject_type, rule.allowed_head) + || !type_allowed(actual_object_type, rule.allowed_tail) + { + return; + } + + let mut chunk_indexes = BTreeSet::new(); + chunk_indexes.insert(chunk_index); + self.relations.push(RawRelation { + subject: subject_name, + subject_type: actual_subject_type.to_string(), + predicate: rule.canonical.to_string(), + object: object_name, + object_type: actual_object_type.to_string(), + confidence, + chunk_indexes, + order_index, + metadata, + }); + } +} + +fn email_header_regex() -> &'static Regex { + static REGEX: OnceLock = OnceLock::new(); + REGEX + .get_or_init(|| Regex::new(r"^(From|To|Cc):\s*(?P.+)$").expect("email header regex")) +} + +fn named_email_regex() -> &'static Regex { + static REGEX: OnceLock = OnceLock::new(); + REGEX.get_or_init(|| { + Regex::new(r"(?P[^,<]+?)\s*<(?P[^>]+)>").expect("named email regex") + }) +} + +fn graph_fact_regex() -> &'static Regex { + static REGEX: OnceLock = OnceLock::new(); + REGEX.get_or_init(|| { + Regex::new( + r"^(?P[A-Za-z0-9][A-Za-z0-9 ._\-/]+?)\s+(?Pworks_on|depends_on|uses|evaluates|owns|prefers)\s+(?P.+)$", + ) + .expect("graph fact regex") + }) +} + +fn explicit_owner_regex() -> &'static Regex { + static REGEX: OnceLock = OnceLock::new(); + REGEX.get_or_init(|| { + Regex::new(r"^(?P[A-Za-z][A-Za-z ._-]+?) owns (?P.+)$") + .expect("explicit owner regex") + }) +} + +fn explicit_preference_regex() -> &'static Regex { + static REGEX: OnceLock = OnceLock::new(); + REGEX.get_or_init(|| { + Regex::new(r"^(?P[A-Za-z][A-Za-z ._-]+?) prefers (?P.+)$") + .expect("explicit preference regex") + }) +} + +fn action_item_regex() -> &'static Regex { + static REGEX: OnceLock = OnceLock::new(); + REGEX.get_or_init(|| { + Regex::new(r"^(?P[A-Za-z][A-Za-z ._-]+?):\s*(?P.+)$") + .expect("action item regex") + }) +} + +fn will_review_regex() -> &'static Regex { + static REGEX: OnceLock = OnceLock::new(); + REGEX.get_or_init(|| { + Regex::new(r"^(?P[A-Za-z][A-Za-z ._-]+?) will review (?P.+)$") + .expect("will review regex") + }) +} + +fn recipient_regex() -> &'static Regex { + static REGEX: OnceLock = OnceLock::new(); + REGEX.get_or_init(|| { + Regex::new( + r"(?i)(?P[A-Z][A-Za-z]+(?: [A-Z][A-Za-z]+)?)\s+(gave|sent|handed|passed)\s+(?P.+?)\s+to\s+(?P[A-Z][A-Za-z]+(?: [A-Z][A-Za-z]+)?)", + ) + .expect("recipient regex") + }) +} + +fn spatial_regex() -> &'static Regex { + static REGEX: OnceLock = OnceLock::new(); + REGEX.get_or_init(|| { + Regex::new( + r"(?i)(?P[A-Za-z][A-Za-z0-9 _-]+?)\s+is\s+(?Pnorth|south|east|west)\s+of\s+(?P[A-Za-z][A-Za-z0-9 _-]+)", + ) + .expect("spatial regex") + }) +} + +fn month_date_regex() -> &'static Regex { + static REGEX: OnceLock = OnceLock::new(); + REGEX.get_or_init(|| { + Regex::new(r"(?i)\b(?:Jan|Feb|Mar|Apr|May|Jun|Jul|Aug|Sep|Sept|Oct|Nov|Dec)[a-z]*\s+\d{1,2},\s+\d{4}\b") + .expect("month date regex") + }) +} + +fn iso_date_regex() -> &'static Regex { + static REGEX: OnceLock = OnceLock::new(); + REGEX.get_or_init(|| Regex::new(r"\b\d{4}-\d{2}-\d{2}\b").expect("iso date regex")) +} + +fn person_name_regex() -> &'static Regex { + static REGEX: OnceLock = OnceLock::new(); + REGEX + .get_or_init(|| Regex::new(r"\b[A-Z][a-z]+(?: [A-Z][a-z]+)+\b").expect("person name regex")) +} + +fn sanitize_entity_name(name: &str) -> String { + let trimmed = name.trim().trim_matches(|ch: char| { + matches!(ch, '-' | ':' | ';' | ',' | '.' | '"' | '\'' | '(' | ')') + }); + if trimmed.is_empty() { + return String::new(); + } + UnifiedMemory::collapse_whitespace(trimmed).to_uppercase() +} + +fn sanitize_fact_text(text: &str) -> String { + let trimmed = text + .trim() + .trim_start_matches('-') + .trim() + .trim_matches(|ch: char| matches!(ch, ':' | ';' | ',' | '.')); + UnifiedMemory::collapse_whitespace(trimmed) +} + +fn classify_entity(name: &str, known_people: &HashMap) -> &'static str { + let upper = sanitize_entity_name(name); + if upper.is_empty() { + return "TOPIC"; + } + if month_date_regex().is_match(name) || iso_date_regex().is_match(name) { + return "DATE"; + } + if upper.contains('@') { + return "ORGANIZATION"; + } + if known_people.contains_key(&upper) || person_name_regex().is_match(name) { + return "PERSON"; + } + if matches!( + upper.as_str(), + "OPENHUMAN" | "JSON-RPC" | "JSON-RPC 2.0" | "NEOCORTEX_V2" | "NEOCORTEX V2" + ) { + return "PRODUCT"; + } + if upper.contains("MODEL") { + return "TOOL"; + } + if upper.contains("MODE") { + return "MODE"; + } + if upper.contains("MILESTONE") + || upper.contains("ROADMAP") + || upper.contains("CONTRACT") + || upper.contains("API") + || upper.contains("MEMORY") + || upper.contains("FIXTURE") + || upper.contains("THREAD") + || upper.contains("WORK") + { + return "WORK_ITEM"; + } + if upper.contains("OFFICE") + || upper.contains("ROOM") + || upper.contains("GARDEN") + || upper.contains("KITCHEN") + { + return "ROOM"; + } + if upper.contains("TINYHUMANS") || upper.ends_with("CORE") { + return "ORGANIZATION"; + } + if (upper.contains('-') || upper.contains('_')) && !upper.contains(' ') { + return "PROJECT"; + } + "TOPIC" +} + +fn split_sentences(text: &str) -> Vec { + let mut out = Vec::new(); + let mut current = String::new(); + for ch in text.chars() { + current.push(ch); + if matches!(ch, '.' | '!' | '?' | '\n') { + let candidate = sanitize_fact_text(¤t); + if !candidate.is_empty() { + out.push(candidate); + } + current.clear(); + } + } + let tail = sanitize_fact_text(¤t); + if !tail.is_empty() { + out.push(tail); + } + let mut merged: Vec = Vec::new(); + for sentence in out { + if sentence.len() < 5 && !merged.is_empty() { + if let Some(last) = merged.last_mut() { + last.push(' '); + last.push_str(&sentence); + } + } else { + merged.push(sentence); + } + } + if merged.is_empty() && !text.trim().is_empty() { + merged.push(sanitize_fact_text(text)); + } + merged +} + +fn build_units(chunks: &[String], mode: ExtractionMode) -> Vec { + let mut units = Vec::new(); + let mut order_index = 0_i64; + for (chunk_index, chunk) in chunks.iter().enumerate() { + match mode { + ExtractionMode::Chunk => { + let text = sanitize_fact_text(chunk); + if text.is_empty() { + continue; + } + units.push(ExtractionUnit { + text, + chunk_index, + order_index, + }); + order_index += 1; + } + ExtractionMode::Sentence => { + for sentence in split_sentences(chunk) { + if sentence.is_empty() { + continue; + } + units.push(ExtractionUnit { + text: sentence, + chunk_index, + order_index, + }); + order_index += 1; + } + } + } + } + units +} + +fn find_chunk_index(chunks: &[String], excerpt: &str, hint: usize) -> usize { + if chunks.is_empty() { + return 0; + } + let needle = UnifiedMemory::normalize_search_text(excerpt); + if needle.is_empty() { + return hint.min(chunks.len().saturating_sub(1)); + } + for index in hint..chunks.len() { + if UnifiedMemory::normalize_search_text(&chunks[index]).contains(&needle) { + return index; + } + } + for index in 0..hint.min(chunks.len()) { + if UnifiedMemory::normalize_search_text(&chunks[index]).contains(&needle) { + return index; + } + } + hint.min(chunks.len().saturating_sub(1)) +} + +fn extract_people_from_header(value: &str, accumulator: &mut ExtractionAccumulator) -> Vec { + let mut people = Vec::new(); + for captures in named_email_regex().captures_iter(value) { + let name = sanitize_fact_text( + captures + .name("name") + .map(|value| value.as_str()) + .unwrap_or(""), + ); + if name.is_empty() { + continue; + } + let canonical = sanitize_entity_name(&name); + let _ = accumulator.add_entity(&canonical, "PERSON", 0.95); + accumulator.remember_person_aliases(&canonical); + people.push(canonical); + } + people +} + +fn detect_primary_subject(text: &str) -> Option { + if text.contains("OpenHuman") { + return Some("OPENHUMAN".to_string()); + } + None +} + +fn apply_model_extraction( + runtime: &relex::RelexRuntime, + unit: &ExtractionUnit, + accumulator: &mut ExtractionAccumulator, + config: &MemoryIngestionConfig, +) { + let Ok(extraction) = runtime.extract( + &unit.text, + config.entity_threshold, + config.relation_threshold, + ) else { + return; + }; + + for entity in extraction.entities { + let _ = accumulator.add_entity(&entity.name, &entity.entity_type, entity.confidence); + } + + for relation in extraction.relations { + let mut metadata = Map::new(); + metadata.insert("extractor".to_string(), json!("gliner_relex_onnx")); + metadata.insert("source_text".to_string(), json!(unit.text)); + accumulator.add_relation( + &relation.subject, + &relation.subject_type, + &relation.predicate, + &relation.object, + &relation.object_type, + relation.confidence, + unit.chunk_index, + unit.order_index, + metadata, + ); + match UnifiedMemory::normalize_graph_predicate(&relation.predicate).as_str() { + "PREFERS" => { + accumulator.preferences.insert(format!( + "{} prefers {}", + sanitize_entity_name(&relation.subject), + sanitize_fact_text(&relation.object) + )); + accumulator.tags.insert("preference".to_string()); + accumulator.doc_kind = Some("profile".to_string()); + } + "HAS_DEADLINE" => { + accumulator.tags.insert("deadline".to_string()); + } + "OWNS" | "REVIEWS" => { + accumulator.tags.insert("owner".to_string()); + } + _ => {} + } + } +} + +fn enrich_document_metadata( + input: &NamespaceDocumentInput, + parsed: &ParsedIngestion, + config: &MemoryIngestionConfig, +) -> (NamespaceDocumentInput, Vec) { + let mut metadata = match input.metadata.clone() { + Value::Object(map) => map, + _ => Map::new(), + }; + for (key, value) in parsed.metadata.as_object().cloned().unwrap_or_default() { + metadata.insert(key, value); + } + metadata.insert( + "ingestion".to_string(), + json!({ + "backend": "openhuman_rust_relex", + "model_name": config.model_name, + "extraction_mode": config.extraction_mode.as_str(), + "entity_count": parsed.entities.len(), + "relation_count": parsed.relations.len(), + "preference_count": parsed.preference_count, + "decision_count": parsed.decision_count, + "chunk_count": parsed.chunk_count, + }), + ); + if parsed.preference_count > 0 || parsed.decision_count > 0 { + metadata.insert("kind".to_string(), json!("profile")); + } + + let mut tags = input.tags.iter().cloned().collect::>(); + tags.extend(parsed.tags.iter().cloned()); + let tags = tags.into_iter().collect::>(); + + ( + NamespaceDocumentInput { + namespace: input.namespace.clone(), + key: input.key.clone(), + title: input.title.clone(), + content: input.content.clone(), + source_type: input.source_type.clone(), + priority: input.priority.clone(), + tags: tags.clone(), + metadata: Value::Object(metadata), + category: input.category.clone(), + session_id: input.session_id.clone(), + document_id: input.document_id.clone(), + }, + tags, + ) +} + +fn reverse_aliases(aliases: &HashMap) -> BTreeMap> { + let mut reverse = BTreeMap::new(); + for (alias, canonical) in aliases { + if alias == canonical { + continue; + } + reverse + .entry(canonical.clone()) + .or_insert_with(Vec::new) + .push(alias.clone()); + } + for values in reverse.values_mut() { + values.sort(); + values.dedup(); + } + reverse +} + +fn build_alias_map(entities: &HashMap) -> HashMap { + let mut by_type = HashMap::>::new(); + for entity in entities.values() { + by_type + .entry(entity.entity_type.clone()) + .or_default() + .push(entity.name.clone()); + } + + let mut aliases = HashMap::new(); + for names in by_type.values_mut() { + names.sort_by_key(|name| std::cmp::Reverse(name.len())); + for short in names.iter() { + for long in names.iter() { + if short == long || long.len() <= short.len() { + continue; + } + if long.starts_with(&format!("{short} ")) || long.ends_with(&format!(" {short}")) { + aliases.entry(short.clone()).or_insert_with(|| long.clone()); + break; + } + } + } + } + aliases +} + +fn resolve_alias(name: &str, aliases: &HashMap) -> String { + let mut current = name.to_string(); + let mut seen = BTreeSet::new(); + while let Some(next) = aliases.get(¤t) { + if !seen.insert(current.clone()) { + break; + } + current = next.clone(); + } + current +} + +async fn parse_document( + content: &str, + title: &str, + config: &MemoryIngestionConfig, +) -> ParsedIngestion { + let chunks = UnifiedMemory::chunk_document_content(content, DEFAULT_CHUNK_TOKENS); + let relex_runtime = relex::runtime(&config.model_name).await; + let model_enabled = relex_runtime.is_some(); + let mut accumulator = ExtractionAccumulator { + document_title: Some(sanitize_entity_name(title)), + primary_subject: detect_primary_subject(title), + ..ExtractionAccumulator::default() + }; + + let mut chunk_hint = 0_usize; + for raw_line in content.lines() { + let line = sanitize_fact_text(raw_line); + if line.is_empty() { + continue; + } + + let chunk_index = find_chunk_index(&chunks, &line, chunk_hint); + chunk_hint = chunk_index; + let order_index = i64::try_from(chunk_index).unwrap_or(i64::MAX); + + if raw_line.trim_start().starts_with('#') { + let heading = sanitize_entity_name(raw_line.trim_start_matches('#')); + if !heading.is_empty() { + if accumulator.document_title.is_none() { + accumulator.document_title = Some(heading.clone()); + } + accumulator.current_subject = Some(heading); + } + continue; + } + + if let Some(captures) = email_header_regex().captures(&line) { + let header_name = captures + .get(1) + .map(|value| value.as_str()) + .unwrap_or_default() + .to_ascii_uppercase(); + let value = captures + .name("value") + .map(|value| value.as_str()) + .unwrap_or(""); + let people = extract_people_from_header(value, &mut accumulator); + if header_name == "FROM" { + accumulator.current_sender = people.first().cloned(); + } else if header_name == "TO" || header_name == "CC" { + if let Some(sender) = accumulator.current_sender.clone() { + for recipient in &people { + accumulator.add_relation( + &sender, + "PERSON", + "communicates_with", + recipient, + "PERSON", + 0.82, + chunk_index, + order_index, + Map::new(), + ); + } + } + } + continue; + } + + if let Some(subject) = line.strip_prefix("Subject:") { + let subject_text = sanitize_fact_text(subject); + if let Some(primary_subject) = detect_primary_subject(&subject_text) { + accumulator.primary_subject = Some(primary_subject); + } + continue; + } + + if let Some(date_text) = line.strip_prefix("Date:") { + let date_text = sanitize_fact_text(date_text); + if let Some(sender) = accumulator.current_sender.clone() { + accumulator.add_relation( + &sender, + "PERSON", + "has_deadline", + &date_text, + "DATE", + 0.75, + chunk_index, + order_index, + Map::new(), + ); + } + continue; + } + + if let Some(value) = line.strip_prefix("Project name:") { + let project = sanitize_entity_name(value); + if !project.is_empty() { + accumulator.primary_subject = Some(project.clone()); + let _ = accumulator.add_entity(&project, "PROJECT", 0.96); + } + continue; + } + + if let Some(value) = line.strip_prefix("Subproject:") { + let subproject = sanitize_entity_name(value); + if !subproject.is_empty() { + let _ = accumulator.add_entity(&subproject, "PROJECT", 0.92); + } + continue; + } + + if let Some(value) = line.strip_prefix("Owner:") { + let owner = sanitize_entity_name(value); + let owned = accumulator + .current_subject + .clone() + .or_else(|| accumulator.primary_subject.clone()) + .or_else(|| accumulator.document_title.clone()) + .unwrap_or_else(|| "DOCUMENT".to_string()); + accumulator.add_relation( + &owner, + "PERSON", + "owns", + &owned, + "WORK_ITEM", + 0.94, + chunk_index, + order_index, + Map::new(), + ); + continue; + } + + if let Some(value) = line.strip_prefix("Name:") { + let name = sanitize_entity_name(value); + if !name.is_empty() { + accumulator.current_subject = Some(name.clone()); + let _ = accumulator.add_entity(&name, "WORK_ITEM", 0.93); + } + continue; + } + + if let Some(value) = line.strip_prefix("Due date:") { + let due_date = sanitize_fact_text(value); + let subject = accumulator + .current_subject + .clone() + .or_else(|| accumulator.primary_subject.clone()) + .unwrap_or_else(|| "DOCUMENT".to_string()); + accumulator.add_relation( + &subject, + "WORK_ITEM", + "has_deadline", + &due_date, + "DATE", + 0.92, + chunk_index, + order_index, + Map::new(), + ); + accumulator.tags.insert("deadline".to_string()); + continue; + } + + if let Some(value) = line.strip_prefix("Target milestone:") { + let due_date = sanitize_fact_text(value); + let subject = accumulator + .primary_subject + .clone() + .or_else(|| accumulator.document_title.clone()) + .unwrap_or_else(|| "DOCUMENT".to_string()); + accumulator.add_relation( + &subject, + "PROJECT", + "has_deadline", + &due_date, + "DATE", + 0.92, + chunk_index, + order_index, + Map::new(), + ); + accumulator.tags.insert("deadline".to_string()); + continue; + } + + if let Some(value) = line.strip_prefix("Preferred embedding model for local experiments:") { + let model = sanitize_fact_text(value); + let subject = accumulator + .primary_subject + .clone() + .or_else(|| accumulator.document_title.clone()) + .unwrap_or_else(|| "DOCUMENT".to_string()); + accumulator.add_relation( + &subject, + "PROJECT", + "uses", + &model, + "TOOL", + 0.88, + chunk_index, + order_index, + Map::new(), + ); + accumulator + .decisions + .insert(format!("{subject} uses {model}")); + accumulator.tags.insert("decision".to_string()); + continue; + } + + if let Some(value) = line.strip_prefix("Preferred extraction mode to try first:") { + let mode = sanitize_fact_text(value); + let subject = accumulator + .primary_subject + .clone() + .or_else(|| accumulator.document_title.clone()) + .unwrap_or_else(|| "DOCUMENT".to_string()); + accumulator.add_relation( + &subject, + "PROJECT", + "uses", + &mode, + "MODE", + 0.88, + chunk_index, + order_index, + Map::new(), + ); + accumulator + .decisions + .insert(format!("{subject} uses {mode}")); + accumulator.tags.insert("decision".to_string()); + continue; + } + + if !model_enabled { + if let Some(captures) = graph_fact_regex().captures(&line) { + let subject = captures + .name("subject") + .map(|value| value.as_str()) + .unwrap_or(""); + let predicate = captures + .name("predicate") + .map(|value| value.as_str()) + .unwrap_or(""); + let object = captures + .name("object") + .map(|value| value.as_str()) + .unwrap_or(""); + let subject_type = classify_entity(subject, &accumulator.known_people); + let object_type = classify_entity(object, &accumulator.known_people); + accumulator.add_relation( + subject, + subject_type, + predicate, + object, + object_type, + 0.87, + chunk_index, + order_index, + Map::new(), + ); + if UnifiedMemory::normalize_graph_predicate(predicate) == "PREFERS" { + accumulator.preferences.insert(format!( + "{} prefers {}", + sanitize_entity_name(subject), + sanitize_fact_text(object) + )); + accumulator.tags.insert("preference".to_string()); + accumulator.doc_kind = Some("profile".to_string()); + } + continue; + } + + } + + if let Some(captures) = explicit_owner_regex().captures(&line) { + let subject = captures + .name("subject") + .map(|value| value.as_str()) + .unwrap_or(""); + let object = captures + .name("object") + .map(|value| value.as_str()) + .unwrap_or(""); + accumulator.add_relation( + subject, + "PERSON", + "owns", + object, + classify_entity(object, &accumulator.known_people), + 0.94, + chunk_index, + order_index, + Map::new(), + ); + accumulator.tags.insert("owner".to_string()); + continue; + } + + if let Some(captures) = will_review_regex().captures(&line) { + let subject = captures + .name("subject") + .map(|value| value.as_str()) + .unwrap_or(""); + let object = captures + .name("object") + .map(|value| value.as_str()) + .unwrap_or(""); + accumulator.add_relation( + subject, + "PERSON", + "reviews", + object, + classify_entity(object, &accumulator.known_people), + 0.80, + chunk_index, + order_index, + Map::new(), + ); + accumulator.tags.insert("owner".to_string()); + continue; + } + + if let Some(captures) = explicit_preference_regex().captures(&line) { + let subject = captures + .name("subject") + .map(|value| value.as_str()) + .unwrap_or(""); + let object = captures + .name("object") + .map(|value| value.as_str()) + .unwrap_or(""); + accumulator.add_relation( + subject, + "PERSON", + "prefers", + object, + classify_entity(object, &accumulator.known_people), + 0.90, + chunk_index, + order_index, + Map::new(), + ); + accumulator.preferences.insert(format!( + "{} prefers {}", + sanitize_entity_name(subject), + sanitize_fact_text(object) + )); + accumulator.tags.insert("preference".to_string()); + accumulator.doc_kind = Some("profile".to_string()); + continue; + } + + if let Some(value) = line.strip_prefix("I prefer ") { + if let Some(subject) = accumulator.current_sender.clone() { + let preference = sanitize_fact_text(value); + accumulator.add_relation( + &subject, + "PERSON", + "prefers", + &preference, + classify_entity(&preference, &accumulator.known_people), + 0.92, + chunk_index, + order_index, + Map::new(), + ); + accumulator + .preferences + .insert(format!("{subject} prefers {preference}")); + accumulator.tags.insert("preference".to_string()); + accumulator.doc_kind = Some("profile".to_string()); + continue; + } + } + + if let Some(captures) = action_item_regex().captures(&line) { + let subject = captures + .name("subject") + .map(|value| value.as_str()) + .unwrap_or(""); + let object = captures + .name("object") + .map(|value| value.as_str()) + .unwrap_or(""); + if accumulator + .known_people + .contains_key(&sanitize_entity_name(subject)) + || classify_entity(subject, &accumulator.known_people) == "PERSON" + { + accumulator.add_relation( + subject, + "PERSON", + "owns", + object, + classify_entity(object, &accumulator.known_people), + 0.83, + chunk_index, + order_index, + Map::new(), + ); + accumulator.tags.insert("owner".to_string()); + continue; + } + } + + let upper = sanitize_entity_name(&line); + let decision_subject = accumulator + .primary_subject + .clone() + .or_else(|| accumulator.document_title.clone()) + .unwrap_or_else(|| "DOCUMENT".to_string()); + if upper.contains("JSON-RPC") { + accumulator.add_relation( + &decision_subject, + "PROJECT", + "uses", + "JSON-RPC", + "PRODUCT", + 0.86, + chunk_index, + order_index, + Map::new(), + ); + accumulator + .decisions + .insert(format!("{decision_subject} uses JSON-RPC")); + accumulator.tags.insert("decision".to_string()); + continue; + } + if upper.contains("SHOULD USE NAMESPACE") + || upper.contains("USE NAMESPACE AS THE STORAGE") + || upper.contains("NAMESPACE AS THE MAIN SCOPE KEY") + { + accumulator.add_relation( + &decision_subject, + "PROJECT", + "uses", + "namespace", + "TOPIC", + 0.84, + chunk_index, + order_index, + Map::new(), + ); + accumulator + .decisions + .insert(format!("{decision_subject} uses namespace")); + accumulator.tags.insert("decision".to_string()); + continue; + } + if upper.contains("USER_ID") && (upper.contains("DO NOT NEED") || upper.contains("AVOID")) { + accumulator.add_relation( + &decision_subject, + "PROJECT", + "avoids", + "user_id", + "TOPIC", + 0.82, + chunk_index, + order_index, + Map::new(), + ); + accumulator + .decisions + .insert(format!("{decision_subject} avoids user_id")); + accumulator.tags.insert("decision".to_string()); + } + } + + for unit in build_units(&chunks, config.extraction_mode) { + if let Some(ref runtime) = relex_runtime { + apply_model_extraction(&runtime, &unit, &mut accumulator, config); + } + + if let Some(captures) = recipient_regex().captures(&unit.text) { + let giver = captures + .name("giver") + .map(|value| value.as_str()) + .unwrap_or(""); + let object = captures + .name("object") + .map(|value| value.as_str()) + .unwrap_or(""); + let recipient = captures + .name("recipient") + .map(|value| value.as_str()) + .unwrap_or(""); + accumulator.add_relation( + giver, + "PERSON", + "uses", + object, + classify_entity(object, &accumulator.known_people), + config.adjacency_threshold.max(0.62), + unit.chunk_index, + unit.order_index, + Map::new(), + ); + accumulator.add_relation( + recipient, + "PERSON", + "uses", + object, + classify_entity(object, &accumulator.known_people), + (config.adjacency_threshold * 0.9).max(0.55), + unit.chunk_index, + unit.order_index, + Map::new(), + ); + } + + if let Some(captures) = spatial_regex().captures(&unit.text) { + let head = captures + .name("head") + .map(|value| value.as_str()) + .unwrap_or(""); + let direction = captures + .name("direction") + .map(|value| value.as_str()) + .unwrap_or(""); + let tail = captures + .name("tail") + .map(|value| value.as_str()) + .unwrap_or(""); + let inverse = match direction.to_ascii_lowercase().as_str() { + "north" => "south_of", + "south" => "north_of", + "east" => "west_of", + "west" => "east_of", + _ => "", + }; + let predicate = format!("{direction}_of"); + accumulator.add_relation( + head, + "ROOM", + &predicate, + tail, + "ROOM", + config.adjacency_threshold.max(0.70), + unit.chunk_index, + unit.order_index, + Map::new(), + ); + if !inverse.is_empty() { + accumulator.add_relation( + tail, + "ROOM", + inverse, + head, + "ROOM", + config.adjacency_threshold.max(0.70), + unit.chunk_index, + unit.order_index, + Map::new(), + ); + } + } + } + + let aliases = build_alias_map(&accumulator.entities); + let reverse_alias = reverse_aliases(&aliases); + let mut canonical_entities = BTreeMap::::new(); + for entity in accumulator.entities.values() { + let canonical = resolve_alias(&entity.name, &aliases); + let entry = canonical_entities + .entry(canonical.clone()) + .or_insert_with(|| RawEntity { + name: canonical.clone(), + entity_type: entity.entity_type.clone(), + confidence: entity.confidence, + }); + if entity.confidence > entry.confidence { + entry.confidence = entity.confidence; + entry.entity_type = entity.entity_type.clone(); + } + } + + let mut aggregated_relations = BTreeMap::<(String, String, String), RawRelation>::new(); + for relation in accumulator.relations { + let subject = resolve_alias(&relation.subject, &aliases); + let object = resolve_alias(&relation.object, &aliases); + if subject == object { + continue; + } + let key = (subject.clone(), relation.predicate.clone(), object.clone()); + let entry = aggregated_relations + .entry(key) + .or_insert_with(|| RawRelation { + subject, + subject_type: relation.subject_type.clone(), + predicate: relation.predicate.clone(), + object, + object_type: relation.object_type.clone(), + confidence: relation.confidence, + chunk_indexes: relation.chunk_indexes.clone(), + order_index: relation.order_index, + metadata: relation.metadata.clone(), + }); + entry.confidence = entry.confidence.max(relation.confidence); + entry.order_index = entry.order_index.min(relation.order_index); + entry.chunk_indexes.extend(relation.chunk_indexes); + } + + let entities = canonical_entities + .into_values() + .filter(|entity| entity.confidence >= config.entity_threshold) + .map(|entity| ExtractedEntity { + name: entity.name.clone(), + entity_type: entity.entity_type, + aliases: reverse_alias.get(&entity.name).cloned().unwrap_or_default(), + }) + .collect::>(); + + let relations = aggregated_relations + .into_values() + .filter(|relation| relation.confidence >= config.relation_threshold) + .map(|relation| ExtractedRelation { + subject: relation.subject, + subject_type: relation.subject_type, + predicate: relation.predicate, + object: relation.object, + object_type: relation.object_type, + confidence: relation.confidence, + evidence_count: u32::try_from(relation.chunk_indexes.len()).unwrap_or(u32::MAX), + chunk_ids: relation + .chunk_indexes + .iter() + .map(|index| format!("chunk:{index}")) + .collect::>(), + order_index: Some(relation.order_index), + metadata: Value::Object(relation.metadata), + }) + .collect::>(); + + let mut tags = accumulator.tags.into_iter().collect::>(); + tags.sort(); + let metadata = json!({ + "kind": accumulator.doc_kind.or_else(|| { + if !accumulator.preferences.is_empty() || !accumulator.decisions.is_empty() { + Some("profile".to_string()) + } else { + None + } + }), + "primary_subject": accumulator.primary_subject, + "decisions": accumulator.decisions.iter().cloned().collect::>(), + "preferences": accumulator.preferences.iter().cloned().collect::>(), + "extracted_entities": entities.iter().map(|entity| { + json!({ + "name": entity.name, + "entity_type": entity.entity_type, + "aliases": entity.aliases, + }) + }).collect::>(), + }); + + ParsedIngestion { + tags, + metadata, + entities, + relations, + chunk_count: chunks.len(), + preference_count: accumulator.preferences.len(), + decision_count: accumulator.decisions.len(), + } +} + +impl UnifiedMemory { + pub async fn ingest_document( + &self, + request: MemoryIngestionRequest, + ) -> Result { + let parsed = parse_document( + &request.document.content, + &request.document.title, + &request.config, + ) + .await; + let (enriched_input, tags) = + enrich_document_metadata(&request.document, &parsed, &request.config); + let namespace = Self::sanitize_namespace(&enriched_input.namespace); + let document_id = self.upsert_document(enriched_input).await?; + + self.graph_remove_document_namespace(&namespace, &document_id) + .await?; + + for relation in &parsed.relations { + let chunk_ids = relation + .chunk_ids + .iter() + .filter_map(|chunk_id| chunk_id.strip_prefix("chunk:")) + .map(|chunk_index| format!("{document_id}:{chunk_index}")) + .collect::>(); + + let attrs = json!({ + "source": "ingestion", + "model_name": request.config.model_name, + "extraction_mode": request.config.extraction_mode.as_str(), + "confidence": relation.confidence, + "evidence_count": relation.evidence_count, + "order_index": relation.order_index, + "document_id": document_id, + "document_ids": [document_id.clone()], + "chunk_ids": chunk_ids, + "entity_types": { + "subject": relation.subject_type, + "object": relation.object_type, + }, + "metadata": relation.metadata, + }); + + self.graph_upsert_namespace( + &namespace, + &relation.subject, + &relation.predicate, + &relation.object, + &attrs, + ) + .await?; + } + + Ok(MemoryIngestionResult { + document_id, + namespace, + model_name: request.config.model_name, + extraction_mode: request.config.extraction_mode.as_str().to_string(), + chunk_count: parsed.chunk_count, + entity_count: parsed.entities.len(), + relation_count: parsed.relations.len(), + preference_count: parsed.preference_count, + decision_count: parsed.decision_count, + tags, + entities: parsed.entities, + relations: parsed.relations, + }) + } +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use serde_json::json; + use tempfile::TempDir; + + use crate::openhuman::memory::{ + embeddings::NoopEmbedding, MemoryIngestionConfig, MemoryIngestionRequest, + NamespaceDocumentInput, UnifiedMemory, + }; + + fn fixture(path: &str) -> String { + std::fs::read_to_string(format!( + "{}\\tests\\fixtures\\ingestion\\{}", + env!("CARGO_MANIFEST_DIR"), + path + )) + .expect("fixture should load") + } + + #[tokio::test] + async fn gmail_fixture_ingestion_recovers_required_signals() { + let tmp = TempDir::new().unwrap(); + let memory = UnifiedMemory::new(tmp.path(), Arc::new(NoopEmbedding), None).unwrap(); + let result = memory + .ingest_document(MemoryIngestionRequest { + document: NamespaceDocumentInput { + namespace: "skill-gmail".to_string(), + key: "gmail-thread-memory-integration".to_string(), + title: "Memory integration plan for OpenHuman desktop".to_string(), + content: fixture("gmail_thread_example.txt"), + source_type: "gmail".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(); + + assert!(result + .entities + .iter() + .any(|entity| entity.name == "SANIL JAIN")); + assert!(result + .entities + .iter() + .any(|entity| entity.name == "RAVI KULKARNI")); + assert!(result + .entities + .iter() + .any(|entity| entity.name == "ASHA MEHTA")); + assert!(result + .entities + .iter() + .any(|entity| entity.name == "OPENHUMAN")); + assert!(result + .relations + .iter() + .any(|relation| relation.subject == "OPENHUMAN" + && relation.predicate == "USES" + && relation.object.contains("JSON-RPC"))); + assert!(result + .relations + .iter() + .any(|relation| relation.subject == "RAVI KULKARNI" && relation.predicate == "OWNS")); + assert!(result.preference_count >= 1); + assert!(result.decision_count >= 1); + + let context = memory + .query_namespace_context_data( + "skill-gmail", + "who owns the rust memory api alignment", + 5, + ) + .await + .unwrap(); + assert!(context + .hits + .iter() + .flat_map(|hit| hit.supporting_relations.iter()) + .any(|relation| relation.subject == "RAVI KULKARNI" && relation.predicate == "OWNS")); + + let recall = memory + .recall_namespace_context_data("skill-gmail", 5) + .await + .unwrap(); + assert!(!recall.context_text.is_empty()); + assert!(recall + .hits + .iter() + .any(|hit| hit.content.contains("OpenHuman") || hit.content.contains("JSON-RPC"))); + assert!(recall + .hits + .iter() + .any(|hit| !hit.supporting_relations.is_empty())); + + let memories = memory + .recall_namespace_memories("skill-gmail", 5) + .await + .unwrap(); + assert!(memories.iter().any(|hit| hit.content.contains("JSON-RPC"))); + assert!(memories + .iter() + .any(|hit| matches!(hit.kind, crate::openhuman::memory::MemoryItemKind::Document))); + assert!(memories + .iter() + .any(|hit| !hit.supporting_relations.is_empty())); + } + + #[tokio::test] + async fn notion_fixture_ingestion_recovers_required_signals() { + let tmp = TempDir::new().unwrap(); + let memory = UnifiedMemory::new(tmp.path(), Arc::new(NoopEmbedding), None).unwrap(); + let result = memory + .ingest_document(MemoryIngestionRequest { + document: NamespaceDocumentInput { + namespace: "skill-notion".to_string(), + key: "notion-roadmap-memory-layer".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(); + + assert!(result + .entities + .iter() + .any(|entity| entity.name == "OPENHUMAN")); + assert!(result + .entities + .iter() + .any(|entity| entity.name == "SANIL JAIN")); + assert!(result + .relations + .iter() + .any(|relation| relation.subject == "OPENHUMAN" + && relation.predicate == "USES" + && relation.object.contains("JSON-RPC"))); + assert!(result + .relations + .iter() + .any(|relation| relation.subject == "CORE CONTRACT LOCKED" + && relation.predicate == "HAS_DEADLINE")); + assert!(result + .relations + .iter() + .any(|relation| relation.subject == "SANIL JAIN" && relation.predicate == "PREFERS")); + assert!(result.preference_count >= 1); + assert!(result.decision_count >= 1); + + let graph_rows = memory + .graph_query_namespace("skill-notion", Some("OPENHUMAN"), Some("USES")) + .await + .unwrap(); + assert!(!graph_rows.is_empty()); + + let context = memory + .query_namespace_context_data( + "skill-notion", + "who prefers core-first delivery over ui-first delivery", + 5, + ) + .await + .unwrap(); + assert!(context + .hits + .iter() + .flat_map(|hit| hit.supporting_relations.iter()) + .any(|relation| relation.subject == "SANIL JAIN" && relation.predicate == "PREFERS")); + + let recall = memory + .recall_namespace_context_data("skill-notion", 5) + .await + .unwrap(); + assert!(!recall.context_text.is_empty()); + assert!(recall + .hits + .iter() + .any(|hit| hit.content.contains("OpenHuman"))); + + let memories = memory + .recall_namespace_memories("skill-notion", 5) + .await + .unwrap(); + assert!(memories + .iter() + .any(|hit| hit.content.contains("OpenHuman") || hit.content.contains("core-first"))); + assert!(memories + .iter() + .any(|hit| matches!(hit.kind, crate::openhuman::memory::MemoryItemKind::Document))); + assert!(memories + .iter() + .any(|hit| !hit.supporting_relations.is_empty())); + } +} diff --git a/src/openhuman/memory/mod.rs b/src/openhuman/memory/mod.rs index 6e00af989..9083420a7 100644 --- a/src/openhuman/memory/mod.rs +++ b/src/openhuman/memory/mod.rs @@ -1,10 +1,16 @@ pub mod chunker; pub mod embeddings; +pub mod ingestion; pub mod ops; +pub(crate) mod relex; pub mod rpc_models; pub mod store; pub mod traits; +pub use ingestion::{ + ExtractedEntity, ExtractedRelation, ExtractionMode, MemoryIngestionConfig, + MemoryIngestionRequest, MemoryIngestionResult, DEFAULT_GLINER_RELEX_MODEL, +}; pub use ops as rpc; pub use ops::*; pub use rpc_models::*; diff --git a/src/openhuman/memory/ops.rs b/src/openhuman/memory/ops.rs index 379795531..68b4c15d5 100644 --- a/src/openhuman/memory/ops.rs +++ b/src/openhuman/memory/ops.rs @@ -668,6 +668,9 @@ pub async fn memory_init( let client = Arc::new(MemoryClient::from_workspace_dir(workspace_dir.clone())?); *lock_memory_client_state()? = Some(client); let memory_dir = workspace_dir.join("memory"); + tokio::spawn(async { + let _ = super::relex::warm_default_bundle().await; + }); Ok(envelope( MemoryInitResponse { initialized: true, diff --git a/src/openhuman/memory/relex.rs b/src/openhuman/memory/relex.rs new file mode 100644 index 000000000..dcf392b31 --- /dev/null +++ b/src/openhuman/memory/relex.rs @@ -0,0 +1,874 @@ +use std::env; +use std::sync::Arc; +use std::path::{Path, PathBuf}; +use std::sync::OnceLock; + +use anyhow::{anyhow, Context, Result}; +use futures_util::TryStreamExt; +use glob::glob; +use ndarray::{Array, Array2, Array3, Array4, Ix2, Ix3, Ix4}; +use ort::{ + session::{builder::GraphOptimizationLevel, Session}, + value::Tensor, +}; +use parking_lot::Mutex; +use serde::Deserialize; +use sha2::{Digest, Sha256}; +use tokenizers::Tokenizer; +use tokio::io::AsyncWriteExt; +use tokio::sync::Mutex as AsyncMutex; + +use crate::openhuman::memory::DEFAULT_GLINER_RELEX_MODEL; + +const DEFAULT_EXPORTED_RELEX_DIR: &str = + "_tmp/gliner-export/artifacts/gliner-relex-large-v0.5-onnx"; +const DEFAULT_MANAGED_RELEX_DIR: &str = ".openhuman/models/gliner-relex-large-v0.5-onnx"; +const DEFAULT_RELEX_RELEASE_BASE_URL: &str = + "https://github.com/sanil-23/GLiNER/releases/download/tinyhumans-gliner-relex-v0.5-onnx.1"; +const MODEL_FILE_NAME: &str = "model_quantized.onnx"; +const FALLBACK_MODEL_FILE_NAME: &str = "model.onnx"; +const TOKENIZER_FILE_NAME: &str = "tokenizer.json"; +const TOKENIZER_CONFIG_FILE_NAME: &str = "tokenizer_config.json"; +const GLINER_CONFIG_FILE_NAME: &str = "gliner_config.json"; +#[cfg(target_os = "windows")] +const ORT_DYLIB_FILE_NAME: &str = "onnxruntime.dll"; +#[cfg(target_os = "macos")] +const ORT_DYLIB_FILE_NAME: &str = "libonnxruntime.dylib"; + +struct BundleAsset { + remote_name: &'static str, + local_name: &'static str, + sha256: &'static str, +} + +const CORE_BUNDLE_ASSETS: &[BundleAsset] = &[ + BundleAsset { + remote_name: MODEL_FILE_NAME, + local_name: MODEL_FILE_NAME, + sha256: "7D4B8D35750D0AEC35DA0EB1EDFE33076C6958B8CD6EEC4560C59822536C9AEF", + }, + BundleAsset { + remote_name: TOKENIZER_FILE_NAME, + local_name: TOKENIZER_FILE_NAME, + sha256: "0FD23B86F1BACEE52F4485FCD4441B923132302BED55BC5E081172CA013E7654", + }, + BundleAsset { + remote_name: TOKENIZER_CONFIG_FILE_NAME, + local_name: TOKENIZER_CONFIG_FILE_NAME, + sha256: "3157274603C17459B0589DBB6818A47714D780718A6D0EB505C10347C466F2CD", + }, + BundleAsset { + remote_name: GLINER_CONFIG_FILE_NAME, + local_name: GLINER_CONFIG_FILE_NAME, + sha256: "FF6D7FEFD65F721515A3822BB074F2A36EC9B66AC75DAA400E2465FFE52F02BA", + }, +]; + +#[cfg(target_os = "windows")] +const PLATFORM_BUNDLE_ASSETS: &[BundleAsset] = &[BundleAsset { + remote_name: ORT_DYLIB_FILE_NAME, + local_name: ORT_DYLIB_FILE_NAME, + sha256: "EF720FC44A4EA48626BFE1EBD29642DE20222D7F104A509EA305D9F3CB3B7850", +}]; +#[cfg(target_os = "macos")] +const PLATFORM_BUNDLE_ASSETS: &[BundleAsset] = &[BundleAsset { + remote_name: ORT_DYLIB_FILE_NAME, + local_name: ORT_DYLIB_FILE_NAME, + sha256: "285C8CD1E53856507B9B2E38EE9AFFC69AA6E90AC30F8670DC8195710CA14B77", +}]; +#[cfg(not(any(target_os = "windows", target_os = "macos")))] +const PLATFORM_BUNDLE_ASSETS: &[BundleAsset] = &[]; + +const ENTITY_LABELS: &[&str] = &[ + "person", + "organization", + "project", + "product", + "tool", + "topic", + "work item", + "mode", + "place", + "room", + "date", +]; + +const RELATION_LABELS: &[&str] = &[ + "owns", + "uses", + "works on", + "responsible for", + "reviews", + "works for", + "depends on", + "prefers", + "has deadline", + "communicates with", + "investigates", + "evaluates", + "north of", + "south of", + "east of", + "west of", + "avoids", +]; + +#[derive(Debug, Clone)] +pub(crate) struct RelexEntity { + pub name: String, + pub entity_type: String, + pub confidence: f32, +} + +#[derive(Debug, Clone)] +pub(crate) struct RelexRelation { + pub subject: String, + pub subject_type: String, + pub predicate: String, + pub object: String, + pub object_type: String, + pub confidence: f32, +} + +#[derive(Debug, Clone, Default)] +pub(crate) struct RelexExtraction { + pub entities: Vec, + pub relations: Vec, +} + +#[derive(Debug)] +pub(crate) struct RelexRuntime { + tokenizer: Tokenizer, + session: Mutex, + config: RelexBundleConfig, +} + +#[derive(Debug, Clone, Deserialize)] +struct RelexBundleConfig { + #[serde(default = "default_ent_token")] + ent_token: String, + #[serde(default = "default_rel_token")] + rel_token: String, + #[serde(default = "default_sep_token")] + sep_token: String, + #[serde(default = "default_max_width")] + max_width: usize, +} + +#[derive(Debug, Clone)] +struct PromptBatch { + input_ids: Array2, + attention_mask: Array2, + words_mask: Array2, + text_lengths: Array2, + num_words: usize, +} + +#[derive(Debug, Clone)] +struct TokenSlice { + start: usize, + end: usize, + text: String, +} + +#[derive(Debug, Clone)] +struct DecodedSpan { + start: usize, + end: usize, + text: String, + class_name: String, + probability: f32, +} + +fn default_ent_token() -> String { + "<>".to_string() +} + +fn default_rel_token() -> String { + "<>".to_string() +} + +fn default_sep_token() -> String { + "<>".to_string() +} + +fn default_max_width() -> usize { + 12 +} + +pub(crate) async fn runtime(model_name: &str) -> Option> { + if !uses_default_bundle(model_name) { + return load_runtime_for_model(model_name).await.ok(); + } + + static DEFAULT_RUNTIME: OnceLock>>> = OnceLock::new(); + static DEFAULT_RUNTIME_BOOTSTRAP: OnceLock> = OnceLock::new(); + + let runtime_cell = DEFAULT_RUNTIME.get_or_init(|| Mutex::new(None)); + if let Some(runtime) = runtime_cell.lock().clone() { + return Some(runtime); + } + + let _guard = DEFAULT_RUNTIME_BOOTSTRAP + .get_or_init(|| AsyncMutex::new(())) + .lock() + .await; + + if let Some(runtime) = runtime_cell.lock().clone() { + return Some(runtime); + } + + let runtime = load_default_runtime().await.ok().map(Arc::new)?; + *runtime_cell.lock() = Some(runtime.clone()); + Some(runtime) +} + +pub(crate) async fn warm_default_bundle() -> Option> { + runtime(DEFAULT_GLINER_RELEX_MODEL).await +} + +impl RelexRuntime { + pub(crate) fn extract( + &self, + text: &str, + entity_threshold: f32, + relation_threshold: f32, + ) -> Result { + let tokens = split_whitespace_tokens(text); + if tokens.is_empty() { + return Ok(RelexExtraction::default()); + } + + let prompt = encode_prompt( + &self.tokenizer, + &self.config, + &tokens, + ENTITY_LABELS, + RELATION_LABELS, + )?; + let (span_idx, span_mask) = make_spans_tensors(prompt.num_words, self.config.max_width); + + let inputs = ort::inputs! { + "input_ids" => Tensor::from_array(prompt.input_ids.clone())?, + "attention_mask" => Tensor::from_array(prompt.attention_mask.clone())?, + "words_mask" => Tensor::from_array(prompt.words_mask.clone())?, + "text_lengths" => Tensor::from_array(prompt.text_lengths.clone())?, + "span_idx" => Tensor::from_array(span_idx.clone())?, + "span_mask" => Tensor::from_array(span_mask.clone())?, + }; + + let mut session = self.session.lock(); + let outputs = session.run(inputs)?; + let logits = extract_f32_4d(outputs.get("logits").context("missing logits output")?)?; + + let spans = decode_entity_spans( + &logits, + text, + &tokens, + ENTITY_LABELS, + self.config.max_width, + entity_threshold, + ); + let entities = spans + .iter() + .map(|span| RelexEntity { + name: span.text.clone(), + entity_type: normalize_entity_label(&span.class_name).to_string(), + confidence: span.probability, + }) + .collect::>(); + + let mut relations = Vec::new(); + let rel_idx = outputs.get("rel_idx").map(extract_i64_3d).transpose()?; + let rel_logits = outputs.get("rel_logits").map(extract_f32_3d).transpose()?; + let rel_mask = outputs.get("rel_mask").map(extract_bool_2d).transpose()?; + + if let (Some(rel_idx), Some(rel_logits), Some(rel_mask)) = (rel_idx, rel_logits, rel_mask) { + let rel_pairs = rel_idx.index_axis(ndarray::Axis(0), 0); + let rel_scores = rel_logits.index_axis(ndarray::Axis(0), 0); + let rel_valid = rel_mask.index_axis(ndarray::Axis(0), 0); + + for pair_idx in 0..rel_valid.shape()[0] { + if !rel_valid[[pair_idx]] { + continue; + } + let head_idx = rel_pairs[[pair_idx, 0]]; + let tail_idx = rel_pairs[[pair_idx, 1]]; + if head_idx < 0 || tail_idx < 0 { + continue; + } + + let head_idx = head_idx as usize; + let tail_idx = tail_idx as usize; + if head_idx >= spans.len() || tail_idx >= spans.len() { + continue; + } + + let head = &spans[head_idx]; + let tail = &spans[tail_idx]; + let class_count = rel_scores.shape()[1].min(RELATION_LABELS.len()); + for class_idx in 0..class_count { + let probability = sigmoid(rel_scores[[pair_idx, class_idx]]); + if probability < relation_threshold { + continue; + } + relations.push(RelexRelation { + subject: head.text.clone(), + subject_type: normalize_entity_label(&head.class_name).to_string(), + predicate: normalize_relation_label(RELATION_LABELS[class_idx]), + object: tail.text.clone(), + object_type: normalize_entity_label(&tail.class_name).to_string(), + confidence: probability, + }); + } + } + } + + Ok(RelexExtraction { + entities, + relations, + }) + } +} + +fn uses_default_bundle(model_name: &str) -> bool { + model_name.trim().is_empty() + || model_name == DEFAULT_GLINER_RELEX_MODEL + || model_name == default_bundle_dir().to_string_lossy() + || model_name == default_managed_bundle_dir().to_string_lossy() +} + +async fn load_default_runtime() -> Result { + let bundle_dir = resolve_bundle_dir(DEFAULT_GLINER_RELEX_MODEL) + .await + .ok_or_else(|| anyhow!("relex bundle directory not found"))?; + load_runtime_from_bundle_dir(&bundle_dir) +} + +async fn load_runtime_for_model(model_name: &str) -> Result> { + let bundle_dir = resolve_bundle_dir(model_name) + .await + .ok_or_else(|| anyhow!("relex bundle directory not found"))?; + load_runtime_from_bundle_dir(&bundle_dir).map(Arc::new) +} + +fn load_runtime_from_bundle_dir(bundle_dir: &Path) -> Result { + ensure_ort_dylib_path(bundle_dir); + + let tokenizer_path = bundle_dir.join(TOKENIZER_FILE_NAME); + let model_path = model_file_path(bundle_dir) + .ok_or_else(|| anyhow!("model file not found in {}", bundle_dir.display()))?; + let config_path = bundle_dir.join(GLINER_CONFIG_FILE_NAME); + + let tokenizer = Tokenizer::from_file(&tokenizer_path).map_err(|err| { + anyhow!( + "failed to load tokenizer from {}: {err}", + tokenizer_path.display() + ) + })?; + let config = serde_json::from_slice::( + &std::fs::read(&config_path) + .with_context(|| format!("failed to read {}", config_path.display()))?, + ) + .with_context(|| format!("failed to parse {}", config_path.display()))?; + + let session = Session::builder()? + .with_optimization_level(GraphOptimizationLevel::Level3)? + .commit_from_file(&model_path) + .with_context(|| format!("failed to load model {}", model_path.display()))?; + + Ok(RelexRuntime { + tokenizer, + session: Mutex::new(session), + config, + }) +} + +async fn resolve_bundle_dir(model_name: &str) -> Option { + if let Ok(path) = env::var("OPENHUMAN_GLINER_RELEX_DIR") { + let bundle_dir = PathBuf::from(path); + if bundle_complete(&bundle_dir) { + return Some(bundle_dir); + } + } + + let requested = PathBuf::from(model_name); + if requested.is_absolute() || model_name.contains('/') || model_name.contains('\\') { + if requested.is_dir() && bundle_complete(&requested) { + return Some(requested); + } + if requested.is_file() + && requested + .file_name() + .is_some_and(|name| name == FALLBACK_MODEL_FILE_NAME || name == MODEL_FILE_NAME) + { + return requested.parent().map(Path::to_path_buf); + } + } + + let managed_dir = default_managed_bundle_dir(); + if bundle_complete(&managed_dir) { + return Some(managed_dir); + } + + let bundle_dir = default_bundle_dir(); + if bundle_complete(&bundle_dir) { + return Some(bundle_dir); + } + + if uses_default_bundle(model_name) { + if ensure_managed_bundle(&managed_dir).await.is_ok() && bundle_complete(&managed_dir) { + return Some(managed_dir); + } + } + + None +} + +fn default_bundle_dir() -> PathBuf { + Path::new(env!("CARGO_MANIFEST_DIR")) + .parent() + .unwrap_or_else(|| Path::new(env!("CARGO_MANIFEST_DIR"))) + .join(DEFAULT_EXPORTED_RELEX_DIR) +} + +fn default_managed_bundle_dir() -> PathBuf { + if let Ok(path) = env::var("OPENHUMAN_GLINER_RELEX_CACHE_DIR") { + return PathBuf::from(path); + } + + directories::UserDirs::new() + .map(|dirs| dirs.home_dir().join(DEFAULT_MANAGED_RELEX_DIR)) + .unwrap_or_else(|| PathBuf::from(DEFAULT_MANAGED_RELEX_DIR)) +} + +fn bundle_complete(bundle_dir: &Path) -> bool { + bundle_dir.join(TOKENIZER_FILE_NAME).exists() + && bundle_dir.join(GLINER_CONFIG_FILE_NAME).exists() + && model_file_path(bundle_dir).is_some() +} + +fn model_file_path(bundle_dir: &Path) -> Option { + for file_name in [MODEL_FILE_NAME, FALLBACK_MODEL_FILE_NAME] { + let candidate = bundle_dir.join(file_name); + if candidate.exists() { + return Some(candidate); + } + } + None +} + +fn ensure_ort_dylib_path(bundle_dir: &Path) { + if env::var_os("ORT_DYLIB_PATH").is_some() { + return; + } + + #[cfg(any(target_os = "windows", target_os = "macos"))] + { + let bundled = bundle_dir.join(ORT_DYLIB_FILE_NAME); + if bundled.exists() { + env::set_var("ORT_DYLIB_PATH", bundled); + return; + } + } + + if let Some(lib_path) = env::var_os("ORT_LIB_LOCATION") { + let candidate = PathBuf::from(lib_path); + if candidate.is_file() { + env::set_var("ORT_DYLIB_PATH", candidate); + return; + } + let runtime_lib = candidate.join(ORT_DYLIB_FILE_NAME); + if runtime_lib.exists() { + env::set_var("ORT_DYLIB_PATH", runtime_lib); + return; + } + } + + #[cfg(target_os = "windows")] + { + let Some(user_profile) = env::var_os("USERPROFILE") else { + return; + }; + let pattern = PathBuf::from(user_profile) + .join("AppData/Local/uv/cache/archive-v0/*/onnxruntime/capi/onnxruntime.dll") + .to_string_lossy() + .replace('\\', "/"); + if let Ok(paths) = glob(&pattern) { + for candidate in paths.flatten() { + if candidate.exists() { + env::set_var("ORT_DYLIB_PATH", &candidate); + break; + } + } + } + } +} + +async fn ensure_managed_bundle(bundle_dir: &Path) -> Result<()> { + static MANAGED_BUNDLE_BOOTSTRAP: OnceLock> = OnceLock::new(); + let _guard = MANAGED_BUNDLE_BOOTSTRAP + .get_or_init(|| AsyncMutex::new(())) + .lock() + .await; + + if bundle_complete(bundle_dir) { + return Ok(()); + } + + tokio::fs::create_dir_all(bundle_dir) + .await + .with_context(|| format!("failed to create {}", bundle_dir.display()))?; + + let client = crate::openhuman::config::build_runtime_proxy_client("memory.relex"); + let base_url = env::var("OPENHUMAN_GLINER_RELEX_BASE_URL") + .ok() + .filter(|value| !value.trim().is_empty()) + .unwrap_or_else(|| DEFAULT_RELEX_RELEASE_BASE_URL.to_string()); + + for asset in CORE_BUNDLE_ASSETS.iter().chain(PLATFORM_BUNDLE_ASSETS.iter()) { + let target = bundle_dir.join(asset.local_name); + download_asset_if_needed(&client, &base_url, asset, &target).await?; + } + + Ok(()) +} + +async fn download_asset_if_needed( + client: &reqwest::Client, + base_url: &str, + asset: &BundleAsset, + target: &Path, +) -> Result<()> { + if file_matches_sha256(target, asset.sha256).await? { + return Ok(()); + } + + let url = format!( + "{}/{}", + base_url.trim_end_matches('/'), + asset.remote_name.trim_start_matches('/') + ); + let tmp = target.with_extension("download"); + if let Some(parent) = target.parent() { + tokio::fs::create_dir_all(parent) + .await + .with_context(|| format!("failed to create {}", parent.display()))?; + } + + let response = client + .get(&url) + .send() + .await + .with_context(|| format!("failed to start relex asset download {url}"))?; + if !response.status().is_success() { + return Err(anyhow!( + "failed to download relex asset {}, status {}", + asset.remote_name, + response.status() + )); + } + + let mut file = tokio::fs::File::create(&tmp) + .await + .with_context(|| format!("failed to create {}", tmp.display()))?; + let mut stream = response.bytes_stream(); + while let Some(chunk) = stream + .try_next() + .await + .with_context(|| format!("download stream error for {}", asset.remote_name))? + { + file.write_all(&chunk) + .await + .with_context(|| format!("failed writing {}", tmp.display()))?; + } + file.flush() + .await + .with_context(|| format!("failed flushing {}", tmp.display()))?; + + if !asset.sha256.is_empty() && !file_matches_sha256(&tmp, asset.sha256).await? { + let _ = tokio::fs::remove_file(&tmp).await; + return Err(anyhow!( + "checksum mismatch for downloaded relex asset {}", + asset.remote_name + )); + } + + tokio::fs::rename(&tmp, target) + .await + .with_context(|| format!("failed to finalize {}", target.display()))?; + Ok(()) +} + +async fn file_matches_sha256(path: &Path, expected: &str) -> Result { + if expected.is_empty() { + return Ok(path.exists()); + } + if !path.exists() { + return Ok(false); + } + let bytes = tokio::fs::read(path) + .await + .with_context(|| format!("failed to read {}", path.display()))?; + let actual = hex::encode(Sha256::digest(bytes)); + Ok(actual.eq_ignore_ascii_case(expected)) +} + +fn split_whitespace_tokens(text: &str) -> Vec { + let mut tokens = Vec::new(); + let mut current_start: Option = None; + + for (idx, ch) in text.char_indices() { + if ch.is_whitespace() { + if let Some(start) = current_start.take() { + tokens.push(TokenSlice { + start, + end: idx, + text: text[start..idx].to_string(), + }); + } + } else if current_start.is_none() { + current_start = Some(idx); + } + } + + if let Some(start) = current_start { + tokens.push(TokenSlice { + start, + end: text.len(), + text: text[start..].to_string(), + }); + } + + tokens +} + +fn encode_prompt( + tokenizer: &Tokenizer, + config: &RelexBundleConfig, + tokens: &[TokenSlice], + entity_labels: &[&str], + relation_labels: &[&str], +) -> Result { + let mut prompt_words = Vec::new(); + for label in entity_labels { + prompt_words.push(config.ent_token.clone()); + prompt_words.push((*label).to_string()); + } + prompt_words.push(config.sep_token.clone()); + for label in relation_labels { + prompt_words.push(config.rel_token.clone()); + prompt_words.push((*label).to_string()); + } + prompt_words.push(config.sep_token.clone()); + + let mut words = prompt_words.clone(); + words.extend(tokens.iter().map(|token| token.text.clone())); + + let mut encoded_words = Vec::with_capacity(words.len()); + let mut total_tokens = 2usize; + let mut prompt_subtokens = 0usize; + + for (index, word) in words.iter().enumerate() { + let encoding = tokenizer + .encode(word.as_str(), false) + .map_err(|err| anyhow!("failed to tokenize prompt word `{word}`: {err}"))?; + let ids = encoding.get_ids().to_vec(); + if index < prompt_words.len() { + prompt_subtokens += ids.len(); + } + total_tokens += ids.len(); + encoded_words.push(ids); + } + + let text_offset = prompt_subtokens + 1; + let mut input_ids = vec![0_i64; total_tokens]; + let mut attention_mask = vec![0_i64; total_tokens]; + let mut words_mask = vec![0_i64; total_tokens]; + + let mut cursor = 0usize; + input_ids[cursor] = 1; + attention_mask[cursor] = 1; + cursor += 1; + + let mut word_id = 0_i64; + for ids in encoded_words { + for (token_index, token_id) in ids.iter().enumerate() { + input_ids[cursor] = i64::from(*token_id); + attention_mask[cursor] = 1; + if cursor >= text_offset && token_index == 0 { + words_mask[cursor] = word_id; + } + cursor += 1; + } + if cursor >= text_offset { + word_id += 1; + } + } + + input_ids[cursor] = 2; + attention_mask[cursor] = 1; + + Ok(PromptBatch { + input_ids: Array2::from_shape_vec((1, total_tokens), input_ids)?, + attention_mask: Array2::from_shape_vec((1, total_tokens), attention_mask)?, + words_mask: Array2::from_shape_vec((1, total_tokens), words_mask)?, + text_lengths: Array2::from_shape_vec((1, 1), vec![tokens.len() as i64])?, + num_words: tokens.len(), + }) +} + +fn make_spans_tensors(num_words: usize, max_width: usize) -> (Array3, Array2) { + let num_spans = num_words * max_width; + let mut span_idx = Array3::::zeros((1, num_spans, 2)); + let mut span_mask = Array2::::from_elem((1, num_spans), false); + + for start in 0..num_words { + let actual_max_width = max_width.min(num_words.saturating_sub(start)); + for width in 0..actual_max_width { + let dim = start * max_width + width; + span_idx[[0, dim, 0]] = start as i64; + span_idx[[0, dim, 1]] = (start + width) as i64; + span_mask[[0, dim]] = true; + } + } + + (span_idx, span_mask) +} + +fn decode_entity_spans( + logits: &Array4, + text: &str, + tokens: &[TokenSlice], + entity_labels: &[&str], + max_width: usize, + threshold: f32, +) -> Vec { + let mut spans = Vec::new(); + let num_words = tokens.len(); + let width_count = logits.shape().get(2).copied().unwrap_or_default() as usize; + let class_count = logits.shape().get(3).copied().unwrap_or_default() as usize; + + for start in 0..num_words { + let actual_max_width = max_width + .min(width_count) + .min(num_words.saturating_sub(start)); + for width in 0..actual_max_width { + let end_word = start + width; + if end_word >= num_words { + continue; + } + for class_idx in 0..class_count.min(entity_labels.len()) { + let probability = sigmoid(logits[[0, start, width, class_idx]]); + if probability < threshold { + continue; + } + let start_offset = tokens[start].start; + let end_offset = tokens[end_word].end; + spans.push(DecodedSpan { + start: start_offset, + end: end_offset, + text: text[start_offset..end_offset].to_string(), + class_name: entity_labels[class_idx].to_string(), + probability, + }); + } + } + } + + spans.sort_unstable_by_key(|span| (span.start, span.end)); + greedy_filter(spans) +} + +fn extract_f32_4d(value: &ort::value::DynValue) -> Result> { + let (shape, data) = value.try_extract_tensor::()?; + Ok(Array::from_shape_vec(shape.to_ixdyn(), data.to_vec())?.into_dimensionality::()?) +} + +fn extract_f32_3d(value: &ort::value::DynValue) -> Result> { + let (shape, data) = value.try_extract_tensor::()?; + Ok(Array::from_shape_vec(shape.to_ixdyn(), data.to_vec())?.into_dimensionality::()?) +} + +fn extract_i64_3d(value: &ort::value::DynValue) -> Result> { + let (shape, data) = value.try_extract_tensor::()?; + Ok(Array::from_shape_vec(shape.to_ixdyn(), data.to_vec())?.into_dimensionality::()?) +} + +fn extract_bool_2d(value: &ort::value::DynValue) -> Result> { + let (shape, data) = value.try_extract_tensor::()?; + Ok(Array::from_shape_vec(shape.to_ixdyn(), data.to_vec())?.into_dimensionality::()?) +} + +fn greedy_filter(spans: Vec) -> Vec { + if spans.is_empty() { + return spans; + } + + let mut selected = Vec::with_capacity(spans.len()); + let mut previous = 0usize; + let mut next = 1usize; + + while next < spans.len() { + let left = &spans[previous]; + let right = &spans[next]; + if disjoint(left, right) { + selected.push(left.clone()); + previous = next; + } else if left.probability < right.probability { + previous = next; + } + next += 1; + } + + selected.push(spans[previous].clone()); + selected +} + +fn disjoint(left: &DecodedSpan, right: &DecodedSpan) -> bool { + right.start >= left.end || right.end <= left.start +} + +fn normalize_entity_label(label: &str) -> &'static str { + match label { + "person" => "PERSON", + "organization" => "ORGANIZATION", + "project" => "PROJECT", + "product" => "PRODUCT", + "tool" => "TOOL", + "topic" => "TOPIC", + "work item" => "WORK_ITEM", + "mode" => "MODE", + "place" => "PLACE", + "room" => "ROOM", + "date" => "DATE", + _ => "TOPIC", + } +} + +fn normalize_relation_label(label: &str) -> String { + match label { + "owns" => "owns".to_string(), + "uses" => "uses".to_string(), + "works on" => "works_on".to_string(), + "responsible for" => "responsible_for".to_string(), + "reviews" => "reviews".to_string(), + "works for" => "works_for".to_string(), + "depends on" => "depends_on".to_string(), + "prefers" => "prefers".to_string(), + "has deadline" => "has_deadline".to_string(), + "communicates with" => "communicates_with".to_string(), + "investigates" => "investigates".to_string(), + "evaluates" => "evaluates".to_string(), + "north of" => "north_of".to_string(), + "south of" => "south_of".to_string(), + "east of" => "east_of".to_string(), + "west of" => "west_of".to_string(), + "avoids" => "avoids".to_string(), + _ => label.to_string(), + } +} + +fn sigmoid(value: f32) -> f32 { + 1.0 / (1.0 + (-value).exp()) +} diff --git a/src/openhuman/memory/store/client.rs b/src/openhuman/memory/store/client.rs index 04e897200..bb2b1e7a9 100644 --- a/src/openhuman/memory/store/client.rs +++ b/src/openhuman/memory/store/client.rs @@ -3,6 +3,7 @@ use std::path::PathBuf; use std::sync::Arc; use crate::openhuman::memory::embeddings::{self, EmbeddingProvider}; +use crate::openhuman::memory::ingestion::{MemoryIngestionRequest, MemoryIngestionResult}; use crate::openhuman::memory::store::types::{ NamespaceDocumentInput, NamespaceMemoryHit, NamespaceRetrievalContext, }; @@ -45,6 +46,13 @@ impl MemoryClient { self.inner.upsert_document(input).await } + pub async fn ingest_doc( + &self, + request: MemoryIngestionRequest, + ) -> Result { + self.inner.ingest_document(request).await + } + #[allow(clippy::too_many_arguments)] pub async fn store_skill_sync( &self, diff --git a/tests/fixtures/ingestion/README.md b/tests/fixtures/ingestion/README.md new file mode 100644 index 000000000..ec11013bf --- /dev/null +++ b/tests/fixtures/ingestion/README.md @@ -0,0 +1,25 @@ +# Ingestion Fixtures + +These fixtures are plain-text source samples for memory ingestion tests. + +They are intentionally written as raw strings rather than strongly typed JSON so +future ingestion tests can exercise the same path used for real imported text. + +Current fixtures: + +- `gmail_thread_example.txt` + Gmail-like thread with headers, quoted replies, task ownership, dates, and + durable user/project facts. + +- `notion_page_example.txt` + Notion-like project page with sections, bullet lists, decisions, owners, + milestones, and operating notes. + +Suggested test usage: + +- Load fixture text as a string. +- Pass it through chunking and extraction. +- Assert that ingestion can recover: + - entities such as people, tools, projects, and dates + - relations such as ownership, dependencies, and responsibilities + - durable memory facts such as preferences, deadlines, and decisions diff --git a/tests/fixtures/ingestion/gmail_thread_example.txt b/tests/fixtures/ingestion/gmail_thread_example.txt new file mode 100644 index 000000000..70c44a145 --- /dev/null +++ b/tests/fixtures/ingestion/gmail_thread_example.txt @@ -0,0 +1,85 @@ +From: Sanil Jain +To: Asha Mehta , Ravi Kulkarni +Cc: OpenHuman Core +Subject: Re: Memory integration plan for OpenHuman desktop +Date: Tue, 12 Mar 2026 09:14:00 +0530 +Thread-Id: memory-integration-2026-03 + +Hi Asha and Ravi, + +Quick summary after today's sync: + +1. We should keep JSON-RPC as the transport for the desktop core. +2. The memory layer in the Rust core should use namespace as the main scope key. +3. We do not need user_id in the local storage contract for the current desktop runtime. +4. The frontend can adapt to richer result payloads as long as they still arrive inside JSON-RPC result. + +Current work items: +- Ravi owns the Rust memory API alignment for list, delete, query, and recall. +- Asha owns the Neocortex v2 ingestion experiment using the GLiNER relex model. +- Sanil will review response models so they follow the Neocortex API style. + +Important project facts: +- Project name: OpenHuman +- Subproject: memory-layer-completion +- Target milestone: March 22, 2026 +- Preferred embedding model for local experiments: text-embedding-3-small +- Preferred extraction mode to try first: sentence + +Known constraints: +- The desktop app is local-first. +- Core RPC currently binds to localhost only. +- We should avoid introducing user_id into every memory request unless we later support multi-user or remote runtimes. + +Action items: +- Ravi: draft typed request/response structs for memory.query_namespace and memory.recall_namespace by Friday. +- Asha: prepare two ingestion fixtures, one Gmail-like and one Notion-like, with enough structure to test entity and relation extraction. +- Sanil: decide whether memory.init becomes a no-op compatibility method or is removed from the frontend wrappers. + +One durable preference to remember: +I prefer keeping the memory core simple first and delaying graph traversal until after ingestion and recall are stable. + +Thanks, +Sanil + +--- + +From: Asha Mehta +To: Sanil Jain , Ravi Kulkarni +Subject: Re: Memory integration plan for OpenHuman desktop +Date: Tue, 12 Mar 2026 08:41:00 +0530 + +Agreed. + +For the Neocortex donor path, I reviewed the neocortex_v2 extractor again: +- It uses a single GLiNER relex model. +- It supports sentence-level and chunk-level extraction. +- It adds recipient and spatial relation heuristics. + +I think we should preserve those heuristics when we port the ingestion flow into OpenHuman. + +Also, please record this: +- Ravi prefers narrower worker ownership to avoid merge conflicts. +- I prefer evaluation fixtures that include dates, owners, and product decisions. + +Regards, +Asha + +--- + +From: Ravi Kulkarni +To: Sanil Jain , Asha Mehta +Subject: Re: Memory integration plan for OpenHuman desktop +Date: Tue, 12 Mar 2026 08:09:00 +0530 + +One more note before I start: + +- I will treat namespace as mandatory for memory query and recall. +- I will treat memory file APIs as optional until the core contract settles. +- I want the Gmail importer to preserve subject, sender, recipients, and sent_at metadata. + +Dependency note: +- The frontend wrapper work depends on finalizing the result shape from the Rust core. +- The ingestion evaluation can run in parallel once the storage mapping is clear. + +Ravi diff --git a/tests/fixtures/ingestion/notion_page_example.txt b/tests/fixtures/ingestion/notion_page_example.txt new file mode 100644 index 000000000..2439210b1 --- /dev/null +++ b/tests/fixtures/ingestion/notion_page_example.txt @@ -0,0 +1,132 @@ +# OpenHuman Memory Layer Roadmap + +Workspace: tinyhumans / engineering +Owner: Sanil Jain +Last edited: 2026-03-14 +Status: In Progress +Tags: memory, rust-core, ingestion, neocortex + +## Overview + +This page tracks the work needed to complete the OpenHuman memory layer in the Rust core. + +The current direction is: +- keep JSON-RPC as the transport +- use namespace as the storage and retrieval scope key +- avoid requiring user_id in local memory APIs +- adopt Neocortex-style typed request and response models inside JSON-RPC result + +## Core Decisions + +### Decision 1: Transport +We will keep JSON-RPC 2.0 as the transport for the desktop core. + +### Decision 2: Scope +Namespace is the primary logical partition for local memory. +Examples: +- conversations +- conscious +- skill-gmail +- skill-notion + +### Decision 3: Ingestion donor +We will use neocortex_v2 as the donor path for better memory extraction. +Important features to preserve: +- joint entity and relation extraction +- sentence-level extraction option +- relation constraints +- recipient relation synthesis +- spatial relation synthesis + +## Deliverables + +### Thread 0: Contract +Owner: Sanil Jain +Deliverables: +- final memory RPC names +- request and response model table +- decision on memory.init +- decision on file APIs + +### Thread 1: Core Memory Domain +Owner: Ravi Kulkarni +Deliverables: +- stable document storage semantics +- stable namespace list and document list behavior +- stable query and recall behavior +- clarified graph and KV scope + +### Thread 3: Ingestion +Owner: Asha Mehta +Deliverables: +- extraction adapter plan +- mapping into memory_docs, vector_chunks, and graph_namespace +- sample-data evaluation + +## Current Data Model Notes + +### Documents +Documents should preserve: +- document_id +- namespace +- title +- content +- metadata +- created_at +- updated_at + +### Graph facts +Graph storage should capture facts like: +- Ravi works_on memory-layer-completion +- Asha evaluates neocortex_v2 +- OpenHuman uses JSON-RPC +- memory-layer-completion depends_on API-contract + +### Durable preferences +Examples of durable user or team memory: +- Sanil prefers core-first delivery over UI-first delivery. +- Ravi prefers strict ownership boundaries for parallel agents. +- Asha prefers evaluation fixtures with realistic semi-structured text. + +## Milestones + +### Milestone A +Name: Core contract locked +Due date: 2026-03-18 +Success criteria: +- final RPC method names agreed +- JSON-RPC transport explicitly retained +- response envelope strategy documented + +### Milestone B +Name: Core memory operational +Due date: 2026-03-22 +Success criteria: +- list, delete, query, and recall work in Rust +- stable outputs exist for frontend adaptation + +### Milestone C +Name: Ingestion quality baseline +Due date: 2026-03-26 +Success criteria: +- Gmail-like and Notion-like fixtures ingest successfully +- extracted entities and relations are reviewed manually + +## Risks + +- The frontend currently expects raw values for some memory methods. +- neocortex_v2 preserves duplicate relation evidence, while OpenHuman may prefer aggregation. +- If we do not define request and response models early, parallel agents may diverge. + +## Testing Notes + +Use these sample source types for ingestion tests: +- Gmail thread as raw imported message text +- Notion page as raw exported document text + +Assertions should check for: +- person names +- project names +- ownership relations +- deadlines and dates +- decisions and preferences