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

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