blek
/
hey
3
3
Fork 0
Code Issues 2 Pull Requests Packages Projects Releases 4 Wiki Activity

add cache

This commit is contained in:
b1ek 2024-04-12 22:36:36 +10:00
parent 5f4dbbbdfd
commit 621f54316a
Signed by: blek
GPG Key ID: 14546221E3595D0C
4 changed files with 198 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

82
Cargo.lock generated
View File

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

2
Cargo.toml
View File

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

85
src/cache.rs Normal file
View File

@ -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: From<String>>() -> 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::<String>(),
None => panic!("Cannot detect your home directory!")
}
}.into()
}
pub fn get_file_name<T: From<String>>() -> 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<Self, Box<dyn Error>> {
let path: PathBuf = Self::get_path();
Self::ensure_dir_exists()?;
let file_path = path.join(Self::get_file_name::<PathBuf>());
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<dyn Error>> {
let path: PathBuf = Self::get_path();
Self::ensure_dir_exists()?;
let file_path = path.join(Self::get_file_name::<PathBuf>());
fs::write(file_path, serde_json::to_string_pretty(self)?)?;
Ok(())
}
pub fn set_last_vqd<T: Into<String>>(self: &mut Self, vqd: T) -> Result<(), Box<dyn Error>> {
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<T> {
if self.last_vqd_time - (chrono::Local::now().timestamp_millis() as u64) < 60000 {
Some((&self.last_vqd).into())
} else {
None
}
}
}

35
src/main.rs
View File

@ -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<String, Box<dyn Error>> {
}
}
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::<String>())
} 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;
}