152 lines
4.0 KiB
Go
152 lines
4.0 KiB
Go
package model
|
|
|
|
import (
|
|
"fmt"
|
|
"strings"
|
|
|
|
"github.com/QuantumNous/new-api/common"
|
|
)
|
|
|
|
// 简化的供应商映射规则
|
|
var defaultVendorRules = map[string]string{
|
|
"gpt": "OpenAI",
|
|
"dall-e": "OpenAI",
|
|
"whisper": "OpenAI",
|
|
"o1": "OpenAI",
|
|
"o3": "OpenAI",
|
|
"claude": "Anthropic",
|
|
"gemini": "Google",
|
|
"moonshot": "Moonshot",
|
|
"kimi": "Moonshot",
|
|
"chatglm": "智谱",
|
|
"glm-": "智谱",
|
|
"qwen": "阿里巴巴",
|
|
"deepseek": "DeepSeek",
|
|
"abab": "MiniMax",
|
|
"minimax": "MiniMax",
|
|
"ernie": "百度",
|
|
"spark": "讯飞",
|
|
"hunyuan": "腾讯",
|
|
"command": "Cohere",
|
|
"@cf/": "Cloudflare",
|
|
"360": "360",
|
|
"yi": "零一万物",
|
|
"jina": "Jina",
|
|
"mistral": "Mistral",
|
|
"grok": "xAI",
|
|
"llama": "Meta",
|
|
"doubao": "字节跳动",
|
|
"seedance": "字节跳动",
|
|
"kling": "快手",
|
|
"jimeng": "即梦",
|
|
"vidu": "Vidu",
|
|
}
|
|
|
|
// 供应商默认图标映射
|
|
var defaultVendorIcons = map[string]string{
|
|
"OpenAI": "OpenAI",
|
|
"Anthropic": "Claude.Color",
|
|
"Google": "Gemini.Color",
|
|
"Moonshot": "Moonshot",
|
|
"智谱": "Zhipu.Color",
|
|
"阿里巴巴": "Qwen.Color",
|
|
"DeepSeek": "DeepSeek.Color",
|
|
"MiniMax": "Minimax.Color",
|
|
"百度": "Wenxin.Color",
|
|
"讯飞": "Spark.Color",
|
|
"腾讯": "Hunyuan.Color",
|
|
"Cohere": "Cohere.Color",
|
|
"Cloudflare": "Cloudflare.Color",
|
|
"360": "Ai360.Color",
|
|
"零一万物": "Yi.Color",
|
|
"Jina": "Jina",
|
|
"Mistral": "Mistral.Color",
|
|
"xAI": "XAI",
|
|
"Meta": "Ollama",
|
|
"字节跳动": "Doubao.Color",
|
|
"快手": "Kling.Color",
|
|
"即梦": "Jimeng.Color",
|
|
"Vidu": "Vidu",
|
|
"微软": "AzureAI",
|
|
"Microsoft": "AzureAI",
|
|
"Azure": "AzureAI",
|
|
}
|
|
|
|
// initDefaultVendorMapping 简化的默认供应商映射
|
|
func initDefaultVendorMapping(metaMap map[string]*Model, vendorMap map[int]*Vendor, enableAbilities []AbilityWithChannel) {
|
|
for _, ability := range enableAbilities {
|
|
modelName := ability.Model
|
|
if existing, exists := metaMap[modelName]; exists {
|
|
// 已有记录但 vendor_id 已赋值 → 跳过
|
|
if existing.VendorID != 0 {
|
|
continue
|
|
}
|
|
// vendor_id=0 → 继续尝试匹配（修复已有 DB 记录但缺供应商的场景）
|
|
}
|
|
|
|
// 匹配供应商
|
|
vendorID := 0
|
|
modelLower := strings.ToLower(modelName)
|
|
for pattern, vendorName := range defaultVendorRules {
|
|
if strings.Contains(modelLower, pattern) {
|
|
vendorID = getOrCreateVendor(vendorName, vendorMap)
|
|
break
|
|
}
|
|
}
|
|
|
|
if vendorID == 0 {
|
|
// 仍未匹配到供应商 → 跳过（不给未匹配的模型写入 vendor_id=0 覆盖已有记录）
|
|
if strings.Contains(modelLower, "seedance") || strings.Contains(modelLower, "minimax") {
|
|
common.SysLog(fmt.Sprintf("initDefaultVendorMapping: model %q matched pattern but vendorID=0 (getOrCreateVendor failed)", modelName))
|
|
}
|
|
continue
|
|
}
|
|
|
|
// 写入或更新供应商 ID
|
|
if existing, exists := metaMap[modelName]; exists {
|
|
existing.VendorID = vendorID
|
|
} else {
|
|
metaMap[modelName] = &Model{
|
|
ModelName: modelName,
|
|
VendorID: vendorID,
|
|
Status: 1,
|
|
NameRule: NameRuleExact,
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
// 查找或创建供应商
|
|
func getOrCreateVendor(vendorName string, vendorMap map[int]*Vendor) int {
|
|
// 查找现有供应商
|
|
for id, vendor := range vendorMap {
|
|
if vendor.Name == vendorName {
|
|
return id
|
|
}
|
|
}
|
|
|
|
// 创建新供应商
|
|
newVendor := &Vendor{
|
|
Name: vendorName,
|
|
Status: 1,
|
|
Icon: getDefaultVendorIcon(vendorName),
|
|
}
|
|
|
|
if err := newVendor.Insert(); err != nil {
|
|
common.SysLog(fmt.Sprintf("getOrCreateVendor: insert vendor %q failed: %v", vendorName, err))
|
|
return 0
|
|
}
|
|
|
|
common.SysLog(fmt.Sprintf("getOrCreateVendor: created new vendor %q with id=%d", vendorName, newVendor.Id))
|
|
vendorMap[newVendor.Id] = newVendor
|
|
return newVendor.Id
|
|
}
|
|
|
|
// 获取供应商默认图标
|
|
func getDefaultVendorIcon(vendorName string) string {
|
|
if icon, exists := defaultVendorIcons[vendorName]; exists {
|
|
return icon
|
|
}
|
|
return ""
|
|
}
|