Sub2API v1.0 - AI API 网关（二开初始版本，基于上游 Wei-Shaw/sub2api）

backend/internal/service/grok_search_count.go

@@ -0,0 +1,233 @@
+package service
+
+import (
+	"strconv"
+	"strings"
+
+	"github.com/tidwall/gjson"
+)
+
+// countGrokNativeSearchCallsFromJSONBytes counts completed native search tool
+// calls in a Responses-style JSON body (output array or nested response.output).
+// Counts: web_search_call, x_search_call, tool_search_call, and function_call
+// named tool_search / web_search / x_search.
+func countGrokNativeSearchCallsFromJSONBytes(body []byte) int {
+	if len(body) == 0 || !gjson.ValidBytes(body) {
+		return 0
+	}
+	// Responses envelopes normally expose either top-level output (JSON mode)
+	// or response.output (terminal SSE payload). Compatibility layers can retain
+	// both copies; counting both would bill the same search twice. Prefer the
+	// canonical nested response when present and fall back to top-level output.
+	if nested := gjson.GetBytes(body, "response.output"); nested.IsArray() {
+		return countGrokNativeSearchCallsInOutputArray(nested)
+	}
+	return countGrokNativeSearchCallsInOutputArray(gjson.GetBytes(body, "output"))
+}
+
+func countGrokNativeSearchCallsFromSSEBody(body string) int {
+	if strings.TrimSpace(body) == "" {
+		return 0
+	}
+	seen := make(map[string]struct{})
+	total := 0
+	forEachOpenAISSEDataPayload(body, func(data []byte) {
+		total += countGrokNativeSearchCallsInSSEDataDedup(data, seen)
+	})
+	return total
+}
+
+// countGrokNativeSearchCallsInSSEData counts search tool calls in one SSE
+// payload without cross-event dedup. Prefer countGrokNativeSearchCallsInSSEDataDedup
+// for live streams so item.done + response.completed do not double-bill.
+func countGrokNativeSearchCallsInSSEData(data []byte) int {
+	n, _ := countGrokNativeSearchCallsInSSEDataWithKeys(data)
+	return n
+}
+
+// countGrokNativeSearchCallsInSSEDataDedup increments only unseen call ids.
+// Callers must reuse the same seen map for the full stream lifetime.
+//
+// When call_id/id is missing, a synthetic key is built from item type + name so
+// item.done + response.completed for the same tool still count once (never fall
+// back to raw multi-event n, which ~2× overbills).
+func countGrokNativeSearchCallsInSSEDataDedup(data []byte, seen map[string]struct{}) int {
+	if seen == nil {
+		return countGrokNativeSearchCallsInSSEData(data)
+	}
+	n, keys := countGrokNativeSearchCallsInSSEDataWithKeys(data)
+	if n <= 0 {
+		return 0
+	}
+	// Prefer stable ids; fill gaps with synthetic keys so we never raw-add n.
+	if len(keys) < n {
+		// Rebuild keys for every item so unkeyed items still get a fingerprint.
+		keys = collectGrokNativeSearchCallKeys(data)
+	}
+	if len(keys) == 0 {
+		// True empty — should not happen when n>0; fail-closed to 0 extra bill.
+		return 0
+	}
+	added := 0
+	local := make(map[string]struct{}, len(keys))
+	isItemDone := strings.TrimSpace(gjson.GetBytes(data, "type").String()) == "response.output_item.done"
+	for _, k := range keys {
+		if k == "" {
+			continue
+		}
+		if _, ok := local[k]; ok {
+			continue
+		}
+		local[k] = struct{}{}
+		if _, ok := seen[k]; ok {
+			if !isItemDone || !strings.HasPrefix(k, "synth:") {
+				continue
+			}
+			// Each id-less item.done is a distinct completed invocation. Advance
+			// its ordinal so interrupted streams remain accurately billable.
+			separator := strings.LastIndexByte(k, ':')
+			if separator < 0 {
+				continue
+			}
+			base := k[:separator]
+			for ordinal := 2; ; ordinal++ {
+				candidate := base + ":" + strconv.Itoa(ordinal)
+				if _, exists := seen[candidate]; !exists {
+					k = candidate
+					break
+				}
+			}
+		}
+		seen[k] = struct{}{}
+		added++
+	}
+	return added
+}
+
+func collectGrokNativeSearchCallKeys(data []byte) []string {
+	if len(data) == 0 || !gjson.ValidBytes(data) {
+		return nil
+	}
+	// An empty type means a bare item object without an SSE envelope; anything
+	// else that is not a completion event carries no billable call.
+	switch strings.TrimSpace(gjson.GetBytes(data, "type").String()) {
+	case "response.output_item.done", "response.completed", "response.done", "":
+	default:
+		return nil
+	}
+	var keys []string
+	syntheticOrdinals := make(map[string]int)
+	consider := func(item gjson.Result) {
+		if !isGrokNativeSearchOutputItem(item) {
+			return
+		}
+		key := firstNonEmpty(
+			strings.TrimSpace(item.Get("call_id").String()),
+			strings.TrimSpace(item.Get("id").String()),
+			strings.TrimSpace(item.Get("item.call_id").String()),
+			strings.TrimSpace(item.Get("item.id").String()),
+		)
+		if key == "" {
+			// Include the ordinal among same-kind calls. A plain type:name key
+			// collapses two id-less web searches in one completed response into
+			// one charge. The ordinal remains stable between ordered item.done
+			// events and response.completed output.
+			base := "synth:" + strings.ToLower(strings.TrimSpace(item.Get("type").String())) +
+				":" + strings.ToLower(strings.TrimSpace(item.Get("name").String()))
+			syntheticOrdinals[base]++
+			key = base + ":" + strconv.Itoa(syntheticOrdinals[base])
+		}
+		keys = append(keys, key)
+	}
+	if item := gjson.GetBytes(data, "item"); item.Exists() {
+		consider(item)
+	}
+	gjson.GetBytes(data, "response.output").ForEach(func(_, item gjson.Result) bool {
+		consider(item)
+		return true
+	})
+	gjson.GetBytes(data, "output").ForEach(func(_, item gjson.Result) bool {
+		consider(item)
+		return true
+	})
+	if len(keys) == 0 && isGrokNativeSearchOutputItem(gjson.ParseBytes(data)) {
+		consider(gjson.ParseBytes(data))
+	}
+	return keys
+}
+
+func countGrokNativeSearchCallsInSSEDataWithKeys(data []byte) (int, []string) {
+	if len(data) == 0 || !gjson.ValidBytes(data) {
+		return 0, nil
+	}
+	// Count once on item completion / completed response, not on every delta.
+	// An empty type is a bare item object without an SSE envelope.
+	switch strings.TrimSpace(gjson.GetBytes(data, "type").String()) {
+	case "response.output_item.done", "response.completed", "response.done", "":
+	default:
+		return 0, nil
+	}
+	var keys []string
+	n := 0
+	consider := func(item gjson.Result) {
+		if !isGrokNativeSearchOutputItem(item) {
+			return
+		}
+		n++
+		key := firstNonEmpty(
+			strings.TrimSpace(item.Get("call_id").String()),
+			strings.TrimSpace(item.Get("id").String()),
+			strings.TrimSpace(item.Get("item.call_id").String()),
+			strings.TrimSpace(item.Get("item.id").String()),
+		)
+		if key != "" {
+			keys = append(keys, key)
+		}
+	}
+	if item := gjson.GetBytes(data, "item"); item.Exists() {
+		consider(item)
+	}
+	gjson.GetBytes(data, "response.output").ForEach(func(_, item gjson.Result) bool {
+		consider(item)
+		return true
+	})
+	gjson.GetBytes(data, "output").ForEach(func(_, item gjson.Result) bool {
+		consider(item)
+		return true
+	})
+	// Bare output item event without nested item key.
+	if n == 0 && isGrokNativeSearchOutputItem(gjson.ParseBytes(data)) {
+		consider(gjson.ParseBytes(data))
+	}
+	return n, keys
+}
+
+func countGrokNativeSearchCallsInOutputArray(output gjson.Result) int {
+	if !output.IsArray() {
+		return 0
+	}
+	count := 0
+	output.ForEach(func(_, item gjson.Result) bool {
+		if isGrokNativeSearchOutputItem(item) {
+			count++
+		}
+		return true
+	})
+	return count
+}
+
+func isGrokNativeSearchOutputItem(item gjson.Result) bool {
+	if !item.Exists() {
+		return false
+	}
+	itemType := strings.ToLower(strings.TrimSpace(item.Get("type").String()))
+	switch itemType {
+	case "web_search_call", "x_search_call", "tool_search_call":
+		return true
+	case "function_call", "custom_tool_call":
+		name := strings.ToLower(strings.TrimSpace(item.Get("name").String()))
+		return name == "web_search" || name == "x_search" || name == "tool_search"
+	default:
+		return false
+	}
+}