use async_trait::async_trait;
use std::sync::Arc;
use tracing::info;
use crate::core::SessionManager;
use crate::llm::Provider;
use crate::storage::{Storage, Message, MessageRole, MessageContent};
use super::{
SummarizationService, SummarizationConfig, SummarizationError,
ConversationMetrics, ConversationSummary, TokenMonitor, ConversationSummarizer, utils,
};
pub struct SummarizationServiceImpl {
monitor: TokenMonitor,
summarizer: ConversationSummarizer,
storage: Arc<dyn Storage>,
session_manager: Arc<SessionManager>,
config: SummarizationConfig,
}
impl SummarizationServiceImpl {
pub fn new(
provider: Box<dyn Provider>,
storage: Arc<dyn Storage>,
session_manager: Arc<SessionManager>,
config: SummarizationConfig,
) -> Self {
let monitor = TokenMonitor::new(config.clone());
let summarizer = ConversationSummarizer::new(provider, config.clone());
Self {
monitor,
summarizer,
storage,
session_manager,
config,
}
}
pub fn config(&self) -> &SummarizationConfig {
&self.config
}
pub fn monitor(&self) -> &TokenMonitor {
&self.monitor
}
pub fn summarizer(&self) -> &ConversationSummarizer {
&self.summarizer
}
fn calculate_total_tokens(&self, messages: &[Message]) -> u32 {
self.summarizer.estimate_message_tokens(messages)
}
}
#[async_trait]
impl SummarizationService for SummarizationServiceImpl {
async fn monitor_session(&self, session_id: &str) -> Result<ConversationMetrics, SummarizationError> {
if !self.config.enabled {
return Err(SummarizationError::Disabled);
}
let session = self.session_manager.get_session(session_id).await
.map_err(|e| SummarizationError::SessionManagerError(e.to_string()))?
.ok_or_else(|| SummarizationError::SessionNotFound(session_id.to_string()))?;
let messages = self.storage.get_messages(session_id).await
.map_err(|e| SummarizationError::StorageError(e.to_string()))?;
let message_count = messages.len() as u32;
let model = &session.model;
self.monitor.get_metrics(session_id, model, message_count).await
}
async fn needs_summarization(&self, session_id: &str) -> Result<bool, SummarizationError> {
if !self.config.enabled {
return Ok(false);
}
let metrics = self.monitor_session(session_id).await?;
Ok(metrics.needs_summarization)
}
async fn summarize_conversation(&self, session_id: &str) -> Result<ConversationSummary, SummarizationError> {
if !self.config.enabled {
return Err(SummarizationError::Disabled);
}
info!("Starting conversation summarization for session: {}", session_id);
let messages = self.storage.get_messages(session_id).await
.map_err(|e| SummarizationError::StorageError(e.to_string()))?;
if messages.is_empty() {
return Err(SummarizationError::EmptyConversation);
}
self.summarizer.validate_messages(&messages)?;
let summary = self.summarizer.summarize(session_id, messages).await?;
info!("Successfully summarized session {} ({} tokens)", session_id, summary.summary_tokens);
Ok(summary)
}
async fn create_continuation_session(&self, summary: ConversationSummary) -> Result<String, SummarizationError> {
info!("Creating continuation session for {}", summary.original_session_id);
let new_session_id = self.session_manager.create_session(
Some(format!("Continued from session {}", summary.original_session_id))
).await
.map_err(|e| SummarizationError::SessionManagerError(e.to_string()))?;
let summary_content = utils::format_summary_for_session(&summary);
let summary_message = Message {
id: uuid::Uuid::new_v4().to_string(),
session_id: new_session_id.clone(),
role: MessageRole::Assistant,
content: MessageContent::Text(summary_content),
timestamp: chrono::Utc::now(),
metadata: serde_json::json!({
"type": "summary",
"original_session_id": summary.original_session_id,
"summary_tokens": summary.summary_tokens,
"summarizer_model": summary.summarizer_model
}),
};
self.storage.save_message(&summary_message).await
.map_err(|e| SummarizationError::StorageError(e.to_string()))?;
for mut message in summary.preserved_messages {
message.session_id = new_session_id.clone();
message.id = uuid::Uuid::new_v4().to_string();
self.storage.save_message(&message).await
.map_err(|e| SummarizationError::StorageError(e.to_string()))?;
}
info!("Created continuation session: {}", new_session_id);
Ok(new_session_id)
}
async fn auto_summarize_if_needed(&self, session_id: &str) -> Result<Option<String>, SummarizationError> {
if !self.config.enabled {
return Ok(None);
}
if !self.needs_summarization(session_id).await? {
return Ok(None);
}
info!("Auto-summarization triggered for session: {}", session_id);
let summary = self.summarize_conversation(session_id).await?;
let new_session_id = self.create_continuation_session(summary).await?;
info!("Auto-summarization completed. New session: {}", new_session_id);
Ok(Some(new_session_id))
}
async fn update_token_usage(&self, session_id: &str, usage: crate::core::TokenUsage) -> Result<(), SummarizationError> {
self.monitor.update_usage(session_id, usage).await
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::llm::providers::MockProvider;
use crate::llm::{Model, ModelCapabilities};
use crate::storage::MemoryStorage;
use crate::core::TokenUsage;
fn create_test_model() -> Model {
Model {
id: "test-model".to_string(),
name: "Test Model".to_string(),
provider: "test".to_string(),
context_length: 4096,
max_output_tokens: 2048,
supports_tools: false,
supports_streaming: false,
supports_vision: false,
cost_per_input_token: 0.0,
cost_per_output_token: 0.0,
capabilities: ModelCapabilities::default(),
}
}
async fn create_test_service() -> SummarizationServiceImpl {
let model = create_test_model();
let provider = Box::new(MockProvider::new());
let storage = Arc::new(MemoryStorage::new());
let session_manager = Arc::new(SessionManager::new(storage.clone()));
let config = SummarizationConfig::default();
SummarizationServiceImpl::new(provider, storage, session_manager, config)
}
#[tokio::test]
async fn test_service_creation() {
let service = create_test_service().await;
assert!(service.config().enabled);
assert_eq!(service.config().summarizer_model, "gpt-4");
}
#[tokio::test]
async fn test_needs_summarization_disabled() {
let mut service = create_test_service().await;
service.config.enabled = false;
let result = service.needs_summarization("test-session").await.unwrap();
assert!(!result);
}
#[tokio::test]
async fn test_summarize_conversation_disabled() {
let mut service = create_test_service().await;
service.config.enabled = false;
let result = service.summarize_conversation("test-session").await;
assert!(matches!(result, Err(SummarizationError::Disabled)));
}
#[tokio::test]
async fn test_auto_summarize_disabled() {
let mut service = create_test_service().await;
service.config.enabled = false;
let result = service.auto_summarize_if_needed("test-session").await.unwrap();
assert!(result.is_none());
}
#[tokio::test]
async fn test_update_token_usage() {
let service = create_test_service().await;
let usage = TokenUsage {
input_tokens: 1000,
output_tokens: 500,
total_tokens: 1500,
cache_creation_tokens: 0,
cache_read_tokens: 0,
};
let result = service.update_token_usage("test-session", usage).await;
assert!(result.is_ok());
let recorded_usage = service.monitor.get_usage("test-session").await;
assert_eq!(recorded_usage.total_tokens, 1500);
}
}