fix: finish GenerateViewModel

ollama-cli/src/tui/generate/mod.rs

@@ -1,6 +1,20 @@
-use crate::error::Result;
+use ratatui::{
+    layout::{Constraint, Layout, Rect},
+    style::Style,
+    widgets::{Block, Paragraph, Wrap},
+    Frame,
+};
 
-use super::{event::Action, input::TextInputViewModel, AppEvent};
+use crate::{
+    error::Result,
+    lm::{Prompt, Response},
+};
+
+use super::{
+    event::Action,
+    input::{InputView, TextInputEvent, TextInputViewModel},
+    AppEvent,
+};
 
 #[derive(Debug, Default, Clone, Copy, PartialEq)]
 enum Pane {
@@ -22,17 +36,74 @@ impl Pane {
     }
 }
 
+#[derive(Clone, Debug, Default)]
 pub struct GenerateViewModel {
     input: TextInputViewModel,
     output: String,
+    scroll_state: u16,
     active_pane: Option<Pane>,
     focused_pane: Pane,
 }
 
 impl GenerateViewModel {
-    pub async fn handle_action(&mut self, action: Action) -> Result<Option<AppEvent>> {
+    pub fn handle_response(&mut self, response: Response) -> Result<()> {
+        match response {
+            Response::Token(arc) => {
+                self.output.push_str(arc.as_ref());
+            }
+            Response::Eos => {}
+            Response::Error(arc) => {
+                self.output = arc.to_string();
+            }
+            _ => {}
+        }
+        Ok(())
+    }
+
+    pub fn handle_action(&mut self, action: Action) -> Result<Option<AppEvent>> {
         if let Some(pane) = &self.active_pane {
-            todo!()
+            match pane {
+                Pane::Input => {
+                    Ok(self.input.handle_action(action)?.and_then(
+                        |input_action| match input_action {
+                            TextInputEvent::InputMode(input_mode) => {
+                                Some(AppEvent::InputMode(input_mode))
+                            }
+                            TextInputEvent::Submit(input) => {
+                                Some(AppEvent::Submit(Prompt::Generate(input)))
+                            }
+                            TextInputEvent::Quit => {
+                                self.active_pane = None;
+                                None
+                            }
+                        },
+                    ))
+                }
+                Pane::Output => match action {
+                    Action::Beginning => {
+                        self.scroll_state = 0;
+                        Ok(None)
+                    }
+                    Action::End => {
+                        self.scroll_state = (self.output.lines().count() - 1) as u16;
+                        Ok(None)
+                    }
+                    Action::Up => {
+                        self.scroll_state = self.scroll_state.saturating_sub(1);
+                        Ok(None)
+                    }
+                    Action::Down => {
+                        self.scroll_state = self.scroll_state.saturating_add(1);
+                        Ok(None)
+                    }
+                    Action::Quit | Action::Escape => {
+                        self.active_pane = None;
+                        Ok(None)
+                    }
+                    Action::Enter => todo!(),
+                    _ => Ok(None),
+                },
+            }
         } else {
             match action {
                 Action::Up => {
@@ -54,3 +125,20 @@ impl GenerateViewModel {
         }
     }
 }
+
+#[extend::ext(name = GenerateView)]
+pub impl<'a> Frame<'a> {
+    fn generate_view(&mut self, parent: Rect, style: Style, view_model: &GenerateViewModel) {
+        let vertical = Layout::vertical([Constraint::Percentage(20), Constraint::Min(1)]);
+
+        let [input_area, output_area] = vertical.areas(parent);
+
+        self.input_view(input_area, style, &view_model.input);
+
+        let output = Paragraph::new(view_model.output.as_str())
+            .wrap(Wrap { trim: true })
+            .block(Block::bordered());
+
+        self.render_widget(output, output_area);
+    }
+}

ollama-cli/src/tui/messages/mod.rs

@@ -1,6 +1,5 @@
 use std::collections::VecDeque;
 
-use crossterm::event::{KeyCode, KeyEvent};
 use ratatui::{
     style::{Style, Stylize},
     text::{Span, Text},

ollama-cli/src/tui/mod.rs

@@ -4,6 +4,7 @@ use chat::ChatViewModel;
 use crossterm::ExecutableCommand as _;
 use event::{Action, EventProcessor, InputMode};
 use futures::StreamExt as _;
+use generate::{GenerateView, GenerateViewModel};
 use model_context::ModelContext;
 use models::{ModelsView, ModelsViewModel};
 use nav::{NavView, NavViewModel};
@@ -44,8 +45,9 @@ pub struct AppContext {
 #[strum_discriminants(strum(serialize_all = "lowercase"))]
 #[strum(serialize_all = "lowercase")]
 pub enum View {
-    Chat(ChatViewModel),
     Models(ModelsViewModel),
+    Chat(ChatViewModel),
+    Generate(GenerateViewModel),
     Nav(NavViewModel),
 }
 
@@ -60,6 +62,7 @@ impl View {
         let result = match self {
             View::Chat(ref mut chat_view_model) => chat_view_model.handle_response(response),
             View::Models(ref mut models_view_model) => models_view_model.handle_response(response),
+            View::Generate(ref mut view_model) => view_model.handle_response(response),
             View::Nav(_nav_view_model) => Ok(()),
         };
 
@@ -78,6 +81,7 @@ impl View {
                 models_view_model.handle_event(Action::Refresh).await
             }
             View::Nav(_nav_view_model) => Ok(None),
+            View::Generate(_generate_view_model) => Ok(None),
         }
     }
 }
@@ -113,6 +117,9 @@ impl AppContext {
             View::Nav(nav_view_model) => {
                 frame.nav_view(frame.area(), Style::active(), nav_view_model)
             }
+            View::Generate(generate_view_model) => {
+                frame.generate_view(frame.area(), Style::default(), generate_view_model)
+            }
         }
     }
 
@@ -181,6 +188,7 @@ impl AppContext {
             View::Chat(ref mut chat_view_model) => chat_view_model.handle_action(action).await?,
             View::Models(models_view_model) => models_view_model.handle_event(action).await?,
             View::Nav(nav_view_model) => nav_view_model.handle_action(action)?,
+            View::Generate(generate_view_model) => generate_view_model.handle_action(action)?,
         };
         Ok(app_event)
     }