Skip to content

Commit 559245b

Browse files
committed
feat: /usage
1 parent 3e4ae79 commit 559245b

File tree

4 files changed

+160
-0
lines changed

4 files changed

+160
-0
lines changed

crates/q_cli/src/cli/chat/command.rs

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -38,6 +38,7 @@ pub enum Command {
3838
Tools {
3939
subcommand: Option<ToolsSubcommand>,
4040
},
41+
Usage,
4142
}
4243

4344
#[derive(Debug, Clone, PartialEq, Eq)]
@@ -538,6 +539,7 @@ impl Command {
538539
},
539540
}
540541
},
542+
"usage" => Self::Usage,
541543
_ => {
542544
return Ok(Self::Ask {
543545
prompt: input.to_string(),

crates/q_cli/src/cli/chat/conversation_state.rs

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -185,6 +185,11 @@ impl ConversationState {
185185
self.conversation_id.as_ref()
186186
}
187187

188+
/// Returns the conversation history.
189+
pub fn get_chat_history(&self) -> Vec<ChatMessage> {
190+
self.history.iter().cloned().collect()
191+
}
192+
188193
/// Returns the message id associated with the last assistant message, if present.
189194
///
190195
/// This is equivalent to `utterance_id` in the Q API.

crates/q_cli/src/cli/chat/mod.rs

Lines changed: 152 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -196,6 +196,7 @@ const HELP_TEXT: &str = color_print::cstr! {"
196196
<em>add</em> <black!>Add file(s) to context [--global] [--force]</black!>
197197
<em>rm</em> <black!>Remove file(s) from context [--global]</black!>
198198
<em>clear</em> <black!>Clear all files from current context [--global]</black!>
199+
<em>/usage</em> <black!>Show current session's context window usage</black!>
199200
200201
<cyan,em>Tips:</cyan,em>
201202
<em>!{command}</em> <black!>Quickly execute a command in your current session</black!>
@@ -205,6 +206,8 @@ const HELP_TEXT: &str = color_print::cstr! {"
205206

206207
const RESPONSE_TIMEOUT_CONTENT: &str = "Response timed out - message took too long to generate";
207208

209+
const CONTEXT_WINDOW_SIZE: usize = 200_000; // tokens
210+
208211
pub async fn chat(
209212
input: Option<String>,
210213
no_interactive: bool,
@@ -1535,6 +1538,155 @@ where
15351538
// during PromptUser.
15361539
execute!(self.output, style::Print("\n\n"),)?;
15371540

1541+
ChatState::PromptUser {
1542+
tool_uses: Some(tool_uses),
1543+
pending_tool_index,
1544+
skip_printing_tools: true,
1545+
}
1546+
},
1547+
Command::Usage => {
1548+
let context_messages = self.conversation_state.context_messages().await;
1549+
let chat_history = self.conversation_state.get_chat_history();
1550+
let assistant_messages = chat_history
1551+
.iter()
1552+
.filter_map(|message| {
1553+
if let fig_api_client::model::ChatMessage::AssistantResponseMessage(msg) = message {
1554+
Some(msg)
1555+
} else {
1556+
None
1557+
}
1558+
})
1559+
.collect::<Vec<_>>();
1560+
1561+
let user_messages = chat_history
1562+
.iter()
1563+
.filter_map(|message| {
1564+
if let fig_api_client::model::ChatMessage::UserInputMessage(msg) = message {
1565+
Some(msg)
1566+
} else {
1567+
None
1568+
}
1569+
})
1570+
.collect::<Vec<_>>();
1571+
1572+
let context_token_count = context_messages
1573+
.iter()
1574+
.map(|msg| TokenCounter::count_tokens(&msg.0.content))
1575+
.sum::<usize>();
1576+
1577+
let assistant_token_count = assistant_messages
1578+
.iter()
1579+
.map(|msg| TokenCounter::count_tokens(&msg.content))
1580+
.sum::<usize>();
1581+
1582+
let user_token_count = user_messages
1583+
.iter()
1584+
.map(|msg| TokenCounter::count_tokens(&msg.content))
1585+
.sum::<usize>();
1586+
1587+
let total_token_used: usize = context_token_count + assistant_token_count + user_token_count;
1588+
1589+
let window_width = self.terminal_width();
1590+
let progress_bar_width = std::cmp::min(window_width, 80); // set a max width for the progress bar for better aesthetic
1591+
1592+
let context_width =
1593+
((context_token_count as f64 / CONTEXT_WINDOW_SIZE as f64) * progress_bar_width as f64) as usize;
1594+
let assistant_width =
1595+
((assistant_token_count as f64 / CONTEXT_WINDOW_SIZE as f64) * progress_bar_width as f64) as usize;
1596+
let user_width =
1597+
((user_token_count as f64 / CONTEXT_WINDOW_SIZE as f64) * progress_bar_width as f64) as usize;
1598+
let left_over_width = progress_bar_width - context_width - assistant_width - user_width;
1599+
1600+
queue!(
1601+
self.output,
1602+
style::Print(format!(
1603+
"\nCurrent context window ({} of {}k tokens used)\n",
1604+
total_token_used,
1605+
CONTEXT_WINDOW_SIZE / 1000
1606+
)),
1607+
style::SetForegroundColor(Color::DarkCyan),
1608+
// add a nice visual to mimic "tiny" progress, so the overral progress bar doesn't look too empty
1609+
style::Print("|".repeat(if context_width == 0 && context_token_count > 0 {
1610+
1
1611+
} else {
1612+
0
1613+
})),
1614+
style::Print("█".repeat(context_width)),
1615+
style::SetForegroundColor(Color::Blue),
1616+
style::Print("|".repeat(if assistant_width == 0 && assistant_token_count > 0 {
1617+
1
1618+
} else {
1619+
0
1620+
})),
1621+
style::Print("█".repeat(assistant_width)),
1622+
style::SetForegroundColor(Color::Magenta),
1623+
style::Print("|".repeat(if user_width == 0 && user_token_count > 0 { 1 } else { 0 })),
1624+
style::Print("█".repeat(user_width)),
1625+
style::SetForegroundColor(Color::DarkGrey),
1626+
style::Print("█".repeat(left_over_width)),
1627+
style::Print(" "),
1628+
style::SetForegroundColor(Color::Reset),
1629+
style::Print(format!(
1630+
"{:.2}%",
1631+
(total_token_used as f32 / CONTEXT_WINDOW_SIZE as f32) * 100.0
1632+
)),
1633+
)?;
1634+
1635+
queue!(self.output, style::Print("\n\n"))?;
1636+
self.output.flush()?;
1637+
1638+
queue!(
1639+
self.output,
1640+
style::SetForegroundColor(Color::DarkCyan),
1641+
style::Print("█ Context files: "),
1642+
style::SetForegroundColor(Color::Reset),
1643+
style::Print(format!(
1644+
"~{} tokens ({:.2}%)\n",
1645+
context_token_count,
1646+
(context_token_count as f32 / CONTEXT_WINDOW_SIZE as f32) * 100.0
1647+
)),
1648+
style::SetForegroundColor(Color::Blue),
1649+
style::Print("█ Q responses: "),
1650+
style::SetForegroundColor(Color::Reset),
1651+
style::Print(format!(
1652+
" ~{} tokens ({:.2}%)\n",
1653+
assistant_token_count,
1654+
(assistant_token_count as f32 / CONTEXT_WINDOW_SIZE as f32) * 100.0
1655+
)),
1656+
style::SetForegroundColor(Color::Magenta),
1657+
style::Print("█ Your prompts: "),
1658+
style::SetForegroundColor(Color::Reset),
1659+
style::Print(format!(
1660+
" ~{} tokens ({:.2}%)\n\n",
1661+
user_token_count,
1662+
(user_token_count as f32 / CONTEXT_WINDOW_SIZE as f32) * 100.0
1663+
)),
1664+
)?;
1665+
1666+
queue!(
1667+
self.output,
1668+
style::SetAttribute(Attribute::Bold),
1669+
style::Print("\n💡 Pro Tips:\n"),
1670+
style::SetAttribute(Attribute::Reset),
1671+
style::SetForegroundColor(Color::DarkGrey),
1672+
style::Print("Run "),
1673+
style::SetForegroundColor(Color::DarkGreen),
1674+
style::Print("/compact"),
1675+
style::SetForegroundColor(Color::DarkGrey),
1676+
style::Print(" to replace the conversation history with its summary\n"),
1677+
style::Print("Run "),
1678+
style::SetForegroundColor(Color::DarkGreen),
1679+
style::Print("/clear"),
1680+
style::SetForegroundColor(Color::DarkGrey),
1681+
style::Print(" to erase the entire chat history\n"),
1682+
style::Print("Run "),
1683+
style::SetForegroundColor(Color::DarkGreen),
1684+
style::Print("/context show"),
1685+
style::SetForegroundColor(Color::DarkGrey),
1686+
style::Print(" to see tokens per context file\n\n"),
1687+
style::SetForegroundColor(Color::Reset),
1688+
)?;
1689+
15381690
ChatState::PromptUser {
15391691
tool_uses: Some(tool_uses),
15401692
pending_tool_index,

crates/q_cli/src/cli/chat/prompt.rs

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -67,6 +67,7 @@ const COMMANDS: &[&str] = &[
6767
"/compact",
6868
"/compact help",
6969
"/compact --summary",
70+
"/usage",
7071
];
7172

7273
pub fn generate_prompt(current_profile: Option<&str>, warning: bool) -> String {

0 commit comments

Comments
 (0)