228 lines
8.3 KiB
Go
228 lines
8.3 KiB
Go
package plugin
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"errors"
|
|
"math/big"
|
|
"net/http"
|
|
"strconv"
|
|
"strings"
|
|
|
|
"cpa-ext/internal/pricing"
|
|
)
|
|
|
|
type priceRatesDTO struct {
|
|
InputPer1M string `json:"input_per_1m"`
|
|
CacheReadPer1M string `json:"cache_read_per_1m"`
|
|
CacheWritePer1M string `json:"cache_write_per_1m"`
|
|
OutputPer1M string `json:"output_per_1m"`
|
|
}
|
|
|
|
type longContextDTO struct {
|
|
ThresholdInputTokens int64 `json:"threshold_input_tokens"`
|
|
Comparison string `json:"comparison"`
|
|
priceRatesDTO
|
|
}
|
|
|
|
type priceDTO struct {
|
|
Model string `json:"model"`
|
|
Base priceRatesDTO `json:"base"`
|
|
LongContext *longContextDTO `json:"long_context,omitempty"`
|
|
FastPricingEnabled bool `json:"fast_pricing_enabled"`
|
|
FastMultiplier string `json:"fast_multiplier"`
|
|
}
|
|
|
|
type priceDeleteRequest struct {
|
|
Model string `json:"model"`
|
|
}
|
|
|
|
func (a *App) listPrices() ManagementResponse {
|
|
a.mu.RLock()
|
|
defer a.mu.RUnlock()
|
|
if a.store == nil {
|
|
return jsonManagementResponse(http.StatusServiceUnavailable, map[string]any{"error": map[string]string{"code": "database_unavailable", "message": "价格数据库尚未初始化"}})
|
|
}
|
|
policies, err := a.store.ListPrices(context.Background())
|
|
if err != nil {
|
|
return jsonManagementResponse(http.StatusInternalServerError, map[string]any{"error": map[string]string{"code": "database_error", "message": err.Error()}})
|
|
}
|
|
prices := make([]priceDTO, 0, len(policies))
|
|
for _, policy := range policies {
|
|
prices = append(prices, policyToDTO(policy))
|
|
}
|
|
return jsonManagementResponse(http.StatusOK, map[string]any{"prices": prices})
|
|
}
|
|
|
|
func (a *App) putPrice(body []byte) ManagementResponse {
|
|
var dto priceDTO
|
|
if err := decodeJSONBody(body, &dto); err != nil {
|
|
return jsonManagementResponse(http.StatusBadRequest, map[string]any{"error": map[string]string{"code": "invalid_price", "message": err.Error()}})
|
|
}
|
|
policy, err := dtoToPolicy(dto)
|
|
if err != nil {
|
|
return jsonManagementResponse(http.StatusBadRequest, map[string]any{"error": map[string]string{"code": "invalid_price", "message": err.Error()}})
|
|
}
|
|
a.mu.Lock()
|
|
defer a.mu.Unlock()
|
|
if a.store == nil {
|
|
return jsonManagementResponse(http.StatusServiceUnavailable, map[string]any{"error": map[string]string{"code": "database_unavailable", "message": "价格数据库尚未初始化"}})
|
|
}
|
|
if err := a.store.UpsertPrice(context.Background(), policy); err != nil {
|
|
return jsonManagementResponse(http.StatusInternalServerError, map[string]any{"error": map[string]string{"code": "database_error", "message": err.Error()}})
|
|
}
|
|
a.prices[normalizeModelName(policy.Model)] = policy
|
|
if _, err := a.store.BackfillMissingCosts(context.Background(), policy); err != nil {
|
|
return jsonManagementResponse(http.StatusInternalServerError, map[string]any{"error": map[string]string{"code": "database_error", "message": err.Error()}})
|
|
}
|
|
return jsonManagementResponse(http.StatusOK, policyToDTO(policy))
|
|
}
|
|
|
|
func (a *App) deletePrice(body []byte) ManagementResponse {
|
|
var request priceDeleteRequest
|
|
if err := decodeJSONBody(body, &request); err != nil {
|
|
return jsonManagementResponse(http.StatusBadRequest, map[string]any{"error": map[string]string{"code": "invalid_price", "message": err.Error()}})
|
|
}
|
|
request.Model = strings.TrimSpace(request.Model)
|
|
if request.Model == "" {
|
|
return jsonManagementResponse(http.StatusBadRequest, map[string]any{"error": map[string]string{"code": "invalid_price", "message": "model 不能为空"}})
|
|
}
|
|
a.mu.Lock()
|
|
defer a.mu.Unlock()
|
|
if a.store == nil {
|
|
return jsonManagementResponse(http.StatusServiceUnavailable, map[string]any{"error": map[string]string{"code": "database_unavailable", "message": "价格数据库尚未初始化"}})
|
|
}
|
|
if err := a.store.DeletePrice(context.Background(), request.Model); err != nil {
|
|
return jsonManagementResponse(http.StatusInternalServerError, map[string]any{"error": map[string]string{"code": "database_error", "message": err.Error()}})
|
|
}
|
|
delete(a.prices, normalizeModelName(request.Model))
|
|
return jsonManagementResponse(http.StatusOK, map[string]any{"deleted": request.Model})
|
|
}
|
|
|
|
func dtoToPolicy(dto priceDTO) (pricing.Policy, error) {
|
|
base, err := dtoRates(dto.Base)
|
|
if err != nil {
|
|
return pricing.Policy{}, err
|
|
}
|
|
multiplierText := strings.TrimSpace(dto.FastMultiplier)
|
|
if multiplierText == "" {
|
|
multiplierText = "2.5"
|
|
}
|
|
multiplier, err := parseRatio(multiplierText)
|
|
if err != nil {
|
|
return pricing.Policy{}, errors.New("fast_multiplier 必须是正数，最多保留六位小数")
|
|
}
|
|
policy := pricing.Policy{Model: strings.TrimSpace(dto.Model), Base: base, FastPricingEnabled: dto.FastPricingEnabled, FastMultiplier: multiplier}
|
|
if dto.LongContext != nil {
|
|
rates, ratesErr := dtoRates(dto.LongContext.priceRatesDTO)
|
|
if ratesErr != nil {
|
|
return pricing.Policy{}, ratesErr
|
|
}
|
|
policy.LongContext = &pricing.LongContext{ThresholdInputTokens: dto.LongContext.ThresholdInputTokens, Comparison: dto.LongContext.Comparison, Rates: rates}
|
|
}
|
|
if err := policy.Validate(); err != nil {
|
|
return pricing.Policy{}, err
|
|
}
|
|
return policy, nil
|
|
}
|
|
|
|
func dtoRates(dto priceRatesDTO) (pricing.Rates, error) {
|
|
values := []*int64{new(int64), new(int64), new(int64), new(int64)}
|
|
texts := []string{dto.InputPer1M, dto.CacheReadPer1M, dto.CacheWritePer1M, dto.OutputPer1M}
|
|
for index, text := range texts {
|
|
value, err := parseDecimalMicros(text)
|
|
if err != nil {
|
|
return pricing.Rates{}, errors.New("价格必须是非负数字，最多保留六位小数")
|
|
}
|
|
*values[index] = value
|
|
}
|
|
return pricing.Rates{InputMicrosPer1M: *values[0], CacheReadMicrosPer1M: *values[1], CacheWriteMicrosPer1M: *values[2], OutputMicrosPer1M: *values[3]}, nil
|
|
}
|
|
|
|
func policyToDTO(policy pricing.Policy) priceDTO {
|
|
dto := priceDTO{Model: policy.Model, Base: ratesToDTO(policy.Base), FastPricingEnabled: policy.FastPricingEnabled, FastMultiplier: formatRatio(policy.FastMultiplier)}
|
|
if policy.LongContext != nil {
|
|
dto.LongContext = &longContextDTO{ThresholdInputTokens: policy.LongContext.ThresholdInputTokens, Comparison: policy.LongContext.Comparison, priceRatesDTO: ratesToDTO(policy.LongContext.Rates)}
|
|
}
|
|
return dto
|
|
}
|
|
|
|
func ratesToDTO(rates pricing.Rates) priceRatesDTO {
|
|
return priceRatesDTO{InputPer1M: formatMicros(rates.InputMicrosPer1M), CacheReadPer1M: formatMicros(rates.CacheReadMicrosPer1M), CacheWritePer1M: formatMicros(rates.CacheWriteMicrosPer1M), OutputPer1M: formatMicros(rates.OutputMicrosPer1M)}
|
|
}
|
|
|
|
func parseDecimalMicros(text string) (int64, error) {
|
|
text = strings.TrimSpace(text)
|
|
if text == "" || strings.HasPrefix(text, "-") || strings.Count(text, ".") > 1 {
|
|
return 0, errors.New("invalid decimal")
|
|
}
|
|
parts := strings.SplitN(text, ".", 2)
|
|
if parts[0] == "" {
|
|
parts[0] = "0"
|
|
}
|
|
fraction := ""
|
|
if len(parts) == 2 {
|
|
fraction = parts[1]
|
|
}
|
|
if len(fraction) > 6 {
|
|
return 0, errors.New("too many decimal places")
|
|
}
|
|
for len(fraction) < 6 {
|
|
fraction += "0"
|
|
}
|
|
whole := new(big.Int)
|
|
if _, ok := whole.SetString(parts[0]+fraction, 10); !ok || !whole.IsInt64() {
|
|
return 0, errors.New("invalid decimal")
|
|
}
|
|
return whole.Int64(), nil
|
|
}
|
|
|
|
func parseRatio(text string) (pricing.Ratio, error) {
|
|
numerator, err := parseDecimalMicros(text)
|
|
if err != nil || numerator <= 0 {
|
|
return pricing.Ratio{}, errors.New("invalid ratio")
|
|
}
|
|
denominator := int64(1_000_000)
|
|
divisor := gcd(numerator, denominator)
|
|
return pricing.Ratio{Numerator: numerator / divisor, Denominator: denominator / divisor}, nil
|
|
}
|
|
|
|
func gcd(a, b int64) int64 {
|
|
for b != 0 {
|
|
a, b = b, a%b
|
|
}
|
|
return a
|
|
}
|
|
|
|
func formatMicros(value int64) string {
|
|
prefix := ""
|
|
if value < 0 {
|
|
prefix = "-"
|
|
value = -value
|
|
}
|
|
whole := value / 1_000_000
|
|
fraction := strconv.FormatInt(value%1_000_000+1_000_000, 10)[1:]
|
|
fraction = strings.TrimRight(fraction, "0")
|
|
if fraction == "" {
|
|
return prefix + strconv.FormatInt(whole, 10)
|
|
}
|
|
return prefix + strconv.FormatInt(whole, 10) + "." + fraction
|
|
}
|
|
|
|
func formatRatio(value pricing.Ratio) string {
|
|
if value.Denominator == 0 {
|
|
return ""
|
|
}
|
|
rational := new(big.Rat).SetFrac(big.NewInt(value.Numerator), big.NewInt(value.Denominator))
|
|
return strings.TrimRight(strings.TrimRight(rational.FloatString(6), "0"), ".")
|
|
}
|
|
|
|
func decodeJSONBody(body []byte, target any) error {
|
|
decoder := json.NewDecoder(strings.NewReader(string(body)))
|
|
decoder.DisallowUnknownFields()
|
|
if err := decoder.Decode(target); err != nil {
|
|
return err
|
|
}
|
|
return nil
|
|
}
|