@@ -10,14 +10,16 @@ use crossterm::{
10
10
} ;
11
11
12
12
use crate :: cli:: chat:: consts:: CONTEXT_WINDOW_SIZE ;
13
- use crate :: cli:: chat:: token_counter:: TokenCount ;
13
+ use crate :: cli:: chat:: token_counter:: {
14
+ CharCount ,
15
+ TokenCount ,
16
+ } ;
14
17
use crate :: cli:: chat:: {
15
18
ChatError ,
16
19
ChatSession ,
17
20
ChatState ,
18
21
} ;
19
22
use crate :: platform:: Context ;
20
-
21
23
#[ deny( missing_docs) ]
22
24
#[ derive( Debug , PartialEq , Args ) ]
23
25
pub struct UsageArgs ;
@@ -43,13 +45,19 @@ impl UsageArgs {
43
45
}
44
46
45
47
let data = state. calculate_conversation_size ( ) ;
46
-
48
+ let tool_specs_json: String = state
49
+ . tools
50
+ . values ( )
51
+ . filter_map ( |s| serde_json:: to_string ( s) . ok ( ) )
52
+ . collect :: < Vec < String > > ( )
53
+ . join ( "" ) ;
47
54
let context_token_count: TokenCount = data. context_messages . into ( ) ;
48
55
let assistant_token_count: TokenCount = data. assistant_messages . into ( ) ;
49
56
let user_token_count: TokenCount = data. user_messages . into ( ) ;
57
+ let tools_char_count: CharCount = tool_specs_json. len ( ) . into ( ) ; // usize → CharCount
58
+ let tools_token_count: TokenCount = tools_char_count. into ( ) ; // CharCount → TokenCount
50
59
let total_token_used: TokenCount =
51
- ( data. context_messages + data. user_messages + data. assistant_messages ) . into ( ) ;
52
-
60
+ ( data. context_messages + data. user_messages + data. assistant_messages + tools_char_count) . into ( ) ;
53
61
let window_width = session. terminal_width ( ) ;
54
62
// set a max width for the progress bar for better aesthetic
55
63
let progress_bar_width = std:: cmp:: min ( window_width, 80 ) ;
@@ -58,13 +66,18 @@ impl UsageArgs {
58
66
( ( context_token_count. value ( ) as f64 / CONTEXT_WINDOW_SIZE as f64 ) * progress_bar_width as f64 ) as usize ;
59
67
let assistant_width =
60
68
( ( assistant_token_count. value ( ) as f64 / CONTEXT_WINDOW_SIZE as f64 ) * progress_bar_width as f64 ) as usize ;
69
+ let tools_width =
70
+ ( ( tools_token_count. value ( ) as f64 / CONTEXT_WINDOW_SIZE as f64 ) * progress_bar_width as f64 ) as usize ;
61
71
let user_width =
62
72
( ( user_token_count. value ( ) as f64 / CONTEXT_WINDOW_SIZE as f64 ) * progress_bar_width as f64 ) as usize ;
63
73
64
- let left_over_width =
65
- progress_bar_width - std:: cmp:: min ( context_width + assistant_width + user_width, progress_bar_width) ;
74
+ let left_over_width = progress_bar_width
75
+ - std:: cmp:: min (
76
+ context_width + assistant_width + user_width + tools_width,
77
+ progress_bar_width,
78
+ ) ;
66
79
67
- let is_overflow = ( context_width + assistant_width + user_width) > progress_bar_width;
80
+ let is_overflow = ( context_width + assistant_width + user_width + tools_width ) > progress_bar_width;
68
81
69
82
if is_overflow {
70
83
queue ! (
@@ -91,6 +104,7 @@ impl UsageArgs {
91
104
total_token_used,
92
105
CONTEXT_WINDOW_SIZE / 1000
93
106
) ) ,
107
+ // Context files
94
108
style:: SetForegroundColor ( Color :: DarkCyan ) ,
95
109
// add a nice visual to mimic "tiny" progress, so the overral progress bar doesn't look too
96
110
// empty
@@ -100,13 +114,23 @@ impl UsageArgs {
100
114
0
101
115
} ) ) ,
102
116
style:: Print ( "█" . repeat( context_width) ) ,
117
+ // Tools
118
+ style:: SetForegroundColor ( Color :: DarkRed ) ,
119
+ style:: Print ( "|" . repeat( if tools_width == 0 && * tools_token_count > 0 {
120
+ 1
121
+ } else {
122
+ 0
123
+ } ) ) ,
124
+ style:: Print ( "█" . repeat( tools_width) ) ,
125
+ // Assistant responses
103
126
style:: SetForegroundColor ( Color :: Blue ) ,
104
127
style:: Print ( "|" . repeat( if assistant_width == 0 && * assistant_token_count > 0 {
105
128
1
106
129
} else {
107
130
0
108
131
} ) ) ,
109
132
style:: Print ( "█" . repeat( assistant_width) ) ,
133
+ // User prompts
110
134
style:: SetForegroundColor ( Color :: Magenta ) ,
111
135
style:: Print ( "|" . repeat( if user_width == 0 && * user_token_count > 0 { 1 } else { 0 } ) ) ,
112
136
style:: Print ( "█" . repeat( user_width) ) ,
@@ -133,6 +157,14 @@ impl UsageArgs {
133
157
context_token_count,
134
158
( context_token_count. value( ) as f32 / CONTEXT_WINDOW_SIZE as f32 ) * 100.0
135
159
) ) ,
160
+ style:: SetForegroundColor ( Color :: DarkRed ) ,
161
+ style:: Print ( "█ Tools: " ) ,
162
+ style:: SetForegroundColor ( Color :: Reset ) ,
163
+ style:: Print ( format!(
164
+ " ~{} tokens ({:.2}%)\n " ,
165
+ tools_token_count,
166
+ ( tools_token_count. value( ) as f32 / CONTEXT_WINDOW_SIZE as f32 ) * 100.0
167
+ ) ) ,
136
168
style:: SetForegroundColor ( Color :: Blue ) ,
137
169
style:: Print ( "█ Q responses: " ) ,
138
170
style:: SetForegroundColor ( Color :: Reset ) ,
0 commit comments