diff --git a/README.md b/README.md index e3b44e9..e802eb1 100644 --- a/README.md +++ b/README.md @@ -60,7 +60,9 @@ The provided configuration file may include environment variable substitution us "authorization": "x-api-key", "apikey": "${ANTHROPIC_API_KEY}", "capability": { - "models": ["claude-*"], + "models": { + "claude-*": {} + }, "user_agents": ["claude-code"] }, "compatibility": { @@ -120,9 +122,20 @@ When `capability` is set on a provider, Lacuna will use it to filter the request "anthropic": { "baseurl": "https://api.anthropic.com", "capability": { - "models": ["claude-*"], + "models": { + "claude-*": {} + }, "user_agents": ["claude-code"] } + }, + "bedrock": { + "baseurl": "https://bedrock-runtime.us-east-1.amazonaws.com", + "capability": { + "models": { + "us.anthropic.claude-opus-4-5*": { "rewrite": "arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/abcd1234567" }, + "us.anthropic.claude-haiku-4-5*": {} + } + } } } } @@ -132,7 +145,12 @@ In the above example: - The `anthropic` provider will only allow requests with a User-Agent that match Lacuna's built-in `claude-code` pattern. - The `anthropic` provider may only be used with models that match the `claude-*` glob pattern. +`capability.models` maps each model glob pattern to a settings object with an optional `rewrite`: +- `{}` ⇒ the model is allowed and forwarded as-is. +- `{ "rewrite": "" }` ⇒ the model is allowed and the upstream is called with `` instead. In the example, requests for `us.anthropic.claude-opus-4-5*` are sent to Bedrock as the configured application-inference-profile ARN. + Notes: +- An empty or omitted `models` allows all models (no rewrite). - Built-in User-Agent patterns can be found in [src/user_agent.rs](src/user_agent.rs). ## 📦 Dev Dependencies diff --git a/examples/lacuna/lacuna.config.json b/examples/lacuna/lacuna.config.json index 1bfd53b..7abde9d 100644 --- a/examples/lacuna/lacuna.config.json +++ b/examples/lacuna/lacuna.config.json @@ -12,6 +12,10 @@ "baseurl": "https://api.anthropic.com", "authorization": "x-api-key", "apikey": "${ANTHROPIC_API_KEY}", + "capability": { + "models": {}, + "user_agents": [] + }, "compatibility": { "anthropic_messages": true } @@ -21,6 +25,10 @@ "baseurl": "https://bedrock-runtime.us-east-1.amazonaws.com", "authorization": "bearer", "apikey": "${BEDROCK_API_KEY}", + "capability": { + "models": {}, + "user_agents": [] + }, "compatibility": { "bedrock_model_invoke": true }, diff --git a/src/config.rs b/src/config.rs index 42380ec..606820a 100644 --- a/src/config.rs +++ b/src/config.rs @@ -29,15 +29,66 @@ pub struct Lacuna { pub user_agents: Vec, } +#[derive(Debug, Clone, PartialEq)] +pub struct ModelRule { + pub pattern: glob::Pattern, + pub rewrite: Option, +} + +/// Exists only so the map value can be (de)serialized with `deny_unknown_fields` +/// (rejecting typos such as `{ "rewite": "x" }`) +#[derive(Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +struct ModelRuleRepr { + #[serde(default, skip_serializing_if = "Option::is_none")] + rewrite: Option, +} + +/// Serialize `Vec` as an ordered object: `pattern => { "rewrite"?: target }`. +pub(crate) fn serialize_model_rules( + rules: &[ModelRule], + serializer: S, +) -> Result { + use serde::ser::SerializeMap; + let mut map = serializer.serialize_map(Some(rules.len()))?; + for rule in rules { + map.serialize_entry( + rule.pattern.as_str(), + &ModelRuleRepr { + rewrite: rule.rewrite.clone(), + }, + )?; + } + map.end() +} + +/// Deserialize the map (`pattern => { "rewrite"?: target }`) into rules sorted +/// by pattern, so iteration order is deterministic regardless of the JSON key order. +pub(crate) fn deserialize_model_rules<'de, D: serde::Deserializer<'de>>( + deserializer: D, +) -> Result, D::Error> { + let map = std::collections::BTreeMap::::deserialize(deserializer)?; + map.into_iter() + .map(|(key, value)| { + let pattern = glob::Pattern::new(&key).map_err(serde::de::Error::custom)?; + Ok(ModelRule { + pattern, + rewrite: value.rewrite, + }) + }) + .collect() +} + #[derive(Debug, Deserialize, Serialize, PartialEq, Default)] #[serde(deny_unknown_fields)] pub struct Capability { #[serde( + rename = "models", default, - serialize_with = "crate::serde_utils::serialize_patterns", - deserialize_with = "crate::serde_utils::deserialize_patterns" + serialize_with = "serialize_model_rules", + deserialize_with = "deserialize_model_rules" )] - pub models: Vec, + pub model_rules: Vec, #[serde( default, @@ -123,6 +174,81 @@ impl Config { mod tests { use super::*; + fn parse_model_rules(json: &str) -> Result, serde_json::Error> { + let mut de = serde_json::Deserializer::from_str(json); + deserialize_model_rules(&mut de) + } + + fn dump_model_rules(rules: &[ModelRule]) -> String { + let mut buf = Vec::new(); + let mut ser = serde_json::Serializer::new(&mut buf); + serialize_model_rules(rules, &mut ser).unwrap(); + String::from_utf8(buf).unwrap() + } + + #[test] + fn deserializes_model_rules_to_deterministic_sorted_order_regardless_of_input_order() { + let a = parse_model_rules(r#"{ "c-*": {}, "a-*": {}, "b-*": {} }"#).unwrap(); + let b = parse_model_rules(r#"{ "b-*": {}, "c-*": {}, "a-*": {} }"#).unwrap(); + assert_eq!(a, b); + let patterns: Vec<&str> = a.iter().map(|r| r.pattern.as_str()).collect(); + assert_eq!(patterns, vec!["a-*", "b-*", "c-*"]); + } + + #[test] + fn duplicate_model_rule_patterns_are_deduped_last_wins() { + let rules = parse_model_rules(r#"{ "a-*": { "rewrite": "X" }, "a-*": {} }"#).unwrap(); + assert_eq!( + rules, + vec![ModelRule { + pattern: glob::Pattern::new("a-*").unwrap(), + rewrite: None, + }] + ); + } + + #[test] + fn deserialization_rejected_when_model_rule_pattern_is_invalid() { + let err = parse_model_rules(r#"{ "[invalid": {} }"#) + .unwrap_err() + .to_string(); + assert!( + err.contains("Pattern syntax error"), + "expected pattern syntax error, got: {err}" + ); + } + + #[test] + fn deserialization_rejected_when_model_rule_has_unknown_field() { + let err = parse_model_rules(r#"{ "a-*": { "something": "x" } }"#) + .unwrap_err() + .to_string(); + assert!( + err.contains("unknown field"), + "expected unknown field error, got: {err}" + ); + } + + #[test] + fn serialization_deserialization_round_trip() { + let rules = vec![ + ModelRule { + pattern: glob::Pattern::new("a-*").unwrap(), + rewrite: Some("X".to_string()), + }, + ModelRule { + pattern: glob::Pattern::new("b-*").unwrap(), + rewrite: None, + }, + ]; + assert_eq!( + dump_model_rules(&rules), + r#"{"a-*":{"rewrite":"X"},"b-*":{}}"# + ); + // And it round-trips back to the same rules. + assert_eq!(parse_model_rules(&dump_model_rules(&rules)).unwrap(), rules); + } + #[test] fn deserialize_json_config() { let json = r#"{ @@ -131,7 +257,10 @@ mod tests { "name": "OpenAI", "baseurl": "https://api.openai.com/v1", "capability": { - "models": ["gpt-4o", "gpt-4o-mini"], + "models": { + "gpt-4o": {}, + "gpt-4o-mini": {} + }, "user_agents": ["claude-code"] }, "apikey": "sk-test", @@ -144,7 +273,11 @@ mod tests { "anthropic": { "name": "Anthropic", "baseurl": "https://api.anthropic.com", - "capability": { "models": ["claude-sonnet-4-20250514"] }, + "capability": { + "models": { + "claude-sonnet-4-20250514": {} + } + }, "apikey": "sk-ant-test", "authorization": "x-api-key", "compatibility": { @@ -155,7 +288,11 @@ mod tests { "gemini": { "name": "Gemini", "baseurl": "https://generativelanguage.googleapis.com", - "capability": { "models": ["gemini-2.0-flash"] }, + "capability": { + "models": { + "gemini-2.0-flash": {} + } + }, "authorization": "x-goog-api-key", "headers": { "x-some-header": "foo" @@ -170,10 +307,16 @@ mod tests { assert_eq!(openai.name, "OpenAI"); assert_eq!(openai.baseurl, "https://api.openai.com/v1"); assert_eq!( - openai.capability.models, + openai.capability.model_rules, vec![ - glob::Pattern::new("gpt-4o").unwrap(), - glob::Pattern::new("gpt-4o-mini").unwrap() + ModelRule { + pattern: glob::Pattern::new("gpt-4o").unwrap(), + rewrite: None, + }, + ModelRule { + pattern: glob::Pattern::new("gpt-4o-mini").unwrap(), + rewrite: None, + }, ] ); assert_eq!( @@ -208,7 +351,11 @@ mod tests { "minimal": { "name": "Minimal", "baseurl": "https://example.com", - "capability": { "models": ["model-1"] } + "capability": { + "models": { + "model-1": {} + } + } } } }"#; @@ -275,7 +422,11 @@ mod tests { "baseurl": "https://api.openai.com/", "apikey": "YOUR_OPENAI_KEY", "capability": { - "models": ["gpt-5", "gpt-5-mini", "gpt-4.1"], + "models": { + "gpt-5": {}, + "gpt-5-mini": {}, + "gpt-4.1": {} + }, "user_agents": ["claude-code"] }, "name": "OpenAI", @@ -290,12 +441,12 @@ mod tests { "apikey": "bedrock-api-key-xxx", "authorization": "bearer", "capability": { - "models": [ - "us.anthropic.claude-haiku-4-5-20251001-v1:0", - "us.anthropic.claude-sonnet-4-5-20250929-v1:0", - "us.anthropic.claude-opus-4-5-20251101-v1:0", - "us.anthropic.claude-opus-4-6-v1" - ] + "models": { + "us.anthropic.claude-haiku-4-5-20251001-v1:0": {}, + "us.anthropic.claude-sonnet-4-5-20250929-v1:0": {}, + "us.anthropic.claude-opus-4-5*": { "rewrite": "arn:aws:bedrock:us-east-1:409905535292:application-inference-profile/11cprf2uimr9" }, + "us.anthropic.claude-opus-4-6-v1": {} + } }, "compatibility": { "bedrock_model_invoke": true @@ -305,7 +456,13 @@ mod tests { "baseurl": "https://api.anthropic.com", "apikey": "YOUR_ANTHROPIC_KEY", "authorization": "x-api-key", - "capability": { "models": ["claude-sonnet-4-5", "claude-haiku-4-5", "claude-opus-4-5"] }, + "capability": { + "models": { + "claude-sonnet-4-5": {}, + "claude-haiku-4-5": {}, + "claude-opus-4-5": {} + } + }, "compatibility": { "openai_chat": false, "anthropic_messages": true @@ -315,7 +472,12 @@ mod tests { "baseurl": "https://generativelanguage.googleapis.com", "apikey": "YOUR_GEMINI_KEY", "authorization": "x-goog-api-key", - "capability": { "models": ["gemini-2.5-flash", "gemini-2.5-pro"] }, + "capability": { + "models": { + "gemini-2.5-flash": {}, + "gemini-2.5-pro": {} + } + }, "name": "Google Gemini", "compatibility": { "openai_chat": false, @@ -327,16 +489,16 @@ mod tests { "authorization": "bearer", "apikey": "keyfile::ba3..3kb.data...67", "capability": { - "models": [ - "gemini-2.0-flash-exp", - "gemini-2.5-flash", - "gemini-2.5-flash-image", - "gemini-2.5-pro", - "claude-opus-4-5@20251101", - "claude-haiku-4-5@20251001", - "claude-sonnet-4-5@20250929", - "claude-opus-4-6" - ] + "models": { + "gemini-2.0-flash-exp": {}, + "gemini-2.5-flash": {}, + "gemini-2.5-flash-image": {}, + "gemini-2.5-pro": {}, + "claude-opus-4-5@20251101": {}, + "claude-haiku-4-5@20251001": {}, + "claude-sonnet-4-5@20250929": {}, + "claude-opus-4-6": {} + } }, "compatibility": { // Gemini model support @@ -349,17 +511,22 @@ mod tests { "baseurl": "https://openrouter.ai/api/", "apikey": "YOUR_OPENROUTER_KEY", "capability": { - "models": [ - "qwen/qwen3-235b-a22b-2507", - "google/gemini-2.5-pro-preview", - "x-ai/grok-code-fast-1" - ] + "models": { + "qwen/qwen3-235b-a22b-2507": {}, + "google/gemini-2.5-pro-preview": {}, + "x-ai/grok-code-fast-1": {} + } } }, "private": { "baseurl": "YOUR_PRIVATE_LLM_URL", "tailnet": true, - "capability": { "models": ["qwen3-coder-30b", "llama-3.1-70b"] } + "capability": { + "models": { + "qwen3-coder-30b": {}, + "llama-3.1-70b": {} + } + } } } } diff --git a/src/http_handlers/proxy.rs b/src/http_handlers/proxy.rs index c852f87..a0984f9 100644 --- a/src/http_handlers/proxy.rs +++ b/src/http_handlers/proxy.rs @@ -211,7 +211,7 @@ mod tests { use crate::provider::ProviderManager; use crate::provider::compatibility::Compatibility; - use crate::test_utils::{make_provider, make_provider_with_models, spawn_echo_server}; + use crate::test_utils::{make_provider, make_provider_with_model_rules, spawn_echo_server}; #[tokio::test] async fn unmatched_path_returns_404() { @@ -412,18 +412,24 @@ mod tests { compat.bedrock_model_invoke = true; let mut manager = ProviderManager::new(); - manager.add(make_provider_with_models( + manager.add(make_provider_with_model_rules( "provider-a", &format!("http://{addr}"), compat.clone(), - vec![glob::Pattern::new("*").unwrap()], + vec![crate::config::ModelRule { + pattern: glob::Pattern::new("*").unwrap(), + rewrite: None, + }], )); - manager.add(make_provider_with_models( + manager.add(make_provider_with_model_rules( "provider-b", &format!("http://{addr}"), compat, - vec![glob::Pattern::new("gpt-4o").unwrap()], + vec![crate::config::ModelRule { + pattern: glob::Pattern::new("gpt-4o").unwrap(), + rewrite: None, + }], )); let app = crate::app::AppBuilder::new().manager(manager).build(); diff --git a/src/provider/mod.rs b/src/provider/mod.rs index 9e82d09..ffbdde9 100644 --- a/src/provider/mod.rs +++ b/src/provider/mod.rs @@ -60,7 +60,7 @@ pub struct Provider { #[allow(dead_code)] pub name: String, pub baseurl: reqwest::Url, - pub models: Vec, + pub model_rules: Vec, pub user_agents: Vec, pub authorizer: Authorization, client: reqwest::Client, @@ -74,10 +74,16 @@ impl Provider { pub fn from_config(key: &str, config: &config::Provider) -> Result { let baseurl = reqwest::Url::parse(&config.baseurl)?; let authenticator = build_authenticator(&config.authorization, &config.apikey); + let model_patterns: Vec = config + .capability + .model_rules + .iter() + .map(|r| r.pattern.clone()) + .collect(); let authorizer = Authorization { rules: vec![Rule { providers: vec![], - models: config.capability.models.clone(), + models: model_patterns, user_agents: config.capability.user_agents.clone(), }], }; @@ -85,7 +91,7 @@ impl Provider { key: key.to_owned(), name: config.name.clone(), baseurl, - models: config.capability.models.clone(), + model_rules: config.capability.model_rules.clone(), user_agents: config.capability.user_agents.clone(), authorizer, client: reqwest::Client::new(), @@ -160,7 +166,10 @@ mod tests { description: String::new(), baseurl: baseurl.to_owned(), capability: config::Capability { - models: vec![glob::Pattern::new("model-1").unwrap()], + model_rules: vec![config::ModelRule { + pattern: glob::Pattern::new("model-1").unwrap(), + rewrite: None, + }], user_agents: vec![], }, apikey: apikey.to_owned(), diff --git a/src/test_utils.rs b/src/test_utils.rs index a863fa8..d29bb91 100644 --- a/src/test_utils.rs +++ b/src/test_utils.rs @@ -37,14 +37,14 @@ pub fn make_provider( baseurl: &str, compat: provider::compatibility::Compatibility, ) -> provider::Provider { - make_provider_with_models(key, baseurl, compat, vec![]) + make_provider_with_model_rules(key, baseurl, compat, vec![]) } -pub fn make_provider_with_models( +pub fn make_provider_with_model_rules( key: &str, baseurl: &str, compat: provider::compatibility::Compatibility, - models: Vec, + model_rules: Vec, ) -> provider::Provider { provider::Provider::from_config( key, @@ -53,7 +53,7 @@ pub fn make_provider_with_models( description: String::new(), baseurl: baseurl.to_owned(), capability: config::Capability { - models, + model_rules, user_agents: vec![], }, apikey: String::new(),