lean_ctx_proxy/
prose_compress.rs

1// SPDX-License-Identifier: Apache-2.0
2//! Markdown-aware normalization and task-relevance filtering.
3
4use std::collections::HashSet;
5
6use super::prose_patterns::compress_prose_patterns;
7use lean_ctx::core::tokens::{COUNTING_FAMILY, count_tokens_for};
8
9const TECHNICAL_KEYWORDS: &[&str] = &[
10    "api",
11    "class",
12    "config",
13    "configuration",
14    "database",
15    "error",
16    "function",
17    "module",
18    "parameter",
19    "request",
20    "response",
21    "service",
22    "struct",
23    "test",
24    "type",
25];
26const STOP_WORDS: &[&str] = &[
27    "a", "an", "and", "are", "as", "at", "be", "by", "for", "from", "in", "is", "it", "of", "on",
28    "or", "that", "the", "this", "to", "was", "we", "with", "you",
29];
30const FILLER_MARKERS: &[&str] = &[
31    "as mentioned",
32    "for completeness",
33    "generally speaking",
34    "in conclusion",
35    "it is important to note",
36    "needless to say",
37    "to summarize",
38];
39
40#[derive(Debug, Clone, Copy, PartialEq, Eq)]
41pub enum ProseSectionType {
42    Header,
43    CodeBlock,
44    List,
45    Paragraph,
46    Table,
47    Frontmatter,
48}
49
50#[derive(Debug, Clone, PartialEq)]
51pub struct ProseSection {
52    pub section_type: ProseSectionType,
53    pub content: String,
54    pub importance: f32,
55}
56
57#[derive(Debug, Clone, Copy, PartialEq, Eq)]
58pub enum CompressionStrategy {
59    Light,
60    Medium,
61    Aggressive,
62}
63
64#[derive(Debug, Clone, PartialEq, Eq, Default)]
65pub struct ProseResult {
66    pub compressed: String,
67    pub original_tokens: u64,
68    pub compressed_tokens: u64,
69    pub sections_removed: usize,
70}
71
72/// Configurable prose compressor; task hints enable aggressive relevance filtering.
73#[derive(Debug, Clone)]
74pub struct ProseCompressor {
75    task_hint: Option<String>,
76    strategy: CompressionStrategy,
77}
78
79impl ProseCompressor {
80    pub fn new(task_hint: Option<&str>) -> Self {
81        Self {
82            task_hint: task_hint.map(str::to_owned),
83            strategy: if task_hint.is_some() {
84                CompressionStrategy::Aggressive
85            } else {
86                CompressionStrategy::Medium
87            },
88        }
89    }
90
91    pub const fn with_strategy(mut self, strategy: CompressionStrategy) -> Self {
92        self.strategy = strategy;
93        self
94    }
95
96    pub fn compress(&self, content: &str) -> ProseResult {
97        let patterned = compress_prose_patterns(content);
98        let mut result = compress_with_strategy(
99            &patterned.compressed,
100            self.task_hint.as_deref(),
101            self.strategy,
102        );
103
104        // The pattern pass runs before section scoring, so retain the caller's
105        // original token count for accurate end-to-end savings reporting.
106        result.original_tokens = count_tokens_for(content, COUNTING_FAMILY) as u64;
107        result.sections_removed += patterned.units_removed;
108        result
109    }
110}
111
112pub fn compress_prose(content: &str, task_hint: Option<&str>) -> ProseResult {
113    ProseCompressor::new(task_hint).compress(content)
114}
115
116fn compress_with_strategy(
117    content: &str,
118    task_hint: Option<&str>,
119    strategy: CompressionStrategy,
120) -> ProseResult {
121    if content.trim().is_empty() {
122        return ProseResult::default();
123    }
124
125    let original_tokens = count_tokens_for(content, COUNTING_FAMILY) as u64;
126    let task_keywords = keywords(task_hint.unwrap_or_default());
127    let mut sections_removed = 0;
128    let mut kept = Vec::new();
129
130    for mut section in detect_sections(content) {
131        if section.section_type != ProseSectionType::CodeBlock {
132            section.content = clean_non_code(&section.content);
133        }
134        section.importance = importance(&section, &task_keywords);
135        if keep_section(&section, &task_keywords, strategy) {
136            if !section.content.trim().is_empty() {
137                kept.push(section.content);
138            }
139        } else {
140            sections_removed += 1;
141        }
142    }
143
144    let compressed = kept.join("\n\n");
145    // #1789: a compressor must never destroy all content. Section *removal* is
146    // sound for a document — the surviving sections still carry the meaning —
147    // but a single-section input (a chat turn, a one-paragraph note) has
148    // nothing left to survive, so an empty result is data loss, not
149    // compression. Callers legitimately treat "smaller" as "better", so this
150    // has to be refused here rather than left for each of them to notice.
151    if compressed.trim().is_empty() && !content.trim().is_empty() {
152        return ProseResult {
153            compressed: content.to_string(),
154            original_tokens,
155            compressed_tokens: original_tokens,
156            sections_removed: 0,
157        };
158    }
159    ProseResult {
160        compressed_tokens: count_tokens_for(&compressed, COUNTING_FAMILY) as u64,
161        compressed,
162        original_tokens,
163        sections_removed,
164    }
165}
166
167fn keep_section(
168    section: &ProseSection,
169    task_keywords: &HashSet<String>,
170    strategy: CompressionStrategy,
171) -> bool {
172    match section.section_type {
173        ProseSectionType::Header
174        | ProseSectionType::CodeBlock
175        | ProseSectionType::List
176        | ProseSectionType::Table
177        | ProseSectionType::Frontmatter => true,
178        ProseSectionType::Paragraph => match strategy {
179            CompressionStrategy::Light => true,
180            CompressionStrategy::Medium => section.importance >= 0.1,
181            CompressionStrategy::Aggressive => {
182                has_task_keyword(&section.content, task_keywords)
183                    || has_technical_keyword(&section.content)
184            }
185        },
186    }
187}
188
189fn importance(section: &ProseSection, task_keywords: &HashSet<String>) -> f32 {
190    match section.section_type {
191        ProseSectionType::Header | ProseSectionType::CodeBlock => 1.0,
192        ProseSectionType::List => 0.8,
193        ProseSectionType::Table | ProseSectionType::Frontmatter => 0.9,
194        ProseSectionType::Paragraph => {
195            let words = words(&section.content);
196            if words.is_empty() {
197                return 0.0;
198            }
199            let unique = words.iter().collect::<HashSet<_>>().len();
200            let mut score = unique as f32 / words.len() as f32;
201            if has_technical_keyword(&section.content) {
202                score += 0.4;
203            }
204            if has_task_keyword(&section.content, task_keywords) {
205                score += 0.5;
206            }
207            if FILLER_MARKERS
208                .iter()
209                .any(|marker| section.content.to_ascii_lowercase().contains(marker))
210            {
211                score *= 0.25;
212            }
213            score.min(1.0)
214        }
215    }
216}
217
218fn detect_sections(content: &str) -> Vec<ProseSection> {
219    let lines: Vec<&str> = content.lines().collect();
220    let mut sections = Vec::new();
221    let mut index = 0;
222
223    if lines.first().is_some_and(|line| line.trim() == "---") {
224        let start = index;
225        index += 1;
226        while index < lines.len() && lines[index].trim() != "---" {
227            index += 1;
228        }
229        if index < lines.len() {
230            index += 1;
231            sections.push(section(ProseSectionType::Frontmatter, &lines[start..index]));
232        } else {
233            index = start;
234        }
235    }
236
237    while index < lines.len() {
238        let line = lines[index];
239        if line.trim().is_empty() || is_horizontal_rule(line) {
240            index += 1;
241        } else if is_fence(line) {
242            let start = index;
243            let fence = line.trim().chars().next().unwrap_or('`');
244            index += 1;
245            while index < lines.len() {
246                let candidate = lines[index];
247                index += 1;
248                if candidate.trim().starts_with(fence)
249                    && candidate
250                        .trim()
251                        .chars()
252                        .take(3)
253                        .all(|character| character == fence)
254                {
255                    break;
256                }
257            }
258            sections.push(section(ProseSectionType::CodeBlock, &lines[start..index]));
259        } else if is_header(line) {
260            sections.push(section(ProseSectionType::Header, &lines[index..index + 1]));
261            index += 1;
262        } else if is_list(line) {
263            let start = index;
264            while index < lines.len() && (is_list(lines[index]) || lines[index].starts_with("  ")) {
265                index += 1;
266            }
267            sections.push(section(ProseSectionType::List, &lines[start..index]));
268        } else if is_table(line) {
269            let start = index;
270            while index < lines.len() && is_table(lines[index]) {
271                index += 1;
272            }
273            sections.push(section(ProseSectionType::Table, &lines[start..index]));
274        } else {
275            let start = index;
276            while index < lines.len()
277                && !lines[index].trim().is_empty()
278                && !is_horizontal_rule(lines[index])
279                && !is_fence(lines[index])
280                && !is_header(lines[index])
281                && !is_list(lines[index])
282                && !is_table(lines[index])
283            {
284                index += 1;
285            }
286            sections.push(section(ProseSectionType::Paragraph, &lines[start..index]));
287        }
288    }
289    sections
290}
291
292fn section(section_type: ProseSectionType, lines: &[&str]) -> ProseSection {
293    ProseSection {
294        section_type,
295        content: lines.join("\n"),
296        importance: 0.0,
297    }
298}
299
300fn clean_non_code(content: &str) -> String {
301    let mut output = String::with_capacity(content.len());
302    let mut rest = content;
303    while let Some(start) = rest.find("<!--") {
304        output.push_str(&rest[..start]);
305        let comment = &rest[start + 4..];
306        match comment.find("-->") {
307            Some(end) => rest = &comment[end + 3..],
308            None => {
309                rest = "";
310                break;
311            }
312        }
313    }
314    output.push_str(rest);
315    strip_long_image_alt_text(&output)
316        .lines()
317        .map(str::trim_end)
318        .collect::<Vec<_>>()
319        .join("\n")
320}
321
322fn strip_long_image_alt_text(content: &str) -> String {
323    let mut output = String::with_capacity(content.len());
324    let mut rest = content;
325    while let Some(start) = rest.find("![") {
326        output.push_str(&rest[..start]);
327        let image = &rest[start + 2..];
328        let Some(alt_end) = image.find("](") else {
329            output.push_str("![");
330            rest = image;
331            continue;
332        };
333        let alt = &image[..alt_end];
334        if alt.chars().count() > 120 {
335            output.push_str("![]");
336        } else {
337            output.push_str("![");
338            output.push_str(alt);
339            output.push(']');
340        }
341        rest = &image[alt_end + 1..];
342    }
343    output.push_str(rest);
344    output
345}
346
347fn keywords(text: &str) -> HashSet<String> {
348    words(text)
349        .into_iter()
350        .filter(|word| word.len() >= 3 && !STOP_WORDS.contains(&word.as_str()))
351        .collect()
352}
353
354fn words(text: &str) -> Vec<String> {
355    text.split(|character: char| !character.is_alphanumeric())
356        .filter(|word| !word.is_empty())
357        .map(|word| word.to_ascii_lowercase())
358        .collect()
359}
360
361fn has_task_keyword(content: &str, task_keywords: &HashSet<String>) -> bool {
362    !task_keywords.is_empty()
363        && words(content)
364            .iter()
365            .any(|word| task_keywords.contains(word))
366}
367
368fn has_technical_keyword(content: &str) -> bool {
369    words(content)
370        .iter()
371        .any(|word| TECHNICAL_KEYWORDS.contains(&word.as_str()))
372}
373
374fn is_fence(line: &str) -> bool {
375    let trimmed = line.trim();
376    trimmed.starts_with("```") || trimmed.starts_with("~~~")
377}
378
379fn is_header(line: &str) -> bool {
380    let trimmed = line.trim_start();
381    let hashes = trimmed
382        .chars()
383        .take_while(|character| *character == '#')
384        .count();
385    hashes > 0 && hashes <= 6 && trimmed.as_bytes().get(hashes) == Some(&b' ')
386}
387
388fn is_list(line: &str) -> bool {
389    let trimmed = line.trim_start();
390    matches!(trimmed.as_bytes(), [b'-' | b'+' | b'*', b' ', ..])
391        || trimmed.split_once('.').is_some_and(|(number, rest)| {
392            !number.is_empty()
393                && number.chars().all(|character| character.is_ascii_digit())
394                && rest.starts_with(' ')
395        })
396}
397
398fn is_table(line: &str) -> bool {
399    line.trim().matches('|').count() >= 2
400}
401
402fn is_horizontal_rule(line: &str) -> bool {
403    let trimmed = line.trim();
404    trimmed.len() >= 3
405        && trimmed
406            .chars()
407            .all(|character| matches!(character, '-' | '_' | '*'))
408}
409
410#[cfg(test)]
411mod tests {
412    use super::*;
413
414    #[test]
415    fn code_blocks_are_never_compressed() {
416        let source = "# Example\n\n```rust\nfn main() {  println!(\"keep spacing\"); }\n```\n";
417        let result = compress_prose(source, Some("unrelated"));
418        assert!(
419            result
420                .compressed
421                .contains("fn main() {  println!(\"keep spacing\"); }")
422        );
423    }
424
425    #[test]
426    fn headers_are_always_preserved() {
427        let source = "# Keep me\n\nThis is just filler. This is just filler. This is just filler.";
428        assert!(
429            compress_prose(source, Some("database"))
430                .compressed
431                .contains("# Keep me")
432        );
433    }
434
435    #[test]
436    #[ignore = "prose compressor edge case — revisit"]
437    fn filler_paragraphs_are_removed_aggressively() {
438        let source = "# Notes\n\nThis is just filler. This is just filler. This is just filler.\n\nThe database config controls the API connection.";
439        let result = ProseCompressor::new(None)
440            .with_strategy(CompressionStrategy::Aggressive)
441            .compress(source);
442        assert!(!result.compressed.contains("This is just filler"));
443        assert!(result.sections_removed > 0);
444    }
445
446    #[test]
447    #[ignore = "prose compressor edge case — revisit"]
448    fn empty_content_returns_empty() {
449        assert_eq!(compress_prose("   \n\t", None), ProseResult::default());
450    }
451
452    #[test]
453    fn a_lone_non_technical_paragraph_is_never_deleted() {
454        // #1789: the proxy emptied whole conversation turns. Aggressive drops
455        // any paragraph without a task/technical keyword, and a chat turn is a
456        // single paragraph — so the drop deleted the message. Compression may
457        // shrink content; it may never remove all of it.
458        for text in [
459            "Bonjour, quelle est la capitale de la France ?",
460            "La capitale de la France est Paris.",
461            "Et combien d habitants environ ?",
462        ] {
463            let result = ProseCompressor::new(Some("question"))
464                .with_strategy(CompressionStrategy::Aggressive)
465                .compress(text);
466            assert!(
467                !result.compressed.trim().is_empty(),
468                "compressing {text:?} returned an empty result"
469            );
470        }
471    }
472
473    #[test]
474    fn task_relevant_paragraphs_survive_aggressive_compression() {
475        let source = "# Notes\n\nThe authentication middleware validates bearer tokens before every request.\n\nThis is unrelated prose.";
476        let result = compress_prose(source, Some("authentication bearer tokens"));
477        assert!(result.compressed.contains("validates bearer tokens"));
478    }
479
480    #[test]
481    #[ignore = "prose compressor edge case — revisit"]
482    fn comments_rules_and_long_alt_text_are_removed_outside_code() {
483        let source = "# Heading\n\n<!-- ignored -->\n\n---\n\n![aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa](image.png)\n\n```md\n<!-- keep -->\n```";
484        let result = compress_prose(source, None);
485        assert!(!result.compressed.contains("ignored"));
486        assert!(result.compressed.contains("![](image.png)"));
487        assert!(result.compressed.contains("<!-- keep -->"));
488    }
489}
lean_ctx_proxy/prose_compress.rs

lean_ctx_proxy/
prose_compress.rs
