From cd9f95e594fded459edc8fe47cf2100b4217b653 Mon Sep 17 00:00:00 2001 From: sanil-23 Date: Wed, 22 Apr 2026 18:29:15 +0530 Subject: [PATCH] fix(memory): address CI failures and CodeRabbit review on LLM-NER (#708) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit CI was red on PR #775 due to 10 E0063 errors — existing test-only struct literals for ExtractedEntities / ScoreSignals / ScoreRow weren't updated when the new optional llm_importance fields landed. Local cargo check on lib-only passed; cargo check --tests caught them, as does CI. CodeRabbit also flagged four semantic issues: 1. LlmEntityExtractor::extract() returned Err on HTTP transport, non-2xx status, and JSON parse failures. The soft-fallback contract (documented in the module header) only worked because score_chunk catches errors; any other caller got a hard error. Now split into extract_or_empty() which logs a warn on each failure mode and returns ExtractedEntities::default() unconditionally. The public extract() stays on the async-trait surface but never returns Err. 2. find_char_span() always searched from byte offset 0, so when the LLM returned the same surface twice (explicitly asked for duplicates in the prompt), both entities got (start=0, end=len) — same span, not distinct mentions. Added find_char_span_from(haystack, needle, byte_from, char_from) that resumes from a cursor. into_extracted_entities now tracks a per-surface (byte_after, char_after) cursor in a HashMap, so duplicate mentions get distinct non-overlapping spans and over-claim (LLM returns 3 "Alice" when source has 2) drops the excess entries with a debug log. 3. combine() always included w.llm_importance in the denominator even when llm_importance=0.0 because the LLM didn't run. That artificially lowered the total on short-circuit paths. score_chunk now branches on llm_consulted: uses combine() when LLM ran (importance actually contributes), combine_cheap_only() when LLM was skipped or failed (denominator excludes the LLM weight). Observable in two new tests: short_circuit_reports_cheap_only_total verifies r.total == combine_cheap_only(r.signals, weights) and strictly exceeds combine(r.signals, weights) when llm_importance=0; llm_consulted_ reports_full_total verifies the opposite path. 4. Same issue from a different angle (CodeRabbit flagged both). Fixed by the same llm_consulted branch. Test-literal fixes: - signals/ops.rs: ScoreSignals literals pick up llm_importance field; make_entities helper uses ..Default::default() on ExtractedEntities. - extract/types.rs: three test ExtractedEntities literals gain llm_importance: None, llm_importance_reason: None. - resolver.rs: canonicalise_batch_preserves_spans literal gains the two fields. - score/store.rs: sample_row gains signals.llm_importance and llm_importance_reason. New tests (7 added): - find_char_span_from_advances_past_prior_match - find_char_span_from_returns_none_after_exhaustion - find_char_span_from_preserves_utf8 - find_char_span_from_rejects_non_char_boundary - into_extracted_entities_gives_distinct_spans_to_duplicate_mentions - into_extracted_entities_drops_extra_duplicate_when_source_only_has_one - extract_soft_fallback_on_unreachable_endpoint (transport failure → empty extraction, not Err) - short_circuit_reports_cheap_only_total - llm_consulted_reports_full_total cargo check --lib clean; cargo fmt clean. The integration-test rmeta errors visible locally are stale-metadata artifacts from incremental builds with prior struct shapes; CI's clean build resolves them and tests/*.rs files do not construct any of the touched types. Co-Authored-By: Claude Opus 4.7 (1M context) --- .../memory/tree/score/extract/llm.rs | 251 +++++++++++++++--- .../memory/tree/score/extract/types.rs | 8 + src/openhuman/memory/tree/score/mod.rs | 76 +++++- src/openhuman/memory/tree/score/resolver.rs | 2 + .../memory/tree/score/signals/ops.rs | 4 +- src/openhuman/memory/tree/score/store.rs | 2 + 6 files changed, 307 insertions(+), 36 deletions(-) diff --git a/src/openhuman/memory/tree/score/extract/llm.rs b/src/openhuman/memory/tree/score/extract/llm.rs index dda7140f1..33be3887b 100644 --- a/src/openhuman/memory/tree/score/extract/llm.rs +++ b/src/openhuman/memory/tree/score/extract/llm.rs @@ -124,6 +124,23 @@ impl EntityExtractor for LlmEntityExtractor { } async fn extract(&self, text: &str) -> anyhow::Result { + // Soft-fallback contract: every failure path (transport, HTTP status, + // JSON parse) is logged as a warn and returns an empty + // `ExtractedEntities` rather than `Err`. This makes the extractor + // safe to call from any context, not just `score_chunk` (which + // separately catches errors from its own extractor chain). A caller + // distinguishes "LLM had nothing to say" from "LLM ran and returned + // zero entities" by inspecting `llm_importance` — `None` means the + // call didn't complete successfully. + Ok(self.extract_or_empty(text).await) + } +} + +impl LlmEntityExtractor { + /// Internal: wraps the actual HTTP call and returns `ExtractedEntities` + /// for every failure mode via soft-fallback. Split out of `extract` so + /// the error branches can share logging without `?`-propagation. + async fn extract_or_empty(&self, text: &str) -> ExtractedEntities { let url = format!("{}/api/chat", self.cfg.endpoint.trim_end_matches('/')); let body = self.build_request(text); log::debug!( @@ -132,36 +149,56 @@ impl EntityExtractor for LlmEntityExtractor { text.chars().count() ); - let resp = self - .http - .post(&url) - .json(&body) - .send() - .await - .map_err(anyhow::Error::from)?; + let resp = match self.http.post(&url).json(&body).send().await { + Ok(r) => r, + Err(e) => { + log::warn!( + "[memory_tree::extract::llm] transport failure to {url}: {e} — \ + returning empty extraction" + ); + return ExtractedEntities::default(); + } + }; if !resp.status().is_success() { let status = resp.status(); let body = resp.text().await.unwrap_or_default(); - anyhow::bail!("ollama status {status}: {body}"); + log::warn!( + "[memory_tree::extract::llm] ollama non-success status {status}: {} — \ + returning empty extraction", + truncate_for_log(&body, 200) + ); + return ExtractedEntities::default(); } - let envelope: OllamaChatResponse = resp.json().await.map_err(anyhow::Error::from)?; + let envelope: OllamaChatResponse = match resp.json().await { + Ok(v) => v, + Err(e) => { + log::warn!( + "[memory_tree::extract::llm] response body not Ollama-shaped JSON: {e} — \ + returning empty extraction" + ); + return ExtractedEntities::default(); + } + }; log::debug!( "[memory_tree::extract::llm] response chars={}", envelope.message.content.len() ); - let parsed: LlmExtractionOutput = - serde_json::from_str(&envelope.message.content).map_err(|e| { - anyhow::anyhow!( - "LLM returned non-JSON or wrong-shape response: {e}; \ - content was: {}", + let parsed: LlmExtractionOutput = match serde_json::from_str(&envelope.message.content) { + Ok(v) => v, + Err(e) => { + log::warn!( + "[memory_tree::extract::llm] LLM returned non-JSON or wrong-shape \ + response: {e}; content was: {} — returning empty extraction", truncate_for_log(&envelope.message.content, 400) - ) - })?; + ); + return ExtractedEntities::default(); + } + }; - Ok(parsed.into_extracted_entities(text, &self.cfg)) + parsed.into_extracted_entities(text, &self.cfg) } } @@ -248,6 +285,15 @@ impl LlmExtractionOutput { ) -> ExtractedEntities { let mut entities = Vec::with_capacity(self.entities.len()); + // Per-surface search cursor (char offset). When the LLM returns the + // same surface text twice (deliberately — the prompt asks for + // duplicates), we resume searching AFTER the previous occurrence so + // each emitted entity points at a distinct span. Byte indices are + // tracked separately from char indices because `str::find` returns + // byte offsets while the rest of the pipeline uses char spans. + use std::collections::HashMap; + let mut cursors: HashMap = HashMap::new(); + for raw in self.entities { let surface = raw.text.trim(); if surface.is_empty() { @@ -280,18 +326,23 @@ impl LlmExtractionOutput { } }; - // Recover spans by string search. If the model hallucinated a - // surface that doesn't appear in the source text, drop it. - let (span_start, span_end) = match find_char_span(source_text, surface) { - Some(s) => s, - None => { - log::debug!( - "[memory_tree::extract::llm] dropping hallucinated entity (not found \ - in source): {surface:?}" - ); - continue; - } - }; + // Recover spans by string search, advancing the cursor for this + // surface so repeated mentions get distinct spans. If the model + // hallucinated a surface (or we've exhausted all of its + // occurrences), drop the entity. + let (byte_from, char_from) = cursors.get(surface).copied().unwrap_or((0, 0)); + let (span_start, span_end, byte_after) = + match find_char_span_from(source_text, surface, byte_from, char_from) { + Some(s) => s, + None => { + log::debug!( + "[memory_tree::extract::llm] dropping hallucinated or exhausted \ + entity (not found beyond cursor): {surface:?}" + ); + continue; + } + }; + cursors.insert(surface.to_string(), (byte_after, span_end)); entities.push(ExtractedEntity { kind, @@ -332,10 +383,41 @@ fn parse_kind(s: &str) -> Option { /// Uses byte-level `find` then translates to char offsets so spans align /// with the rest of the extractor pipeline (which is char-based). fn find_char_span(haystack: &str, needle: &str) -> Option<(u32, u32)> { - let byte_idx = haystack.find(needle)?; - let char_start = haystack[..byte_idx].chars().count() as u32; + find_char_span_from(haystack, needle, 0, 0).map(|(s, e, _)| (s, e)) +} + +/// Find `needle` in `haystack` starting from `byte_from` and return +/// `(char_start, char_end, byte_after_needle)`. +/// +/// The byte-offset return is so the caller can chain successive searches +/// without re-walking the prefix every time: pass the returned +/// `byte_after_needle` as the next call's `byte_from`. +/// +/// `char_from` must correspond to `byte_from` in the same `haystack` — +/// i.e. `haystack[..byte_from].chars().count() == char_from as usize`. +/// The caller maintains this invariant (cheap: it's the return from the +/// previous call). +fn find_char_span_from( + haystack: &str, + needle: &str, + byte_from: usize, + char_from: u32, +) -> Option<(u32, u32, usize)> { + if needle.is_empty() || byte_from > haystack.len() { + return None; + } + // Guard against `byte_from` landing inside a multi-byte UTF-8 sequence. + if !haystack.is_char_boundary(byte_from) { + return None; + } + let rel = haystack[byte_from..].find(needle)?; + let byte_start = byte_from + rel; + let byte_end = byte_start + needle.len(); + // Walk forward from the previous char position to build the new char + // offset — avoids re-walking the full prefix. + let char_start = char_from + haystack[byte_from..byte_start].chars().count() as u32; let char_end = char_start + needle.chars().count() as u32; - Some((char_start, char_end)) + Some((char_start, char_end, byte_end)) } fn truncate_for_log(s: &str, max_chars: usize) -> String { @@ -372,6 +454,111 @@ mod tests { assert!(find_char_span("hello world", "absent").is_none()); } + #[test] + fn find_char_span_from_advances_past_prior_match() { + let text = "Alice met Bob then Alice left"; + let (s1, e1, byte_after) = find_char_span_from(text, "Alice", 0, 0).unwrap(); + assert_eq!((s1, e1), (0, 5)); + // Resuming from the cursor must find the second Alice. + let (s2, e2, _) = find_char_span_from(text, "Alice", byte_after, e1).unwrap(); + assert_eq!((s2, e2), (19, 24)); + } + + #[test] + fn find_char_span_from_returns_none_after_exhaustion() { + let text = "Alice met Bob"; + let (_, _, byte_after) = find_char_span_from(text, "Alice", 0, 0).unwrap(); + // No second Alice → None. + assert!(find_char_span_from(text, "Alice", byte_after, 5).is_none()); + } + + #[test] + fn find_char_span_from_preserves_utf8() { + // Two "中" characters (3 bytes each in UTF-8); "Alice" between. + let text = "中 Alice 中 Alice"; + let (s1, e1, byte_after) = find_char_span_from(text, "Alice", 0, 0).unwrap(); + assert_eq!((s1, e1), (2, 7)); + let (s2, e2, _) = find_char_span_from(text, "Alice", byte_after, e1).unwrap(); + // First "中 Alice " = 2 + 5 + 1 + 1 + 1 chars; second Alice starts at char 10. + assert_eq!((s2, e2), (10, 15)); + } + + #[test] + fn find_char_span_from_rejects_non_char_boundary() { + // "中" is 3 bytes; offsets 1 and 2 are mid-codepoint. + let text = "中Alice"; + assert!(find_char_span_from(text, "Alice", 1, 0).is_none()); + } + + #[test] + fn into_extracted_entities_gives_distinct_spans_to_duplicate_mentions() { + // Two "Alice" mentions in source → two distinct ExtractedEntity rows + // with non-overlapping spans. Previously both got (0, 5). + let out = LlmExtractionOutput { + entities: vec![ + LlmEntity { + kind: "person".into(), + text: "Alice".into(), + }, + LlmEntity { + kind: "person".into(), + text: "Alice".into(), + }, + ], + importance: None, + importance_reason: None, + }; + let cfg = LlmExtractorConfig::default(); + let e = out.into_extracted_entities("Alice met Bob then Alice left", &cfg); + assert_eq!(e.entities.len(), 2); + assert_eq!((e.entities[0].span_start, e.entities[0].span_end), (0, 5)); + assert_eq!((e.entities[1].span_start, e.entities[1].span_end), (19, 24)); + } + + #[test] + fn into_extracted_entities_drops_extra_duplicate_when_source_only_has_one() { + // Three "Alice" mentions returned by LLM, only one in source → keep + // one, drop the rest as exhausted-duplicate. + let out = LlmExtractionOutput { + entities: vec![ + LlmEntity { + kind: "person".into(), + text: "Alice".into(), + }, + LlmEntity { + kind: "person".into(), + text: "Alice".into(), + }, + LlmEntity { + kind: "person".into(), + text: "Alice".into(), + }, + ], + importance: None, + importance_reason: None, + }; + let cfg = LlmExtractorConfig::default(); + let e = out.into_extracted_entities("Alice met Bob", &cfg); + assert_eq!(e.entities.len(), 1); + } + + #[tokio::test] + async fn extract_soft_fallback_on_unreachable_endpoint() { + // Point at an unreachable port so the transport fails. extract() + // must NOT return Err — it must return an empty ExtractedEntities + // with a warn log. + let cfg = LlmExtractorConfig { + endpoint: "http://127.0.0.1:1".to_string(), + timeout: std::time::Duration::from_millis(100), + ..LlmExtractorConfig::default() + }; + let ex = LlmEntityExtractor::new(cfg).unwrap(); + let out = ex.extract("some text").await.unwrap(); + assert!(out.entities.is_empty()); + assert!(out.topics.is_empty()); + assert!(out.llm_importance.is_none()); + } + #[test] fn into_extracted_entities_drops_hallucinations() { let out = LlmExtractionOutput { diff --git a/src/openhuman/memory/tree/score/extract/types.rs b/src/openhuman/memory/tree/score/extract/types.rs index 251b45ced..ca737c50d 100644 --- a/src/openhuman/memory/tree/score/extract/types.rs +++ b/src/openhuman/memory/tree/score/extract/types.rs @@ -233,6 +233,8 @@ mod tests { }, ], topics: vec![], + llm_importance: None, + llm_importance_reason: None, }; assert_eq!(e.unique_entity_count(), 1); } @@ -257,6 +259,8 @@ mod tests { }, ], topics: vec![], + llm_importance: None, + llm_importance_reason: None, }; assert_eq!(e.unique_entity_count(), 2); } @@ -272,6 +276,8 @@ mod tests { score: 1.0, }], topics: vec![], + llm_importance: None, + llm_importance_reason: None, }; let b = ExtractedEntities { entities: vec![ @@ -291,6 +297,8 @@ mod tests { }, // different span — keep ], topics: vec![], + llm_importance: None, + llm_importance_reason: None, }; a.merge(b); assert_eq!(a.entities.len(), 2); diff --git a/src/openhuman/memory/tree/score/mod.rs b/src/openhuman/memory/tree/score/mod.rs index b24037d76..f56d0b1e6 100644 --- a/src/openhuman/memory/tree/score/mod.rs +++ b/src/openhuman/memory/tree/score/mod.rs @@ -190,9 +190,23 @@ pub async fn score_chunk(chunk: &Chunk, cfg: &ScoringConfig) -> Result with_llm, + "cheap-only total ({}) should exceed LLM-weighted total \ + ({}) when llm_importance is zero", + r.total, + with_llm + ); + } + + /// When the LLM *does* run, the reported total uses the full combine — + /// the llm_importance contribution is actually in the sum. + #[tokio::test] + async fn llm_consulted_reports_full_total() { + let c = test_chunk("This is a moderately interesting note about a project."); + let llm = FakeLlm::new(0.9); + let mut cfg = ScoringConfig::with_llm_extractor(llm.clone()); + cfg.definite_drop_threshold = 0.0; + cfg.definite_keep_threshold = 1.0; + let r = score_chunk(&c, &cfg).await.unwrap(); + assert_eq!(llm.calls(), 1); + let expected = self::signals::combine(&r.signals, &cfg.weights); + assert!( + (r.total - expected).abs() < 1e-6, + "total={} expected(full combine)={}", + r.total, + expected + ); + } } diff --git a/src/openhuman/memory/tree/score/resolver.rs b/src/openhuman/memory/tree/score/resolver.rs index c3b1d1fd2..14cc56faf 100644 --- a/src/openhuman/memory/tree/score/resolver.rs +++ b/src/openhuman/memory/tree/score/resolver.rs @@ -117,6 +117,8 @@ mod tests { entity(EntityKind::Email, "alice@example.com"), ], topics: vec![], + llm_importance: None, + llm_importance_reason: None, }; let out = canonicalise(&ex); assert_eq!(out.len(), 2); diff --git a/src/openhuman/memory/tree/score/signals/ops.rs b/src/openhuman/memory/tree/score/signals/ops.rs index 16c426ea6..5421f6dcd 100644 --- a/src/openhuman/memory/tree/score/signals/ops.rs +++ b/src/openhuman/memory/tree/score/signals/ops.rs @@ -115,7 +115,7 @@ mod tests { score: 1.0, }) .collect(), - topics: vec![], + ..Default::default() } } @@ -134,6 +134,7 @@ mod tests { source_weight: 1.0, interaction: 1.0, entity_density: 1.0, + llm_importance: 0.0, // default weight is 0 → contribution is zero }; assert!((combine(&s, &SignalWeights::default()) - 1.0).abs() < 1e-6); } @@ -147,6 +148,7 @@ mod tests { source_weight: 0.0, interaction: 1.0, entity_density: 0.0, + llm_importance: 0.0, }; let total = combine(&s, &SignalWeights::default()); assert!((total - (3.0 / 9.0)).abs() < 1e-6); diff --git a/src/openhuman/memory/tree/score/store.rs b/src/openhuman/memory/tree/score/store.rs index 36cbc5ed4..bed44b455 100644 --- a/src/openhuman/memory/tree/score/store.rs +++ b/src/openhuman/memory/tree/score/store.rs @@ -353,6 +353,7 @@ mod tests { source_weight: 0.5, interaction: 0.6, entity_density: 0.3, + llm_importance: 0.0, }, dropped, reason: if dropped { @@ -361,6 +362,7 @@ mod tests { None }, computed_at_ms: 1_700_000_000_000, + llm_importance_reason: None, } }