use super::{Options, RenderedOutput, SyncTarget};
use crate::args;
use crate::io::{home_directory, read_file, ApiResult};
use crate::prelude::{OsString, PathBuf};
use crate::schema::agent::ModelDetails;
use crate::util::cmd::command_string;
use crate::util::constants::app::DEFAULT_LLAMA_SWAP_CONFIG_PATH;
use crate::util::StringConversion;
use crate::util::ToStrings;
use alloc::collections::{BTreeMap, BTreeSet};
use alloc::string::String;
use alloc::vec::Vec;
use color_eyre::eyre::eyre;
use core::{fmt, iter::once};
use serde::{Deserialize, Serialize};
use serde_norway::{Mapping, Number, Value};
use serde_with::skip_serializing_none;
use validator::{Validate, ValidationError, ValidationErrors};
#[derive(Clone, Debug, Deserialize, Serialize, Validate)]
#[serde(transparent)]
pub struct Alias {
#[validate(length(min = 1))]
value: String,
}
#[derive(Clone, Debug, Deserialize, Serialize)]
#[serde(from = "String", into = "String")]
pub enum Argument {
Reserved(String),
Value(String),
}
#[skip_serializing_none]
#[derive(Clone, Debug, Default, Deserialize, Serialize, Validate)]
#[serde(rename_all = "camelCase")]
pub struct Config {
#[serde(skip_serializing)]
#[validate(length(min = 1))]
pub path: Option<String>,
pub models_directory: Option<String>,
#[validate(length(min = 1))]
pub executable: Option<String>,
#[validate(range(min = 1))]
pub context_size: Option<u64>,
#[validate(range(min = 0))]
pub ttl: Option<i64>,
#[validate(nested)]
pub extra_args: Option<Vec<Argument>>,
#[validate(nested)]
pub environment: Option<Vec<EnvironmentVariable>>,
#[validate(nested)]
pub models: Option<BTreeMap<String, ModelOverride>>,
}
#[derive(Clone, Debug, Deserialize, Serialize, Validate)]
#[serde(transparent)]
pub struct EnvironmentVariable {
#[validate(custom(function = "is_keyvalue_string"))]
value: String,
}
#[skip_serializing_none]
#[derive(Clone, Debug, Default, Deserialize, Serialize, Validate)]
#[serde(deny_unknown_fields, rename_all = "camelCase")]
pub struct ModelOverride {
#[validate(range(min = 1))]
pub context_size: Option<u64>,
#[validate(range(min = 0))]
pub ttl: Option<i64>,
#[validate(nested)]
pub aliases: Option<Vec<Alias>>,
#[validate(nested)]
pub extra_args: Option<Vec<Argument>>,
#[validate(nested)]
pub environment: Option<Vec<EnvironmentVariable>>,
#[validate(length(min = 1))]
pub executable: Option<String>,
}
pub(super) struct ModelValidation<'a> {
config: &'a Config,
model_ids: &'a [String],
}
impl SyncTarget for Config {
const COMMAND: &'static str = "llama-swap";
fn merge(self, overrides: Self) -> Self {
Self {
path: overrides.path.or(self.path),
models_directory: overrides.models_directory.or(self.models_directory),
executable: overrides.executable.or(self.executable),
context_size: overrides.context_size.or(self.context_size),
ttl: overrides.ttl.or(self.ttl),
extra_args: overrides.extra_args.or(self.extra_args),
environment: overrides.environment.or(self.environment),
models: overrides.models.or(self.models),
}
}
fn merge_cli_overrides(self, overrides: Self) -> Self {
Self {
path: overrides.path.or(self.path),
models_directory: overrides.models_directory.or(self.models_directory),
..self
}
}
fn resolve_path(explicit: Option<&str>) -> ApiResult<PathBuf> {
explicit
.map(|path| PathBuf::from(path.to_string().to_cross_platform_path()))
.map_or_else(|| home_directory(DEFAULT_LLAMA_SWAP_CONFIG_PATH), Ok)
}
fn render(&self, options: Options<'_>) -> ApiResult<RenderedOutput> {
Self::resolve_path(self.path.as_deref()).and_then(|path| {
path.is_file()
.then(|| read_file(path.clone()))
.transpose()
.map(|content| content.unwrap_or_default())
.and_then(|before| {
match before.is_empty() {
| true => Ok(Value::Mapping(Default::default())),
| false => serde_norway::from_str(&before).map_err(|why| eyre!("Failed to parse existing llama-swap config: {why}")),
}
.and_then(|existing| self.upsert(existing, options.models, options.prune))
.and_then(|merged| {
serde_norway::to_string(&merged)
.map_err(|why| eyre!("Failed to serialize llama-swap config: {why}"))
.and_then(|content| {
serde_norway::from_str::<Value>(&content)
.map(|_| RenderedOutput {
target: "llama-swap",
path,
before,
content,
})
.map_err(|why| eyre!("Generated llama-swap YAML is invalid: {why}"))
})
})
})
})
}
}
impl Config {
pub fn build_command(
executable: &str,
gguf_path: Option<&str>,
extra_args: Option<&[Argument]>,
environment: Option<&[EnvironmentVariable]>,
context_size: Option<u64>,
) -> String {
let line = |command: &str, arguments: Vec<OsString>| command_string(command, &arguments);
let extra = extra_args.into_iter().flatten().map(ToString::to_string).collect::<Vec<_>>();
let arguments = args![
"--offline",
"--jinja",
("--batch-size", "2048"),
("--host", "0.0.0.0"),
("--sleep-idle-seconds", "600"),
("--tools", "all"),
("--ubatch-size", "2048")
];
let defaults = arguments.iter().enumerate().filter_map(|(index, option)| {
let option = option.to_string_lossy();
option.starts_with("--").then(|| {
let values = index
.checked_add(1)
.and_then(|index| arguments.get(index))
.filter(|value| !value.to_string_lossy().starts_with("--"))
.cloned()
.into_iter()
.collect();
line(&option, values)
})
});
let model = gguf_path.map(|path| line("--model", args![path]));
let context = context_size.map(|size| line("--ctx-size", args![size.to_string()]));
let extra = extra.into_iter().map(OsString::from).collect::<Vec<_>>();
let extra = extra
.split_first()
.map(|(command, arguments)| command_string(&command.to_string_lossy(), arguments));
let environment = environment
.into_iter()
.flatten()
.map(|entry| command_string(&format!("env:{entry}"), &[]));
once(command_string(executable, &[]))
.chain(model)
.chain(once(line("--port", args!["${PORT}"])))
.chain(context)
.chain(defaults)
.chain(extra)
.chain(environment)
.collect::<Vec<_>>()
.join("\n ")
}
pub fn model_entry(&self, model: &ModelDetails) -> Value {
let model_name = model.name.as_deref().or(model.id.as_deref()).unwrap_or("unknown");
let overrides = self.models.as_ref().and_then(|models| models.get(model_name));
let aliases = overrides
.and_then(|config| config.aliases.as_ref())
.map(|aliases| Value::Sequence(aliases.iter().map(|alias| Value::String(alias.as_str().to_string())).collect()));
let context_size = overrides.and_then(|config| config.context_size).or(self.context_size);
let command = Self::build_command(
overrides
.and_then(|config| config.executable.as_deref())
.or(self.executable.as_deref())
.unwrap_or("llama-server"),
model.path.as_deref(),
overrides.and_then(|config| config.extra_args.as_deref()).or(self.extra_args.as_deref()),
overrides.and_then(|config| config.environment.as_deref()).or(self.environment.as_deref()),
context_size,
);
let metadata = Value::Mapping(once((Value::String("acorn".to_string()), Value::Bool(true))).collect());
let mapping = once((Value::String("proxy".to_string()), Value::String("http://127.0.0.1:${PORT}".to_string())))
.chain(once((Value::String("cmd".to_string()), Value::String(command))))
.chain(
overrides
.and_then(|config| config.ttl)
.or(self.ttl)
.map(|ttl| (Value::String("ttl".to_string()), Value::Number(Number::from(ttl)))),
)
.chain(aliases.map(|value| (Value::String("aliases".to_string()), value)))
.chain(once((Value::String("metadata".to_string()), metadata)))
.collect();
Value::Mapping(mapping)
}
pub fn upsert(&self, existing: Value, models: &[ModelDetails], prune: bool) -> ApiResult<Value> {
match existing {
| Value::Null => self.upsert(Value::Mapping(Default::default()), models, prune),
| Value::Mapping(root) => {
let mut root = merge_mappings(Self::base_config(), root);
let models_key = Value::String("models".to_string());
root.remove(Value::String("modelsDir".to_string()));
let existing_models = root
.remove(&models_key)
.and_then(|value| match value {
| Value::Mapping(models) => Some(models),
| _ => None,
})
.unwrap_or_default();
let current_ids = models
.iter()
.filter_map(|model| model.name.as_ref().or(model.id.as_ref()))
.cloned()
.collect::<BTreeSet<_>>();
let retained = existing_models
.into_iter()
.filter(|(identifier, model)| {
!prune || identifier.as_str().is_some_and(|identifier| current_ids.contains(identifier)) || !is_managed(model)
})
.collect::<Mapping>();
let merged = models.iter().fold(retained, |mut entries, model| {
let identifier = model
.name
.as_ref()
.or(model.id.as_ref())
.cloned()
.unwrap_or_else(|| "unknown".to_string());
let key = Value::String(identifier);
let existing = entries.remove(&key).unwrap_or_else(|| Value::Mapping(Default::default()));
entries.insert(key, merge_model(existing, self.model_entry(model)));
entries
});
root.insert(models_key, Value::Mapping(merged));
Ok(Value::Mapping(root))
}
| _ => Err(eyre!("Existing llama-swap configuration root must be a mapping")),
}
}
fn base_config() -> Mapping {
let proxy = once((Value::String("listen".to_string()), Value::String("http://0.0.0.0:10732".to_string()))).collect();
let macros = [
("context_size", "${env.LLAMA_ARG_CTX_SIZE}"),
("parallel", "${env.LLAMA_ARG_N_PARALLEL}"),
("models_dir", "${env.HOME}/.models"),
("ssl_key", "${env.HOME}/certs/my.key"),
("ssl_cert", "${env.HOME}/certs/my.pem"),
]
.into_iter()
.map(|(key, value)| (Value::String(key.to_string()), Value::String(value.to_string())))
.collect();
[
(Value::String("healthCheckTimeout".to_string()), Value::Number(Number::from(500))),
(Value::String("proxy".to_string()), Value::Mapping(proxy)),
(Value::String("macros".to_string()), Value::Mapping(macros)),
]
.into_iter()
.collect()
}
}
impl Alias {
fn as_str(&self) -> &str {
&self.value
}
}
impl From<&str> for Alias {
fn from(value: &str) -> Self {
Self { value: value.to_string() }
}
}
impl fmt::Display for Argument {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
| Self::Reserved(value) | Self::Value(value) => value.fmt(formatter),
}
}
}
impl From<String> for Argument {
fn from(value: String) -> Self {
match value.split_once('=').map_or(value.as_str(), |(option, _)| option) {
| "--port" | "--model" | "-m" => Self::Reserved(value),
| _ => Self::Value(value),
}
}
}
impl From<&str> for Argument {
fn from(value: &str) -> Self {
Self::from(value.to_string())
}
}
impl From<Argument> for String {
fn from(value: Argument) -> Self {
value.to_string()
}
}
impl<'a> From<(&'a Config, &'a [String])> for ModelValidation<'a> {
fn from((config, model_ids): (&'a Config, &'a [String])) -> Self {
Self { config, model_ids }
}
}
impl Validate for Argument {
fn validate(&self) -> Result<(), ValidationErrors> {
match self {
| Self::Reserved(_) => {
let mut errors = ValidationErrors::new();
errors.add(
"argument",
ValidationError::new("reserved").with_message("Cannot contain --port, --model, or -m".into()),
);
Err(errors)
}
| Self::Value(_) => Ok(()),
}
}
}
impl Validate for ModelValidation<'_> {
fn validate(&self) -> Result<(), ValidationErrors> {
self.config
.models
.as_ref()
.into_iter()
.flat_map(|models| models.iter())
.try_fold(BTreeSet::new(), |mut aliases, (model_id, config)| {
match self.model_ids.iter().any(|candidate| candidate == model_id) {
| false => Err(validation_errors(
"unknown_model",
format!("llamaSwap.models contains unknown model override '{model_id}'"),
)),
| true => config
.aliases
.as_ref()
.into_iter()
.flatten()
.try_for_each(|alias| {
let alias = alias.as_str();
match (
self.model_ids.iter().any(|candidate| candidate == alias),
aliases.insert(alias.to_string()),
) {
| (true, _) => Err(validation_errors(
"model_alias",
format!("llamaSwap alias '{alias}' conflicts with a configured model ID"),
)),
| (_, false) => Err(validation_errors("duplicate_alias", format!("Duplicate llamaSwap alias '{alias}'"))),
| _ => Ok(()),
}
})
.map(|()| aliases),
}
})
.map(|_| ())
}
}
impl fmt::Display for EnvironmentVariable {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
self.value.fmt(formatter)
}
}
impl From<&str> for EnvironmentVariable {
fn from(value: &str) -> Self {
Self { value: value.to_string() }
}
}
fn is_keyvalue_string(value: &str) -> Result<(), ValidationError> {
let is_valid = value
.split_once('=')
.is_some_and(|(key, _)| !key.trim().is_empty() && !key.chars().any(char::is_whitespace));
match is_valid {
| true => Ok(()),
| false => Err(ValidationError::new("keyvalue").with_message("Provide a valid KEY=VALUE entry".into())),
}
}
fn is_managed(model: &Value) -> bool {
model
.as_mapping()
.and_then(|model| model.get(Value::String("metadata".to_string())))
.and_then(Value::as_mapping)
.and_then(|metadata| metadata.get(Value::String("acorn".to_string())))
.and_then(Value::as_bool)
.unwrap_or(false)
}
fn merge_mappings(defaults: Mapping, overrides: Mapping) -> Mapping {
defaults.into_iter().fold(overrides, |mut merged, (key, default_value)| {
let value = match (merged.get(&key), default_value) {
| (Some(Value::Mapping(overrides)), Value::Mapping(defaults)) => Some(Value::Mapping(merge_mappings(defaults, overrides.clone()))),
| (None, value) => Some(value),
| _ => None,
};
if let Some(value) = value {
merged.insert(key, value);
}
merged
})
}
fn merge_model(existing: Value, generated: Value) -> Value {
let owned = vec!["cmd", "command", "ttl", "aliases"]
.to_strings()
.into_iter()
.map(Value::String)
.collect::<Vec<_>>();
let mut existing = match existing {
| Value::Mapping(mapping) => mapping,
| _ => Default::default(),
};
let mut generated = match generated {
| Value::Mapping(mapping) => mapping,
| _ => Default::default(),
};
let proxy_key = Value::String("proxy".to_string());
if existing.contains_key(&proxy_key) {
generated.remove(&proxy_key);
}
let metadata_key = Value::String("metadata".to_string());
let metadata = existing
.remove(&metadata_key)
.and_then(|value| match value {
| Value::Mapping(mapping) => Some(mapping),
| _ => None,
})
.unwrap_or_default()
.into_iter()
.chain(
generated
.remove(&metadata_key)
.and_then(|value| match value {
| Value::Mapping(mapping) => Some(mapping),
| _ => None,
})
.unwrap_or_default(),
)
.collect();
Value::Mapping(
existing
.into_iter()
.filter(|(key, _)| !owned.contains(key))
.chain(generated)
.chain(once((metadata_key, Value::Mapping(metadata))))
.collect(),
)
}
fn validation_errors(code: &'static str, message: String) -> ValidationErrors {
let mut errors = ValidationErrors::new();
errors.add("models", ValidationError::new(code).with_message(message.into()));
errors
}
#[cfg(test)]
mod tests {
#![allow(clippy::unwrap_used, clippy::expect_used, clippy::panic)]
use super::*;
#[test]
fn test_argument_validation_rejects_reserved_options() {
["--port", "--port=9000", "--model", "--model=/tmp/model.gguf", "-m=/tmp/model.gguf"]
.into_iter()
.for_each(|argument| assert!(Argument::from(argument).validate().is_err()));
assert!(Argument::from("--flash-attn").validate().is_ok());
}
#[test]
fn test_build_command_simple() {
let cmd = Config::build_command("llama-server", Some("/models/qwen.gguf"), None, None, Some(8192));
assert_eq!(
cmd,
"llama-server\n --model /models/qwen.gguf\n --port ${PORT}\n --ctx-size 8192\n --offline\n --jinja\n --batch-size 2048\n --host 0.0.0.0\n --sleep-idle-seconds 600\n --tools all\n --ubatch-size 2048"
);
}
#[test]
fn test_build_command_with_environment() {
let env = vec![EnvironmentVariable::from("CUDA_VISIBLE_DEVICES=0")];
let cmd = Config::build_command("llama-server", Some("/models/qwen.gguf"), None, Some(&env), None);
assert!(cmd.contains("env:CUDA_VISIBLE_DEVICES=0"));
}
#[test]
fn test_build_command_with_extra_args() {
let extras = vec![
Argument::from("--batch-size"),
Argument::from("1024"),
Argument::from("--flash-attn"),
Argument::from("on"),
];
let cmd = Config::build_command("llama-server", Some("/models/qwen.gguf"), Some(&extras), None, None);
assert!(cmd.contains("--batch-size 1024"));
assert!(cmd
.find("--batch-size 2048")
.zip(cmd.rfind("--batch-size 1024"))
.is_some_and(|(default, override_)| default < override_));
assert!(cmd.contains("--flash-attn"));
assert!(cmd.contains("on"));
}
#[test]
fn test_build_command_with_optional_values() {
let cmd = Config::build_command("llama-server", None, None, None, None);
assert!(cmd.starts_with("llama-server\n --port ${PORT}\n --offline"));
assert!(cmd.ends_with("--ubatch-size 2048"));
}
#[test]
fn test_build_command_with_spaces_in_path() {
let cmd = Config::build_command("llama-server", Some("/path with spaces/model.gguf"), None, None, None);
assert!(cmd.contains('"'));
}
#[test]
fn test_resolve_path_uses_llama_swap_user_config() {
assert_eq!(
Config::resolve_path(None).unwrap(),
home_directory(DEFAULT_LLAMA_SWAP_CONFIG_PATH).unwrap()
);
}
#[test]
fn test_upsert_preserves_unrelated_values_and_prunes_only_managed_models() {
let existing: Value = serde_norway::from_str(
r#"modelsDir: /legacy/models
healthCheckTimeout: 120
models:
qwen:
cmd: old
proxy: http://127.0.0.1:9000
metadata:
owner: user
stale:
cmd: stale
metadata:
acorn: true
custom:
cmd: custom
"#,
)
.unwrap();
let models = [ModelDetails::init().id("qwen").name("qwen").path("/models/qwen.gguf").build()];
let config = Config::default();
let additive = config.upsert(existing.clone(), &models, false).unwrap();
let additive_text = serde_norway::to_string(&additive).unwrap();
assert!(additive_text.contains("healthCheckTimeout: 120"));
assert!(additive_text.contains("proxy: http://127.0.0.1:9000"));
assert!(additive_text.contains("owner: user"));
assert!(additive_text.contains("acorn: true"));
assert!(additive_text.contains("stale:"));
assert!(additive_text.contains("custom:"));
assert!(additive_text.contains("cmd: |-"));
assert!(additive_text.contains("llama-server\n"));
assert!(additive_text.contains("--port ${PORT}\n"));
assert!(additive_text.contains("--model /models/qwen.gguf\n"));
assert!(!additive_text.contains("modelsDir"));
let pruned = config.upsert(existing, &models, true).unwrap();
let pruned_text = serde_norway::to_string(&pruned).unwrap();
assert!(!pruned_text.contains("stale:"));
assert!(pruned_text.contains("custom:"));
let repeated = config.upsert(pruned, &models, true).unwrap();
assert_eq!(pruned_text, serde_norway::to_string(&repeated).unwrap());
}
#[test]
fn test_upsert_adds_gold_base_defaults() {
let models = [ModelDetails::init().id("qwen").name("qwen").path("/models/qwen.gguf").build()];
let output = Config::default()
.upsert(Value::Mapping(Default::default()), &models, false)
.and_then(|value| serde_norway::to_string(&value).map_err(Into::into))
.unwrap();
assert!(output.contains("healthCheckTimeout: 500"));
assert!(output.contains("listen: http://0.0.0.0:10732"));
assert!(output.contains("context_size: ${env.LLAMA_ARG_CTX_SIZE}"));
assert!(output.contains("parallel: ${env.LLAMA_ARG_N_PARALLEL}"));
assert!(output.contains("models_dir: ${env.HOME}/.models"));
assert!(output.contains("ssl_key: ${env.HOME}/certs/my.key"));
assert!(output.contains("ssl_cert: ${env.HOME}/certs/my.pem"));
assert!(output.contains("proxy: http://127.0.0.1:${PORT}"));
assert!(output.contains("--offline\n"));
assert!(output.contains("--tools all\n"));
}
#[test]
fn test_validate_rejects_invalid_context_ttl_environment_and_aliases() {
let model_ids = vec!["qwen".to_string(), "gemma".to_string()];
let invalid = [
Config {
context_size: Some(0),
..Default::default()
},
Config {
ttl: Some(-1),
..Default::default()
},
Config {
environment: Some(vec![EnvironmentVariable::from("INVALID")]),
..Default::default()
},
Config {
models: Some(
[
(
"qwen".to_string(),
ModelOverride {
aliases: Some(vec![Alias::from("shared")]),
..Default::default()
},
),
(
"gemma".to_string(),
ModelOverride {
aliases: Some(vec![Alias::from("shared")]),
..Default::default()
},
),
]
.into_iter()
.collect(),
),
..Default::default()
},
];
invalid.iter().for_each(|config| {
assert!(Validate::validate(config).is_err() || ModelValidation::from((config, model_ids.as_slice())).validate().is_err());
});
}
}