add models task
This commit is contained in:
@@ -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
|
||||
}
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user