new-api/setting/operation_setting/tools.go

package operation_setting

import (
	"encoding/json"
	"math"
	"sort"
	"strings"
	"sync/atomic"

	"github.com/QuantumNous/new-api/common"
	"github.com/QuantumNous/new-api/setting/config"
)

// ---------------------------------------------------------------------------
// Tool call prices ($/1K calls, admin-configurable)
// DB key: tool_price_setting.prices
//
// Key format:
//   - "tool_name"              → default price for all models
//   - "tool_name:model_prefix*" → override for models matching the prefix
//
// Effective index: hardcoded defaults → hardcoded model overrides → valid
// operator values. Lookup uses the longest model prefix before the tool
// default, and a matched numeric zero is terminal.
//
// Built-in keys are written in the vendor's list currency.
// ---------------------------------------------------------------------------

const ToolPriceOptionKey = "tool_price_setting.prices"

const (
	defaultWebSearchToolPrice        = 10.0
	defaultWebSearchPreviewToolPrice = 10.0
	defaultFileSearchToolPrice       = 2.5
	defaultGoogleSearchToolPrice     = 14.0
	defaultImageGenerationToolPrice  = 150.0
	defaultSearchPreviewModelPrice   = 25.0
)

// seedHardcodedToolPrices injects compile-time built-in fallbacks (tool
// defaults and model-prefix overrides) into the destination. The source is
// constants, not a mutable package map or operator configuration.
func seedHardcodedToolPrices(prices map[string]float64) {
	prices["web_search"] = defaultWebSearchToolPrice
	prices["web_search_preview"] = defaultWebSearchPreviewToolPrice
	prices["file_search"] = defaultFileSearchToolPrice
	prices["google_search"] = defaultGoogleSearchToolPrice
	prices["image_generation"] = defaultImageGenerationToolPrice
	prices["web_search_preview:gpt-4o*"] = defaultSearchPreviewModelPrice
	prices["web_search_preview:gpt-4.1*"] = defaultSearchPreviewModelPrice
	prices["web_search_preview:gpt-4o-mini*"] = defaultSearchPreviewModelPrice
	prices["web_search_preview:gpt-4.1-mini*"] = defaultSearchPreviewModelPrice

	// Google Search grounding (USD per 1K, ai.google.dev/gemini-api/docs/pricing
	// and cloud.google.com/vertex-ai/generative-ai/pricing): Gemini 3 bills each
	// search query at $14 (the google_search default); Gemini 2.5 and older bill
	// $35 per grounded prompt, so their query count is free and the grounded
	// prompt carries the price.
	prices["google_search:gemini-2.5*"] = 0
	prices["google_search:gemini-2.0*"] = 0
	prices["google_search:gemini-1.5*"] = 0
	prices["google_search_grounded_prompt:gemini-2.5*"] = 35
	prices["google_search_grounded_prompt:gemini-2.0*"] = 35
	prices["google_search_grounded_prompt:gemini-1.5*"] = 35

	// Vendor search tiers, verbatim in the vendor's list currency per 1K
	// calls. CNY: Zhipu search engines, 0.01/0.03/0.05 CNY per call
	// (https://docs.bigmodel.cn/cn/guide/tools/web-search), and Alibaba Model
	// Studio search strategies at Beijing prices
	// (https://help.aliyun.com/zh/model-studio/web-search); the Singapore
	// agent price (73.392381 CNY) is set by the operator. USD: Azure
	// Responses web search, billed as Bing transactions at $14 per 1,000
	// (https://www.microsoft.com/en-us/bing/apis), and xAI web search $5 per
	// 1K calls, x_search $5 per 1K posts and $10 per 1K profiles
	// (https://docs.x.ai/developers/pricing).
	prices["search_std"] = 10
	prices["search_pro"] = 30
	prices["search_pro_sogou"] = 50
	prices["search_pro_quark"] = 50
	prices["search_strategy_turbo"] = 3
	prices["search_strategy_max"] = 4
	prices["search_strategy_agent"] = 4
	prices["search_strategy_agent_max"] = 4
	prices["bing_web_search"] = 14
	prices["web_search:grok*"] = 5
	prices["x_search_posts"] = 5
	prices["x_search_profiles"] = 10

	// DeepSeek Search and Fetch APIs, $0.75 per 1,000 successful requests
	// each (https://api-docs.deepseek.com/quick_start/pricing), for the
	// searches and page fetches DeepSeek runs server-side for Claude Code and
	// Codex traffic. /v1/tools/search and /v1/tools/fetch bill through the
	// deepseek-search and deepseek-fetch model prices instead.
	prices["deepseek_search"] = 0.75
	prices["deepseek_fetch"] = 0.75
}

