llm/testing/
llm_response.rs1use crate::LlmResponse;
2
3pub fn llm_response(message_id: &str) -> LlmResponseBuilder {
4 LlmResponseBuilder::new(message_id)
5}
6
7pub struct LlmResponseBuilder {
8 chunks: Vec<LlmResponse>,
9}
10
11impl LlmResponseBuilder {
12 pub fn new(message_id: &str) -> Self {
13 Self { chunks: vec![LlmResponse::start(message_id)] }
14 }
15
16 pub fn text(mut self, chunks: &[&str]) -> Self {
17 for chunk in chunks {
18 self.chunks.push(LlmResponse::text(chunk));
19 }
20
21 self
22 }
23
24 pub fn tool_call(mut self, id: &str, name: &str, argument_chunks: &[&str]) -> Self {
25 self.chunks.push(LlmResponse::tool_request_start(id, name));
26
27 for chunk in argument_chunks {
28 self.chunks.push(LlmResponse::tool_request_arg(id, chunk));
29 }
30
31 self.chunks.push(LlmResponse::tool_request_complete(id, name, &argument_chunks.join("")));
32
33 self
34 }
35
36 pub fn tool_call_with_invalid_json(mut self, id: &str, name: &str) -> Self {
37 self.chunks.push(LlmResponse::tool_request_start(id, name));
38 self.chunks.push(LlmResponse::tool_request_complete(id, name, "invalid json"));
39
40 self
41 }
42
43 pub fn build(mut self) -> Vec<LlmResponse> {
44 self.chunks.push(LlmResponse::done());
45 self.chunks
46 }
47}