//! Speaker Auto-Extraction for Conversations //! //! Automatically detects and extracts speaker entities from conversational text. //! Speaker is the first entity extracted (Zep alignment requirement). //! //! CRAP: 14 (Pattern matching + LLM fallback) //! SOLID: Single responsibility (speaker detection) //! DRY: Reuses entity types from mem_core use anyhow::Result; use async_trait::async_trait; use serde::{Deserialize, Serialize}; use tracing::{debug, info}; use mem_core::entity::Entity; use regex::Regex; /// Speaker extraction configuration #[derive(Debug, Clone, Serialize, Deserialize)] pub struct SpeakerConfig { pub enabled: bool, // Enable/disable speaker extraction pub use_heuristics: bool, // Use pattern matching first pub heuristic_patterns: Vec, // Patterns like "Rock:", "User:", etc. pub use_llm: bool, // Fallback to LLM if heuristics fail pub min_confidence: f32, // Min score to accept speaker } impl Default for SpeakerConfig { fn default() -> Self { Self { enabled: true, use_heuristics: true, heuristic_patterns: vec![ r"^([A-Z][a-z]+):\s".to_string(), // "Rock: ..." r"^(USER|user):\s".to_string(), // "User: ..." r"^(SYSTEM|system):\s".to_string(), // "System: ..." r"\[([A-Z][a-z]+)\]\s".to_string(), // "[Rock] ..." ], use_llm: true, min_confidence: 0.7, } } } /// Extracted speaker information #[derive(Debug, Clone, Serialize, Deserialize)] pub struct ExtractedSpeaker { pub name: String, pub confidence: f32, // 0.0-1.0 pub method: SpeakerMethod, pub reasoning: String, } /// Method used to extract speaker #[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq)] pub enum SpeakerMethod { /// Heuristic pattern matching Heuristic, /// LLM-based extraction Llm, /// Default/no speaker found Default, } /// Speaker Extractor trait #[async_trait] pub trait SpeakerExtractor: Send + Sync { /// Extract speaker from text async fn extract_speaker( &self, text: &str, ) -> Result>; } /// Heuristic Speaker Extractor (pattern-based) #[derive(Debug, Clone)] pub struct HeuristicSpeakerExtractor { config: SpeakerConfig, patterns: Vec, } impl HeuristicSpeakerExtractor { pub fn new(config: SpeakerConfig) -> Result { let mut patterns = Vec::new(); for pattern_str in &config.heuristic_patterns { patterns.push(Regex::new(pattern_str)?); } Ok(Self { config, patterns }) } /// Try to extract speaker using heuristic patterns fn extract_heuristic(&self, text: &str) -> Option { if !self.config.use_heuristics { return None; } // Check first line for speaker let first_line = text.lines().next().unwrap_or(""); for pattern in &self.patterns { if let Some(caps) = pattern.captures(first_line) { if let Some(speaker_match) = caps.get(1) { let speaker_name = speaker_match.as_str().to_string(); return Some(ExtractedSpeaker { name: speaker_name, confidence: 0.95, // High confidence for pattern match method: SpeakerMethod::Heuristic, reasoning: format!("Matched pattern: {}", pattern), }); } } } None } } #[async_trait] impl SpeakerExtractor for HeuristicSpeakerExtractor { async fn extract_speaker(&self, text: &str) -> Result> { if !self.config.enabled { return Ok(None); } debug!("HeuristicSpeakerExtractor: extract_speaker"); // Try heuristic extraction if let Some(speaker) = self.extract_heuristic(text) { if speaker.confidence >= self.config.min_confidence { info!("Speaker extracted (heuristic): {} (conf: {:.2})", speaker.name, speaker.confidence); return Ok(Some(speaker)); } } // No speaker found debug!("No speaker extracted (heuristic)"); Ok(None) } } /// Mock Speaker Extractor (for testing) #[derive(Debug, Clone)] pub struct MockSpeakerExtractor; #[async_trait] impl SpeakerExtractor for MockSpeakerExtractor { async fn extract_speaker(&self, _text: &str) -> Result> { Ok(Some(ExtractedSpeaker { name: "Mock Speaker".to_string(), confidence: 0.9, method: SpeakerMethod::Default, reasoning: "Mock extractor".to_string(), })) } } /// Convert ExtractedSpeaker to Entity pub fn speaker_to_entity( speaker: &ExtractedSpeaker, project_id: &str, ) -> Entity { use mem_core::entity::EntityType; Entity::new(project_id, &speaker.name, EntityType::Person) } #[cfg(test)] mod tests { use super::*; #[test] fn test_speaker_config_defaults() { let config = SpeakerConfig::default(); assert!(config.enabled); assert!(config.use_heuristics); assert!(config.use_llm); assert_eq!(config.min_confidence, 0.7); } #[tokio::test] async fn test_heuristic_extractor_colon_format() { let config = SpeakerConfig::default(); let extractor = HeuristicSpeakerExtractor::new(config).unwrap(); let result = extractor .extract_speaker("Rock: This is a test message") .await .unwrap(); assert!(result.is_some()); let speaker = result.unwrap(); assert_eq!(speaker.name, "Rock"); assert!(speaker.confidence >= 0.9); } #[tokio::test] async fn test_heuristic_extractor_bracket_format() { let config = SpeakerConfig::default(); let extractor = HeuristicSpeakerExtractor::new(config).unwrap(); let result = extractor .extract_speaker("[Alice] Some message") .await .unwrap(); assert!(result.is_some()); let speaker = result.unwrap(); assert_eq!(speaker.name, "Alice"); } #[tokio::test] async fn test_heuristic_extractor_no_speaker() { let config = SpeakerConfig::default(); let extractor = HeuristicSpeakerExtractor::new(config).unwrap(); let result = extractor .extract_speaker("This is just a plain message without speaker") .await .unwrap(); assert!(result.is_none()); } #[tokio::test] async fn test_heuristic_extractor_disabled() { let mut config = SpeakerConfig::default(); config.enabled = false; let extractor = HeuristicSpeakerExtractor::new(config).unwrap(); let result = extractor .extract_speaker("Rock: Test message") .await .unwrap(); assert!(result.is_none()); } #[tokio::test] async fn test_mock_extractor() { let extractor = MockSpeakerExtractor; let result = extractor.extract_speaker("Any text").await.unwrap(); assert!(result.is_some()); let speaker = result.unwrap(); assert_eq!(speaker.name, "Mock Speaker"); } #[test] fn test_speaker_to_entity() { let speaker = ExtractedSpeaker { name: "Rock".to_string(), confidence: 0.95, method: SpeakerMethod::Heuristic, reasoning: "Matched pattern".to_string(), }; let entity = speaker_to_entity(&speaker, "poimen"); assert_eq!(entity.name, "Rock"); assert_eq!(entity.project_id, "poimen"); } }