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ů.
 
 
 

369 řádky
10 KiB

  1. package model
  2. import (
  3. "fmt"
  4. "sort"
  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. const (
  38. QuotaTypeByTokens = 0
  39. QuotaTypeByCall = 1
  40. )
  41. var (
  42. pricingMap []Pricing
  43. pricingByModel map[string]*Pricing
  44. vendorsList []PricingVendor
  45. supportedEndpointMap map[string]common.EndpointInfo
  46. lastGetPricingTime time.Time
  47. updatePricingLock sync.Mutex
  48. modelEnableGroups = make(map[string][]string)
  49. modelQuotaTypeMap = make(map[string]int)
  50. modelEnableGroupsLock = sync.RWMutex{}
  51. )
  52. var (
  53. modelSupportEndpointTypes = make(map[string][]constant.EndpointType)
  54. modelSupportEndpointsLock = sync.RWMutex{}
  55. )
  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 || len(pricingMap) == 0 {
  69. updatePricingLock.Lock()
  70. defer updatePricingLock.Unlock()
  71. if time.Since(lastGetPricingTime) > time.Minute || len(pricingMap) == 0 {
  72. modelSupportEndpointsLock.Lock()
  73. defer modelSupportEndpointsLock.Unlock()
  74. updatePricing()
  75. }
  76. }
  77. return pricingMap
  78. }
  79. func GetPricingByModel(modelName string) *Pricing {
  80. GetPricing()
  81. modelEnableGroupsLock.RLock()
  82. p := pricingByModel[modelName]
  83. modelEnableGroupsLock.RUnlock()
  84. return p
  85. }
  86. func GetVendors() []PricingVendor {
  87. if time.Since(lastGetPricingTime) > time.Minute || len(pricingMap) == 0 {
  88. GetPricing()
  89. }
  90. return vendorsList
  91. }
  92. func GetModelSupportEndpointTypes(model string) []constant.EndpointType {
  93. if model == "" {
  94. return make([]constant.EndpointType, 0)
  95. }
  96. modelSupportEndpointsLock.RLock()
  97. defer modelSupportEndpointsLock.RUnlock()
  98. if endpoints, ok := modelSupportEndpointTypes[model]; ok {
  99. return endpoints
  100. }
  101. return make([]constant.EndpointType, 0)
  102. }
  103. func updatePricing() {
  104. enableAbilities, err := GetAllEnableAbilityWithChannels()
  105. if err != nil {
  106. common.SysLog(fmt.Sprintf("GetAllEnableAbilityWithChannels error: %v", err))
  107. return
  108. }
  109. var allMeta []Model
  110. _ = DB.Find(&allMeta).Error
  111. metaMap := make(map[string]*Model)
  112. prefixList := make([]*Model, 0)
  113. suffixList := make([]*Model, 0)
  114. containsList := make([]*Model, 0)
  115. for i := range allMeta {
  116. m := &allMeta[i]
  117. switch m.NameRule {
  118. case NameRuleExact:
  119. metaMap[m.ModelName] = m
  120. case NameRulePrefix:
  121. prefixList = append(prefixList, m)
  122. case NameRuleSuffix:
  123. suffixList = append(suffixList, m)
  124. case NameRuleContains:
  125. containsList = append(containsList, m)
  126. }
  127. }
  128. for _, m := range prefixList {
  129. for _, pricingModel := range enableAbilities {
  130. if strings.HasPrefix(pricingModel.Model, m.ModelName) {
  131. if _, exists := metaMap[pricingModel.Model]; !exists {
  132. metaMap[pricingModel.Model] = m
  133. }
  134. }
  135. }
  136. }
  137. for _, m := range suffixList {
  138. for _, pricingModel := range enableAbilities {
  139. if strings.HasSuffix(pricingModel.Model, m.ModelName) {
  140. if _, exists := metaMap[pricingModel.Model]; !exists {
  141. metaMap[pricingModel.Model] = m
  142. }
  143. }
  144. }
  145. }
  146. for _, m := range containsList {
  147. for _, pricingModel := range enableAbilities {
  148. if strings.Contains(pricingModel.Model, m.ModelName) {
  149. if _, exists := metaMap[pricingModel.Model]; !exists {
  150. metaMap[pricingModel.Model] = m
  151. }
  152. }
  153. }
  154. }
  155. var vendors []Vendor
  156. _ = DB.Find(&vendors).Error
  157. vendorMap := make(map[int]*Vendor)
  158. for i := range vendors {
  159. vendorMap[vendors[i].Id] = &vendors[i]
  160. }
  161. initDefaultVendorMapping(metaMap, vendorMap, enableAbilities)
  162. vendorOrderMap := make(map[int]int)
  163. for _, v := range vendorMap {
  164. vendorOrderMap[v.Id] = v.SortOrder
  165. }
  166. vendorsList = make([]PricingVendor, 0, len(vendorMap))
  167. for _, v := range vendorMap {
  168. vendorsList = append(vendorsList, PricingVendor{
  169. ID: v.Id,
  170. Name: v.Name,
  171. Description: v.Description,
  172. Icon: v.Icon,
  173. })
  174. }
  175. sort.Slice(vendorsList, func(i, j int) bool {
  176. oi := vendorOrderMap[vendorsList[i].ID]
  177. oj := vendorOrderMap[vendorsList[j].ID]
  178. if oi != oj {
  179. return oi < oj
  180. }
  181. return vendorsList[i].ID < vendorsList[j].ID
  182. })
  183. modelGroupsMap := make(map[string]*types.Set[string])
  184. for _, ability := range enableAbilities {
  185. groups, ok := modelGroupsMap[ability.Model]
  186. if !ok {
  187. groups = types.NewSet[string]()
  188. modelGroupsMap[ability.Model] = groups
  189. }
  190. groups.Add(ability.Group)
  191. }
  192. modelSupportEndpointsStr := make(map[string][]string)
  193. for _, ability := range enableAbilities {
  194. endpoints := modelSupportEndpointsStr[ability.Model]
  195. channelTypes := common.GetEndpointTypesByChannelType(ability.ChannelType, ability.Model)
  196. for _, channelType := range channelTypes {
  197. if !common.StringsContains(endpoints, string(channelType)) {
  198. endpoints = append(endpoints, string(channelType))
  199. }
  200. }
  201. modelSupportEndpointsStr[ability.Model] = endpoints
  202. }
  203. for modelName, meta := range metaMap {
  204. if strings.TrimSpace(meta.Endpoints) == "" {
  205. continue
  206. }
  207. var raw map[string]interface{}
  208. if err := common.Unmarshal([]byte(meta.Endpoints), &raw); err == nil {
  209. endpoints := make([]string, 0, len(raw))
  210. for k, v := range raw {
  211. switch v.(type) {
  212. case string, map[string]interface{}:
  213. if !common.StringsContains(endpoints, k) {
  214. endpoints = append(endpoints, k)
  215. }
  216. }
  217. }
  218. if len(endpoints) > 0 {
  219. modelSupportEndpointsStr[modelName] = endpoints
  220. }
  221. }
  222. }
  223. modelSupportEndpointTypes = make(map[string][]constant.EndpointType)
  224. for model, endpoints := range modelSupportEndpointsStr {
  225. supportedEndpoints := make([]constant.EndpointType, 0)
  226. for _, endpointStr := range endpoints {
  227. supportedEndpoints = append(supportedEndpoints, constant.EndpointType(endpointStr))
  228. }
  229. modelSupportEndpointTypes[model] = supportedEndpoints
  230. }
  231. supportedEndpointMap = make(map[string]common.EndpointInfo)
  232. for _, endpoints := range modelSupportEndpointTypes {
  233. for _, et := range endpoints {
  234. if info, ok := common.GetDefaultEndpointInfo(et); ok {
  235. if _, exists := supportedEndpointMap[string(et)]; !exists {
  236. supportedEndpointMap[string(et)] = info
  237. }
  238. }
  239. }
  240. }
  241. for _, meta := range metaMap {
  242. if strings.TrimSpace(meta.Endpoints) == "" {
  243. continue
  244. }
  245. var raw map[string]interface{}
  246. if err := common.Unmarshal([]byte(meta.Endpoints), &raw); err == nil {
  247. for k, v := range raw {
  248. switch val := v.(type) {
  249. case string:
  250. supportedEndpointMap[k] = common.EndpointInfo{Path: val, Method: "POST"}
  251. case map[string]interface{}:
  252. ep := common.EndpointInfo{Method: "POST"}
  253. if p, ok := val["path"].(string); ok {
  254. ep.Path = p
  255. }
  256. if m, ok := val["method"].(string); ok {
  257. ep.Method = strings.ToUpper(m)
  258. }
  259. supportedEndpointMap[k] = ep
  260. }
  261. }
  262. }
  263. }
  264. pricingMap = make([]Pricing, 0, len(modelGroupsMap))
  265. for model, groups := range modelGroupsMap {
  266. pricing := Pricing{
  267. ModelName: model,
  268. EnableGroup: groups.Items(),
  269. SupportedEndpointTypes: modelSupportEndpointTypes[model],
  270. }
  271. if meta, ok := metaMap[model]; ok {
  272. if meta.Status == 0 {
  273. continue
  274. }
  275. pricing.Description = meta.Description
  276. pricing.Icon = meta.Icon
  277. pricing.Tags = meta.Tags
  278. pricing.VendorID = meta.VendorID
  279. pricing.Type = meta.Type
  280. }
  281. applyGlobalDefault(&pricing, model)
  282. pricingMap = append(pricingMap, pricing)
  283. }
  284. sort.Slice(pricingMap, func(i, j int) bool {
  285. mi, okI := metaMap[pricingMap[i].ModelName]
  286. mj, okJ := metaMap[pricingMap[j].ModelName]
  287. si := 999999
  288. sj := 999999
  289. if okI {
  290. si = mi.SortOrder
  291. }
  292. if okJ {
  293. sj = mj.SortOrder
  294. }
  295. if si != sj {
  296. return si < sj
  297. }
  298. return pricingMap[i].ModelName < pricingMap[j].ModelName
  299. })
  300. if len(pricingMap) > 0 {
  301. pricingMap[0].PricingVersion = "82c4a357505fff6fee8462c3f7ec8a645bb95532669cb73b2cabee6a416ec24f"
  302. }
  303. modelEnableGroupsLock.Lock()
  304. modelEnableGroups = make(map[string][]string)
  305. modelQuotaTypeMap = make(map[string]int)
  306. byModel := make(map[string]*Pricing, len(pricingMap))
  307. for i := range pricingMap {
  308. p := &pricingMap[i]
  309. modelEnableGroups[p.ModelName] = p.EnableGroup
  310. modelQuotaTypeMap[p.ModelName] = p.QuotaType
  311. byModel[p.ModelName] = p
  312. }
  313. pricingByModel = byModel
  314. modelEnableGroupsLock.Unlock()
  315. lastGetPricingTime = time.Now()
  316. }
  317. func GetSupportedEndpointMap() map[string]common.EndpointInfo {
  318. return supportedEndpointMap
  319. }
  320. func applyGlobalDefault(pricing *Pricing, model string) {
  321. modelPrice, findPrice := ratio_setting.GetModelPrice(model, false)
  322. if findPrice {
  323. pricing.ModelPrice = modelPrice
  324. pricing.ModelRatio = 0
  325. pricing.CompletionRatio = 0
  326. pricing.QuotaType = QuotaTypeByCall
  327. } else {
  328. modelRatio, _, _ := ratio_setting.GetModelRatio(model)
  329. pricing.ModelPrice = 0
  330. pricing.ModelRatio = modelRatio
  331. pricing.CompletionRatio = ratio_setting.GetCompletionRatio(model)
  332. pricing.QuotaType = QuotaTypeByTokens
  333. }
  334. pricing.CacheRatio, _ = ratio_setting.GetCacheRatio(model)
  335. pricing.CacheCreationRatio, _ = ratio_setting.GetCreateCacheRatio(model)
  336. }