//! Infrastructure adapters for local model servers. use std::collections::HashMap; use std::path::{Path, PathBuf}; use std::sync::{Arc, Mutex}; use std::time::Duration; use async_trait::async_trait; use futures_util::StreamExt; use serde::{Deserialize, Serialize}; use tokio::io::AsyncWriteExt; use tokio::process::{Child, Command}; use tokio::sync::Mutex as AsyncMutex; use domain::model_server::{ ExecutablePath, HfModelRef, LlamaCppOptions, LocalModelRef, LocalModelServerConfig, LocalModelServerKind, ModelPath, ModelServerEndpoint, ModelSource, }; use domain::ports::{ FileSystem, ManagedProcess, ManagedProcessHandle, ModelArtifactCancel, ModelArtifactDownloader, ModelArtifactProgress, ModelArtifactResolution, ModelArtifactState, ModelServerArgv, ModelServerError, ModelServerProbe, ModelServerRegistry, ModelServerRuntime, ProcessStatus, RemotePath, SpawnSpec, }; use domain::{LocalModelServerId, ProjectPath, StopPolicy}; /// HTTP readiness probe for OpenAI-compatible servers. #[derive(Clone)] pub struct HttpOpenAiCompatibleProbe { client: reqwest::Client, } impl Default for HttpOpenAiCompatibleProbe { fn default() -> Self { Self::new(Duration::from_secs(2)) } } impl HttpOpenAiCompatibleProbe { /// Builds the probe with a request timeout. #[must_use] pub fn new(timeout: Duration) -> Self { let client = reqwest::Client::builder() .timeout(timeout) .build() .expect("reqwest client with timeout builds"); Self { client } } } #[async_trait] impl ModelServerProbe for HttpOpenAiCompatibleProbe { async fn probe( &self, endpoint: &ModelServerEndpoint, ) -> Result { let models = format!("{}/models", endpoint.base_url.trim_end_matches('/')); if is_ready(self.client.get(models).send().await) { return Ok(domain::ModelServerStatus::ReadyReused); } let root = endpoint.base_url.trim_end_matches("/v1").to_owned(); if is_ready(self.client.get(root).send().await) { return Ok(domain::ModelServerStatus::ReadyReused); } Ok(domain::ModelServerStatus::Unreachable) } } fn is_ready(result: Result) -> bool { result .map(|response| response.status().is_success()) .unwrap_or(false) } /// Hugging Face artifact resolver/downloader backed by an IdeA-owned cache dir. #[derive(Clone)] pub struct HfModelArtifactDownloader { cache_dir: PathBuf, client: reqwest::Client, repo_locks: Arc>>>>, } impl HfModelArtifactDownloader { /// Builds a downloader that stores artifacts under `cache_dir`. #[must_use] pub fn new(cache_dir: impl Into) -> Self { Self { cache_dir: cache_dir.into(), client: reqwest::Client::new(), repo_locks: Arc::new(Mutex::new(HashMap::new())), } } /// Deterministic cache path for a Hugging Face model reference. /// /// This is the local path used for a single **merged** (non-sharded) GGUF file, renamed to a /// synthetic `{quant}.gguf` name. Sharded files are never renamed this way: llama.cpp /// discovers multi-part siblings by pattern-matching the *original* HF filename in the same /// directory, so shard files keep their remote name (see `local_path_for_filename`). #[must_use] pub fn cache_path_for(&self, repo: &HfModelRef) -> PathBuf { let (base, quant) = split_hf_ref(repo); let file_stem = quant.unwrap_or("model"); self.cache_dir .join(base.replace('/', "--")) .join(format!("{}.gguf", safe_cache_component(file_stem))) } fn cache_dir_for(&self, repo: &HfModelRef) -> PathBuf { let (base, _) = split_hf_ref(repo); self.cache_dir.join(base.replace('/', "--")) } /// Local path for a remote filename that is part of a multi-shard set: the original name is /// preserved (sanitized) so llama.cpp's shard autoload can find siblings by pattern. fn local_path_for_shard(&self, repo: &HfModelRef, remote_filename: &str) -> PathBuf { self.cache_dir_for(repo) .join(safe_cache_component(remote_filename)) } fn manifest_path_for(&self, repo: &HfModelRef) -> PathBuf { let (_, quant) = split_hf_ref(repo); let file_stem = quant.unwrap_or("model"); self.cache_dir_for(repo) .join(format!("{}.manifest.json", safe_cache_component(file_stem))) } /// Reads a manifest of a previously-downloaded shard set, if present and complete on disk. fn cached_shard_set(&self, repo: &HfModelRef) -> Option> { let manifest_path = self.manifest_path_for(repo); let raw = std::fs::read(&manifest_path).ok()?; let manifest: ShardManifest = serde_json::from_slice(&raw).ok()?; let dir = self.cache_dir_for(repo); let paths: Vec = manifest .files .iter() .map(|filename| dir.join(filename)) .collect(); if paths.iter().all(|path| path.is_file()) { Some(paths) } else { None } } fn write_shard_manifest(&self, repo: &HfModelRef, filenames: &[String]) -> std::io::Result<()> { let manifest_path = self.manifest_path_for(repo); let manifest = ShardManifest { files: filenames.to_vec(), }; let json = serde_json::to_vec(&manifest)?; std::fs::write(manifest_path, json) } fn lock_for(&self, repo: &HfModelRef) -> Arc> { let mut locks = self.repo_locks.lock().expect("repo locks mutex poisoned"); Arc::clone( locks .entry(repo.as_str().to_owned()) .or_insert_with(|| Arc::new(AsyncMutex::new(()))), ) } fn cached_state(&self, repo: &HfModelRef) -> Result { let merged_path = self.cache_path_for(repo); if merged_path.is_file() { return Ok(ModelArtifactState::Downloaded { size_bytes: Some(file_size(&merged_path)?), path: model_path_from_pathbuf(merged_path)?, }); } if let Some(paths) = self.cached_shard_set(repo) { if let Some(first) = paths.first() { return Ok(ModelArtifactState::Downloaded { size_bytes: Some(paths_size(&paths)?), path: model_path_from_pathbuf(first.clone())?, }); } } Ok(ModelArtifactState::Missing) } fn delete_cached(&self, repo: &HfModelRef) -> Result<(), ModelServerError> { let merged_path = self.cache_path_for(repo); if merged_path.is_file() { remove_file_if_exists(&merged_path)?; } let manifest_path = self.manifest_path_for(repo); if let Ok(raw) = std::fs::read(&manifest_path) { let manifest: ShardManifest = serde_json::from_slice(&raw).map_err(|e| ModelServerError::Store(e.to_string()))?; let dir = self.cache_dir_for(repo); for filename in manifest.files { remove_file_if_exists(&dir.join(filename))?; } remove_file_if_exists(&manifest_path)?; } let repo_dir = self.cache_dir_for(repo); if repo_dir.is_dir() && std::fs::read_dir(&repo_dir) .map(is_empty_dir) .unwrap_or(false) { std::fs::remove_dir(&repo_dir).map_err(|e| ModelServerError::Store(e.to_string()))?; } Ok(()) } async fn resolve_remote_filenames( &self, repo: &HfModelRef, ) -> Result, ModelServerError> { let (base, quant) = split_hf_ref(repo); let url = format!("https://huggingface.co/api/models/{base}"); let response = self .client .get(url) .send() .await .map_err(|e| ModelServerError::Probe(e.to_string()))?; if !response.status().is_success() { return Err(ModelServerError::Probe(format!( "huggingface model metadata returned {}", response.status() ))); } let metadata: HfModelMetadata = response .json() .await .map_err(|e| ModelServerError::Probe(e.to_string()))?; let files = select_gguf_files(&metadata.siblings, quant); if files.is_empty() { return Err(ModelServerError::PathNotAccessible(format!( "no matching .gguf artifact found for {}", repo.as_str() ))); } Ok(files) } } #[derive(Debug, Default, Serialize, Deserialize)] struct ShardManifest { files: Vec, } #[async_trait] impl ModelArtifactDownloader for HfModelArtifactDownloader { async fn hf_model_state( &self, repo: &HfModelRef, ) -> Result { let lock = self.lock_for(repo); let _guard = lock.lock().await; self.cached_state(repo) } async fn resolve_hf_model( &self, repo: &HfModelRef, progress: std::sync::Arc, cancel: ModelArtifactCancel, ) -> Result { let lock = self.lock_for(repo); let _guard = lock.lock().await; if cancel.is_cancelled() { return Err(ModelServerError::Cancelled); } // Fast path 1: a single merged file was already downloaded (synthetic name), no network needed. let merged_path = self.cache_path_for(repo); if merged_path.is_file() { return Ok(ModelArtifactResolution { path: model_path_from_pathbuf(merged_path)?, cache_hit: true, }); } // Fast path 2: a previously-downloaded shard set is fully present on disk, no network needed. if let Some(paths) = self.cached_shard_set(repo) { if let Some(first) = paths.into_iter().next() { return Ok(ModelArtifactResolution { path: model_path_from_pathbuf(first)?, cache_hit: true, }); } } let filenames = self.resolve_remote_filenames(repo).await?; if cancel.is_cancelled() { return Err(ModelServerError::Cancelled); } let is_shard_set = filenames.len() > 1; let targets: Vec = if is_shard_set { filenames .iter() .map(|filename| self.local_path_for_shard(repo, filename)) .collect() } else { vec![merged_path.clone()] }; let (base, _) = split_hf_ref(repo); let total_files = filenames.len() as u64; for (index, (filename, target)) in filenames.iter().zip(targets.iter()).enumerate() { if cancel.is_cancelled() { return Err(ModelServerError::Cancelled); } let url = format!( "https://huggingface.co/{base}/resolve/main/{}", url_path_segment(filename) ); let response = self .client .get(url) .send() .await .map_err(|e| ModelServerError::Probe(e.to_string()))?; if !response.status().is_success() { return Err(ModelServerError::Probe(format!( "huggingface artifact download returned {}", response.status() ))); } if let Some(parent) = target.parent() { tokio::fs::create_dir_all(parent) .await .map_err(|e| ModelServerError::Store(e.to_string()))?; } let tmp_path = PathBuf::from(format!("{}.part", target.to_string_lossy())); let mut file = tokio::fs::File::create(&tmp_path) .await .map_err(|e| ModelServerError::Store(e.to_string()))?; let total = response.content_length(); let mut downloaded = 0_u64; let mut stream = response.bytes_stream(); while let Some(chunk) = stream.next().await { if cancel.is_cancelled() { let _ = tokio::fs::remove_file(&tmp_path).await; return Err(ModelServerError::Cancelled); } let chunk = chunk.map_err(|e| ModelServerError::Probe(e.to_string()))?; file.write_all(&chunk) .await .map_err(|e| ModelServerError::Store(e.to_string()))?; downloaded += chunk.len() as u64; progress(ModelArtifactProgress { downloaded_bytes: Some(downloaded + index as u64 * total.unwrap_or(0)), total_bytes: total.map(|t| t * total_files), source: Some(repo.as_str().to_owned()), }); } file.flush() .await .map_err(|e| ModelServerError::Store(e.to_string()))?; drop(file); tokio::fs::rename(&tmp_path, target) .await .map_err(|e| ModelServerError::Store(e.to_string()))?; } if is_shard_set { self.write_shard_manifest(repo, &filenames) .map_err(|e| ModelServerError::Store(e.to_string()))?; } Ok(ModelArtifactResolution { path: model_path_from_pathbuf(targets[0].clone())?, cache_hit: false, }) } async fn delete_hf_model(&self, repo: &HfModelRef) -> Result<(), ModelServerError> { let lock = self.lock_for(repo); let _guard = lock.lock().await; self.delete_cached(repo) } } #[derive(Debug, Deserialize)] struct HfModelMetadata { #[serde(default)] siblings: Vec, } #[derive(Debug, Deserialize)] struct HfSibling { #[serde(rename = "rfilename")] filename: String, } fn split_hf_ref(repo: &HfModelRef) -> (&str, Option<&str>) { repo.as_str() .split_once(':') .map_or((repo.as_str(), None), |(base, quant)| (base, Some(quant))) } /// Matches the llama.cpp / HF multi-part GGUF naming convention, e.g. /// `model-q5_k_m-00001-of-00002.gguf`. Captures the shard index and total shard count. fn shard_suffix() -> &'static regex::Regex { static RE: std::sync::OnceLock = std::sync::OnceLock::new(); RE.get_or_init(|| regex::Regex::new(r"-(\d{5})-of-(\d{5})\.gguf$").expect("valid regex")) } fn shard_index(filename: &str) -> Option { shard_suffix() .captures(filename) .and_then(|caps| caps.get(1)) .and_then(|m| m.as_str().parse().ok()) } /// Resolves the ordered set of remote filenames to download for a quantization: the merged /// (non-sharded) file if one matches, otherwise all matching shards sorted by index. A merged /// file always takes priority over a shard set when both exist in the repo, since it is what /// `llama-server --model` expects and a shard set would otherwise be picked first by naive /// alphabetical sort ('-' sorts before '.' in ASCII). fn select_gguf_files(siblings: &[HfSibling], quant: Option<&str>) -> Vec { let ggufs: Vec<&str> = siblings .iter() .map(|sibling| sibling.filename.as_str()) .filter(|filename| filename.ends_with(".gguf")) .collect(); let matching: Vec<&str> = match quant { None => ggufs, Some(quant) => { let quant = quant.to_ascii_lowercase(); ggufs .into_iter() .filter(|filename| filename.to_ascii_lowercase().contains(&quant)) .collect() } }; let (mut merged, mut shards): (Vec<&str>, Vec<&str>) = (Vec::new(), Vec::new()); for filename in matching { if shard_index(filename).is_some() { shards.push(filename); } else { merged.push(filename); } } if let Some(first_merged) = { merged.sort_unstable(); merged.first() } { return vec![(*first_merged).to_owned()]; } shards.sort_unstable_by_key(|filename| shard_index(filename).unwrap_or(u32::MAX)); shards.into_iter().map(str::to_owned).collect() } fn safe_cache_component(raw: &str) -> String { raw.chars() .map(|ch| match ch { 'A'..='Z' | 'a'..='z' | '0'..='9' | '.' | '_' | '-' => ch, _ => '_', }) .collect() } fn url_path_segment(raw: &str) -> String { raw.split('/') .map(url_component) .collect::>() .join("/") } fn url_component(raw: &str) -> String { raw.bytes() .flat_map(|byte| match byte { b'A'..=b'Z' | b'a'..=b'z' | b'0'..=b'9' | b'.' | b'_' | b'-' => { vec![byte as char] } other => format!("%{other:02X}").chars().collect(), }) .collect() } fn model_path_from_pathbuf(path: PathBuf) -> Result { ModelPath::new(path.to_string_lossy().into_owned()) .map_err(|e| ModelServerError::Invalid(e.to_string())) } fn file_size(path: &Path) -> Result { std::fs::metadata(path) .map(|metadata| metadata.len()) .map_err(|e| ModelServerError::Store(e.to_string())) } fn paths_size(paths: &[PathBuf]) -> Result { paths .iter() .map(|path| file_size(path)) .try_fold(0_u64, |acc, size| size.map(|size| acc.saturating_add(size))) } fn remove_file_if_exists(path: &Path) -> Result<(), ModelServerError> { match std::fs::remove_file(path) { Ok(()) => Ok(()), Err(err) if err.kind() == std::io::ErrorKind::NotFound => Ok(()), Err(err) => Err(ModelServerError::Store(err.to_string())), } } fn is_empty_dir(entries: std::fs::ReadDir) -> bool { entries.into_iter().next().is_none() } /// Builds `llama-server` argv without shell interpolation. #[derive(Debug, Default, Clone, Copy)] pub struct LlamaCppRuntime; impl LlamaCppRuntime { /// Creates the runtime. #[must_use] pub const fn new() -> Self { Self } } impl ModelServerRuntime for LlamaCppRuntime { fn build_argv( &self, config: &LocalModelServerConfig, ) -> Result { if config.kind != LocalModelServerKind::LlamaCpp { return Err(ModelServerError::Invalid( "only LlamaCpp local model servers are supported".to_owned(), )); } let command = resolve_binary( config .binary .as_ref() .map(|binary| binary.as_str()) .unwrap_or("llama-server"), )?; let mut args = Vec::new(); match config .model .source .as_ref() .ok_or_else(|| ModelServerError::PathNotAccessible("model.source missing".to_owned()))? { ModelSource::LocalPath { path } => { args.push("--model".to_owned()); args.push(path.as_str().to_owned()); } ModelSource::HuggingFace { repo } => { args.push("-hf".to_owned()); args.push(repo.as_str().to_owned()); } } args.extend([ "--port".to_owned(), config.endpoint.port.to_string(), "--host".to_owned(), config.options.host.clone(), ]); if let Some(gpu_layers) = config.options.gpu_layers { args.extend(["-ngl".to_owned(), gpu_layers.to_string()]); } if let Some(context_size) = config.options.context_size { args.extend(["-c".to_owned(), context_size.to_string()]); } if config.options.jinja { args.push("--jinja".to_owned()); } args.extend(config.args.clone()); Ok(ModelServerArgv { command, args }) } fn build_spawn_spec( &self, config: &LocalModelServerConfig, ) -> Result { let argv = self.build_argv(config)?; Ok(SpawnSpec { command: argv.command, args: argv.args, cwd: ProjectPath::new("/").map_err(|e| ModelServerError::Invalid(e.to_string()))?, env: Vec::new(), context_plan: None, sandbox: None, }) } } fn resolve_binary(raw: &str) -> Result { if is_path_like(raw) { let path = Path::new(raw); if path.is_file() { return Ok(raw.to_owned()); } return Err(ModelServerError::PathNotAccessible(raw.to_owned())); } if let Some(path) = find_in_path(raw) { return Ok(path.to_string_lossy().into_owned()); } Err(ModelServerError::PathNotAccessible(format!( "{raw} not found in PATH" ))) } fn is_path_like(raw: &str) -> bool { raw.contains('/') || raw.contains('\\') || Path::new(raw).is_absolute() } fn find_in_path(command: &str) -> Option { std::env::var_os("PATH").and_then(|path| { std::env::split_paths(&path) .map(|dir| dir.join(command)) .find(|candidate| candidate.is_file()) }) } /// Local child-process manager for long-lived model servers. #[derive(Default)] pub struct LocalManagedProcess { children: Mutex>, } impl LocalManagedProcess { /// Creates an empty process manager. #[must_use] pub fn new() -> Self { Self::default() } } #[async_trait] impl ManagedProcess for LocalManagedProcess { async fn spawn(&self, spec: SpawnSpec) -> Result { let mut command = Command::new(&spec.command); command.args(&spec.args); if spec.cwd.as_str() != "/" { command.current_dir(spec.cwd.as_str()); } for (key, value) in &spec.env { command.env(key, value); } let child = command .spawn() .map_err(|e| ModelServerError::Process(format!("{}: {e}", spec.command)))?; let id = uuid::Uuid::new_v4().to_string(); self.children.lock().unwrap().insert(id.clone(), child); Ok(ManagedProcessHandle { id }) } async fn kill(&self, handle: &ManagedProcessHandle) -> Result<(), ModelServerError> { let Some(mut child) = self.children.lock().unwrap().remove(&handle.id) else { return Ok(()); }; child .start_kill() .map_err(|e| ModelServerError::Process(e.to_string())) } async fn status( &self, handle: &ManagedProcessHandle, ) -> Result { let mut children = self.children.lock().unwrap(); let Some(child) = children.get_mut(&handle.id) else { return Ok(ProcessStatus::Unknown); }; match child .try_wait() .map_err(|e| ModelServerError::Process(e.to_string()))? { Some(status) => Ok(ProcessStatus::Exited { code: status.code(), }), None => Ok(ProcessStatus::Running), } } } /// File name of the global local-model-server registry. const MODEL_SERVERS_FILE: &str = "model-servers.json"; const MODEL_SERVERS_VERSION: u32 = 2; #[derive(Debug, Clone, Serialize, Deserialize)] #[serde(rename_all = "camelCase")] struct ModelServersDoc { version: u32, servers: Vec, } #[derive(Debug, Clone, Deserialize)] #[serde(untagged)] enum PersistedModelServersDoc { V2(ModelServersDoc), V1(ModelServersDocV1), } #[derive(Debug, Clone, Deserialize)] #[serde(rename_all = "camelCase")] struct ModelServersDocV1 { servers: Vec, } #[derive(Debug, Clone, Deserialize)] #[serde(rename_all = "camelCase")] struct LocalModelServerConfigV1 { id: LocalModelServerId, kind: LocalModelServerKind, name: String, endpoint: ModelServerEndpoint, model: LocalModelRefV1, #[serde(default)] binary: Option, #[serde(default)] args: Vec, auto_start: bool, #[serde(default = "default_stop_policy")] stop_policy: StopPolicy, } #[derive(Debug, Clone, Deserialize)] #[serde(rename_all = "camelCase")] struct LocalModelRefV1 { id: String, label: String, #[serde(default)] path: Option, served_name: String, } fn default_stop_policy() -> StopPolicy { StopPolicy::StopOnAppExit } impl From for ModelServersDoc { fn from(doc: ModelServersDocV1) -> Self { Self { version: MODEL_SERVERS_VERSION, servers: doc .servers .into_iter() .map(|server| LocalModelServerConfig { id: server.id, kind: server.kind, name: server.name, endpoint: server.endpoint, model: LocalModelRef { id: server.model.id, label: server.model.label, source: server .model .path .map(|path| ModelSource::LocalPath { path }), served_name: server.model.served_name, }, binary: server.binary, options: LlamaCppOptions::default(), args: server.args, auto_start: server.auto_start, stop_policy: server.stop_policy, warmup_deadline_secs: None, }) .collect(), } } } impl Default for ModelServersDoc { fn default() -> Self { Self { version: MODEL_SERVERS_VERSION, servers: Vec::new(), } } } /// Filesystem-backed global model-server registry. #[derive(Clone)] pub struct FsModelServerRegistry { fs: std::sync::Arc, app_data_dir: String, } impl FsModelServerRegistry { /// Builds the registry from a filesystem and global app-data dir. #[must_use] pub fn new(fs: std::sync::Arc, app_data_dir: impl Into) -> Self { Self { fs, app_data_dir: app_data_dir.into(), } } fn path(&self) -> RemotePath { let base = self.app_data_dir.trim_end_matches(['/', '\\']); RemotePath::new(format!("{base}/{MODEL_SERVERS_FILE}")) } async fn read_doc(&self) -> Result { match self.fs.read(&self.path()).await { Ok(bytes) => match serde_json::from_slice::(&bytes) .map_err(|e| ModelServerError::Store(format!("serialization failed: {e}")))? { PersistedModelServersDoc::V2(doc) => Ok(ModelServersDoc { version: MODEL_SERVERS_VERSION, servers: doc.servers, }), PersistedModelServersDoc::V1(doc) => Ok(doc.into()), }, Err(domain::ports::FsError::NotFound(_)) => Ok(ModelServersDoc::default()), Err(domain::ports::FsError::PermissionDenied(p)) => { Err(ModelServerError::PermissionDenied(p)) } Err(e) => Err(ModelServerError::Store(e.to_string())), } } async fn write_doc(&self, doc: &ModelServersDoc) -> Result<(), ModelServerError> { let dir = RemotePath::new(self.app_data_dir.trim_end_matches(['/', '\\']).to_owned()); self.fs.create_dir_all(&dir).await.map_err(|e| match e { domain::ports::FsError::PermissionDenied(p) => ModelServerError::PermissionDenied(p), other => ModelServerError::Store(other.to_string()), })?; let bytes = serde_json::to_vec_pretty(doc) .map_err(|e| ModelServerError::Store(format!("serialization failed: {e}")))?; self.fs .write(&self.path(), &bytes) .await .map_err(|e| match e { domain::ports::FsError::PermissionDenied(p) => { ModelServerError::PermissionDenied(p) } other => ModelServerError::Store(other.to_string()), }) } } #[async_trait] impl ModelServerRegistry for FsModelServerRegistry { async fn get( &self, id: &LocalModelServerId, ) -> Result, ModelServerError> { Ok(self .read_doc() .await? .servers .into_iter() .find(|server| &server.id == id)) } async fn list(&self) -> Result, ModelServerError> { Ok(self.read_doc().await?.servers) } async fn save(&self, config: LocalModelServerConfig) -> Result<(), ModelServerError> { config .options .validate() .map_err(|e| ModelServerError::Invalid(e.to_string()))?; domain::model_server::validate_free_args(&config.args) .map_err(|e| ModelServerError::Invalid(e.to_string()))?; let mut doc = self.read_doc().await?; doc.version = MODEL_SERVERS_VERSION; if let Some(slot) = doc.servers.iter_mut().find(|server| server.id == config.id) { *slot = config; } else { doc.servers.push(config); } self.write_doc(&doc).await } async fn delete(&self, id: LocalModelServerId) -> Result<(), ModelServerError> { let mut doc = self.read_doc().await?; doc.servers.retain(|server| server.id != id); self.write_doc(&doc).await } } #[cfg(test)] mod gguf_selection_tests { use super::{select_gguf_files, HfSibling}; fn sibling(name: &str) -> HfSibling { HfSibling { filename: name.to_owned(), } } #[test] fn prefers_merged_file_over_colliding_shards() { let siblings = vec![ sibling("qwen2.5-coder-7b-instruct-q5_k_m-00001-of-00002.gguf"), sibling("qwen2.5-coder-7b-instruct-q5_k_m-00002-of-00002.gguf"), sibling("qwen2.5-coder-7b-instruct-q5_k_m.gguf"), ]; let files = select_gguf_files(&siblings, Some("q5_k_m")); assert_eq!(files, vec!["qwen2.5-coder-7b-instruct-q5_k_m.gguf"]); } #[test] fn selects_all_shards_sorted_by_index_when_no_merged_file() { let siblings = vec![ sibling("model-q4_k_m-00002-of-00003.gguf"), sibling("model-q4_k_m-00001-of-00003.gguf"), sibling("model-q4_k_m-00003-of-00003.gguf"), sibling("model-q8_0-00001-of-00002.gguf"), ]; let files = select_gguf_files(&siblings, Some("q4_k_m")); assert_eq!( files, vec![ "model-q4_k_m-00001-of-00003.gguf", "model-q4_k_m-00002-of-00003.gguf", "model-q4_k_m-00003-of-00003.gguf", ] ); } #[test] fn returns_empty_when_no_match() { let siblings = vec![sibling("model-q4_k_m.gguf")]; assert!(select_gguf_files(&siblings, Some("q5_k_m")).is_empty()); } }