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 separator chunk (newline) emitted between tool-call rounds.
40    pub fn separator() -> Self {
41        Self {
42            choices: vec![Choice {
43                delta: Delta {
44                    content: Some("\n".into()),
45                    ..Default::default()
46                },
47                ..Default::default()
48            }],
49            ..Default::default()
50        }
51    }
52
53    /// Get the content of the first choice
54    pub fn content(&self) -> Option<&str> {
55        self.choices
56            .first()
57            .and_then(|c| c.delta.content.as_deref())
58            .filter(|s| !s.is_empty())
59    }
60
61    /// Get the reasoning content of the first choice
62    pub fn reasoning_content(&self) -> Option<&str> {
63        self.choices
64            .first()
65            .and_then(|c| c.delta.reasoning_content.as_deref())
66            .filter(|s| !s.is_empty())
67    }
68
69    /// Get the tool calls of the first choice
70    pub fn tool_calls(&self) -> Option<&[ToolCall]> {
71        self.choices
72            .first()
73            .and_then(|choice| choice.delta.tool_calls.as_deref())
74    }
75
76    /// Get the reason the model stopped generating
77    pub fn reason(&self) -> Option<&FinishReason> {
78        self.choices
79            .first()
80            .and_then(|choice| choice.finish_reason.as_ref())
81    }
82}