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"
|
"sync"
|
||||||
"time"
|
"time"
|
||||||
|
|
||||||
|
"github.com/bluele/gcache"
|
||||||
"github.com/gin-gonic/gin"
|
"github.com/gin-gonic/gin"
|
||||||
"github.com/lib/pq"
|
"github.com/lib/pq"
|
||||||
"github.com/tidwall/gjson"
|
"github.com/tidwall/gjson"
|
||||||
@@ -34,6 +35,7 @@ type Proxy struct {
|
|||||||
usageChan chan *model.Usage // 用于异步处理的channel
|
usageChan chan *model.Usage // 用于异步处理的channel
|
||||||
apikey *model.ApiKey
|
apikey *model.ApiKey
|
||||||
httpClient *http.Client
|
httpClient *http.Client
|
||||||
|
cache gcache.Cache
|
||||||
|
|
||||||
userDAO *dao.UserDAO
|
userDAO *dao.UserDAO
|
||||||
apiKeyDao *dao.ApiKeyDAO
|
apiKeyDao *dao.ApiKeyDAO
|
||||||
@@ -53,12 +55,14 @@ func NewProxy(ctx context.Context, cfg *config.Config, db *gorm.DB, wg *sync.Wai
|
|||||||
client.Transport = tr
|
client.Transport = tr
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
np := &Proxy{
|
np := &Proxy{
|
||||||
ctx: ctx,
|
ctx: ctx,
|
||||||
cfg: cfg,
|
cfg: cfg,
|
||||||
db: db,
|
db: db,
|
||||||
wg: wg,
|
wg: wg,
|
||||||
httpClient: client,
|
httpClient: client,
|
||||||
|
cache: gcache.New(1).Build(),
|
||||||
usageChan: make(chan *model.Usage, cfg.UsageChanSize),
|
usageChan: make(chan *model.Usage, cfg.UsageChanSize),
|
||||||
userDAO: userDAO,
|
userDAO: userDAO,
|
||||||
apiKeyDao: apiKeyDAO,
|
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("/v2/*proxypath", router.HandleProxy)
|
||||||
v1.POST("/*proxypath", proxy.HandleProxy)
|
v1.POST("/*proxypath", proxy.HandleProxy)
|
||||||
// v1.GET("/models", dashboard.HandleModels)
|
v1.GET("/models", proxy.HandleModels)
|
||||||
}
|
}
|
||||||
|
|
||||||
idxFS, err := fs.Sub(web, "dist")
|
idxFS, err := fs.Sub(web, "dist")
|
||||||
|
|||||||
Reference in New Issue
Block a user