|
| 1 | +//! Deterministic entity extraction for associative (spreading-activation) recall. |
| 2 | +//! |
| 3 | +//! SAG-inspired: a chunk's high-signal tokens (error codes, CLI flags, paths, |
| 4 | +//! code symbols) become index entities in `chunk_entities`. Two chunks that share |
| 5 | +//! an entity are associatively linked even when their prose is dissimilar, which |
| 6 | +//! lets `recall` spread activation across that link — the ACT-R associative term |
| 7 | +//! (`S_ji`) that base-level activation alone cannot supply. |
| 8 | +//! |
| 9 | +//! Deterministic on purpose — no LLM, no `regex` dependency (manual char |
| 10 | +//! scanning) — so it runs on the no-LLM capture hot path *and* in migration |
| 11 | +//! backfill. Mirrors the project constraint that capture never depends on the LLM |
| 12 | +//! (see `ResilientDistiller`). |
| 13 | +//! |
| 14 | +//! Deliberately ignores plain lowercase words: those are already recoverable via |
| 15 | +//! the lexical/BM25 channel, and entities must stay **discriminative** so the |
| 16 | +//! ACT-R fan term keeps promiscuous tokens from dominating the spread. |
| 17 | +
|
| 18 | +use std::collections::HashSet; |
| 19 | + |
| 20 | +/// One extracted entity: the normalized join key plus a coarse type tag. The tag |
| 21 | +/// is diagnostic only — recall keys on `entity`, never on `etype`. |
| 22 | +#[derive(Debug, Clone, PartialEq, Eq)] |
| 23 | +pub struct ExtractedEntity { |
| 24 | + pub entity: String, |
| 25 | + pub etype: &'static str, |
| 26 | +} |
| 27 | + |
| 28 | +/// Cap per chunk so a pathological blob can't bloat the index. |
| 29 | +const MAX_ENTITIES_PER_CHUNK: usize = 64; |
| 30 | +const MIN_ENTITY_LEN: usize = 2; |
| 31 | + |
| 32 | +/// Punctuation trimmed from both ends of a token. Excludes `-` `/` `_` `:` so |
| 33 | +/// flags (`--release`), paths (`core/src/x.rs`), snake_case and `::` paths keep |
| 34 | +/// their structure; `.` is included so trailing sentence dots drop while internal |
| 35 | +/// dots (`file.rs`) survive (trim is end-anchored, not internal). |
| 36 | +const TRIM: &[char] = &[ |
| 37 | + '"', '\'', '`', ',', ';', '.', '(', ')', '[', ']', '{', '}', '?', '!', '<', '>', '*', '|', |
| 38 | + '=', '@', '#', '%', '\\', |
| 39 | +]; |
| 40 | + |
| 41 | +/// Extract high-signal entities from a chunk's content and optional trigger desc. |
| 42 | +pub fn extract_entities(content: &str, trigger_desc: Option<&str>) -> Vec<ExtractedEntity> { |
| 43 | + let mut raw: Vec<(String, &'static str)> = Vec::new(); |
| 44 | + for text in [Some(content), trigger_desc].into_iter().flatten() { |
| 45 | + collect_backtick_spans(text, &mut raw); |
| 46 | + for tok in text.split_whitespace() { |
| 47 | + if let Some(cl) = classify(tok) { |
| 48 | + raw.push(cl); |
| 49 | + } |
| 50 | + } |
| 51 | + } |
| 52 | + |
| 53 | + let mut seen: HashSet<String> = HashSet::new(); |
| 54 | + let mut out: Vec<ExtractedEntity> = Vec::new(); |
| 55 | + for (entity, etype) in raw { |
| 56 | + if entity.len() >= MIN_ENTITY_LEN && seen.insert(entity.clone()) { |
| 57 | + out.push(ExtractedEntity { entity, etype }); |
| 58 | + if out.len() >= MAX_ENTITIES_PER_CHUNK { |
| 59 | + break; |
| 60 | + } |
| 61 | + } |
| 62 | + } |
| 63 | + out |
| 64 | +} |
| 65 | + |
| 66 | +/// Backtick-delimited spans are author-marked code, so their tokens clear a lower |
| 67 | +/// bar: any identifier-like piece becomes a `symbol` even if it would otherwise |
| 68 | +/// look like a plain word (e.g. `` `recall` ``). |
| 69 | +fn collect_backtick_spans(text: &str, raw: &mut Vec<(String, &'static str)>) { |
| 70 | + let mut in_span = false; |
| 71 | + let mut buf = String::new(); |
| 72 | + for ch in text.chars() { |
| 73 | + if ch == '`' { |
| 74 | + if in_span { |
| 75 | + for piece in buf.split_whitespace() { |
| 76 | + if let Some(cl) = classify(piece) { |
| 77 | + raw.push(cl); |
| 78 | + } else if let Some(sym) = as_identifier(piece) { |
| 79 | + raw.push((sym, "symbol")); |
| 80 | + } |
| 81 | + } |
| 82 | + buf.clear(); |
| 83 | + } |
| 84 | + in_span = !in_span; |
| 85 | + } else if in_span { |
| 86 | + buf.push(ch); |
| 87 | + } |
| 88 | + } |
| 89 | +} |
| 90 | + |
| 91 | +/// Classify a whitespace-delimited token into a high-signal entity, or `None` for |
| 92 | +/// plain prose. Normalization lowercases for join stability — both the stored |
| 93 | +/// chunk text and the recall query run through this same function. |
| 94 | +fn classify(tok: &str) -> Option<(String, &'static str)> { |
| 95 | + // Flags keep their leading dashes (TRIM excludes '-'); detect before trimming. |
| 96 | + let pre = tok.trim_matches(TRIM); |
| 97 | + if pre.starts_with("--") && pre.len() >= 4 && pre[2..].chars().all(is_ident_char) { |
| 98 | + return Some((pre.to_lowercase(), "flag")); |
| 99 | + } |
| 100 | + let t = pre; |
| 101 | + if t.len() < MIN_ENTITY_LEN { |
| 102 | + return None; |
| 103 | + } |
| 104 | + |
| 105 | + // Error code: short alpha prefix + >=3 digits, nothing else (E0277, GH1234). |
| 106 | + if let Some((alpha, digits)) = split_alpha_digits(t) { |
| 107 | + if (1..=4).contains(&alpha.len()) && digits.len() >= 3 { |
| 108 | + return Some((t.to_lowercase(), "error")); |
| 109 | + } |
| 110 | + } |
| 111 | + |
| 112 | + // Rust-style path / namespaced symbol. |
| 113 | + if t.contains("::") && t.chars().all(is_ident_char) { |
| 114 | + return Some((t.to_lowercase(), "path")); |
| 115 | + } |
| 116 | + // Filesystem path: a slash plus a dotted file component. |
| 117 | + if t.contains('/') && t.contains('.') && t.chars().all(is_ident_char) { |
| 118 | + return Some((t.trim_end_matches('/').to_lowercase(), "path")); |
| 119 | + } |
| 120 | + |
| 121 | + // snake_case / kebab identifier. |
| 122 | + if (t.contains('_') || t.contains('-')) && is_identifier_token(t) { |
| 123 | + return Some((t.to_lowercase(), "symbol")); |
| 124 | + } |
| 125 | + // CamelCase / mixedCase identifier. |
| 126 | + if is_camel_case(t) { |
| 127 | + return Some((t.to_lowercase(), "symbol")); |
| 128 | + } |
| 129 | + None |
| 130 | +} |
| 131 | + |
| 132 | +/// Lower bar used inside backtick spans: accept any identifier-like token. |
| 133 | +fn as_identifier(tok: &str) -> Option<String> { |
| 134 | + let t = tok.trim_matches(TRIM); |
| 135 | + if t.len() >= MIN_ENTITY_LEN && is_identifier_token(t) { |
| 136 | + Some(t.to_lowercase()) |
| 137 | + } else { |
| 138 | + None |
| 139 | + } |
| 140 | +} |
| 141 | + |
| 142 | +fn is_ident_char(c: char) -> bool { |
| 143 | + c.is_alphanumeric() || matches!(c, '_' | '-' | ':' | '.' | '/') |
| 144 | +} |
| 145 | + |
| 146 | +/// All chars are identifier chars and at least one is alphabetic (excludes bare |
| 147 | +/// numbers / version strings like `1.2.3`). |
| 148 | +fn is_identifier_token(t: &str) -> bool { |
| 149 | + t.chars().all(is_ident_char) && t.chars().any(|c| c.is_alphabetic()) |
| 150 | +} |
| 151 | + |
| 152 | +/// Split a token into a leading alphabetic run and a trailing all-digit run, |
| 153 | +/// rejecting anything with other characters in between. |
| 154 | +fn split_alpha_digits(t: &str) -> Option<(&str, &str)> { |
| 155 | + let split = t.find(|c: char| c.is_ascii_digit())?; |
| 156 | + let (alpha, digits) = t.split_at(split); |
| 157 | + if !alpha.is_empty() |
| 158 | + && alpha.chars().all(|c| c.is_ascii_alphabetic()) |
| 159 | + && !digits.is_empty() |
| 160 | + && digits.chars().all(|c| c.is_ascii_digit()) |
| 161 | + { |
| 162 | + Some((alpha, digits)) |
| 163 | + } else { |
| 164 | + None |
| 165 | + } |
| 166 | +} |
| 167 | + |
| 168 | +/// Has an internal lower→upper transition (mixedCase / CamelCase) and is otherwise |
| 169 | +/// alphanumeric — captures `recallCandidates`, `KnowledgeBase`, `BM25`. |
| 170 | +fn is_camel_case(t: &str) -> bool { |
| 171 | + if !t.chars().all(|c| c.is_alphanumeric()) { |
| 172 | + return false; |
| 173 | + } |
| 174 | + let bytes: Vec<char> = t.chars().collect(); |
| 175 | + bytes |
| 176 | + .windows(2) |
| 177 | + .any(|w| w[0].is_lowercase() && w[1].is_uppercase()) |
| 178 | +} |
| 179 | + |
| 180 | +#[cfg(test)] |
| 181 | +mod tests { |
| 182 | + use super::*; |
| 183 | + |
| 184 | + fn ents(content: &str) -> Vec<String> { |
| 185 | + extract_entities(content, None) |
| 186 | + .into_iter() |
| 187 | + .map(|e| e.entity) |
| 188 | + .collect() |
| 189 | + } |
| 190 | + |
| 191 | + #[test] |
| 192 | + fn extracts_error_codes() { |
| 193 | + assert!(ents("hit error E0277 while building").contains(&"e0277".to_string())); |
| 194 | + // Bare numbers and version strings are not entities. |
| 195 | + assert!(!ents("bumped to 1.2.3 in 2026").contains(&"1.2.3".to_string())); |
| 196 | + } |
| 197 | + |
| 198 | + #[test] |
| 199 | + fn extracts_flags_and_paths() { |
| 200 | + let e = ents("run cargo build --release and edit core/src/kb/recall.rs"); |
| 201 | + assert!(e.contains(&"--release".to_string())); |
| 202 | + assert!(e.contains(&"core/src/kb/recall.rs".to_string())); |
| 203 | + } |
| 204 | + |
| 205 | + #[test] |
| 206 | + fn extracts_symbols_not_plain_words() { |
| 207 | + let e = ents("call KnowledgeBase::recall via the snake_case helper get_deps"); |
| 208 | + assert!(e.contains(&"knowledgebase::recall".to_string())); |
| 209 | + assert!(e.contains(&"snake_case".to_string())); |
| 210 | + assert!(e.contains(&"get_deps".to_string())); |
| 211 | + // plain prose is left to the lexical channel |
| 212 | + assert!(!e.contains(&"call".to_string())); |
| 213 | + assert!(!e.contains(&"via".to_string())); |
| 214 | + } |
| 215 | + |
| 216 | + #[test] |
| 217 | + fn backtick_spans_lower_the_bar() { |
| 218 | + let e = ents("the `recall` method matters"); |
| 219 | + assert!(e.contains(&"recall".to_string())); |
| 220 | + // outside backticks, the same lowercase word is ignored |
| 221 | + assert!(!ents("recall the method").contains(&"recall".to_string())); |
| 222 | + } |
| 223 | + |
| 224 | + #[test] |
| 225 | + fn dedups_and_caps() { |
| 226 | + let e = extract_entities("E0277 E0277 E0277", None); |
| 227 | + assert_eq!(e.iter().filter(|x| x.entity == "e0277").count(), 1); |
| 228 | + } |
| 229 | +} |
0 commit comments