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

llm/testing/
llm_response.rs
