llm/testing/
llm_response.rs

1use crate::{LlmError, LlmResponse, StopReason};
2
3pub fn llm_response() -> LlmResponseBuilder {
4    LlmResponseBuilder::new()
5}
6
7/// A turn whose call fails before the provider emits any frames.
8pub fn failed_call(error: impl Into<LlmError>) -> Vec<Result<LlmResponse, LlmError>> {
9    vec![Err(error.into())]
10}
11
12pub struct LlmResponseBuilder {
13    chunks: Vec<LlmResponse>,
14}
15
16impl Default for LlmResponseBuilder {
17    fn default() -> Self {
18        Self::new()
19    }
20}
21
22impl LlmResponseBuilder {
23    pub fn new() -> Self {
24        Self { chunks: vec![LlmResponse::Start] }
25    }
26
27    pub fn text(mut self, chunks: &[&str]) -> Self {
28        for chunk in chunks {
29            self.chunks.push(LlmResponse::text(chunk));
30        }
31
32        self
33    }
34
35    pub fn reasoning(mut self, chunks: &[&str]) -> Self {
36        for chunk in chunks {
37            self.chunks.push(LlmResponse::reasoning(chunk));
38        }
39
40        self
41    }
42
43    pub fn encrypted_reasoning(mut self, id: &str, content: &str) -> Self {
44        self.chunks.push(LlmResponse::encrypted_reasoning(id, content));
45        self
46    }
47
48    pub fn tool_call(mut self, id: &str, name: &str, argument_chunks: &[&str]) -> Self {
49        self.chunks.push(LlmResponse::tool_request_start(id, name));
50
51        for chunk in argument_chunks {
52            self.chunks.push(LlmResponse::tool_request_arg(id, chunk));
53        }
54
55        self.chunks.push(LlmResponse::tool_request_complete(id, name, &argument_chunks.join("")));
56
57        self
58    }
59
60    pub fn usage(mut self, input_tokens: u64, output_tokens: u64) -> Self {
61        self.chunks.push(LlmResponse::usage(input_tokens, output_tokens));
62        self
63    }
64
65    pub fn tool_call_without_deltas(mut self, id: &str, name: &str, arguments: &str) -> Self {
66        self.chunks.push(LlmResponse::tool_request_start(id, name));
67        self.chunks.push(LlmResponse::tool_request_complete(id, name, arguments));
68
69        self
70    }
71
72    pub fn build(mut self) -> Vec<LlmResponse> {
73        self.chunks.push(LlmResponse::done());
74        self.chunks
75    }
76
77    pub fn build_with_stop_reason(mut self, stop_reason: StopReason) -> Vec<LlmResponse> {
78        self.chunks.push(LlmResponse::done_with_stop_reason(stop_reason));
79        self.chunks
80    }
81
82    pub fn build_results(self) -> Vec<Result<LlmResponse, LlmError>> {
83        self.build().into_iter().map(Ok).collect()
84    }
85
86    /// The stream surfaces `error` after the frames built so far, then closes
87    /// with `Done` — a provider that reports a failure before ending cleanly.
88    pub fn build_with_error(self, error: impl Into<LlmError>) -> Vec<Result<LlmResponse, LlmError>> {
89        let mut results = self.build_results();
90        results.insert(results.len() - 1, Err(error.into()));
91        results
92    }
93
94    /// The stream dies on `error` instead of delivering `Done` — a connection
95    /// lost mid-flight.
96    pub fn build_interrupted(self, error: impl Into<LlmError>) -> Vec<Result<LlmResponse, LlmError>> {
97        let mut results: Vec<_> = self.chunks.into_iter().map(Ok).collect();
98        results.push(Err(error.into()));
99        results
100    }
101}
102
103#[cfg(test)]
104mod tests {
105    use super::*;
106    use crate::ProviderError;
107
108    #[test]
109    fn build_with_stop_reason_preserves_response_chunks() {
110        let response = llm_response().text(&["hello"]).usage(10, 2).build_with_stop_reason(StopReason::Length);
111
112        assert!(matches!(
113            response.as_slice(),
114            [
115                LlmResponse::Start,
116                LlmResponse::Text { .. },
117                LlmResponse::Usage { .. },
118                LlmResponse::Done { stop_reason: Some(StopReason::Length) },
119            ]
120        ));
121    }
122
123    #[test]
124    fn reasoning_appends_reasoning_frames() {
125        let frames = llm_response().reasoning(&["thinking", "harder"]).text(&["answer"]).build();
126
127        assert!(matches!(
128            frames.as_slice(),
129            [
130                LlmResponse::Start,
131                LlmResponse::Reasoning { .. },
132                LlmResponse::Reasoning { .. },
133                LlmResponse::Text { .. },
134                LlmResponse::Done { .. },
135            ]
136        ));
137    }
138
139    #[test]
140    fn build_results_wraps_success_frames_in_ok() {
141        let results = llm_response().text(&["hi"]).build_results();
142
143        assert!(matches!(
144            results.as_slice(),
145            [Ok(LlmResponse::Start), Ok(LlmResponse::Text { .. }), Ok(LlmResponse::Done { .. })]
146        ));
147    }
148
149    #[test]
150    fn build_with_error_surfaces_error_before_done() {
151        let results = llm_response().usage(9, 1).build_with_error(ProviderError::api("HTTP 500"));
152
153        assert!(matches!(
154            results.as_slice(),
155            [Ok(LlmResponse::Start), Ok(LlmResponse::Usage { .. }), Err(_), Ok(LlmResponse::Done { .. }),]
156        ));
157    }
158
159    #[test]
160    fn build_interrupted_ends_with_error_and_no_done() {
161        let results = llm_response().text(&["partial"]).build_interrupted(ProviderError::stream_interrupted("boom"));
162
163        assert!(matches!(results.as_slice(), [Ok(LlmResponse::Start), Ok(LlmResponse::Text { .. }), Err(_)]));
164    }
165
166    #[test]
167    fn failed_call_contains_only_the_error() {
168        let results = failed_call(ProviderError::server("boom"));
169
170        assert!(matches!(results.as_slice(), [Err(_)]));
171    }
172}
llm/testing/llm_response.rs

llm/testing/
llm_response.rs
