mirror of
https://github.com/sipeed/picoclaw.git
synced 2026-06-12 18:08:54 +00:00
feat(model): llm rate limiting (#2198)
* feat(model): rate limiting * fix(agent): preserve per-model identity in rate limiting and fallback * fix test
This commit is contained in:
+116
-43
@@ -8,44 +8,102 @@ import (
|
||||
"github.com/sipeed/picoclaw/pkg/providers"
|
||||
)
|
||||
|
||||
func buildModelListResolver(cfg *config.Config) func(raw string) (string, bool) {
|
||||
ensureProtocol := func(model string) string {
|
||||
model = strings.TrimSpace(model)
|
||||
if model == "" {
|
||||
return ""
|
||||
}
|
||||
if strings.Contains(model, "/") {
|
||||
return model
|
||||
}
|
||||
return "openai/" + model
|
||||
func ensureProtocolModel(model string) string {
|
||||
model = strings.TrimSpace(model)
|
||||
if model == "" {
|
||||
return ""
|
||||
}
|
||||
if strings.Contains(model, "/") {
|
||||
return model
|
||||
}
|
||||
return "openai/" + model
|
||||
}
|
||||
|
||||
func modelConfigIdentityKey(mc *config.ModelConfig) string {
|
||||
if mc == nil {
|
||||
return ""
|
||||
}
|
||||
if name := strings.TrimSpace(mc.ModelName); name != "" {
|
||||
return "model_name:" + name
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
func candidateFromModelConfig(
|
||||
defaultProvider string,
|
||||
mc *config.ModelConfig,
|
||||
) (providers.FallbackCandidate, bool) {
|
||||
if mc == nil {
|
||||
return providers.FallbackCandidate{}, false
|
||||
}
|
||||
|
||||
return func(raw string) (string, bool) {
|
||||
raw = strings.TrimSpace(raw)
|
||||
if raw == "" || cfg == nil {
|
||||
return "", false
|
||||
}
|
||||
|
||||
if mc, err := cfg.GetModelConfig(raw); err == nil && mc != nil && strings.TrimSpace(mc.Model) != "" {
|
||||
return ensureProtocol(mc.Model), true
|
||||
}
|
||||
|
||||
for i := range cfg.ModelList {
|
||||
fullModel := strings.TrimSpace(cfg.ModelList[i].Model)
|
||||
if fullModel == "" {
|
||||
continue
|
||||
}
|
||||
if fullModel == raw {
|
||||
return ensureProtocol(fullModel), true
|
||||
}
|
||||
_, modelID := providers.ExtractProtocol(fullModel)
|
||||
if modelID == raw {
|
||||
return ensureProtocol(fullModel), true
|
||||
}
|
||||
}
|
||||
|
||||
return "", false
|
||||
ref := providers.ParseModelRef(ensureProtocolModel(mc.Model), defaultProvider)
|
||||
if ref == nil {
|
||||
return providers.FallbackCandidate{}, false
|
||||
}
|
||||
|
||||
return providers.FallbackCandidate{
|
||||
Provider: ref.Provider,
|
||||
Model: ref.Model,
|
||||
RPM: mc.RPM,
|
||||
IdentityKey: modelConfigIdentityKey(mc),
|
||||
}, true
|
||||
}
|
||||
|
||||
func lookupModelConfigByRef(cfg *config.Config, raw string) *config.ModelConfig {
|
||||
raw = strings.TrimSpace(raw)
|
||||
if raw == "" || cfg == nil {
|
||||
return nil
|
||||
}
|
||||
|
||||
if mc, err := cfg.GetModelConfig(raw); err == nil && mc != nil && strings.TrimSpace(mc.Model) != "" {
|
||||
return mc
|
||||
}
|
||||
|
||||
for i := range cfg.ModelList {
|
||||
mc := cfg.ModelList[i]
|
||||
if mc == nil {
|
||||
continue
|
||||
}
|
||||
fullModel := strings.TrimSpace(mc.Model)
|
||||
if fullModel == "" {
|
||||
continue
|
||||
}
|
||||
if fullModel == raw {
|
||||
return mc
|
||||
}
|
||||
_, modelID := providers.ExtractProtocol(fullModel)
|
||||
if modelID == raw {
|
||||
return mc
|
||||
}
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func resolveModelCandidate(
|
||||
cfg *config.Config,
|
||||
defaultProvider string,
|
||||
raw string,
|
||||
) (providers.FallbackCandidate, bool) {
|
||||
raw = strings.TrimSpace(raw)
|
||||
if raw == "" {
|
||||
return providers.FallbackCandidate{}, false
|
||||
}
|
||||
|
||||
if mc := lookupModelConfigByRef(cfg, raw); mc != nil {
|
||||
return candidateFromModelConfig(defaultProvider, mc)
|
||||
}
|
||||
|
||||
ref := providers.ParseModelRef(raw, defaultProvider)
|
||||
if ref == nil {
|
||||
return providers.FallbackCandidate{}, false
|
||||
}
|
||||
|
||||
return providers.FallbackCandidate{
|
||||
Provider: ref.Provider,
|
||||
Model: ref.Model,
|
||||
}, true
|
||||
}
|
||||
|
||||
func resolveModelCandidates(
|
||||
@@ -54,14 +112,29 @@ func resolveModelCandidates(
|
||||
primary string,
|
||||
fallbacks []string,
|
||||
) []providers.FallbackCandidate {
|
||||
return providers.ResolveCandidatesWithLookup(
|
||||
providers.ModelConfig{
|
||||
Primary: primary,
|
||||
Fallbacks: fallbacks,
|
||||
},
|
||||
defaultProvider,
|
||||
buildModelListResolver(cfg),
|
||||
)
|
||||
seen := make(map[string]bool)
|
||||
candidates := make([]providers.FallbackCandidate, 0, 1+len(fallbacks))
|
||||
|
||||
addCandidate := func(raw string) {
|
||||
candidate, ok := resolveModelCandidate(cfg, defaultProvider, raw)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
|
||||
key := candidate.StableKey()
|
||||
if seen[key] {
|
||||
return
|
||||
}
|
||||
seen[key] = true
|
||||
candidates = append(candidates, candidate)
|
||||
}
|
||||
|
||||
addCandidate(primary)
|
||||
for _, fallback := range fallbacks {
|
||||
addCandidate(fallback)
|
||||
}
|
||||
|
||||
return candidates
|
||||
}
|
||||
|
||||
func resolvedCandidateModel(candidates []providers.FallbackCandidate, fallback string) string {
|
||||
|
||||
Reference in New Issue
Block a user