use crate::llm::{CallOptions, ChatModel, LlmError};
use async_trait::async_trait;
use juncture_core::state::messages::Message;
use thiserror::Error;
#[derive(Debug, Error)]
pub enum MemoryError {
#[error("summarization failed: {0}")]
SummarizationFailed(String),
#[error("store error: {0}")]
StoreError(String),
#[error("fact extraction failed: {0}")]
ExtractionFailed(String),
#[error("conversation too short to summarize")]
InsufficientContent,
}
#[async_trait]
pub trait Summarizer: Send + Sync + 'static {
async fn summarize(&self, text: &str) -> Result<String, MemoryError>;
}
#[derive(Clone, Debug)]
pub struct LlmSummarizer<M: ChatModel> {
model: M,
max_output_tokens: u32,
}
impl<M: ChatModel> LlmSummarizer<M> {
#[must_use]
pub const fn new(model: M, max_output_tokens: u32) -> Self {
Self {
model,
max_output_tokens,
}
}
}
#[async_trait]
impl<M: ChatModel> Summarizer for LlmSummarizer<M> {
async fn summarize(&self, text: &str) -> Result<String, MemoryError> {
let trimmed = text.trim();
if trimmed.len() < 50 {
return Err(MemoryError::InsufficientContent);
}
let system_msg = Message::system(
"Summarize the following text concisely, preserving key facts and information. Output only the summary.",
);
let user_msg = Message::human(trimmed);
let options = CallOptions {
max_tokens: Some(self.max_output_tokens),
..Default::default()
};
let response = juncture_core::wasm_send::force_send(
self.model.invoke(&[system_msg, user_msg], Some(&options)),
)
.await
.map_err(|e| match e {
LlmError::InvalidResponse(msg) => {
MemoryError::SummarizationFailed(format!("invalid response: {msg}"))
}
#[cfg(any(feature = "anthropic", feature = "openai", feature = "ollama"))]
LlmError::NetworkError(e) => {
MemoryError::SummarizationFailed(format!("network error: {e}"))
}
_ => MemoryError::SummarizationFailed(format!("LLM error: {e}")),
})?;
let summary = response.content_text();
if summary.trim().is_empty() {
return Err(MemoryError::SummarizationFailed(
"LLM returned empty summary".to_string(),
));
}
Ok(summary.to_string())
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_memory_error_display() {
let err = MemoryError::SummarizationFailed("test error".to_string());
assert_eq!(err.to_string(), "summarization failed: test error");
let err = MemoryError::StoreError("db error".to_string());
assert_eq!(err.to_string(), "store error: db error");
let err = MemoryError::ExtractionFailed("parse error".to_string());
assert_eq!(err.to_string(), "fact extraction failed: parse error");
let err = MemoryError::InsufficientContent;
assert_eq!(err.to_string(), "conversation too short to summarize");
}
#[test]
fn test_llm_summarizer_construction() {
let model = crate::llm::MockChatModel::new("test-model");
let summarizer = LlmSummarizer::new(model, 500);
assert_eq!(summarizer.max_output_tokens, 500);
}
}