//! Entity Linking (Phase 5.1) //! //! Identifies co-references, links text spans to entities, detects aliases, //! and suggests entity merges. use std::collections::{HashMap, HashSet}; use sqlx::PgPool; use serde::{Deserialize, Serialize}; use tracing::{debug, warn}; /// Result of linking a text mention to an entity #[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] pub struct MentionLink { /// The text span that was linked pub mention_text: String, /// Start offset in original text pub start_offset: usize, /// End offset in original text pub end_offset: usize, /// Entity ID it was linked to pub entity_id: String, /// Entity name pub entity_name: String, /// Confidence of link (0.0-1.0) pub confidence: f32, /// Why it was linked (semantic, lexical, alias, etc.) pub reason: LinkReason, } /// Reason for linking #[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Hash)] pub enum LinkReason { /// Semantic similarity (high embedding match) SemanticMatch, /// Lexical match (exact or near-exact string) LexicalMatch, /// Known alias AliasMatch, /// Acronym expansion (e.g., "k8s" → "Kubernetes") AcronymMatch, /// Partial/substring match PartialMatch, } /// Alias suggestion #[derive(Debug, Clone, Serialize, Deserialize)] pub struct AliasSuggestion { /// Entity ID pub entity_id: String, /// Entity name (canonical) pub canonical_name: String, /// Suggested alias pub alias: String, /// Confidence (0.0-1.0) pub confidence: f32, /// How often this alias appears in text pub frequency: usize, } /// Entity merge candidate #[derive(Debug, Clone, Serialize, Deserialize)] pub struct MergeSuggestion { /// Entity 1 ID pub entity1_id: String, /// Entity 1 name pub entity1_name: String, /// Entity 2 ID pub entity2_id: String, /// Entity 2 name pub entity2_name: String, /// Confidence they're the same (0.0-1.0) pub confidence: f32, /// Reasons for merge pub reasons: Vec, } /// Co-reference cluster (multiple mentions of same entity) #[derive(Debug, Clone, Serialize, Deserialize)] pub struct CoreferenceCluster { /// Representative entity ID pub entity_id: String, /// All mention texts in this cluster pub mentions: Vec, /// Mention count pub mention_count: usize, /// Confidence this is correct clustering pub confidence: f32, } /// Entity Linking Engine pub struct EntityLinker { pool: PgPool, } impl EntityLinker { pub fn new(pool: PgPool) -> Self { EntityLinker { pool } } /// Link mentions in text to existing entities /// /// Returns: /// - Vec: Successful links /// - Vec: Unlinked mentions pub async fn link_mentions( &self, text: &str, project_id: &str, ) -> Result<(Vec, Vec), String> { if text.is_empty() { return Ok((vec![], vec![])); } // Extract potential mentions (noun phrases, capitalized sequences) let mentions = self.extract_mentions(text)?; debug!("Extracted {} potential mentions from text", mentions.len()); // Get all entities from database let entities = self.fetch_entities(project_id).await?; debug!("Loaded {} entities from database", entities.len()); let mut links = Vec::new(); let mut unlinked = Vec::new(); for mention in mentions { match self.find_best_link(&mention.text, &entities).await? { Some((entity_id, entity_name, confidence, reason)) => { links.push(MentionLink { mention_text: mention.text.clone(), start_offset: mention.start, end_offset: mention.end, entity_id, entity_name, confidence, reason, }); } None => { unlinked.push(mention.text); } } } Ok((links, unlinked)) } /// Detect aliases for an entity pub async fn detect_aliases( &self, entity_id: &str, entity_name: &str, text_sample: &[String], ) -> Result, String> { let mut aliases = HashMap::new(); for text in text_sample { let mentions = self.extract_mentions(text)?; for mention in mentions { if self.is_similar(&mention.text, entity_name) { let entry = aliases.entry(mention.text.clone()).or_insert((0, 0.5)); entry.0 += 1; } } } // Convert to suggestions, only include frequent ones let suggestions: Vec<_> = aliases .into_iter() .filter(|(_, (count, _))| *count > 1) // At least 2 occurrences .map(|(alias, (frequency, confidence))| AliasSuggestion { entity_id: entity_id.to_string(), canonical_name: entity_name.to_string(), alias, confidence: (confidence * (frequency as f32 / 10.0).min(1.0)).min(1.0), frequency, }) .collect(); Ok(suggestions) } /// Suggest entity merges based on similarity pub async fn suggest_merges( &self, project_id: &str, similarity_threshold: f32, ) -> Result, String> { let entities = self.fetch_entities(project_id).await?; let mut suggestions = Vec::new(); for (i, ent1) in entities.iter().enumerate() { for ent2 in &entities[(i + 1)..] { let similarity = self.compute_similarity(&ent1.name, &ent2.name); if similarity >= similarity_threshold { let mut reasons = Vec::new(); if ent1.name.contains(&ent2.name) || ent2.name.contains(&ent1.name) { reasons.push("Substring match".to_string()); } if self.edit_distance(&ent1.name, &ent2.name) <= 2 { reasons.push("Near edit distance".to_string()); } if self.have_common_relations(&ent1.id, &ent2.id) { reasons.push("Common relations".to_string()); } suggestions.push(MergeSuggestion { entity1_id: ent1.id.clone(), entity1_name: ent1.name.clone(), entity2_id: ent2.id.clone(), entity2_name: ent2.name.clone(), confidence: similarity, reasons, }); } } } Ok(suggestions) } /// Identify coreference clusters pub async fn detect_coreferences( &self, texts: &[String], project_id: &str, ) -> Result, String> { let mut clusters: HashMap> = HashMap::new(); let entities = self.fetch_entities(project_id).await?; for text in texts { let (links, _) = self.link_mentions(text, project_id).await?; for link in links { clusters .entry(link.entity_id) .or_insert_with(Vec::new) .push(link.mention_text); } } let mut result = Vec::new(); for (entity_id, mentions) in clusters { if let Some(entity) = entities.iter().find(|e| e.id == entity_id) { let unique_mentions: Vec<_> = mentions.iter().cloned().collect::>().into_iter().collect(); result.push(CoreferenceCluster { entity_id: entity_id.clone(), mention_count: mentions.len(), confidence: 0.85, // Confidence from linking process mentions: unique_mentions, }); } } Ok(result) } // ========== Private Helper Methods ========== /// Extract potential entity mentions from text fn extract_mentions(&self, text: &str) -> Result, String> { let mut mentions = Vec::new(); // Simple mention extraction: capitalized sequences, quoted text let words: Vec<&str> = text.split_whitespace().collect(); let mut i = 0; while i < words.len() { let word = words[i]; // Capitalized word (potential entity) if word.chars().next().map_or(false, |c| c.is_uppercase()) && word.len() > 2 { let start_pos = text.find(word).unwrap_or(0); let end_pos = start_pos + word.len(); mentions.push(Mention { text: word.to_string(), start: start_pos, end: end_pos, }); // Multi-word entity (consecutive capitalized words) let mut j = i + 1; let mut multi_text = word.to_string(); while j < words.len() && words[j].chars().next().map_or(false, |c| c.is_uppercase()) { multi_text.push(' '); multi_text.push_str(words[j]); j += 1; } if j > i + 1 { let start_pos = text.find(&multi_text).unwrap_or(0); let end_pos = start_pos + multi_text.len(); mentions.push(Mention { text: multi_text, start: start_pos, end: end_pos, }); i = j - 1; } } i += 1; } Ok(mentions) } /// Find best link for a mention async fn find_best_link( &self, mention: &str, entities: &[EntityInfo], ) -> Result, String> { let mut best: Option<(String, String, f32, LinkReason)> = None; for entity in entities { // Check exact match first (highest confidence) if entity.name.eq_ignore_ascii_case(mention) { return Ok(Some(( entity.id.clone(), entity.name.clone(), 0.99, LinkReason::LexicalMatch, ))); } // Check semantic similarity let similarity = self.compute_similarity(mention, &entity.name); if similarity > 0.7 { if best.is_none() || similarity > best.as_ref().unwrap().2 { best = Some(( entity.id.clone(), entity.name.clone(), similarity, LinkReason::SemanticMatch, )); } } // Check acronym (e.g., "k8s" for "Kubernetes") if self.is_acronym(mention, &entity.name) { return Ok(Some(( entity.id.clone(), entity.name.clone(), 0.95, LinkReason::AcronymMatch, ))); } } Ok(best) } /// Fetch all entities for a project async fn fetch_entities(&self, project_id: &str) -> Result, String> { // Stub: would query database // For now, return empty Ok(vec![]) } /// Compute string similarity (Jaro-Winkler style) fn compute_similarity(&self, s1: &str, s2: &str) -> f32 { let s1_lower = s1.to_lowercase(); let s2_lower = s2.to_lowercase(); if s1_lower == s2_lower { return 1.0; } if s1_lower.contains(&s2_lower) || s2_lower.contains(&s1_lower) { return 0.85; } // Simple Levenshtein-based similarity let distance = self.edit_distance(&s1_lower, &s2_lower); let max_len = s1_lower.len().max(s2_lower.len()); 1.0 - (distance as f32 / max_len as f32) } /// Edit distance (Levenshtein) fn edit_distance(&self, s1: &str, s2: &str) -> usize { let len1 = s1.len(); let len2 = s2.len(); let mut dp = vec![vec![0; len2 + 1]; len1 + 1]; for i in 0..=len1 { dp[i][0] = i; } for j in 0..=len2 { dp[0][j] = j; } for (i, c1) in s1.chars().enumerate() { for (j, c2) in s2.chars().enumerate() { let cost = if c1 == c2 { 0 } else { 1 }; dp[i + 1][j + 1] = (dp[i][j + 1] + 1).min(dp[i + 1][j] + 1).min(dp[i][j] + cost); } } dp[len1][len2] } /// Check if s1 is acronym of s2 fn is_acronym(&self, s1: &str, s2: &str) -> bool { if s1.len() > s2.len() || s1.is_empty() { return false; } let words: Vec<&str> = s2.split_whitespace().collect(); let acronym: String = words.iter().filter_map(|w| w.chars().next()).collect(); acronym.to_lowercase() == s1.to_lowercase() } /// Check if two strings are similar fn is_similar(&self, s1: &str, s2: &str) -> bool { self.compute_similarity(s1, s2) > 0.7 } /// Check if two entities have common relations (stub) fn have_common_relations(&self, _id1: &str, _id2: &str) -> bool { // TODO: Query edge table for common neighbors false } } /// Internal mention structure struct Mention { text: String, start: usize, end: usize, } /// Entity info for linking struct EntityInfo { id: String, name: String, } #[cfg(test)] mod tests { use super::*; fn create_linker_mock() -> EntityLinker { // Create with in-memory pool (stub for testing) let pool = sqlx::postgres::PgPoolOptions::new() .max_connections(1) .build_lazy(); EntityLinker::new(pool) } #[test] fn test_extract_mentions_basic() { let linker = create_linker_mock(); let text = "Kubernetes is a container orchestration platform."; let mentions = linker.extract_mentions(text).unwrap(); assert!(mentions.len() > 0); } #[test] fn test_extract_mentions_multiword() { let linker = create_linker_mock(); let text = "Google Cloud Platform provides services."; let mentions = linker.extract_mentions(text).unwrap(); assert!(mentions.iter().any(|m| m.text.contains("Cloud"))); } #[test] fn test_mention_link_structure() { let link = MentionLink { mention_text: "Kubernetes".to_string(), start_offset: 0, end_offset: 10, entity_id: "e1".to_string(), entity_name: "Kubernetes".to_string(), confidence: 0.95, reason: LinkReason::LexicalMatch, }; assert_eq!(link.confidence, 0.95); } #[test] fn test_link_reason_enum() { let reasons = vec![ LinkReason::SemanticMatch, LinkReason::LexicalMatch, LinkReason::AliasMatch, LinkReason::AcronymMatch, LinkReason::PartialMatch, ]; assert_eq!(reasons.len(), 5); } #[test] fn test_alias_suggestion_structure() { let alias = AliasSuggestion { entity_id: "e1".to_string(), canonical_name: "Kubernetes".to_string(), alias: "k8s".to_string(), confidence: 0.9, frequency: 5, }; assert_eq!(alias.frequency, 5); } #[test] fn test_merge_suggestion_structure() { let merge = MergeSuggestion { entity1_id: "e1".to_string(), entity1_name: "Kubernetes".to_string(), entity2_id: "e2".to_string(), entity2_name: "K8s".to_string(), confidence: 0.85, reasons: vec!["Acronym match".to_string()], }; assert_eq!(merge.confidence, 0.85); assert_eq!(merge.reasons.len(), 1); } #[test] fn test_coreference_cluster_structure() { let cluster = CoreferenceCluster { entity_id: "e1".to_string(), mentions: vec!["Kubernetes".to_string(), "k8s".to_string()], mention_count: 2, confidence: 0.85, }; assert_eq!(cluster.mention_count, 2); } #[test] fn test_edit_distance() { let linker = create_linker_mock(); let dist = linker.edit_distance("Kubernetes", "kubernetes"); assert_eq!(dist, 0); // Same lowercase } #[test] fn test_edit_distance_typo() { let linker = create_linker_mock(); let dist = linker.edit_distance("Kubernetes", "Kubenetes"); assert!(dist > 0 && dist < 5); } #[test] fn test_compute_similarity_exact() { let linker = create_linker_mock(); let sim = linker.compute_similarity("test", "test"); assert_eq!(sim, 1.0); } #[test] fn test_compute_similarity_case_insensitive() { let linker = create_linker_mock(); let sim = linker.compute_similarity("Test", "test"); assert_eq!(sim, 1.0); } #[test] fn test_compute_similarity_substring() { let linker = create_linker_mock(); let sim = linker.compute_similarity("Kubernetes", "kubernetes"); assert!(sim > 0.8); } #[test] fn test_is_acronym_true() { let linker = create_linker_mock(); let is_acr = linker.is_acronym("k8s", "Kubernetes"); assert!(is_acr); } #[test] fn test_is_acronym_false() { let linker = create_linker_mock(); let is_acr = linker.is_acronym("test", "Kubernetes"); assert!(!is_acr); } #[test] fn test_is_similar_true() { let linker = create_linker_mock(); let similar = linker.is_similar("Kubernetes", "kubernetes"); assert!(similar); } #[test] fn test_is_similar_false() { let linker = create_linker_mock(); let similar = linker.is_similar("test", "completely different"); assert!(!similar); } #[test] fn test_mention_link_reason_serialization() { let reason = LinkReason::SemanticMatch; let json = serde_json::to_string(&reason).unwrap(); assert!(json.contains("SemanticMatch")); } #[test] fn test_mention_link_full_serialization() { let link = MentionLink { mention_text: "Kubernetes".to_string(), start_offset: 0, end_offset: 10, entity_id: "e1".to_string(), entity_name: "Kubernetes".to_string(), confidence: 0.95, reason: LinkReason::LexicalMatch, }; let json = serde_json::to_string(&link).unwrap(); assert!(json.contains("Kubernetes")); assert!(json.contains("0.95")); } }