From 73e53c233322bc61d16fee4fda7c26cf37e3e644 Mon Sep 17 00:00:00 2001
From: Sakurasan <26715255+Sakurasan@users.noreply.github.com>
Date: Mon, 21 Apr 2025 01:40:06 +0800
Subject: [PATCH] add models task

---
 internal/controller/proxy/models.go | 6 +++---
 internal/controller/proxy/proxy.go  | 9 ++++++---
 2 files changed, 9 insertions(+), 6 deletions(-)

diff --git a/internal/controller/proxy/models.go b/internal/controller/proxy/models.go
index df5fceb..497bab9 100644
--- a/internal/controller/proxy/models.go
+++ b/internal/controller/proxy/models.go
@@ -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
 }
diff --git a/internal/controller/proxy/proxy.go b/internal/controller/proxy/proxy.go
index 351ac8c..b11f81c 100644
--- a/internal/controller/proxy/proxy.go
+++ b/internal/controller/proxy/proxy.go
@@ -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