1 files changed, 8 insertions, 2 deletions
diff --git a/src/config/role.rs b/src/config/role.rs
index 515810c..1b49bad 100644
--- a/src/config/role.rs
+++ b/src/config/role.rs
@@ -1,7 +1,7 @@
-use super::Input;
+use super::{Config, Input};
 
 use crate::{
-    client::{Message, MessageContent, MessageRole, Model},
+    client::{list_chat_models, Message, MessageContent, MessageRole, Model},
     utils::{detect_os, detect_shell},
 };
 
@@ -99,6 +99,12 @@ async function timeout(ms) {
         self.prompt.contains(INPUT_PLACEHOLDER)
     }
 
+    pub fn retrieve_model(&self, config: &Config) -> Option<Model> {
+        self.model_id
+            .as_ref()
+            .and_then(|model_id| Model::find(&list_chat_models(config), model_id))
+    }
+
     pub fn set_model(&mut self, model: &Model) {
         self.model_id = Some(model.id());
     }