Stream token usage on every agent message (#5342)

This commit is contained in:
David Katz
2025-10-29 23:23:25 -04:00
committed by GitHub
parent c875f13757
commit 37e1bb1d37
18 changed files with 214 additions and 63 deletions
+24 -14
View File
@@ -825,9 +825,11 @@ impl Agent {
}
}
Err(e) => {
yield AgentEvent::Message(Message::assistant().with_text(
format!("Ran into this error trying to compact: {e}.\n\nPlease try again or create a new session")
));
yield AgentEvent::Message(
Message::assistant().with_text(
format!("Ran into this error trying to compact: {e}.\n\nPlease try again or create a new session")
)
);
}
}
}))
@@ -917,7 +919,7 @@ impl Agent {
if let Some(final_output_tool) = self.final_output_tool.lock().await.as_ref() {
if final_output_tool.final_output.is_some() {
let final_event = AgentEvent::Message(
Message::assistant().with_text(final_output_tool.final_output.clone().unwrap()),
Message::assistant().with_text(final_output_tool.final_output.clone().unwrap())
);
yield final_event;
break;
@@ -926,9 +928,11 @@ impl Agent {
turns_taken += 1;
if turns_taken > max_turns {
yield AgentEvent::Message(Message::assistant().with_text(
"I've reached the maximum number of actions I can do without user input. Would you like me to continue?"
));
yield AgentEvent::Message(
Message::assistant().with_text(
"I've reached the maximum number of actions I can do without user input. Would you like me to continue?"
)
);
break;
}
@@ -1178,18 +1182,22 @@ impl Agent {
}
Err(e) => {
error!("Error: {}", e);
yield AgentEvent::Message(Message::assistant().with_text(
yield AgentEvent::Message(
Message::assistant().with_text(
format!("Ran into this error trying to compact: {e}.\n\nPlease retry if you think this is a transient or recoverable error.")
));
)
);
break;
}
}
}
Err(e) => {
error!("Error: {}", e);
yield AgentEvent::Message(Message::assistant().with_text(
yield AgentEvent::Message(
Message::assistant().with_text(
format!("Ran into this error: {e}.\n\nPlease retry if you think this is a transient or recoverable error.")
));
)
);
break;
}
}
@@ -1224,9 +1232,11 @@ impl Agent {
}
Err(e) => {
error!("Retry logic failed: {}", e);
yield AgentEvent::Message(Message::assistant().with_text(
format!("Retry logic encountered an error: {}", e)
));
yield AgentEvent::Message(
Message::assistant().with_text(
format!("Retry logic encountered an error: {}", e)
)
);
exit_chat = true;
}
}