Compare commits

..

3 Commits

4 changed files with 82 additions and 9 deletions

View File

@ -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) {
@ -155,9 +167,10 @@ pub async fn get_res<'a>(cli: &Client, query: String, vqd: String, config: &Conf
}
}
}
println!("\n");
if let Some(err) = error {
eprintln!("Error while writing to session: {err:#?}");
eprintln!("Session may be broken.");
}
println!("\n");
}

27
src/history.rs Normal file
View File

@ -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}")
}
}

View File

@ -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";
@ -42,6 +43,13 @@ async fn main() {
let args = Args::parse();
let query = args.query.join(" ").trim().to_string();
{
let session = Session::create_or_restore("");
if session.is_restored() {
session.print_history();
}
}
if query.len() == 0 {
exit(0);
}
@ -57,8 +65,6 @@ async fn main() {
config.check_tos();
println!("{GREEN}Contacting DuckDuckGo chat AI...{RESET}");
let cli = Client::new();
simulate_browser_reqs(&cli).await.unwrap();
@ -67,8 +73,6 @@ async fn main() {
None => get_vqd(&cli).await.unwrap()
};
println!("{vqd:?}");
get_res(&cli, query, vqd, &config).await;
}

View File

@ -7,13 +7,18 @@ use chrono::{DateTime, Local};
use serde::{Deserialize, Serialize};
use crate::api::{ChatChunk, ChatMessagePayload};
use crate::{GRAY, RESET};
use crate::history::HistoryObject;
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Session {
session_id: String,
ttl: DateTime<Local>,
messages: Vec<ChatMessagePayload>,
vqd: String
vqd: String,
#[serde(skip)]
restored: bool
}
impl Session {
@ -68,6 +73,7 @@ impl Session {
match inner(Self::path_for_id(id)) {
Ok(mut session) => {
session.restored = true;
if session.is_expired() {
session.destroy().expect("Couldn't destroy expired session");
None
@ -94,11 +100,16 @@ impl Session {
session_id: id.into(),
ttl: Self::get_ttl(),
messages: vec![],
vqd: vqd.into()
vqd: vqd.into(),
restored: false
}
}
fn save(&self) -> Result<(), Box<dyn std::error::Error>> {
if self.messages.len() == 0 {
return Ok(());
}
fs::write(self.path(), serde_json::to_string_pretty(self)?)?;
Ok(())
}
@ -107,6 +118,24 @@ impl Session {
self.ttl < Local::now()
}
pub fn is_restored(&self) -> bool {
self.restored
}
pub fn print_history(&self) {
println!("{GRAY}* start restored conversation *{RESET}");
println!();
for message in self.messages.iter() {
println!("{}", message.display_as_history());
if message.role == "assistant" {
println!("\n");
}
}
println!("{GRAY}* end restored conversation *{RESET}");
}
pub fn destroy(&self) -> Result<(), Box<dyn std::error::Error>> {
Ok(fs::remove_file(self.path())?)
}
@ -114,7 +143,7 @@ impl Session {
pub fn create_or_restore<T: Into<String>>(vqd: T) -> Self {
let session_id: String = Self::terminal_session_id();
match Self::restore_with_id(&session_id) {
Some(session) => session,
Some(session) => { session },
None => {
let session = Self::new(&session_id, vqd);
session.save().expect("Couldn't save new session");