llm/testing/
llm_response.rs

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

llm/testing/
llm_response.rs
