Fixed some stuff, and added initial fetch to get the current conversation of the VIVO runtime

src/main.rs

@@ -8,13 +8,14 @@ use crossterm::event::{Event, KeyCode, KeyModifiers};
 use crossterm::terminal::{enable_raw_mode, EnterAlternateScreen, disable_raw_mode, LeaveAlternateScreen};
 use ratatui::{DefaultTerminal, Frame, Terminal};
 use ratatui::backend::CrosstermBackend;
-use ratatui::layout::{Alignment, Constraint, Layout, Position};
+use ratatui::layout::{Alignment, Constraint, Layout, Position, Rect};
 use ratatui::style::{Color, Style};
 use ratatui::text::Line;
 use ratatui::widgets::{Block, Borders, List, ListItem, ListState, Paragraph, Wrap};
-use crate::rest::REST;
+use crate::rest::{Message, OllamaMessageRole, REST};
 use textwrap::wrap;
 use tokio::sync::mpsc;
+use crate::rest::OllamaMessageRole::{Assistant, System, User};
 
 #[tokio::main]
 async fn main() -> color_eyre::Result<()> {
@@ -49,16 +50,18 @@ pub struct App {
     input: String,
     should_quit: bool,
     pos: usize,
-    messages: Arc<Mutex<Vec<String>>>,
+    messages: Arc<Mutex<Vec<Message>>>,
     rest: REST,
     list_state: ListState,
-    reply_rx: mpsc::UnboundedReceiver<String>,
+    reply_rx: mpsc::UnboundedReceiver<Message>,
-    reply_tx: mpsc::UnboundedSender<String>,
+    reply_tx: mpsc::UnboundedSender<Message>,
+    chat_area: Rect,
+    input_area: Rect,
 }
 
 impl App {
     pub fn new() -> Self {
-        let (tx, rx) = mpsc::unbounded_channel();
+        let (tx, rx) = mpsc::unbounded_channel::<Message>();
         App{
             input: String::new(),
             should_quit: false,
@@ -68,10 +71,28 @@ impl App {
             list_state: ListState::default(),
             reply_rx: rx,
             reply_tx: tx,
+            chat_area: Rect::default(),
+            input_area: Rect::default(),
         }
     }
 
     fn run(&mut self, terminal: &mut DefaultTerminal) -> std::io::Result<()> {
+        let tx = self.reply_tx.clone();
+        let mut rest = self.rest.clone();
+
+        tokio::spawn(async move {
+            match rest.init().await {
+                Ok(messages) => {
+                    for message in messages {
+                        let _ = tx.send(message);
+                    }
+                },
+                Err(e) => {
+                    eprintln!("{}", e);
+                }
+            }
+        });
+
         terminal.clear()?;
         while !self.should_quit {
             terminal.draw(|frame| self.render(frame))?;
@@ -111,17 +132,18 @@ impl App {
                 KeyCode::Enter => {
                     let str_message = self.input.clone();
 
-                    self.push_message(str_message.clone());
-
-                    self.reset_cursor();
-                    self.input = String::new();
                     let message = rest::Message::new(
                         str_message.clone(),
-                        "USER".parse().unwrap(),
+                        User,
                         false,
                         true,
                         false,
-                        );
+                    );
+
+                    self.push_message(message.clone());
+
+                    self.reset_cursor();
+                    self.input = String::new();
 
                     let tx = self.reply_tx.clone();
                     let mut rest = self.rest.clone();
@@ -129,9 +151,9 @@ impl App {
                     tokio::spawn(async move {
                         match rest.send(message).await {
                             Ok(reply) => {
-                                if reply.role == "ASSISTANT"
+                                if reply.role == Assistant
                                 {
-                                    let _ = tx.send(reply.message);
+                                    let _ = tx.send(reply);
                                 }
                             },
                             Err(e) => {
@@ -143,6 +165,17 @@ impl App {
                 _ => {}
             }
         }
+        /*if let Event::Mouse(mouse) = event::read()? {
+            match mouse.kind {
+                MouseEventKind::ScrollUp => {
+                    self.move_view_up();
+                },
+                MouseEventKind::ScrollDown => {
+                    self.move_view_down();
+                }
+                _ => {}
+            }
+        }*/
         Ok(())
     }
 
@@ -152,23 +185,38 @@ impl App {
         }
     }
 
+    fn total_lines(&self) -> usize {
+        self.messages
+            .lock()
+            .unwrap()
+            .iter()
+            .filter(|msg| msg.role != OllamaMessageRole::System)
+            .map(|msg| {
+                let wrapped = wrap(msg.message.as_str(), self.chat_area.width.saturating_sub(2) as usize).len();
+                wrapped + 1 // +1 for the separator row (skip this for the very last message if you want to be exact)
+            })
+            .sum()
+    }
+
     fn move_view_down(&mut self) {
+        let total = self.total_lines();
         let i = match self.list_state.selected() {
-            Some(i) => (i+1).min(self.messages.lock().unwrap().len().saturating_sub(1)),
+            Some(i) => (i + 1).min(total.saturating_sub(1)),
             None => 0,
         };
         self.list_state.select(Some(i));
     }
 
     fn move_view_up(&mut self) {
+        let total = self.total_lines();
         let i = match self.list_state.selected() {
             Some(i) => i.saturating_sub(1),
-            None => self.messages.lock().unwrap().len().saturating_sub(1),
+            None => total.saturating_sub(1),
         };
         self.list_state.select(Some(i));
     }
 
-    fn push_message(&mut self, message: String) {
+    fn push_message(&mut self, message: Message) {
         let mut messages = self.messages.lock().unwrap();
         messages.push(message);
         let last = messages.len().saturating_sub(1);
@@ -241,6 +289,8 @@ impl App {
             ]);
 
         let [chat_area, input_area] = frame.area().layout(&layout);
+        self.chat_area = chat_area;
+        self.input_area = input_area;
 
         if chat_area.height < 10 {
             // too small — render a warning instead of your normal UI
@@ -255,19 +305,34 @@ impl App {
         let items: Vec<ListItem> = messages
             .iter()
             .enumerate()
-            .map(|(i, msg)| {
+            .flat_map(|(i, msg)| {
-                let mut lines: Vec<Line>= wrap(msg.as_str(), chat_area.width.saturating_sub(2) as usize)
+                if msg.role == System
+                {
+                    return Vec::new();
+                }
+
+                let mut lines: Vec<ListItem>= wrap(msg.message.as_str(), chat_area.width.saturating_sub(2) as usize)
                     .iter()
-                    .map(|s| Line::from(s.to_string()))
+                    .map(|s| {
+                        ListItem::new(Line::styled(
+                            s.to_string(),
+                            Style::default().fg(match msg.role {
+                                User => Color::White,
+                                Assistant => Color::LightCyan,
+                                OllamaMessageRole::Tool => Color::Yellow,
+                                OllamaMessageRole::System => Color::Red,
+                            }),
+                        ))
+                    })
                     .collect();
 
                 if i != messages.len() - 1 {
                     // separator row — width matches the area you render into
                     let sep = "─".repeat(chat_area.width as usize);
-                    lines.push(Line::styled(sep, Style::default().fg(Color::DarkGray)));
+                    lines.push(ListItem::new(Line::styled(sep, Style::default().fg(Color::DarkGray))));
                 }
 
-                ListItem::new(lines)
+                lines
             })
             .collect();
 

src/rest.rs

@@ -1,18 +1,57 @@
 use reqwest::Client;
 use serde::{Deserialize, Serialize};
 
-#[derive(Serialize, Deserialize, Debug)]
+#[derive(Serialize, Deserialize, Debug, Clone)]
 #[serde(rename_all = "camelCase")]
 pub struct Message {
     pub message: String,
-    pub role: String,
+    pub role: OllamaMessageRole,
     pub finished: bool,
     pub no_voice: bool,
     pub stream: bool,
 }
 
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
+#[serde(rename_all = "SCREAMING_SNAKE_CASE")]
+pub enum OllamaMessageRole {
+    /// Represents a user message.
+    User,
+    /// Represents an assistant message.
+    Assistant,
+    /// Represents a tool message.
+    Tool,
+    /// Represents a system message.
+    System,
+}
+
+impl OllamaMessageRole {
+    /// Gets the role as its string representation.
+    pub fn as_str(&self) -> &'static str {
+        match self {
+            OllamaMessageRole::User => "user",
+            OllamaMessageRole::Assistant => "assistant",
+            OllamaMessageRole::Tool => "tool",
+            OllamaMessageRole::System => "system",
+        }
+    }
+}
+
+impl std::str::FromStr for OllamaMessageRole {
+    type Err = String;
+
+    fn from_str(s: &str) -> Result<Self, Self::Err> {
+        match s.to_lowercase().as_str() {
+            "user" => Ok(OllamaMessageRole::User),
+            "assistant" => Ok(OllamaMessageRole::Assistant),
+            "tool" => Ok(OllamaMessageRole::Tool),
+            "system" => Ok(OllamaMessageRole::System),
+            other => Err(format!("Invalid role: {}", other)),
+        }
+    }
+}
+
 impl Message {
-    pub fn new(message: String, role: String, finished: bool, no_voice: bool, stream: bool) -> Self {
+    pub fn new(message: String, role: OllamaMessageRole, finished: bool, no_voice: bool, stream: bool) -> Self {
         Message {
             message,
             role,
@@ -48,4 +87,15 @@ impl REST {
             .await?;
         Ok(res)
     }
+
+    pub async fn init(&mut self) -> Result<Vec<Message>, reqwest::Error> {
+        let res = self.client
+            .get("http://localhost:8080/v1/message/init")
+            .send()
+            .await?
+            .error_for_status()?
+            .json::<Vec<Message>>()
+            .await?;
+        Ok(res)
+    }
 }