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

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