diff --git a/crates/service/web/storyteller_web/src/main.rs b/crates/service/web/storyteller_web/src/main.rs index 9f28e84a45..387e47edb6 100644 --- a/crates/service/web/storyteller_web/src/main.rs +++ b/crates/service/web/storyteller_web/src/main.rs @@ -23,6 +23,7 @@ pub const RESERVED_SUBSTRINGS : &'static str = include_str!("../../../../../db/r pub mod billing; pub mod configs; pub mod http_server; +pub mod memory_cache; pub mod model; pub mod routes; pub mod server_state; diff --git a/crates/service/web/storyteller_web/src/memory_cache/mod.rs b/crates/service/web/storyteller_web/src/memory_cache/mod.rs new file mode 100644 index 0000000000..93062b0a1d --- /dev/null +++ b/crates/service/web/storyteller_web/src/memory_cache/mod.rs @@ -0,0 +1 @@ +pub mod model_token_to_info_cache; \ No newline at end of file diff --git a/crates/service/web/storyteller_web/src/memory_cache/model_token_to_info_cache.rs b/crates/service/web/storyteller_web/src/memory_cache/model_token_to_info_cache.rs new file mode 100644 index 0000000000..4e4abb3d81 --- /dev/null +++ b/crates/service/web/storyteller_web/src/memory_cache/model_token_to_info_cache.rs @@ -0,0 +1,57 @@ +use anyhow::anyhow; +use enums::by_table::generic_inference_jobs::inference_category::InferenceCategory; +use enums::by_table::generic_inference_jobs::inference_model_type::InferenceModelType; +use errors::AnyhowResult; +use std::collections::HashMap; +use std::sync::{Arc, LockResult, RwLock}; + +#[derive(Clone)] +pub struct ModelInfoLite { + pub inference_category: InferenceCategory, + pub model_type: InferenceModelType, +} + +/// Model token to type will never change, so we can cache them indefinitely. +/// This is for generic inference jobs only. +pub struct ModelTokenToInfoCache { + database: Arc>>, +} + +impl ModelTokenToInfoCache { + pub fn new() -> Self { + Self { + database: Arc::new(RwLock::new(HashMap::new())), + } + } + + pub fn insert_one(&self, token: &str, info: ModelInfoLite) -> AnyhowResult<()> { + match self.database.write() { + Err(err) => return Err(anyhow!("lock err: {:?}", err)), + Ok(mut lock) => { + lock.insert(token.to_string(), info); + } + } + Ok(()) + } + + pub fn insert_many(&self, records: Vec<(&str, ModelInfoLite)>) -> AnyhowResult<()> { + match self.database.write() { + Err(err) => return Err(anyhow!("lock err: {:?}", err)), + Ok(mut lock) => { + for (token, info) in records.into_iter() { + lock.insert(token.to_string(), info); + } + } + } + Ok(()) + } + + pub fn get_info(&self, token: &str) -> AnyhowResult> { + match self.database.read() { + Err(err) => return Err(anyhow!("lock err: {:?}", err)), + Ok(lock) => { + Ok(lock.get(token).map(|item| item.clone())) + } + } + } +}