Skip to main content

vectorless/index/summary/
strategy.rs

1// Copyright (c) 2026 vectorless developers
2// SPDX-License-Identifier: Apache-2.0
3
4//! Summary generation strategies.
5
6use async_trait::async_trait;
7
8use crate::domain::{DocumentTree, NodeId};
9use crate::llm::{LlmClient, LlmResult};
10
11/// Configuration for summary strategies.
12#[derive(Debug, Clone)]
13pub struct SummaryStrategyConfig {
14    /// Maximum tokens for a summary.
15    pub max_tokens: usize,
16
17    /// Minimum content tokens to generate summary.
18    pub min_content_tokens: usize,
19
20    /// Whether to persist lazy-generated summaries.
21    pub persist_lazy: bool,
22}
23
24impl Default for SummaryStrategyConfig {
25    fn default() -> Self {
26        Self {
27            max_tokens: 200,
28            min_content_tokens: 50,
29            persist_lazy: false,
30        }
31    }
32}
33
34/// Strategy for generating summaries.
35#[derive(Debug, Clone)]
36pub enum SummaryStrategy {
37    /// No summary generation.
38    None,
39
40    /// Generate for all nodes.
41    Full {
42        /// Strategy configuration.
43        config: SummaryStrategyConfig,
44    },
45
46    /// Generate selectively.
47    Selective {
48        /// Minimum tokens threshold.
49        min_tokens: usize,
50
51        /// Only generate for branch nodes (non-leaves).
52        branch_only: bool,
53
54        /// Strategy configuration.
55        config: SummaryStrategyConfig,
56    },
57
58    /// Generate on-demand at query time.
59    Lazy {
60        /// Whether to persist generated summaries.
61        persist: bool,
62
63        /// Strategy configuration.
64        config: SummaryStrategyConfig,
65    },
66}
67
68impl Default for SummaryStrategy {
69    fn default() -> Self {
70        Self::Selective {
71            min_tokens: 100,
72            branch_only: true,
73            config: SummaryStrategyConfig::default(),
74        }
75    }
76}
77
78impl SummaryStrategy {
79    /// Create a "none" strategy.
80    pub fn none() -> Self {
81        Self::None
82    }
83
84    /// Create a "full" strategy.
85    pub fn full() -> Self {
86        Self::Full {
87            config: SummaryStrategyConfig::default(),
88        }
89    }
90
91    /// Create a "selective" strategy.
92    pub fn selective(min_tokens: usize, branch_only: bool) -> Self {
93        Self::Selective {
94            min_tokens,
95            branch_only,
96            config: SummaryStrategyConfig::default(),
97        }
98    }
99
100    /// Create a "lazy" strategy.
101    pub fn lazy(persist: bool) -> Self {
102        Self::Lazy {
103            persist,
104            config: SummaryStrategyConfig::default(),
105        }
106    }
107
108    /// Check if we should generate a summary for a node.
109    pub fn should_generate(
110        &self,
111        tree: &DocumentTree,
112        node_id: NodeId,
113        token_count: usize,
114    ) -> bool {
115        match self {
116            Self::None => false,
117            Self::Full { .. } => token_count > 0,
118            Self::Selective {
119                min_tokens,
120                branch_only,
121                ..
122            } => {
123                let is_branch = !tree.is_leaf(node_id);
124                let enough_tokens = token_count >= *min_tokens;
125
126                if *branch_only {
127                    is_branch && enough_tokens
128                } else {
129                    enough_tokens
130                }
131            }
132            Self::Lazy { .. } => false, // Generated on-demand
133        }
134    }
135
136    /// Check if lazy strategy is enabled.
137    pub fn is_lazy(&self) -> bool {
138        matches!(self, Self::Lazy { .. })
139    }
140
141    /// Get the config.
142    pub fn config(&self) -> SummaryStrategyConfig {
143        match self {
144            Self::None => SummaryStrategyConfig::default(),
145            Self::Full { config } => config.clone(),
146            Self::Selective { config, .. } => config.clone(),
147            Self::Lazy { config, .. } => config.clone(),
148        }
149    }
150}
151
152/// Summary generator trait.
153#[async_trait]
154pub trait SummaryGenerator: Send + Sync {
155    /// Generate a summary for the given content.
156    async fn generate(&self, title: &str, content: &str) -> LlmResult<String>;
157}
158
159/// LLM-based summary generator.
160pub struct LlmSummaryGenerator {
161    client: LlmClient,
162    max_tokens: usize,
163}
164
165impl LlmSummaryGenerator {
166    /// Create a new summary generator.
167    pub fn new(client: LlmClient) -> Self {
168        Self {
169            client,
170            max_tokens: 200,
171        }
172    }
173
174    /// Set max tokens.
175    pub fn with_max_tokens(mut self, max_tokens: usize) -> Self {
176        self.max_tokens = max_tokens;
177        self
178    }
179}
180
181#[async_trait]
182impl SummaryGenerator for LlmSummaryGenerator {
183    async fn generate(&self, title: &str, content: &str) -> LlmResult<String> {
184        let system_prompt = "You are a document summarization assistant. \
185            Generate a concise summary (2-3 sentences) of the given section. \
186            Focus on the main topics and key information. \
187            Respond with only the summary, no additional text.";
188
189        let user_prompt = format!("Title: {}\n\nContent:\n{}", title, content);
190
191        self.client
192            .complete_with_max_tokens(&system_prompt, &user_prompt, self.max_tokens as u16)
193            .await
194    }
195}