From 85bc4265cf0639c11cc354f574775e4b81942b7c Mon Sep 17 00:00:00 2001
From: b1ek <me@blek.codes>
Date: Tue, 22 Oct 2024 14:26:40 +1000
Subject: [PATCH] print new messages with history trait

---
 src/api.rs     | 14 +++++++++++++-
 src/history.rs | 27 +++++++++++++++++++++++++++
 src/main.rs    |  3 +--
 src/session.rs | 25 +++++--------------------
 4 files changed, 46 insertions(+), 23 deletions(-)
 create mode 100644 src/history.rs

diff --git a/src/api.rs b/src/api.rs
index 762a9d9..bb0ed5d 100644
--- a/src/api.rs
+++ b/src/api.rs
@@ -5,7 +5,7 @@ use std::process::exit;
 use reqwest::{header::{HeaderMap, HeaderValue}, Client};
 use serde::{Deserialize, Serialize};
 
-use crate::{config::Config, session::Session};
+use crate::{config::Config, history::HistoryObject, session::Session};
 use crate::{WARN, RED, RESET};
 
 #[derive(Debug, Clone, Serialize, Deserialize)]
@@ -14,6 +14,15 @@ pub struct ChatMessagePayload {
     pub content: String,
 }
 
+impl ChatMessagePayload {
+    pub fn empty(role: String) -> Self {
+        Self {
+            role,
+            content: String::new()
+        }
+    }
+}
+
 #[derive(Debug, Clone, Serialize, Deserialize)]
 pub struct ChatPayload {
     pub model: String,
@@ -114,6 +123,8 @@ pub async fn get_res<'a>(cli: &Client, query: String, vqd: String, config: &Conf
     };
     let payload = serde_json::to_string(&payload).unwrap();
 
+    println!("{}", init_msg.display_as_history());
+
     let req = cli.post("https://duckduckgo.com/duckchat/v1/chat")
         .header("Content-Type", "application/json")
         .header("User-Agent", "Mozilla/5.0 (X11; Linux x86_64; rv:124.0) Gecko/20100101 Firefox/124.0")
@@ -135,6 +146,7 @@ pub async fn get_res<'a>(cli: &Client, query: String, vqd: String, config: &Conf
     }
 
     let mut error = None;
+    println!("{}", ChatMessagePayload::empty("assistant".into()).display_as_history());
     while let Some(chunk) = res.chunk().await.unwrap() {
 
         if let Ok(obj) = serde_json::from_slice::<ErrChatChunk>(&chunk) {
diff --git a/src/history.rs b/src/history.rs
new file mode 100644
index 0000000..13fe78b
--- /dev/null
+++ b/src/history.rs
@@ -0,0 +1,27 @@
+use crate::api::ChatMessagePayload;
+use crate::{RESET, GREEN, GRAY, BLUE};
+
+pub trait HistoryObject {
+    fn display_as_history(&self) -> String;
+}
+
+impl HistoryObject for ChatMessagePayload {
+    fn display_as_history(&self) -> String {
+        let user = &self.role;
+        let msg = &self.content;
+
+        let role_color = {
+            if user == "user" { GREEN }
+            else if user == "assistant" { BLUE }
+            else { GRAY }
+        };
+
+        let short_user = {
+            if user == "user" { "you" }
+            else if user == "assistant" { "ai" }
+            else { user }
+        }.to_string();
+
+        format!("{role_color}{short_user}{RESET}\t: {msg}{RESET}")
+    }
+}
\ No newline at end of file
diff --git a/src/main.rs b/src/main.rs
index 8f1ca8f..3a3ecb6 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -12,6 +12,7 @@ use crate::api::{get_res, get_vqd, simulate_browser_reqs};
 mod config;
 mod api;
 mod session;
+mod history;
 
 pub const GREEN: &str = "\x1b[1;32m";
 pub const RED:   &str = "\x1b[1;31m";
@@ -72,8 +73,6 @@ async fn main() {
         None => get_vqd(&cli).await.unwrap()
     };
 
-    println!("{vqd:?}");
-
     get_res(&cli, query, vqd, &config).await;
 
 }
diff --git a/src/session.rs b/src/session.rs
index 3e8217e..91c89a2 100644
--- a/src/session.rs
+++ b/src/session.rs
@@ -7,7 +7,8 @@ use chrono::{DateTime, Local};
 use serde::{Deserialize, Serialize};
 
 use crate::api::{ChatChunk, ChatMessagePayload};
-use crate::{BLUE, GRAY, GREEN, RESET};
+use crate::{GRAY, RESET};
+use crate::history::HistoryObject;
 
 #[derive(Debug, Clone, Serialize, Deserialize)]
 pub struct Session {
@@ -126,25 +127,9 @@ impl Session {
         println!();
         
         for message in self.messages.iter() {
-            let user = &message.role;
-            let msg = &message.content;
-
-            let role_color = {
-                if user == "user" { GREEN }
-                else if user == "assistant" { BLUE }
-                else { GRAY }
-            };
-
-            let short_user = {
-                if user == "user" { "you" }
-                else if user == "assistant" { "ai" }
-                else { user }
-            }.to_string();
-
-            println!("{role_color}{short_user}{RESET}\t: {msg}{RESET}");
-
-            if user == "assistant" {
-                println!();
+            println!("{}", message.display_as_history());
+            if message.role == "assistant" {
+                println!("\n");
             }
         }