127 lines
2.8 KiB
Go
127 lines
2.8 KiB
Go
package probe
|
|
|
|
import (
|
|
"strings"
|
|
"unicode"
|
|
)
|
|
|
|
type AliasResult struct {
|
|
Raw string
|
|
Normalized string
|
|
Canonical string
|
|
}
|
|
|
|
func NormalizeModelID(raw string) string {
|
|
trimmed := strings.TrimSpace(strings.ToLower(raw))
|
|
if trimmed == "" {
|
|
return ""
|
|
}
|
|
|
|
if idx := strings.LastIndex(trimmed, "/"); idx >= 0 {
|
|
trimmed = trimmed[idx+1:]
|
|
}
|
|
|
|
replacer := strings.NewReplacer("_", "-", " ", "-", ".", ".", "--", "-")
|
|
normalized := replacer.Replace(trimmed)
|
|
for strings.Contains(normalized, "--") {
|
|
normalized = strings.ReplaceAll(normalized, "--", "-")
|
|
}
|
|
return strings.Trim(normalized, "-")
|
|
}
|
|
|
|
func CanonicalModelID(raw string) string {
|
|
return NormalizeModelID(raw)
|
|
}
|
|
|
|
func CanonicalModelFamily(raw string) string {
|
|
normalized := NormalizeModelID(raw)
|
|
switch {
|
|
case strings.HasPrefix(normalized, "kimi-k2."):
|
|
return strings.Replace(normalized, "kimi-k2.", "kimi-2.", 1)
|
|
case strings.HasPrefix(normalized, "kimi-k2-"):
|
|
return strings.Replace(normalized, "kimi-k2-", "kimi-2-", 1)
|
|
default:
|
|
return normalized
|
|
}
|
|
}
|
|
|
|
func BuildAliasTable(rawModels []string) map[string]AliasResult {
|
|
table := make(map[string]AliasResult, len(rawModels)*4)
|
|
for _, rawModel := range rawModels {
|
|
rawModel = strings.TrimSpace(rawModel)
|
|
if rawModel == "" {
|
|
continue
|
|
}
|
|
|
|
result := AliasResult{
|
|
Raw: rawModel,
|
|
Normalized: NormalizeModelID(rawModel),
|
|
Canonical: CanonicalModelFamily(rawModel),
|
|
}
|
|
|
|
keys := []string{
|
|
rawModel,
|
|
result.Normalized,
|
|
result.Canonical,
|
|
CanonicalModelID(rawModel),
|
|
lookupKey(rawModel),
|
|
lookupKey(result.Normalized),
|
|
lookupKey(result.Canonical),
|
|
lookupKey(CanonicalModelID(rawModel)),
|
|
}
|
|
for _, key := range keys {
|
|
if key == "" {
|
|
continue
|
|
}
|
|
if _, exists := table[key]; !exists {
|
|
table[key] = result
|
|
}
|
|
}
|
|
}
|
|
return table
|
|
}
|
|
|
|
func ResolveRequestedModel(requested string, rawModels []string) (resolved string, ok bool) {
|
|
result, ok := BuildAliasTable(rawModels)[lookupKey(requested)]
|
|
if !ok {
|
|
return "", false
|
|
}
|
|
return result.Raw, true
|
|
}
|
|
|
|
func RecommendModels(requested []string, rawModels []string) []string {
|
|
table := BuildAliasTable(rawModels)
|
|
recommended := make([]string, 0, len(requested))
|
|
seen := make(map[string]struct{}, len(requested))
|
|
|
|
for _, requestedModel := range requested {
|
|
result, ok := table[lookupKey(requestedModel)]
|
|
if !ok {
|
|
continue
|
|
}
|
|
if _, exists := seen[result.Raw]; exists {
|
|
continue
|
|
}
|
|
seen[result.Raw] = struct{}{}
|
|
recommended = append(recommended, result.Raw)
|
|
}
|
|
|
|
return recommended
|
|
}
|
|
|
|
func lookupKey(raw string) string {
|
|
canonical := CanonicalModelFamily(raw)
|
|
if canonical == "" {
|
|
return ""
|
|
}
|
|
|
|
var builder strings.Builder
|
|
builder.Grow(len(canonical))
|
|
for _, r := range canonical {
|
|
if unicode.IsLetter(r) || unicode.IsDigit(r) {
|
|
builder.WriteRune(unicode.ToLower(r))
|
|
}
|
|
}
|
|
return builder.String()
|
|
}
|