llm/testing/
chat_message.rs

1use crate::types::IsoString;
2use crate::{
3    AssistantReply, ChatMessage, LlmModel, LlmResponse, MessageId, ToolCallError, ToolCallRequest, ToolCallResult,
4};
5
6pub fn assistant_message() -> AssistantMessageBuilder {
7    AssistantMessageBuilder { message_id: MessageId::new(), model: None, reply: AssistantReply::default() }
8}
9
10pub fn tool_result(id: &str, name: &str, arguments: &str, result: &str) -> ChatMessage {
11    ChatMessage::ToolCallResult(Ok(ToolCallResult {
12        id: id.to_string(),
13        name: name.to_string(),
14        arguments: arguments.to_string(),
15        result: result.to_string(),
16    }))
17}
18
19pub fn tool_error(id: &str, name: &str, error: &str) -> ChatMessage {
20    ChatMessage::ToolCallResult(Err(ToolCallError {
21        id: id.to_string(),
22        name: name.to_string(),
23        arguments: None,
24        error: error.to_string(),
25    }))
26}
27
28#[derive(Clone)]
29pub struct AssistantMessageBuilder {
30    message_id: MessageId,
31    model: Option<LlmModel>,
32    reply: AssistantReply,
33}
34
35impl AssistantMessageBuilder {
36    pub fn id(mut self, message_id: &str) -> Self {
37        self.message_id = message_id.into();
38        self
39    }
40
41    pub fn model(mut self, model: &str) -> Self {
42        self.model = Some(model.parse().unwrap_or_else(|error| panic!("invalid model '{model}': {error}")));
43        self
44    }
45
46    pub fn text(self, text: &str) -> Self {
47        self.push(&LlmResponse::text(text))
48    }
49
50    pub fn tool_call(self, id: &str, name: &str, arguments: &str) -> Self {
51        self.push(&LlmResponse::ToolRequestComplete {
52            tool_call: ToolCallRequest { id: id.to_string(), name: name.to_string(), arguments: arguments.to_string() },
53        })
54    }
55
56    pub fn reasoning(self, summary: &str) -> Self {
57        self.push(&LlmResponse::reasoning(summary))
58    }
59
60    pub fn encrypted_reasoning(self, id: &str, content: &str) -> Self {
61        self.push(&LlmResponse::encrypted_reasoning(id, content))
62    }
63
64    pub fn build(self) -> ChatMessage {
65        ChatMessage::Assistant {
66            message_id: self.message_id,
67            model: self.model,
68            content: self.reply,
69            timestamp: IsoString::now(),
70        }
71    }
72
73    pub fn build_reply(self) -> AssistantReply {
74        self.reply
75    }
76
77    fn push(mut self, response: &LlmResponse) -> Self {
78        self.reply.push(response);
79        self
80    }
81}
llm/testing/chat_message.rs

llm/testing/
chat_message.rs
