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:
@@ -857,6 +857,10 @@ type ModelConfig struct {
|
||||
secModelName string
|
||||
apiKeys []string
|
||||
secDirty bool
|
||||
|
||||
// isVirtual marks this model as a virtual model generated from multi-key expansion.
|
||||
// Virtual models should not be persisted to config files.
|
||||
isVirtual bool
|
||||
}
|
||||
|
||||
// APIKey returns the first API key from apiKeys
|
||||
@@ -867,6 +871,11 @@ func (c *ModelConfig) APIKey() string {
|
||||
return ""
|
||||
}
|
||||
|
||||
// IsVirtual returns true if this model was generated from multi-key expansion.
|
||||
func (c *ModelConfig) IsVirtual() bool {
|
||||
return c.isVirtual
|
||||
}
|
||||
|
||||
// Validate checks if the ModelConfig has all required fields.
|
||||
func (c *ModelConfig) Validate() error {
|
||||
if c.ModelName == "" {
|
||||
@@ -1800,7 +1809,20 @@ func SaveConfig(path string, cfg *Config) error {
|
||||
return err
|
||||
}
|
||||
|
||||
// Filter out virtual models before serializing to config file
|
||||
nonVirtualModels := make([]*ModelConfig, 0, len(cfg.ModelList))
|
||||
for _, m := range cfg.ModelList {
|
||||
if !m.isVirtual {
|
||||
nonVirtualModels = append(nonVirtualModels, m)
|
||||
}
|
||||
}
|
||||
// Temporarily replace ModelList with filtered version for serialization
|
||||
originalModelList := cfg.ModelList
|
||||
cfg.ModelList = nonVirtualModels
|
||||
|
||||
data, err := json.MarshalIndent(cfg, "", " ")
|
||||
// Restore original ModelList after serialization
|
||||
cfg.ModelList = originalModelList
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
@@ -2018,6 +2040,7 @@ func expandMultiKeyModels(models []*ModelConfig) []*ModelConfig {
|
||||
RequestTimeout: m.RequestTimeout,
|
||||
ThinkingLevel: m.ThinkingLevel,
|
||||
ExtraBody: m.ExtraBody,
|
||||
isVirtual: true,
|
||||
}
|
||||
expanded = append(expanded, additionalEntry)
|
||||
fallbackNames = append(fallbackNames, expandedName)
|
||||
|
||||
@@ -391,6 +391,65 @@ func TestSaveConfig_PreservesDisabledTelegramPlaceholder(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
// TestSaveConfig_FiltersVirtualModels verifies that SaveConfig does not write
|
||||
// virtual models (generated by expandMultiKeyModels) to the config file.
|
||||
func TestSaveConfig_FiltersVirtualModels(t *testing.T) {
|
||||
tmpDir := t.TempDir()
|
||||
path := filepath.Join(tmpDir, "config.json")
|
||||
|
||||
cfg := DefaultConfig()
|
||||
|
||||
// Manually add a virtual model to ModelList (simulating what expandMultiKeyModels does)
|
||||
primaryModel := &ModelConfig{
|
||||
ModelName: "gpt-4",
|
||||
Model: "openai/gpt-4o",
|
||||
apiKeys: []string{"key1"},
|
||||
}
|
||||
virtualModel := &ModelConfig{
|
||||
ModelName: "gpt-4__key_1",
|
||||
Model: "openai/gpt-4o",
|
||||
apiKeys: []string{"key2"},
|
||||
isVirtual: true,
|
||||
}
|
||||
cfg.ModelList = []*ModelConfig{primaryModel, virtualModel}
|
||||
|
||||
// SaveConfig should filter out virtual models
|
||||
if err := SaveConfig(path, cfg); err != nil {
|
||||
t.Fatalf("SaveConfig failed: %v", err)
|
||||
}
|
||||
|
||||
// Reload and verify
|
||||
reloaded, err := LoadConfig(path)
|
||||
if err != nil {
|
||||
t.Fatalf("LoadConfig failed: %v", err)
|
||||
}
|
||||
|
||||
// Should only have the primary model, not the virtual one
|
||||
if len(reloaded.ModelList) != 1 {
|
||||
t.Fatalf("expected 1 model after reload, got %d", len(reloaded.ModelList))
|
||||
}
|
||||
|
||||
if reloaded.ModelList[0].ModelName != "gpt-4" {
|
||||
t.Errorf("expected model_name 'gpt-4', got %q", reloaded.ModelList[0].ModelName)
|
||||
}
|
||||
|
||||
// Verify virtual model was not persisted
|
||||
for _, m := range reloaded.ModelList {
|
||||
if m.ModelName == "gpt-4__key_1" {
|
||||
t.Errorf("virtual model gpt-4__key_1 should not have been saved")
|
||||
}
|
||||
}
|
||||
|
||||
// Verify the saved file does not contain the virtual model name
|
||||
data, err := os.ReadFile(path)
|
||||
if err != nil {
|
||||
t.Fatalf("ReadFile failed: %v", err)
|
||||
}
|
||||
if strings.Contains(string(data), "gpt-4__key_1") {
|
||||
t.Errorf("saved config should not contain virtual model name 'gpt-4__key_1'")
|
||||
}
|
||||
}
|
||||
|
||||
// TestConfig_Complete verifies all config fields are set
|
||||
func TestConfig_Complete(t *testing.T) {
|
||||
cfg := DefaultConfig()
|
||||
|
||||
@@ -232,6 +232,78 @@ func TestExpandMultiKeyModels_PreservesOtherFields(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestExpandMultiKeyModels_IsVirtualFlag(t *testing.T) {
|
||||
models := []*ModelConfig{
|
||||
{
|
||||
ModelName: "gpt-4",
|
||||
Model: "openai/gpt-4o",
|
||||
apiKeys: []string{"key1", "key2", "key3"},
|
||||
},
|
||||
}
|
||||
|
||||
result := expandMultiKeyModels(models)
|
||||
|
||||
// Should expand to 3 models
|
||||
if len(result) != 3 {
|
||||
t.Fatalf("expected 3 models, got %d", len(result))
|
||||
}
|
||||
|
||||
// Primary model should NOT be virtual
|
||||
primary := result[2]
|
||||
if primary.isVirtual {
|
||||
t.Errorf("primary model should not be virtual")
|
||||
}
|
||||
if primary.ModelName != "gpt-4" {
|
||||
t.Errorf("expected primary model_name 'gpt-4', got %q", primary.ModelName)
|
||||
}
|
||||
|
||||
// Virtual models should have isVirtual = true
|
||||
virtual1 := result[0]
|
||||
if !virtual1.isVirtual {
|
||||
t.Errorf("gpt-4__key_1 should be virtual")
|
||||
}
|
||||
if virtual1.ModelName != "gpt-4__key_1" {
|
||||
t.Errorf("expected virtual model_name 'gpt-4__key_1', got %q", virtual1.ModelName)
|
||||
}
|
||||
|
||||
virtual2 := result[1]
|
||||
if !virtual2.isVirtual {
|
||||
t.Errorf("gpt-4__key_2 should be virtual")
|
||||
}
|
||||
if virtual2.ModelName != "gpt-4__key_2" {
|
||||
t.Errorf("expected virtual model_name 'gpt-4__key_2', got %q", virtual2.ModelName)
|
||||
}
|
||||
|
||||
// IsVirtual() method should work
|
||||
if !virtual1.IsVirtual() {
|
||||
t.Errorf("IsVirtual() should return true for virtual model")
|
||||
}
|
||||
if primary.IsVirtual() {
|
||||
t.Errorf("IsVirtual() should return false for primary model")
|
||||
}
|
||||
}
|
||||
|
||||
func TestExpandMultiKeyModels_SingleKey_NotVirtual(t *testing.T) {
|
||||
models := []*ModelConfig{
|
||||
{
|
||||
ModelName: "gpt-4",
|
||||
Model: "openai/gpt-4o",
|
||||
apiKeys: []string{"single-key"},
|
||||
},
|
||||
}
|
||||
|
||||
result := expandMultiKeyModels(models)
|
||||
|
||||
if len(result) != 1 {
|
||||
t.Fatalf("expected 1 model, got %d", len(result))
|
||||
}
|
||||
|
||||
// Single key model should NOT be virtual
|
||||
if result[0].isVirtual {
|
||||
t.Errorf("single key model should not be virtual")
|
||||
}
|
||||
}
|
||||
|
||||
func TestMergeAPIKeys(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -21,6 +21,7 @@ export interface ModelInfo {
|
||||
// Meta
|
||||
configured: boolean
|
||||
is_default: boolean
|
||||
is_virtual: boolean
|
||||
}
|
||||
|
||||
interface ModelsListResponse {
|
||||
|
||||
@@ -28,7 +28,7 @@ export function ModelCard({
|
||||
}: ModelCardProps) {
|
||||
const { t } = useTranslation()
|
||||
const isOAuth = model.auth_method === "oauth"
|
||||
const canSetDefault = model.configured && !model.is_default
|
||||
const canSetDefault = model.configured && !model.is_default && !model.is_virtual
|
||||
|
||||
return (
|
||||
<div
|
||||
@@ -64,6 +64,11 @@ export function ModelCard({
|
||||
{t("models.badge.default")}
|
||||
</span>
|
||||
)}
|
||||
{model.is_virtual && (
|
||||
<span className="bg-muted text-muted-foreground shrink-0 rounded px-1.5 py-0.5 text-[10px] leading-none font-medium">
|
||||
{t("models.badge.virtual")}
|
||||
</span>
|
||||
)}
|
||||
</div>
|
||||
|
||||
<div className="flex shrink-0 items-center gap-0.5">
|
||||
|
||||
@@ -154,7 +154,8 @@
|
||||
"unconfigured": "Not configured"
|
||||
},
|
||||
"badge": {
|
||||
"default": "Default"
|
||||
"default": "Default",
|
||||
"virtual": "Virtual"
|
||||
},
|
||||
"action": {
|
||||
"edit": "Edit API key",
|
||||
|
||||
@@ -154,7 +154,8 @@
|
||||
"unconfigured": "未配置"
|
||||
},
|
||||
"badge": {
|
||||
"default": "默认"
|
||||
"default": "默认",
|
||||
"virtual": "虚拟"
|
||||
},
|
||||
"action": {
|
||||
"edit": "编辑 API Key",
|
||||
|
||||
Reference in New Issue
Block a user