Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
17 changes: 17 additions & 0 deletions crates/frontend/src/api.rs
Original file line number Diff line number Diff line change
Expand Up @@ -6132,6 +6132,8 @@ pub struct AdminLlmGatewayKeyView {
#[serde(default = "default_anthropic_upstream_pool_mode")]
pub kiro_anthropic_upstream_pool_mode: String,
pub model_name_map: Option<BTreeMap<String, String>>,
#[serde(default)]
pub kiro_model_group_preferences: BTreeMap<String, String>,
pub request_max_concurrency: Option<u64>,
pub request_min_start_interval_ms: Option<u64>,
#[serde(default = "default_true")]
Expand Down Expand Up @@ -9030,6 +9032,7 @@ pub async fn create_admin_llm_gateway_key(
preferred_pool_strategy: default_kiro_pool_strategy(),
kiro_anthropic_upstream_pool_mode: default_anthropic_upstream_pool_mode(),
model_name_map: None,
kiro_model_group_preferences: BTreeMap::new(),
request_max_concurrency,
request_min_start_interval_ms,
kiro_request_validation_enabled: true,
Expand Down Expand Up @@ -9096,6 +9099,7 @@ pub struct PatchAdminLlmGatewayKeyRequest<'a> {
pub preferred_pool_strategy: Option<&'a str>,
pub kiro_anthropic_upstream_pool_mode: Option<&'a str>,
pub model_name_map: Option<&'a BTreeMap<String, String>>,
pub kiro_model_group_preferences: Option<&'a BTreeMap<String, String>>,
pub request_max_concurrency: Option<u64>,
pub request_min_start_interval_ms: Option<u64>,
pub codex_fast_enabled: Option<bool>,
Expand Down Expand Up @@ -9136,6 +9140,7 @@ pub async fn patch_admin_llm_gateway_key(
request.preferred_pool_strategy,
request.kiro_anthropic_upstream_pool_mode,
request.model_name_map,
request.kiro_model_group_preferences,
request.request_max_concurrency,
request.request_min_start_interval_ms,
request.codex_fast_enabled,
Expand Down Expand Up @@ -9236,6 +9241,11 @@ pub async fn patch_admin_llm_gateway_key(
.map_err(|e| format!("Serialize error: {:?}", e))?;
body.insert("model_name_map".to_string(), value);
}
if let Some(preferences) = request.kiro_model_group_preferences {
let value = serde_json::to_value(preferences)
.map_err(|e| format!("Serialize error: {:?}", e))?;
body.insert("kiro_model_group_preferences".to_string(), value);
}
if let Some(request_max_concurrency) = request.request_max_concurrency {
body.insert(
"request_max_concurrency".to_string(),
Expand Down Expand Up @@ -11445,6 +11455,7 @@ pub async fn create_admin_kiro_key(
preferred_pool_strategy: default_kiro_pool_strategy(),
kiro_anthropic_upstream_pool_mode: default_anthropic_upstream_pool_mode(),
model_name_map: None,
kiro_model_group_preferences: BTreeMap::new(),
request_max_concurrency: None,
request_min_start_interval_ms: None,
kiro_request_validation_enabled: true,
Expand Down Expand Up @@ -11513,6 +11524,7 @@ pub async fn patch_admin_kiro_key(
request.preferred_pool_strategy,
request.kiro_anthropic_upstream_pool_mode,
request.model_name_map,
request.kiro_model_group_preferences,
request.request_max_concurrency,
request.request_min_start_interval_ms,
request.kiro_request_validation_enabled,
Expand Down Expand Up @@ -11608,6 +11620,11 @@ pub async fn patch_admin_kiro_key(
.map_err(|e| format!("Serialize error: {:?}", e))?;
body.insert("model_name_map".to_string(), value);
}
if let Some(preferences) = request.kiro_model_group_preferences {
let value = serde_json::to_value(preferences)
.map_err(|e| format!("Serialize error: {:?}", e))?;
body.insert("kiro_model_group_preferences".to_string(), value);
}
if let Some(kiro_request_validation_enabled) = request.kiro_request_validation_enabled {
body.insert(
"kiro_request_validation_enabled".to_string(),
Expand Down
228 changes: 228 additions & 0 deletions crates/frontend/src/pages/admin_kiro_gateway.rs

Large diffs are not rendered by default.

1 change: 1 addition & 0 deletions crates/frontend/src/pages/admin_llm_gateway.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1489,6 +1489,7 @@ fn key_editor_card(props: &KeyEditorCardProps) -> Html {
preferred_pool_strategy: None,
kiro_anthropic_upstream_pool_mode: None,
model_name_map: None,
kiro_model_group_preferences: None,
request_max_concurrency: request_max_concurrency_value,
request_min_start_interval_ms: request_min_start_interval_ms_value,
codex_fast_enabled: Some(codex_fast_enabled_value),
Expand Down
1 change: 1 addition & 0 deletions crates/llm-access-core/src/store/empty.rs
Original file line number Diff line number Diff line change
Expand Up @@ -431,6 +431,7 @@ impl AdminKeyStore for EmptyAdminKeyStore {
auto_account_names: None,
preferred_pool_strategy: default_kiro_pool_strategy(),
model_name_map: None,
kiro_model_group_preferences: std::collections::BTreeMap::new(),
request_max_concurrency: key.request_max_concurrency,
request_min_start_interval_ms: key.request_min_start_interval_ms,
codex_fast_enabled: true,
Expand Down
6 changes: 6 additions & 0 deletions crates/llm-access-core/src/store/keys.rs
Original file line number Diff line number Diff line change
Expand Up @@ -70,6 +70,9 @@ pub struct AdminKey {
pub preferred_pool_strategy: String,
/// Model name mapping.
pub model_name_map: Option<BTreeMap<String, String>>,
/// Exact Kiro request-model to preferred account-group map.
#[serde(default)]
pub kiro_model_group_preferences: BTreeMap<String, String>,
/// Per-key request concurrency cap.
pub request_max_concurrency: Option<u64>,
/// Per-key request pacing interval.
Expand Down Expand Up @@ -382,6 +385,8 @@ pub struct AdminKeyPatch {
pub preferred_pool_strategy: Option<String>,
/// New model name map.
pub model_name_map: Option<Option<BTreeMap<String, String>>>,
/// New exact Kiro request-model to preferred account-group map.
pub kiro_model_group_preferences: Option<BTreeMap<String, String>>,
/// New per-key request concurrency cap.
pub request_max_concurrency: Option<Option<u64>>,
/// New per-key request pacing interval.
Expand Down Expand Up @@ -455,6 +460,7 @@ mod tests {
auto_account_names: None,
preferred_pool_strategy: "balanced".to_string(),
model_name_map: None,
kiro_model_group_preferences: BTreeMap::new(),
request_max_concurrency: None,
request_min_start_interval_ms: None,
codex_fast_enabled: true,
Expand Down
68 changes: 68 additions & 0 deletions crates/llm-access-core/src/store/kiro_model_routing.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,68 @@
//! Kiro model -> account-group preference helpers.
//!
//! ```text
//! raw request model
//! |
//! +-- exact preference hit --> preferred account set
//! |
//! +-- no hit ---------------> existing affinity/order rules
//! ```

use std::collections::BTreeMap;

/// Per-key exact model-name to Kiro account-group preference map.
pub type KiroModelGroupPreferences = BTreeMap<String, String>;

/// Normalize admin-supplied model preference rules.
pub fn normalize_kiro_model_group_preferences(
input: KiroModelGroupPreferences,
) -> KiroModelGroupPreferences {
input
.into_iter()
.filter_map(|(model, group_id)| {
let model = model.trim();
let group_id = group_id.trim();
(!model.is_empty() && !group_id.is_empty())
.then(|| (model.to_string(), group_id.to_string()))
})
.collect()
}

/// Resolve one exact model-name preference.
pub fn kiro_model_group_preference<'a>(
preferences: &'a KiroModelGroupPreferences,
model: &str,
) -> Option<&'a str> {
preferences.get(model.trim()).map(String::as_str)
}

#[cfg(test)]
mod tests {
use super::*;

#[test]
fn normalizes_model_group_preferences_by_trimming_and_dropping_empty_entries() {
let normalized = normalize_kiro_model_group_preferences(BTreeMap::from([
(" claude-sonnet-4 ".to_string(), " group-sonnet ".to_string()),
("".to_string(), "group-empty-model".to_string()),
("claude-opus-4".to_string(), " ".to_string()),
]));

assert_eq!(
normalized,
BTreeMap::from([("claude-sonnet-4".to_string(), "group-sonnet".to_string())])
);
}

#[test]
fn resolves_model_group_preferences_by_exact_model_name() {
let preferences =
BTreeMap::from([("claude-sonnet-4".to_string(), "group-sonnet".to_string())]);

assert_eq!(
kiro_model_group_preference(&preferences, "claude-sonnet-4"),
Some("group-sonnet")
);
assert_eq!(kiro_model_group_preference(&preferences, "claude-sonnet"), None);
}
}
4 changes: 4 additions & 0 deletions crates/llm-access-core/src/store/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,7 @@ mod empty;
mod groups;
mod keys;
mod kiro_account;
mod kiro_model_routing;
mod proxy;
mod public;
mod routes;
Expand Down Expand Up @@ -71,6 +72,9 @@ pub use kiro_account::{
ADMIN_KIRO_ACCOUNT_ISSUE_ABNORMAL, ADMIN_KIRO_ACCOUNT_ISSUE_AUTH_401,
ADMIN_KIRO_ACCOUNT_ISSUE_DISABLED, ADMIN_KIRO_ACCOUNT_ISSUE_ERROR,
};
pub use kiro_model_routing::{
kiro_model_group_preference, normalize_kiro_model_group_preferences, KiroModelGroupPreferences,
};
pub use proxy::{
default_proxy_bindings, AdminProxyBinding, AdminProxyConfig, AdminProxyConfigPatch,
AdminProxyEndpointCheck, AdminProxyEndpointCheckUpdate, AdminProxyTrafficSnapshot,
Expand Down
4 changes: 4 additions & 0 deletions crates/llm-access-core/src/store/routes.rs
Original file line number Diff line number Diff line change
@@ -1,6 +1,8 @@
//! Provider routing contracts: authenticated key, provider proxy config,
//! Codex/Kiro route + auth-update views, and JWT/auth-error helpers.

use std::collections::BTreeMap;

use base64::{engine::general_purpose::URL_SAFE_NO_PAD, Engine as _};
use serde_json::Value;

Expand Down Expand Up @@ -234,6 +236,8 @@ pub struct ProviderKiroRoute {
pub cctest_proxy_api_key: Option<String>,
/// JSON object mapping public model names to upstream Kiro model names.
pub model_name_map_json: String,
/// Exact request-model to preferred account names resolved from key config.
pub model_group_preferred_account_names: BTreeMap<String, Vec<String>>,
/// Effective Kiro cache k-model JSON for this key.
pub cache_kmodels_json: String,
/// Effective Kiro cache policy JSON for this key.
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,9 @@ CREATE TABLE IF NOT EXISTS llm_key_route_config (
preferred_pool_strategy IN ('balanced', 'credit_first')
),
model_name_map_json JSONB,
kiro_model_group_preferences_json JSONB NOT NULL DEFAULT '{}'::jsonb CHECK (
jsonb_typeof(kiro_model_group_preferences_json) = 'object'
),
request_max_concurrency BIGINT,
request_min_start_interval_ms BIGINT,
codex_fast_enabled BOOLEAN NOT NULL DEFAULT TRUE,
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,21 @@
ALTER TABLE IF EXISTS llm_key_route_config
ADD COLUMN IF NOT EXISTS kiro_model_group_preferences_json JSONB NOT NULL DEFAULT '{}'::jsonb;

UPDATE llm_key_route_config
SET kiro_model_group_preferences_json = '{}'::jsonb
WHERE kiro_model_group_preferences_json IS NULL
OR jsonb_typeof(kiro_model_group_preferences_json) <> 'object';

DO $$
BEGIN
IF NOT EXISTS (
SELECT 1
FROM pg_constraint
WHERE conname = 'ck_llm_key_route_config_kiro_model_group_preferences_object'
AND conrelid = 'llm_key_route_config'::regclass
) THEN
ALTER TABLE llm_key_route_config
ADD CONSTRAINT ck_llm_key_route_config_kiro_model_group_preferences_object
CHECK (jsonb_typeof(kiro_model_group_preferences_json) = 'object');
END IF;
END $$;
5 changes: 5 additions & 0 deletions crates/llm-access-migrations/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -178,6 +178,11 @@ const POSTGRES_MIGRATIONS: &[SqlMigration] = &[
name: "anthropic_upstream_probe_state",
sql: include_str!("../migrations/postgres/0034_anthropic_upstream_probe_state.sql"),
},
SqlMigration {
version: 35,
name: "kiro_model_group_preferences",
sql: include_str!("../migrations/postgres/0035_kiro_model_group_preferences.sql"),
},
];

/// Return target DuckDB migrations in execution order.
Expand Down
Loading
Loading