From bc74b114eceed5cede548dc9f213591266d5bcdc Mon Sep 17 00:00:00 2001 From: Alex Jones Date: Tue, 31 Oct 2023 20:05:46 +0000 Subject: [PATCH] feat: working caching Signed-off-by: Alex Jones --- Cargo.lock | 1 + Cargo.toml | 3 +- src/analyze/mod.rs | 80 +++++++++++++++++++++++++++++---------------- src/bedrock/mod.rs | 1 + src/config/cache.rs | 49 +++++++++++++++++++++++++++ src/config/mod.rs | 9 ++++- 6 files changed, 112 insertions(+), 31 deletions(-) create mode 100644 src/config/cache.rs diff --git a/Cargo.lock b/Cargo.lock index a76b8e4..e479513 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1154,6 +1154,7 @@ dependencies = [ "aws-sdk-s3", "aws-sdk-sts 0.33.0", "aws-smithy-runtime-api", + "base64", "clap", "colored", "futures", diff --git a/Cargo.toml b/Cargo.toml index 07c90a3..bfb7f83 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "isotope" -version = "0.0.2" +version = "0.0.3" repository = "https://github.com/isotope-rs/isotope.git" edition = "2021" description = "Isotope scans AWS services and makes suggestions on how to improve them using Artificial Intelligence." @@ -37,6 +37,7 @@ aws-smithy-runtime-api = { version = "0.56.1", features = ["client"] } unescape = "0.1.0" aws-sdk-bedrock = "0.3.0" aws-sdk-bedrockruntime = "0.3.0" +base64 = { version = "0.21.4", features = [] } # Config for 'cargo dist' [workspace.metadata.dist] diff --git a/src/analyze/mod.rs b/src/analyze/mod.rs index c8cbf89..b8ab800 100644 --- a/src/analyze/mod.rs +++ b/src/analyze/mod.rs @@ -1,19 +1,20 @@ -use std::collections::hash_map::Entry; -use std::collections::HashMap; -use std::ops::DerefMut; use crate::analyzer::analyzer_trait::Analyzer; use crate::analyzer::types::AnalysisResults; use crate::config::Conf; -use crate::{analyzer, Args, bedrock}; +use crate::{analyzer, bedrock, Args}; use crate::{config, outputs}; use aws_config::meta::region::{ProvideRegion, RegionProviderChain}; +use colored::Colorize; +use std::collections::hash_map::Entry; +use std::collections::HashMap; +use std::ops::DerefMut; use std::sync::mpsc; use std::sync::mpsc::{Receiver, Sender}; -use colored::Colorize; pub async fn run_analysis(args: &Args) { let mut conf: Conf = config::Conf { cloud: String::new(), + stored_advice: HashMap::new(), }; let c = config::get_or_create_config(); match c { @@ -26,14 +27,24 @@ pub async fn run_analysis(args: &Args) { // Setup bedrock let bedrockClient = bedrock::BedrockClient::new(config.clone()); - println!("Current AWS region: {}", RegionProviderChain::default_provider().region().await.unwrap().as_ref().yellow()); + println!( + "Current AWS region: {}", + RegionProviderChain::default_provider() + .region() + .await + .unwrap() + .as_ref() + .yellow() + ); // Create channels let (tx, rx): (Sender>, Receiver>) = mpsc::channel(); let analyzers: Vec> = analyzer::generate_analyzers(config.clone()); match &args.analyzer { Some(analyzer_arg) => { - let filtered_analyzer = &analyzers.iter().find(|x| x.get_name().as_str() == analyzer_arg); + let filtered_analyzer = &analyzers + .iter() + .find(|x| x.get_name().as_str() == analyzer_arg); match filtered_analyzer { Some(x) => { let thread_tx = tx.clone(); @@ -80,39 +91,50 @@ pub async fn run_analysis(args: &Args) { task.await.unwrap(); } - let mut processed_results: HashMap> = HashMap::new(); + let mut processed_results: HashMap> = HashMap::new(); // generate Vectors aligned to each analyzer type - // Feed results into Bedrock for mut res in results { if !res.message.is_empty() { - let result = bedrockClient.enrich(res.message.clone()).await; - // TODO: missing step to copy the bedrock result into res - match result { - Ok(x) => { - res.advice = x.clone(); - // pass ownership over of advice - // check if the processed results analyzer exists as key - // upsert the analysis result into the vector - match processed_results.entry(x) { - Entry::Occupied(mut e) => { - e.get_mut().push(res); - }, - Entry::Vacant(e) => { - e.insert(vec![res]); - }, - } + // Check if the data is in the cache + match conf.fetch_from_cache(&res.message) { + Some(x) => { + res.advice = x.clone() }, - Err(e) => ( + None => { + let result = bedrockClient.enrich(res.message.clone()).await; + // TODO: missing step to copy the bedrock result into res + match result { + Ok(x) => { + res.advice = x.clone(); + // upsert into the cache for next time + conf.clone().upsert_into_cache(&res.message,&x); + // pass ownership over of advice + // check if the processed results analyzer exists as key + // upsert the analysis result into the vector - ), + } + Err(e) => (), + } + } + } + match processed_results.entry(res.analyzer_name.clone()) { + Entry::Occupied(mut e) => { + e.get_mut().push(res); + } + Entry::Vacant(e) => { + e.insert(vec![res]); + } } } } - // + match args.json { Some(x) => { - let mut p = outputs::Processor::new(processed_results, Some(outputs::Configuration::new(x))); + let mut p = outputs::Processor::new( + processed_results, + Some(outputs::Configuration::new(x)), + ); p.print(); } None => { diff --git a/src/bedrock/mod.rs b/src/bedrock/mod.rs index 6001dcc..ef034cb 100644 --- a/src/bedrock/mod.rs +++ b/src/bedrock/mod.rs @@ -8,6 +8,7 @@ use std::io::{stdout, Write}; use aws_sdk_config::config::Region; mod prompt; + #[derive(Serialize)] struct ClaudParams { prompt: String, diff --git a/src/config/cache.rs b/src/config/cache.rs new file mode 100644 index 0000000..b501704 --- /dev/null +++ b/src/config/cache.rs @@ -0,0 +1,49 @@ +use std::error::Error; +use std::io::Bytes; +use crate::config::{Conf, save_config}; +use base64::{Engine as _, engine::{self, general_purpose}, alphabet}; + +impl Conf { + fn encode_cache_key(&self,raw_cache_key: &str) -> String { + general_purpose::STANDARD.encode(raw_cache_key) + } + + pub fn upsert_into_cache(mut self, raw_cache_key: &str, payload: &str) -> Result<(), Box> { + + let encoded_key = self.encode_cache_key(raw_cache_key); + let encoded_payload = self.encode_cache_key(payload); + // append the new data to the hashmap + self.stored_advice.insert(encoded_key, + encoded_payload); + + // Save the config + save_config(self)?; + Ok(()) + } + pub fn fetch_from_cache(&self,raw_cache_key: &str) -> Option { + let encoded_key = self.encode_cache_key(raw_cache_key); + let found_keys: Vec<&String> = self.stored_advice.iter().filter_map(|(k,v)| + if k == encoded_key.as_str() {Some(v)} else {None}).collect(); + + // TODO: verify this + if found_keys.is_empty() { + return None + } + + let mut found = String::new(); + let decoded_bytes = match general_purpose::STANDARD.decode(found_keys.first().unwrap().to_string()) { + Ok(bytes) => { + let decoded_string = match String::from_utf8(bytes) { + Ok(s) => found = s, + Err(e) => { + eprintln!("Error converting bytes to String: {}", e); + } + }; + }, + Err(e) => { + eprintln!("Error decoding Base64: {}", e); + } + }; + Some(found) + } +} \ No newline at end of file diff --git a/src/config/mod.rs b/src/config/mod.rs index a33dabb..a155dfc 100644 --- a/src/config/mod.rs +++ b/src/config/mod.rs @@ -1,3 +1,6 @@ +mod cache; + +use std::collections::HashMap; use serde::{Deserialize, Serialize}; use simple_home_dir::*; use std::error::Error; @@ -7,9 +10,12 @@ use std::path::Path; pub const CONFFILE: &str = "isotope.config"; -#[derive(Serialize, Deserialize, Debug)] +#[derive(Serialize, Deserialize, Debug, Clone)] pub struct Conf { pub cloud: String, + // This is stored as a hashed string representing the issue e.g. S3 bucket public + // With a subsequent value (also b64 encoded) + pub stored_advice: HashMap } pub fn get_conf_path() -> String { let home = home_dir().unwrap(); @@ -33,6 +39,7 @@ pub fn get_or_create_config() -> Result> { let p = get_conf_path(); let c = Conf { cloud: String::new(), + stored_advice: HashMap::new(), }; if !Path::new(&p).exists() { let mut f = File::create(&p)?;