300 lines
11 KiB
Go
300 lines
11 KiB
Go
package plugin
|
|
|
|
import (
|
|
"context"
|
|
"net/http"
|
|
"net/url"
|
|
"sort"
|
|
"strconv"
|
|
"strings"
|
|
"time"
|
|
|
|
"billing/internal/modelcatalog"
|
|
"billing/internal/pricing"
|
|
"billing/internal/repository"
|
|
)
|
|
|
|
type catalogEntryDTO struct {
|
|
ID string `json:"id"`
|
|
Provider string `json:"provider"`
|
|
ProviderName string `json:"provider_name"`
|
|
Model string `json:"model"`
|
|
ModelName string `json:"model_name"`
|
|
Base priceRatesDTO `json:"base"`
|
|
LongContext *longContextDTO `json:"long_context,omitempty"`
|
|
}
|
|
|
|
type catalogChangeDTO struct {
|
|
Model string `json:"model"`
|
|
CatalogID string `json:"catalog_id"`
|
|
Status string `json:"status"`
|
|
Before priceDTO `json:"before"`
|
|
After *priceDTO `json:"after,omitempty"`
|
|
FetchedAt *time.Time `json:"fetched_at,omitempty"`
|
|
}
|
|
|
|
type catalogImportRequest struct {
|
|
Model string `json:"model"`
|
|
CatalogID string `json:"catalog_id"`
|
|
}
|
|
|
|
type catalogApplyRequest struct {
|
|
Revision string `json:"revision"`
|
|
Models []string `json:"models"`
|
|
}
|
|
|
|
func (a *App) searchPriceCatalog(query url.Values) ManagementResponse {
|
|
a.mu.RLock()
|
|
catalog := a.catalog
|
|
a.mu.RUnlock()
|
|
if catalog == nil {
|
|
return catalogManagementError(http.StatusServiceUnavailable, "catalog_unavailable", "价格目录尚未初始化")
|
|
}
|
|
limit := 20
|
|
if raw := strings.TrimSpace(query.Get("limit")); raw != "" {
|
|
parsed, err := strconv.Atoi(raw)
|
|
if err != nil || parsed < 1 || parsed > 50 {
|
|
return catalogManagementError(http.StatusBadRequest, "invalid_limit", "limit 必须是 1-50 的整数")
|
|
}
|
|
limit = parsed
|
|
}
|
|
entries, info, loaded := catalog.Search(query.Get("q"), limit)
|
|
_, lastError, _ := catalog.Status()
|
|
results := make([]catalogEntryDTO, 0, len(entries))
|
|
for _, entry := range entries {
|
|
results = append(results, catalogEntryToDTO(entry))
|
|
}
|
|
return jsonManagementResponse(http.StatusOK, map[string]any{
|
|
"loaded": loaded, "catalog": info, "models": results, "last_error": lastError,
|
|
})
|
|
}
|
|
|
|
func (a *App) refreshPriceCatalog() ManagementResponse {
|
|
a.mu.RLock()
|
|
catalog := a.catalog
|
|
a.mu.RUnlock()
|
|
if catalog == nil {
|
|
return catalogManagementError(http.StatusServiceUnavailable, "catalog_unavailable", "价格目录尚未初始化")
|
|
}
|
|
ctx, cancel := context.WithTimeout(context.Background(), 35*time.Second)
|
|
defer cancel()
|
|
info, err := catalog.Refresh(ctx)
|
|
if err != nil {
|
|
return catalogManagementError(http.StatusBadGateway, "catalog_download_failed", err.Error())
|
|
}
|
|
a.mu.RLock()
|
|
defer a.mu.RUnlock()
|
|
if a.catalog != catalog || a.store == nil {
|
|
return catalogManagementError(http.StatusServiceUnavailable, "catalog_reconfigured", "价格目录配置已经变化，请重新刷新")
|
|
}
|
|
changes, err := previewCatalogChanges(context.Background(), a.store, catalog)
|
|
if err != nil {
|
|
return catalogManagementError(http.StatusInternalServerError, "database_error", err.Error())
|
|
}
|
|
changed, missing := catalogChangeCounts(changes)
|
|
return jsonManagementResponse(http.StatusOK, map[string]any{
|
|
"catalog": info, "changes": changes, "changed": changed, "missing": missing,
|
|
})
|
|
}
|
|
|
|
func (a *App) importCatalogPrice(body []byte) ManagementResponse {
|
|
var request catalogImportRequest
|
|
if err := decodeJSONBody(body, &request); err != nil {
|
|
return catalogManagementError(http.StatusBadRequest, "invalid_catalog_import", err.Error())
|
|
}
|
|
request.Model = strings.TrimSpace(request.Model)
|
|
request.CatalogID = strings.TrimSpace(request.CatalogID)
|
|
if request.Model == "" || request.CatalogID == "" {
|
|
return catalogManagementError(http.StatusBadRequest, "invalid_catalog_import", "model 和 catalog_id 不能为空")
|
|
}
|
|
a.priceMu.Lock()
|
|
defer a.priceMu.Unlock()
|
|
a.mu.RLock()
|
|
catalog := a.catalog
|
|
current, currentExists := a.prices[normalizeModelName(request.Model)]
|
|
a.mu.RUnlock()
|
|
if catalog == nil {
|
|
return catalogManagementError(http.StatusServiceUnavailable, "catalog_unavailable", "价格目录尚未初始化")
|
|
}
|
|
entry, info, found := catalog.Lookup(request.CatalogID)
|
|
if !found {
|
|
return catalogManagementError(http.StatusNotFound, "catalog_price_not_found", "参考价格不存在或当前目录尚未下载")
|
|
}
|
|
fastEnabled := false
|
|
multiplier := pricing.Ratio{Numerator: 5, Denominator: 2}
|
|
if currentExists {
|
|
fastEnabled = current.FastPricingEnabled
|
|
multiplier = current.FastMultiplier
|
|
}
|
|
policy := entry.Policy(request.Model, fastEnabled, multiplier)
|
|
if err := policy.Validate(); err != nil {
|
|
return catalogManagementError(http.StatusBadRequest, "invalid_catalog_price", err.Error())
|
|
}
|
|
record := repository.PriceRecord{
|
|
Policy: policy,
|
|
Source: repository.PriceSource{Kind: repository.PriceSourceModelsDev, CatalogID: entry.ID, Revision: info.Revision, FetchedAt: info.FetchedAt},
|
|
}
|
|
return a.persistPriceRecordLocked(record, catalog)
|
|
}
|
|
|
|
func (a *App) applyCatalogChanges(body []byte) ManagementResponse {
|
|
var request catalogApplyRequest
|
|
if err := decodeJSONBody(body, &request); err != nil {
|
|
return catalogManagementError(http.StatusBadRequest, "invalid_catalog_apply", err.Error())
|
|
}
|
|
request.Revision = strings.TrimSpace(request.Revision)
|
|
requested := make(map[string]struct{}, len(request.Models))
|
|
for _, model := range request.Models {
|
|
model = normalizeModelName(model)
|
|
if model != "" {
|
|
requested[model] = struct{}{}
|
|
}
|
|
}
|
|
if request.Revision == "" || len(requested) == 0 {
|
|
return catalogManagementError(http.StatusBadRequest, "invalid_catalog_apply", "revision 和 models 不能为空")
|
|
}
|
|
|
|
a.priceMu.Lock()
|
|
defer a.priceMu.Unlock()
|
|
a.mu.Lock()
|
|
catalog, store := a.catalog, a.store
|
|
if catalog == nil || store == nil {
|
|
a.mu.Unlock()
|
|
return catalogManagementError(http.StatusServiceUnavailable, "catalog_unavailable", "价格目录或数据库尚未初始化")
|
|
}
|
|
info, _, loaded := catalog.Status()
|
|
if !loaded || info.Revision != request.Revision {
|
|
a.mu.Unlock()
|
|
return catalogManagementError(http.StatusConflict, "catalog_changed", "价格目录已经变化，请重新查看差异")
|
|
}
|
|
records, err := store.ListPriceRecords(context.Background())
|
|
if err != nil {
|
|
a.mu.Unlock()
|
|
return catalogManagementError(http.StatusInternalServerError, "database_error", err.Error())
|
|
}
|
|
updates := make([]repository.PriceRecord, 0, len(requested))
|
|
seen := make(map[string]struct{}, len(requested))
|
|
for _, record := range records {
|
|
normalized := normalizeModelName(record.Policy.Model)
|
|
if _, wanted := requested[normalized]; !wanted {
|
|
continue
|
|
}
|
|
if record.Source.Kind != repository.PriceSourceModelsDev {
|
|
a.mu.Unlock()
|
|
return catalogManagementError(http.StatusConflict, "price_source_changed", record.Policy.Model+" 已经改为手动价格")
|
|
}
|
|
entry, _, found := catalog.Lookup(record.Source.CatalogID)
|
|
if !found {
|
|
a.mu.Unlock()
|
|
return catalogManagementError(http.StatusConflict, "catalog_price_missing", record.Policy.Model+" 的参考价格已不存在")
|
|
}
|
|
policy := entry.Policy(record.Policy.Model, record.Policy.FastPricingEnabled, record.Policy.FastMultiplier)
|
|
updates = append(updates, repository.PriceRecord{
|
|
Policy: policy,
|
|
Source: repository.PriceSource{Kind: repository.PriceSourceModelsDev, CatalogID: entry.ID, Revision: info.Revision, FetchedAt: info.FetchedAt},
|
|
})
|
|
seen[normalized] = struct{}{}
|
|
}
|
|
if len(seen) != len(requested) {
|
|
a.mu.Unlock()
|
|
return catalogManagementError(http.StatusConflict, "price_source_changed", "部分待更新价格已经被删除或改变")
|
|
}
|
|
if err := store.UpsertPriceRecords(context.Background(), updates); err != nil {
|
|
a.mu.Unlock()
|
|
return catalogManagementError(http.StatusInternalServerError, "database_error", err.Error())
|
|
}
|
|
for _, update := range updates {
|
|
a.prices[normalizeModelName(update.Policy.Model)] = update.Policy
|
|
}
|
|
a.mu.Unlock()
|
|
|
|
for _, update := range updates {
|
|
if _, err := store.BackfillMissingCosts(context.Background(), update.Policy); err != nil {
|
|
a.mu.RLock()
|
|
reconfigured := a.store != store
|
|
a.mu.RUnlock()
|
|
if reconfigured {
|
|
return catalogManagementError(http.StatusServiceUnavailable, "database_reconfigured", "价格已经更新，重新配置中止了历史补算")
|
|
}
|
|
return catalogManagementError(http.StatusInternalServerError, "database_error", err.Error())
|
|
}
|
|
}
|
|
models := make([]string, 0, len(updates))
|
|
for _, update := range updates {
|
|
models = append(models, update.Policy.Model)
|
|
}
|
|
sort.Strings(models)
|
|
return jsonManagementResponse(http.StatusOK, map[string]any{"updated": len(updates), "models": models, "catalog": info})
|
|
}
|
|
|
|
func previewCatalogChanges(ctx context.Context, store *repository.SQLiteUsageRepository, catalog *modelcatalog.Manager) ([]catalogChangeDTO, error) {
|
|
records, err := store.ListPriceRecords(ctx)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
changes := make([]catalogChangeDTO, 0)
|
|
for _, record := range records {
|
|
if record.Source.Kind != repository.PriceSourceModelsDev {
|
|
continue
|
|
}
|
|
entry, info, found := catalog.Lookup(record.Source.CatalogID)
|
|
before := priceRecordToDTO(record)
|
|
if !found {
|
|
changes = append(changes, catalogChangeDTO{Model: record.Policy.Model, CatalogID: record.Source.CatalogID, Status: "missing", Before: before})
|
|
continue
|
|
}
|
|
policy := entry.Policy(record.Policy.Model, record.Policy.FastPricingEnabled, record.Policy.FastMultiplier)
|
|
if sameCatalogPolicy(record.Policy, policy) {
|
|
continue
|
|
}
|
|
afterRecord := repository.PriceRecord{
|
|
Policy: policy,
|
|
Source: repository.PriceSource{Kind: repository.PriceSourceModelsDev, CatalogID: entry.ID, Revision: info.Revision, FetchedAt: info.FetchedAt},
|
|
}
|
|
after := priceRecordToDTO(afterRecord)
|
|
fetchedAt := info.FetchedAt
|
|
changes = append(changes, catalogChangeDTO{Model: record.Policy.Model, CatalogID: entry.ID, Status: "changed", Before: before, After: &after, FetchedAt: &fetchedAt})
|
|
}
|
|
sort.Slice(changes, func(i, j int) bool { return changes[i].Model < changes[j].Model })
|
|
return changes, nil
|
|
}
|
|
|
|
func catalogChangeCounts(changes []catalogChangeDTO) (changed, missing int) {
|
|
for _, change := range changes {
|
|
if change.Status == "changed" {
|
|
changed++
|
|
} else if change.Status == "missing" {
|
|
missing++
|
|
}
|
|
}
|
|
return changed, missing
|
|
}
|
|
|
|
func sameCatalogPolicy(left, right pricing.Policy) bool {
|
|
if left.Model != right.Model || left.Base != right.Base || left.FastPricingEnabled != right.FastPricingEnabled || left.FastMultiplier != right.FastMultiplier {
|
|
return false
|
|
}
|
|
if left.LongContext == nil || right.LongContext == nil {
|
|
return left.LongContext == nil && right.LongContext == nil
|
|
}
|
|
return *left.LongContext == *right.LongContext
|
|
}
|
|
|
|
func catalogEntryToDTO(entry modelcatalog.Entry) catalogEntryDTO {
|
|
dto := catalogEntryDTO{
|
|
ID: entry.ID, Provider: entry.Provider, ProviderName: entry.ProviderName,
|
|
Model: entry.Model, ModelName: entry.ModelName, Base: ratesToDTO(entry.Base),
|
|
}
|
|
if entry.LongContext != nil {
|
|
dto.LongContext = &longContextDTO{
|
|
ThresholdInputTokens: entry.LongContext.ThresholdInputTokens,
|
|
Comparison: entry.LongContext.Comparison, priceRatesDTO: ratesToDTO(entry.LongContext.Rates),
|
|
}
|
|
}
|
|
return dto
|
|
}
|
|
|
|
func catalogManagementError(status int, code, message string) ManagementResponse {
|
|
return jsonManagementResponse(status, map[string]any{"error": map[string]string{"code": code, "message": message}})
|
|
}
|