Skip to content

Commit

Permalink
ensure profile creation is done in background process, improve ui
Browse files Browse the repository at this point in the history
  • Loading branch information
@aprxi
aprxi committed Aug 31, 2024
1 parent a0be871 commit a1f0a59
Show file tree
Hide file tree
Showing 16 changed files with 450 additions and 171 deletions.
37 changes: 29 additions & 8 deletions lumni/src/apps/builtin/llm/prompt/src/app.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -99,11 +99,23 @@ pub async fn run_cli(
}
}

let prompt_instruction = PromptInstructionBuilder::new(db_conn.clone())
.with_matches(matches.as_ref().expect("Clap matches not found"))
.await?
.build()
.await?;
let prompt_instruction =
match PromptInstructionBuilder::new(db_conn.clone())
.with_matches(matches.as_ref().expect("Clap matches not found"))
.await
{
Ok(builder) => match builder.build().await {
Ok(instruction) => Some(instruction),
Err(e) => {
log::warn!("{}", e);
None
}
},
Err(e) => {
log::warn!("{}", e);
None
}
};

match input {
Some(question) => {
Expand Down Expand Up @@ -134,7 +146,7 @@ pub async fn run_cli(
}

async fn interactive_mode(
prompt_instruction: PromptInstruction,
prompt_instruction: Option<PromptInstruction>,
db_conn: Arc<ConversationDatabase>,
) -> Result<(), ApplicationError> {
let app = App::new(prompt_instruction, Arc::clone(&db_conn)).await?;
Expand Down Expand Up @@ -180,12 +192,21 @@ async fn interactive_mode(
}

async fn process_non_interactive_input(
prompt_instruction: PromptInstruction,
prompt_instruction: Option<PromptInstruction>,
db_conn: Arc<ConversationDatabase>,
question: Option<String>,
) -> Result<(), ApplicationError> {
let instruction = match prompt_instruction {
Some(instruction) => instruction,
None => {
return Err(ApplicationError::InvalidInput(
"No prompt instruction provided".to_string(),
));
}
};

let chat = Arc::new(Mutex::new(ThreadedChatSession::new(
prompt_instruction,
instruction,
db_conn.clone(),
)));

Expand Down
12 changes: 6 additions & 6 deletions lumni/src/apps/builtin/llm/prompt/src/chat/db/user_profile/mod.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -92,14 +92,13 @@ impl UserProfileDbHandler {
.get_profile_settings(&profile, MaskMode::Unmask)
.await?;

let model_server = settings
let model_server = match settings
.get("__TEMPLATE.__MODEL_SERVER")
.and_then(|v| v.as_str())
.ok_or_else(|| {
ApplicationError::InvalidInput(
"MODEL_SERVER not found in profile".to_string(),
)
})?;
{
Some(server) => server,
None => return Ok(None),
};

let server = ModelServer::from_str(model_server)?;

Expand Down Expand Up @@ -212,6 +211,7 @@ impl UserProfileDbHandler {
tx.execute_batch(
"
DELETE FROM user_profiles;
DELETE FROM provider_configs;
DELETE FROM encryption_keys;
",
)?;
Expand Down
3 changes: 3 additions & 0 deletions lumni/src/apps/builtin/llm/prompt/src/chat/db/user_profile/profile_operations.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@ use lumni::api::error::ApplicationError;
use lumni::Timestamp;
use rusqlite::{params, OptionalExtension, Transaction};
use serde_json::Value as JsonValue;
use tokio::time::{sleep, Duration};

use super::{
DatabaseOperationError, EncryptionHandler, EncryptionMode, MaskMode,
Expand All @@ -18,6 +19,8 @@ impl UserProfileDbHandler {
profile_name: &str,
settings: &JsonValue,
) -> Result<UserProfile, ApplicationError> {
// Simulate a 3-second delay
// sleep(Duration::from_secs(3)).await;
let timestamp = Timestamp::from_system_time().unwrap().as_millis();

let encryption_key_id = self.get_or_create_encryption_key().await?;
Expand Down
1 change: 0 additions & 1 deletion lumni/src/apps/builtin/llm/prompt/src/chat/prompt/builder.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -66,7 +66,6 @@ impl PromptInstructionBuilder {
// If no profile selector is provided, use the default profile
self = self.from_default().await?;
}

Ok(self)
}

Expand Down
91 changes: 52 additions & 39 deletions lumni/src/apps/builtin/llm/prompt/src/chat/session/chat_session_manager.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -25,61 +25,70 @@ pub struct SessionInfo {

pub struct ChatSessionManager {
sessions: HashMap<Uuid, ThreadedChatSession>,
pub active_session_info: SessionInfo,
pub active_session_info: Option<SessionInfo>,
}

#[allow(dead_code)]
impl ChatSessionManager {
pub async fn new(
initial_prompt_instruction: PromptInstruction,
initial_prompt_instruction: Option<PromptInstruction>,
db_conn: Arc<ConversationDatabase>,
) -> Self {
let session_id = Uuid::new_v4();
let conversation_id = initial_prompt_instruction.get_conversation_id();
let server_name = initial_prompt_instruction
.get_completion_options()
.model_server
.as_ref()
.map(|s| s.to_string());

let initial_session = ThreadedChatSession::new(
initial_prompt_instruction,
db_conn.clone(),
);

let mut sessions = HashMap::new();
sessions.insert(session_id, initial_session);

Self {
sessions,
active_session_info: SessionInfo {
let active_session_info = if let Some(prompt_instruction) =
initial_prompt_instruction
{
let session_id = Uuid::new_v4();
let conversation_id = prompt_instruction.get_conversation_id();
let server_name = prompt_instruction
.get_completion_options()
.model_server
.as_ref()
.map(|s| s.to_string());
let initial_session =
ThreadedChatSession::new(prompt_instruction, db_conn.clone());
sessions.insert(session_id, initial_session);
Some(SessionInfo {
id: session_id,
conversation_id,
server_name,
},
})
} else {
None
};

Self {
sessions,
active_session_info,
}
}

pub fn get_active_session(
&mut self,
) -> Result<&mut ThreadedChatSession, ApplicationError> {
self.sessions
.get_mut(&self.active_session_info.id)
.ok_or_else(|| {
ApplicationError::Runtime(
) -> Result<Option<&mut ThreadedChatSession>, ApplicationError> {
if let Some(ref session_info) = self.active_session_info {
if let Some(session) = self.sessions.get_mut(&session_info.id) {
Ok(Some(session))
} else {
Err(ApplicationError::Runtime(
"Active session not found".to_string(),
)
})
))
}
} else {
Ok(None)
}
}

pub fn get_conversation_id_for_active_session(
&self,
) -> Option<ConversationId> {
self.active_session_info.conversation_id
self.active_session_info
.as_ref()
.and_then(|info| info.conversation_id)
}

pub fn get_active_session_id(&self) -> Uuid {
self.active_session_info.id
pub fn get_active_session_id(&self) -> Option<Uuid> {
self.active_session_info.as_ref().map(|info| info.id)
}

pub async fn stop_session(
Expand Down Expand Up @@ -126,11 +135,11 @@ impl ChatSessionManager {
.as_ref()
.map(|s| s.to_string());

self.active_session_info = SessionInfo {
self.active_session_info = Some(SessionInfo {
id,
conversation_id,
server_name,
};
});
Ok(())
} else {
Err(ApplicationError::InvalidInput(
Expand All @@ -140,14 +149,18 @@ impl ChatSessionManager {
}

pub fn stop_active_chat_session(&mut self) -> Result<(), ApplicationError> {
if let Some(session) =
self.sessions.get_mut(&self.active_session_info.id)
{
session.stop();
Ok(())
if let Some(ref session_info) = self.active_session_info {
if let Some(session) = self.sessions.get_mut(&session_info.id) {
session.stop();
Ok(())
} else {
Err(ApplicationError::Runtime(
"Active session not found".to_string(),
))
}
} else {
Err(ApplicationError::Runtime(
"Active session not found".to_string(),
"No active session to stop".to_string(),
))
}
}
Expand Down
47 changes: 27 additions & 20 deletions lumni/src/apps/builtin/llm/prompt/src/chat/session/conversation_loop.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -47,20 +47,21 @@ pub async fn prompt_app<B: Backend>(
}
_ = async {
// Process chat events
let events = app.chat_manager.get_active_session()?.subscribe().recv().await;
if let Ok(event) = events {
match event {
ChatEvent::ResponseUpdate(content) => {
app.ui.response.text_append(&content, Some(color_scheme.get_secondary_style()))?;
redraw_ui = true;
},
ChatEvent::FinalResponse => {
app.ui.response.text_append("\n\n", Some(color_scheme.get_secondary_style()))?;
redraw_ui = true;
},
ChatEvent::Error(error) => {
log::error!("Chat session error: {}", error);
redraw_ui = true;
if let Ok(Some(active_session)) = app.chat_manager.get_active_session() {
if let Ok(event) = active_session.subscribe().recv().await {
match event {
ChatEvent::ResponseUpdate(content) => {
app.ui.response.text_append(&content, Some(color_scheme.get_secondary_style()))?;
redraw_ui = true;
},
ChatEvent::FinalResponse => {
app.ui.response.text_append("\n\n", Some(color_scheme.get_secondary_style()))?;
redraw_ui = true;
},
ChatEvent::Error(error) => {
log::error!("Chat session error: {}", error);
redraw_ui = true;
}
}
}
}
Expand Down Expand Up @@ -136,11 +137,13 @@ async fn handle_key_event(
let mut conversation_handler = db_conn.get_conversation_handler(
app.get_conversation_id_for_active_session(),
);

let active_session = app.chat_manager.get_active_session()?;
let new_window_event = key_event_handler
.process_key(
key_event,
&mut app.ui,
app.chat_manager.get_active_session()?,
active_session,
mode,
keep_running.clone(),
&mut conversation_handler,
Expand Down Expand Up @@ -338,11 +341,15 @@ async fn send_prompt<'a>(
// prompt should end with single newline
let color_scheme = app.color_scheme.clone();
let formatted_prompt = format!("{}\n", prompt.trim_end());
let result = app
.chat_manager
.get_active_session()?
.message(&formatted_prompt)
.await;

let active_session =
app.chat_manager.get_active_session()?.ok_or_else(|| {
ApplicationError::NotReady(
"No active session available".to_string(),
)
})?;
let result = active_session.message(&formatted_prompt).await;

match result {
Ok(_) => {
// clear prompt
Expand Down
Loading

0 comments on commit a1f0a59

Please sign in to comment.