Skip to content

Commit b2ab96c

Browse files
evanliu048evanliu048
committed
refactor
1 parent 7818bc7 commit b2ab96c

File tree

7 files changed

+270
-256
lines changed

7 files changed

+270
-256
lines changed

‎crates/chat-cli/src/cli/chat/cli/context/context_data_provider.rs‎

Lines changed: 64 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,64 @@
1+
//! Provides context window usage data and analysis
2+
3+
use crate::cli::chat::cli::model::context_window_tokens;
4+
use crate::cli::chat::token_counter::{
5+
CharCount,
6+
TokenCount,
7+
};
8+
use crate::cli::chat::{
9+
ChatError,
10+
ChatSession,
11+
};
12+
use crate::os::Os;
13+
14+
/// Detailed context window usage data
15+
#[derive(Debug)]
16+
pub struct ContextWindowData {
17+
pub total_tokens: TokenCount,
18+
pub context_tokens: TokenCount,
19+
pub assistant_tokens: TokenCount,
20+
pub user_tokens: TokenCount,
21+
pub tools_tokens: TokenCount,
22+
pub context_window_size: usize,
23+
pub dropped_context_files: Vec<(String, String)>,
24+
}
25+
26+
/// Get detailed usage data for context window analysis
27+
pub(super) async fn get_detailed_context_data(
28+
session: &mut ChatSession,
29+
os: &Os,
30+
) -> Result<ContextWindowData, ChatError> {
31+
let context_window_size = context_window_tokens(session.conversation.model_info.as_ref());
32+
33+
let state = session
34+
.conversation
35+
.backend_conversation_state(os, true, &mut std::io::stderr())
36+
.await?;
37+
38+
let data = state.calculate_conversation_size();
39+
let tool_specs_json: String = state
40+
.tools
41+
.values()
42+
.filter_map(|s| serde_json::to_string(s).ok())
43+
.collect::<Vec<String>>()
44+
.join("");
45+
let tools_char_count: CharCount = tool_specs_json.len().into();
46+
let total_tokens: TokenCount =
47+
(data.context_messages + data.user_messages + data.assistant_messages + tools_char_count).into();
48+
49+
Ok(ContextWindowData {
50+
total_tokens,
51+
context_tokens: data.context_messages.into(),
52+
assistant_tokens: data.assistant_messages.into(),
53+
user_tokens: data.user_messages.into(),
54+
tools_tokens: tools_char_count.into(),
55+
context_window_size,
56+
dropped_context_files: state.dropped_context_files,
57+
})
58+
}
59+
60+
/// Get total context usage percentage (external API)
61+
pub async fn get_total_context_usage_percentage(session: &mut ChatSession, os: &Os) -> Result<f32, ChatError> {
62+
let data = get_detailed_context_data(session, os).await?;
63+
Ok((data.total_tokens.value() as f32 / data.context_window_size as f32) * 100.0)
64+
}

‎crates/chat-cli/src/cli/chat/cli/context/context_renderer.rs‎

