use lellm_core::TokenUsage;
use lellm_graph::StreamChunk;
use lellm_provider::ProviderEvent;
#[derive(Debug, Clone)]
pub enum AgentStreamEvent {
TextDelta { text: String },
ThinkingDelta {
thinking: String,
redacted: Option<String>,
},
ToolResult {
call_id: String,
tool_name: String,
content: String,
is_error: bool,
},
Usage(TokenUsage),
ResponseComplete,
}
impl AgentStreamEvent {
pub fn to_chunk(&self) -> Option<StreamChunk> {
match self {
AgentStreamEvent::TextDelta { text } => Some(StreamChunk::TextDelta(text.clone())),
AgentStreamEvent::ThinkingDelta { thinking, redacted } => {
Some(StreamChunk::ThinkingDelta {
text: thinking.clone(),
redacted: redacted.clone(),
})
}
AgentStreamEvent::ToolResult { .. } => None,
AgentStreamEvent::Usage(_) => None,
AgentStreamEvent::ResponseComplete => None,
}
}
}
pub fn transform(event: &ProviderEvent) -> Vec<AgentStreamEvent> {
match event {
ProviderEvent::Token { token } => {
vec![AgentStreamEvent::TextDelta {
text: token.clone(),
}]
}
ProviderEvent::ThinkingDelta { thinking, redacted } => {
vec![AgentStreamEvent::ThinkingDelta {
thinking: thinking.clone(),
redacted: redacted.clone(),
}]
}
ProviderEvent::ResponseComplete { tool_calls, usage } => {
let mut events = Vec::new();
let _ = tool_calls; if let Some(u) = usage {
events.push(AgentStreamEvent::Usage(u.clone()));
}
events.push(AgentStreamEvent::ResponseComplete);
events
}
ProviderEvent::Start { .. } => vec![], }
}
#[derive(Debug)]
pub enum TranslationResult {
Emit(StreamChunk),
EmitWithText { chunk: StreamChunk, delta: String },
EmitWithThinking {
chunk: StreamChunk,
delta: String,
redacted: Option<String>,
},
Usage(TokenUsage),
Finished,
Ignore,
}
pub fn translate_provider_event(event: &ProviderEvent) -> TranslationResult {
for agent_event in transform(event) {
match agent_event {
AgentStreamEvent::TextDelta { text } => {
return TranslationResult::EmitWithText {
chunk: StreamChunk::TextDelta(text.clone()),
delta: text,
};
}
AgentStreamEvent::ThinkingDelta { thinking, redacted } => {
return TranslationResult::EmitWithThinking {
chunk: StreamChunk::ThinkingDelta {
text: thinking.clone(),
redacted: redacted.clone(),
},
delta: thinking,
redacted,
};
}
AgentStreamEvent::Usage(u) => return TranslationResult::Usage(u),
AgentStreamEvent::ResponseComplete => return TranslationResult::Finished,
AgentStreamEvent::ToolResult { .. } => {}
}
}
TranslationResult::Ignore
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_transform_token() {
let event = ProviderEvent::Token {
token: "hello".into(),
};
let result = transform(&event);
assert_eq!(result.len(), 1);
match &result[0] {
AgentStreamEvent::TextDelta { text } => assert_eq!(text, "hello"),
other => panic!("expected TextDelta, got {:?}", other),
}
}
#[test]
fn test_transform_thinking_with_redacted() {
let event = ProviderEvent::ThinkingDelta {
thinking: "visible".into(),
redacted: Some("sensitive".into()),
};
let result = transform(&event);
assert_eq!(result.len(), 1);
match &result[0] {
AgentStreamEvent::ThinkingDelta { thinking, redacted } => {
assert_eq!(thinking, "visible");
assert_eq!(redacted, &Some("sensitive".into()));
}
other => panic!("expected ThinkingDelta, got {:?}", other),
}
}
#[test]
fn test_transform_thinking_without_redacted() {
let event = ProviderEvent::ThinkingDelta {
thinking: "thinking".into(),
redacted: None,
};
let result = transform(&event);
assert_eq!(result.len(), 1);
match &result[0] {
AgentStreamEvent::ThinkingDelta { thinking, redacted } => {
assert_eq!(thinking, "thinking");
assert!(redacted.is_none());
}
other => panic!("expected ThinkingDelta, got {:?}", other),
}
}
#[test]
fn test_transform_response_complete_with_usage() {
let event = ProviderEvent::ResponseComplete {
tool_calls: vec![],
usage: Some(TokenUsage::default()),
};
let result = transform(&event);
assert_eq!(result.len(), 2);
assert!(matches!(result[0], AgentStreamEvent::Usage(_)));
assert!(matches!(result[1], AgentStreamEvent::ResponseComplete));
}
#[test]
fn test_transform_response_complete_without_usage() {
let event = ProviderEvent::ResponseComplete {
tool_calls: vec![],
usage: None,
};
let result = transform(&event);
assert_eq!(result.len(), 1);
assert!(matches!(result[0], AgentStreamEvent::ResponseComplete));
}
#[test]
fn test_transform_start_ignored() {
let event = ProviderEvent::Start {
model: "test".into(),
};
let result = transform(&event);
assert!(result.is_empty());
}
#[test]
fn test_to_chunk_text_delta() {
let event = AgentStreamEvent::TextDelta {
text: "hello".into(),
};
let chunk = event.to_chunk();
assert!(chunk.is_some());
match chunk.unwrap() {
StreamChunk::TextDelta(t) => assert_eq!(t, "hello"),
other => panic!("expected TextDelta, got {:?}", other),
}
}
#[test]
fn test_to_chunk_thinking_delta() {
let event = AgentStreamEvent::ThinkingDelta {
thinking: "think".into(),
redacted: Some("secret".into()),
};
let chunk = event.to_chunk();
assert!(chunk.is_some());
match chunk.unwrap() {
StreamChunk::ThinkingDelta { text, redacted } => {
assert_eq!(text, "think");
assert_eq!(redacted, Some("secret".into()));
}
other => panic!("expected ThinkingDelta, got {:?}", other),
}
}
#[test]
fn test_to_chunk_usage_none() {
let event = AgentStreamEvent::Usage(TokenUsage::default());
assert!(event.to_chunk().is_none());
}
#[test]
fn test_to_chunk_response_complete_none() {
let event = AgentStreamEvent::ResponseComplete;
assert!(event.to_chunk().is_none());
}
#[test]
fn test_translate_provider_event_compat() {
let event = ProviderEvent::Token {
token: "hello".into(),
};
match translate_provider_event(&event) {
TranslationResult::EmitWithText {
chunk: StreamChunk::TextDelta(t),
delta,
} => {
assert_eq!(t, "hello");
assert_eq!(delta, "hello");
}
other => panic!("expected EmitWithText(TextDelta), got {:?}", other),
}
}
#[test]
fn test_ignore_protocol_events() {
let event = ProviderEvent::Start {
model: "test".into(),
};
assert!(matches!(
translate_provider_event(&event),
TranslationResult::Ignore
));
}
}