Вы не можете выбрать более 25 тем Темы должны начинаться с буквы или цифры, могут содержать дефисы(-) и должны содержать не более 35 символов.
 
 
 

469 строки
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. // 0 排到最后
  170. if oi == 0 && oj == 0 {
  171. return vendorsList[i].ID < vendorsList[j].ID
  172. }
  173. if oi == 0 {
  174. return false
  175. }
  176. if oj == 0 {
  177. return true
  178. }
  179. if oi != oj {
  180. return oi < oj
  181. }
  182. return vendorsList[i].ID < vendorsList[j].ID
  183. })
  184. modelGroupsMap := make(map[string]*types.Set[string])
  185. for _, ability := range enableAbilities {
  186. groups, ok := modelGroupsMap[ability.Model]
  187. if !ok {
  188. groups = types.NewSet[string]()
  189. modelGroupsMap[ability.Model] = groups
  190. }
  191. groups.Add(ability.Group)
  192. }
  193. //这里使用切片而不是Set,因为一个模型可能支持多个端点类型,并且第一个端点是优先使用端点
  194. modelSupportEndpointsStr := make(map[string][]string)
  195. // 先根据已有能力填充原生端点
  196. for _, ability := range enableAbilities {
  197. endpoints := modelSupportEndpointsStr[ability.Model]
  198. channelTypes := common.GetEndpointTypesByChannelType(ability.ChannelType, ability.Model)
  199. for _, channelType := range channelTypes {
  200. if !common.StringsContains(endpoints, string(channelType)) {
  201. endpoints = append(endpoints, string(channelType))
  202. }
  203. }
  204. modelSupportEndpointsStr[ability.Model] = endpoints
  205. }
  206. // 再补充模型自定义端点:若配置有效则替换默认端点,不做合并
  207. for modelName, meta := range metaMap {
  208. if strings.TrimSpace(meta.Endpoints) == "" {
  209. continue
  210. }
  211. var raw map[string]interface{}
  212. if err := json.Unmarshal([]byte(meta.Endpoints), &raw); err == nil {
  213. endpoints := make([]string, 0, len(raw))
  214. for k, v := range raw {
  215. switch v.(type) {
  216. case string, map[string]interface{}:
  217. if !common.StringsContains(endpoints, k) {
  218. endpoints = append(endpoints, k)
  219. }
  220. }
  221. }
  222. if len(endpoints) > 0 {
  223. modelSupportEndpointsStr[modelName] = endpoints
  224. }
  225. }
  226. }
  227. modelSupportEndpointTypes = make(map[string][]constant.EndpointType)
  228. for model, endpoints := range modelSupportEndpointsStr {
  229. supportedEndpoints := make([]constant.EndpointType, 0)
  230. for _, endpointStr := range endpoints {
  231. endpointType := constant.EndpointType(endpointStr)
  232. supportedEndpoints = append(supportedEndpoints, endpointType)
  233. }
  234. modelSupportEndpointTypes[model] = supportedEndpoints
  235. }
  236. // 构建全局 supportedEndpointMap(默认 + 自定义覆盖)
  237. supportedEndpointMap = make(map[string]common.EndpointInfo)
  238. // 1. 默认端点
  239. for _, endpoints := range modelSupportEndpointTypes {
  240. for _, et := range endpoints {
  241. if info, ok := common.GetDefaultEndpointInfo(et); ok {
  242. if _, exists := supportedEndpointMap[string(et)]; !exists {
  243. supportedEndpointMap[string(et)] = info
  244. }
  245. }
  246. }
  247. }
  248. // 2. 自定义端点(models 表)覆盖默认
  249. for _, meta := range metaMap {
  250. if strings.TrimSpace(meta.Endpoints) == "" {
  251. continue
  252. }
  253. var raw map[string]interface{}
  254. if err := json.Unmarshal([]byte(meta.Endpoints), &raw); err == nil {
  255. for k, v := range raw {
  256. switch val := v.(type) {
  257. case string:
  258. supportedEndpointMap[k] = common.EndpointInfo{Path: val, Method: "POST"}
  259. case map[string]interface{}:
  260. ep := common.EndpointInfo{Method: "POST"}
  261. if p, ok := val["path"].(string); ok {
  262. ep.Path = p
  263. }
  264. if m, ok := val["method"].(string); ok {
  265. ep.Method = strings.ToUpper(m)
  266. }
  267. supportedEndpointMap[k] = ep
  268. default:
  269. // ignore unsupported types
  270. }
  271. }
  272. }
  273. }
  274. // 从渠道定价表加载实际定价数据(仅启用渠道),同时获取渠道名称
  275. var allCPs []struct {
  276. ChannelPricing
  277. ChannelName string
  278. }
  279. DB.Table("channel_pricings").
  280. Select("channel_pricings.*, channels.name as channel_name").
  281. Joins("JOIN channels ON channel_pricings.channel_id = channels.id").
  282. Where("channels.status = 1 AND channel_pricings.deleted_at IS NULL").
  283. Find(&allCPs)
  284. cpMap := make(map[string][]ChannelPricing)
  285. channelNameMap := make(map[int]string)
  286. for i := range allCPs {
  287. cpMap[allCPs[i].ModelName] = append(cpMap[allCPs[i].ModelName], allCPs[i].ChannelPricing)
  288. channelNameMap[allCPs[i].ChannelId] = allCPs[i].ChannelName
  289. }
  290. pricingMap = make([]Pricing, 0)
  291. for model, groups := range modelGroupsMap {
  292. pricing := Pricing{
  293. ModelName: model,
  294. EnableGroup: groups.Items(),
  295. SupportedEndpointTypes: modelSupportEndpointTypes[model],
  296. }
  297. // 补充模型元数据(描述、标签、供应商、状态)
  298. if meta, ok := metaMap[model]; ok {
  299. // 若模型被禁用(status==0),则直接跳过,不返回给前端
  300. if meta.Status == 0 {
  301. continue
  302. }
  303. pricing.Description = meta.Description
  304. pricing.Icon = meta.Icon
  305. pricing.Tags = meta.Tags
  306. pricing.VendorID = meta.VendorID
  307. pricing.Type = meta.Type
  308. }
  309. // 使用渠道定价表中的实际数据,选取最便宜的渠道
  310. applyBestChannelPricing(&pricing, cpMap[model], model)
  311. // 填充默认通道名称
  312. if chId, ok := GetDefaultChannelId(model); ok {
  313. if name, found := channelNameMap[chId]; found {
  314. pricing.DefaultChannelName = name
  315. }
  316. }
  317. pricingMap = append(pricingMap, pricing)
  318. }
  319. // 按 sort_order 排序 pricingMap,0 排到最后
  320. sort.Slice(pricingMap, func(i, j int) bool {
  321. mi, okI := metaMap[pricingMap[i].ModelName]
  322. mj, okJ := metaMap[pricingMap[j].ModelName]
  323. if okI && okJ {
  324. si, sj := mi.SortOrder, mj.SortOrder
  325. if si == 0 && sj == 0 {
  326. return pricingMap[i].ModelName < pricingMap[j].ModelName
  327. }
  328. if si == 0 {
  329. return false
  330. }
  331. if sj == 0 {
  332. return true
  333. }
  334. if si != sj {
  335. return si < sj
  336. }
  337. } else if okI {
  338. return true
  339. } else if okJ {
  340. return false
  341. }
  342. return pricingMap[i].ModelName < pricingMap[j].ModelName
  343. })
  344. // 防止大更新后数据不通用
  345. if len(pricingMap) > 0 {
  346. pricingMap[0].PricingVersion = "82c4a357505fff6fee8462c3f7ec8a645bb95532669cb73b2cabee6a416ec24f"
  347. }
  348. // 刷新缓存映射,供高并发快速查询
  349. modelEnableGroupsLock.Lock()
  350. modelEnableGroups = make(map[string][]string)
  351. modelQuotaTypeMap = make(map[string]int)
  352. for _, p := range pricingMap {
  353. modelEnableGroups[p.ModelName] = p.EnableGroup
  354. modelQuotaTypeMap[p.ModelName] = p.QuotaType
  355. }
  356. modelEnableGroupsLock.Unlock()
  357. lastGetPricingTime = time.Now()
  358. }
  359. // GetSupportedEndpointMap 返回全局端点到路径的映射
  360. func GetSupportedEndpointMap() map[string]common.EndpointInfo {
  361. return supportedEndpointMap
  362. }
  363. // applyGlobalDefault 用全局默认值填充 Pricing(无渠道定价时的回退)
  364. func applyGlobalDefault(pricing *Pricing, model string) {
  365. modelPrice, findPrice := ratio_setting.GetModelPrice(model, false)
  366. if findPrice {
  367. pricing.ModelPrice = modelPrice
  368. pricing.QuotaType = 1
  369. } else {
  370. modelRatio, _, _ := ratio_setting.GetModelRatio(model)
  371. pricing.ModelRatio = modelRatio
  372. pricing.CompletionRatio = ratio_setting.GetCompletionRatio(model)
  373. pricing.QuotaType = 0
  374. }
  375. pricing.CacheRatio, _ = ratio_setting.GetCacheRatio(model)
  376. pricing.CacheCreationRatio, _ = ratio_setting.GetCreateCacheRatio(model)
  377. }
  378. // applyBestChannelPricing 从渠道定价中选取最优(最便宜)的价格填充 Pricing
  379. // 优先选按量计费(quota_type=0)的渠道,因为缓存价格仅对按量计费有意义
  380. // 扩展比率字段(cache_ratio 等)为 0 表示未设置,需回退到全局默认值
  381. func applyBestChannelPricing(pricing *Pricing, cps []ChannelPricing, model string) {
  382. if len(cps) == 0 {
  383. applyGlobalDefault(pricing, model)
  384. return
  385. }
  386. // 优先选按量计费 (quota_type=0) 中 model_ratio 最低的渠道
  387. var bestPerToken *ChannelPricing
  388. for i := range cps {
  389. cp := &cps[i]
  390. if cp.QuotaType == 0 {
  391. if bestPerToken == nil || cp.ModelRatio < bestPerToken.ModelRatio {
  392. bestPerToken = cp
  393. }
  394. }
  395. }
  396. if bestPerToken != nil {
  397. pricing.QuotaType = 0
  398. pricing.ModelRatio = bestPerToken.ModelRatio
  399. pricing.CompletionRatio = bestPerToken.CompletionRatio
  400. if bestPerToken.CacheRatio > 0 {
  401. pricing.CacheRatio = bestPerToken.CacheRatio
  402. } else {
  403. pricing.CacheRatio, _ = ratio_setting.GetCacheRatio(model)
  404. }
  405. if bestPerToken.CacheCreationRatio > 0 {
  406. pricing.CacheCreationRatio = bestPerToken.CacheCreationRatio
  407. } else {
  408. pricing.CacheCreationRatio, _ = ratio_setting.GetCreateCacheRatio(model)
  409. }
  410. return
  411. }
  412. // 没有按量渠道,选按次计费 (quota_type=1) 中 model_price 最低的渠道
  413. var bestPerCall *ChannelPricing
  414. for i := range cps {
  415. cp := &cps[i]
  416. if cp.QuotaType == 1 {
  417. if bestPerCall == nil || cp.ModelPrice < bestPerCall.ModelPrice {
  418. bestPerCall = cp
  419. }
  420. }
  421. }
  422. if bestPerCall != nil {
  423. pricing.QuotaType = 1
  424. pricing.ModelPrice = bestPerCall.ModelPrice
  425. return
  426. }
  427. applyGlobalDefault(pricing, model)
  428. }