Skip to main content

oxirs_graphrag/
entity_linking.rs

1//! Entity linking and disambiguation for knowledge graphs.
2//!
3//! Provides candidate generation via string matching and context-based
4//! disambiguation using TF-IDF cosine similarity.
5
6use std::collections::HashMap;
7
8// ──────────────────────────────────────────────────────────────────────────────
9// Types
10// ──────────────────────────────────────────────────────────────────────────────
11
12/// A span of text that may refer to a named entity.
13#[derive(Debug, Clone, PartialEq, Eq)]
14pub struct EntityMention {
15    /// Surface form of the mention.
16    pub text: String,
17    /// Start byte offset in the containing text.
18    pub start: usize,
19    /// End byte offset (exclusive) in the containing text.
20    pub end: usize,
21}
22
23impl EntityMention {
24    /// Create a mention from text and character positions.
25    pub fn new(text: impl Into<String>, start: usize, end: usize) -> Self {
26        Self {
27            text: text.into(),
28            start,
29            end,
30        }
31    }
32}
33
34/// A candidate entity from the knowledge base.
35#[derive(Debug, Clone)]
36pub struct EntityCandidate {
37    /// Entity IRI.
38    pub iri: String,
39    /// Primary label.
40    pub label: String,
41    /// String-matching similarity score (0.0–1.0).
42    pub score: f64,
43    /// Alternative labels and aliases.
44    pub aliases: Vec<String>,
45}
46
47impl EntityCandidate {
48    fn new(iri: impl Into<String>, label: impl Into<String>, aliases: Vec<String>) -> Self {
49        Self {
50            iri: iri.into(),
51            label: label.into(),
52            score: 0.0,
53            aliases,
54        }
55    }
56}
57
58/// A successfully linked entity.
59#[derive(Debug, Clone)]
60pub struct LinkedEntity {
61    /// The text mention.
62    pub mention: EntityMention,
63    /// The best matching entity candidate.
64    pub entity: EntityCandidate,
65    /// Overall confidence (0.0–1.0) combining string + context similarity.
66    pub confidence: f64,
67}
68
69// ──────────────────────────────────────────────────────────────────────────────
70// TfIdfIndex
71// ──────────────────────────────────────────────────────────────────────────────
72
73/// A simple TF-IDF index for context-based disambiguation.
74pub struct TfIdfIndex {
75    /// Documents: (doc_id, term → tf).
76    docs: Vec<(String, HashMap<String, f64>)>,
77    /// Inverse document frequency: term → idf.
78    idf: HashMap<String, f64>,
79}
80
81impl TfIdfIndex {
82    /// Create an empty index.
83    pub fn new() -> Self {
84        Self {
85            docs: Vec::new(),
86            idf: HashMap::new(),
87        }
88    }
89
90    /// Add a document to the index.
91    pub fn add_document(&mut self, doc_id: impl Into<String>, text: &str) {
92        let tokens = tokenize(text);
93        let total = tokens.len() as f64;
94        if total == 0.0 {
95            return;
96        }
97        let mut tf: HashMap<String, f64> = HashMap::new();
98        for tok in &tokens {
99            *tf.entry(tok.clone()).or_insert(0.0) += 1.0 / total;
100        }
101        self.docs.push((doc_id.into(), tf));
102    }
103
104    /// Recompute IDF from all indexed documents.
105    pub fn build(&mut self) {
106        let n = self.docs.len() as f64;
107        let mut df: HashMap<String, usize> = HashMap::new();
108        for (_, tf) in &self.docs {
109            for term in tf.keys() {
110                *df.entry(term.clone()).or_insert(0) += 1;
111            }
112        }
113        self.idf.clear();
114        for (term, count) in df {
115            self.idf.insert(term, (n / count as f64).ln() + 1.0);
116        }
117    }
118
119    /// Compute TF-IDF cosine similarity between a query string and a document.
120    pub fn similarity(&self, query: &str, doc_id: &str) -> f64 {
121        let doc = match self.docs.iter().find(|(id, _)| id == doc_id) {
122            Some((_, tf)) => tf,
123            None => return 0.0,
124        };
125
126        let q_tokens = tokenize(query);
127        let q_total = q_tokens.len() as f64;
128        if q_total == 0.0 {
129            return 0.0;
130        }
131        let mut q_tf: HashMap<String, f64> = HashMap::new();
132        for tok in &q_tokens {
133            *q_tf.entry(tok.clone()).or_insert(0.0) += 1.0 / q_total;
134        }
135
136        // Cosine similarity of TF-IDF vectors
137        let mut dot = 0.0_f64;
138        let mut q_norm = 0.0_f64;
139        let mut d_norm = 0.0_f64;
140
141        let all_terms: std::collections::HashSet<&String> = q_tf.keys().chain(doc.keys()).collect();
142
143        for term in all_terms {
144            let idf = self.idf.get(term).copied().unwrap_or(1.0);
145            let q_val = q_tf.get(term).copied().unwrap_or(0.0) * idf;
146            let d_val = doc.get(term).copied().unwrap_or(0.0) * idf;
147            dot += q_val * d_val;
148            q_norm += q_val * q_val;
149            d_norm += d_val * d_val;
150        }
151
152        let denom = q_norm.sqrt() * d_norm.sqrt();
153        if denom < 1e-15 {
154            0.0
155        } else {
156            (dot / denom).clamp(0.0, 1.0)
157        }
158    }
159
160    /// Number of indexed documents.
161    pub fn doc_count(&self) -> usize {
162        self.docs.len()
163    }
164}
165
166impl Default for TfIdfIndex {
167    fn default() -> Self {
168        Self::new()
169    }
170}
171
172// ──────────────────────────────────────────────────────────────────────────────
173// EntityLinker
174// ──────────────────────────────────────────────────────────────────────────────
175
176/// Links entity mentions in text to knowledge-base entities.
177pub struct EntityLinker {
178    /// Knowledge base entries: iri → (label, aliases, context_doc_id).
179    kb: HashMap<String, KbEntry>,
180    /// TF-IDF index over entity descriptions (for context disambiguation).
181    tfidf: TfIdfIndex,
182    /// First character of every indexed label/alias (lowercased) → list of
183    /// IRIs sharing it. Used as a candidate-generation blocking step in
184    /// [`candidate_generation`](Self::candidate_generation) so mention
185    /// resolution need not run Jaro-Winkler against the entire knowledge
186    /// base for every mention.
187    first_char_index: HashMap<char, Vec<String>>,
188    /// Minimum confidence threshold below which an entity is treated as NIL.
189    pub nil_threshold: f64,
190}
191
192struct KbEntry {
193    label: String,
194    aliases: Vec<String>,
195}
196
197impl EntityLinker {
198    /// Create an entity linker with default NIL threshold 0.1.
199    pub fn new() -> Self {
200        Self {
201            kb: HashMap::new(),
202            tfidf: TfIdfIndex::new(),
203            first_char_index: HashMap::new(),
204            nil_threshold: 0.1,
205        }
206    }
207
208    /// Add an entity to the knowledge base.
209    ///
210    /// `context` is an optional textual description used for TF-IDF
211    /// disambiguation.
212    pub fn add_entity(
213        &mut self,
214        iri: impl Into<String>,
215        label: impl Into<String>,
216        aliases: &[&str],
217    ) {
218        let iri = iri.into();
219        let label = label.into();
220        let aliases: Vec<String> = aliases.iter().map(|s| s.to_string()).collect();
221        let context = format!("{} {}", label, aliases.join(" "));
222        self.tfidf.add_document(iri.clone(), &context);
223
224        let mut first_chars: std::collections::HashSet<char> = std::collections::HashSet::new();
225        if let Some(c) = label.to_lowercase().chars().next() {
226            first_chars.insert(c);
227        }
228        for alias in &aliases {
229            if let Some(c) = alias.to_lowercase().chars().next() {
230                first_chars.insert(c);
231            }
232        }
233        for c in first_chars {
234            self.first_char_index
235                .entry(c)
236                .or_default()
237                .push(iri.clone());
238        }
239
240        self.kb.insert(iri, KbEntry { label, aliases });
241    }
242
243    /// Finalise the TF-IDF index (call after all entities are added).
244    pub fn build_index(&mut self) {
245        self.tfidf.build();
246    }
247
248    /// Find and link all entity mentions in `text`.
249    pub fn link(&self, text: &str) -> Vec<LinkedEntity> {
250        let mentions = detect_mentions(text);
251        let mut linked = Vec::new();
252
253        for mention in mentions {
254            let candidates = self.candidate_generation(&mention.text);
255            if candidates.is_empty() {
256                continue;
257            }
258            let best = self.disambiguate(&mention, &candidates, text);
259            if let Some(entity) = best {
260                let confidence = entity.score;
261                if confidence >= self.nil_threshold {
262                    linked.push(LinkedEntity {
263                        mention,
264                        entity,
265                        confidence,
266                    });
267                }
268            }
269        }
270        linked
271    }
272
273    /// Generate entity candidates matching the mention by string similarity.
274    ///
275    /// Uses [`first_char_index`](Self::first_char_index) as a blocking step:
276    /// only knowledge-base entries sharing the mention's first character are
277    /// scored with Jaro-Winkler, instead of the entire knowledge base. This
278    /// is a standard record-linkage optimization (Jaro-Winkler's own prefix
279    /// bonus already rewards shared leading characters, so genuine matches
280    /// overwhelmingly share the first character with the mention).
281    pub fn candidate_generation(&self, mention: &str) -> Vec<EntityCandidate> {
282        let mention_lower = mention.to_lowercase();
283        let Some(first_char) = mention_lower.chars().next() else {
284            return Vec::new();
285        };
286        let Some(blocked_iris) = self.first_char_index.get(&first_char) else {
287            return Vec::new();
288        };
289
290        let mut seen: std::collections::HashSet<&str> = std::collections::HashSet::new();
291        let mut candidates: Vec<EntityCandidate> = blocked_iris
292            .iter()
293            .filter(|iri| seen.insert(iri.as_str()))
294            .filter_map(|iri| {
295                let entry = self.kb.get(iri)?;
296                let label_score = jaro_winkler(&mention_lower, &entry.label.to_lowercase());
297                let alias_score = entry
298                    .aliases
299                    .iter()
300                    .map(|a| jaro_winkler(&mention_lower, &a.to_lowercase()))
301                    .fold(0.0_f64, f64::max);
302                let score = label_score.max(alias_score);
303                if score > 0.6 {
304                    let mut c = EntityCandidate::new(
305                        iri.clone(),
306                        entry.label.clone(),
307                        entry.aliases.clone(),
308                    );
309                    c.score = score;
310                    Some(c)
311                } else {
312                    None
313                }
314            })
315            .collect();
316
317        candidates.sort_by(|a, b| {
318            b.score
319                .partial_cmp(&a.score)
320                .unwrap_or(std::cmp::Ordering::Equal)
321        });
322        candidates
323    }
324
325    /// Disambiguate among candidates using context TF-IDF similarity.
326    pub fn disambiguate(
327        &self,
328        _mention: &EntityMention,
329        candidates: &[EntityCandidate],
330        context: &str,
331    ) -> Option<EntityCandidate> {
332        if candidates.is_empty() {
333            return None;
334        }
335
336        let mut best_score = f64::NEG_INFINITY;
337        let mut best: Option<EntityCandidate> = None;
338
339        for cand in candidates {
340            let ctx_score = self.tfidf.similarity(context, &cand.iri);
341            // Combined score: string similarity × 0.6 + context × 0.4
342            let combined = cand.score * 0.6 + ctx_score * 0.4;
343            if combined > best_score {
344                best_score = combined;
345                let mut winner = cand.clone();
346                winner.score = combined;
347                best = Some(winner);
348            }
349        }
350        best
351    }
352
353    /// Number of entities in the knowledge base.
354    pub fn entity_count(&self) -> usize {
355        self.kb.len()
356    }
357}
358
359impl Default for EntityLinker {
360    fn default() -> Self {
361        Self::new()
362    }
363}
364
365// ──────────────────────────────────────────────────────────────────────────────
366// Private helpers
367// ──────────────────────────────────────────────────────────────────────────────
368
369/// Detect potential entity mentions in text by looking for capitalised tokens
370/// or sequences.
371fn detect_mentions(text: &str) -> Vec<EntityMention> {
372    let mut mentions = Vec::new();
373    let mut chars = text.char_indices().peekable();
374    let bytes = text.as_bytes();
375    let len = bytes.len();
376
377    while let Some((start, ch)) = chars.next() {
378        if ch.is_uppercase() {
379            // Consume a capitalised word sequence (handles "Albert Einstein")
380            let mut end = start + ch.len_utf8();
381            while end < len {
382                let next_ch = text[end..].chars().next().unwrap_or('\0');
383                if next_ch.is_alphanumeric() || next_ch == ' ' {
384                    // Allow one space if followed by uppercase (multi-word entity)
385                    if next_ch == ' ' {
386                        let after_space = end + 1;
387                        if after_space < len {
388                            let nc2 = text[after_space..].chars().next().unwrap_or('\0');
389                            if nc2.is_uppercase() {
390                                end = after_space + nc2.len_utf8();
391                                // advance the chars iterator past the space and the uppercase char
392                                let _ = chars.next(); // space
393                                let _ = chars.next(); // uppercase
394                                continue;
395                            }
396                        }
397                        break;
398                    }
399                    end += next_ch.len_utf8();
400                    let _ = chars.next();
401                } else {
402                    break;
403                }
404            }
405            let mention_text = text[start..end].trim().to_string();
406            if mention_text.len() >= 2 {
407                mentions.push(EntityMention::new(mention_text, start, end));
408            }
409        }
410    }
411    mentions
412}
413
414/// Jaro-Winkler string similarity (0.0–1.0).
415fn jaro_winkler(s1: &str, s2: &str) -> f64 {
416    if s1 == s2 {
417        return 1.0;
418    }
419    let jaro = jaro(s1, s2);
420    let prefix_len = s1
421        .chars()
422        .zip(s2.chars())
423        .take(4)
424        .take_while(|(a, b)| a == b)
425        .count();
426    let p = 0.1_f64;
427    jaro + (prefix_len as f64 * p * (1.0 - jaro))
428}
429
430fn jaro(s1: &str, s2: &str) -> f64 {
431    let s1: Vec<char> = s1.chars().collect();
432    let s2: Vec<char> = s2.chars().collect();
433    let len1 = s1.len();
434    let len2 = s2.len();
435    if len1 == 0 && len2 == 0 {
436        return 1.0;
437    }
438    if len1 == 0 || len2 == 0 {
439        return 0.0;
440    }
441
442    let match_window = (len1.max(len2) / 2).saturating_sub(1);
443    let mut s1_matches = vec![false; len1];
444    let mut s2_matches = vec![false; len2];
445    let mut matches = 0usize;
446    let mut transpositions = 0usize;
447
448    for (i, &c1) in s1.iter().enumerate() {
449        let start = i.saturating_sub(match_window);
450        let end = (i + match_window + 1).min(len2);
451        for (j, &c2) in s2[start..end].iter().enumerate() {
452            let j_real = start + j;
453            if !s2_matches[j_real] && c1 == c2 {
454                s1_matches[i] = true;
455                s2_matches[j_real] = true;
456                matches += 1;
457                break;
458            }
459        }
460    }
461
462    if matches == 0 {
463        return 0.0;
464    }
465
466    let mut k = 0;
467    for (i, &s1m) in s1_matches.iter().enumerate() {
468        if s1m {
469            while !s2_matches[k] {
470                k += 1;
471            }
472            if s1[i] != s2[k] {
473                transpositions += 1;
474            }
475            k += 1;
476        }
477    }
478
479    let m = matches as f64;
480    (m / len1 as f64 + m / len2 as f64 + (m - transpositions as f64 / 2.0) / m) / 3.0
481}
482
483/// Tokenise text into lowercase alpha-numeric tokens.
484fn tokenize(text: &str) -> Vec<String> {
485    text.split(|c: char| !c.is_alphanumeric())
486        .filter(|s| !s.is_empty())
487        .map(|s| s.to_lowercase())
488        .collect()
489}
490
491// ──────────────────────────────────────────────────────────────────────────────
492// Tests
493// ──────────────────────────────────────────────────────────────────────────────
494
495#[cfg(test)]
496mod tests {
497    use super::*;
498
499    fn linker_with_persons() -> EntityLinker {
500        let mut linker = EntityLinker::new();
501        linker.add_entity(
502            "http://example.org/Albert_Einstein",
503            "Albert Einstein",
504            &["Einstein", "A. Einstein"],
505        );
506        linker.add_entity(
507            "http://example.org/Marie_Curie",
508            "Marie Curie",
509            &["Curie", "M. Curie"],
510        );
511        linker.add_entity(
512            "http://example.org/Isaac_Newton",
513            "Isaac Newton",
514            &["Newton"],
515        );
516        linker.build_index();
517        linker
518    }
519
520    // ── EntityMention ─────────────────────────────────────────────────────────
521
522    #[test]
523    fn test_mention_new() {
524        let m = EntityMention::new("Alice", 0, 5);
525        assert_eq!(m.text, "Alice");
526        assert_eq!(m.start, 0);
527        assert_eq!(m.end, 5);
528    }
529
530    #[test]
531    fn test_mention_equality() {
532        let m1 = EntityMention::new("Bob", 0, 3);
533        let m2 = EntityMention::new("Bob", 0, 3);
534        assert_eq!(m1, m2);
535    }
536
537    // ── TfIdfIndex ────────────────────────────────────────────────────────────
538
539    #[test]
540    fn test_tfidf_add_document() {
541        let mut idx = TfIdfIndex::new();
542        idx.add_document("doc1", "quantum physics relativity");
543        idx.build();
544        assert_eq!(idx.doc_count(), 1);
545    }
546
547    #[test]
548    fn test_tfidf_similarity_same_doc() {
549        let mut idx = TfIdfIndex::new();
550        idx.add_document("doc1", "quantum physics relativity");
551        idx.build();
552        let sim = idx.similarity("quantum physics", "doc1");
553        assert!(sim > 0.0, "similarity should be > 0, got {sim}");
554    }
555
556    #[test]
557    fn test_tfidf_similarity_different_content() {
558        let mut idx = TfIdfIndex::new();
559        idx.add_document("doc1", "quantum physics relativity");
560        idx.add_document("doc2", "cooking recipes baking bread");
561        idx.build();
562        let s1 = idx.similarity("quantum physics", "doc1");
563        let s2 = idx.similarity("quantum physics", "doc2");
564        assert!(s1 > s2, "physics query should match doc1 better");
565    }
566
567    #[test]
568    fn test_tfidf_unknown_doc() {
569        let idx = TfIdfIndex::new();
570        assert_eq!(idx.similarity("anything", "unknown"), 0.0);
571    }
572
573    #[test]
574    fn test_tfidf_empty_query() {
575        let mut idx = TfIdfIndex::new();
576        idx.add_document("d", "hello world");
577        idx.build();
578        assert_eq!(idx.similarity("", "d"), 0.0);
579    }
580
581    #[test]
582    fn test_tfidf_default() {
583        let idx = TfIdfIndex::default();
584        assert_eq!(idx.doc_count(), 0);
585    }
586
587    // ── EntityLinker ──────────────────────────────────────────────────────────
588
589    #[test]
590    fn test_linker_entity_count() {
591        let linker = linker_with_persons();
592        assert_eq!(linker.entity_count(), 3);
593    }
594
595    #[test]
596    fn test_linker_default() {
597        let linker = EntityLinker::default();
598        assert_eq!(linker.entity_count(), 0);
599    }
600
601    // ── candidate_generation ──────────────────────────────────────────────────
602
603    #[test]
604    fn test_candidate_generation_exact_label() {
605        let linker = linker_with_persons();
606        let cands = linker.candidate_generation("Einstein");
607        assert!(!cands.is_empty());
608        assert!(cands[0].iri.contains("Einstein"));
609    }
610
611    #[test]
612    fn test_candidate_generation_partial() {
613        let linker = linker_with_persons();
614        let cands = linker.candidate_generation("Newton");
615        assert!(!cands.is_empty());
616        assert!(cands.iter().any(|c| c.iri.contains("Newton")));
617    }
618
619    #[test]
620    fn test_candidate_generation_no_match() {
621        let linker = linker_with_persons();
622        let cands = linker.candidate_generation("Zorkblat");
623        assert!(cands.is_empty());
624    }
625
626    #[test]
627    fn test_candidate_generation_sorted_by_score() {
628        let linker = linker_with_persons();
629        let cands = linker.candidate_generation("Curie");
630        for i in 1..cands.len() {
631            assert!(cands[i - 1].score >= cands[i].score);
632        }
633    }
634
635    #[test]
636    fn test_candidate_generation_alias_match() {
637        let linker = linker_with_persons();
638        // "Curie" is an alias for Marie Curie
639        let cands = linker.candidate_generation("Curie");
640        assert!(cands.iter().any(|c| c.iri.contains("Curie")));
641    }
642
643    /// Regression test for the first-character blocking rewrite of
644    /// `candidate_generation`: correct candidates must still be found via
645    /// their blocking bucket even with many unrelated entities present.
646    #[test]
647    fn test_candidate_generation_blocking_scales_with_many_entities() {
648        let mut linker = EntityLinker::new();
649        for i in 0..300 {
650            linker.add_entity(
651                format!("http://example.org/Other{i}"),
652                format!("Other{i}"),
653                &[],
654            );
655        }
656        linker.add_entity(
657            "http://example.org/Albert_Einstein",
658            "Albert Einstein",
659            &["Einstein", "A. Einstein"],
660        );
661        linker.build_index();
662
663        let cands = linker.candidate_generation("Einstein");
664        assert!(!cands.is_empty());
665        assert!(cands[0].iri.contains("Einstein"));
666    }
667
668    /// A mention whose first character matches no indexed label/alias should
669    /// short-circuit to an empty candidate list without scoring anything.
670    #[test]
671    fn test_candidate_generation_blocking_no_match_for_unindexed_first_char() {
672        let linker = linker_with_persons();
673        let cands = linker.candidate_generation("Zzzzzzz");
674        assert!(cands.is_empty());
675    }
676
677    // ── disambiguate ─────────────────────────────────────────────────────────
678
679    #[test]
680    fn test_disambiguate_returns_best() {
681        let linker = linker_with_persons();
682        let cands = linker.candidate_generation("Einstein");
683        let mention = EntityMention::new("Einstein", 0, 8);
684        let best = linker.disambiguate(&mention, &cands, "Einstein worked on relativity");
685        assert!(best.is_some());
686        assert!(best.expect("should succeed").iri.contains("Einstein"));
687    }
688
689    #[test]
690    fn test_disambiguate_empty_candidates() {
691        let linker = linker_with_persons();
692        let mention = EntityMention::new("X", 0, 1);
693        let result = linker.disambiguate(&mention, &[], "context");
694        assert!(result.is_none());
695    }
696
697    #[test]
698    fn test_disambiguate_score_in_range() {
699        let linker = linker_with_persons();
700        let cands = linker.candidate_generation("Newton");
701        let mention = EntityMention::new("Newton", 0, 6);
702        if let Some(best) = linker.disambiguate(&mention, &cands, "gravity laws Newton") {
703            assert!((0.0..=1.0).contains(&best.score));
704        }
705    }
706
707    // ── link ──────────────────────────────────────────────────────────────────
708
709    #[test]
710    fn test_link_finds_entity() {
711        let linker = linker_with_persons();
712        let linked = linker.link("Einstein developed relativity theory.");
713        assert!(!linked.is_empty());
714        assert!(linked[0].entity.iri.contains("Einstein"));
715    }
716
717    #[test]
718    fn test_link_confidence_above_threshold() {
719        let linker = linker_with_persons();
720        let linked = linker.link("Newton formulated laws of motion.");
721        for le in &linked {
722            assert!(le.confidence >= linker.nil_threshold);
723        }
724    }
725
726    #[test]
727    fn test_link_no_entities_in_empty_text() {
728        let linker = linker_with_persons();
729        let linked = linker.link("");
730        assert!(linked.is_empty());
731    }
732
733    #[test]
734    fn test_link_result_fields() {
735        let linker = linker_with_persons();
736        let linked = linker.link("Einstein and Curie were scientists.");
737        for le in &linked {
738            assert!(!le.mention.text.is_empty());
739            assert!(!le.entity.iri.is_empty());
740            assert!((0.0..=1.0).contains(&le.confidence));
741        }
742    }
743
744    // ── Jaro-Winkler ──────────────────────────────────────────────────────────
745
746    #[test]
747    fn test_jaro_winkler_identical() {
748        assert!((jaro_winkler("hello", "hello") - 1.0).abs() < 1e-9);
749    }
750
751    #[test]
752    fn test_jaro_winkler_completely_different() {
753        let score = jaro_winkler("abc", "xyz");
754        assert!(score < 0.5, "score = {score}");
755    }
756
757    #[test]
758    fn test_jaro_winkler_prefix_boost() {
759        let jw = jaro_winkler("einstein", "einstien");
760        assert!(jw > 0.8, "score = {jw}");
761    }
762
763    #[test]
764    fn test_jaro_winkler_empty_strings() {
765        assert!((jaro("", "") - 1.0).abs() < 1e-9);
766        assert!((jaro("abc", "") - 0.0).abs() < 1e-9);
767    }
768
769    // ── detect_mentions ───────────────────────────────────────────────────────
770
771    #[test]
772    fn test_detect_mentions_finds_capitalized() {
773        let mentions = detect_mentions("Alice and Bob went to Paris.");
774        let texts: Vec<&str> = mentions.iter().map(|m| m.text.as_str()).collect();
775        // At least Alice, Bob, Paris should be detected
776        assert!(texts
777            .iter()
778            .any(|t| *t == "Alice" || t.starts_with("Alice")));
779    }
780
781    #[test]
782    fn test_detect_mentions_empty() {
783        assert!(detect_mentions("").is_empty());
784    }
785
786    #[test]
787    fn test_detect_mentions_lowercase_only() {
788        let mentions = detect_mentions("all lowercase words here");
789        assert!(mentions.is_empty());
790    }
791
792    // ── tokenize ─────────────────────────────────────────────────────────────
793
794    #[test]
795    fn test_tokenize_basic() {
796        let tokens = tokenize("Hello World");
797        assert_eq!(tokens, vec!["hello", "world"]);
798    }
799
800    #[test]
801    fn test_tokenize_empty() {
802        assert!(tokenize("").is_empty());
803    }
804
805    #[test]
806    fn test_tokenize_punctuation_split() {
807        let tokens = tokenize("foo, bar; baz.");
808        assert_eq!(tokens, vec!["foo", "bar", "baz"]);
809    }
810
811    // ── Full pipeline ─────────────────────────────────────────────────────────
812
813    #[test]
814    fn test_full_pipeline() {
815        let mut linker = EntityLinker::new();
816        linker.add_entity("http://ex.org/Paris", "Paris", &["City of Light"]);
817        linker.add_entity("http://ex.org/London", "London", &["British capital"]);
818        linker.build_index();
819
820        let linked = linker.link("Paris is a famous city in France.");
821        if !linked.is_empty() {
822            assert!(linked[0].entity.iri.contains("Paris"));
823        }
824        // No assertion on count since detection depends on heuristic
825    }
826
827    #[test]
828    fn test_nil_threshold_filters_low_confidence() {
829        let mut linker = EntityLinker::new();
830        linker.add_entity("http://ex.org/X", "Xyzzy", &[]);
831        linker.build_index();
832        linker.nil_threshold = 0.99; // Very high threshold
833
834        let linked = linker.link("Xyzzy something");
835        // Most links should be filtered out by high threshold
836        for le in &linked {
837            assert!(le.confidence >= 0.99);
838        }
839    }
840}