89 lines
2.1 KiB
Go
89 lines
2.1 KiB
Go
package static
|
|
|
|
import (
|
|
"context"
|
|
"strconv"
|
|
"strings"
|
|
|
|
"github.com/coze-dev/coze-studio/backend/infra/contract/modelmgr"
|
|
"github.com/coze-dev/coze-studio/backend/pkg/lang/ptr"
|
|
"github.com/coze-dev/coze-studio/backend/pkg/lang/sets"
|
|
)
|
|
|
|
func NewModelMgr(staticModels []*modelmgr.Model) (modelmgr.Manager, error) {
|
|
mapping := make(map[int64]*modelmgr.Model, len(staticModels))
|
|
for i := range staticModels {
|
|
mapping[staticModels[i].ID] = staticModels[i]
|
|
}
|
|
return &staticModelManager{
|
|
models: staticModels,
|
|
modelMapping: mapping,
|
|
}, nil
|
|
}
|
|
|
|
type staticModelManager struct {
|
|
models []*modelmgr.Model
|
|
modelMapping map[int64]*modelmgr.Model
|
|
}
|
|
|
|
func (s *staticModelManager) ListInUseModel(ctx context.Context, limit int, Cursor *string) (*modelmgr.ListModelResponse, error) {
|
|
return s.ListModel(ctx, &modelmgr.ListModelRequest{
|
|
Status: []modelmgr.ModelStatus{modelmgr.StatusInUse},
|
|
Limit: limit,
|
|
Cursor: Cursor,
|
|
})
|
|
}
|
|
|
|
func (s *staticModelManager) ListModel(_ context.Context, req *modelmgr.ListModelRequest) (*modelmgr.ListModelResponse, error) {
|
|
startIdx := 0
|
|
if req.Cursor != nil {
|
|
start, err := strconv.ParseInt(*req.Cursor, 10, 64)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
startIdx = int(start)
|
|
}
|
|
|
|
limit := req.Limit
|
|
if limit == 0 {
|
|
limit = 100
|
|
}
|
|
|
|
var (
|
|
i int
|
|
respList []*modelmgr.Model
|
|
statSet = sets.FromSlice(req.Status)
|
|
)
|
|
|
|
for i = startIdx; i < len(s.models) && len(respList) < limit; i++ {
|
|
m := s.models[i]
|
|
if req.FuzzyModelName != nil && !strings.Contains(m.Name, *req.FuzzyModelName) {
|
|
continue
|
|
}
|
|
if len(statSet) > 0 && !statSet.Contains(m.Meta.Status) {
|
|
continue
|
|
}
|
|
respList = append(respList, m)
|
|
}
|
|
|
|
resp := &modelmgr.ListModelResponse{
|
|
ModelList: respList,
|
|
}
|
|
resp.HasMore = i != len(s.models)
|
|
if resp.HasMore {
|
|
resp.NextCursor = ptr.Of(strconv.FormatInt(int64(i), 10))
|
|
}
|
|
|
|
return resp, nil
|
|
}
|
|
|
|
func (s *staticModelManager) MGetModelByID(_ context.Context, req *modelmgr.MGetModelRequest) ([]*modelmgr.Model, error) {
|
|
resp := make([]*modelmgr.Model, 0, len(s.models))
|
|
for _, id := range req.IDs {
|
|
if m, found := s.modelMapping[id]; found {
|
|
resp = append(resp, m)
|
|
}
|
|
}
|
|
return resp, nil
|
|
}
|