rig_core/test_utils/
streaming.rs

1//! Streaming helpers for [`MockCompletionModel`](super::MockCompletionModel).
2
3use crate::{
4    completion::{CompletionError, Usage},
5    message::ReasoningContent,
6    streaming::{RawStreamingChoice, RawStreamingToolCall, StreamFinal, ToolCallDeltaContent},
7};
8
9/// Provider descriptor name reported by the test doubles.
10pub const MOCK_PROVIDER: &str = "mock";
11
12/// Build the terminal record the mock model yields, carrying `usage`.
13pub fn mock_final(usage: Usage) -> StreamFinal {
14    StreamFinal::new(MOCK_PROVIDER, usage)
15}
16
17/// Convert a fixture JSON value into canonical params: `null`/`{}` mean
18/// "none", any other non-object is a scripting mistake surfaced as a stream
19/// error.
20fn fixture_additional_params(
21    value: serde_json::Value,
22) -> Result<Option<crate::message::AdditionalParams>, CompletionError> {
23    crate::message::AdditionalParams::try_from_value(value).map_err(|other| {
24        CompletionError::ProviderError(format!(
25            "mock stream fixture `additional_params` must be a JSON object, got: {other}"
26        ))
27    })
28}
29
30/// Build a terminal record whose usage has only `total_tokens` set.
31pub fn mock_final_with_total_tokens(total_tokens: u64) -> StreamFinal {
32    let mut usage = Usage::new();
33    usage.total_tokens = total_tokens;
34    mock_final(usage)
35}
36
37/// Scripted streaming event yielded by [`MockCompletionModel`](super::MockCompletionModel).
38#[derive(Clone, Debug)]
39pub enum MockStreamEvent {
40    /// Text chunk.
41    Text(String),
42    /// Start a new text content block with optional provider metadata.
43    TextStart {
44        id: String,
45        additional_params: Option<serde_json::Value>,
46    },
47    /// Provider-specific metadata for the current text content block.
48    TextAdditionalParams(serde_json::Value),
49    /// Complete tool call event.
50    ToolCall {
51        id: String,
52        name: String,
53        arguments: serde_json::Value,
54        call_id: Option<String>,
55    },
56    /// Tool call delta event.
57    ToolCallDelta {
58        id: String,
59        content: ToolCallDeltaContent,
60    },
61    /// Complete reasoning event.
62    Reasoning {
63        id: String,
64        content: ReasoningContent,
65    },
66    /// Reasoning delta event.
67    ReasoningDelta { id: String, reasoning: String },
68    /// Provider-assigned message ID.
69    MessageId(String),
70    /// Provider-native output item that Rig does not model.
71    Unknown(serde_json::Value),
72    /// Final raw response carrying optional usage.
73    FinalResponse(StreamFinal),
74    /// Stream error.
75    Error(MockError),
76}
77
78use super::completion::MockError;
79
80/// Fixture-syntax decoding of a part identity.
81///
82/// Corpus fixtures are plain data and spell identities as strings; the
83/// legacy minted renderings (`reasoning-0`, `block-3`, `output-1`, `tool-2`,
84/// `text-0`) are the fixture syntax for a `StreamPartId::Minted` of that kind
85/// and index, and anything else is a wire id. This is *fixture encoding*,
86/// not provenance recovery: production code never parses an id string —
87/// provenance travels in [`StreamPartId`] itself.
88fn fixture_part_id(id: String) -> crate::streaming::StreamPartId {
89    use crate::streaming::MintKind;
90    for (namespace, kind) in [
91        ("reasoning-", MintKind::Reasoning),
92        ("block-", MintKind::Block),
93        ("output-", MintKind::Output),
94        ("tool-", MintKind::Tool),
95        ("text-", MintKind::Text),
96    ] {
97        if let Some(rest) = id.strip_prefix(namespace)
98            && let Ok(index) = rest.parse::<u64>()
99        {
100            return kind.for_wire_index(index);
101        }
102    }
103    crate::streaming::StreamPartId::wire(id)
104}
105
106impl MockStreamEvent {
107    /// Create a text chunk.
108    pub fn text(text: impl Into<String>) -> Self {
109        Self::Text(text.into())
110    }
111
112    /// Start a new text content block identified by `id`.
113    pub fn text_start(id: impl Into<String>, additional_params: Option<serde_json::Value>) -> Self {
114        Self::TextStart {
115            id: id.into(),
116            additional_params,
117        }
118    }
119
120    /// Add provider-specific metadata to the current text content block.
121    pub fn text_additional_params(additional_params: serde_json::Value) -> Self {
122        Self::TextAdditionalParams(additional_params)
123    }
124
125    /// Create a complete tool call event.
126    pub fn tool_call(
127        id: impl Into<String>,
128        name: impl Into<String>,
129        arguments: serde_json::Value,
130    ) -> Self {
131        Self::ToolCall {
132            id: id.into(),
133            name: name.into(),
134            arguments,
135            call_id: None,
136        }
137    }
138
139    /// Attach a provider-specific call ID to a complete tool call event.
140    pub fn with_call_id(mut self, call_id: impl Into<String>) -> Self {
141        if let Self::ToolCall { call_id: id, .. } = &mut self {
142            *id = Some(call_id.into());
143        }
144        self
145    }
146
147    /// Create a tool call name delta.
148    pub fn tool_call_name_delta(id: impl Into<String>, name: impl Into<String>) -> Self {
149        Self::ToolCallDelta {
150            id: id.into(),
151            content: ToolCallDeltaContent::Name(name.into()),
152        }
153    }
154
155    /// Create a tool call arguments delta.
156    pub fn tool_call_arguments_delta(id: impl Into<String>, arguments: impl Into<String>) -> Self {
157        Self::ToolCallDelta {
158            id: id.into(),
159            content: ToolCallDeltaContent::Delta(arguments.into()),
160        }
161    }
162
163    /// Create a complete reasoning event with the default mock id
164    /// (`"reasoning-0"`). Use [`Self::with_reasoning_id`] for tests that
165    /// need distinct reasoning items.
166    pub fn reasoning(reasoning: impl Into<String>) -> Self {
167        Self::Reasoning {
168            id: "reasoning-0".to_string(),
169            content: ReasoningContent::Text {
170                text: reasoning.into(),
171                signature: None,
172            },
173        }
174    }
175
176    /// Attach a provider-specific reasoning ID to a complete reasoning event.
177    pub fn with_reasoning_id(mut self, reasoning_id: impl Into<String>) -> Self {
178        if let Self::Reasoning { id, .. } = &mut self {
179            *id = reasoning_id.into();
180        }
181        self
182    }
183
184    /// Create a reasoning delta event with the default mock id
185    /// (`"reasoning-0"`). Use [`Self::reasoning_delta_with_id`] for tests
186    /// that need distinct reasoning items.
187    pub fn reasoning_delta(reasoning: impl Into<String>) -> Self {
188        Self::reasoning_delta_with_id("reasoning-0", reasoning)
189    }
190
191    /// Create a reasoning delta event with an explicit reasoning item id.
192    pub fn reasoning_delta_with_id(id: impl Into<String>, reasoning: impl Into<String>) -> Self {
193        Self::ReasoningDelta {
194            id: id.into(),
195            reasoning: reasoning.into(),
196        }
197    }
198
199    /// Create a provider-assigned message ID event.
200    pub fn message_id(id: impl Into<String>) -> Self {
201        Self::MessageId(id.into())
202    }
203
204    /// Create an unmodeled provider output item.
205    pub fn unknown(value: serde_json::Value) -> Self {
206        Self::Unknown(value)
207    }
208
209    /// Create a final response event with usage.
210    pub fn final_response(usage: Usage) -> Self {
211        Self::FinalResponse(mock_final(usage))
212    }
213
214    /// Create a final response event with default zero usage.
215    pub fn final_response_with_default_usage() -> Self {
216        Self::FinalResponse(mock_final(Usage::new()))
217    }
218
219    /// Create a final response event whose usage has only `total_tokens` set.
220    pub fn final_response_with_total_tokens(total_tokens: u64) -> Self {
221        Self::FinalResponse(mock_final_with_total_tokens(total_tokens))
222    }
223
224    /// Create a stream error event.
225    pub fn error(message: impl Into<String>) -> Self {
226        Self::Error(MockError::provider(message))
227    }
228
229    pub(crate) fn into_raw_choice(self) -> Result<RawStreamingChoice, CompletionError> {
230        match self {
231            Self::Text(text) => Ok(RawStreamingChoice::Message(text)),
232            Self::TextStart {
233                id,
234                additional_params,
235            } => Ok(RawStreamingChoice::TextStart {
236                id: fixture_part_id(id),
237                additional_params: additional_params
238                    .map(fixture_additional_params)
239                    .transpose()?
240                    .flatten(),
241            }),
242            Self::TextAdditionalParams(additional_params) => {
243                match fixture_additional_params(additional_params)? {
244                    // The real variant is non-empty by construction; an empty
245                    // fixture object is a scripting mistake, not a no-op.
246                    None => Err(CompletionError::ProviderError(
247                        "mock stream fixture `TextAdditionalParams` carries no data — \
248                         drop the event instead"
249                            .to_string(),
250                    )),
251                    Some(params) => Ok(RawStreamingChoice::TextAdditionalParams(params)),
252                }
253            }
254            Self::ToolCall {
255                id,
256                name,
257                arguments,
258                call_id,
259            } => {
260                let mut tool_call = RawStreamingToolCall::new(fixture_part_id(id), name, arguments);
261                if let Some(call_id) = call_id {
262                    tool_call = tool_call.with_call_id(call_id);
263                }
264                Ok(RawStreamingChoice::ToolCall(tool_call))
265            }
266            Self::ToolCallDelta { id, content } => Ok(RawStreamingChoice::ToolCallDelta {
267                id: fixture_part_id(id),
268                content,
269            }),
270            Self::Reasoning { id, content } => {
271                // Fixture syntax: a wire-shaped id is both the key and the
272                // durable handle; a legacy minted rendering is a key only.
273                let key = fixture_part_id(id.clone());
274                let provider_id = match &key {
275                    key_is_wire if key_is_wire.wire_str().is_some() => {
276                        crate::streaming::WireId::new(id)
277                    }
278                    _ => None,
279                };
280                Ok(RawStreamingChoice::Reasoning {
281                    id: key,
282                    provider_id,
283                    content,
284                })
285            }
286            Self::ReasoningDelta { id, reasoning } => {
287                let key = fixture_part_id(id.clone());
288                let provider_id = match &key {
289                    key_is_wire if key_is_wire.wire_str().is_some() => {
290                        crate::streaming::WireId::new(id)
291                    }
292                    _ => None,
293                };
294                Ok(RawStreamingChoice::ReasoningDelta {
295                    id: key,
296                    provider_id,
297                    reasoning,
298                })
299            }
300            Self::MessageId(id) => Ok(RawStreamingChoice::MessageId(id)),
301            Self::Unknown(value) => Ok(RawStreamingChoice::Unknown(value.into())),
302            Self::FinalResponse(response) => Ok(RawStreamingChoice::FinalResponse(response)),
303            Self::Error(error) => Err(error.into_completion_error()),
304        }
305    }
306}
rig_core/test_utils/streaming.rs

rig_core/test_utils/
streaming.rs
