support fetch models
This commit is contained in:
58
internal/controller/proxy/models.go
Normal file
58
internal/controller/proxy/models.go
Normal file
@@ -0,0 +1,58 @@
|
||||
package controller
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"opencatd-open/internal/dto"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
func (p *Proxy) HandleModels(c *gin.Context) {
|
||||
models, err := p.getCache()
|
||||
if err != nil {
|
||||
dto.Fail(c, http.StatusBadGateway, err.Error())
|
||||
return
|
||||
}
|
||||
type _model struct {
|
||||
ID string `json:"id"`
|
||||
}
|
||||
var ms []_model
|
||||
for _, model := range models {
|
||||
ms = append(ms, _model{ID: model})
|
||||
}
|
||||
dto.Success(c, ms)
|
||||
}
|
||||
|
||||
func (p *Proxy) setCache() error {
|
||||
apikeys, err := p.apiKeyDao.FindKeys(nil)
|
||||
models := make(map[string]bool)
|
||||
if err == nil && len(apikeys) > 0 {
|
||||
for _, k := range apikeys {
|
||||
if len(k.SupportModelsArray) > 0 {
|
||||
for _, sm := range k.SupportModelsArray {
|
||||
models[sm] = true
|
||||
}
|
||||
} else {
|
||||
var sma []string
|
||||
json.Unmarshal([]byte(*k.SupportModels), &sma) // nolint:errCheck
|
||||
for _, sm := range sma {
|
||||
models[sm] = true
|
||||
}
|
||||
}
|
||||
}
|
||||
} else {
|
||||
return fmt.Errorf("empty data")
|
||||
}
|
||||
var support_models []string
|
||||
for m, _ := range models {
|
||||
support_models = append(support_models, m)
|
||||
}
|
||||
return p.cache.Set("models", support_models)
|
||||
}
|
||||
|
||||
func (p *Proxy) getCache() ([]string, error) {
|
||||
models, err := p.cache.Get("models")
|
||||
return models.([]string), err
|
||||
}
|
||||
@@ -19,6 +19,7 @@ import (
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/bluele/gcache"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/lib/pq"
|
||||
"github.com/tidwall/gjson"
|
||||
@@ -34,6 +35,7 @@ type Proxy struct {
|
||||
usageChan chan *model.Usage // 用于异步处理的channel
|
||||
apikey *model.ApiKey
|
||||
httpClient *http.Client
|
||||
cache gcache.Cache
|
||||
|
||||
userDAO *dao.UserDAO
|
||||
apiKeyDao *dao.ApiKeyDAO
|
||||
@@ -53,12 +55,14 @@ func NewProxy(ctx context.Context, cfg *config.Config, db *gorm.DB, wg *sync.Wai
|
||||
client.Transport = tr
|
||||
}
|
||||
}
|
||||
|
||||
np := &Proxy{
|
||||
ctx: ctx,
|
||||
cfg: cfg,
|
||||
db: db,
|
||||
wg: wg,
|
||||
httpClient: client,
|
||||
cache: gcache.New(1).Build(),
|
||||
usageChan: make(chan *model.Usage, cfg.UsageChanSize),
|
||||
userDAO: userDAO,
|
||||
apiKeyDao: apiKeyDAO,
|
||||
|
||||
@@ -124,7 +124,7 @@ func SetRouter(cfg *config.Config, db *gorm.DB, web *embed.FS) {
|
||||
{
|
||||
// v1.POST("/v2/*proxypath", router.HandleProxy)
|
||||
v1.POST("/*proxypath", proxy.HandleProxy)
|
||||
// v1.GET("/models", dashboard.HandleModels)
|
||||
v1.GET("/models", proxy.HandleModels)
|
||||
}
|
||||
|
||||
idxFS, err := fs.Sub(web, "dist")
|
||||
|
||||
Reference in New Issue
Block a user