//! Persists model and provider configuration. use std::{collections::HashSet, str::FromStr}; use sqlx::{Row, Sqlite, Transaction}; use crate::{ model::{model_hash, normalize_model_input, ModelConfig, ModelConfigInput, ModelType}, Error, Result, }; use super::{now_ms, Store}; const MODEL_COLUMNS: &str = r#" model_hash, sort_order, display_name, group_name, model_type, base_url, use_full_url, api_key, tooltip_data, model_id, reasoning_effort, openai_endpoint, openai_extra_params_enabled, openai_extra_params_json, custom_headers_enabled, custom_headers_json, anthropic_extra_params_enabled, anthropic_extra_params_json, context_window_tokens, max_completion_tokens, anthropic_max_tokens, anthropic_thinking_effort, thinking_budget_tokens, created_at_ms, updated_at_ms "#; impl Store { pub async fn models(&self) -> Result> { let query = format!("SELECT {MODEL_COLUMNS} FROM model_configs ORDER BY sort_order, display_name"); sqlx::query(&query) .fetch_all(&self.pool) .await? .into_iter() .map(model_from_row) .collect() } pub async fn model(&self, hash: &str) -> Result> { let query = format!("SELECT {MODEL_COLUMNS} FROM model_configs WHERE model_hash = ?"); sqlx::query(&query) .bind(hash) .fetch_optional(&self.pool) .await? .map(model_from_row) .transpose() } pub async fn create_model(&self, input: &ModelConfigInput) -> Result { let mut models = self.create_models(std::slice::from_ref(input)).await?; Ok(models.remove(0)) } pub async fn create_models(&self, inputs: &[ModelConfigInput]) -> Result> { if inputs.is_empty() { return Err(Error::Config("at least one model is required".into())); } let mut normalized = Vec::with_capacity(inputs.len()); let mut hashes = HashSet::with_capacity(inputs.len()); for input in inputs { let input = normalize_model_input(input)?; let hash = model_hash(&input)?; if !hashes.insert(hash.clone()) { return Err(Error::Config("model configurations must be unique".into())); } normalized.push((hash, input)); } let now = now_ms(); let _write = self.writes.lock().await; let mut transaction = self.pool.begin().await?; for (hash, input) in &normalized { insert_model(&mut transaction, hash, input, now).await?; } transaction.commit().await?; let mut saved = Vec::with_capacity(normalized.len()); for (hash, _) in normalized { saved.push(self.model(&hash).await?.expect("inserted model must exist")); } Ok(saved) } pub(super) async fn create_models_if_missing( &self, inputs: &[ModelConfigInput], ) -> Result { let mut normalized = Vec::with_capacity(inputs.len()); let mut hashes = HashSet::with_capacity(inputs.len()); for input in inputs { let input = normalize_model_input(input)?; let hash = model_hash(&input)?; if hashes.insert(hash.clone()) { normalized.push((hash, input)); } } let now = now_ms(); let _write = self.writes.lock().await; let mut transaction = self.pool.begin().await?; let mut inserted = 0; for (hash, input) in &normalized { inserted += usize::from( insert_model_with_conflict(&mut transaction, hash, input, now, true).await?, ); } transaction.commit().await?; Ok(inserted) } pub async fn update_model( &self, current_hash: &str, input: &ModelConfigInput, ) -> Result { let current = self .model(current_hash) .await? .ok_or_else(|| Error::RunNotFound(format!("model {current_hash}")))?; let input = normalize_model_input(input)?; let next_hash = model_hash(&input)?; let now = now_ms(); let _write = self.writes.lock().await; let mut transaction = self.pool.begin().await?; if next_hash != current.model_hash { sqlx::query("UPDATE llm_calls SET model_hash = NULL WHERE model_hash = ?") .bind(¤t.model_hash) .execute(&mut *transaction) .await?; } let result = sqlx::query( r#"UPDATE model_configs SET model_hash = ?, sort_order = ?, display_name = ?, group_name = ?, model_type = ?, base_url = ?, use_full_url = ?, api_key = ?, tooltip_data = ?, model_id = ?, reasoning_effort = ?, openai_endpoint = ?, openai_extra_params_enabled = ?, openai_extra_params_json = ?, custom_headers_enabled = ?, custom_headers_json = ?, anthropic_extra_params_enabled = ?, anthropic_extra_params_json = ?, context_window_tokens = ?, max_completion_tokens = ?, anthropic_max_tokens = ?, anthropic_thinking_effort = ?, thinking_budget_tokens = ?, updated_at_ms = ? WHERE model_hash = ?"#, ) .bind(&next_hash) .bind(input.sort_order) .bind(&input.display_name) .bind(&input.group_name) .bind(input.model_type.as_str()) .bind(&input.base_url) .bind(input.use_full_url) .bind(&input.api_key) .bind(&input.tooltip_data) .bind(&input.model_id) .bind(&input.reasoning_effort) .bind(&input.openai_endpoint) .bind(input.openai_extra_params_enabled) .bind(serde_json::to_string(&input.openai_extra_params)?) .bind(input.custom_headers_enabled) .bind(serde_json::to_string(&input.custom_headers)?) .bind(input.anthropic_extra_params_enabled) .bind(serde_json::to_string(&input.anthropic_extra_params)?) .bind(input.context_window_tokens.map(to_i64).transpose()?) .bind(input.max_completion_tokens.map(to_i64).transpose()?) .bind(input.anthropic_max_tokens.map(to_i64).transpose()?) .bind(&input.anthropic_thinking_effort) .bind(input.thinking_budget_tokens.map(to_i64).transpose()?) .bind(now) .bind(current_hash) .execute(&mut *transaction) .await?; if result.rows_affected() != 1 { return Err(Error::RunNotFound(format!("model {current_hash}"))); } transaction.commit().await?; Ok(self .model(&next_hash) .await? .expect("updated model must exist")) } pub async fn delete_model(&self, hash: &str) -> Result<()> { let _write = self.writes.lock().await; let mut transaction = self.pool.begin().await?; sqlx::query("UPDATE llm_calls SET model_hash = NULL WHERE model_hash = ?") .bind(hash) .execute(&mut *transaction) .await?; let result = sqlx::query("DELETE FROM model_configs WHERE model_hash = ?") .bind(hash) .execute(&mut *transaction) .await?; if result.rows_affected() != 1 { return Err(Error::RunNotFound(format!("model {hash}"))); } transaction.commit().await?; Ok(()) } pub async fn reorder_models(&self, model_hashes: &[String]) -> Result> { let current = self.models().await?; let current_hashes = current .iter() .map(|model| model.model_hash.as_str()) .collect::>(); let requested_hashes = model_hashes .iter() .map(String::as_str) .collect::>(); if model_hashes.len() != current.len() || requested_hashes.len() != current.len() || requested_hashes != current_hashes { return Err(Error::Config( "model configuration changed; refresh and try sorting again".into(), )); } let now = now_ms(); let _write = self.writes.lock().await; let mut transaction = self.pool.begin().await?; for (index, hash) in model_hashes.iter().enumerate() { sqlx::query( "UPDATE model_configs SET sort_order = ?, updated_at_ms = ? WHERE model_hash = ?", ) .bind(i64::try_from(index + 1).expect("model order fits in i64")) .bind(now) .bind(hash) .execute(&mut *transaction) .await?; } transaction.commit().await?; self.models().await } } async fn insert_model( transaction: &mut Transaction<'_, Sqlite>, hash: &str, input: &ModelConfigInput, now: i64, ) -> Result<()> { insert_model_with_conflict(transaction, hash, input, now, false).await?; Ok(()) } async fn insert_model_with_conflict( transaction: &mut Transaction<'_, Sqlite>, hash: &str, input: &ModelConfigInput, now: i64, ignore_existing: bool, ) -> Result { let mut statement = String::from( r#"INSERT INTO model_configs( model_hash, sort_order, display_name, group_name, model_type, base_url, use_full_url, api_key, tooltip_data, model_id, reasoning_effort, openai_endpoint, openai_extra_params_enabled, openai_extra_params_json, custom_headers_enabled, custom_headers_json, anthropic_extra_params_enabled, anthropic_extra_params_json, context_window_tokens, max_completion_tokens, anthropic_max_tokens, anthropic_thinking_effort, thinking_budget_tokens, created_at_ms, updated_at_ms ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)"#, ); if ignore_existing { statement.push_str(" ON CONFLICT(model_hash) DO NOTHING"); } let result = sqlx::query(&statement) .bind(hash) .bind(input.sort_order) .bind(&input.display_name) .bind(&input.group_name) .bind(input.model_type.as_str()) .bind(&input.base_url) .bind(input.use_full_url) .bind(&input.api_key) .bind(&input.tooltip_data) .bind(&input.model_id) .bind(&input.reasoning_effort) .bind(&input.openai_endpoint) .bind(input.openai_extra_params_enabled) .bind(serde_json::to_string(&input.openai_extra_params)?) .bind(input.custom_headers_enabled) .bind(serde_json::to_string(&input.custom_headers)?) .bind(input.anthropic_extra_params_enabled) .bind(serde_json::to_string(&input.anthropic_extra_params)?) .bind(input.context_window_tokens.map(to_i64).transpose()?) .bind(input.max_completion_tokens.map(to_i64).transpose()?) .bind(input.anthropic_max_tokens.map(to_i64).transpose()?) .bind(&input.anthropic_thinking_effort) .bind(input.thinking_budget_tokens.map(to_i64).transpose()?) .bind(now) .bind(now) .execute(&mut **transaction) .await?; Ok(result.rows_affected() == 1) } fn model_from_row(row: sqlx::sqlite::SqliteRow) -> Result { Ok(ModelConfig { model_hash: row.try_get("model_hash")?, sort_order: row.try_get("sort_order")?, display_name: row.try_get("display_name")?, group_name: row.try_get("group_name")?, model_type: ModelType::from_str(row.try_get("model_type")?)?, base_url: row.try_get("base_url")?, use_full_url: row.try_get("use_full_url")?, api_key: row.try_get("api_key")?, tooltip_data: row.try_get("tooltip_data")?, model_id: row.try_get("model_id")?, reasoning_effort: row.try_get("reasoning_effort")?, openai_endpoint: row.try_get("openai_endpoint")?, openai_extra_params_enabled: row.try_get("openai_extra_params_enabled")?, openai_extra_params: serde_json::from_str( row.try_get::("openai_extra_params_json")? .as_str(), )?, custom_headers_enabled: row.try_get("custom_headers_enabled")?, custom_headers: serde_json::from_str( row.try_get::("custom_headers_json")?.as_str(), )?, anthropic_extra_params_enabled: row.try_get("anthropic_extra_params_enabled")?, anthropic_extra_params: serde_json::from_str( row.try_get::("anthropic_extra_params_json")? .as_str(), )?, context_window_tokens: optional_u64(&row, "context_window_tokens")?, max_completion_tokens: optional_u64(&row, "max_completion_tokens")?, anthropic_max_tokens: optional_u64(&row, "anthropic_max_tokens")?, anthropic_thinking_effort: row.try_get("anthropic_thinking_effort")?, thinking_budget_tokens: optional_u64(&row, "thinking_budget_tokens")?, created_at_ms: row.try_get("created_at_ms")?, updated_at_ms: row.try_get("updated_at_ms")?, }) } fn optional_u64(row: &sqlx::sqlite::SqliteRow, column: &str) -> Result> { row.try_get::, _>(column)? .map(|value| { u64::try_from(value).map_err(|_| Error::Config(format!("{column} cannot be negative"))) }) .transpose() } fn to_i64(value: u64) -> Result { i64::try_from(value).map_err(|_| Error::Config("token value is too large".into())) } #[cfg(test)] mod tests { use super::*; fn model_input(group_name: Option<&str>) -> ModelConfigInput { ModelConfigInput { sort_order: 0, display_name: "Test Model".into(), group_name: group_name.map(String::from), model_type: ModelType::OpenAi, base_url: "https://example.com/v1/chat/completions".into(), use_full_url: true, api_key: "test-key".into(), tooltip_data: "Test Model".into(), model_id: "test-model".into(), reasoning_effort: None, openai_endpoint: crate::model::OPENAI_CHAT_ENDPOINT.into(), openai_extra_params_enabled: false, openai_extra_params: serde_json::json!({}), custom_headers_enabled: false, custom_headers: serde_json::json!({}), anthropic_extra_params_enabled: false, anthropic_extra_params: serde_json::json!({}), context_window_tokens: None, max_completion_tokens: None, anthropic_max_tokens: None, anthropic_thinking_effort: None, thinking_budget_tokens: None, } } /// 分组名是纯展示字段:入库时去除首尾空白、空串归一为 NULL, /// 更新分组名不得改变模型身份哈希。 #[tokio::test] async fn group_name_round_trips_without_changing_model_identity() { let directory = tempfile::tempdir().unwrap(); let store = Store::connect(&format!( "sqlite://{}", directory.path().join("test.db").display() )) .await .unwrap(); let created = store .create_model(&model_input(Some(" My Group "))) .await .unwrap(); assert_eq!(created.group_name.as_deref(), Some("My Group")); let renamed = store .update_model(&created.model_hash, &model_input(Some("Renamed"))) .await .unwrap(); assert_eq!(renamed.model_hash, created.model_hash); assert_eq!(renamed.group_name.as_deref(), Some("Renamed")); let cleared = store .update_model(&created.model_hash, &model_input(Some(" "))) .await .unwrap(); assert_eq!(cleared.model_hash, created.model_hash); assert_eq!(cleared.group_name, None); } }