119 lines
2.5 KiB
Go
119 lines
2.5 KiB
Go
package chat
|
|
|
|
import (
|
|
"context"
|
|
"sync"
|
|
"time"
|
|
|
|
"cursor-api-proxy/internal/svc"
|
|
apitypes "cursor-api-proxy/internal/types"
|
|
"cursor-api-proxy/pkg/domain/types"
|
|
|
|
"github.com/zeromicro/go-zero/core/logx"
|
|
)
|
|
|
|
const modelCacheTTLMs = 5 * 60 * 1000
|
|
|
|
type ModelCache struct {
|
|
At int64
|
|
Models []types.CursorCliModel
|
|
}
|
|
|
|
type ModelCacheRef struct {
|
|
mu sync.Mutex
|
|
cache *ModelCache
|
|
inflight bool
|
|
waiters []chan struct{}
|
|
}
|
|
|
|
var globalModelCache = &ModelCacheRef{}
|
|
|
|
type ModelsLogic struct {
|
|
logx.Logger
|
|
ctx context.Context
|
|
svcCtx *svc.ServiceContext
|
|
}
|
|
|
|
func NewModelsLogic(ctx context.Context, svcCtx *svc.ServiceContext) *ModelsLogic {
|
|
return &ModelsLogic{
|
|
Logger: logx.WithContext(ctx),
|
|
ctx: ctx,
|
|
svcCtx: svcCtx,
|
|
}
|
|
}
|
|
|
|
func (l *ModelsLogic) Models() (resp *apitypes.ModelsResponse, err error) {
|
|
now := time.Now().UnixMilli()
|
|
|
|
globalModelCache.mu.Lock()
|
|
if globalModelCache.cache != nil && now-globalModelCache.cache.At <= modelCacheTTLMs {
|
|
cache := globalModelCache.cache
|
|
globalModelCache.mu.Unlock()
|
|
return buildModelsResponse(cache.Models), nil
|
|
}
|
|
|
|
if globalModelCache.inflight {
|
|
ch := make(chan struct{}, 1)
|
|
globalModelCache.waiters = append(globalModelCache.waiters, ch)
|
|
globalModelCache.mu.Unlock()
|
|
<-ch
|
|
globalModelCache.mu.Lock()
|
|
cache := globalModelCache.cache
|
|
globalModelCache.mu.Unlock()
|
|
return buildModelsResponse(cache.Models), nil
|
|
}
|
|
|
|
globalModelCache.inflight = true
|
|
globalModelCache.mu.Unlock()
|
|
|
|
fetched, err := types.ListCursorCliModels(l.svcCtx.Config.AgentBin, l.svcCtx.Config.TimeoutMs)
|
|
|
|
globalModelCache.mu.Lock()
|
|
globalModelCache.inflight = false
|
|
if err == nil {
|
|
globalModelCache.cache = &ModelCache{At: time.Now().UnixMilli(), Models: fetched}
|
|
}
|
|
waiters := globalModelCache.waiters
|
|
globalModelCache.waiters = nil
|
|
globalModelCache.mu.Unlock()
|
|
|
|
for _, ch := range waiters {
|
|
ch <- struct{}{}
|
|
}
|
|
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
return buildModelsResponse(fetched), nil
|
|
}
|
|
|
|
func buildModelsResponse(mods []types.CursorCliModel) *apitypes.ModelsResponse {
|
|
models := make([]apitypes.ModelData, len(mods))
|
|
for i, m := range mods {
|
|
models[i] = apitypes.ModelData{
|
|
Id: m.ID,
|
|
Object: "model",
|
|
OwnedBy: "cursor",
|
|
}
|
|
}
|
|
|
|
ids := make([]string, len(mods))
|
|
for i, m := range mods {
|
|
ids[i] = m.ID
|
|
}
|
|
aliases := types.GetAnthropicModelAliases(ids)
|
|
for _, a := range aliases {
|
|
models = append(models, apitypes.ModelData{
|
|
Id: a.ID,
|
|
Object: "model",
|
|
OwnedBy: "cursor",
|
|
})
|
|
}
|
|
|
|
return &apitypes.ModelsResponse{
|
|
Object: "list",
|
|
Data: models,
|
|
}
|
|
}
|