use std::path::Path;
use crate::config::Config;
use crate::embedding::EmbeddingEngine;
use crate::errors::Error;
use crate::sqlite::Database;
pub const MAX_INPUT_LENGTH: usize = 100_000;
pub const MAX_SEARCH_LIMIT: usize = 10_000;
pub(crate) fn validate_limit(limit: usize) -> Result<(), Error> {
if limit == 0 {
return Err(Error::InvalidInput(
"Limit must be greater than 0".to_string(),
));
}
if limit > MAX_SEARCH_LIMIT {
return Err(Error::InvalidInput(format!(
"Limit {} exceeds maximum allowed ({})",
limit, MAX_SEARCH_LIMIT
)));
}
Ok(())
}
pub struct MemoryStore {
pub(crate) db: Database,
pub(crate) embedder: Option<EmbeddingEngine>,
pub(crate) model_id: String,
pub(crate) config: Config,
#[cfg(test)]
pub(crate) test_embedder: Option<TestEmbedder>,
}
#[allow(dead_code)]
pub(crate) type TestEmbedder = Box<dyn Fn(&str) -> Result<Vec<f32>, Error> + Send + Sync>;
impl MemoryStore {
pub fn new(db_path: &Path, model_id: &str, config: Config) -> Result<Self, Error> {
use std::path::Component;
for component in db_path.components() {
if matches!(component, Component::ParentDir) {
return Err(Error::Config(
"Invalid database path: contains '..' which may escape the intended directory"
.to_string(),
));
}
}
let db_real_path = if db_path.exists() {
std::fs::canonicalize(db_path).map_err(|e| {
Error::Config(format!(
"Invalid database path: cannot canonicalize existing path: {}",
e
))
})?
} else {
let parent = db_path.parent().ok_or_else(|| {
Error::Config("Invalid database path: no parent directory".to_string())
})?;
let canonical_parent = std::fs::canonicalize(parent).map_err(|e| {
Error::Config(format!(
"Invalid database path: parent directory not accessible: {}",
e
))
})?;
let filename = db_path
.file_name()
.ok_or_else(|| Error::Config("Invalid database path: no filename".to_string()))?;
canonical_parent.join(filename)
};
let db = Database::open(&db_real_path)?;
Ok(MemoryStore {
db,
embedder: None,
model_id: model_id.to_string(),
config,
#[cfg(test)]
test_embedder: None,
})
}
pub(crate) fn embedder(&mut self) -> Result<&mut EmbeddingEngine, Error> {
if self.embedder.is_none() {
self.embedder = Some(EmbeddingEngine::new(&self.model_id)?);
}
Ok(self.embedder.as_mut().unwrap())
}
#[allow(dead_code)]
pub(crate) fn set_preinitialized_embedder(&mut self, engine: EmbeddingEngine) {
self.embedder = Some(engine);
}
pub(crate) fn validate_input_length(text: &str) -> Result<(), Error> {
if text.trim().is_empty() {
return Err(Error::EmptyInput);
}
if text.len() > MAX_INPUT_LENGTH {
return Err(Error::InputTooLong {
max_length: MAX_INPUT_LENGTH,
actual_length: text.len(),
});
}
Ok(())
}
#[cfg(test)]
pub(crate) fn from_db(db: Database, config: Config) -> Self {
MemoryStore {
db,
embedder: None,
model_id: String::new(),
config,
test_embedder: None,
}
}
#[cfg(test)]
pub(crate) fn from_db_with_test_embedder(db: Database) -> Self {
let mut store = Self::from_db(db, Config::default());
store.test_embedder = Some(Box::new(crate::memory::crud::test_fake_embedder));
store
}
#[cfg(test)]
pub(crate) fn test_store() -> Self {
use tempfile::TempDir;
let dir = TempDir::new().expect("create temp dir");
let path = dir.path().join("test.db");
std::mem::forget(dir);
let db = Database::open(&path).expect("open test database");
Self::from_db_with_test_embedder(db)
}
#[cfg(test)]
pub(crate) fn test_db_path() -> std::path::PathBuf {
use tempfile::TempDir;
let dir = TempDir::new().expect("create temp dir");
let path = dir.path().join("test.db");
std::mem::forget(dir);
path
}
}