use std::collections::HashMap;
use std::sync::Arc;
use tokio::sync::RwLock;
use crate::llm::LlmClient;
use super::{SummaryGenerator, SummaryStrategyConfig};
pub struct LazyStrategy {
generator: Arc<RwLock<Box<dyn SummaryGenerator>>>,
cache: Arc<RwLock<HashMap<String, String>>>,
persist: bool,
config: SummaryStrategyConfig,
}
impl LazyStrategy {
pub fn new(client: LlmClient) -> Self {
Self {
generator: Arc::new(RwLock::new(Box::new(super::LlmSummaryGenerator::new(
client,
)))),
cache: Arc::new(RwLock::new(HashMap::new())),
persist: false,
config: SummaryStrategyConfig::default(),
}
}
pub fn with_persist(client: LlmClient, persist: bool) -> Self {
Self {
generator: Arc::new(RwLock::new(Box::new(super::LlmSummaryGenerator::new(
client,
)))),
cache: Arc::new(RwLock::new(HashMap::new())),
persist,
config: SummaryStrategyConfig::default(),
}
}
pub fn with_generator(generator: Box<dyn SummaryGenerator>) -> Self {
Self {
generator: Arc::new(RwLock::new(generator)),
cache: Arc::new(RwLock::new(HashMap::new())),
persist: false,
config: SummaryStrategyConfig::default(),
}
}
pub fn with_persist_mode(mut self, persist: bool) -> Self {
self.persist = persist;
self
}
pub fn with_config(mut self, config: SummaryStrategyConfig) -> Self {
self.config = config;
self
}
pub async fn has_cached(&self, node_id: &str) -> bool {
let cache = self.cache.read().await;
cache.contains_key(node_id)
}
pub async fn get_cached(&self, node_id: &str) -> Option<String> {
let cache = self.cache.read().await;
cache.get(node_id).cloned()
}
pub async fn get_or_generate(
&self,
node_id: &str,
title: &str,
content: &str,
) -> crate::llm::LlmResult<String> {
if self.persist {
if let Some(cached) = self.get_cached(node_id).await {
return Ok(cached);
}
}
let generator = self.generator.read().await;
let summary = generator.generate(title, content).await?;
if self.persist {
let mut cache = self.cache.write().await;
cache.insert(node_id.to_string(), summary.clone());
}
Ok(summary)
}
pub async fn populate_cache(&self, summaries: HashMap<String, String>) {
let mut cache = self.cache.write().await;
cache.extend(summaries);
}
pub async fn clear_cache(&self) {
let mut cache = self.cache.write().await;
cache.clear();
}
pub async fn cache_size(&self) -> usize {
let cache = self.cache.read().await;
cache.len()
}
pub fn is_persist_enabled(&self) -> bool {
self.persist
}
pub fn config(&self) -> &SummaryStrategyConfig {
&self.config
}
}
impl std::fmt::Debug for LazyStrategy {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("LazyStrategy")
.field("persist", &self.persist)
.field("config", &self.config)
.finish()
}
}