Skip to content

fix(providers): allow replacing preferred models #1104

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
penso wants to merge 1 commit into
base: main
Choose a base branch
from
Open

fix(providers): allow replacing preferred models #1104

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
116 changes: 109 additions & 7 deletions crates/provider-setup/src/service/implementation/credentials.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -232,13 +232,10 @@ impl LiveProviderSetupService {
.and_then(|v| v.as_str())
.ok_or_else(|| "missing 'provider' parameter".to_string())?;

let models: Vec<String> = params
.get("models")
.and_then(|v| v.as_array())
.ok_or_else(|| "missing 'models' array parameter".to_string())?
.iter()
.filter_map(|v| v.as_str().map(String::from))
.collect();
if !params.get("models").is_some_and(Value::is_array) {
return Err("missing 'models' array parameter".into());
}
let models = parse_models_param(&params);

// Validate provider exists (known or custom).
if !is_custom_provider(provider_name) {
Expand All @@ -248,13 +245,22 @@ impl LiveProviderSetupService {
}
}

let previous_models = self
.key_store
.load_config(provider_name)
.map(|config| config.models)
.unwrap_or_default();

self.key_store
.save_config(provider_name, None, None, Some(models.clone()))
.map_err(ServiceError::message)?;

// Update the cross-provider priority list.
if let Some(ref priority) = self.priority_models {
let mut list = priority.write().await;
for previous in previous_models {
list.retain(|existing| existing != &previous);
}
for m in models.iter().rev() {
list.retain(|existing| existing != m);
list.insert(0, m.clone());
Expand All @@ -271,3 +277,99 @@ impl LiveProviderSetupService {
Ok(serde_json::json!({ "ok": true }))
}
}

#[allow(clippy::unwrap_used, clippy::expect_used)]
#[cfg(test)]
mod tests {
use std::sync::Arc;

use tokio::sync::RwLock;

use {
super::LiveProviderSetupService, crate::key_store::KeyStore,
moltis_config::schema::ProvidersConfig, moltis_providers::ProviderRegistry,
};

fn test_service(
key_store: KeyStore,
priority: Arc<RwLock<Vec<String>>>,
) -> LiveProviderSetupService {
let mut service = LiveProviderSetupService::new(
Arc::new(RwLock::new(ProviderRegistry::empty())),
ProvidersConfig::default(),
None,
);
service.key_store = key_store;
service.set_priority_models(priority);
service
}

#[tokio::test]
async fn save_models_replaces_previous_provider_priorities() {
let dir = tempfile::tempdir().unwrap();
let key_store = KeyStore::with_path(dir.path().join("provider_keys.json"));
key_store
.save_config(
"openai",
Some("sk-test".to_string()),
None,
Some(vec!["old-a".to_string(), "old-b".to_string()]),
)
.unwrap();
let priority = Arc::new(RwLock::new(vec![
"other-provider-model".to_string(),
"old-a".to_string(),
"old-b".to_string(),
]));
let service = test_service(key_store.clone(), Arc::clone(&priority));

service
.save_models_inner(serde_json::json!({
"provider": "openai",
"models": ["openai::new-a"]
}))
.await
.unwrap();

assert_eq!(key_store.load_config("openai").unwrap().models, vec![
"new-a"
]);
assert_eq!(*priority.read().await, vec![
"new-a".to_string(),
"other-provider-model".to_string()
]);
}

#[tokio::test]
async fn save_models_empty_selection_clears_previous_provider_priorities() {
let dir = tempfile::tempdir().unwrap();
let key_store = KeyStore::with_path(dir.path().join("provider_keys.json"));
key_store
.save_config(
"openai",
Some("sk-test".to_string()),
None,
Some(vec!["old-a".to_string(), "old-b".to_string()]),
)
.unwrap();
let priority = Arc::new(RwLock::new(vec![
"old-a".to_string(),
"other-provider-model".to_string(),
"old-b".to_string(),
]));
let service = test_service(key_store.clone(), Arc::clone(&priority));

service
.save_models_inner(serde_json::json!({
"provider": "openai",
"models": []
}))
.await
.unwrap();

assert!(key_store.load_config("openai").unwrap().models.is_empty());
assert_eq!(*priority.read().await, vec![
"other-provider-model".to_string()
]);
}
}
86 changes: 86 additions & 0 deletions crates/web/ui/e2e/specs/providers.spec.js
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -186,4 +186,90 @@ test.describe("Provider setup page", () => {
.toContain("Custom.gguf configured successfully!");
expect(pageErrors).toEqual([]);
});

test("preferred model selector replaces existing selections", async ({ page }) => {
const pageErrors = watchPageErrors(page);
await openProvidersPage(page);

await page.evaluate(async () => {
const [providers, state] = await Promise.all([import("/assets/js/providers.js"), import("/assets/js/state.js")]);
window.__providerSaveModelsCalls = [];

state.setWs({
readyState: WebSocket.OPEN,
send(raw) {
const frame = JSON.parse(raw);
const respond = (payload) => {
queueMicrotask(() => {
const pending = state.pending[frame.id];
if (!pending) return;
pending(payload);
delete state.pending[frame.id];
});
};

if (frame.method === "models.list") {
respond({
ok: true,
payload: [
{
id: "openai::gpt-preferred-a",
provider: "openai",
displayName: "GPT Preferred A",
supportsTools: true,
preferred: true,
},
{
id: "openai::gpt-preferred-b",
provider: "openai",
displayName: "GPT Preferred B",
supportsTools: true,
preferred: true,
},
],
});
return;
}

if (frame.method === "providers.available") {
respond({
ok: true,
payload: [
{
name: "openai",
displayName: "OpenAI",
authType: "api-key",
models: ["gpt-preferred-a", "gpt-preferred-b"],
},
],
});
return;
}

if (frame.method === "providers.save_models") {
window.__providerSaveModelsCalls.push(frame.params);
respond({ ok: true, payload: { ok: true } });
return;
}

respond({ ok: true, payload: {} });
},
});

providers.openModelSelectorForProvider("openai", "OpenAI");
});

await expect(page.locator("#providerModal")).toBeVisible();
await expect(page.locator("#providerModalBody .model-card.selected")).toHaveCount(2);

await page.locator("#providerModalBody .model-card", { hasText: "GPT Preferred B" }).click();
await expect(page.locator("#providerModalBody .model-card.selected")).toHaveCount(1);
await page.getByRole("button", { name: "Save", exact: true }).click();

await expect
.poll(() => page.evaluate(() => window.__providerSaveModelsCalls || []))
.toEqual([{ provider: "openai", models: ["openai::gpt-preferred-a"] }]);

expect(pageErrors).toEqual([]);
});
});
7 changes: 6 additions & 1 deletion crates/web/ui/src/providers/auth-flow.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -802,7 +802,12 @@ function showMultiModelSelector(
m.title.textContent = `${providerDisplayName} \u2014 Preferred Models`;
m.body.textContent = "";

const selectedIds: Set<string> = new Set(savedModels);
const savedModelTokens = new Set(Array.from(savedModels, stripModelNamespace));
const selectedIds: Set<string> = new Set(
models
.filter((model) => savedModels.has(model.id) || savedModelTokens.has(stripModelNamespace(model.id)))
.map((model) => model.id),
);

// Track per-model probe state: "probing" | "ok" | { error: string }
const probeResults: Map<string, string | ProbeResult> = new Map();
Expand Down
Loading