vectorless/index/summary/
strategy.rs1use async_trait::async_trait;
7
8use crate::domain::{DocumentTree, NodeId};
9use crate::llm::{LlmClient, LlmResult};
10
11#[derive(Debug, Clone)]
13pub struct SummaryStrategyConfig {
14 pub max_tokens: usize,
16
17 pub min_content_tokens: usize,
19
20 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#[derive(Debug, Clone)]
36pub enum SummaryStrategy {
37 None,
39
40 Full {
42 config: SummaryStrategyConfig,
44 },
45
46 Selective {
48 min_tokens: usize,
50
51 branch_only: bool,
53
54 config: SummaryStrategyConfig,
56 },
57
58 Lazy {
60 persist: bool,
62
63 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 pub fn none() -> Self {
81 Self::None
82 }
83
84 pub fn full() -> Self {
86 Self::Full {
87 config: SummaryStrategyConfig::default(),
88 }
89 }
90
91 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 pub fn lazy(persist: bool) -> Self {
102 Self::Lazy {
103 persist,
104 config: SummaryStrategyConfig::default(),
105 }
106 }
107
108 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, }
134 }
135
136 pub fn is_lazy(&self) -> bool {
138 matches!(self, Self::Lazy { .. })
139 }
140
141 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#[async_trait]
154pub trait SummaryGenerator: Send + Sync {
155 async fn generate(&self, title: &str, content: &str) -> LlmResult<String>;
157}
158
159pub struct LlmSummaryGenerator {
161 client: LlmClient,
162 max_tokens: usize,
163}
164
165impl LlmSummaryGenerator {
166 pub fn new(client: LlmClient) -> Self {
168 Self {
169 client,
170 max_tokens: 200,
171 }
172 }
173
174 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}