use crate::types::*;
use ahash::RandomState;
use dashmap::DashMap;
use std::collections::HashMap;
use std::sync::Arc;
use tokio::sync::RwLock;
use tracing::{debug, info};
use uuid::Uuid;
#[derive(Clone)]
pub struct IntentClassifier {
training_data: Arc<RwLock<Vec<TrainingExample>>>,
vocabulary: Arc<DashMap<String, usize, RandomState>>,
intent_patterns: Arc<DashMap<IntentId, Vec<String>, RandomState>>,
config: ClassifierConfig,
}
impl IntentClassifier {
pub async fn new() -> Result<Self> {
Self::with_config(ClassifierConfig::default()).await
}
pub async fn with_config(config: ClassifierConfig) -> Result<Self> {
let classifier = Self {
training_data: Arc::new(RwLock::new(Vec::new())),
vocabulary: Arc::new(DashMap::with_hasher(RandomState::new())),
intent_patterns: Arc::new(DashMap::with_hasher(RandomState::new())),
config,
};
classifier.load_bootstrap_data().await?;
if classifier.config.debug_mode {
info!("Intent classifier initialized with {} dimensions", classifier.config.feature_dimensions);
}
Ok(classifier)
}
pub async fn predict_intent(&self, text: &str) -> Result<IntentPrediction> {
let start_time = std::time::Instant::now();
if self.config.debug_mode {
debug!("Classifying intent for text: '{}'", text);
}
if let Some(exact_match) = self.find_exact_match(text).await? {
return Ok(exact_match);
}
let features = self.extract_features(text).await?;
let intent_scores = self.calculate_intent_scores(&features).await?;
let (best_intent, best_confidence) = self.find_best_intent(&intent_scores)?;
let alternative_intents = self.get_alternative_intents(&intent_scores, &best_intent);
let reasoning = self.generate_reasoning(text, &best_intent, &features).await;
let prediction = IntentPrediction {
intent: best_intent,
confidence: best_confidence,
alternative_intents,
reasoning,
};
if self.config.debug_mode {
let elapsed = start_time.elapsed();
info!("Intent prediction: {} (confidence: {:.3}, time: {:?})",
prediction.intent, prediction.confidence.value(), elapsed);
}
Ok(prediction)
}
pub async fn classify(&self, request: ClassificationRequest) -> Result<ClassificationResponse> {
let start_time = std::time::Instant::now();
let request_id = Uuid::new_v4();
let mut prediction = self.predict_intent(&request.text).await?;
if !request.include_alternatives {
prediction.alternative_intents.clear();
}
if !request.include_reasoning {
prediction.reasoning = String::new();
}
let processing_time_ms = start_time.elapsed().as_millis() as f64;
Ok(ClassificationResponse {
prediction,
processing_time_ms,
request_id,
})
}
pub async fn add_training_example(&self, example: TrainingExample) -> Result<()> {
if example.text.trim().is_empty() {
return Err(IntentError::InvalidParameter {
parameter: "text".to_string(),
message: "Training example text cannot be empty".to_string(),
});
}
if !(0.0..=1.0).contains(&example.confidence) {
return Err(IntentError::InvalidParameter {
parameter: "confidence".to_string(),
message: format!("Confidence must be between 0.0 and 1.0, got {}", example.confidence),
});
}
{
let mut training_data = self.training_data.write().await;
training_data.push(example.clone());
}
self.update_intent_patterns(&example.intent, &example.text).await?;
self.update_vocabulary(&example.text).await;
if self.config.debug_mode {
info!("Added training example: '{}' -> {}", example.text, example.intent);
}
Ok(())
}
pub async fn add_feedback(&self, feedback: IntentFeedback) -> Result<()> {
if self.config.debug_mode {
info!("Adding feedback: '{}' -> {} (predicted: {}, satisfaction: {})",
feedback.text, feedback.actual_intent, feedback.predicted_intent, feedback.satisfaction_score);
}
let confidence = feedback.satisfaction_score / 5.0; let example = TrainingExample {
text: feedback.text,
intent: feedback.actual_intent,
confidence,
source: TrainingSource::UserFeedback,
};
self.add_training_example(example).await?;
if self.should_retrain().await {
self.retrain().await?;
}
Ok(())
}
pub async fn get_stats(&self) -> ClassifierStats {
let training_data = self.training_data.read().await;
ClassifierStats {
training_examples: training_data.len(),
vocabulary_size: self.vocabulary.len(),
intent_count: self.intent_patterns.len(),
feedback_examples: training_data
.iter()
.filter(|e| matches!(e.source, TrainingSource::UserFeedback))
.count(),
last_updated: Some(chrono::Utc::now()),
}
}
pub async fn export_training_data(&self) -> Result<String> {
let training_data = self.training_data.read().await;
serde_json::to_string_pretty(&*training_data)
.map_err(IntentError::SerializationError)
}
pub async fn import_training_data(&self, json_data: &str) -> Result<()> {
let examples: Vec<TrainingExample> = serde_json::from_str(json_data)
.map_err(IntentError::SerializationError)?;
for example in examples {
self.add_training_example(example).await?;
}
Ok(())
}
pub async fn clear_training_data(&self) -> Result<()> {
{
let mut training_data = self.training_data.write().await;
training_data.clear();
}
self.vocabulary.clear();
self.intent_patterns.clear();
self.load_bootstrap_data().await?;
if self.config.debug_mode {
info!("Cleared all training data and reloaded bootstrap data");
}
Ok(())
}
async fn find_exact_match(&self, text: &str) -> Result<Option<IntentPrediction>> {
let training_data = self.training_data.read().await;
for example in training_data.iter() {
if example.text == text {
let confidence = Confidence::new(example.confidence)
.unwrap_or_else(|_| Confidence::default());
return Ok(Some(IntentPrediction {
intent: example.intent.clone(),
confidence,
alternative_intents: vec![],
reasoning: "Exact match found in training data".to_string(),
}));
}
}
Ok(None)
}
async fn extract_features(&self, text: &str) -> Result<FeatureVector> {
let cleaned_text = self.preprocess_text(text);
let text_features = self.extract_text_features(&cleaned_text).await?;
let context_features = self.extract_context_features(&cleaned_text);
let mut metadata = HashMap::new();
metadata.insert("text_length".to_string(), cleaned_text.len() as f64);
metadata.insert("word_count".to_string(), cleaned_text.split_whitespace().count() as f64);
Ok(FeatureVector {
text_features,
context_features,
metadata,
})
}
async fn calculate_intent_scores(&self, features: &FeatureVector) -> Result<HashMap<IntentId, f64>> {
let mut scores = HashMap::new();
for entry in self.intent_patterns.iter() {
let (intent, pattern_texts) = (entry.key(), entry.value());
let mut intent_score: f64 = 0.0;
for pattern_text in pattern_texts {
let pattern_features = self.extract_text_features(pattern_text).await?;
let similarity = self.cosine_similarity(&features.text_features, &pattern_features);
intent_score = intent_score.max(similarity);
}
intent_score += self.calculate_context_boost(intent, features);
scores.insert(intent.clone(), intent_score.min(1.0));
}
if scores.values().all(|&score| score < self.config.min_confidence_threshold) {
let fallback_scores = self.rule_based_classification(features).await;
for (intent, score) in fallback_scores {
scores.entry(intent).or_insert(score);
}
}
Ok(scores)
}
fn find_best_intent(&self, scores: &HashMap<IntentId, f64>) -> Result<(IntentId, Confidence)> {
let (best_intent, best_score) = scores
.iter()
.max_by(|a, b| a.1.partial_cmp(b.1).unwrap_or(std::cmp::Ordering::Equal))
.ok_or_else(|| IntentError::ClassificationFailed("No intents found".to_string()))?;
let confidence = Confidence::new(*best_score)
.unwrap_or_else(|_| Confidence::default());
Ok((best_intent.clone(), confidence))
}
fn get_alternative_intents(&self, scores: &HashMap<IntentId, f64>, best_intent: &IntentId) -> Vec<(IntentId, Confidence)> {
let mut alternatives: Vec<(IntentId, Confidence)> = scores
.iter()
.filter(|(intent, _)| *intent != best_intent)
.filter_map(|(intent, score)| {
Confidence::new(*score)
.ok()
.map(|confidence| (intent.clone(), confidence))
})
.collect();
alternatives.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
alternatives.truncate(3);
alternatives
}
async fn generate_reasoning(&self, _text: &str, intent: &IntentId, features: &FeatureVector) -> String {
if let Some(intent_patterns) = self.intent_patterns.get(intent) {
if let Some(best_pattern) = intent_patterns.first() {
return format!(
"Classified as '{}' based on similarity to pattern: '{}' (using {} text features)",
intent, best_pattern, features.text_features.len()
);
}
}
format!("Classified as '{}' using rule-based analysis", intent)
}
async fn load_bootstrap_data(&self) -> Result<()> {
let bootstrap_examples = self.get_bootstrap_examples();
for (text, intent_str) in bootstrap_examples {
let example = TrainingExample {
text: text.to_string(),
intent: IntentId::from(intent_str),
confidence: 1.0,
source: TrainingSource::Bootstrap,
};
self.add_training_example(example).await?;
}
if self.config.debug_mode {
info!("Loaded {} bootstrap training examples", self.get_bootstrap_examples().len());
}
Ok(())
}
fn get_bootstrap_examples(&self) -> Vec<(&'static str, &'static str)> {
vec![
("merge these JSON files together", "data_merge"),
("combine multiple JSON documents", "data_merge"),
("join several data files into one", "data_merge"),
("consolidate JSON objects", "data_merge"),
("split this large JSON file", "data_split"),
("break apart this data into smaller pieces", "data_split"),
("divide this file into multiple parts", "data_split"),
("convert JSON to CSV format", "data_transform"),
("transform this data structure", "data_transform"),
("change the format of this file", "data_transform"),
("analyze this dataset for patterns", "data_analyze"),
("examine the data for insights", "data_analyze"),
("what trends do you see in this data", "data_analyze"),
("give me statistics about this data", "data_analyze"),
("read the contents of this file", "file_read"),
("load this document", "file_read"),
("open and parse this file", "file_read"),
("save this data to a file", "file_write"),
("write this content to disk", "file_write"),
("create a new file with this data", "file_write"),
("convert PDF to markdown", "file_convert"),
("change this file format", "file_convert"),
("export as different format", "file_convert"),
("compare these two files", "file_compare"),
("what's different between these documents", "file_compare"),
("find differences in these files", "file_compare"),
("make an API request to this URL", "network_request"),
("call this REST endpoint", "network_request"),
("send HTTP request", "network_request"),
("download this file from the internet", "network_download"),
("fetch data from this URL", "network_download"),
("retrieve file from web", "network_download"),
("check if this website is up", "network_monitor"),
("monitor API endpoint", "network_monitor"),
("test connectivity to server", "network_monitor"),
("extract text from this document", "extraction"),
("pull out specific information", "extraction"),
("get the important parts from this", "extraction"),
("validate this data against schema", "validation"),
("check if this data is correct", "validation"),
("verify the format of this file", "validation"),
("generate a report from this data", "generation"),
("create summary of this information", "generation"),
("produce documentation", "generation"),
("classify this content", "classification"),
("categorize this data", "classification"),
("determine the type of this file", "classification"),
("analyze this code for issues", "code_analyze"),
("review this source code", "code_analyze"),
("check code quality", "code_analyze"),
("process this text document", "text_process"),
("clean up this text", "text_process"),
("parse natural language", "text_process"),
]
}
fn preprocess_text(&self, text: &str) -> String {
text.to_lowercase()
.chars()
.filter(|c| c.is_alphanumeric() || c.is_whitespace())
.collect::<String>()
.split_whitespace()
.collect::<Vec<_>>()
.join(" ")
}
async fn extract_text_features(&self, text: &str) -> Result<Vec<f64>> {
let mut features = vec![0.0; self.config.feature_dimensions];
let words: Vec<&str> = text.split_whitespace().collect();
let word_count = words.len() as f64;
if word_count == 0.0 {
return Ok(features);
}
for word in words {
if let Some(index) = self.vocabulary.get(word) {
if *index < features.len() {
features[*index] += 1.0 / word_count;
}
}
}
Ok(features)
}
fn extract_context_features(&self, text: &str) -> Vec<f64> {
vec![
text.len() as f64 / 100.0, text.split_whitespace().count() as f64 / 20.0, if text.contains('?') { 1.0 } else { 0.0 }, if text.contains("file") { 1.0 } else { 0.0 }, if text.contains("data") { 1.0 } else { 0.0 }, ]
}
fn cosine_similarity(&self, a: &[f64], b: &[f64]) -> f64 {
if a.len() != b.len() {
return 0.0;
}
let dot_product: f64 = a.iter().zip(b.iter()).map(|(x, y)| x * y).sum();
let norm_a: f64 = a.iter().map(|x| x * x).sum::<f64>().sqrt();
let norm_b: f64 = b.iter().map(|x| x * x).sum::<f64>().sqrt();
if norm_a == 0.0 || norm_b == 0.0 {
0.0
} else {
dot_product / (norm_a * norm_b)
}
}
fn calculate_context_boost(&self, intent: &IntentId, features: &FeatureVector) -> f64 {
let mut boost = 0.0;
if intent.0.contains("file") && features.context_features.get(3).unwrap_or(&0.0) > &0.0 {
boost += 0.1;
}
if intent.0.contains("data") && features.context_features.get(4).unwrap_or(&0.0) > &0.0 {
boost += 0.1;
}
boost
}
async fn rule_based_classification(&self, _features: &FeatureVector) -> HashMap<IntentId, f64> {
let mut scores = HashMap::new();
scores.insert(IntentId::from("general_processing"), 0.5);
scores
}
async fn update_intent_patterns(&self, intent: &IntentId, text: &str) -> Result<()> {
self.intent_patterns
.entry(intent.clone())
.or_insert_with(Vec::new)
.push(text.to_string());
Ok(())
}
async fn update_vocabulary(&self, text: &str) {
for word in text.split_whitespace() {
let vocab_len = self.vocabulary.len();
if vocab_len < self.config.max_vocabulary_size && !self.vocabulary.contains_key(word) {
self.vocabulary.insert(word.to_string(), vocab_len);
}
}
}
async fn should_retrain(&self) -> bool {
let training_data = self.training_data.read().await;
let feedback_count = training_data
.iter()
.filter(|example| matches!(example.source, TrainingSource::UserFeedback))
.count();
feedback_count >= self.config.retraining_threshold
}
async fn retrain(&self) -> Result<()> {
if self.config.debug_mode {
info!("Retraining intent classification model");
}
let training_data = self.training_data.read().await;
self.vocabulary.clear();
for example in training_data.iter() {
self.update_vocabulary(&example.text).await;
}
if self.config.debug_mode {
info!("Model retraining completed. Vocabulary size: {}", self.vocabulary.len());
}
Ok(())
}
}
impl Default for IntentClassifier {
fn default() -> Self {
Self {
training_data: Arc::new(RwLock::new(Vec::new())),
vocabulary: Arc::new(DashMap::with_hasher(RandomState::new())),
intent_patterns: Arc::new(DashMap::with_hasher(RandomState::new())),
config: ClassifierConfig::default(),
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[tokio::test]
async fn test_intent_classification() {
let classifier = IntentClassifier::new().await.unwrap();
let prediction = classifier
.predict_intent("merge these JSON files together")
.await
.unwrap();
assert_eq!(prediction.intent.0, "data_merge");
assert!(prediction.confidence.value() > 0.5);
}
#[tokio::test]
async fn test_feedback_learning() {
let classifier = IntentClassifier::new().await.unwrap();
let feedback = IntentFeedback {
text: "combine data files".to_string(),
predicted_intent: IntentId::from("data_transform"),
actual_intent: IntentId::from("data_merge"),
satisfaction_score: 5.0,
notes: None,
timestamp: chrono::Utc::now(),
};
classifier.add_feedback(feedback).await.unwrap();
let stats = classifier.get_stats().await;
assert!(stats.feedback_examples > 0);
}
#[tokio::test]
async fn test_training_data_export_import() {
let classifier = IntentClassifier::new().await.unwrap();
let example = TrainingExample {
text: "test example".to_string(),
intent: IntentId::from("test_intent"),
confidence: 0.9,
source: TrainingSource::Programmatic,
};
classifier.add_training_example(example).await.unwrap();
let exported = classifier.export_training_data().await.unwrap();
let new_classifier = IntentClassifier::new().await.unwrap();
new_classifier.import_training_data(&exported).await.unwrap();
let stats = new_classifier.get_stats().await;
assert!(stats.training_examples > 0);
}
}
#[cfg(test)]
impl IntentClassifier {
pub fn test_cosine_similarity(&self, a: &[f64], b: &[f64]) -> f64 {
self.cosine_similarity(a, b)
}
pub fn test_extract_context_features(&self, text: &str) -> Vec<f64> {
self.extract_context_features(text)
}
pub fn test_preprocess_text(&self, text: &str) -> String {
self.preprocess_text(text)
}
}