You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
 
 
 

407 rivejä
12 KiB

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