1use crate::types::{RequestBuilder, CompletionResponse, MessageContent, ContentBlock};
2use crate::error::Result;
3use crate::client::LowLevelClient;
4
5#[derive(Debug, Clone, PartialEq, Eq)]
7pub enum ChatterId {
8 User,
10 System,
12 Me,
14 Agent,
16 Custom(String),
18 Multiple(Vec<ChatterId>),
20}
21
22impl ChatterId {
23 pub fn to_role(&self) -> String {
25 match self {
26 ChatterId::User | ChatterId::Me => "user".to_string(),
27 ChatterId::System => "system".to_string(),
28 ChatterId::Agent => "assistant".to_string(),
29 ChatterId::Custom(name) => name.clone(),
30 ChatterId::Multiple(ids) => {
31 ids.first()
33 .map(|id| id.to_role())
34 .unwrap_or_else(|| "user".to_string())
35 }
36 }
37 }
38
39 pub fn custom(name: &str) -> Self {
41 ChatterId::Custom(name.to_string())
42 }
43}
44
45#[derive(Debug, Clone)]
47pub struct ChatMessage {
48 pub chatter_id: ChatterId,
49 pub content: MessageContent,
50}
51
52pub struct ChatBuilder {
54 conversation: Vec<ChatMessage>,
55 request_builder: RequestBuilder,
56 current_chatter: ChatterId,
57}
58
59impl ChatBuilder {
60 pub(crate) fn new(request_builder: RequestBuilder, initial_chatter: ChatterId) -> Self {
62 Self {
63 conversation: Vec::new(),
64 request_builder,
65 current_chatter: initial_chatter,
66 }
67 }
68
69 pub fn add_message(mut self, chatter_id: ChatterId, content: &str) -> Self {
71 let message = ChatMessage {
72 chatter_id: chatter_id.clone(),
73 content: MessageContent::Text {
74 role: chatter_id.to_role(),
75 content: content.to_string(),
76 },
77 };
78 self.conversation.push(message);
79 self
80 }
81
82 pub fn add_multimodal_message(mut self, chatter_id: ChatterId, content_blocks: Vec<ContentBlock>) -> Self {
84 let message = ChatMessage {
85 chatter_id: chatter_id.clone(),
86 content: MessageContent::Multimodal {
87 role: chatter_id.to_role(),
88 content: content_blocks,
89 },
90 };
91 self.conversation.push(message);
92 self
93 }
94
95 pub fn system(mut self, prompt: &str) -> Self {
97 self.request_builder.request.system_prompt = Some(prompt.to_string());
98 self
99 }
100
101 pub async fn send(mut self, chatter_id: ChatterId, content: &str) -> Result<CompletionResponse> {
103 self = self.add_message(chatter_id, content);
105
106 let messages: Vec<MessageContent> = self.conversation.iter()
108 .map(|msg| msg.content.clone())
109 .collect();
110
111 self.request_builder.request.messages = messages;
113
114 let client = crate::client::HttpClient::from_model_id(
116 &*self.request_builder.config_provider,
117 &self.request_builder.model_id
118 )?;
119
120 let response = client.complete(&self.request_builder.request).await?;
121
122 Ok(response)
127 }
128
129 pub async fn send_multimodal(mut self, chatter_id: ChatterId, content_blocks: Vec<ContentBlock>) -> Result<CompletionResponse> {
131 self = self.add_multimodal_message(chatter_id, content_blocks);
133
134 let messages: Vec<MessageContent> = self.conversation.iter()
136 .map(|msg| msg.content.clone())
137 .collect();
138
139 self.request_builder.request.messages = messages;
141
142 let client = crate::client::HttpClient::from_model_id(
144 &*self.request_builder.config_provider,
145 &self.request_builder.model_id
146 )?;
147
148 client.complete(&self.request_builder.request).await
149 }
150
151 pub async fn continue_conversation(self, content: &str) -> Result<CompletionResponse> {
153 let current_chatter = self.current_chatter.clone();
154 self.send(current_chatter, content).await
155 }
156
157 pub fn conversation(&self) -> &[ChatMessage] {
159 &self.conversation
160 }
161
162 pub fn clear_history(mut self) -> Self {
164 self.conversation.clear();
165 self
166 }
167
168 pub fn message_count(&self) -> usize {
170 self.conversation.len()
171 }
172
173 pub fn temperature(mut self, temp: f64) -> Self {
175 self.request_builder = self.request_builder.temperature(temp);
176 self
177 }
178
179 pub fn max_tokens(mut self, tokens: u32) -> Self {
180 self.request_builder = self.request_builder.max_tokens(tokens);
181 self
182 }
183
184 pub fn top_p(mut self, top_p: f64) -> Self {
185 self.request_builder = self.request_builder.top_p(top_p);
186 self
187 }
188}
189
190#[cfg(test)]
191mod tests {
192 use super::*;
193
194 #[test]
195 fn test_chatter_id_to_role() {
196 assert_eq!(ChatterId::User.to_role(), "user");
197 assert_eq!(ChatterId::Me.to_role(), "user");
198 assert_eq!(ChatterId::System.to_role(), "system");
199 assert_eq!(ChatterId::Agent.to_role(), "assistant");
200 assert_eq!(ChatterId::Custom("narrator".to_string()).to_role(), "narrator");
201
202 let custom = ChatterId::custom("teacher");
203 assert_eq!(custom.to_role(), "teacher");
204
205 let multiple = ChatterId::Multiple(vec![ChatterId::User, ChatterId::Agent]);
206 assert_eq!(multiple.to_role(), "user");
207 }
208
209 #[test]
210 fn test_chatter_id_equality() {
211 assert_eq!(ChatterId::User, ChatterId::User);
212 assert_eq!(ChatterId::Custom("alice".to_string()), ChatterId::Custom("alice".to_string()));
213 assert_ne!(ChatterId::User, ChatterId::Agent);
214 assert_ne!(ChatterId::Custom("alice".to_string()), ChatterId::Custom("bob".to_string()));
215 }
216}