Du kannst nicht mehr als 25 Themen auswählen Themen müssen entweder mit einem Buchstaben oder einer Ziffer beginnen. Sie können Bindestriche („-“) enthalten und bis zu 35 Zeichen lang sein.
 
 
 

451 Zeilen
14 KiB

  1. package model
  2. import (
  3. "encoding/json"
  4. "fmt"
  5. "sort"
  6. "strings"
  7. "sync"
  8. "time"
  9. "github.com/QuantumNous/new-api/common"
  10. "github.com/QuantumNous/new-api/constant"
  11. "github.com/QuantumNous/new-api/setting/ratio_setting"
  12. "github.com/QuantumNous/new-api/types"
  13. )
  14. type Pricing struct {
  15. ModelName string `json:"model_name"`
  16. Description string `json:"description,omitempty"`
  17. Icon string `json:"icon,omitempty"`
  18. Tags string `json:"tags,omitempty"`
  19. VendorID int `json:"vendor_id,omitempty"`
  20. QuotaType int `json:"quota_type"`
  21. ModelRatio float64 `json:"model_ratio"`
  22. ModelPrice float64 `json:"model_price"`
  23. OwnerBy string `json:"owner_by"`
  24. CompletionRatio float64 `json:"completion_ratio"`
  25. CacheRatio float64 `json:"cache_ratio"`
  26. CacheCreationRatio float64 `json:"cache_creation_ratio"`
  27. EnableGroup []string `json:"enable_groups"`
  28. SupportedEndpointTypes []constant.EndpointType `json:"supported_endpoint_types"`
  29. PricingVersion string `json:"pricing_version,omitempty"`
  30. Type int `json:"type"`
  31. DefaultChannelName string `json:"default_channel_name,omitempty"`
  32. }
  33. type PricingVendor struct {
  34. ID int `json:"id"`
  35. Name string `json:"name"`
  36. Description string `json:"description,omitempty"`
  37. Icon string `json:"icon,omitempty"`
  38. }
  39. var (
  40. pricingMap []Pricing
  41. vendorsList []PricingVendor
  42. supportedEndpointMap map[string]common.EndpointInfo
  43. lastGetPricingTime time.Time
  44. updatePricingLock sync.Mutex
  45. // 缓存映射:模型名 -> 启用分组 / 计费类型
  46. modelEnableGroups = make(map[string][]string)
  47. modelQuotaTypeMap = make(map[string]int)
  48. modelEnableGroupsLock = sync.RWMutex{}
  49. )
  50. var (
  51. modelSupportEndpointTypes = make(map[string][]constant.EndpointType)
  52. modelSupportEndpointsLock = sync.RWMutex{}
  53. )
  54. func GetPricing() []Pricing {
  55. if time.Since(lastGetPricingTime) > time.Minute*1 || len(pricingMap) == 0 {
  56. updatePricingLock.Lock()
  57. defer updatePricingLock.Unlock()
  58. // Double check after acquiring the lock
  59. if time.Since(lastGetPricingTime) > time.Minute*1 || len(pricingMap) == 0 {
  60. modelSupportEndpointsLock.Lock()
  61. defer modelSupportEndpointsLock.Unlock()
  62. updatePricing()
  63. }
  64. }
  65. return pricingMap
  66. }
  67. // GetVendors 返回当前定价接口使用到的供应商信息
  68. func GetVendors() []PricingVendor {
  69. if time.Since(lastGetPricingTime) > time.Minute*1 || len(pricingMap) == 0 {
  70. // 保证先刷新一次
  71. GetPricing()
  72. }
  73. return vendorsList
  74. }
  75. func GetModelSupportEndpointTypes(model string) []constant.EndpointType {
  76. if model == "" {
  77. return make([]constant.EndpointType, 0)
  78. }
  79. modelSupportEndpointsLock.RLock()
  80. defer modelSupportEndpointsLock.RUnlock()
  81. if endpoints, ok := modelSupportEndpointTypes[model]; ok {
  82. return endpoints
  83. }
  84. return make([]constant.EndpointType, 0)
  85. }
  86. func updatePricing() {
  87. //modelRatios := common.GetModelRatios()
  88. enableAbilities, err := GetAllEnableAbilityWithChannels()
  89. if err != nil {
  90. common.SysLog(fmt.Sprintf("GetAllEnableAbilityWithChannels error: %v", err))
  91. return
  92. }
  93. // 预加载模型元数据与供应商一次,避免循环查询
  94. var allMeta []Model
  95. _ = DB.Find(&allMeta).Error
  96. metaMap := make(map[string]*Model)
  97. prefixList := make([]*Model, 0)
  98. suffixList := make([]*Model, 0)
  99. containsList := make([]*Model, 0)
  100. for i := range allMeta {
  101. m := &allMeta[i]
  102. if m.NameRule == NameRuleExact {
  103. metaMap[m.ModelName] = m
  104. } else {
  105. switch m.NameRule {
  106. case NameRulePrefix:
  107. prefixList = append(prefixList, m)
  108. case NameRuleSuffix:
  109. suffixList = append(suffixList, m)
  110. case NameRuleContains:
  111. containsList = append(containsList, m)
  112. }
  113. }
  114. }
  115. // 将非精确规则模型匹配到 metaMap
  116. for _, m := range prefixList {
  117. for _, pricingModel := range enableAbilities {
  118. if strings.HasPrefix(pricingModel.Model, m.ModelName) {
  119. if _, exists := metaMap[pricingModel.Model]; !exists {
  120. metaMap[pricingModel.Model] = m
  121. }
  122. }
  123. }
  124. }
  125. for _, m := range suffixList {
  126. for _, pricingModel := range enableAbilities {
  127. if strings.HasSuffix(pricingModel.Model, m.ModelName) {
  128. if _, exists := metaMap[pricingModel.Model]; !exists {
  129. metaMap[pricingModel.Model] = m
  130. }
  131. }
  132. }
  133. }
  134. for _, m := range containsList {
  135. for _, pricingModel := range enableAbilities {
  136. if strings.Contains(pricingModel.Model, m.ModelName) {
  137. if _, exists := metaMap[pricingModel.Model]; !exists {
  138. metaMap[pricingModel.Model] = m
  139. }
  140. }
  141. }
  142. }
  143. // 预加载供应商
  144. var vendors []Vendor
  145. _ = DB.Find(&vendors).Error
  146. vendorMap := make(map[int]*Vendor)
  147. for i := range vendors {
  148. vendorMap[vendors[i].Id] = &vendors[i]
  149. }
  150. // 初始化默认供应商映射
  151. initDefaultVendorMapping(metaMap, vendorMap, enableAbilities)
  152. // 构建对前端友好的供应商列表
  153. vendorOrderMap := make(map[int]int)
  154. for _, v := range vendorMap {
  155. vendorOrderMap[v.Id] = v.SortOrder
  156. }
  157. vendorsList = make([]PricingVendor, 0, len(vendorMap))
  158. for _, v := range vendorMap {
  159. vendorsList = append(vendorsList, PricingVendor{
  160. ID: v.Id,
  161. Name: v.Name,
  162. Description: v.Description,
  163. Icon: v.Icon,
  164. })
  165. }
  166. sort.Slice(vendorsList, func(i, j int) bool {
  167. oi := vendorOrderMap[vendorsList[i].ID]
  168. oj := vendorOrderMap[vendorsList[j].ID]
  169. if oi != oj {
  170. return oi < oj
  171. }
  172. return vendorsList[i].ID < vendorsList[j].ID
  173. })
  174. modelGroupsMap := make(map[string]*types.Set[string])
  175. for _, ability := range enableAbilities {
  176. groups, ok := modelGroupsMap[ability.Model]
  177. if !ok {
  178. groups = types.NewSet[string]()
  179. modelGroupsMap[ability.Model] = groups
  180. }
  181. groups.Add(ability.Group)
  182. }
  183. //这里使用切片而不是Set,因为一个模型可能支持多个端点类型,并且第一个端点是优先使用端点
  184. modelSupportEndpointsStr := make(map[string][]string)
  185. // 先根据已有能力填充原生端点
  186. for _, ability := range enableAbilities {
  187. endpoints := modelSupportEndpointsStr[ability.Model]
  188. channelTypes := common.GetEndpointTypesByChannelType(ability.ChannelType, ability.Model)
  189. for _, channelType := range channelTypes {
  190. if !common.StringsContains(endpoints, string(channelType)) {
  191. endpoints = append(endpoints, string(channelType))
  192. }
  193. }
  194. modelSupportEndpointsStr[ability.Model] = endpoints
  195. }
  196. // 再补充模型自定义端点:若配置有效则替换默认端点,不做合并
  197. for modelName, meta := range metaMap {
  198. if strings.TrimSpace(meta.Endpoints) == "" {
  199. continue
  200. }
  201. var raw map[string]interface{}
  202. if err := json.Unmarshal([]byte(meta.Endpoints), &raw); err == nil {
  203. endpoints := make([]string, 0, len(raw))
  204. for k, v := range raw {
  205. switch v.(type) {
  206. case string, map[string]interface{}:
  207. if !common.StringsContains(endpoints, k) {
  208. endpoints = append(endpoints, k)
  209. }
  210. }
  211. }
  212. if len(endpoints) > 0 {
  213. modelSupportEndpointsStr[modelName] = endpoints
  214. }
  215. }
  216. }
  217. modelSupportEndpointTypes = make(map[string][]constant.EndpointType)
  218. for model, endpoints := range modelSupportEndpointsStr {
  219. supportedEndpoints := make([]constant.EndpointType, 0)
  220. for _, endpointStr := range endpoints {
  221. endpointType := constant.EndpointType(endpointStr)
  222. supportedEndpoints = append(supportedEndpoints, endpointType)
  223. }
  224. modelSupportEndpointTypes[model] = supportedEndpoints
  225. }
  226. // 构建全局 supportedEndpointMap(默认 + 自定义覆盖)
  227. supportedEndpointMap = make(map[string]common.EndpointInfo)
  228. // 1. 默认端点
  229. for _, endpoints := range modelSupportEndpointTypes {
  230. for _, et := range endpoints {
  231. if info, ok := common.GetDefaultEndpointInfo(et); ok {
  232. if _, exists := supportedEndpointMap[string(et)]; !exists {
  233. supportedEndpointMap[string(et)] = info
  234. }
  235. }
  236. }
  237. }
  238. // 2. 自定义端点(models 表)覆盖默认
  239. for _, meta := range metaMap {
  240. if strings.TrimSpace(meta.Endpoints) == "" {
  241. continue
  242. }
  243. var raw map[string]interface{}
  244. if err := json.Unmarshal([]byte(meta.Endpoints), &raw); err == nil {
  245. for k, v := range raw {
  246. switch val := v.(type) {
  247. case string:
  248. supportedEndpointMap[k] = common.EndpointInfo{Path: val, Method: "POST"}
  249. case map[string]interface{}:
  250. ep := common.EndpointInfo{Method: "POST"}
  251. if p, ok := val["path"].(string); ok {
  252. ep.Path = p
  253. }
  254. if m, ok := val["method"].(string); ok {
  255. ep.Method = strings.ToUpper(m)
  256. }
  257. supportedEndpointMap[k] = ep
  258. default:
  259. // ignore unsupported types
  260. }
  261. }
  262. }
  263. }
  264. // 从渠道定价表加载实际定价数据(仅启用渠道),同时获取渠道名称
  265. var allCPs []struct {
  266. ChannelPricing
  267. ChannelName string
  268. }
  269. DB.Table("channel_pricings").
  270. Select("channel_pricings.*, channels.name as channel_name").
  271. Joins("JOIN channels ON channel_pricings.channel_id = channels.id").
  272. Where("channels.status = 1 AND channel_pricings.deleted_at IS NULL").
  273. Find(&allCPs)
  274. cpMap := make(map[string][]ChannelPricing)
  275. channelNameMap := make(map[int]string)
  276. for i := range allCPs {
  277. cpMap[allCPs[i].ModelName] = append(cpMap[allCPs[i].ModelName], allCPs[i].ChannelPricing)
  278. channelNameMap[allCPs[i].ChannelId] = allCPs[i].ChannelName
  279. }
  280. pricingMap = make([]Pricing, 0)
  281. for model, groups := range modelGroupsMap {
  282. pricing := Pricing{
  283. ModelName: model,
  284. EnableGroup: groups.Items(),
  285. SupportedEndpointTypes: modelSupportEndpointTypes[model],
  286. }
  287. // 补充模型元数据(描述、标签、供应商、状态)
  288. if meta, ok := metaMap[model]; ok {
  289. // 若模型被禁用(status==0),则直接跳过,不返回给前端
  290. if meta.Status == 0 {
  291. continue
  292. }
  293. pricing.Description = meta.Description
  294. pricing.Icon = meta.Icon
  295. pricing.Tags = meta.Tags
  296. pricing.VendorID = meta.VendorID
  297. pricing.Type = meta.Type
  298. }
  299. // 使用渠道定价表中的实际数据,选取最便宜的渠道
  300. applyBestChannelPricing(&pricing, cpMap[model], model)
  301. // 填充默认通道名称
  302. if chId, ok := GetDefaultChannelId(model); ok {
  303. if name, found := channelNameMap[chId]; found {
  304. pricing.DefaultChannelName = name
  305. }
  306. }
  307. pricingMap = append(pricingMap, pricing)
  308. }
  309. // 按 sort_order 排序 pricingMap,999999 视为未设置
  310. sort.Slice(pricingMap, func(i, j int) bool {
  311. mi, okI := metaMap[pricingMap[i].ModelName]
  312. mj, okJ := metaMap[pricingMap[j].ModelName]
  313. si := 999999
  314. sj := 999999
  315. if okI {
  316. si = mi.SortOrder
  317. }
  318. if okJ {
  319. sj = mj.SortOrder
  320. }
  321. if si != sj {
  322. return si < sj
  323. }
  324. return pricingMap[i].ModelName < pricingMap[j].ModelName
  325. })
  326. // 防止大更新后数据不通用
  327. if len(pricingMap) > 0 {
  328. pricingMap[0].PricingVersion = "82c4a357505fff6fee8462c3f7ec8a645bb95532669cb73b2cabee6a416ec24f"
  329. }
  330. // 刷新缓存映射,供高并发快速查询
  331. modelEnableGroupsLock.Lock()
  332. modelEnableGroups = make(map[string][]string)
  333. modelQuotaTypeMap = make(map[string]int)
  334. for _, p := range pricingMap {
  335. modelEnableGroups[p.ModelName] = p.EnableGroup
  336. modelQuotaTypeMap[p.ModelName] = p.QuotaType
  337. }
  338. modelEnableGroupsLock.Unlock()
  339. lastGetPricingTime = time.Now()
  340. }
  341. // GetSupportedEndpointMap 返回全局端点到路径的映射
  342. func GetSupportedEndpointMap() map[string]common.EndpointInfo {
  343. return supportedEndpointMap
  344. }
  345. // applyGlobalDefault 用全局默认值填充 Pricing(无渠道定价时的回退)
  346. func applyGlobalDefault(pricing *Pricing, model string) {
  347. modelPrice, findPrice := ratio_setting.GetModelPrice(model, false)
  348. if findPrice {
  349. pricing.ModelPrice = modelPrice
  350. pricing.QuotaType = 1
  351. } else {
  352. modelRatio, _, _ := ratio_setting.GetModelRatio(model)
  353. pricing.ModelRatio = modelRatio
  354. pricing.CompletionRatio = ratio_setting.GetCompletionRatio(model)
  355. pricing.QuotaType = 0
  356. }
  357. pricing.CacheRatio, _ = ratio_setting.GetCacheRatio(model)
  358. pricing.CacheCreationRatio, _ = ratio_setting.GetCreateCacheRatio(model)
  359. }
  360. // applyBestChannelPricing 从渠道定价中选取最优(最便宜)的价格填充 Pricing
  361. // 优先选按量计费(quota_type=0)的渠道,因为缓存价格仅对按量计费有意义
  362. // 扩展比率字段(cache_ratio 等)为 0 表示未设置,需回退到全局默认值
  363. func applyBestChannelPricing(pricing *Pricing, cps []ChannelPricing, model string) {
  364. if len(cps) == 0 {
  365. applyGlobalDefault(pricing, model)
  366. return
  367. }
  368. // 优先选按量计费 (quota_type=0) 中 model_ratio 最低的渠道
  369. var bestPerToken *ChannelPricing
  370. for i := range cps {
  371. cp := &cps[i]
  372. if cp.QuotaType == 0 {
  373. if bestPerToken == nil || cp.ModelRatio < bestPerToken.ModelRatio {
  374. bestPerToken = cp
  375. }
  376. }
  377. }
  378. if bestPerToken != nil {
  379. pricing.QuotaType = 0
  380. pricing.ModelRatio = bestPerToken.ModelRatio
  381. pricing.CompletionRatio = bestPerToken.CompletionRatio
  382. if bestPerToken.CacheRatio > 0 {
  383. pricing.CacheRatio = bestPerToken.CacheRatio
  384. } else {
  385. pricing.CacheRatio, _ = ratio_setting.GetCacheRatio(model)
  386. }
  387. if bestPerToken.CacheCreationRatio > 0 {
  388. pricing.CacheCreationRatio = bestPerToken.CacheCreationRatio
  389. } else {
  390. pricing.CacheCreationRatio, _ = ratio_setting.GetCreateCacheRatio(model)
  391. }
  392. return
  393. }
  394. // 没有按量渠道,选按次计费 (quota_type=1) 中 model_price 最低的渠道
  395. var bestPerCall *ChannelPricing
  396. for i := range cps {
  397. cp := &cps[i]
  398. if cp.QuotaType == 1 {
  399. if bestPerCall == nil || cp.ModelPrice < bestPerCall.ModelPrice {
  400. bestPerCall = cp
  401. }
  402. }
  403. }
  404. if bestPerCall != nil {
  405. pricing.QuotaType = 1
  406. pricing.ModelPrice = bestPerCall.ModelPrice
  407. return
  408. }
  409. applyGlobalDefault(pricing, model)
  410. }