Persist discovered AI endpoint models in the database per endpoint.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
@@ -0,0 +1 @@
|
||||
DROP TABLE ai_endpoint_models;
|
||||
@@ -0,0 +1,11 @@
|
||||
CREATE TABLE ai_endpoint_models (
|
||||
kind TEXT NOT NULL,
|
||||
model_id TEXT NOT NULL,
|
||||
label TEXT NOT NULL,
|
||||
context_window INTEGER,
|
||||
max_output_tokens INTEGER,
|
||||
supports_tools INTEGER NOT NULL DEFAULT 0,
|
||||
supports_vision INTEGER NOT NULL DEFAULT 0,
|
||||
updated_at BIGINT NOT NULL,
|
||||
PRIMARY KEY (kind, model_id)
|
||||
);
|
||||
@@ -15,7 +15,8 @@ mod tests {
|
||||
use super::*;
|
||||
use crate::db::Database;
|
||||
use crate::db::schema::{
|
||||
ai_catalog_meta, ai_model_modalities, ai_models, ai_providers, chat_conversations,
|
||||
ai_catalog_meta, ai_endpoint_models, ai_model_modalities, ai_models, ai_providers,
|
||||
chat_conversations,
|
||||
chat_messages, db_notifications, dismissed_duplicate_pairs, embedding_keys,
|
||||
generated_file_hashes, import_definitions, mcp_proposals, media, media_translations,
|
||||
post_links, post_media, post_translations, posts, projects, scripts, settings, tags,
|
||||
@@ -36,7 +37,7 @@ mod tests {
|
||||
let applied = db
|
||||
.conn()
|
||||
.with_migrations(|conn| conn.applied_migrations().unwrap().len());
|
||||
assert_eq!(applied, 7);
|
||||
assert_eq!(applied, 8);
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -73,6 +74,7 @@ mod tests {
|
||||
ai_models::table,
|
||||
ai_model_modalities::table,
|
||||
ai_catalog_meta::table,
|
||||
ai_endpoint_models::table,
|
||||
embedding_keys::table,
|
||||
dismissed_duplicate_pairs::table,
|
||||
import_definitions::table,
|
||||
|
||||
@@ -7,6 +7,19 @@ diesel::table! {
|
||||
}
|
||||
}
|
||||
|
||||
diesel::table! {
|
||||
ai_endpoint_models (kind, model_id) {
|
||||
kind -> Text,
|
||||
model_id -> Text,
|
||||
label -> Text,
|
||||
context_window -> Nullable<Integer>,
|
||||
max_output_tokens -> Nullable<Integer>,
|
||||
supports_tools -> Integer,
|
||||
supports_vision -> Integer,
|
||||
updated_at -> BigInt,
|
||||
}
|
||||
}
|
||||
|
||||
diesel::table! {
|
||||
ai_model_modalities (rowid) {
|
||||
rowid -> Integer,
|
||||
@@ -64,9 +77,9 @@ diesel::table! {
|
||||
title -> Text,
|
||||
model -> Nullable<Text>,
|
||||
copilot_session_id -> Nullable<Text>,
|
||||
surface_state -> Nullable<Text>,
|
||||
created_at -> BigInt,
|
||||
updated_at -> BigInt,
|
||||
surface_state -> Nullable<Text>,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -79,10 +92,10 @@ diesel::table! {
|
||||
tool_call_id -> Nullable<Text>,
|
||||
tool_calls -> Nullable<Text>,
|
||||
created_at -> BigInt,
|
||||
cache_read_tokens -> Nullable<Integer>,
|
||||
cache_write_tokens -> Nullable<Integer>,
|
||||
token_usage_input -> Nullable<Integer>,
|
||||
token_usage_output -> Nullable<Integer>,
|
||||
cache_read_tokens -> Nullable<Integer>,
|
||||
cache_write_tokens -> Nullable<Integer>,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -351,6 +364,7 @@ diesel::joinable!(templates -> projects (project_id));
|
||||
|
||||
diesel::allow_tables_to_appear_in_same_query!(
|
||||
ai_catalog_meta,
|
||||
ai_endpoint_models,
|
||||
ai_model_modalities,
|
||||
ai_models,
|
||||
ai_providers,
|
||||
|
||||
@@ -65,6 +65,7 @@ pub struct AiModeSettings {
|
||||
pub image_model: Option<String>,
|
||||
pub chat_supports_tools: Option<bool>,
|
||||
pub image_supports_vision: Option<bool>,
|
||||
pub models: Vec<AiModelInfo>,
|
||||
}
|
||||
|
||||
impl AiSettings {
|
||||
@@ -301,6 +302,78 @@ pub fn save_model_preferences(
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub fn save_endpoint_models(
|
||||
conn: &Connection,
|
||||
kind: AiEndpointKind,
|
||||
models: &[AiModelInfo],
|
||||
) -> EngineResult<()> {
|
||||
use crate::db::schema::ai_endpoint_models::dsl;
|
||||
use diesel::prelude::*;
|
||||
|
||||
let updated_at = now_unix_ms();
|
||||
conn.with(|connection| {
|
||||
connection.transaction(|connection| {
|
||||
diesel::delete(dsl::ai_endpoint_models.filter(dsl::kind.eq(kind.as_str())))
|
||||
.execute(connection)?;
|
||||
for model in models {
|
||||
diesel::insert_into(dsl::ai_endpoint_models)
|
||||
.values((
|
||||
dsl::kind.eq(kind.as_str()),
|
||||
dsl::model_id.eq(&model.id),
|
||||
dsl::label.eq(&model.name),
|
||||
dsl::context_window.eq(model.context_window.map(|value| value as i32)),
|
||||
dsl::max_output_tokens
|
||||
.eq(model.max_output_tokens.map(|value| value as i32)),
|
||||
dsl::supports_tools.eq(i32::from(model.supports_tools)),
|
||||
dsl::supports_vision.eq(i32::from(model.supports_vision)),
|
||||
dsl::updated_at.eq(updated_at),
|
||||
))
|
||||
.execute(connection)?;
|
||||
}
|
||||
diesel::QueryResult::Ok(())
|
||||
})
|
||||
})?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub fn load_endpoint_models(
|
||||
conn: &Connection,
|
||||
kind: AiEndpointKind,
|
||||
) -> EngineResult<Vec<AiModelInfo>> {
|
||||
use crate::db::schema::ai_endpoint_models::dsl;
|
||||
use diesel::prelude::*;
|
||||
|
||||
let rows = conn.with(|connection| {
|
||||
dsl::ai_endpoint_models
|
||||
.filter(dsl::kind.eq(kind.as_str()))
|
||||
.order(dsl::label.asc())
|
||||
.select((
|
||||
dsl::model_id,
|
||||
dsl::label,
|
||||
dsl::context_window,
|
||||
dsl::max_output_tokens,
|
||||
dsl::supports_tools,
|
||||
dsl::supports_vision,
|
||||
))
|
||||
.load::<(String, String, Option<i32>, Option<i32>, i32, i32)>(connection)
|
||||
})?;
|
||||
Ok(rows
|
||||
.into_iter()
|
||||
.map(
|
||||
|(id, label, context_window, max_output_tokens, supports_tools, supports_vision)| {
|
||||
AiModelInfo {
|
||||
id,
|
||||
name: label,
|
||||
context_window: context_window.map(|value| value.max(0) as u64),
|
||||
max_output_tokens: max_output_tokens.map(|value| value.max(0) as u64),
|
||||
supports_tools: supports_tools != 0,
|
||||
supports_vision: supports_vision != 0,
|
||||
}
|
||||
},
|
||||
)
|
||||
.collect())
|
||||
}
|
||||
|
||||
pub fn save_system_prompt(conn: &Connection, system_prompt: &str) -> EngineResult<()> {
|
||||
set_setting(conn, "ai.system_prompt", system_prompt, now_unix_ms())
|
||||
}
|
||||
@@ -550,6 +623,7 @@ fn load_mode_settings(
|
||||
conn,
|
||||
&endpoint_setting_key(kind, "image_supports_vision"),
|
||||
)?,
|
||||
models: load_endpoint_models(conn, kind)?,
|
||||
})
|
||||
}
|
||||
|
||||
@@ -1033,6 +1107,47 @@ mod tests {
|
||||
assert!(models[1].supports_vision);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn endpoint_models_persist_per_endpoint_and_overwrite_on_refresh() {
|
||||
let db = setup();
|
||||
let online = vec![AiModelInfo {
|
||||
id: "gpt-4.1".to_string(),
|
||||
name: "GPT 4.1".to_string(),
|
||||
context_window: Some(128_000),
|
||||
max_output_tokens: Some(8_192),
|
||||
supports_tools: true,
|
||||
supports_vision: true,
|
||||
}];
|
||||
let airplane = vec![AiModelInfo {
|
||||
id: "llama3.2".to_string(),
|
||||
name: "Llama 3.2".to_string(),
|
||||
context_window: None,
|
||||
max_output_tokens: None,
|
||||
supports_tools: false,
|
||||
supports_vision: false,
|
||||
}];
|
||||
save_endpoint_models(db.conn(), AiEndpointKind::Online, &online).unwrap();
|
||||
save_endpoint_models(db.conn(), AiEndpointKind::Airplane, &airplane).unwrap();
|
||||
|
||||
let settings = load_ai_settings(db.conn(), false).unwrap();
|
||||
assert_eq!(settings.online.models, online);
|
||||
assert_eq!(settings.airplane.models, airplane);
|
||||
|
||||
let refreshed = vec![AiModelInfo {
|
||||
id: "gpt-5".to_string(),
|
||||
name: "GPT 5".to_string(),
|
||||
context_window: Some(256_000),
|
||||
max_output_tokens: Some(16_384),
|
||||
supports_tools: true,
|
||||
supports_vision: false,
|
||||
}];
|
||||
save_endpoint_models(db.conn(), AiEndpointKind::Online, &refreshed).unwrap();
|
||||
|
||||
let settings = load_ai_settings(db.conn(), false).unwrap();
|
||||
assert_eq!(settings.online.models, refreshed);
|
||||
assert_eq!(settings.airplane.models, airplane);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn model_preferences_are_independent_per_endpoint() {
|
||||
let db = setup();
|
||||
|
||||
@@ -3563,34 +3563,27 @@ impl BdsApp {
|
||||
}
|
||||
|
||||
fn chat_model_options(&self) -> Vec<ChatModelChoice> {
|
||||
let mut models = self
|
||||
.settings_state
|
||||
.as_ref()
|
||||
.map(|state| {
|
||||
let mode = if self.offline_mode {
|
||||
&state.airplane_ai
|
||||
} else {
|
||||
&state.online_ai
|
||||
};
|
||||
mode.model_options
|
||||
.iter()
|
||||
.map(|model| ChatModelChoice {
|
||||
id: model.id.clone(),
|
||||
label: model.label.clone(),
|
||||
})
|
||||
.collect::<Vec<_>>()
|
||||
let Some(db) = &self.db else {
|
||||
return Vec::new();
|
||||
};
|
||||
let Ok(settings) = ai::load_ai_settings(db.conn(), self.offline_mode) else {
|
||||
return Vec::new();
|
||||
};
|
||||
let active = settings.active();
|
||||
let mut models = active
|
||||
.models
|
||||
.iter()
|
||||
.map(|model| ChatModelChoice {
|
||||
id: model.id.clone(),
|
||||
label: model.name.clone(),
|
||||
})
|
||||
.unwrap_or_default();
|
||||
if let Some(db) = &self.db
|
||||
&& let Ok(settings) = ai::load_ai_settings(db.conn(), self.offline_mode)
|
||||
{
|
||||
let model = settings.active().endpoint.model.clone();
|
||||
if !model.trim().is_empty() && !models.iter().any(|choice| choice.id == model) {
|
||||
models.push(ChatModelChoice {
|
||||
id: model.clone(),
|
||||
label: model,
|
||||
});
|
||||
}
|
||||
.collect::<Vec<_>>();
|
||||
let model = active.endpoint.model.clone();
|
||||
if !model.trim().is_empty() && !models.iter().any(|choice| choice.id == model) {
|
||||
models.push(ChatModelChoice {
|
||||
id: model.clone(),
|
||||
label: model,
|
||||
});
|
||||
}
|
||||
models.sort_by(|left, right| left.label.cmp(&right.label));
|
||||
models.dedup_by(|left, right| left.id == right.id);
|
||||
@@ -8150,12 +8143,13 @@ impl BdsApp {
|
||||
}
|
||||
|
||||
fn refresh_ai_models(
|
||||
_db: &Database,
|
||||
db: &Database,
|
||||
state: &mut SettingsViewState,
|
||||
kind: AiEndpointKind,
|
||||
) -> Result<(), String> {
|
||||
let endpoint = Self::compose_ai_endpoint(state, kind)?;
|
||||
let models = ai::refresh_model_catalog(&endpoint).map_err(|error| error.to_string())?;
|
||||
ai::save_endpoint_models(db.conn(), kind, &models).map_err(|error| error.to_string())?;
|
||||
let options = models
|
||||
.into_iter()
|
||||
.map(|model| AiModelOption {
|
||||
@@ -8273,6 +8267,16 @@ impl BdsApp {
|
||||
api_key_configured: settings.endpoint.api_key_configured,
|
||||
chat_supports_tools: settings.chat_supports_tools.unwrap_or(false),
|
||||
image_supports_vision: settings.image_supports_vision.unwrap_or(false),
|
||||
model_options: settings
|
||||
.models
|
||||
.into_iter()
|
||||
.map(|model| AiModelOption {
|
||||
id: model.id,
|
||||
label: model.name,
|
||||
supports_tools: model.supports_tools,
|
||||
supports_vision: model.supports_vision,
|
||||
})
|
||||
.collect(),
|
||||
..Default::default()
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user