poimen-memory/crates/mem-ingest/src/speaker_extractor.rs

//! 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<String>,         // 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<Option<ExtractedSpeaker>>;
}

/// Heuristic Speaker Extractor (pattern-based)
#[derive(Debug, Clone)]
pub struct HeuristicSpeakerExtractor {
    config: SpeakerConfig,
    patterns: Vec<Regex>,
}

impl HeuristicSpeakerExtractor {
    pub fn new(config: SpeakerConfig) -> Result<Self> {
        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<ExtractedSpeaker> {
        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<Option<ExtractedSpeaker>> {
        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<Option<ExtractedSpeaker>> {
        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");
    }
}