use std::fs;
use std::path::PathBuf;
use crate::model_source::domain::{ModelInfo, ModelSource};
use crate::model_source::dtos::ModelInfoDto;
use crate::model_source::errors::ModelSourceError;
use crate::model_source::errors::ModelSourceError::ModelDownloadError;
use crate::model_source::file_utils::download_file_with_integrity;
pub struct Gpt4AllModelSource {
pub url: String,
pub folder: PathBuf
}
impl Gpt4AllModelSource {
async fn _fetch_available_models(&self) -> Result<Vec<ModelInfoDto>, ModelSourceError> {
let response = reqwest::get(&self.url).await.map_err(|_| ModelSourceError::FetchError)?;
let model_info_dto: Vec<ModelInfoDto> = response
.json().await
.map_err(|_| ModelSourceError::ModelInfoParsingError)?;
Ok(model_info_dto)
}
}
impl ModelSource for Gpt4AllModelSource {
async fn fetch_all_models(&self) -> Result<Vec<ModelInfo>, ModelSourceError> {
let model_info_dto = self._fetch_available_models().await?;
let model_info = model_info_dto
.into_iter()
.map(ModelInfo::from)
.collect();
Ok(model_info)
}
async fn download(&self, model_file_name: &str) -> Result<String, ModelSourceError> {
let available_models = self._fetch_available_models().await?;
let model_info = available_models
.iter().find(|&model_info| { model_info.file_name == model_file_name })
.ok_or(ModelSourceError::ModelNotFound)?;
let path = self.folder.join(&model_info.file_name);
let path_str = path
.to_str()
.ok_or(ModelDownloadError)?
.to_string();
if path.exists() { return Ok(path_str) }
download_file_with_integrity(&model_info.url, &path_str, &model_info.md5sum)
.await.map_err(|_| { ModelDownloadError })?;
Ok(path_str)
}
async fn get_local_models(&self) -> Result<Vec<String>, ModelSourceError> {
let mut local_models = Vec::new();
if let Ok(entries) = fs::read_dir(&self.folder) {
for entry in entries {
if let Ok(entry) = entry {
if let Some(file_name) = entry.file_name().to_str() {
if entry.file_type().map(|ft| ft.is_file()).unwrap_or(false)
&& file_name.ends_with(".gguf") {
local_models.push(file_name.to_string());
}
}
}
}
} else {
return Err(ModelSourceError::FetchError);
}
Ok(local_models)
}
}