feat: merge 2nd round from all rcc/* sessions
- api: tool_use parsing, message_delta, request_id tracking, retry logic - tools: extended tool suite (WebSearch, WebFetch, Agent, etc.) - cli: live streamed conversations, session restore, compact commands - runtime: config loading, system prompt builder, token usage, compaction
This commit is contained in:
121
rust/crates/runtime/src/usage.rs
Normal file
121
rust/crates/runtime/src/usage.rs
Normal file
@@ -0,0 +1,121 @@
|
||||
use crate::session::Session;
|
||||
|
||||
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
|
||||
pub struct TokenUsage {
|
||||
pub input_tokens: u32,
|
||||
pub output_tokens: u32,
|
||||
pub cache_creation_input_tokens: u32,
|
||||
pub cache_read_input_tokens: u32,
|
||||
}
|
||||
|
||||
impl TokenUsage {
|
||||
#[must_use]
|
||||
pub fn total_tokens(self) -> u32 {
|
||||
self.input_tokens
|
||||
+ self.output_tokens
|
||||
+ self.cache_creation_input_tokens
|
||||
+ self.cache_read_input_tokens
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Default, PartialEq, Eq)]
|
||||
pub struct UsageTracker {
|
||||
latest_turn: TokenUsage,
|
||||
cumulative: TokenUsage,
|
||||
turns: u32,
|
||||
}
|
||||
|
||||
impl UsageTracker {
|
||||
#[must_use]
|
||||
pub fn new() -> Self {
|
||||
Self::default()
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn from_session(session: &Session) -> Self {
|
||||
let mut tracker = Self::new();
|
||||
for message in &session.messages {
|
||||
if let Some(usage) = message.usage {
|
||||
tracker.record(usage);
|
||||
}
|
||||
}
|
||||
tracker
|
||||
}
|
||||
|
||||
pub fn record(&mut self, usage: TokenUsage) {
|
||||
self.latest_turn = usage;
|
||||
self.cumulative.input_tokens += usage.input_tokens;
|
||||
self.cumulative.output_tokens += usage.output_tokens;
|
||||
self.cumulative.cache_creation_input_tokens += usage.cache_creation_input_tokens;
|
||||
self.cumulative.cache_read_input_tokens += usage.cache_read_input_tokens;
|
||||
self.turns += 1;
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn current_turn_usage(&self) -> TokenUsage {
|
||||
self.latest_turn
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn cumulative_usage(&self) -> TokenUsage {
|
||||
self.cumulative
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn turns(&self) -> u32 {
|
||||
self.turns
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::{TokenUsage, UsageTracker};
|
||||
use crate::session::{ContentBlock, ConversationMessage, MessageRole, Session};
|
||||
|
||||
#[test]
|
||||
fn tracks_true_cumulative_usage() {
|
||||
let mut tracker = UsageTracker::new();
|
||||
tracker.record(TokenUsage {
|
||||
input_tokens: 10,
|
||||
output_tokens: 4,
|
||||
cache_creation_input_tokens: 2,
|
||||
cache_read_input_tokens: 1,
|
||||
});
|
||||
tracker.record(TokenUsage {
|
||||
input_tokens: 20,
|
||||
output_tokens: 6,
|
||||
cache_creation_input_tokens: 3,
|
||||
cache_read_input_tokens: 2,
|
||||
});
|
||||
|
||||
assert_eq!(tracker.turns(), 2);
|
||||
assert_eq!(tracker.current_turn_usage().input_tokens, 20);
|
||||
assert_eq!(tracker.current_turn_usage().output_tokens, 6);
|
||||
assert_eq!(tracker.cumulative_usage().output_tokens, 10);
|
||||
assert_eq!(tracker.cumulative_usage().input_tokens, 30);
|
||||
assert_eq!(tracker.cumulative_usage().total_tokens(), 48);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn reconstructs_usage_from_session_messages() {
|
||||
let session = Session {
|
||||
version: 1,
|
||||
messages: vec![ConversationMessage {
|
||||
role: MessageRole::Assistant,
|
||||
blocks: vec![ContentBlock::Text {
|
||||
text: "done".to_string(),
|
||||
}],
|
||||
usage: Some(TokenUsage {
|
||||
input_tokens: 5,
|
||||
output_tokens: 2,
|
||||
cache_creation_input_tokens: 1,
|
||||
cache_read_input_tokens: 0,
|
||||
}),
|
||||
}],
|
||||
};
|
||||
|
||||
let tracker = UsageTracker::from_session(&session);
|
||||
assert_eq!(tracker.turns(), 1);
|
||||
assert_eq!(tracker.cumulative_usage().total_tokens(), 8);
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user