mirrors2/opencatd-open
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

add models task

Browse Source
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
Download Patch File Download Diff File Expand all files Collapse all files

6
internal/controller/proxy/models.go
View File

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

9
internal/controller/proxy/proxy.go
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.ProcessUsage()
go np.ScheduleTask() go np.ScheduleTask()
np.setModelCache()
return np return np
} }
@@ -189,7 +189,7 @@ func (p *Proxy) Do(usage *model.Usage) error {
func (p *Proxy) SelectApiKey(model string) error { func (p *Proxy) SelectApiKey(model string) error {
akpikeys, err := p.apiKeyDao.FindApiKeysBySupportModel(p.db, model) akpikeys, err := p.apiKeyDao.FindApiKeysBySupportModel(p.db, model)
fmt.Println(len(akpikeys), err)
if err != nil || len(akpikeys) == 0 { if err != nil || len(akpikeys) == 0 {
if strings.HasPrefix(model, "gpt") || strings.HasPrefix(model, "o1") || strings.HasPrefix(model, "o3") || strings.HasPrefix(model, "o4") { 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"}) keys, err := p.apiKeyDao.FindKeys(map[string]any{"active = ?": true, "apitype = ?": "openai"})
@@ -275,7 +275,10 @@ func (p *Proxy) ScheduleTask() {
select { select {
case <-time.After(time.Duration(p.cfg.TaskTimeInterval) * time.Minute): case <-time.After(time.Duration(p.cfg.TaskTimeInterval) * time.Minute):
p.updateSupportModel() 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(): case <-p.ctx.Done():
fmt.Println("schedule task done") fmt.Println("schedule task done")
return return