Skip to main content

walrus_core/model/
stream.rs

1//! Streaming response abstractions for the unified LLM Interfaces
2
3use crate::model::{
4    FinishReason,
5    response::{Choice, CompletionMeta, Delta},
6    tool::ToolCall,
7};
8use serde::Deserialize;
9
10/// A streaming chat completion chunk
11#[derive(Debug, Clone, Deserialize, Default)]
12pub struct StreamChunk {
13    /// Completion metadata
14    #[serde(flatten)]
15    pub meta: CompletionMeta,
16
17    /// The list of completion choices (with delta content)
18    pub choices: Vec<Choice>,
19
20    /// Token usage statistics (only in final chunk)
21    pub usage: Option<crate::model::Usage>,
22}
23
24impl StreamChunk {
25    /// Create a new tool chunk
26    pub fn tool(calls: &[ToolCall]) -> Self {
27        Self {
28            choices: vec![Choice {
29                delta: Delta {
30                    tool_calls: Some(calls.to_vec()),
31                    ..Default::default()
32                },
33                ..Default::default()
34            }],
35            ..Default::default()
36        }
37    }
38
39    /// Create a content-only chunk with the given text.
40    pub fn text(content: String) -> Self {
41        Self {
42            choices: vec![Choice {
43                delta: Delta {
44                    content: Some(content),
45                    ..Default::default()
46                },
47                ..Default::default()
48            }],
49            ..Default::default()
50        }
51    }
52
53    /// Create a separator chunk (newline) emitted between tool-call rounds.
54    pub fn separator() -> Self {
55        Self {
56            choices: vec![Choice {
57                delta: Delta {
58                    content: Some("\n".into()),
59                    ..Default::default()
60                },
61                ..Default::default()
62            }],
63            ..Default::default()
64        }
65    }
66
67    /// Get the content of the first choice
68    pub fn content(&self) -> Option<&str> {
69        self.choices
70            .first()
71            .and_then(|c| c.delta.content.as_deref())
72            .filter(|s| !s.is_empty())
73    }
74
75    /// Get the reasoning content of the first choice
76    pub fn reasoning_content(&self) -> Option<&str> {
77        self.choices
78            .first()
79            .and_then(|c| c.delta.reasoning_content.as_deref())
80            .filter(|s| !s.is_empty())
81    }
82
83    /// Get the tool calls of the first choice
84    pub fn tool_calls(&self) -> Option<&[ToolCall]> {
85        self.choices
86            .first()
87            .and_then(|choice| choice.delta.tool_calls.as_deref())
88    }
89
90    /// Get the reason the model stopped generating
91    pub fn reason(&self) -> Option<&FinishReason> {
92        self.choices
93            .first()
94            .and_then(|choice| choice.finish_reason.as_ref())
95    }
96}