Nelze vybrat více než 25 témat Téma musí začínat písmenem nebo číslem, může obsahovat pomlčky („-“) a může být dlouhé až 35 znaků.
 
 
 

327 řádky
9.9 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. EnableGroup []string `json:"enable_groups"`
  25. SupportedEndpointTypes []constant.EndpointType `json:"supported_endpoint_types"`
  26. PricingVersion string `json:"pricing_version,omitempty"`
  27. Type int `json:"type"`
  28. }
  29. type PricingVendor struct {
  30. ID int `json:"id"`
  31. Name string `json:"name"`
  32. Description string `json:"description,omitempty"`
  33. Icon string `json:"icon,omitempty"`
  34. }
  35. var (
  36. pricingMap []Pricing
  37. vendorsList []PricingVendor
  38. supportedEndpointMap map[string]common.EndpointInfo
  39. lastGetPricingTime time.Time
  40. updatePricingLock sync.Mutex
  41. // 缓存映射:模型名 -> 启用分组 / 计费类型
  42. modelEnableGroups = make(map[string][]string)
  43. modelQuotaTypeMap = make(map[string]int)
  44. modelEnableGroupsLock = sync.RWMutex{}
  45. )
  46. var (
  47. modelSupportEndpointTypes = make(map[string][]constant.EndpointType)
  48. modelSupportEndpointsLock = sync.RWMutex{}
  49. )
  50. func GetPricing() []Pricing {
  51. if time.Since(lastGetPricingTime) > time.Minute*1 || len(pricingMap) == 0 {
  52. updatePricingLock.Lock()
  53. defer updatePricingLock.Unlock()
  54. // Double check after acquiring the lock
  55. if time.Since(lastGetPricingTime) > time.Minute*1 || len(pricingMap) == 0 {
  56. modelSupportEndpointsLock.Lock()
  57. defer modelSupportEndpointsLock.Unlock()
  58. updatePricing()
  59. }
  60. }
  61. return pricingMap
  62. }
  63. // GetVendors 返回当前定价接口使用到的供应商信息
  64. func GetVendors() []PricingVendor {
  65. if time.Since(lastGetPricingTime) > time.Minute*1 || len(pricingMap) == 0 {
  66. // 保证先刷新一次
  67. GetPricing()
  68. }
  69. return vendorsList
  70. }
  71. func GetModelSupportEndpointTypes(model string) []constant.EndpointType {
  72. if model == "" {
  73. return make([]constant.EndpointType, 0)
  74. }
  75. modelSupportEndpointsLock.RLock()
  76. defer modelSupportEndpointsLock.RUnlock()
  77. if endpoints, ok := modelSupportEndpointTypes[model]; ok {
  78. return endpoints
  79. }
  80. return make([]constant.EndpointType, 0)
  81. }
  82. func updatePricing() {
  83. //modelRatios := common.GetModelRatios()
  84. enableAbilities, err := GetAllEnableAbilityWithChannels()
  85. if err != nil {
  86. common.SysLog(fmt.Sprintf("GetAllEnableAbilityWithChannels error: %v", err))
  87. return
  88. }
  89. // 预加载模型元数据与供应商一次,避免循环查询
  90. var allMeta []Model
  91. _ = DB.Find(&allMeta).Error
  92. metaMap := make(map[string]*Model)
  93. prefixList := make([]*Model, 0)
  94. suffixList := make([]*Model, 0)
  95. containsList := make([]*Model, 0)
  96. for i := range allMeta {
  97. m := &allMeta[i]
  98. if m.NameRule == NameRuleExact {
  99. metaMap[m.ModelName] = m
  100. } else {
  101. switch m.NameRule {
  102. case NameRulePrefix:
  103. prefixList = append(prefixList, m)
  104. case NameRuleSuffix:
  105. suffixList = append(suffixList, m)
  106. case NameRuleContains:
  107. containsList = append(containsList, m)
  108. }
  109. }
  110. }
  111. // 将非精确规则模型匹配到 metaMap
  112. for _, m := range prefixList {
  113. for _, pricingModel := range enableAbilities {
  114. if strings.HasPrefix(pricingModel.Model, m.ModelName) {
  115. if _, exists := metaMap[pricingModel.Model]; !exists {
  116. metaMap[pricingModel.Model] = m
  117. }
  118. }
  119. }
  120. }
  121. for _, m := range suffixList {
  122. for _, pricingModel := range enableAbilities {
  123. if strings.HasSuffix(pricingModel.Model, m.ModelName) {
  124. if _, exists := metaMap[pricingModel.Model]; !exists {
  125. metaMap[pricingModel.Model] = m
  126. }
  127. }
  128. }
  129. }
  130. for _, m := range containsList {
  131. for _, pricingModel := range enableAbilities {
  132. if strings.Contains(pricingModel.Model, m.ModelName) {
  133. if _, exists := metaMap[pricingModel.Model]; !exists {
  134. metaMap[pricingModel.Model] = m
  135. }
  136. }
  137. }
  138. }
  139. // 预加载供应商
  140. var vendors []Vendor
  141. _ = DB.Find(&vendors).Error
  142. vendorMap := make(map[int]*Vendor)
  143. for i := range vendors {
  144. vendorMap[vendors[i].Id] = &vendors[i]
  145. }
  146. // 初始化默认供应商映射
  147. initDefaultVendorMapping(metaMap, vendorMap, enableAbilities)
  148. // 构建对前端友好的供应商列表
  149. vendorsList = make([]PricingVendor, 0, len(vendorMap))
  150. for _, v := range vendorMap {
  151. vendorsList = append(vendorsList, PricingVendor{
  152. ID: v.Id,
  153. Name: v.Name,
  154. Description: v.Description,
  155. Icon: v.Icon,
  156. })
  157. }
  158. modelGroupsMap := make(map[string]*types.Set[string])
  159. for _, ability := range enableAbilities {
  160. groups, ok := modelGroupsMap[ability.Model]
  161. if !ok {
  162. groups = types.NewSet[string]()
  163. modelGroupsMap[ability.Model] = groups
  164. }
  165. groups.Add(ability.Group)
  166. }
  167. //这里使用切片而不是Set,因为一个模型可能支持多个端点类型,并且第一个端点是优先使用端点
  168. modelSupportEndpointsStr := make(map[string][]string)
  169. // 先根据已有能力填充原生端点
  170. for _, ability := range enableAbilities {
  171. endpoints := modelSupportEndpointsStr[ability.Model]
  172. channelTypes := common.GetEndpointTypesByChannelType(ability.ChannelType, ability.Model)
  173. for _, channelType := range channelTypes {
  174. if !common.StringsContains(endpoints, string(channelType)) {
  175. endpoints = append(endpoints, string(channelType))
  176. }
  177. }
  178. modelSupportEndpointsStr[ability.Model] = endpoints
  179. }
  180. // 再补充模型自定义端点:若配置有效则替换默认端点,不做合并
  181. for modelName, meta := range metaMap {
  182. if strings.TrimSpace(meta.Endpoints) == "" {
  183. continue
  184. }
  185. var raw map[string]interface{}
  186. if err := json.Unmarshal([]byte(meta.Endpoints), &raw); err == nil {
  187. endpoints := make([]string, 0, len(raw))
  188. for k, v := range raw {
  189. switch v.(type) {
  190. case string, map[string]interface{}:
  191. if !common.StringsContains(endpoints, k) {
  192. endpoints = append(endpoints, k)
  193. }
  194. }
  195. }
  196. if len(endpoints) > 0 {
  197. modelSupportEndpointsStr[modelName] = endpoints
  198. }
  199. }
  200. }
  201. modelSupportEndpointTypes = make(map[string][]constant.EndpointType)
  202. for model, endpoints := range modelSupportEndpointsStr {
  203. supportedEndpoints := make([]constant.EndpointType, 0)
  204. for _, endpointStr := range endpoints {
  205. endpointType := constant.EndpointType(endpointStr)
  206. supportedEndpoints = append(supportedEndpoints, endpointType)
  207. }
  208. modelSupportEndpointTypes[model] = supportedEndpoints
  209. }
  210. // 构建全局 supportedEndpointMap(默认 + 自定义覆盖)
  211. supportedEndpointMap = make(map[string]common.EndpointInfo)
  212. // 1. 默认端点
  213. for _, endpoints := range modelSupportEndpointTypes {
  214. for _, et := range endpoints {
  215. if info, ok := common.GetDefaultEndpointInfo(et); ok {
  216. if _, exists := supportedEndpointMap[string(et)]; !exists {
  217. supportedEndpointMap[string(et)] = info
  218. }
  219. }
  220. }
  221. }
  222. // 2. 自定义端点(models 表)覆盖默认
  223. for _, meta := range metaMap {
  224. if strings.TrimSpace(meta.Endpoints) == "" {
  225. continue
  226. }
  227. var raw map[string]interface{}
  228. if err := json.Unmarshal([]byte(meta.Endpoints), &raw); err == nil {
  229. for k, v := range raw {
  230. switch val := v.(type) {
  231. case string:
  232. supportedEndpointMap[k] = common.EndpointInfo{Path: val, Method: "POST"}
  233. case map[string]interface{}:
  234. ep := common.EndpointInfo{Method: "POST"}
  235. if p, ok := val["path"].(string); ok {
  236. ep.Path = p
  237. }
  238. if m, ok := val["method"].(string); ok {
  239. ep.Method = strings.ToUpper(m)
  240. }
  241. supportedEndpointMap[k] = ep
  242. default:
  243. // ignore unsupported types
  244. }
  245. }
  246. }
  247. }
  248. pricingMap = make([]Pricing, 0)
  249. for model, groups := range modelGroupsMap {
  250. pricing := Pricing{
  251. ModelName: model,
  252. EnableGroup: groups.Items(),
  253. SupportedEndpointTypes: modelSupportEndpointTypes[model],
  254. }
  255. // 补充模型元数据(描述、标签、供应商、状态)
  256. if meta, ok := metaMap[model]; ok {
  257. // 若模型被禁用(status!=1),则直接跳过,不返回给前端
  258. if meta.Status != 1 || meta.Type == 0 {
  259. continue
  260. }
  261. pricing.Description = meta.Description
  262. pricing.Icon = meta.Icon
  263. pricing.Tags = meta.Tags
  264. pricing.VendorID = meta.VendorID
  265. pricing.Type = meta.Type
  266. }
  267. modelPrice, findPrice := ratio_setting.GetModelPrice(model, false)
  268. if findPrice {
  269. pricing.ModelPrice = modelPrice
  270. pricing.QuotaType = 1
  271. } else {
  272. modelRatio, _, _ := ratio_setting.GetModelRatio(model)
  273. pricing.ModelRatio = modelRatio
  274. pricing.CompletionRatio = ratio_setting.GetCompletionRatio(model)
  275. pricing.QuotaType = 0
  276. }
  277. pricingMap = append(pricingMap, pricing)
  278. }
  279. // 防止大更新后数据不通用
  280. if len(pricingMap) > 0 {
  281. pricingMap[0].PricingVersion = "82c4a357505fff6fee8462c3f7ec8a645bb95532669cb73b2cabee6a416ec24f"
  282. }
  283. // 刷新缓存映射,供高并发快速查询
  284. modelEnableGroupsLock.Lock()
  285. modelEnableGroups = make(map[string][]string)
  286. modelQuotaTypeMap = make(map[string]int)
  287. for _, p := range pricingMap {
  288. modelEnableGroups[p.ModelName] = p.EnableGroup
  289. modelQuotaTypeMap[p.ModelName] = p.QuotaType
  290. }
  291. modelEnableGroupsLock.Unlock()
  292. lastGetPricingTime = time.Now()
  293. }
  294. // GetSupportedEndpointMap 返回全局端点到路径的映射
  295. func GetSupportedEndpointMap() map[string]common.EndpointInfo {
  296. return supportedEndpointMap
  297. }