292 lines
9.3 KiB
Go
292 lines
9.3 KiB
Go
package common
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"fmt"
|
|
"math"
|
|
"strings"
|
|
"time"
|
|
|
|
"github.com/gogf/gf/v2/frame/g"
|
|
"github.com/gogf/gf/v2/os/gcache"
|
|
)
|
|
|
|
const (
|
|
litellmPricingURL = "https://raw.githubusercontent.com/BerriAI/litellm/main/model_prices_and_context_window.json"
|
|
litellmCacheKey = "litellm_pricing_data"
|
|
litellmCacheTTL = 24 * time.Hour
|
|
)
|
|
|
|
// LiteLLMModelEntry represents a single model entry in the LiteLLM pricing JSON.
|
|
type LiteLLMModelEntry struct {
|
|
LitellmProvider string `json:"litellm_provider"`
|
|
Mode string `json:"mode"`
|
|
MaxInputTokens int `json:"max_input_tokens"`
|
|
MaxOutputTokens int `json:"max_output_tokens"`
|
|
InputCostPerToken float64 `json:"input_cost_per_token"`
|
|
OutputCostPerToken float64 `json:"output_cost_per_token"`
|
|
CacheReadInputTokenCost float64 `json:"cache_read_input_token_cost"`
|
|
CacheCreationInputTokenCost float64 `json:"cache_creation_input_token_cost"`
|
|
InputCostPerCharacter float64 `json:"input_cost_per_character"`
|
|
OutputCostPerCharacter float64 `json:"output_cost_per_character"`
|
|
OutputCostPerImage float64 `json:"output_cost_per_image"`
|
|
SupportsVision bool `json:"supports_vision"`
|
|
SupportsFunctionCalling bool `json:"supports_function_calling"`
|
|
SupportsParallelFuncCalling bool `json:"supports_parallel_function_calling"`
|
|
SupportsToolChoice bool `json:"supports_tool_choice"`
|
|
SupportsResponseSchema bool `json:"supports_response_schema"`
|
|
SupportsSystemMessages bool `json:"supports_system_messages"`
|
|
SupportsPromptCaching bool `json:"supports_prompt_caching"`
|
|
SupportsAudioInput bool `json:"supports_audio_input"`
|
|
SupportsAudioOutput bool `json:"supports_audio_output"`
|
|
SupportsPdfInput bool `json:"supports_pdf_input"`
|
|
SupportsEmbeddingImage bool `json:"supports_embedding_image"`
|
|
SupportsReasoning bool `json:"supports_reasoning"`
|
|
SupportsWebSearch bool `json:"supports_web_search"`
|
|
DeprecationDate string `json:"deprecation_date"`
|
|
}
|
|
|
|
var litellmCache = gcache.New()
|
|
|
|
// FetchLiteLLMPricing fetches and caches the LiteLLM pricing data.
|
|
// Returns the full map of model_name → LiteLLMModelEntry.
|
|
func FetchLiteLLMPricing(ctx context.Context) (map[string]*LiteLLMModelEntry, error) {
|
|
cached, err := litellmCache.Get(ctx, litellmCacheKey)
|
|
if err == nil && cached != nil {
|
|
if data, ok := cached.Val().(map[string]*LiteLLMModelEntry); ok {
|
|
return data, nil
|
|
}
|
|
}
|
|
|
|
g.Log().Info(ctx, "[PricingSource] fetching LiteLLM pricing data from remote...")
|
|
|
|
resp, err := g.Client().SetHeaderMap(map[string]string{
|
|
"User-Agent": "github.com/qianfree/team-api/1.0",
|
|
}).Get(ctx, litellmPricingURL)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("获取远程数据失败: %w", err)
|
|
}
|
|
defer resp.Close()
|
|
|
|
if resp.StatusCode != 200 {
|
|
return nil, fmt.Errorf("获取远程数据失败: HTTP %d", resp.StatusCode)
|
|
}
|
|
|
|
body := resp.ReadAll()
|
|
if len(body) == 0 {
|
|
return nil, fmt.Errorf("获取远程数据失败: 响应为空")
|
|
}
|
|
|
|
var rawEntries map[string]json.RawMessage
|
|
if err := json.Unmarshal(body, &rawEntries); err != nil {
|
|
return nil, fmt.Errorf("获取远程数据失败: JSON 解析失败: %w", err)
|
|
}
|
|
|
|
delete(rawEntries, "sample_spec")
|
|
|
|
raw := make(map[string]*LiteLLMModelEntry, len(rawEntries))
|
|
for name, rawMsg := range rawEntries {
|
|
var entry LiteLLMModelEntry
|
|
if err := json.Unmarshal(rawMsg, &entry); err != nil {
|
|
continue
|
|
}
|
|
raw[name] = &entry
|
|
}
|
|
|
|
if err := litellmCache.Set(ctx, litellmCacheKey, raw, litellmCacheTTL); err != nil {
|
|
g.Log().Warning(ctx, "[PricingSource] failed to cache pricing data: %v", err)
|
|
}
|
|
|
|
g.Log().Infof(ctx, "[PricingSource] loaded %d models from LiteLLM", len(raw))
|
|
return raw, nil
|
|
}
|
|
|
|
// FindLiteLLMModel searches for a model in the LiteLLM pricing data.
|
|
// Tries exact match first, then strips common prefixes (provider/).
|
|
func FindLiteLLMModel(data map[string]*LiteLLMModelEntry, modelName string) (string, *LiteLLMModelEntry) {
|
|
// Exact match
|
|
if entry, ok := data[modelName]; ok {
|
|
return modelName, entry
|
|
}
|
|
|
|
// Try common provider prefixes that LiteLLM uses
|
|
providers := []string{"openai/", "anthropic/", "vertex_ai/", "bedrock/", "azure/"}
|
|
for _, prefix := range providers {
|
|
if entry, ok := data[prefix+modelName]; ok {
|
|
return prefix + modelName, entry
|
|
}
|
|
}
|
|
|
|
// Try stripping our custom prefixes (in case model name has a provider prefix)
|
|
for _, prefix := range providers {
|
|
if strings.HasPrefix(modelName, prefix) {
|
|
stripped := strings.TrimPrefix(modelName, prefix)
|
|
if entry, ok := data[stripped]; ok {
|
|
return stripped, entry
|
|
}
|
|
}
|
|
}
|
|
|
|
return "", nil
|
|
}
|
|
|
|
// ---------- models.dev 数据源 ----------
|
|
|
|
const (
|
|
modelsDevURL = "https://models.dev/api.json"
|
|
modelsDevCacheKey = "models_dev_pricing_data"
|
|
modelsDevCacheTTL = 24 * time.Hour
|
|
)
|
|
|
|
// ModelsDevProvider represents a provider entry in the models.dev JSON.
|
|
type ModelsDevProvider struct {
|
|
Models map[string]ModelsDevModel `json:"models"`
|
|
}
|
|
|
|
// ModelsDevModel represents a model entry with cost info.
|
|
type ModelsDevModel struct {
|
|
Cost ModelsDevCost `json:"cost"`
|
|
}
|
|
|
|
// ModelsDevCost represents pricing in USD per 1M tokens.
|
|
type ModelsDevCost struct {
|
|
Input *float64 `json:"input"`
|
|
Output *float64 `json:"output"`
|
|
CacheRead *float64 `json:"cache_read"`
|
|
}
|
|
|
|
// ModelsDevModelEntry is a flattened, validated model entry used for lookup.
|
|
type ModelsDevModelEntry struct {
|
|
Provider string
|
|
Input float64 // USD/1M tokens, 0 = free
|
|
Output *float64 // USD/1M tokens, nil = not available
|
|
CacheRead *float64 // USD/1M tokens, nil = not available
|
|
}
|
|
|
|
var modelsDevCache = gcache.New()
|
|
|
|
// FetchModelsDevPricing fetches and caches the models.dev pricing data.
|
|
// Returns a map of model_name → ModelsDevModelEntry (cheapest non-zero input across providers).
|
|
func FetchModelsDevPricing(ctx context.Context) (map[string]*ModelsDevModelEntry, error) {
|
|
cached, err := modelsDevCache.Get(ctx, modelsDevCacheKey)
|
|
if err == nil && cached != nil {
|
|
if data, ok := cached.Val().(map[string]*ModelsDevModelEntry); ok {
|
|
return data, nil
|
|
}
|
|
}
|
|
|
|
g.Log().Info(ctx, "[PricingSource] fetching models.dev pricing data from remote...")
|
|
|
|
resp, err := g.Client().SetHeaderMap(map[string]string{
|
|
"User-Agent": "github.com/qianfree/team-api/1.0",
|
|
}).Get(ctx, modelsDevURL)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("获取远程数据失败: %w", err)
|
|
}
|
|
defer resp.Close()
|
|
|
|
if resp.StatusCode != 200 {
|
|
return nil, fmt.Errorf("获取远程数据失败: HTTP %d", resp.StatusCode)
|
|
}
|
|
|
|
body := resp.ReadAll()
|
|
if len(body) == 0 {
|
|
return nil, fmt.Errorf("获取远程数据失败: 响应为空")
|
|
}
|
|
|
|
var rawProviders map[string]ModelsDevProvider
|
|
if err := json.Unmarshal(body, &rawProviders); err != nil {
|
|
return nil, fmt.Errorf("获取远程数据失败: JSON 解析失败: %w", err)
|
|
}
|
|
|
|
// Build map: model_name → cheapest ModelsDevModelEntry across providers
|
|
result := make(map[string]*ModelsDevModelEntry, 256)
|
|
for providerName, provider := range rawProviders {
|
|
for modelName, model := range provider.Models {
|
|
if model.Cost.Input == nil {
|
|
continue
|
|
}
|
|
input := *model.Cost.Input
|
|
if input < 0 || math.IsNaN(input) || math.IsInf(input, 0) {
|
|
continue
|
|
}
|
|
|
|
var output *float64
|
|
if model.Cost.Output != nil && *model.Cost.Output >= 0 {
|
|
output = model.Cost.Output
|
|
}
|
|
|
|
// input=0 but output>0 cannot derive ratios
|
|
if input == 0 && output != nil && *output > 0 {
|
|
continue
|
|
}
|
|
|
|
var cacheRead *float64
|
|
if model.Cost.CacheRead != nil && *model.Cost.CacheRead >= 0 {
|
|
cacheRead = model.Cost.CacheRead
|
|
}
|
|
|
|
candidate := &ModelsDevModelEntry{
|
|
Provider: providerName,
|
|
Input: input,
|
|
Output: output,
|
|
CacheRead: cacheRead,
|
|
}
|
|
|
|
existing, exists := result[modelName]
|
|
if !exists || shouldReplaceModelsDevCandidate(existing, candidate) {
|
|
result[modelName] = candidate
|
|
}
|
|
}
|
|
}
|
|
|
|
if err := modelsDevCache.Set(ctx, modelsDevCacheKey, result, modelsDevCacheTTL); err != nil {
|
|
g.Log().Warning(ctx, "[PricingSource] failed to cache models.dev data: %v", err)
|
|
}
|
|
|
|
g.Log().Infof(ctx, "[PricingSource] loaded %d models from models.dev", len(result))
|
|
return result, nil
|
|
}
|
|
|
|
func shouldReplaceModelsDevCandidate(current, next *ModelsDevModelEntry) bool {
|
|
currentNonZero := current.Input > 0
|
|
nextNonZero := next.Input > 0
|
|
if currentNonZero != nextNonZero {
|
|
return nextNonZero
|
|
}
|
|
if nextNonZero && next.Input != current.Input {
|
|
return next.Input < current.Input
|
|
}
|
|
return next.Provider < current.Provider
|
|
}
|
|
|
|
// FindModelsDevModel searches for a model in the models.dev pricing data.
|
|
// Tries exact match first, then common provider prefixes.
|
|
func FindModelsDevModel(data map[string]*ModelsDevModelEntry, modelName string) (string, *ModelsDevModelEntry) {
|
|
// Exact match
|
|
if entry, ok := data[modelName]; ok {
|
|
return modelName, entry
|
|
}
|
|
|
|
// Try common provider prefixes
|
|
prefixes := []string{"openai/", "anthropic/", "google/", "vertex_ai/", "bedrock/", "azure/"}
|
|
for _, prefix := range prefixes {
|
|
if entry, ok := data[prefix+modelName]; ok {
|
|
return prefix + modelName, entry
|
|
}
|
|
}
|
|
|
|
// Try stripping provider prefix
|
|
for _, prefix := range prefixes {
|
|
if strings.HasPrefix(modelName, prefix) {
|
|
stripped := strings.TrimPrefix(modelName, prefix)
|
|
if entry, ok := data[stripped]; ok {
|
|
return stripped, entry
|
|
}
|
|
}
|
|
}
|
|
|
|
return "", nil
|
|
}
|