Lines changed: 197 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,197 @@
1+
//! Rendering for context window information
2+
3+
use crossterm::style::Attribute;
4+
use crossterm::{
5+
execute,
6+
queue,
7+
style,
8+
};
9+
10+
use super::context_data_provider::ContextWindowData;
11+
use crate::cli::chat::token_counter::TokenCount;
12+
use crate::cli::chat::{
13+
ChatError,
14+
ChatSession,
15+
};
16+
use crate::theme::StyledText;
17+
18+
/// Calculate context percentage from token counts (private utility)
19+
fn calculate_context_percentage(tokens: TokenCount, context_window_size: usize) -> f32 {
20+
(tokens.value() as f32 / context_window_size as f32) * 100.0
21+
}
22+
23+
/// Render context window information section
24+
pub async fn render_context_window(
25+
context_data: &ContextWindowData,
26+
session: &mut ChatSession,
27+
) -> Result<(), ChatError> {
28+
if !context_data.dropped_context_files.is_empty() {
29+
execute!(
30+
session.stderr,
31+
StyledText::warning_fg(),
32+
style::Print("\nSome context files are dropped due to size limit, please run "),
33+
StyledText::success_fg(),
34+
style::Print("/context show "),
35+
StyledText::warning_fg(),
36+
style::Print("to learn more.\n"),
37+
StyledText::reset(),
38+
)?;
39+
}
40+
41+
let window_width = session.terminal_width();
42+
// set a max width for the progress bar for better aesthetic
43+
let progress_bar_width = std::cmp::min(window_width, 80);
44+
45+
let context_width = ((context_data.context_tokens.value() as f64 / context_data.context_window_size as f64)
46+
* progress_bar_width as f64) as usize;
47+
let assistant_width = ((context_data.assistant_tokens.value() as f64 / context_data.context_window_size as f64)
48+
* progress_bar_width as f64) as usize;
49+
let tools_width = ((context_data.tools_tokens.value() as f64 / context_data.context_window_size as f64)
50+
* progress_bar_width as f64) as usize;
51+
let user_width = ((context_data.user_tokens.value() as f64 / context_data.context_window_size as f64)
52+
* progress_bar_width as f64) as usize;
53+
54+
let left_over_width = progress_bar_width
55+
- std::cmp::min(
56+
context_width + assistant_width + user_width + tools_width,
57+
progress_bar_width,
58+
);
59+
60+
let is_overflow = (context_width + assistant_width + user_width + tools_width) > progress_bar_width;
61+
62+
let total_percentage = calculate_context_percentage(context_data.total_tokens, context_data.context_window_size);
63+
64+
if is_overflow {
65+
queue!(
66+
session.stderr,
67+
style::Print(format!(
68+
"\nCurrent context window ({} of {}k tokens used)\n",
69+
context_data.total_tokens,
70+
context_data.context_window_size / 1000
71+
)),
72+
StyledText::error_fg(),
73+
style::Print("█".repeat(progress_bar_width)),
74+
StyledText::reset(),
75+
style::Print(" "),
76+
style::Print(format!("{total_percentage:.2}%")),
77+
)?;
78+
} else {
79+
queue!(
80+
session.stderr,
81+
style::Print(format!(
82+
"\nCurrent context window ({} of {}k tokens used)\n",
83+
context_data.total_tokens,
84+
context_data.context_window_size / 1000
85+
)),
86+
// Context files
87+
StyledText::brand_fg(),
88+
// add a nice visual to mimic "tiny" progress, so the overrall progress bar doesn't look too
89+
// empty
90+
style::Print(
91+
"|".repeat(if context_width == 0 && context_data.context_tokens.value() > 0 {
92+
1
93+
} else {
94+
0
95+
})
96+
),
97+
style::Print("█".repeat(context_width)),
98+
// Tools
99+
StyledText::error_fg(),
100+
style::Print(
101+
"|".repeat(if tools_width == 0 && context_data.tools_tokens.value() > 0 {
102+
1
103+
} else {
104+
0
105+
})
106+
),
107+
style::Print("█".repeat(tools_width)),
108+
// Assistant responses
109+
StyledText::info_fg(),
110+
style::Print(
111+
"|".repeat(if assistant_width == 0 && context_data.assistant_tokens.value() > 0 {
112+
1
113+
} else {
114+
0
115+
})
116+
),
117+
style::Print("█".repeat(assistant_width)),
118+
// User prompts
119+
StyledText::emphasis_fg(),
120+
style::Print("|".repeat(if user_width == 0 && context_data.user_tokens.value() > 0 {
121+
1
122+
} else {
123+
0
124+
})),
125+
style::Print("█".repeat(user_width)),
126+
StyledText::secondary_fg(),
127+
style::Print("█".repeat(left_over_width)),
128+
style::Print(" "),
129+
StyledText::reset(),
130+
style::Print(format!("{total_percentage:.2}%")),
131+
)?;
132+
}
133+
134+
execute!(session.stderr, style::Print("\n\n"))?;
135+
136+
queue!(
137+
session.stderr,
138+
StyledText::brand_fg(),
139+
style::Print("█ Context files: "),
140+
StyledText::reset(),
141+
style::Print(format!(
142+
"~{} tokens ({:.2}%)\n",
143+
context_data.context_tokens,
144+
calculate_context_percentage(context_data.context_tokens, context_data.context_window_size)
145+
)),
146+
StyledText::error_fg(),
147+
style::Print("█ Tools: "),
148+
StyledText::reset(),
149+
style::Print(format!(
150+
" ~{} tokens ({:.2}%)\n",
151+
context_data.tools_tokens,
152+
calculate_context_percentage(context_data.tools_tokens, context_data.context_window_size)
153+
)),
154+
StyledText::info_fg(),
155+
style::Print("█ Kiro responses: "),
156+
StyledText::reset(),
157+
style::Print(format!(
158+
" ~{} tokens ({:.2}%)\n",
159+
context_data.assistant_tokens,
160+
calculate_context_percentage(context_data.assistant_tokens, context_data.context_window_size)
161+
)),
162+
StyledText::emphasis_fg(),
163+
style::Print("█ Your prompts: "),
164+
StyledText::reset(),
165+
style::Print(format!(
166+
" ~{} tokens ({:.2}%)\n\n",
167+
context_data.user_tokens,
168+
calculate_context_percentage(context_data.user_tokens, context_data.context_window_size)
169+
)),
170+
)?;
171+
172+
queue!(
173+
session.stderr,
174+
style::SetAttribute(Attribute::Bold),
175+
style::Print("\n💡 Pro Tips:\n"),
176+
StyledText::reset_attributes(),
177+
StyledText::secondary_fg(),
178+
style::Print("Run "),
179+
StyledText::success_fg(),
180+
style::Print("/compact"),
181+
StyledText::secondary_fg(),
182+
style::Print(" to replace the conversation history with its summary\n"),
183+
style::Print("Run "),
184+
StyledText::success_fg(),
185+
style::Print("/clear"),
186+
StyledText::secondary_fg(),
187+
style::Print(" to erase the entire chat history\n"),
188+
style::Print("Run "),
189+
StyledText::success_fg(),
190+
style::Print("/context show"),
191+
StyledText::secondary_fg(),
192+
style::Print(" to see tokens per context file\n\n"),
193+
StyledText::reset(),
194+
)?;
195+
196+
Ok(())
197+
}