// ToolPriceSetting is managed by config.GlobalConfig.Register.
// Prices holds operator overrides only; hardcoded fallbacks live in the index.
type ToolPriceSetting struct {
	Prices map[string]float64 `json:"prices"`
}

var toolPriceSetting = ToolPriceSetting{
	Prices: make(map[string]float64),
}

// builtInToolNames are the tool names (the part before ":") that
// seedHardcodedToolPrices prices, built once from the constant seed and never
// from operator configuration.
var builtInToolNames = make(map[string]struct{})

func init() {
	config.GlobalConfig.Register("tool_price_setting", &toolPriceSetting)
	seed := make(map[string]float64)
	seedHardcodedToolPrices(seed)
	for key := range seed {
		name, _, _ := strings.Cut(key, ":")
		builtInToolNames[name] = struct{}{}
	}
	RebuildToolPriceIndex()
}

// IsBuiltInToolPriceKey reports whether the built-in price seed prices the
// tool name, as a default or for a model prefix. Operator prices never count.
func IsBuiltInToolPriceKey(name string) bool {
	_, ok := builtInToolNames[name]
	return ok
}

// ---------------------------------------------------------------------------
// Precomputed price index (atomic, lock-free on read path)
// ---------------------------------------------------------------------------

type prefixEntry struct {
	prefix string
	price  float64
}

type toolPriceIndex struct {
	defaults map[string]float64
	prefixes map[string][]prefixEntry
}

var currentIndex atomic.Pointer[toolPriceIndex]

func isValidToolPrice(price float64) bool {
	return price >= 0 && !math.IsNaN(price) && !math.IsInf(price, 0)
}

func decodeToolPricesJSON(value string, ignoreInvalidEntries bool) (map[string]float64, error) {
	rawValue := json.RawMessage(strings.TrimSpace(value))
	if common.GetJsonType(rawValue) != "object" {
		return nil, common.NewMessage("Tool prices must be a JSON object")
	}

	var rawPrices map[string]json.RawMessage
	if err := common.Unmarshal(rawValue, &rawPrices); err != nil {
		return nil, common.NewMessage("Failed to parse tool prices: {{error}}", map[string]any{"error": err.Error()})
	}

	prices := make(map[string]float64, len(rawPrices))
	for name, rawPrice := range rawPrices {
		var entryErr error
		if common.GetJsonType(rawPrice) != "number" {
			entryErr = common.NewMessage(`Tool price "{{name}}" must be a non-negative number`, map[string]any{"name": name})
		} else {
			var price float64
			if err := common.Unmarshal(rawPrice, &price); err != nil {
				entryErr = common.NewMessage(`Failed to parse tool price "{{name}}": {{error}}`, map[string]any{"name": name, "error": err.Error()})
			} else if !isValidToolPrice(price) {
				entryErr = common.NewMessage(`Tool price "{{name}}" must be a finite non-negative number`, map[string]any{"name": name})
			} else {
				prices[name] = price
			}
		}

		if entryErr == nil {
			continue
		}
		if !ignoreInvalidEntries {
			return nil, entryErr
		}
		common.SysError(entryErr.Error())
	}
	return prices, nil
}

// ValidateToolPricesJSON validates an operator-supplied complete price map.
// A numeric zero is valid and intentionally disables the matching rule.
func ValidateToolPricesJSON(value string) error {
	_, err := decodeToolPricesJSON(value, false)
	return err
}

// LoadToolPricesFromJSONString replaces the complete operator price map.
// Invalid legacy entries are ignored individually so valid sibling overrides
// survive, while missing built-in keys continue to use hardcoded fallbacks.
func LoadToolPricesFromJSONString(value string) {
	prices, err := decodeToolPricesJSON(value, true)
	if err != nil {
		common.SysError(common.LogText("failed to load tool prices, using hardcoded fallbacks: %s", err.Error()))
		prices = make(map[string]float64)
	}
	toolPriceSetting.Prices = prices
	RebuildToolPriceIndex()
}

