implement server switch

serverlessnext · Jul 4, 2024 · 6b7e6b9 · 6b7e6b9
1 parent bfb4312
commit 6b7e6b9
Show file tree

Hide file tree

Showing 12 changed files with 79 additions and 49 deletions.
diff --git a/lumni/src/apps/builtin/llm/prompt/src/app.rs b/lumni/src/apps/builtin/llm/prompt/src/app.rs
@@ -343,24 +343,7 @@ pub async fn run_cli(
 
     // create new (un-initialized) server from requested server name
     let server = ModelServer::from_str(&server_name)?;
-
-    // get default model from server - if available
-    let default_model = match server.list_models().await {
-        Ok(models) => {
-            if models.is_empty() {
-                log::warn!("Received empty model list");
-                None
-            } else {
-                log::debug!("Available models: {:?}", models);
-                Some(models[0].to_owned())
-            }
-        }
-        Err(e) => {
-            log::error!("Failed to list models: {}", e);
-            None
-        }
-    };
-
+    let default_model = server.get_default_model().await;
     // setup prompt, server and chat session
     let prompt_instruction =
         PromptInstruction::new(instruction, assistant, options)?;

diff --git a/lumni/src/apps/builtin/llm/prompt/src/chat/mod.rs b/lumni/src/apps/builtin/llm/prompt/src/chat/mod.rs
@@ -19,7 +19,9 @@ pub use session::ChatSession;
 
 pub use super::defaults::*;
 pub use super::model::PromptRole;
-pub use super::server::{LLMDefinition, ServerManager};
+pub use super::server::{
+    LLMDefinition, ModelServer, ServerManager, ServerTrait,
+};
 
 // gets PERSONAS from the generated code
 include!(concat!(env!("OUT_DIR"), "/llm/prompt/templates.rs"));

diff --git a/lumni/src/apps/builtin/llm/prompt/src/chat/session.rs b/lumni/src/apps/builtin/llm/prompt/src/chat/session.rs
@@ -6,7 +6,9 @@ use tokio::sync::{mpsc, oneshot, Mutex};
 
 use super::exchange::ChatExchange;
 use super::history::ChatHistory;
-use super::{LLMDefinition, PromptInstruction, ServerManager};
+use super::{
+    LLMDefinition, ModelServer, PromptInstruction, ServerManager, ServerTrait,
+};
 use crate::api::error::ApplicationError;
 
 pub struct ChatSession {
@@ -40,9 +42,25 @@ impl ChatSession {
         self.server.server_name()
     }
 
-    pub fn select_endpoint(&mut self, endpoint: &str) {
-        //self.server.select_endpoint(endpoint);
-        log::debug!("Selected endpoint: {}", endpoint);
+    pub async fn select_server(
+        &mut self,
+        server_name: &str,
+    ) -> Result<(), ApplicationError> {
+
+        if self.server_name() != server_name {
+            log::debug!("switching server: {}", server_name);
+            self.stop();
+
+            let mut server = ModelServer::from_str(server_name)?;
+            let model = server.get_default_model().await;
+            if let Some(model) = model {
+                server
+                    .setup_and_initialize(model, &mut self.prompt_instruction)
+                    .await?;
+            }
+            self.server = Box::new(server);
+        }
+        Ok(())
     }
 
     pub fn stop(&mut self) {

diff --git a/lumni/src/apps/builtin/llm/prompt/src/server/bedrock/mod.rs b/lumni/src/apps/builtin/llm/prompt/src/server/bedrock/mod.rs
@@ -19,7 +19,7 @@ use url::Url;
 
 use super::{
     http_post, ChatExchange, ChatHistory, ChatMessage, Endpoints,
-    LLMDefinition, PromptInstruction, ServerTrait, ServerSpecTrait,
+    LLMDefinition, PromptInstruction, ServerSpecTrait, ServerTrait,
 };
 pub use crate::external as lumni;
 

diff --git a/lumni/src/apps/builtin/llm/prompt/src/server/llama/mod.rs b/lumni/src/apps/builtin/llm/prompt/src/server/llama/mod.rs
@@ -12,8 +12,8 @@ use url::Url;
 use super::{
     http_get_with_response, http_post, ChatCompletionOptions, ChatExchange,
     ChatHistory, Endpoints, HttpClient, LLMDefinition, PromptInstruction,
-    PromptRole, ServerTrait, TokenResponse, DEFAULT_CONTEXT_SIZE,
-    ServerSpecTrait,
+    PromptRole, ServerSpecTrait, ServerTrait, TokenResponse,
+    DEFAULT_CONTEXT_SIZE,
 };
 use crate::external as lumni;
 

diff --git a/lumni/src/apps/builtin/llm/prompt/src/server/mod.rs b/lumni/src/apps/builtin/llm/prompt/src/server/mod.rs
@@ -19,14 +19,14 @@ use lumni::api::error::ApplicationError;
 pub use lumni::HttpClient;
 pub use ollama::Ollama;
 pub use openai::OpenAI;
+pub use spec::ServerSpecTrait;
 use tokio::sync::{mpsc, oneshot};
 
 pub use super::chat::{
     http_get_with_response, http_post, http_post_with_response,
     ChatCompletionOptions, ChatExchange, ChatHistory, ChatMessage,
     PromptInstruction, TokenResponse,
 };
-pub use spec::ServerSpecTrait;
 pub use super::defaults::*;
 pub use super::model::{ModelFormatter, ModelFormatterTrait, PromptRole};
 use crate::external as lumni;
@@ -247,6 +247,24 @@ pub trait ServerTrait: Send + Sync {
         }
     }
 
+    async fn get_default_model(&self) -> Option<LLMDefinition> {
+        match self.list_models().await {
+            Ok(models) => {
+                if models.is_empty() {
+                    log::warn!("Received empty model list");
+                    None
+                } else {
+                    log::debug!("Available models: {:?}", models);
+                    Some(models[0].to_owned())
+                }
+            }
+            Err(e) => {
+                log::error!("Failed to list models: {}", e);
+                None
+            }
+        }
+    }
+
     fn process_response(
         &mut self,
         response: Bytes,
@@ -288,7 +306,7 @@ pub trait ServerTrait: Send + Sync {
 }
 
 #[async_trait]
-pub trait ServerManager: ServerTrait  {
+pub trait ServerManager: ServerTrait {
     async fn setup_and_initialize(
         &mut self,
         model: LLMDefinition,

diff --git a/lumni/src/apps/builtin/llm/prompt/src/server/ollama/mod.rs b/lumni/src/apps/builtin/llm/prompt/src/server/ollama/mod.rs
@@ -10,7 +10,7 @@ use url::Url;
 use super::{
     http_get_with_response, http_post, http_post_with_response, ChatExchange,
     ChatHistory, ChatMessage, Endpoints, HttpClient, LLMDefinition,
-    PromptInstruction, ServerTrait, ServerSpecTrait,
+    PromptInstruction, ServerSpecTrait, ServerTrait,
 };
 use crate::external as lumni;
 

diff --git a/lumni/src/apps/builtin/llm/prompt/src/server/openai/mod.rs b/lumni/src/apps/builtin/llm/prompt/src/server/openai/mod.rs
@@ -20,13 +20,14 @@ use url::Url;
 
 use super::{
     http_post, ChatExchange, ChatHistory, ChatMessage, Endpoints,
-    LLMDefinition, PromptInstruction, ServerTrait, ServerSpecTrait,
+    LLMDefinition, PromptInstruction, ServerSpecTrait, ServerTrait,
 };
 pub use crate::external as lumni;
 
-const OPENAI_COMPLETION_ENDPOINT: &str = "https://api.openai.com/v1/chat/completions";
+const OPENAI_COMPLETION_ENDPOINT: &str =
+    "https://api.openai.com/v1/chat/completions";
 
-define_and_impl_server_spec!(OpenAISpec);   //, "OpenAI");
+define_and_impl_server_spec!(OpenAISpec); //, "OpenAI");
 
 pub struct OpenAI {
     spec: OpenAISpec,

diff --git a/lumni/src/apps/builtin/llm/prompt/src/server/spec.rs b/lumni/src/apps/builtin/llm/prompt/src/server/spec.rs
@@ -1,4 +1,4 @@
-
+// add debug
 pub trait ServerSpecTrait {
     fn name(&self) -> &str;
 }

diff --git a/lumni/src/apps/builtin/llm/prompt/src/tui/events/key_event.rs b/lumni/src/apps/builtin/llm/prompt/src/tui/events/key_event.rs
@@ -199,20 +199,25 @@ impl KeyEventHandler {
                     Some(WindowEvent::PromptWindow)
                 } else {
                     if let Some(modal) = tab_ui.modal.as_mut() {
-                        let new_window_event = match modal.handle_key_event(&mut self.key_track, tab_chat) {
+                        let new_window_event = match modal
+                            .handle_key_event(&mut self.key_track, tab_chat)
+                            .await
+                        {
                             Some(WindowEvent::Modal(next_window_type)) => {
                                 if next_window_type == window_type {
                                     // window remains un-changed
-                                    return Some(WindowEvent::Modal(window_type));
+                                    return Some(WindowEvent::Modal(
+                                        window_type,
+                                    ));
                                 }
                                 WindowEvent::Modal(next_window_type)
-                            },
+                            }
                             Some(new_window_event) => new_window_event,
-                            None => WindowEvent::PromptWindow,    // default
+                            None => WindowEvent::PromptWindow, // default
                         };
                         // window change
                         // close existing modal window
-                        tab_ui.clear_modal();   
+                        tab_ui.clear_modal();
                         return Some(new_window_event);
                     } else {
                         Some(WindowEvent::Modal(window_type))

diff --git a/lumni/src/apps/builtin/llm/prompt/src/tui/mod.rs b/lumni/src/apps/builtin/llm/prompt/src/tui/mod.rs
@@ -18,6 +18,6 @@ pub use modal::{ModalConfigWindow, ModalWindowTrait, ModalWindowType};
 pub use ui::TabUi;
 pub use windows::{CommandLine, PromptWindow, ResponseWindow};
 
+pub use super::chat::ChatSession;
 pub use super::server::SUPPORTED_MODEL_ENDPOINTS;
 pub use super::session::TabSession;
-pub use super::chat::ChatSession;
diff --git a/lumni/src/apps/builtin/llm/prompt/src/tui/modal.rs b/lumni/src/apps/builtin/llm/prompt/src/tui/modal.rs
@@ -1,3 +1,4 @@
+use async_trait::async_trait;
 use crossterm::event::KeyCode;
 use ratatui::layout::Rect;
 use ratatui::widgets::Clear;
@@ -6,20 +7,21 @@ use ratatui::Frame;
 use super::components::Scroller;
 use super::events::KeyTrack;
 use super::widgets::SelectEndpoint;
-use super::{WindowEvent, ChatSession};
+use super::{ChatSession, WindowEvent};
 
 #[derive(Debug, Clone, Copy, PartialEq)]
 pub enum ModalWindowType {
     Config,
 }
 
+#[async_trait]
 pub trait ModalWindowTrait {
     fn get_type(&self) -> ModalWindowType;
     fn render_on_frame(&mut self, frame: &mut Frame, area: Rect);
-    fn handle_key_event(
-        &mut self,
-        key_event: &mut KeyTrack,
-        tab_chat: &mut ChatSession,
+    async fn handle_key_event<'a>(
+        &'a mut self,
+        key_event: &'a mut KeyTrack,
+        tab_chat: &'a mut ChatSession,
     ) -> Option<WindowEvent>;
 }
 
@@ -37,6 +39,7 @@ impl ModalConfigWindow {
     }
 }
 
+#[async_trait]
 impl ModalWindowTrait for ModalConfigWindow {
     fn get_type(&self) -> ModalWindowType {
         ModalWindowType::Config
@@ -55,17 +58,17 @@ impl ModalWindowTrait for ModalConfigWindow {
         frame.render_widget(&mut self.widget, area);
     }
 
-    fn handle_key_event(
-        &mut self,
-        key_event: &mut KeyTrack,
-        tab_chat: &mut ChatSession,
+    async fn handle_key_event<'a>(
+        &'a mut self,
+        key_event: &'a mut KeyTrack,
+        tab_chat: &'a mut ChatSession,
     ) -> Option<WindowEvent> {
         match key_event.current_key().code {
             KeyCode::Up => self.widget.key_up(),
             KeyCode::Down => self.widget.key_down(),
             KeyCode::Enter => {
                 let endpoint = self.widget.current_endpoint();
-                tab_chat.select_endpoint(endpoint);
+                tab_chat.select_server(endpoint).await;
                 return Some(WindowEvent::PromptWindow);
             }
             _ => {} // Ignore other keys