diff --git a/Cargo.toml b/Cargo.toml index 4fb9363..81b93ab 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -4,10 +4,17 @@ version = "0.1.0" edition = "2021" [dependencies] -gtk4 = { version = "0.9", features = ["v4_6"] } -tokio = { version = "1.0", features = ["full"] } +# GTK4 and GLib +gtk4 = { version = "0.9", features = ["v4_12"] } + +# HTTP client with JSON support +reqwest = { version = "0.12", default-features = false, features = ["rustls-tls", "json", "stream"] } + +# Async runtime and utilities +tokio = { version = "1.0", features = ["rt-multi-thread", "time"] } +futures-util = "0.3" +async-channel = "2.1" + +# Serialization serde = { version = "1.0", features = ["derive"] } -serde_json = "1.0" -reqwest = { version = "0.12", features = ["json"] } -pango = "0.20" -regex = "1.0" +serde_json = "1.0" \ No newline at end of file diff --git a/src/api.rs b/src/api.rs index 19fafeb..69d5fe9 100644 --- a/src/api.rs +++ b/src/api.rs @@ -1,18 +1,24 @@ use std::sync::{Arc, Mutex}; -use crate::types::{ChatMessage, ChatRequest, ChatResponse, ModelInfo, ModelsResponse}; +use futures_util::StreamExt; +use tokio::time::{timeout, Duration}; +use crate::types::{ChatMessage, ChatRequest, ModelInfo, ModelsResponse, StreamResponse}; -pub async fn fetch_models(base_url: &str) -> Result, Box> { +pub async fn fetch_models(base_url: &str) -> Result, Box> { let url = format!("{}/api/tags", base_url); - let response = reqwest::get(&url).await?; + + // Add timeout to prevent hanging + let response = timeout(Duration::from_secs(10), reqwest::get(&url)).await??; let models_response: ModelsResponse = response.json().await?; Ok(models_response.models) } -pub async fn send_chat_request( +pub async fn send_chat_request_streaming( base_url: &str, model: &str, conversation: &Arc>>, -) -> Result> { + token_sender: async_channel::Sender, + thinking_enabled: bool, +) -> Result> { let messages = { let conversation = conversation.lock().unwrap(); conversation.iter().cloned().collect::>() @@ -21,10 +27,14 @@ pub async fn send_chat_request( let request = ChatRequest { model: model.to_string(), messages, - stream: false, + stream: true, + think: if thinking_enabled { Some(true) } else { Some(false) }, }; - let client = reqwest::Client::new(); + let client = reqwest::Client::builder() + .timeout(Duration::from_secs(120)) // 2 minute timeout + .build()?; + let url = format!("{}/api/chat", base_url); let response = client @@ -33,6 +43,82 @@ pub async fn send_chat_request( .send() .await?; - let chat_response: ChatResponse = response.json().await?; - Ok(chat_response.message.content) -} + if !response.status().is_success() { + return Err(format!("API request failed with status: {}", response.status()).into()); + } + + let mut stream = response.bytes_stream(); + let mut full_response = String::new(); + let mut current_batch = String::new(); + let mut tokens_since_last_send = 0; + const BATCH_SIZE: usize = 20; // Reasonable batch size + const BATCH_TIMEOUT: Duration = Duration::from_millis(100); + + let mut last_send = tokio::time::Instant::now(); + + while let Some(chunk_result) = stream.next().await { + let chunk = chunk_result?; + let text = String::from_utf8_lossy(&chunk); + + for line in text.lines() { + if line.trim().is_empty() { + continue; + } + + match serde_json::from_str::(line) { + Ok(stream_response) => { + let token = stream_response.message.content; + + if !token.is_empty() { + full_response.push_str(&token); + current_batch.push_str(&token); + tokens_since_last_send += 1; + } + + // Send batch if conditions are met + let should_send = tokens_since_last_send >= BATCH_SIZE + || last_send.elapsed() >= BATCH_TIMEOUT + || stream_response.done; + + if should_send && !current_batch.is_empty() { + // Use non-blocking send with async-channel + match token_sender.send(current_batch.clone()).await { + Ok(_) => { + current_batch.clear(); + tokens_since_last_send = 0; + last_send = tokio::time::Instant::now(); + } + Err(_) => { + // Channel closed, stop sending but continue processing + break; + } + } + } + + if stream_response.done { + break; + } + } + Err(parse_error) => { + // Log parse errors but continue processing + eprintln!("Failed to parse streaming response: {} (line: {})", parse_error, line); + continue; + } + } + } + } + + // Send any remaining tokens in the batch + if !current_batch.is_empty() { + let _ = token_sender.send(current_batch).await; + } + + // async_channel automatically closes when sender is dropped + drop(token_sender); + + if full_response.is_empty() { + return Err("No response received from the model".into()); + } + + Ok(full_response) +} \ No newline at end of file diff --git a/src/main.rs b/src/main.rs index 994e532..aa81e45 100644 --- a/src/main.rs +++ b/src/main.rs @@ -5,14 +5,12 @@ mod ui; mod api; mod types; mod state; - -use state::AppState; +mod markdown_processor; const APP_ID: &str = "com.example.ollama-chat"; -#[tokio::main] -async fn main() -> glib::ExitCode { +fn main() -> glib::ExitCode { let app = Application::builder().application_id(APP_ID).build(); app.connect_activate(ui::build_ui); app.run() -} +} \ No newline at end of file diff --git a/src/markdown_processor.rs b/src/markdown_processor.rs new file mode 100644 index 0000000..c677a3d --- /dev/null +++ b/src/markdown_processor.rs @@ -0,0 +1,351 @@ +use gtk4::prelude::*; +use gtk4::{TextBuffer, TextTag, TextTagTable}; +use std::collections::HashMap; + +pub struct MarkdownProcessor { + h1_tag: TextTag, + h2_tag: TextTag, + h3_tag: TextTag, + h4_tag: TextTag, + h5_tag: TextTag, + h6_tag: TextTag, + bold_tag: TextTag, + italic_tag: TextTag, + code_tag: TextTag, +} + +impl MarkdownProcessor { + pub fn new() -> Self { + // Heading tags + let h1_tag = TextTag::new(Some("h1")); + h1_tag.set_weight(700); + h1_tag.set_line_height(1.4); + h1_tag.set_scale(1.5); + + let h2_tag = TextTag::new(Some("h2")); + h2_tag.set_weight(700); + h2_tag.set_line_height(1.4); + h2_tag.set_scale(1.4); + + let h3_tag = TextTag::new(Some("h3")); + h3_tag.set_weight(700); + h3_tag.set_line_height(1.2); + h3_tag.set_scale(1.3); + + let h4_tag = TextTag::new(Some("h4")); + h4_tag.set_weight(700); + h4_tag.set_line_height(1.2); + h4_tag.set_scale(1.2); + + let h5_tag = TextTag::new(Some("h5")); + h5_tag.set_weight(700); + h5_tag.set_line_height(1.2); + h5_tag.set_scale(1.1); + + let h6_tag = TextTag::new(Some("h6")); + h6_tag.set_weight(700); + h6_tag.set_line_height(1.2); + h6_tag.set_scale(1.0); + + // Inline formatting tags + let bold_tag = TextTag::new(Some("bold")); + bold_tag.set_weight(700); + + let italic_tag = TextTag::new(Some("italic")); + italic_tag.set_style(gtk4::pango::Style::Italic); + + let code_tag = TextTag::new(Some("code")); + code_tag.set_property("font", &"monospace"); + code_tag.set_scale(0.9); + + Self { + h1_tag, + h2_tag, + h3_tag, + h4_tag, + h5_tag, + h6_tag, + bold_tag, + italic_tag, + code_tag, + } + } + + pub fn setup_tags(&self, buffer: &TextBuffer) { + let tag_table = buffer.tag_table(); + tag_table.add(&self.h1_tag); + tag_table.add(&self.h2_tag); + tag_table.add(&self.h3_tag); + tag_table.add(&self.h4_tag); + tag_table.add(&self.h5_tag); + tag_table.add(&self.h6_tag); + tag_table.add(&self.bold_tag); + tag_table.add(&self.italic_tag); + tag_table.add(&self.code_tag); + } + + pub fn insert_formatted_text(&self, buffer: &TextBuffer, text: &str, iter: &mut gtk4::TextIter) { + // Process the text line by line to handle headings and inline formatting + let lines: Vec<&str> = text.lines().collect(); + + for (i, line) in lines.iter().enumerate() { + // Check if this line is a heading + if let Some(heading) = self.parse_heading(line) { + let tag = match heading.level { + 1 => &self.h1_tag, + 2 => &self.h2_tag, + 3 => &self.h3_tag, + 4 => &self.h4_tag, + 5 => &self.h5_tag, + _ => &self.h6_tag, + }; + + buffer.insert(iter, "\n"); + buffer.insert_with_tags(iter, &heading.content, &[tag]); + buffer.insert(iter, "\n"); + } else { + // Process inline formatting for non-heading lines + self.insert_formatted_line(buffer, line, iter); + + // Add newline if not the last line + if i < lines.len() - 1 { + buffer.insert(iter, "\n"); + } + } + } + } + + fn insert_formatted_line(&self, buffer: &TextBuffer, line: &str, iter: &mut gtk4::TextIter) { + let segments = self.parse_inline_formatting(line); + + for segment in segments { + match segment.format_type { + FormatType::Plain => { + buffer.insert(iter, &segment.content); + } + FormatType::Bold => { + buffer.insert_with_tags(iter, &segment.content, &[&self.bold_tag]); + } + FormatType::Italic => { + buffer.insert_with_tags(iter, &segment.content, &[&self.italic_tag]); + } + FormatType::Code => { + buffer.insert_with_tags(iter, &segment.content, &[&self.code_tag]); + } + FormatType::BoldItalic => { + buffer.insert_with_tags(iter, &segment.content, &[&self.bold_tag, &self.italic_tag]); + } + FormatType::Heading(_) => { + // This shouldn't happen in inline processing + buffer.insert(iter, &segment.content); + } + } + } + } + + fn parse_heading(&self, line: &str) -> Option { + if line.starts_with('#') { + let hash_count = line.chars().take_while(|&c| c == '#').count(); + + if hash_count <= 6 && line.len() > hash_count + 1 && line.chars().nth(hash_count) == Some(' ') { + let content = &line[hash_count + 1..].trim(); + return Some(HeadingInfo { + level: hash_count, + content: content.to_string(), + }); + } + } + None + } + + fn parse_inline_formatting(&self, text: &str) -> Vec { + let mut segments = Vec::new(); + let mut current_pos = 0; + let chars: Vec = text.chars().collect(); + + while current_pos < chars.len() { + // Look for the next formatting marker + if let Some((marker_pos, marker_type, marker_len)) = self.find_next_marker(&chars, current_pos) { + // Add any plain text before the marker + if marker_pos > current_pos { + let plain_text: String = chars[current_pos..marker_pos].iter().collect(); + if !plain_text.is_empty() { + segments.push(TextSegment { + content: plain_text, + format_type: FormatType::Plain, + }); + } + } + + // Find the closing marker + if let Some((close_pos, close_len)) = self.find_closing_marker(&chars, marker_pos + marker_len, &marker_type) { + let content_start = marker_pos + marker_len; + let content_end = close_pos; + + if content_start < content_end { + let content: String = chars[content_start..content_end].iter().collect(); + segments.push(TextSegment { + content, + format_type: marker_type, + }); + } + + current_pos = close_pos + close_len; + } else { + // No closing marker found, treat as plain text + let plain_char: String = chars[marker_pos..marker_pos + marker_len].iter().collect(); + segments.push(TextSegment { + content: plain_char, + format_type: FormatType::Plain, + }); + current_pos = marker_pos + marker_len; + } + } else { + // No more markers, add the rest as plain text + let remaining: String = chars[current_pos..].iter().collect(); + if !remaining.is_empty() { + segments.push(TextSegment { + content: remaining, + format_type: FormatType::Plain, + }); + } + break; + } + } + + segments + } + + fn find_next_marker(&self, chars: &[char], start_pos: usize) -> Option<(usize, FormatType, usize)> { + let mut earliest_pos = None; + let mut earliest_type = FormatType::Plain; + let mut earliest_len = 0; + + for pos in start_pos..chars.len() { + // Check for inline code (backticks) + if chars[pos] == '`' { + if earliest_pos.is_none() || pos < earliest_pos.unwrap() { + earliest_pos = Some(pos); + earliest_type = FormatType::Code; + earliest_len = 1; + } + break; // Prioritize code as it can contain other markers + } + + // Check for bold/italic markers + if pos + 1 < chars.len() { + // Check for ** (bold) or *** (bold+italic) + if chars[pos] == '*' && chars[pos + 1] == '*' { + if pos + 2 < chars.len() && chars[pos + 2] == '*' { + // *** bold+italic + if earliest_pos.is_none() || pos < earliest_pos.unwrap() { + earliest_pos = Some(pos); + earliest_type = FormatType::BoldItalic; + earliest_len = 3; + } + } else { + // ** bold + if earliest_pos.is_none() || pos < earliest_pos.unwrap() { + earliest_pos = Some(pos); + earliest_type = FormatType::Bold; + earliest_len = 2; + } + } + break; + } + + // Check for __ (bold) + if chars[pos] == '_' && chars[pos + 1] == '_' { + if earliest_pos.is_none() || pos < earliest_pos.unwrap() { + earliest_pos = Some(pos); + earliest_type = FormatType::Bold; + earliest_len = 2; + } + break; + } + } + + // Check for single * or _ (italic) + if chars[pos] == '*' || chars[pos] == '_' { + // Make sure it's not part of a double marker + let is_single = (pos == 0 || chars[pos - 1] != chars[pos]) && + (pos + 1 >= chars.len() || chars[pos + 1] != chars[pos]); + + if is_single && (earliest_pos.is_none() || pos < earliest_pos.unwrap()) { + earliest_pos = Some(pos); + earliest_type = FormatType::Italic; + earliest_len = 1; + } + } + } + + earliest_pos.map(|pos| (pos, earliest_type, earliest_len)) + } + + fn find_closing_marker(&self, chars: &[char], start_pos: usize, marker_type: &FormatType) -> Option<(usize, usize)> { + match marker_type { + FormatType::Code => { + // Look for closing backtick + for pos in start_pos..chars.len() { + if chars[pos] == '`' { + return Some((pos, 1)); + } + } + } + FormatType::Bold => { + // Look for closing ** or __ + for pos in start_pos..chars.len().saturating_sub(1) { + if (chars[pos] == '*' && chars[pos + 1] == '*') || + (chars[pos] == '_' && chars[pos + 1] == '_') { + return Some((pos, 2)); + } + } + } + FormatType::Italic => { + // Look for closing * or _ + for pos in start_pos..chars.len() { + if chars[pos] == '*' || chars[pos] == '_' { + // Make sure it's not part of a double marker + let is_single = (pos == 0 || chars[pos - 1] != chars[pos]) && + (pos + 1 >= chars.len() || chars[pos + 1] != chars[pos]); + if is_single { + return Some((pos, 1)); + } + } + } + } + FormatType::BoldItalic => { + // Look for closing *** + for pos in start_pos..chars.len().saturating_sub(2) { + if chars[pos] == '*' && chars[pos + 1] == '*' && chars[pos + 2] == '*' { + return Some((pos, 3)); + } + } + } + _ => {} + } + None + } +} + +#[derive(Debug, Clone)] +struct TextSegment { + content: String, + format_type: FormatType, +} + +#[derive(Debug, Clone)] +struct HeadingInfo { + level: usize, + content: String, +} + +#[derive(Debug, Clone, PartialEq)] +enum FormatType { + Plain, + Heading(usize), + Bold, + Italic, + Code, + BoldItalic, +} \ No newline at end of file diff --git a/src/state.rs b/src/state.rs index bc2f647..cbb9392 100644 --- a/src/state.rs +++ b/src/state.rs @@ -5,6 +5,8 @@ use crate::types::ChatMessage; pub struct AppState { pub conversation: Arc>>, pub ollama_url: String, + pub thinking_enabled: Arc>, + pub current_request_handle: Arc>>>, } impl Default for AppState { @@ -12,6 +14,8 @@ impl Default for AppState { Self { conversation: Arc::new(Mutex::new(Vec::new())), ollama_url: "http://localhost:11434".to_string(), + thinking_enabled: Arc::new(Mutex::new(false)), + current_request_handle: Arc::new(Mutex::new(None)), } } -} +} \ No newline at end of file diff --git a/src/types.rs b/src/types.rs index 282bf79..a454324 100644 --- a/src/types.rs +++ b/src/types.rs @@ -11,6 +11,8 @@ pub struct ChatRequest { pub model: String, pub messages: Vec, pub stream: bool, + #[serde(skip_serializing_if = "Option::is_none")] + pub think: Option, } #[derive(Debug, Serialize, Deserialize)] @@ -20,6 +22,14 @@ pub struct ChatResponse { pub done: bool, } +#[derive(Debug, Serialize, Deserialize)] +pub struct StreamResponse { + pub model: String, + pub created_at: String, + pub message: ChatMessage, + pub done: bool, +} + #[derive(Debug, Serialize, Deserialize)] pub struct ModelInfo { pub name: String, @@ -28,4 +38,4 @@ pub struct ModelInfo { #[derive(Debug, Serialize, Deserialize)] pub struct ModelsResponse { pub models: Vec, -} +} \ No newline at end of file diff --git a/src/ui.rs b/src/ui.rs index 8096651..d37ea5d 100644 --- a/src/ui.rs +++ b/src/ui.rs @@ -1,11 +1,21 @@ use gtk4::prelude::*; -use gtk4::{glib, Application, ApplicationWindow, Button, ComboBoxText, Label, ScrolledWindow, TextView, TextBuffer, TextTag, TextTagTable, Orientation, PolicyType, WrapMode, Align}; +use gtk4::{glib, Application, ApplicationWindow, Button, DropDown, Label, ScrolledWindow, TextView, TextBuffer, StringList, Orientation, PolicyType, WrapMode, Align, CheckButton}; use gtk4::Box as GtkBox; -use glib::spawn_future_local; +use glib::{spawn_future_local, clone}; +use std::sync::{OnceLock, Arc, Mutex}; +use tokio::runtime::Runtime; use crate::api; use crate::state::AppState; use crate::types::ChatMessage; +use crate::markdown_processor::MarkdownProcessor; + +// enum to track button state +#[derive(Clone, PartialEq)] +enum ButtonState { + Send, + Stop, +} pub fn build_ui(app: &Application) { let window = ApplicationWindow::builder() @@ -15,30 +25,47 @@ pub fn build_ui(app: &Application) { .default_height(700) .build(); - // Apply minimal CSS for larger fonts and spacing let css_provider = gtk4::CssProvider::new(); - css_provider.load_from_data( + css_provider.load_from_string( r#" window { font-size: 16px; } .chat-text { - font-size: 16px; + font-size: 18px; padding: 24px; + border-radius: 12px; } .input-text { font-size: 16px; padding: 16px; + border-radius: 12px; } button { font-size: 16px; padding: 16px 24px; + border-radius: 12px; } - combobox { + .stop-button { + background-color: #dc3545; + color: white; + } + + .send-button { + background-color: #007bff; + color: white; + } + + dropdown { + font-size: 16px; + border-radius: 12px; + } + + checkbutton { font-size: 16px; } "# @@ -62,7 +89,7 @@ pub fn build_ui(app: &Application) { chat_scroll.set_policy(PolicyType::Never, PolicyType::Automatic); chat_scroll.set_vexpand(true); - let (chat_view, chat_buffer) = create_chat_view(); + let (chat_view, chat_buffer, markdown_processor) = create_chat_view(); chat_scroll.set_child(Some(&chat_view)); // Input area @@ -83,23 +110,37 @@ pub fn build_ui(app: &Application) { let input_buffer = input_view.buffer(); input_scroll.set_child(Some(&input_view)); - let send_button = Button::with_label("Send"); - send_button.set_valign(Align::End); + // Single button that changes between Send and Stop + let action_button = Button::with_label("Send"); + action_button.add_css_class("send-button"); + action_button.set_valign(Align::End); + + // Track button state + let button_state = Arc::new(Mutex::new(ButtonState::Send)); input_area_container.append(&input_scroll); - input_area_container.append(&send_button); + input_area_container.append(&action_button); // Bottom controls let controls_container = GtkBox::new(Orientation::Horizontal, 16); let model_label = Label::new(Some("Model:")); - let model_combo = ComboBoxText::new(); + + // Create StringList to hold model names + let model_list = StringList::new(&[]); + let model_dropdown = DropDown::new(Some(model_list.clone()), None::); + + // Add thinking checkbox + let thinking_checkbox = CheckButton::with_label("Think"); + thinking_checkbox.set_tooltip_text(Some("Enable thinking mode for reasoning models (e.g., deepseek-r1)")); + let status_label = Label::new(Some("Ready")); status_label.set_hexpand(true); status_label.set_halign(Align::End); controls_container.append(&model_label); - controls_container.append(&model_combo); + controls_container.append(&model_dropdown); + controls_container.append(&thinking_checkbox); controls_container.append(&status_label); input_container.append(&input_area_container); @@ -114,120 +155,184 @@ pub fn build_ui(app: &Application) { let app_state = AppState::default(); // Load available models - load_models(model_combo.clone(), status_label.clone(), app_state.clone()); + load_models(model_list.clone(), model_dropdown.clone(), status_label.clone(), app_state.clone()); // Set up event handlers - setup_send_handler( - send_button.clone(), + setup_action_button_handler( + action_button.clone(), + button_state.clone(), input_buffer, - chat_buffer, - model_combo, - status_label, - app_state, + chat_buffer.clone(), + model_dropdown, + model_list, + thinking_checkbox.clone(), + status_label.clone(), + app_state.clone(), + markdown_processor, + chat_scroll.clone(), ); - setup_keyboard_shortcut(input_view, send_button); + // Connect thinking checkbox to app state + setup_thinking_checkbox_handler(thinking_checkbox, app_state.clone()); + + setup_keyboard_shortcut(input_view, action_button.clone(), button_state.clone()); window.present(); } -fn create_chat_view() -> (TextView, TextBuffer) { - let chat_view = TextView::new(); - chat_view.set_editable(false); - chat_view.set_cursor_visible(false); - chat_view.set_wrap_mode(WrapMode::WordChar); - chat_view.add_css_class("chat-text"); - - let chat_buffer = TextBuffer::new(None); - chat_view.set_buffer(Some(&chat_buffer)); - - (chat_view, chat_buffer) -} - -fn load_models(combo: ComboBoxText, status_label: Label, app_state: AppState) { - status_label.set_text("Loading models..."); - - let combo_weak = combo.downgrade(); - let status_weak = status_label.downgrade(); - - spawn_future_local(async move { - match api::fetch_models(&app_state.ollama_url).await { - Ok(models) => { - if let (Some(combo), Some(status_label)) = (combo_weak.upgrade(), status_weak.upgrade()) { - combo.remove_all(); - for model in models { - combo.append_text(&model.name); - } - if combo.active().is_none() && combo.model().unwrap().iter_n_children(None) > 0 { - combo.set_active(Some(0)); - } - status_label.set_text("Ready"); - } - } - Err(e) => { - if let Some(status_label) = status_weak.upgrade() { - status_label.set_text(&format!("Error loading models: {}", e)); - } - } +// Helper function to update button appearance +fn update_button_state(button: &Button, state: ButtonState) { + match state { + ButtonState::Send => { + button.set_label("Send"); + button.remove_css_class("stop-button"); + button.add_css_class("send-button"); } - }); + ButtonState::Stop => { + button.set_label("Stop"); + button.remove_css_class("send-button"); + button.add_css_class("stop-button"); + } + } } -fn setup_send_handler( - send_button: Button, +fn setup_action_button_handler( + action_button: Button, + button_state: Arc>, input_buffer: TextBuffer, chat_buffer: TextBuffer, - model_combo: ComboBoxText, + model_dropdown: DropDown, + model_list: StringList, + thinking_checkbox: CheckButton, status_label: Label, app_state: AppState, + markdown_processor: Arc, + chat_scroll: ScrolledWindow, ) { - send_button.connect_clicked(move |_| { - let start_iter = input_buffer.start_iter(); - let end_iter = input_buffer.end_iter(); - let text = input_buffer.text(&start_iter, &end_iter, false); - - if text.trim().is_empty() { - return; + action_button.connect_clicked(clone!( + #[weak] input_buffer, + #[weak] chat_buffer, + #[weak] model_dropdown, + #[weak] model_list, + #[weak] thinking_checkbox, + #[weak] status_label, + #[weak] action_button, + #[strong] button_state, + #[strong] markdown_processor, + move |_| { + let current_state = { + let state = button_state.lock().unwrap(); + state.clone() + }; + + match current_state { + ButtonState::Send => { + // Handle send logic + let start_iter = input_buffer.start_iter(); + let end_iter = input_buffer.end_iter(); + let text = input_buffer.text(&start_iter, &end_iter, false); + + if text.trim().is_empty() { + return; + } + + // Get selected model + let selected_idx = model_dropdown.selected(); + if selected_idx == gtk4::INVALID_LIST_POSITION { + status_label.set_text("Please select a model first"); + return; + } + + let model = match model_list.string(selected_idx) { + Some(m) => m.to_string(), + None => { + status_label.set_text("Invalid model selection"); + return; + } + }; + + // Get thinking checkbox state + let thinking_enabled = thinking_checkbox.is_active(); + + input_buffer.delete(&mut input_buffer.start_iter(), &mut input_buffer.end_iter()); + + // Change button to Stop state + { + let mut state = button_state.lock().unwrap(); + *state = ButtonState::Stop; + } + update_button_state(&action_button, ButtonState::Stop); + + send_message( + text.to_string(), + model, + thinking_enabled, + chat_buffer.clone(), + chat_scroll.clone(), + status_label.clone(), + action_button.clone(), + button_state.clone(), + app_state.clone(), + markdown_processor.clone(), + ); + } + ButtonState::Stop => { + // Handle stop logic + let mut handle = app_state.current_request_handle.lock().unwrap(); + if let Some(task) = handle.take() { + task.abort(); + status_label.set_text("Generation stopped"); + + // Change button back to Send state + { + let mut state = button_state.lock().unwrap(); + *state = ButtonState::Send; + } + update_button_state(&action_button, ButtonState::Send); + } + } + } } - - let selected_model = model_combo.active_text(); - if selected_model.is_none() { - status_label.set_text("Please select a model first"); - return; - } - - let model = selected_model.unwrap().to_string(); - input_buffer.delete(&mut input_buffer.start_iter(), &mut input_buffer.end_iter()); - - send_message( - text.to_string(), - model, - chat_buffer.clone(), - status_label.clone(), - app_state.clone(), - ); - }); + )); } -fn setup_keyboard_shortcut(input_view: TextView, send_button: Button) { +fn setup_keyboard_shortcut(input_view: TextView, action_button: Button, button_state: Arc>) { let input_controller = gtk4::EventControllerKey::new(); - input_controller.connect_key_pressed(move |_, key, _, modifier| { - if key == gtk4::gdk::Key::Return && modifier.contains(gtk4::gdk::ModifierType::CONTROL_MASK) { - send_button.emit_clicked(); - glib::Propagation::Stop - } else { - glib::Propagation::Proceed + input_controller.connect_key_pressed(clone!( + #[weak] action_button, + #[strong] button_state, + #[upgrade_or] glib::Propagation::Proceed, + move |_, key, _, modifier| { + if key == gtk4::gdk::Key::Return && modifier.contains(gtk4::gdk::ModifierType::CONTROL_MASK) { + // Only trigger if in Send state (don't allow Ctrl+Enter to stop) + let current_state = { + let state = button_state.lock().unwrap(); + state.clone() + }; + + if current_state == ButtonState::Send { + action_button.emit_clicked(); + } + glib::Propagation::Stop + } else { + glib::Propagation::Proceed + } } - }); + )); input_view.add_controller(input_controller); } fn send_message( message: String, model: String, + thinking_enabled: bool, chat_buffer: TextBuffer, + chat_scroll: ScrolledWindow, status_label: Label, + action_button: Button, + button_state: Arc>, app_state: AppState, + markdown_processor: Arc ) { // Add user message to conversation { @@ -238,39 +343,213 @@ fn send_message( }); } - append_to_chat(&chat_buffer, "You", &message); - status_label.set_text("Sending message..."); - - let buffer_weak = chat_buffer.downgrade(); - let status_weak = status_label.downgrade(); + append_to_chat(&chat_buffer, "You", &message, &markdown_processor); - spawn_future_local(async move { - match api::send_chat_request(&app_state.ollama_url, &model, &app_state.conversation).await { - Ok(response_text) => { - // Add assistant response to conversation - { - let mut conversation = app_state.conversation.lock().unwrap(); - conversation.push(ChatMessage { - role: "assistant".to_string(), - content: response_text.clone(), - }); + // Add placeholder for assistant message + let mut end_iter = chat_buffer.end_iter(); + let assistant_prefix = if thinking_enabled { + "\n\nAssistant (thinking enabled):\n" + } else { + "\n\nAssistant:\n" + }; + chat_buffer.insert(&mut end_iter, assistant_prefix); + + // Create mark where assistant response will be inserted + let assistant_start_mark = chat_buffer.create_mark(None, &chat_buffer.end_iter(), true); + + let status_text = if thinking_enabled { + "Assistant is thinking..." + } else { + "Assistant is typing..." + }; + status_label.set_text(status_text); + + // Create channels for streaming communication + let (stream_sender, stream_receiver) = async_channel::bounded::(50); + let (result_sender, result_receiver) = async_channel::bounded::>>(1); + + // Spawn tokio task for API streaming + let app_state_clone = app_state.clone(); + let model_clone = model.clone(); + let task_handle = runtime().spawn(async move { + let result = api::send_chat_request_streaming( + &app_state_clone.ollama_url, + &model_clone, + &app_state_clone.conversation, + stream_sender, + thinking_enabled, + ).await; + let _ = result_sender.send(result).await; + }); + + // Store the handle for potential cancellation + { + let mut handle = app_state.current_request_handle.lock().unwrap(); + *handle = Some(task_handle); + } + + // Handle streaming updates on the main loop + spawn_future_local(clone!( + #[weak] chat_buffer, + #[weak] assistant_start_mark, + #[weak] chat_scroll, + async move { + let mut accumulated_text = String::new(); + + while let Ok(token_batch) = stream_receiver.recv().await { + accumulated_text.push_str(&token_batch); + + // Update UI with accumulated text (plain text during streaming) + let mut start_iter = chat_buffer.iter_at_mark(&assistant_start_mark); + let mut end_iter = chat_buffer.end_iter(); + + // Replace content from mark to end + chat_buffer.delete(&mut start_iter, &mut end_iter); + let mut insert_iter = chat_buffer.iter_at_mark(&assistant_start_mark); + chat_buffer.insert(&mut insert_iter, &accumulated_text); + + // Auto-scroll to bottom after each update + let adjustment = chat_scroll.vadjustment(); + adjustment.set_value(adjustment.upper() - adjustment.page_size()); + } + } + )); + + // Handle final result - apply markdown formatting when streaming completes + let app_state_final = app_state.clone(); + spawn_future_local(clone!( + #[weak] status_label, + #[weak] chat_buffer, + #[weak] assistant_start_mark, + #[weak] action_button, + #[strong] button_state, + #[strong] markdown_processor, + async move { + if let Ok(result) = result_receiver.recv().await { + match result { + Ok(response_text) => { + // Clear the plain streaming text + let mut start_iter = chat_buffer.iter_at_mark(&assistant_start_mark); + let mut end_iter = chat_buffer.end_iter(); + chat_buffer.delete(&mut start_iter, &mut end_iter); + + // Insert formatted markdown text using the shared processor + let mut insert_iter = chat_buffer.iter_at_mark(&assistant_start_mark); + markdown_processor.insert_formatted_text(&chat_buffer, &response_text, &mut insert_iter); + + // Add complete response to conversation + { + let mut conversation = app_state_final.conversation.lock().unwrap(); + conversation.push(ChatMessage { + role: "assistant".to_string(), + content: response_text, + }); + } + + status_label.set_text("Ready"); + } + Err(e) => { + status_label.set_text(&format!("Error: {}", e)); + + // Show error in chat + let mut start_iter = chat_buffer.iter_at_mark(&assistant_start_mark); + let mut end_iter = chat_buffer.end_iter(); + chat_buffer.delete(&mut start_iter, &mut end_iter); + let mut insert_iter = chat_buffer.iter_at_mark(&assistant_start_mark); + chat_buffer.insert(&mut insert_iter, &format!("[Error: {}]", e)); + } } - if let (Some(chat_buffer), Some(status_label)) = (buffer_weak.upgrade(), status_weak.upgrade()) { - append_to_chat(&chat_buffer, "Assistant", &response_text); - status_label.set_text("Ready"); + // Change button back to Send state when generation completes + { + let mut state = button_state.lock().unwrap(); + *state = ButtonState::Send; } + update_button_state(&action_button, ButtonState::Send); } - Err(e) => { - if let Some(status_label) = status_weak.upgrade() { - status_label.set_text(&format!("Error: {}", e)); + } + )); +} + +// Helper functions +fn runtime() -> &'static Runtime { + static RUNTIME: OnceLock = OnceLock::new(); + RUNTIME.get_or_init(|| { + Runtime::new().expect("Setting up tokio runtime needs to succeed.") + }) +} + +fn create_chat_view() -> (TextView, TextBuffer, Arc) { + let chat_view = TextView::new(); + chat_view.set_editable(false); + chat_view.set_cursor_visible(false); + chat_view.set_wrap_mode(WrapMode::WordChar); + chat_view.add_css_class("chat-text"); + + let chat_buffer = TextBuffer::new(None); + chat_view.set_buffer(Some(&chat_buffer)); + + // Set up markdown formatting tags + let markdown_processor = Arc::new(MarkdownProcessor::new()); + markdown_processor.setup_tags(&chat_buffer); + + (chat_view, chat_buffer, markdown_processor) +} + +fn load_models(model_list: StringList, dropdown: DropDown, status_label: Label, app_state: AppState) { + status_label.set_text("Loading models..."); + + // Create communication channel + let (sender, receiver) = async_channel::bounded(1); + + // Spawn tokio task for API call + runtime().spawn(async move { + let result = api::fetch_models(&app_state.ollama_url).await; + let _ = sender.send(result).await; + }); + + // Handle response on main loop + spawn_future_local(clone!( + #[weak] model_list, + #[weak] dropdown, + #[weak] status_label, + async move { + if let Ok(result) = receiver.recv().await { + match result { + Ok(models) => { + // Clear existing items and add new ones + model_list.splice(0, model_list.n_items(), + &models.iter().map(|m| m.name.as_str()).collect::>()); + + // Select first model if available + if model_list.n_items() > 0 && dropdown.selected() == gtk4::INVALID_LIST_POSITION { + dropdown.set_selected(0); + } + + status_label.set_text("Ready"); + } + Err(e) => { + status_label.set_text(&format!("Error loading models: {}", e)); + } } } } - }); + )); } -fn append_to_chat(buffer: &TextBuffer, sender: &str, message: &str) { +fn setup_thinking_checkbox_handler(thinking_checkbox: CheckButton, app_state: AppState) { + thinking_checkbox.connect_toggled(clone!( + #[strong] app_state, + move |checkbox| { + let is_active = checkbox.is_active(); + if let Ok(mut thinking_enabled) = app_state.thinking_enabled.lock() { + *thinking_enabled = is_active; + } + } + )); +} + +fn append_to_chat(buffer: &TextBuffer, sender: &str, message: &str, markdown_processor: &MarkdownProcessor) { let mut end_iter = buffer.end_iter(); // Add spacing if buffer is not empty @@ -279,6 +558,15 @@ fn append_to_chat(buffer: &TextBuffer, sender: &str, message: &str) { end_iter = buffer.end_iter(); } - // Add sender label and message - buffer.insert(&mut end_iter, &format!("{}:\n{}", sender, message)); + // Add sender label + buffer.insert(&mut end_iter, &format!("{}:\n", sender)); + end_iter = buffer.end_iter(); + + // Add message - user messages are always plain text + if sender == "You" { + buffer.insert(&mut end_iter, message); + } else { + // For assistant messages, use markdown formatting + markdown_processor.insert_formatted_text(buffer, message, &mut end_iter); + } } \ No newline at end of file