add models task

This commit is contained in:
Sakurasan
2025-04-21 01:40:06 +08:00
parent 470e49b850
commit 73e53c2333
2 changed files with 9 additions and 6 deletions

View File

@@ -10,7 +10,7 @@ import (
)
func (p *Proxy) HandleModels(c *gin.Context) {
models, err := p.getCache()
models, err := p.getModelCache()
if err != nil {
dto.Fail(c, http.StatusBadGateway, err.Error())
return
@@ -25,7 +25,7 @@ func (p *Proxy) HandleModels(c *gin.Context) {
dto.Success(c, ms)
}
func (p *Proxy) setCache() error {
func (p *Proxy) setModelCache() error {
apikeys, err := p.apiKeyDao.FindKeys(nil)
models := make(map[string]bool)
if err == nil && len(apikeys) > 0 {
@@ -52,7 +52,7 @@ func (p *Proxy) setCache() error {
return p.cache.Set("models", support_models)
}
func (p *Proxy) getCache() ([]string, error) {
func (p *Proxy) getModelCache() ([]string, error) {
models, err := p.cache.Get("models")
return models.([]string), err
}

View File

@@ -73,7 +73,7 @@ func NewProxy(ctx context.Context, cfg *config.Config, db *gorm.DB, wg *sync.Wai
go np.ProcessUsage()
go np.ScheduleTask()
np.setModelCache()
return np
}
@@ -189,7 +189,7 @@ func (p *Proxy) Do(usage *model.Usage) error {
func (p *Proxy) SelectApiKey(model string) error {
akpikeys, err := p.apiKeyDao.FindApiKeysBySupportModel(p.db, model)
fmt.Println(len(akpikeys), err)
if err != nil || len(akpikeys) == 0 {
if strings.HasPrefix(model, "gpt") || strings.HasPrefix(model, "o1") || strings.HasPrefix(model, "o3") || strings.HasPrefix(model, "o4") {
keys, err := p.apiKeyDao.FindKeys(map[string]any{"active = ?": true, "apitype = ?": "openai"})
@@ -275,7 +275,10 @@ func (p *Proxy) ScheduleTask() {
select {
case <-time.After(time.Duration(p.cfg.TaskTimeInterval) * time.Minute):
p.updateSupportModel()
case <-time.After(time.Hour * 12):
if err := p.setModelCache(); err != nil {
fmt.Println("refrash model cache err:", err)
}
case <-p.ctx.Done():
fmt.Println("schedule task done")
return