// RebuildToolPriceIndex rebuilds the lookup index from the current config.
// Called on init and after config updates. Not on the billing hot path.
func RebuildToolPriceIndex() {
	merged := make(map[string]float64, 9+len(toolPriceSetting.Prices))
	seedHardcodedToolPrices(merged)
	for k, v := range toolPriceSetting.Prices {
		if !isValidToolPrice(v) {
			continue
		}
		merged[k] = v
	}

	idx := &toolPriceIndex{
		defaults: make(map[string]float64),
		prefixes: make(map[string][]prefixEntry),
	}

	for key, price := range merged {
		before, after, ok := strings.Cut(key, ":")
		if !ok {
			idx.defaults[key] = price
			continue
		}
		toolName := before
		modelPart := after
		prefix := strings.TrimSuffix(modelPart, "*")
		idx.prefixes[toolName] = append(idx.prefixes[toolName], prefixEntry{prefix: prefix, price: price})
	}

	for tool := range idx.prefixes {
		entries := idx.prefixes[tool]
		sort.Slice(entries, func(i, j int) bool {
			if len(entries[i].prefix) == len(entries[j].prefix) {
				return entries[i].prefix < entries[j].prefix
			}
			return len(entries[i].prefix) > len(entries[j].prefix)
		})
		idx.prefixes[tool] = entries
	}

	currentIndex.Store(idx)
}

// GetToolPriceForModel returns the price ($/1K calls) for a tool given a model name.
// Lookup: longest prefix match → tool default → 0.
func GetToolPriceForModel(toolName, modelName string) float64 {
	idx := currentIndex.Load()
	if idx == nil {
		RebuildToolPriceIndex()
		idx = currentIndex.Load()
		if idx == nil {
			return 0
		}
	}

	if entries, ok := idx.prefixes[toolName]; ok && modelName != "" {
		for _, e := range entries {
			if strings.HasPrefix(modelName, e.prefix) {
				return e.price
			}
		}
	}

	if p, ok := idx.defaults[toolName]; ok {
		return p
	}
	return 0
}

// GetToolPrice is a convenience wrapper when no model name is needed.
func GetToolPrice(toolName string) float64 {
	return GetToolPriceForModel(toolName, "")
}

// SetToolPriceForTest injects a tool price and rebuilds the lookup index. Tests only.
func SetToolPriceForTest(name string, price float64) {
	if toolPriceSetting.Prices == nil {
		toolPriceSetting.Prices = make(map[string]float64)
	}
	toolPriceSetting.Prices[name] = price
	RebuildToolPriceIndex()
}

// DeleteToolPriceForTest removes an injected tool price and rebuilds the index. Tests only.
func DeleteToolPriceForTest(name string) {
	delete(toolPriceSetting.Prices, name)
	RebuildToolPriceIndex()
}

// ---------------------------------------------------------------------------
// Gemini audio input pricing (per-million tokens, model-specific)
// ---------------------------------------------------------------------------

const (
	Gemini25FlashPreviewInputAudioPrice     = 1.00
	Gemini25FlashProductionInputAudioPrice  = 1.00
	Gemini25FlashLitePreviewInputAudioPrice = 0.50
	Gemini25FlashNativeAudioInputAudioPrice = 3.00
	Gemini20FlashInputAudioPrice            = 0.70
	GeminiRoboticsER15InputAudioPrice       = 1.00
)

func GetGeminiInputAudioPricePerMillionTokens(modelName string) float64 {
	if strings.HasPrefix(modelName, "gemini-2.5-flash-preview-native-audio") {
		return Gemini25FlashNativeAudioInputAudioPrice
	}
	if strings.HasPrefix(modelName, "gemini-2.5-flash-preview-lite") {
		return Gemini25FlashLitePreviewInputAudioPrice
	}
	if strings.HasPrefix(modelName, "gemini-2.5-flash-preview") {
		return Gemini25FlashPreviewInputAudioPrice
	}
	if strings.HasPrefix(modelName, "gemini-2.5-flash") {
		return Gemini25FlashProductionInputAudioPrice
	}
	if strings.HasPrefix(modelName, "gemini-2.0-flash") {
		return Gemini20FlashInputAudioPrice
	}
	if strings.HasPrefix(modelName, "gemini-robotics-er-1.5") {
		return GeminiRoboticsER15InputAudioPrice
	}
	return 0
}