From c4526de216ec5b8fd473327f3285e063abd376c7 Mon Sep 17 00:00:00 2001 From: Zacharias Date: Sat, 25 Jul 2026 21:24:04 +0200 Subject: [PATCH] Fixed some stuff, and added initial fetch to get the current conversation of the VIVO runtime --- src/main.rs | 109 +++++++++++++++++++++++++++++++++++++++++----------- src/rest.rs | 56 +++++++++++++++++++++++++-- 2 files changed, 140 insertions(+), 25 deletions(-) diff --git a/src/main.rs b/src/main.rs index 00e6194..096a0bc 100644 --- a/src/main.rs +++ b/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>>, + messages: Arc>>, rest: REST, list_state: ListState, - reply_rx: mpsc::UnboundedReceiver, - reply_tx: mpsc::UnboundedSender, + reply_rx: mpsc::UnboundedReceiver, + reply_tx: mpsc::UnboundedSender, + chat_area: Rect, + input_area: Rect, } impl App { pub fn new() -> Self { - let (tx, rx) = mpsc::unbounded_channel(); + let (tx, rx) = mpsc::unbounded_channel::(); 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 = messages .iter() .enumerate() - .map(|(i, msg)| { - let mut lines: Vec= wrap(msg.as_str(), chat_area.width.saturating_sub(2) as usize) + .flat_map(|(i, msg)| { + if msg.role == System + { + return Vec::new(); + } + + let mut lines: Vec= 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(); diff --git a/src/rest.rs b/src/rest.rs index 7b2021b..6626aef 100644 --- a/src/rest.rs +++ b/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 { + 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, reqwest::Error> { + let res = self.client + .get("http://localhost:8080/v1/message/init") + .send() + .await? + .error_for_status()? + .json::>() + .await?; + Ok(res) + } }