Skip to main content

oxirs_graphrag/
entity_linker.rs

1//! String-to-RDF entity linking: mention detection and candidate ranking.
2//!
3//! Provides exact and fuzzy matching of text mentions to entities in a
4//! knowledge base, with configurable score thresholds.
5
6use std::collections::HashMap;
7
8// ---------------------------------------------------------------------------
9// Domain types
10// ---------------------------------------------------------------------------
11
12/// An entity in the knowledge base.
13#[derive(Debug, Clone)]
14pub struct Entity {
15    /// Canonical IRI identifier.
16    pub iri: String,
17    /// Primary label.
18    pub label: String,
19    /// Alternative surface forms / aliases.
20    pub aliases: Vec<String>,
21    /// Entity type string (e.g. "Person", "Organization", "Place").
22    pub entity_type: String,
23    /// Optional short description.
24    pub description: Option<String>,
25    /// Popularity score in \[0.0, 1.0\] used as a prior.
26    pub popularity: f64,
27}
28
29/// A mention of an entity surface form found in text.
30#[derive(Debug, Clone)]
31pub struct EntityMention {
32    /// The surface form detected in the input text.
33    pub text: String,
34    /// Byte offset of the first character.
35    pub start_char: usize,
36    /// Byte offset one past the last character.
37    pub end_char: usize,
38    /// Ranked candidate links for this mention.
39    pub candidates: Vec<LinkCandidate>,
40}
41
42/// A single candidate link returned for a mention.
43#[derive(Debug, Clone)]
44pub struct LinkCandidate {
45    /// The candidate entity.
46    pub entity: Entity,
47    /// Combined ranking score.
48    pub score: f64,
49    /// Normalised string similarity to the mention text.
50    pub string_similarity: f64,
51    /// Entity popularity used as prior.
52    pub prior_probability: f64,
53}
54
55/// A resolved link between a text mention and its best-matching entity.
56#[derive(Debug, Clone)]
57pub struct LinkedEntity {
58    /// The text mention.
59    pub mention: EntityMention,
60    /// Highest-scoring candidate (if any exceeds the linker threshold).
61    pub best_candidate: Option<LinkCandidate>,
62    /// Confidence score of the best candidate (0.0 if none).
63    pub confidence: f64,
64}
65
66// ---------------------------------------------------------------------------
67// LinkerError
68// ---------------------------------------------------------------------------
69
70/// Errors returned by `EntityLinker`.
71#[derive(Debug)]
72pub enum LinkerError {
73    /// The supplied text is empty.
74    EmptyText,
75    /// The knowledge base has no entities.
76    EmptyKnowledgeBase,
77    /// An entity in the knowledge base is invalid.
78    InvalidEntity(String),
79}
80
81impl std::fmt::Display for LinkerError {
82    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
83        match self {
84            LinkerError::EmptyText => write!(f, "Input text is empty"),
85            LinkerError::EmptyKnowledgeBase => write!(f, "Knowledge base contains no entities"),
86            LinkerError::InvalidEntity(msg) => write!(f, "Invalid entity: {}", msg),
87        }
88    }
89}
90
91impl std::error::Error for LinkerError {}
92
93// ---------------------------------------------------------------------------
94// EntityLinker
95// ---------------------------------------------------------------------------
96
97/// String-to-entity linker backed by an in-memory knowledge base.
98pub struct EntityLinker {
99    entities: Vec<Entity>,
100    /// label (lowercased) → list of entity indices
101    label_index: HashMap<String, Vec<usize>>,
102    /// First character of every indexed surface form → list of entity
103    /// indices, used as a cheap candidate-generation blocking step for fuzzy
104    /// matching in [`link_mention`](Self::link_mention) so it need not score
105    /// every entity in the knowledge base.
106    first_char_index: HashMap<char, Vec<usize>>,
107    /// Longest indexed label/alias, in `char`s. Bounds the substring window
108    /// scanned at each text position in [`detect_mentions`](Self::detect_mentions).
109    max_label_chars: usize,
110    /// Minimum character length of a text span to be considered as a mention.
111    pub min_mention_len: usize,
112    /// Minimum combined score for a candidate to be returned.
113    pub threshold: f64,
114}
115
116impl EntityLinker {
117    // -----------------------------------------------------------------------
118    // Construction
119    // -----------------------------------------------------------------------
120
121    /// Create an empty linker with the given minimum combined score threshold.
122    pub fn new(threshold: f64) -> Self {
123        Self {
124            entities: Vec::new(),
125            label_index: HashMap::new(),
126            first_char_index: HashMap::new(),
127            max_label_chars: 0,
128            min_mention_len: 2,
129            threshold,
130        }
131    }
132
133    /// Add a single entity to the knowledge base and update the index.
134    pub fn add_entity(&mut self, entity: Entity) {
135        let idx = self.entities.len();
136        // Index the label and all aliases
137        let mut keys: Vec<String> = entity.aliases.iter().map(|a| a.to_lowercase()).collect();
138        keys.push(entity.label.to_lowercase());
139
140        let mut first_chars: std::collections::HashSet<char> = std::collections::HashSet::new();
141        for key in &keys {
142            self.max_label_chars = self.max_label_chars.max(key.chars().count());
143            if let Some(c) = key.chars().next() {
144                first_chars.insert(c);
145            }
146        }
147        for c in first_chars {
148            self.first_char_index.entry(c).or_default().push(idx);
149        }
150
151        for key in keys {
152            self.label_index.entry(key).or_default().push(idx);
153        }
154        self.entities.push(entity);
155    }
156
157    /// Add multiple entities in bulk and rebuild the index.
158    pub fn add_entities(&mut self, entities: Vec<Entity>) {
159        for entity in entities {
160            self.add_entity(entity);
161        }
162    }
163
164    /// Total number of entities in the knowledge base.
165    pub fn entity_count(&self) -> usize {
166        self.entities.len()
167    }
168
169    // -----------------------------------------------------------------------
170    // Mention detection
171    // -----------------------------------------------------------------------
172
173    /// Scan `text` for all known entity labels / aliases and return mentions.
174    ///
175    /// Rather than scanning the whole text once per indexed label (which is
176    /// `O(vocabulary × text length)`), this walks the text once and, at every
177    /// character position, performs bounded `O(1)` hash lookups into
178    /// [`label_index`](Self::label_index) for substrings up to
179    /// [`max_label_chars`](Self::max_label_chars) long — giving
180    /// `O(text length × longest label)` overall, independent of vocabulary
181    /// size. Overlapping spans from the same starting position are
182    /// de-duplicated in favour of the longer match.
183    pub fn detect_mentions(&self, text: &str) -> Vec<EntityMention> {
184        let text_lower = text.to_lowercase();
185        let mut mentions: Vec<EntityMention> = Vec::new();
186
187        if self.label_index.is_empty() || self.max_label_chars == 0 {
188            return mentions;
189        }
190
191        // Byte offset of every character boundary, plus one past the end, so
192        // substrings can be sliced without ever splitting a multi-byte UTF-8
193        // character.
194        let boundaries: Vec<usize> = text_lower
195            .char_indices()
196            .map(|(byte_idx, _)| byte_idx)
197            .chain(std::iter::once(text_lower.len()))
198            .collect();
199        let num_chars = boundaries.len() - 1;
200
201        for start in 0..num_chars {
202            let start_byte = boundaries[start];
203            let max_len = self.max_label_chars.min(num_chars - start);
204
205            for len in self.min_mention_len..=max_len {
206                let end_char = start + len;
207                let end_byte = boundaries[end_char];
208                let surface = &text_lower[start_byte..end_byte];
209
210                let Some(indices) = self.label_index.get(surface) else {
211                    continue;
212                };
213
214                let candidates = self.candidates_for_surface(surface, indices);
215                if !candidates.is_empty() {
216                    // Use the original casing from the source text
217                    let original_text = &text[start_byte..end_byte];
218                    mentions.push(EntityMention {
219                        text: original_text.to_string(),
220                        start_char: start_byte,
221                        end_char: end_byte,
222                        candidates,
223                    });
224                }
225            }
226        }
227
228        // Sort by position, then deduplicate overlapping spans (keep longest)
229        mentions.sort_by_key(|m| (m.start_char, usize::MAX - (m.end_char - m.start_char)));
230        let mut deduped: Vec<EntityMention> = Vec::new();
231        for mention in mentions {
232            // Skip if fully contained in the last kept mention
233            if let Some(last) = deduped.last() {
234                if mention.start_char >= last.start_char && mention.end_char <= last.end_char {
235                    continue;
236                }
237            }
238            deduped.push(mention);
239        }
240        deduped
241    }
242
243    // -----------------------------------------------------------------------
244    // Linking
245    // -----------------------------------------------------------------------
246
247    /// Link all mentions found in `text` to their best entity candidates.
248    ///
249    /// # Errors
250    /// - `LinkerError::EmptyText` if `text` is empty after trimming.
251    /// - `LinkerError::EmptyKnowledgeBase` if no entities have been added.
252    pub fn link(&self, text: &str) -> Result<Vec<LinkedEntity>, LinkerError> {
253        if text.trim().is_empty() {
254            return Err(LinkerError::EmptyText);
255        }
256        if self.entities.is_empty() {
257            return Err(LinkerError::EmptyKnowledgeBase);
258        }
259
260        let mentions = self.detect_mentions(text);
261        let linked = mentions
262            .into_iter()
263            .map(|mention| {
264                let best = mention
265                    .candidates
266                    .iter()
267                    .find(|c| c.score >= self.threshold)
268                    .cloned();
269                let confidence = best.as_ref().map(|c| c.score).unwrap_or(0.0);
270                LinkedEntity {
271                    mention,
272                    best_candidate: best,
273                    confidence,
274                }
275            })
276            .collect();
277        Ok(linked)
278    }
279
280    /// Find all entity candidates for an arbitrary mention string.
281    ///
282    /// Uses [`first_char_index`](Self::first_char_index) as a candidate-
283    /// generation blocking step — a standard record-linkage technique — so
284    /// only entities sharing the mention's first character are scored,
285    /// instead of computing edit-distance similarity against the entire
286    /// knowledge base. Results are sorted by `score` descending.
287    pub fn link_mention(&self, mention: &str) -> Vec<LinkCandidate> {
288        let mention_lower = mention.to_lowercase();
289        let Some(first_char) = mention_lower.chars().next() else {
290            return Vec::new();
291        };
292        let Some(blocked_indices) = self.first_char_index.get(&first_char) else {
293            return Vec::new();
294        };
295
296        let mut seen: std::collections::HashSet<usize> = std::collections::HashSet::new();
297        let mut candidates: Vec<LinkCandidate> = Vec::new();
298
299        for &idx in blocked_indices {
300            if !seen.insert(idx) {
301                continue;
302            }
303            let Some(entity) = self.entities.get(idx) else {
304                continue;
305            };
306
307            // Collect all surface forms of this entity
308            let mut surfaces: Vec<String> =
309                entity.aliases.iter().map(|a| a.to_lowercase()).collect();
310            surfaces.push(entity.label.to_lowercase());
311
312            let best_sim = surfaces
313                .iter()
314                .map(|s| Self::string_similarity(&mention_lower, s))
315                .fold(0.0_f64, f64::max);
316
317            if best_sim > 0.0 {
318                let score = Self::combined_score(best_sim, entity.popularity);
319                candidates.push(LinkCandidate {
320                    entity: entity.clone(),
321                    score,
322                    string_similarity: best_sim,
323                    prior_probability: entity.popularity,
324                });
325            }
326        }
327
328        // Sort by score descending
329        candidates.sort_by(|a, b| {
330            b.score
331                .partial_cmp(&a.score)
332                .unwrap_or(std::cmp::Ordering::Equal)
333        });
334        candidates
335    }
336
337    /// Look up an entity by its exact IRI.
338    pub fn find_by_iri(&self, iri: &str) -> Option<&Entity> {
339        self.entities.iter().find(|e| e.iri == iri)
340    }
341
342    // -----------------------------------------------------------------------
343    // String similarity
344    // -----------------------------------------------------------------------
345
346    /// Normalised edit-distance similarity: 1 − edit_distance(a, b) / max(|a|, |b|).
347    ///
348    /// Returns 1.0 for identical strings and 0.0 when the edit distance equals
349    /// the length of the longer string.
350    pub fn string_similarity(a: &str, b: &str) -> f64 {
351        if a == b {
352            return 1.0;
353        }
354        let len_a = a.chars().count();
355        let len_b = b.chars().count();
356        let max_len = len_a.max(len_b);
357        if max_len == 0 {
358            return 1.0; // both empty
359        }
360        let dist = Self::edit_distance(a, b);
361        1.0 - (dist as f64 / max_len as f64)
362    }
363
364    /// Levenshtein edit distance between two strings.
365    fn edit_distance(a: &str, b: &str) -> usize {
366        let a_chars: Vec<char> = a.chars().collect();
367        let b_chars: Vec<char> = b.chars().collect();
368        let la = a_chars.len();
369        let lb = b_chars.len();
370
371        if la == 0 {
372            return lb;
373        }
374        if lb == 0 {
375            return la;
376        }
377
378        let mut dp = vec![vec![0usize; lb + 1]; la + 1];
379        for (i, row) in dp.iter_mut().enumerate() {
380            row[0] = i;
381        }
382        for (j, cell) in dp[0].iter_mut().enumerate() {
383            *cell = j;
384        }
385        for i in 1..=la {
386            for j in 1..=lb {
387                let cost = if a_chars[i - 1] == b_chars[j - 1] {
388                    0
389                } else {
390                    1
391                };
392                dp[i][j] = (dp[i - 1][j] + 1)
393                    .min(dp[i][j - 1] + 1)
394                    .min(dp[i - 1][j - 1] + cost);
395            }
396        }
397        dp[la][lb]
398    }
399
400    // -----------------------------------------------------------------------
401    // Internal helpers
402    // -----------------------------------------------------------------------
403
404    /// Rebuild the label index (and the derived blocking indices) from scratch.
405    #[allow(dead_code)]
406    fn rebuild_index(&mut self) {
407        self.label_index.clear();
408        self.first_char_index.clear();
409        self.max_label_chars = 0;
410        for (idx, entity) in self.entities.iter().enumerate() {
411            let mut keys: Vec<String> = entity.aliases.iter().map(|a| a.to_lowercase()).collect();
412            keys.push(entity.label.to_lowercase());
413
414            let mut first_chars: std::collections::HashSet<char> = std::collections::HashSet::new();
415            for key in &keys {
416                self.max_label_chars = self.max_label_chars.max(key.chars().count());
417                if let Some(c) = key.chars().next() {
418                    first_chars.insert(c);
419                }
420            }
421            for c in first_chars {
422                self.first_char_index.entry(c).or_default().push(idx);
423            }
424
425            for key in keys {
426                self.label_index.entry(key).or_default().push(idx);
427            }
428        }
429    }
430
431    /// Build candidates for a known surface form and entity index list.
432    fn candidates_for_surface(&self, surface: &str, indices: &[usize]) -> Vec<LinkCandidate> {
433        let mut candidates: Vec<LinkCandidate> = indices
434            .iter()
435            .filter_map(|&idx| {
436                let entity = self.entities.get(idx)?;
437                let sim = Self::string_similarity(surface, &entity.label.to_lowercase());
438                let score = Self::combined_score(sim, entity.popularity);
439                Some(LinkCandidate {
440                    entity: entity.clone(),
441                    score,
442                    string_similarity: sim,
443                    prior_probability: entity.popularity,
444                })
445            })
446            .collect();
447        candidates.sort_by(|a, b| {
448            b.score
449                .partial_cmp(&a.score)
450                .unwrap_or(std::cmp::Ordering::Equal)
451        });
452        candidates
453    }
454
455    /// Combined ranking score: 0.7 × string_similarity + 0.3 × popularity.
456    fn combined_score(sim: f64, popularity: f64) -> f64 {
457        0.7 * sim + 0.3 * popularity
458    }
459}
460
461// ---------------------------------------------------------------------------
462// Tests
463// ---------------------------------------------------------------------------
464
465#[cfg(test)]
466mod tests {
467    use super::*;
468
469    fn make_entity(iri: &str, label: &str, aliases: &[&str], pop: f64) -> Entity {
470        Entity {
471            iri: iri.to_string(),
472            label: label.to_string(),
473            aliases: aliases.iter().map(|s| s.to_string()).collect(),
474            entity_type: "Thing".to_string(),
475            description: None,
476            popularity: pop,
477        }
478    }
479
480    fn sample_linker() -> EntityLinker {
481        let mut linker = EntityLinker::new(0.3);
482        linker.add_entity(make_entity(
483            "http://example.org/einstein",
484            "Albert Einstein",
485            &["Einstein", "A. Einstein"],
486            0.95,
487        ));
488        linker.add_entity(make_entity(
489            "http://example.org/curie",
490            "Marie Curie",
491            &["Curie", "M. Curie"],
492            0.90,
493        ));
494        linker.add_entity(make_entity(
495            "http://example.org/berlin",
496            "Berlin",
497            &["Berlin City"],
498            0.80,
499        ));
500        linker
501    }
502
503    // --- entity_count -------------------------------------------------------
504
505    #[test]
506    fn test_entity_count_empty() {
507        let linker = EntityLinker::new(0.5);
508        assert_eq!(linker.entity_count(), 0);
509    }
510
511    #[test]
512    fn test_entity_count_after_add() {
513        let mut linker = EntityLinker::new(0.5);
514        linker.add_entity(make_entity("http://x.org/a", "Alpha", &[], 0.5));
515        assert_eq!(linker.entity_count(), 1);
516        linker.add_entity(make_entity("http://x.org/b", "Beta", &[], 0.5));
517        assert_eq!(linker.entity_count(), 2);
518    }
519
520    #[test]
521    fn test_add_entities_bulk() {
522        let mut linker = EntityLinker::new(0.5);
523        linker.add_entities(vec![
524            make_entity("http://x.org/a", "Alpha", &[], 0.5),
525            make_entity("http://x.org/b", "Beta", &[], 0.6),
526            make_entity("http://x.org/c", "Gamma", &[], 0.7),
527        ]);
528        assert_eq!(linker.entity_count(), 3);
529    }
530
531    // --- find_by_iri --------------------------------------------------------
532
533    #[test]
534    fn test_find_by_iri_exists() {
535        let linker = sample_linker();
536        let entity = linker.find_by_iri("http://example.org/einstein");
537        assert!(entity.is_some());
538        assert_eq!(entity.expect("some").label, "Albert Einstein");
539    }
540
541    #[test]
542    fn test_find_by_iri_not_found() {
543        let linker = sample_linker();
544        assert!(linker.find_by_iri("http://example.org/nobody").is_none());
545    }
546
547    // --- string_similarity --------------------------------------------------
548
549    #[test]
550    fn test_string_similarity_exact() {
551        assert!((EntityLinker::string_similarity("hello", "hello") - 1.0).abs() < 1e-9);
552    }
553
554    #[test]
555    fn test_string_similarity_completely_different() {
556        let sim = EntityLinker::string_similarity("abc", "xyz");
557        assert!(sim < 1.0);
558    }
559
560    #[test]
561    fn test_string_similarity_both_empty() {
562        assert!((EntityLinker::string_similarity("", "") - 1.0).abs() < 1e-9);
563    }
564
565    #[test]
566    fn test_string_similarity_one_empty() {
567        let sim = EntityLinker::string_similarity("", "hello");
568        assert!((sim - 0.0).abs() < 1e-9);
569    }
570
571    #[test]
572    fn test_string_similarity_near_match() {
573        let sim = EntityLinker::string_similarity("Einstein", "Einsten");
574        assert!(sim > 0.8, "sim = {}", sim);
575    }
576
577    #[test]
578    fn test_string_similarity_range() {
579        let sim = EntityLinker::string_similarity("kitten", "sitting");
580        assert!((0.0..=1.0).contains(&sim));
581    }
582
583    // --- edit_distance -------------------------------------------------------
584
585    #[test]
586    fn test_edit_distance_identical() {
587        assert_eq!(EntityLinker::edit_distance("abc", "abc"), 0);
588    }
589
590    #[test]
591    fn test_edit_distance_one_empty() {
592        assert_eq!(EntityLinker::edit_distance("", "abc"), 3);
593        assert_eq!(EntityLinker::edit_distance("abc", ""), 3);
594    }
595
596    #[test]
597    fn test_edit_distance_kitten_sitting() {
598        assert_eq!(EntityLinker::edit_distance("kitten", "sitting"), 3);
599    }
600
601    #[test]
602    fn test_edit_distance_sunday_saturday() {
603        assert_eq!(EntityLinker::edit_distance("sunday", "saturday"), 3);
604    }
605
606    #[test]
607    fn test_edit_distance_single_char() {
608        assert_eq!(EntityLinker::edit_distance("a", "b"), 1);
609        assert_eq!(EntityLinker::edit_distance("a", "a"), 0);
610    }
611
612    // --- detect_mentions ----------------------------------------------------
613
614    #[test]
615    fn test_detect_mentions_exact_label() {
616        let linker = sample_linker();
617        let mentions = linker.detect_mentions("Albert Einstein was a physicist.");
618        assert!(!mentions.is_empty(), "expected at least one mention");
619        let texts: Vec<&str> = mentions.iter().map(|m| m.text.as_str()).collect();
620        assert!(
621            texts.iter().any(|t| t.to_lowercase().contains("einstein")),
622            "mentions = {:?}",
623            texts
624        );
625    }
626
627    #[test]
628    fn test_detect_mentions_alias() {
629        let linker = sample_linker();
630        let mentions = linker.detect_mentions("Curie discovered radium.");
631        let texts: Vec<&str> = mentions.iter().map(|m| m.text.as_str()).collect();
632        assert!(
633            texts.iter().any(|t| t.to_lowercase() == "curie"),
634            "mentions = {:?}",
635            texts
636        );
637    }
638
639    #[test]
640    fn test_detect_mentions_case_insensitive() {
641        let linker = sample_linker();
642        let mentions = linker.detect_mentions("berlin is a great city.");
643        assert!(!mentions.is_empty(), "expected mention of Berlin");
644    }
645
646    #[test]
647    fn test_detect_mentions_multiple() {
648        let linker = sample_linker();
649        let mentions = linker.detect_mentions("Einstein visited Berlin.");
650        // Should find at least two distinct mentions
651        assert!(mentions.len() >= 2, "mentions = {:?}", mentions.len());
652    }
653
654    #[test]
655    fn test_detect_mentions_no_match() {
656        let linker = sample_linker();
657        let mentions = linker.detect_mentions("The quick brown fox jumps.");
658        assert!(
659            mentions.is_empty(),
660            "expected no mentions, got {:?}",
661            mentions
662        );
663    }
664
665    // --- link ---------------------------------------------------------------
666
667    #[test]
668    fn test_link_basic() {
669        let linker = sample_linker();
670        let linked = linker
671            .link("Albert Einstein won the Nobel Prize.")
672            .expect("ok");
673        assert!(!linked.is_empty());
674    }
675
676    #[test]
677    fn test_link_empty_text_error() {
678        let linker = sample_linker();
679        assert!(linker.link("").is_err());
680        assert!(linker.link("   ").is_err());
681    }
682
683    #[test]
684    fn test_link_empty_kb_error() {
685        let linker = EntityLinker::new(0.5);
686        assert!(matches!(
687            linker.link("some text"),
688            Err(LinkerError::EmptyKnowledgeBase)
689        ));
690    }
691
692    #[test]
693    fn test_link_confidence_populated() {
694        let linker = sample_linker();
695        let linked = linker.link("Curie was born in Poland.").expect("ok");
696        for le in &linked {
697            if le.best_candidate.is_some() {
698                assert!(le.confidence > 0.0);
699            }
700        }
701    }
702
703    // --- link_mention -------------------------------------------------------
704
705    #[test]
706    fn test_link_mention_exact() {
707        let linker = sample_linker();
708        let candidates = linker.link_mention("Einstein");
709        assert!(!candidates.is_empty());
710        // Top candidate should be Einstein
711        assert_eq!(candidates[0].entity.iri, "http://example.org/einstein");
712    }
713
714    #[test]
715    fn test_link_mention_sorted_descending() {
716        let linker = sample_linker();
717        let candidates = linker.link_mention("Berlin");
718        for window in candidates.windows(2) {
719            assert!(
720                window[0].score >= window[1].score,
721                "candidates not sorted: {} < {}",
722                window[0].score,
723                window[1].score
724            );
725        }
726    }
727
728    #[test]
729    fn test_link_mention_returns_all_above_zero() {
730        let linker = sample_linker();
731        let candidates = linker.link_mention("Curie");
732        // All returned candidates should have positive score
733        for c in &candidates {
734            assert!(c.score > 0.0);
735        }
736    }
737
738    // --- threshold ----------------------------------------------------------
739
740    #[test]
741    fn test_threshold_filters_low_confidence() {
742        let mut linker = EntityLinker::new(0.99); // very high threshold
743        linker.add_entity(make_entity(
744            "http://x.org/z",
745            "Zephyr",
746            &[],
747            0.1, // low popularity
748        ));
749        let linked = linker.link("There is a zephyr wind.").expect("ok");
750        // Either no mentions or no best_candidate passes threshold
751        for le in &linked {
752            assert!(
753                le.best_candidate.is_none()
754                    || le.best_candidate.as_ref().expect("some").score >= 0.99
755            );
756        }
757    }
758
759    #[test]
760    fn test_min_mention_len_filter() {
761        let mut linker = EntityLinker::new(0.0);
762        linker.add_entity(make_entity("http://x.org/a", "AI", &[], 0.5));
763        linker.min_mention_len = 5;
764        // "AI" is 2 chars — below threshold
765        let mentions = linker.detect_mentions("AI is transforming the world.");
766        assert!(
767            mentions.is_empty() || mentions.iter().all(|m| m.text.len() >= 5),
768            "unexpected short mention found"
769        );
770    }
771
772    // --- Error display -------------------------------------------------------
773
774    #[test]
775    fn test_linker_error_display() {
776        assert!(LinkerError::EmptyText.to_string().contains("empty"));
777        assert!(LinkerError::EmptyKnowledgeBase
778            .to_string()
779            .contains("no entities"));
780        assert!(LinkerError::InvalidEntity("bad".to_string())
781            .to_string()
782            .contains("bad"));
783    }
784
785    // --- Combined score ------------------------------------------------------
786
787    #[test]
788    fn test_combined_score_perfect() {
789        // 0.7 * 1.0 + 0.3 * 1.0 = 1.0
790        let linker = sample_linker();
791        let candidates = linker.link_mention("Albert Einstein");
792        // Best candidate should have high score
793        if let Some(c) = candidates.first() {
794            assert!(c.score > 0.5, "score = {}", c.score);
795        }
796    }
797
798    // --- Alias detection ----------------------------------------------------
799
800    #[test]
801    fn test_alias_detection_full_alias() {
802        let linker = sample_linker();
803        let mentions = linker.detect_mentions("A. Einstein changed physics.");
804        let texts: Vec<String> = mentions.iter().map(|m| m.text.to_lowercase()).collect();
805        assert!(
806            texts.iter().any(|t| t.contains("einstein")),
807            "texts = {:?}",
808            texts
809        );
810    }
811
812    // --- Regression: indexed lookups (perf fix) ------------------------------
813
814    /// Regression test for the windowed hash-lookup rewrite of
815    /// `detect_mentions`: with a large number of unrelated entities in the
816    /// knowledge base, a mention of an entity added last must still be found
817    /// (this previously relied on scanning the whole text once per indexed
818    /// label; now it scans the text once and hash-probes the index).
819    #[test]
820    fn test_detect_mentions_indexed_lookup_with_many_entities() {
821        let mut linker = EntityLinker::new(0.3);
822        for i in 0..500 {
823            linker.add_entity(make_entity(
824                &format!("http://example.org/e{i}"),
825                &format!("Entity{i}"),
826                &[],
827                0.5,
828            ));
829        }
830        linker.add_entity(make_entity(
831            "http://example.org/einstein",
832            "Albert Einstein",
833            &["Einstein"],
834            0.95,
835        ));
836
837        let mentions = linker.detect_mentions("Albert Einstein was a physicist.");
838        let texts: Vec<&str> = mentions.iter().map(|m| m.text.as_str()).collect();
839        assert!(
840            texts.iter().any(|t| t.to_lowercase().contains("einstein")),
841            "mentions = {:?}",
842            texts
843        );
844    }
845
846    /// Regression test for the first-character blocking rewrite of
847    /// `link_mention`: the correct entity must still be found via its
848    /// blocking bucket.
849    #[test]
850    fn test_link_mention_blocking_matches_correct_entity() {
851        let linker = sample_linker();
852        let candidates = linker.link_mention("Curie");
853        assert!(
854            candidates
855                .iter()
856                .any(|c| c.entity.iri == "http://example.org/curie"),
857            "candidates = {:?}",
858            candidates.iter().map(|c| &c.entity.iri).collect::<Vec<_>>()
859        );
860    }
861
862    /// A mention whose first character matches no indexed surface form
863    /// should short-circuit to an empty result without scoring every entity.
864    #[test]
865    fn test_link_mention_blocking_no_match_for_unindexed_first_char() {
866        let linker = sample_linker();
867        let candidates = linker.link_mention("Zzzzz nonexistent name");
868        assert!(candidates.is_empty(), "candidates = {:?}", candidates);
869    }
870}