mirror of
https://github.com/sipeed/picoclaw.git
synced 2026-06-12 18:08:54 +00:00
Add virtual model support for multi-key expansion
Virtual models generated from multi-key expansion are now marked and filtered during config persistence. Virtual models display with a badge in the UI and cannot be set as default.
This commit is contained in:
@@ -42,6 +42,7 @@ type modelResponse struct {
|
||||
// Meta
|
||||
Configured bool `json:"configured"`
|
||||
IsDefault bool `json:"is_default"`
|
||||
IsVirtual bool `json:"is_virtual"`
|
||||
}
|
||||
|
||||
// handleListModels returns all model_list entries with masked API keys.
|
||||
@@ -86,6 +87,7 @@ func (h *Handler) handleListModels(w http.ResponseWriter, r *http.Request) {
|
||||
ExtraBody: m.ExtraBody,
|
||||
Configured: configured[i],
|
||||
IsDefault: m.ModelName == defaultModel,
|
||||
IsVirtual: m.IsVirtual(),
|
||||
})
|
||||
}
|
||||
|
||||
@@ -288,11 +290,13 @@ func (h *Handler) handleSetDefaultModel(w http.ResponseWriter, r *http.Request)
|
||||
return
|
||||
}
|
||||
|
||||
// Verify the model_name exists in model_list
|
||||
// Verify the model_name exists in model_list and is not a virtual model
|
||||
found := false
|
||||
isVirtual := false
|
||||
for _, m := range cfg.ModelList {
|
||||
if m.ModelName == req.ModelName {
|
||||
found = true
|
||||
isVirtual = m.IsVirtual()
|
||||
break
|
||||
}
|
||||
}
|
||||
@@ -300,6 +304,10 @@ func (h *Handler) handleSetDefaultModel(w http.ResponseWriter, r *http.Request)
|
||||
http.Error(w, fmt.Sprintf("Model %q not found in model_list", req.ModelName), http.StatusNotFound)
|
||||
return
|
||||
}
|
||||
if isVirtual {
|
||||
http.Error(w, fmt.Sprintf("Cannot set virtual model %q as default", req.ModelName), http.StatusBadRequest)
|
||||
return
|
||||
}
|
||||
|
||||
cfg.Agents.Defaults.ModelName = req.ModelName
|
||||
|
||||
|
||||
@@ -356,6 +356,46 @@ func TestHandleAddModel_PersistsAPIKey(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
// TestHandleSetDefaultModel_RejectsNonexistentModel tests that setting a non-existent
|
||||
// model as default returns 404. This covers the case where virtual models (which are
|
||||
// filtered by SaveConfig) cannot be set as default.
|
||||
func TestHandleSetDefaultModel_RejectsNonexistentModel(t *testing.T) {
|
||||
configPath, cleanup := setupOAuthTestEnv(t)
|
||||
defer cleanup()
|
||||
|
||||
// First save a valid config with a primary model
|
||||
cfg, err := config.LoadConfig(configPath)
|
||||
if err != nil {
|
||||
t.Fatalf("LoadConfig() error = %v", err)
|
||||
}
|
||||
cfg.ModelList = []*config.ModelConfig{
|
||||
{ModelName: "gpt-4", Model: "openai/gpt-4o"},
|
||||
}
|
||||
if err := config.SaveConfig(configPath, cfg); err != nil {
|
||||
t.Fatalf("SaveConfig() error = %v", err)
|
||||
}
|
||||
|
||||
// Try to set a non-existent model (like a virtual model name) as default
|
||||
h := NewHandler(configPath)
|
||||
mux := http.NewServeMux()
|
||||
h.RegisterRoutes(mux)
|
||||
|
||||
rec := httptest.NewRecorder()
|
||||
req := httptest.NewRequest(http.MethodPost, "/api/models/default", bytes.NewBufferString(`{
|
||||
"model_name": "gpt-4__key_1"
|
||||
}`))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
mux.ServeHTTP(rec, req)
|
||||
|
||||
// Should return 404 because the virtual model doesn't exist in the persisted config
|
||||
if rec.Code != http.StatusNotFound {
|
||||
t.Fatalf("status = %d, want %d, body=%s", rec.Code, http.StatusNotFound, rec.Body.String())
|
||||
}
|
||||
if !strings.Contains(rec.Body.String(), "not found") {
|
||||
t.Fatalf("error message should mention 'not found', got: %s", rec.Body.String())
|
||||
}
|
||||
}
|
||||
|
||||
func TestMaskAPIKey(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
|
||||
Reference in New Issue
Block a user