Skip to main content

cllient/
chat.rs

1use crate::types::{RequestBuilder, CompletionResponse, MessageContent, ContentBlock};
2use crate::error::Result;
3use crate::client::LowLevelClient;
4
5/// Identifier for different types of chatters in a conversation
6#[derive(Debug, Clone, PartialEq, Eq)]
7pub enum ChatterId {
8    /// The end user of the application
9    User,
10    /// System instructions/prompts
11    System,
12    /// The current user (alias for User, more personal)
13    Me,
14    /// An AI agent/assistant
15    Agent,
16    /// A custom chatter with a specific name
17    Custom(String),
18    /// Multiple chatters (for group conversations)
19    Multiple(Vec<ChatterId>),
20}
21
22impl ChatterId {
23    /// Convert ChatterId to the role string used in the API
24    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                // For multiple IDs, use the first one or default to "user"
32                ids.first()
33                    .map(|id| id.to_role())
34                    .unwrap_or_else(|| "user".to_string())
35            }
36        }
37    }
38    
39    /// Create a custom chatter with a specific name
40    pub fn custom(name: &str) -> Self {
41        ChatterId::Custom(name.to_string())
42    }
43}
44
45/// A single message in a conversation
46#[derive(Debug, Clone)]
47pub struct ChatMessage {
48    pub chatter_id: ChatterId,
49    pub content: MessageContent,
50}
51
52/// Builder for managing conversations with multiple participants
53pub struct ChatBuilder {
54    conversation: Vec<ChatMessage>,
55    request_builder: RequestBuilder,
56    current_chatter: ChatterId,
57}
58
59impl ChatBuilder {
60    /// Create a new chat builder
61    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    /// Add a text message to the conversation
70    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    /// Add a multimodal message to the conversation
83    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    /// Set the system prompt for the conversation
96    pub fn system(mut self, prompt: &str) -> Self {
97        self.request_builder.request.system_prompt = Some(prompt.to_string());
98        self
99    }
100    
101    /// Send a message from the specified chatter and get a response
102    pub async fn send(mut self, chatter_id: ChatterId, content: &str) -> Result<CompletionResponse> {
103        // Add the new message to the conversation
104        self = self.add_message(chatter_id, content);
105        
106        // Convert conversation to the format expected by the API
107        let messages: Vec<MessageContent> = self.conversation.iter()
108            .map(|msg| msg.content.clone())
109            .collect();
110        
111        // Update the request with the full conversation
112        self.request_builder.request.messages = messages;
113        
114        // Send the request
115        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        // Note: The assistant's response is not added to the conversation here
123        // because `self` is consumed. In practice, the user would need to create
124        // a new ChatBuilder or we'd need to redesign to return an updated builder.
125
126        Ok(response)
127    }
128    
129    /// Send a multimodal message and get a response
130    pub async fn send_multimodal(mut self, chatter_id: ChatterId, content_blocks: Vec<ContentBlock>) -> Result<CompletionResponse> {
131        // Add the new multimodal message to the conversation
132        self = self.add_multimodal_message(chatter_id, content_blocks);
133        
134        // Convert conversation to the format expected by the API
135        let messages: Vec<MessageContent> = self.conversation.iter()
136            .map(|msg| msg.content.clone())
137            .collect();
138        
139        // Update the request with the full conversation
140        self.request_builder.request.messages = messages;
141        
142        // Send the request
143        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    /// Continue the conversation with the current chatter
152    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    /// Get the current conversation history
158    pub fn conversation(&self) -> &[ChatMessage] {
159        &self.conversation
160    }
161    
162    /// Clear the conversation history
163    pub fn clear_history(mut self) -> Self {
164        self.conversation.clear();
165        self
166    }
167    
168    /// Get the number of messages in the conversation
169    pub fn message_count(&self) -> usize {
170        self.conversation.len()
171    }
172    
173    /// Set parameters for the underlying request
174    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}