feat: Enhance token usage tracking and improve chat UI with emojis
This commit is contained in:
asepharyana
@@ -107,6 +107,9 @@ impl SseParser {
|
||||
return vec![];
|
||||
}
|
||||
};
|
||||
|
||||
let mut events = Vec::new();
|
||||
|
||||
if let Some(usage) = value.get("usage") {
|
||||
if !usage.is_null() {
|
||||
let prompt_tokens = usage.get("prompt_tokens").and_then(serde_json::Value::as_u64).unwrap_or_else(|| {
|
||||
@@ -122,84 +125,73 @@ impl SseParser {
|
||||
tracing::warn!("[stream] total_tokens missing in usage chunk");
|
||||
prompt_tokens + completion_tokens
|
||||
});
|
||||
// Only emit Usage as a standalone event if this chunk
|
||||
// contains nothing else (no choices, no delta). Some
|
||||
// non-standard providers may bundle usage WITH content
|
||||
// in the same chunk; emitting both prevents content loss.
|
||||
let has_other_content = value.get("choices")
|
||||
.and_then(|c| c.as_array())
|
||||
.is_some_and(|arr| arr.iter().any(|ch| {
|
||||
ch.get("delta").and_then(|d| d.get("content")).is_some()
|
||||
|| ch.get("delta").and_then(|d| d.get("reasoning_content")).is_some()
|
||||
|| ch.get("delta").and_then(|d| d.get("tool_calls")).is_some()
|
||||
}));
|
||||
if !has_other_content {
|
||||
return vec![StreamEvent::Usage { prompt_tokens, completion_tokens, total_tokens }];
|
||||
}
|
||||
events.push(StreamEvent::Usage { prompt_tokens, completion_tokens, total_tokens });
|
||||
}
|
||||
}
|
||||
match event_type.as_str() {
|
||||
|
||||
let mut other_events = match event_type.as_str() {
|
||||
"message.stop" => vec![StreamEvent::Done],
|
||||
"message.delta" | "" => {
|
||||
let Some(delta) = value.get("delta").or_else(|| value.get("choices")) else { return vec![] };
|
||||
if let Some(choices) = delta.as_array() {
|
||||
let Some(choice) = choices.first() else { return vec![] };
|
||||
let Some(d) = choice.get("delta") else { return vec![] };
|
||||
let mut d_events = Vec::new();
|
||||
if let Some(delta) = value.get("delta").or_else(|| value.get("choices")) {
|
||||
if let Some(choices) = delta.as_array() {
|
||||
if let Some(choice) = choices.first() {
|
||||
if let Some(d) = choice.get("delta") {
|
||||
// Content token
|
||||
if let Some(content) = d.get("content").and_then(|c| c.as_str()) {
|
||||
d_events.push(StreamEvent::Token(content.to_string()));
|
||||
}
|
||||
|
||||
// Content token
|
||||
if let Some(content) = d.get("content").and_then(|c| c.as_str()) {
|
||||
return vec![StreamEvent::Token(content.to_string())];
|
||||
}
|
||||
// Reasoning token
|
||||
if let Some(reasoning) = d.get("reasoning_content").and_then(|r| r.as_str()) {
|
||||
d_events.push(StreamEvent::Reasoning(reasoning.to_string()));
|
||||
}
|
||||
|
||||
// Reasoning token
|
||||
if let Some(reasoning) = d.get("reasoning_content").and_then(|r| r.as_str()) {
|
||||
return vec![StreamEvent::Reasoning(reasoning.to_string())];
|
||||
}
|
||||
// Tool calls — iterate ALL entries, not just first()
|
||||
if let Some(tool_calls) = d.get("tool_calls").and_then(|tc| tc.as_array()) {
|
||||
for tc in tool_calls {
|
||||
let index = tc.get("index").and_then(serde_json::Value::as_u64).unwrap_or_else(|| {
|
||||
tracing::warn!("[stream] tool call delta missing index, defaulting to 0");
|
||||
0
|
||||
}) as usize;
|
||||
let id = tc.get("id").and_then(|i| i.as_str()).map(std::string::ToString::to_string);
|
||||
let name = tc.get("function")
|
||||
.and_then(|f| f.get("name"))
|
||||
.and_then(|n| n.as_str())
|
||||
.map(std::string::ToString::to_string);
|
||||
let args_delta = tc.get("function")
|
||||
.and_then(|f| f.get("arguments"))
|
||||
.and_then(|a| a.as_str())
|
||||
.unwrap_or("")
|
||||
.to_string();
|
||||
d_events.push(StreamEvent::ToolCallDelta {
|
||||
index,
|
||||
id,
|
||||
name,
|
||||
arguments_delta: args_delta,
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
// Tool calls — iterate ALL entries, not just first()
|
||||
if let Some(tool_calls) = d.get("tool_calls").and_then(|tc| tc.as_array()) {
|
||||
let mut events = Vec::with_capacity(tool_calls.len());
|
||||
for tc in tool_calls {
|
||||
let index = tc.get("index").and_then(serde_json::Value::as_u64).unwrap_or_else(|| {
|
||||
tracing::warn!("[stream] tool call delta missing index, defaulting to 0");
|
||||
0
|
||||
}) as usize;
|
||||
let id = tc.get("id").and_then(|i| i.as_str()).map(std::string::ToString::to_string);
|
||||
let name = tc.get("function")
|
||||
.and_then(|f| f.get("name"))
|
||||
.and_then(|n| n.as_str())
|
||||
.map(std::string::ToString::to_string);
|
||||
let args_delta = tc.get("function")
|
||||
.and_then(|f| f.get("arguments"))
|
||||
.and_then(|a| a.as_str())
|
||||
.unwrap_or("")
|
||||
.to_string();
|
||||
events.push(StreamEvent::ToolCallDelta {
|
||||
index,
|
||||
id,
|
||||
name,
|
||||
arguments_delta: args_delta,
|
||||
});
|
||||
}
|
||||
if !events.is_empty() {
|
||||
return events;
|
||||
}
|
||||
}
|
||||
|
||||
// Finish reason
|
||||
if let Some(reason) = choice.get("finish_reason").and_then(|r| r.as_str()) {
|
||||
if reason == "stop" || reason == "tool_calls" {
|
||||
return vec![StreamEvent::Done];
|
||||
// Finish reason
|
||||
if let Some(reason) = choice.get("finish_reason").and_then(|r| r.as_str()) {
|
||||
if reason == "stop" || reason == "tool_calls" {
|
||||
d_events.push(StreamEvent::Done);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
} else if let Some(content) = delta.get("content").and_then(|c| c.as_str()) {
|
||||
d_events.push(StreamEvent::Token(content.to_string()));
|
||||
}
|
||||
}
|
||||
if let Some(content) = delta.get("content").and_then(|c| c.as_str()) {
|
||||
return vec![StreamEvent::Token(content.to_string())];
|
||||
}
|
||||
vec![]
|
||||
d_events
|
||||
}
|
||||
_ => vec![],
|
||||
}
|
||||
};
|
||||
|
||||
events.append(&mut other_events);
|
||||
events
|
||||
}
|
||||
|
||||
/// Clears any partially-buffered SSE frame. Reserved for reconnect/retry flows that
|
||||
@@ -350,6 +342,27 @@ mod tests {
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn feed_parses_usage_and_content_bundled_chunk() {
|
||||
let mut p = SseParser::new();
|
||||
let events = p.feed(
|
||||
"data: {\"choices\":[{\"delta\":{\"content\":\"hello\"}}],\"usage\":{\"prompt_tokens\":10,\"completion_tokens\":5,\"total_tokens\":15}}\n\n",
|
||||
);
|
||||
assert_eq!(events.len(), 2);
|
||||
match (&events[0], &events[1]) {
|
||||
(
|
||||
StreamEvent::Usage { prompt_tokens, completion_tokens, total_tokens },
|
||||
StreamEvent::Token(t),
|
||||
) => {
|
||||
assert_eq!(*prompt_tokens, 10);
|
||||
assert_eq!(*completion_tokens, 5);
|
||||
assert_eq!(*total_tokens, 15);
|
||||
assert_eq!(t, "hello");
|
||||
}
|
||||
other => panic!("expected [Usage, Token], got {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn feed_ignores_empty_data_lines() {
|
||||
let mut p = SseParser::new();
|
||||
|
||||
Reference in New Issue
Block a user