From 621f54316a5154a29deb61f48f70bc2c95c0850b Mon Sep 17 00:00:00 2001 From: b1ek Date: Fri, 12 Apr 2024 22:36:36 +1000 Subject: [PATCH] add cache --- Cargo.lock | 82 ++++++++++++++++++++++++++++++++++++++++++++++++++ Cargo.toml | 2 ++ src/cache.rs | 85 ++++++++++++++++++++++++++++++++++++++++++++++++++++ src/main.rs | 35 ++++++++++++++++++---- 4 files changed, 198 insertions(+), 6 deletions(-) create mode 100644 src/cache.rs diff --git a/Cargo.lock b/Cargo.lock index 4ac88bf..a854ab8 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -17,6 +17,21 @@ version = "1.0.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f26201604c87b1e01bd3d98f8d5d9a8fcbb815e8cedb41ffccbeb4bf593a35fe" +[[package]] +name = "android-tzdata" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e999941b234f3131b00bc13c22d06e8c5ff726d1b6318ac7eb276997bbb4fef0" + +[[package]] +name = "android_system_properties" +version = "0.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "819e7219dbd41043ac279b19830f2efc897156490d7fd6ea916720117ee66311" +dependencies = [ + "libc", +] + [[package]] name = "anstream" version = "0.6.13" @@ -128,6 +143,21 @@ version = "1.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "baf1de4339761588bc0619e3cbc0120ee582ebb74b53b4efbf79117bd2da40fd" +[[package]] +name = "chrono" +version = "0.4.37" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8a0d04d43504c61aa6c7531f1871dd0d418d91130162063b789da00fd7057a5e" +dependencies = [ + "android-tzdata", + "iana-time-zone", + "js-sys", + "num-traits", + "serde", + "wasm-bindgen", + "windows-targets 0.52.4", +] + [[package]] name = "clap" version = "4.5.4" @@ -362,8 +392,10 @@ checksum = "d231dfb89cfffdbc30e7fc41579ed6066ad03abda9e567ccafae602b97ec5024" name = "hey" version = "0.1.0" dependencies = [ + "chrono", "clap", "femme", + "home", "log", "reqwest", "serde", @@ -371,6 +403,15 @@ dependencies = [ "tokio", ] +[[package]] +name = "home" +version = "0.5.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e3d1354bf6b7235cb4a0576c2619fd4ed18183f689b12b006a0ee7329eeff9a5" +dependencies = [ + "windows-sys 0.52.0", +] + [[package]] name = "http" version = "1.1.0" @@ -467,6 +508,29 @@ dependencies = [ "tracing", ] +[[package]] +name = "iana-time-zone" +version = "0.1.60" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e7ffbb5a1b541ea2561f8c41c087286cc091e21e556a4f09a8f6cbf17b69b141" +dependencies = [ + "android_system_properties", + "core-foundation-sys", + "iana-time-zone-haiku", + "js-sys", + "wasm-bindgen", + "windows-core", +] + +[[package]] +name = "iana-time-zone-haiku" +version = "0.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f31827a206f56af32e590ba56d5d2d085f558508192593743f16b2306495269f" +dependencies = [ + "cc", +] + [[package]] name = "idna" version = "0.5.0" @@ -596,6 +660,15 @@ dependencies = [ "tempfile", ] +[[package]] +name = "num-traits" +version = "0.2.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "da0df0e5185db44f69b44f26786fe401b6c293d1907744beaa7fa62b2e5a517a" +dependencies = [ + "autocfg", +] + [[package]] name = "num_cpus" version = "1.16.0" @@ -1396,6 +1469,15 @@ dependencies = [ "wasm-bindgen", ] +[[package]] +name = "windows-core" +version = "0.52.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "33ab640c8d7e35bf8ba19b884ba838ceb4fba93a4e8c65a9059d08afcfc683d9" +dependencies = [ + "windows-targets 0.52.4", +] + [[package]] name = "windows-sys" version = "0.48.0" diff --git a/Cargo.toml b/Cargo.toml index 8b8f3e2..b06b362 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -6,8 +6,10 @@ edition = "2021" # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html [dependencies] +chrono = { version = "0.4.37", features = ["serde"] } clap = { version = "4.5.4", features = ["derive"] } femme = "2.2.1" +home = "0.5.9" log = "0.4.21" reqwest = "0.12.3" serde = { version = "1.0.197", features = ["derive"] } diff --git a/src/cache.rs b/src/cache.rs new file mode 100644 index 0000000..eb6915c --- /dev/null +++ b/src/cache.rs @@ -0,0 +1,85 @@ + +use std::{env, error::Error, fs, io, path::PathBuf}; +use home::home_dir; +use serde::{Deserialize, Serialize}; + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct Cache { + pub last_vqd: String, + pub last_vqd_time: u64 +} + +impl Default for Cache { + fn default() -> Self { + Self { + last_vqd: "".into(), + last_vqd_time: 0 + } + } +} + +impl Cache { + + pub fn get_path>() -> T { + match env::var("HEY_CACHE_PATH") { + Ok(v) => v, + Err(_) => + match home_dir() { + Some(home) => home.join(".cache/hey").as_os_str().as_encoded_bytes().iter().map(|x| char::from(*x)).collect::(), + None => panic!("Cannot detect your home directory!") + } + }.into() + } + + pub fn get_file_name>() -> T { + match env::var("HEY_CACHE_PATH") { + Ok(v) => v, + Err(_) => "cache.json".into() + }.into() + } + + fn ensure_dir_exists() -> io::Result<()> { + let path: PathBuf = Self::get_path(); + if ! path.is_dir() { fs::create_dir_all(path)? } + Ok(()) + } + + pub fn load() -> Result> { + let path: PathBuf = Self::get_path(); + Self::ensure_dir_exists()?; + + let file_path = path.join(Self::get_file_name::()); + if ! file_path.is_file() { + let def = Self::default(); + def.save()?; + Ok(def) + } else { + let file = fs::read_to_string(file_path)?; + Ok(serde_json::from_str(&file)?) + } + } + + pub fn save(self: &Self) -> Result<(), Box> { + let path: PathBuf = Self::get_path(); + Self::ensure_dir_exists()?; + + let file_path = path.join(Self::get_file_name::()); + fs::write(file_path, serde_json::to_string_pretty(self)?)?; + Ok(()) + } + + pub fn set_last_vqd>(self: &mut Self, vqd: T) -> Result<(), Box> { + self.last_vqd = vqd.into(); + self.last_vqd_time = chrono::Local::now().timestamp_millis() as u64; + self.save()?; + Ok(()) + } + + pub fn get_last_vqd<'a, T: From<&'a String>>(self: &'a Self) -> Option { + if self.last_vqd_time - (chrono::Local::now().timestamp_millis() as u64) < 60000 { + Some((&self.last_vqd).into()) + } else { + None + } + } +} \ No newline at end of file diff --git a/src/main.rs b/src/main.rs index 6949221..999f6ee 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,12 +1,17 @@ -use std::{collections::HashMap, error::Error, hash::Hash, process::exit}; +use std::{error::Error, process::exit}; use reqwest::{header::{HeaderMap, HeaderValue}, Client}; use serde::{Deserialize, Serialize}; use clap::Parser; +use crate::cache::Cache; + +mod cache; + const GREEN: &str = "\x1b[1;32m"; const RED: &str = "\x1b[1;31m"; +const WARN: &str = "\x1b[1;34m"; const RESET: &str = "\x1b[0m"; #[derive(Debug, Clone, Serialize, Deserialize)] @@ -88,7 +93,7 @@ async fn get_vqd(cli: &Client) -> Result> { } } -async fn get_res(cli: &Client, query: String, vqd: String) { +async fn get_res<'a>(cli: &Client, query: String, vqd: String, cache: &'a mut Cache) { let payload = ChatPayload { model: "claude-instant-1.2".into(), messages: vec![ ChatMessagePayload { role: "user".into(), content: query } ] @@ -101,11 +106,22 @@ async fn get_res(cli: &Client, query: String, vqd: String) { // .headers(get_headers()) .header("Content-Type", "application/json") .header("User-Agent", "Mozilla/5.0 (X11; Linux x86_64; rv:124.0) Gecko/20100101 Firefox/124.0") - .header("x-vqd-4", vqd) + .header("x-vqd-4", vqd.clone()) .body(payload) .build().unwrap(); let mut res = cli.execute(req).await.unwrap(); + let new_vqd = res.headers().iter().find(|x| x.0 == "x-vqd-4"); + let vqd_set_res = + if let Some(new_vqd) = new_vqd { + cache.set_last_vqd(new_vqd.1.as_bytes().iter().map(|x| char::from(*x)).collect::()) + } else { + eprintln!("{WARN}Warn: DuckDuckGo did not return new VQD. Ignore this if everything else is ok.{RESET}"); + cache.set_last_vqd(vqd.clone()) + }; + if let Err(err) = vqd_set_res { + eprintln!("{WARN}Warn: Could not save VQD to cache: {err}{RESET}"); + } while let Some(chunk) = res.chunk().await.unwrap() { @@ -144,9 +160,16 @@ async fn main() { let args = Args::parse(); let query = args.query.join(" "); + let mut cache = Cache::load().unwrap(); + let cli = Client::new(); - // simulate_browser_reqs(&cli).await.unwrap(); - let vqd = get_vqd(&cli).await.unwrap(); - get_res(&cli, query, vqd).await; + simulate_browser_reqs(&cli).await.unwrap(); + + let vqd = match cache.get_last_vqd() { + Some(v) => { println!("using cached vqd"); v}, + None => get_vqd(&cli).await.unwrap() + }; + + get_res(&cli, query, vqd, &mut cache).await; }