crates/chat-cli/src/cli/chat/cli/context.rs renamed to crates/chat-cli/src/cli/chat/cli/context/mod.rs

Lines changed: 5 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -10,10 +10,9 @@ use crossterm::{
1010
style,
1111
};
1212

13-
use crate::cli::chat::cli::usage::{
14-
usage_data_provider,
15-
usage_renderer,
16-
};
13+
pub mod context_data_provider;
14+
pub mod context_renderer;
15+
1716
use crate::cli::chat::consts::AGENT_FORMAT_HOOKS_DOC_URL;
1817
use crate::cli::chat::context::{
1918
ContextFilePath,
@@ -54,8 +53,8 @@ impl ContextArgs {
5453
Some(subcommand) => subcommand.execute(os, session).await,
5554
None => {
5655
// No subcommand - show context window usage (replaces /usage --context)
57-
let usage_data = usage_data_provider::get_detailed_usage_data(session, os).await?;
58-
usage_renderer::render_context_window(&usage_data, session).await?;
56+
let contexet_data = context_data_provider::get_detailed_context_data(session, os).await?;
57+
context_renderer::render_context_window(&contexet_data, session).await?;
5958

6059
Ok(ChatState::PromptUser {
6160
skip_printing_tools: true,

‎crates/chat-cli/src/cli/chat/cli/usage/mod.rs‎

Lines changed: 0 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,5 @@
11
use clap::Args;
22

3-
use crate::cli::chat::token_counter::TokenCount;
43
use crate::cli::chat::{
54
ChatError,
65
ChatSession,
@@ -11,18 +10,6 @@ use crate::os::Os;
1110
pub mod usage_data_provider;
1211
pub mod usage_renderer;
1312

14-
/// Detailed usage data for context window analysis
15-
#[derive(Debug)]
16-
pub struct DetailedUsageData {
17-
pub total_tokens: TokenCount,
18-
pub context_tokens: TokenCount,
19-
pub assistant_tokens: TokenCount,
20-
pub user_tokens: TokenCount,
21-
pub tools_tokens: TokenCount,
22-
pub context_window_size: usize,
23-
pub dropped_context_files: Vec<(String, String)>,
24-
}
25-
2613
/// Billing usage data from API
2714
#[derive(Debug)]
2815
pub struct BillingUsageData {

‎crates/chat-cli/src/cli/chat/cli/usage/usage_data_provider.rs‎

Lines changed: 1 addition & 49 deletions
Original file line numberDiff line numberDiff line change
@@ -7,51 +7,9 @@ use crate::api_client::error_utils::{
77
GetUsageLimitsErrorType,
88
classify_get_usage_limits_error,
99
};
10-
use crate::cli::chat::cli::model::context_window_tokens;
11-
use crate::cli::chat::token_counter::{
12-
CharCount,
13-
TokenCount,
14-
};
15-
use crate::cli::chat::{
16-
ChatError,
17-
ChatSession,
18-
};
10+
use crate::cli::chat::ChatError;
1911
use crate::os::Os;
2012

21-
/// Get detailed usage data for context window analysis
22-
pub async fn get_detailed_usage_data(
23-
session: &mut ChatSession,
24-
os: &Os,
25-
) -> Result<super::DetailedUsageData, ChatError> {
26-
let context_window_size = context_window_tokens(session.conversation.model_info.as_ref());
27-
28-
let state = session
29-
.conversation
30-
.backend_conversation_state(os, true, &mut std::io::stderr())
31-
.await?;
32-
33-
let data = state.calculate_conversation_size();
34-
let tool_specs_json: String = state
35-
.tools
36-
.values()
37-
.filter_map(|s| serde_json::to_string(s).ok())
38-
.collect::<Vec<String>>()
39-
.join("");
40-
let tools_char_count: CharCount = tool_specs_json.len().into();
41-
let total_tokens: TokenCount =
42-
(data.context_messages + data.user_messages + data.assistant_messages + tools_char_count).into();
43-
44-
Ok(super::DetailedUsageData {
45-
total_tokens,
46-
context_tokens: data.context_messages.into(),
47-
assistant_tokens: data.assistant_messages.into(),
48-
user_tokens: data.user_messages.into(),
49-
tools_tokens: tools_char_count.into(),
50-
context_window_size,
51-
dropped_context_files: state.dropped_context_files,
52-
})
53-
}
54-
5513
/// Get billing usage data from API
5614
pub(super) async fn get_billing_usage_data(os: &Os) -> Result<super::BillingUsageData, ChatError> {
5715
match os.client.get_usage_limits().await {
@@ -163,9 +121,3 @@ pub(super) async fn get_billing_usage_data(os: &Os) -> Result<super::BillingUsag
163121
},
164122
}
165123
}
166-
167-
/// Get total usage percentage (external API)
168-
pub async fn get_total_usage_percentage(session: &mut ChatSession, os: &Os) -> Result<f32, ChatError> {
169-
let data = get_detailed_usage_data(session, os).await?;
170-
Ok((data.total_tokens.value() as f32 / data.context_window_size as f32) * 100.0)
171-
}

0 commit comments

Comments
 (0)