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

403 строки
11 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 parseModelEndpointTypes(rawValue string) []string {
  104. trimmed := strings.TrimSpace(rawValue)
  105. if trimmed == "" {
  106. return nil
  107. }
  108. var rawObject map[string]interface{}
  109. if err := common.Unmarshal([]byte(trimmed), &rawObject); err == nil {
  110. endpoints := make([]string, 0, len(rawObject))
  111. for k, v := range rawObject {
  112. switch v.(type) {
  113. case string, map[string]interface{}:
  114. if !common.StringsContains(endpoints, k) {
  115. endpoints = append(endpoints, k)
  116. }
  117. }
  118. }
  119. if len(endpoints) == 0 {
  120. return nil
  121. }
  122. return endpoints
  123. }
  124. var rawList []string
  125. if err := common.Unmarshal([]byte(trimmed), &rawList); err == nil {
  126. endpoints := make([]string, 0, len(rawList))
  127. for _, endpoint := range rawList {
  128. endpoint = strings.TrimSpace(endpoint)
  129. if endpoint == "" || common.StringsContains(endpoints, endpoint) {
  130. continue
  131. }
  132. endpoints = append(endpoints, endpoint)
  133. }
  134. if len(endpoints) == 0 {
  135. return nil
  136. }
  137. return endpoints
  138. }
  139. return nil
  140. }
  141. func updatePricing() {
  142. enableAbilities, err := GetAllEnableAbilityWithChannels()
  143. if err != nil {
  144. common.SysLog(fmt.Sprintf("GetAllEnableAbilityWithChannels error: %v", err))
  145. return
  146. }
  147. var allMeta []Model
  148. _ = DB.Find(&allMeta).Error
  149. metaMap := make(map[string]*Model)
  150. prefixList := make([]*Model, 0)
  151. suffixList := make([]*Model, 0)
  152. containsList := make([]*Model, 0)
  153. for i := range allMeta {
  154. m := &allMeta[i]
  155. if m.NameRule == NameRuleExact {
  156. metaMap[m.ModelName] = m
  157. continue
  158. }
  159. switch m.NameRule {
  160. case NameRulePrefix:
  161. prefixList = append(prefixList, m)
  162. case NameRuleSuffix:
  163. suffixList = append(suffixList, m)
  164. case NameRuleContains:
  165. containsList = append(containsList, m)
  166. }
  167. }
  168. for _, m := range prefixList {
  169. for _, ability := range enableAbilities {
  170. if strings.HasPrefix(ability.Model, m.ModelName) {
  171. if _, exists := metaMap[ability.Model]; !exists {
  172. metaMap[ability.Model] = m
  173. }
  174. }
  175. }
  176. }
  177. for _, m := range suffixList {
  178. for _, ability := range enableAbilities {
  179. if strings.HasSuffix(ability.Model, m.ModelName) {
  180. if _, exists := metaMap[ability.Model]; !exists {
  181. metaMap[ability.Model] = m
  182. }
  183. }
  184. }
  185. }
  186. for _, m := range containsList {
  187. for _, ability := range enableAbilities {
  188. if strings.Contains(ability.Model, m.ModelName) {
  189. if _, exists := metaMap[ability.Model]; !exists {
  190. metaMap[ability.Model] = m
  191. }
  192. }
  193. }
  194. }
  195. var vendors []Vendor
  196. _ = DB.Find(&vendors).Error
  197. vendorMap := make(map[int]*Vendor, len(vendors))
  198. for i := range vendors {
  199. vendorMap[vendors[i].Id] = &vendors[i]
  200. }
  201. initDefaultVendorMapping(metaMap, vendorMap, enableAbilities)
  202. vendorOrderMap := make(map[int]int, len(vendorMap))
  203. for _, v := range vendorMap {
  204. vendorOrderMap[v.Id] = v.SortOrder
  205. }
  206. vendorsList = make([]PricingVendor, 0, len(vendorMap))
  207. for _, v := range vendorMap {
  208. vendorsList = append(vendorsList, PricingVendor{
  209. ID: v.Id,
  210. Name: v.Name,
  211. Description: v.Description,
  212. Icon: v.Icon,
  213. })
  214. }
  215. sort.Slice(vendorsList, func(i, j int) bool {
  216. oi := vendorOrderMap[vendorsList[i].ID]
  217. oj := vendorOrderMap[vendorsList[j].ID]
  218. if oi != oj {
  219. return oi < oj
  220. }
  221. return vendorsList[i].ID < vendorsList[j].ID
  222. })
  223. modelGroupsMap := make(map[string]*types.Set[string])
  224. for _, ability := range enableAbilities {
  225. groups, ok := modelGroupsMap[ability.Model]
  226. if !ok {
  227. groups = types.NewSet[string]()
  228. modelGroupsMap[ability.Model] = groups
  229. }
  230. groups.Add(ability.Group)
  231. }
  232. modelSupportEndpointsStr := make(map[string][]string)
  233. for _, ability := range enableAbilities {
  234. endpoints := modelSupportEndpointsStr[ability.Model]
  235. channelTypes := common.GetEndpointTypesByChannelType(ability.ChannelType, ability.Model)
  236. for _, channelType := range channelTypes {
  237. if !common.StringsContains(endpoints, string(channelType)) {
  238. endpoints = append(endpoints, string(channelType))
  239. }
  240. }
  241. modelSupportEndpointsStr[ability.Model] = endpoints
  242. }
  243. for modelName, meta := range metaMap {
  244. if strings.TrimSpace(meta.Endpoints) == "" {
  245. continue
  246. }
  247. if endpoints := parseModelEndpointTypes(meta.Endpoints); len(endpoints) > 0 {
  248. modelSupportEndpointsStr[modelName] = endpoints
  249. }
  250. }
  251. modelSupportEndpointTypes = make(map[string][]constant.EndpointType)
  252. for model, endpoints := range modelSupportEndpointsStr {
  253. supportedEndpoints := make([]constant.EndpointType, 0, len(endpoints))
  254. for _, endpointStr := range endpoints {
  255. supportedEndpoints = append(supportedEndpoints, constant.EndpointType(endpointStr))
  256. }
  257. modelSupportEndpointTypes[model] = supportedEndpoints
  258. }
  259. supportedEndpointMap = make(map[string]common.EndpointInfo)
  260. for _, endpoints := range modelSupportEndpointTypes {
  261. for _, et := range endpoints {
  262. if info, ok := common.GetDefaultEndpointInfo(et); ok {
  263. if _, exists := supportedEndpointMap[string(et)]; !exists {
  264. supportedEndpointMap[string(et)] = info
  265. }
  266. }
  267. }
  268. }
  269. for _, meta := range metaMap {
  270. if strings.TrimSpace(meta.Endpoints) == "" {
  271. continue
  272. }
  273. var raw map[string]interface{}
  274. if err := common.Unmarshal([]byte(meta.Endpoints), &raw); err == nil {
  275. for k, v := range raw {
  276. switch val := v.(type) {
  277. case string:
  278. supportedEndpointMap[k] = common.EndpointInfo{Path: val, Method: "POST"}
  279. case map[string]interface{}:
  280. ep := common.EndpointInfo{Method: "POST"}
  281. if p, ok := val["path"].(string); ok {
  282. ep.Path = p
  283. }
  284. if m, ok := val["method"].(string); ok {
  285. ep.Method = strings.ToUpper(m)
  286. }
  287. supportedEndpointMap[k] = ep
  288. }
  289. }
  290. }
  291. }
  292. pricingMap = make([]Pricing, 0, len(modelGroupsMap))
  293. for model, groups := range modelGroupsMap {
  294. pricing := Pricing{
  295. ModelName: model,
  296. EnableGroup: groups.Items(),
  297. SupportedEndpointTypes: modelSupportEndpointTypes[model],
  298. }
  299. if meta, ok := metaMap[model]; ok {
  300. if meta.Status == 0 {
  301. continue
  302. }
  303. pricing.Description = meta.Description
  304. pricing.Icon = meta.Icon
  305. pricing.Tags = meta.Tags
  306. pricing.VendorID = meta.VendorID
  307. pricing.Type = meta.Type
  308. }
  309. applyGlobalDefault(&pricing, model)
  310. pricingMap = append(pricingMap, pricing)
  311. }
  312. sort.Slice(pricingMap, func(i, j int) bool {
  313. mi, okI := metaMap[pricingMap[i].ModelName]
  314. mj, okJ := metaMap[pricingMap[j].ModelName]
  315. si := 999999
  316. sj := 999999
  317. if okI {
  318. si = mi.SortOrder
  319. }
  320. if okJ {
  321. sj = mj.SortOrder
  322. }
  323. if si != sj {
  324. return si < sj
  325. }
  326. return pricingMap[i].ModelName < pricingMap[j].ModelName
  327. })
  328. if len(pricingMap) > 0 {
  329. pricingMap[0].PricingVersion = "82c4a357505fff6fee8462c3f7ec8a645bb95532669cb73b2cabee6a416ec24f"
  330. }
  331. modelEnableGroupsLock.Lock()
  332. modelEnableGroups = make(map[string][]string)
  333. modelQuotaTypeMap = make(map[string]int)
  334. byModel := make(map[string]*Pricing, len(pricingMap))
  335. for i := range pricingMap {
  336. p := &pricingMap[i]
  337. modelEnableGroups[p.ModelName] = p.EnableGroup
  338. modelQuotaTypeMap[p.ModelName] = p.QuotaType
  339. byModel[p.ModelName] = p
  340. }
  341. pricingByModel = byModel
  342. modelEnableGroupsLock.Unlock()
  343. lastGetPricingTime = time.Now()
  344. }
  345. func GetSupportedEndpointMap() map[string]common.EndpointInfo {
  346. return supportedEndpointMap
  347. }
  348. func applyGlobalDefault(pricing *Pricing, model string) {
  349. modelPrice, findPrice := ratio_setting.GetModelPrice(model, false)
  350. if findPrice {
  351. pricing.ModelPrice = modelPrice
  352. pricing.ModelRatio = 0
  353. pricing.CompletionRatio = 0
  354. pricing.QuotaType = QuotaTypeByCall
  355. } else {
  356. modelRatio, _, _ := ratio_setting.GetModelRatio(model)
  357. pricing.ModelPrice = 0
  358. pricing.ModelRatio = modelRatio
  359. pricing.CompletionRatio = ratio_setting.GetCompletionRatio(model)
  360. pricing.QuotaType = QuotaTypeByTokens
  361. }
  362. pricing.CacheRatio, _ = ratio_setting.GetCacheRatio(model)
  363. pricing.CacheCreationRatio, _ = ratio_setting.GetCreateCacheRatio(model)
  364. }