use serde::{Deserialize, Serialize};
use uuid::Uuid;
use super::{InputMessage, StreamChunkKind, StreamSegment};
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct LlmStream {
#[serde(default)]
pub(crate) segments: Vec<StreamSegment>,
}
fn normalise_source(source: &str) -> Option<String> {
let trimmed = source.trim();
(!trimmed.is_empty()).then(|| trimmed.to_string())
}
fn normalise_detail(detail: Option<&str>) -> Option<String> {
detail
.map(str::trim)
.filter(|d| !d.is_empty())
.map(str::to_string)
}
impl LlmStream {
pub fn empty() -> Self {
Self {
segments: Vec::new(),
}
}
pub fn prepend_input(&mut self, messages: Vec<InputMessage>) {
if messages.is_empty() {
return;
}
self.segments.insert(
0,
StreamSegment::Input {
messages,
message_id: None,
ts: None,
},
);
}
pub fn push_notice(&mut self, source: &str, detail: Option<&str>, ts: Option<u64>) -> bool {
let Some(source) = normalise_source(source) else {
return false;
};
self.segments.push(StreamSegment::PromptNotice {
source,
detail: normalise_detail(detail),
ts,
});
true
}
pub fn push_notice_after_input(
&mut self,
source: &str,
detail: Option<&str>,
ts: Option<u64>,
) -> bool {
let Some(source) = normalise_source(source) else {
return false;
};
let pos = usize::from(
self.segments
.first()
.is_some_and(|seg| matches!(seg, StreamSegment::Input { .. })),
);
self.segments.insert(
pos,
StreamSegment::PromptNotice {
source,
detail: normalise_detail(detail),
ts,
},
);
true
}
pub fn from_plain_text(text: &str) -> Self {
let segments = if text.is_empty() {
Vec::new()
} else {
vec![StreamSegment::Text {
text: text.to_string(),
message_id: None,
ts: None,
}]
};
Self { segments }
}
pub fn raw_text(&self) -> String {
self.segments.iter().map(|s| s.text_or_json()).collect()
}
pub fn segments(&self) -> &[StreamSegment] {
&self.segments
}
pub fn segments_mut(&mut self) -> &mut [StreamSegment] {
&mut self.segments
}
pub fn patch_tool_result_for_call(
&mut self,
call_id: Uuid,
result_data: serde_json::Value,
success: bool,
duration_ms: Option<u64>,
) -> bool {
let sealed_marker = "[stream sealed before result arrived]";
let sealed_value = serde_json::Value::String(sealed_marker.to_string());
for seg in &mut self.segments {
if let StreamSegment::ToolResult {
call_id: seg_call_id,
data,
success: seg_success,
duration_ms: seg_duration,
..
} = seg
&& *seg_call_id == call_id
&& *data == sealed_value
{
*data = result_data;
*seg_success = success;
*seg_duration = duration_ms;
return true;
}
}
let has_call = self
.segments
.iter()
.any(|s| matches!(s, StreamSegment::ToolCall { call_id: cid, .. } if *cid == call_id));
if has_call {
let tool_name = self
.segments
.iter()
.find_map(|s| match s {
StreamSegment::ToolCall {
call_id: cid,
tool_name,
..
} if *cid == call_id => Some(tool_name.clone()),
_ => None,
})
.unwrap_or_default();
let agent_type = self
.segments
.iter()
.find_map(|s| match s {
StreamSegment::ToolCall {
call_id: cid,
agent_type,
..
} if *cid == call_id => Some(agent_type.clone()),
_ => None,
})
.flatten();
let pos = self
.segments
.iter()
.rposition(
|s| matches!(s, StreamSegment::ToolCall { call_id: cid, .. } if *cid == call_id),
)
.map(|i| i + 1)
.unwrap_or(self.segments.len());
self.segments.insert(
pos,
StreamSegment::ToolResult {
tool_name,
call_id,
success,
data: result_data,
duration_ms,
agent_type,
message_id: None,
ts: None,
},
);
return true;
}
false
}
pub fn has_tool_call_with_id(&self, call_id: Uuid) -> bool {
self.segments
.iter()
.any(|s| matches!(s, StreamSegment::ToolCall { call_id: cid, .. } if *cid == call_id))
}
pub fn is_empty(&self) -> bool {
self.segments.is_empty()
|| self.segments.iter().all(|s| match s {
StreamSegment::Text { text, .. }
| StreamSegment::Thinking { text, .. }
| StreamSegment::DeepThinking { text, .. } => text.is_empty(),
StreamSegment::ToolCall { params, .. } => params.is_null(),
StreamSegment::ToolResult { data, .. } => data.is_null(),
StreamSegment::Input { .. } => false,
StreamSegment::PromptNotice { .. } => false,
})
}
pub fn text_segments(&self) -> impl Iterator<Item = (StreamChunkKind, &str)> {
self.segments.iter().filter_map(|seg| {
let kind = seg.chunk_kind()?;
Some((kind, seg.text()))
})
}
pub fn tool_segments(&self) -> impl Iterator<Item = &StreamSegment> {
self.segments.iter().filter(|s| s.is_tool())
}
pub fn thinking_text(&self) -> String {
self.text_segments()
.filter(|(kind, _)| {
matches!(
kind,
StreamChunkKind::Thinking | StreamChunkKind::DeepThinking
)
})
.map(|(_, text)| text)
.collect::<Vec<_>>()
.join("")
}
pub fn display_text(&self) -> String {
self.text_segments()
.filter(|(kind, _)| matches!(kind, StreamChunkKind::Text))
.map(|(_, text)| text)
.collect::<Vec<_>>()
.join("")
}
pub fn len(&self) -> usize {
self.segments.iter().map(|s| s.text_or_json().len()).sum()
}
}
impl std::fmt::Display for LlmStream {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
for seg in &self.segments {
f.write_str(&seg.text_or_json())?;
}
Ok(())
}
}
impl Default for LlmStream {
fn default() -> Self {
Self::empty()
}
}
impl From<String> for LlmStream {
fn from(s: String) -> Self {
Self::from_plain_text(&s)
}
}
impl From<&str> for LlmStream {
fn from(s: &str) -> Self {
Self::from_plain_text(s)
}
}
impl PartialEq for LlmStream {
fn eq(&self, other: &Self) -> bool {
self.segments == other.segments
}
}
impl Eq for LlmStream {}