Skip to main content

stynx_code_compact/
grouping.rs

1use stynx_code_types::{ContentBlock, Message, Role};
2
3#[derive(Debug, Clone, PartialEq)]
4pub enum GroupType {
5
6    UserAssistant,
7
8    ToolExchange,
9
10    System,
11}
12
13#[derive(Debug, Clone)]
14pub struct MessageGroup {
15    pub messages: Vec<Message>,
16    pub total_tokens: u64,
17    pub group_type: GroupType,
18}
19
20pub fn group_messages(messages: &[Message]) -> Vec<MessageGroup> {
21    let mut groups = Vec::new();
22    let mut i = 0;
23
24    while i < messages.len() {
25        let msg = &messages[i];
26
27        if msg.role == Role::Assistant && has_tool_use(&msg.content) {
28            if i + 1 < messages.len()
29                && messages[i + 1].role == Role::User
30                && has_tool_result(&messages[i + 1].content)
31            {
32                let token_estimate = estimate_tokens(&msg.content) + estimate_tokens(&messages[i + 1].content);
33                groups.push(MessageGroup {
34                    messages: vec![msg.clone(), messages[i + 1].clone()],
35                    total_tokens: token_estimate,
36                    group_type: GroupType::ToolExchange,
37                });
38                i += 2;
39                continue;
40            }
41        }
42
43        if msg.role == Role::User && !has_tool_result(&msg.content) {
44            if i + 1 < messages.len() && messages[i + 1].role == Role::Assistant {
45                let token_estimate = estimate_tokens(&msg.content) + estimate_tokens(&messages[i + 1].content);
46                groups.push(MessageGroup {
47                    messages: vec![msg.clone(), messages[i + 1].clone()],
48                    total_tokens: token_estimate,
49                    group_type: GroupType::UserAssistant,
50                });
51                i += 2;
52                continue;
53            }
54        }
55
56        let token_estimate = estimate_tokens(&msg.content);
57        groups.push(MessageGroup {
58            messages: vec![msg.clone()],
59            total_tokens: token_estimate,
60            group_type: GroupType::System,
61        });
62        i += 1;
63    }
64
65    groups
66}
67
68fn has_tool_use(blocks: &[ContentBlock]) -> bool {
69    blocks.iter().any(|b| matches!(b, ContentBlock::ToolUse { .. }))
70}
71
72fn has_tool_result(blocks: &[ContentBlock]) -> bool {
73    blocks.iter().any(|b| matches!(b, ContentBlock::ToolResult { .. }))
74}
75
76fn estimate_tokens(blocks: &[ContentBlock]) -> u64 {
77    let chars: usize = blocks
78        .iter()
79        .map(|b| match b {
80            ContentBlock::Text { text } => text.len(),
81            ContentBlock::ToolUse { name, input, .. } => name.len() + input.to_string().len(),
82            ContentBlock::ToolResult { content, .. } => content.len(),
83            ContentBlock::Thinking { thinking } => thinking.len(),
84            ContentBlock::Image { data, .. } => data.len(),
85        })
86        .sum();
87    (chars as u64) / 4
88}