1 files changed, 16 insertions, 2 deletions
diff --git a/src/config/conversation.rs b/src/config/conversation.rs
index 4472781..7fe843d 100644
--- a/src/config/conversation.rs
+++ b/src/config/conversation.rs
@@ -2,7 +2,7 @@ use super::message::{num_tokens_from_messages, Message, MessageRole};
 use super::role::Role;
 use super::MAX_TOKENS;
 
-use anyhow::Result;
+use anyhow::{bail, Result};
 use serde::{Deserialize, Serialize};
 
 #[derive(Debug, Clone, Deserialize, Serialize)]
@@ -19,10 +19,24 @@ impl Conversation {
             role,
             messages: vec![],
         };
-        value.tokens = num_tokens_from_messages(&value.build_emssages(""));
+        value.update_tokens();
         value
     }
 
+    pub fn update_role(&mut self, role: &Role) -> Result<()> {
+        if self.messages.is_empty() {
+            self.role = Some(role.clone());
+            self.update_tokens();
+        } else {
+            bail!("Error: Cannot perform this action in the middle of conversation")
+        }
+        Ok(())
+    }
+
+    pub fn update_tokens(&mut self) {
+        self.tokens = num_tokens_from_messages(&self.build_emssages(""));
+    }
+
     pub fn add_message(&mut self, input: &str, output: &str) -> Result<()> {
         let mut need_add_msg = true;
         if self.messages.is_empty() {