From c26594c3c85fb99ad5b0db1c535b337585d3c103 Mon Sep 17 00:00:00 2001
From: sigoden <sigoden@gmail.com>
Date: Thu, 6 Jun 2024 05:27:21 +0800
Subject: refactor: config refinement (#570)

---
 src/config/role.rs | 10 ++++++++--
 1 file changed, 8 insertions(+), 2 deletions(-)

(limited to 'src/config/role.rs')

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());
     }
-- 
cgit v1.2.3