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

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