use crate::document::{DocumentTree, NodeId};
use crate::llm::LlmClient;
use super::{SummaryGenerator, SummaryStrategyConfig};
pub struct SelectiveStrategy {
generator: Box<dyn SummaryGenerator>,
min_tokens: usize,
branch_only: bool,
config: SummaryStrategyConfig,
}
impl SelectiveStrategy {
pub fn new(client: LlmClient) -> Self {
Self {
generator: Box::new(super::LlmSummaryGenerator::new(client)),
min_tokens: 100,
branch_only: true,
config: SummaryStrategyConfig::default(),
}
}
pub fn with_thresholds(client: LlmClient, min_tokens: usize, branch_only: bool) -> Self {
Self {
generator: Box::new(super::LlmSummaryGenerator::new(client)),
min_tokens,
branch_only,
config: SummaryStrategyConfig::default(),
}
}
pub fn with_generator(generator: Box<dyn SummaryGenerator>) -> Self {
Self {
generator,
min_tokens: 100,
branch_only: true,
config: SummaryStrategyConfig::default(),
}
}
pub fn with_min_tokens(mut self, min_tokens: usize) -> Self {
self.min_tokens = min_tokens;
self
}
pub fn with_branch_only(mut self, branch_only: bool) -> Self {
self.branch_only = branch_only;
self
}
pub fn with_config(mut self, config: SummaryStrategyConfig) -> Self {
self.config = config;
self
}
pub fn should_generate(
&self,
tree: &DocumentTree,
node_id: NodeId,
token_count: usize,
) -> bool {
let enough_tokens = token_count >= self.min_tokens;
if self.branch_only {
let is_branch = !tree.is_leaf(node_id);
is_branch && enough_tokens
} else {
enough_tokens
}
}
pub async fn generate(&self, title: &str, content: &str) -> crate::llm::LlmResult<String> {
self.generator.generate(title, content).await
}
pub fn min_tokens(&self) -> usize {
self.min_tokens
}
pub fn is_branch_only(&self) -> bool {
self.branch_only
}
pub fn config(&self) -> &SummaryStrategyConfig {
&self.config
}
}
impl std::fmt::Debug for SelectiveStrategy {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("SelectiveStrategy")
.field("min_tokens", &self.min_tokens)
.field("branch_only", &self.branch_only)
.field("config", &self.config)
.finish()
}
}