mirror of
https://gitcode.com/JianFeeeee/ModelRouter.git
synced 2026-09-20 08:57:57 +00:00
feat(auto): quota-aware AUTO slot scheduling (model tiers, hour/week/month resets) + key scope periods; fix key/slot create/delete UX
This commit is contained in:
@ -9,6 +9,7 @@ import (
|
||||
"fmt"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"sort"
|
||||
"time"
|
||||
|
||||
"llmsproxy/internal/config"
|
||||
@ -50,12 +51,48 @@ func NewFromConfig(cfg *config.Config) (*Core, error) {
|
||||
if err := c.seedKeys(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if err := c.seedAuto(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if err := c.rebuildRegistry(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return c, nil
|
||||
}
|
||||
|
||||
// seedAuto migrates the legacy per-source model priority into flat AUTO
|
||||
// scheduling slots (one slot per model, priority order) the first time no
|
||||
// explicit auto rules exist.
|
||||
func (c *Core) seedAuto() error {
|
||||
if len(c.store.AutoRules()) > 0 {
|
||||
return nil
|
||||
}
|
||||
type item struct {
|
||||
model string
|
||||
prio int
|
||||
}
|
||||
var flat []item
|
||||
for _, s := range c.mergedSources() {
|
||||
for _, m := range s.Models {
|
||||
if m.Kind == "image" {
|
||||
continue
|
||||
}
|
||||
flat = append(flat, item{m.ID, m.Priority})
|
||||
}
|
||||
}
|
||||
sort.SliceStable(flat, func(i, j int) bool {
|
||||
if flat[i].prio != flat[j].prio {
|
||||
return flat[i].prio > flat[j].prio
|
||||
}
|
||||
return flat[i].model < flat[j].model
|
||||
})
|
||||
entries := make([]config.ModelScope, 0, len(flat))
|
||||
for _, it := range flat {
|
||||
entries = append(entries, config.ModelScope{Model: it.model})
|
||||
}
|
||||
return c.store.SaveAutoRules(entries)
|
||||
}
|
||||
|
||||
// seedKeys migrates the static config gateway_keys into the runtime store as
|
||||
// admin keys (once), so later UI-created keys can share the same store.
|
||||
func (c *Core) seedKeys() error {
|
||||
@ -114,7 +151,7 @@ func (c *Core) ListKeys() []config.GWKey { return c.store.ListKeys() }
|
||||
func (c *Core) FindKey(key string) (config.GWKey, bool) { return c.store.KeyByValue(key) }
|
||||
|
||||
// CreateKey builds a new random gateway key and persists it.
|
||||
func (c *Core) CreateKey(name, role string, models []string, note string) (config.GWKey, error) {
|
||||
func (c *Core) CreateKey(name, role string, models []config.ModelScope, note string) (config.GWKey, error) {
|
||||
key := make([]byte, 16)
|
||||
if _, err := rand.Read(key); err != nil {
|
||||
return config.GWKey{}, err
|
||||
@ -137,7 +174,7 @@ func (c *Core) CreateKey(name, role string, models []string, note string) (confi
|
||||
}
|
||||
|
||||
// UpdateKey mutates a key's name/role/model scope and persists it.
|
||||
func (c *Core) UpdateKey(key, name, role string, models []string, note string) (config.GWKey, error) {
|
||||
func (c *Core) UpdateKey(key, name, role string, models []config.ModelScope, note string) (config.GWKey, error) {
|
||||
rec, ok := c.store.KeyByValue(key)
|
||||
if !ok {
|
||||
return config.GWKey{}, fmt.Errorf("key not found")
|
||||
@ -159,6 +196,29 @@ func (c *Core) UpdateKey(key, name, role string, models []string, note string) (
|
||||
// DeleteKey removes a key record; returns false if it did not exist.
|
||||
func (c *Core) DeleteKey(key string) (bool, error) { return c.store.DeleteKey(key) }
|
||||
|
||||
// ---- AUTO scheduling slots (web UI canvas) ----
|
||||
|
||||
// AutoRules returns the AUTO scheduling slots in priority order (slot 0 =
|
||||
// highest priority).
|
||||
func (c *Core) AutoRules() []config.ModelScope { return c.store.AutoRules() }
|
||||
|
||||
// SaveAutoRules persists the AUTO scheduling slots.
|
||||
func (c *Core) SaveAutoRules(entries []config.ModelScope) error {
|
||||
clean := make([]config.ModelScope, 0, len(entries))
|
||||
for _, e := range entries {
|
||||
if e.Model == "" {
|
||||
continue
|
||||
}
|
||||
clean = append(clean, e)
|
||||
}
|
||||
return c.store.SaveAutoRules(clean)
|
||||
}
|
||||
|
||||
// Registry resolves model -> owning provider.
|
||||
func (c *Core) ProviderForModel(model string) *provider.Provider {
|
||||
return c.registry.ProviderForModel(model)
|
||||
}
|
||||
|
||||
// Config exposes the underlying configuration (read-only usage).
|
||||
func (c *Core) Config() *config.Config { return c.cfg }
|
||||
|
||||
|
||||
Reference in New Issue
Block a user